tocmo0nlord/axolotl
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

fix

Browse Source
This commit is contained in:
Dan Saunders
2025-09-23 18:13:53 -04:00
parent 4db7a21ff7
commit d578c53603
1 changed files with 128 additions and 120 deletions
Download Patch File Download Diff File Expand all files Collapse all files

38
src/axolotl/kernels/moe/tt_mg_gemm/mg_grouped_gemm.py
View File

@@ -81,14 +81,14 @@ def _kernel_mg_forward_hopper(
m_size = tl.load(m_sizes + g) m_size = tl.load(m_sizes + g)
M_end = M_start + m_size M_end = M_start + m_size
if m_size <= 0: if m_size > 0:
continue
num_m_tiles = tl.cdiv(m_size, BLOCK_SIZE_M) num_m_tiles = tl.cdiv(m_size, BLOCK_SIZE_M)
num_n_tiles = tl.cdiv(n_size, BLOCK_SIZE_N) num_n_tiles = tl.cdiv(n_size, BLOCK_SIZE_N)
group_num_tiles = num_m_tiles * num_n_tiles group_num_tiles = num_m_tiles * num_n_tiles
while tbidx >= processed_tiles and tbidx < processed_tiles + group_num_tiles: while (
tbidx >= processed_tiles and tbidx < processed_tiles + group_num_tiles
):
group_index = tbidx - processed_tiles group_index = tbidx - processed_tiles
tile_m_index = group_index % num_m_tiles tile_m_index = group_index % num_m_tiles
@@ -127,7 +127,9 @@ def _kernel_mg_forward_hopper(
row_store_mask = local_row_offsets < m_size row_store_mask = local_row_offsets < m_size
global_row = (M_start + local_row_offsets).to(tl.int32) global_row = (M_start + local_row_offsets).to(tl.int32)
local_col_offsets = tile_n_index * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) local_col_offsets = tile_n_index * BLOCK_SIZE_N + tl.arange(
0, BLOCK_SIZE_N
)
col_store_mask = local_col_offsets < n_size col_store_mask = local_col_offsets < n_size
store_mask = row_store_mask[:, None] & col_store_mask[None, :] store_mask = row_store_mask[:, None] & col_store_mask[None, :]
@@ -155,7 +157,11 @@ def _kernel_mg_forward_hopper(
mask=row_store_mask[:, None] & col_mask1[None, :], mask=row_store_mask[:, None] & col_mask1[None, :],
) )
else: else:
ptr = c_ptr + global_row[:, None] * n_size + local_col_offsets[None, :] ptr = (
c_ptr
+ global_row[:, None] * n_size
+ local_col_offsets[None, :]
)
tl.store(ptr, accumulator.to(c_dtype), mask=store_mask) tl.store(ptr, accumulator.to(c_dtype), mask=store_mask)
tbidx += NUM_SMS tbidx += NUM_SMS
@@ -220,14 +226,14 @@ def _kernel_mg_dx_tma(
m_size = tl.load(m_sizes + g) m_size = tl.load(m_sizes + g)
M_end = M_start + m_size M_end = M_start + m_size
if m_size <= 0: if m_size > 0:
continue
num_m_tiles = tl.cdiv(m_size, BLOCK_SIZE_M) num_m_tiles = tl.cdiv(m_size, BLOCK_SIZE_M)
num_k_tiles = tl.cdiv(K, BLOCK_SIZE_K) num_k_tiles = tl.cdiv(K, BLOCK_SIZE_K)
group_num_tiles = num_m_tiles * num_k_tiles group_num_tiles = num_m_tiles * num_k_tiles
while tbidx >= processed_tiles and tbidx < processed_tiles + group_num_tiles: while (
tbidx >= processed_tiles and tbidx < processed_tiles + group_num_tiles
):
group_index = tbidx - processed_tiles group_index = tbidx - processed_tiles
tile_m_index = group_index % num_m_tiles tile_m_index = group_index % num_m_tiles
@@ -259,7 +265,9 @@ def _kernel_mg_dx_tma(
accumulator += tl.dot(grad_y, w_tile) accumulator += tl.dot(grad_y, w_tile)
local_row_offsets = tile_m_index * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) local_row_offsets = tile_m_index * BLOCK_SIZE_M + tl.arange(
0, BLOCK_SIZE_M
)
row_store_mask = local_row_offsets < m_size row_store_mask = local_row_offsets < m_size
global_row = (M_start + local_row_offsets).to(tl.int32) global_row = (M_start + local_row_offsets).to(tl.int32)
@@ -342,9 +350,7 @@ def _kernel_mg_dw_tma(
m_size = tl.load(m_sizes + g) m_size = tl.load(m_sizes + g)
M_end = M_start + m_size M_end = M_start + m_size
if m_size <= 0: if m_size > 0:
continue
for m_offset_local in range(0, m_size, BLOCK_SIZE_M): for m_offset_local in range(0, m_size, BLOCK_SIZE_M):
rows_remaining = m_size - m_offset_local rows_remaining = m_size - m_offset_local
rows_remaining = tl.maximum(rows_remaining, 0) rows_remaining = tl.maximum(rows_remaining, 0)
@@ -358,7 +364,9 @@ def _kernel_mg_dw_tma(
grad_block = grad_output_desc.load([m_offset, n_offset]) grad_block = grad_output_desc.load([m_offset, n_offset])
grad_mask = row_mask[:, None] & n_mask[None, :] grad_mask = row_mask[:, None] & n_mask[None, :]
grad_block = tl.where(grad_mask, grad_block, tl.zeros_like(grad_block)) grad_block = tl.where(
grad_mask, grad_block, tl.zeros_like(grad_block)
)
contribution = tl.dot( contribution = tl.dot(
grad_block.to(tl.float32).T, grad_block.to(tl.float32).T,