From ae992da9b064a877a5062836b9a7693a6fc12bcc Mon Sep 17 00:00:00 2001 From: David Sidler Date: Thu, 5 Nov 2026 17:20:29 -0600 Subject: [PATCH] Fix all_gather_matmul_hbm_buffer for partial M groups, small K, and int64 offsets - Partial M groups: when the M-tiles of a fetch stage (or of the problem) are not a multiple of group_size_m, the fetch decode left some K flags unset (the GEMM workgroups then spin forever) and the GEMM group extended into the next stage. Clamp the group size to the stage in both the fetch and the GEMM decode, so any group_size_m works. Previously this hung. - Halve the default k_per_flag until it divides the K-blocks, so small and single-rank problems work instead of hitting an assert. - Use int64 offsets for C and staged_a (all_gather_matmul and hbm_buffer) to avoid int32 overflow on large problems. Adds tests for partial M groups (several group/stage combinations) and for the default k_per_flag on a small problem. Co-Authored-By: Claude --- iris/ops/all_gather_matmul.py | 2 +- iris/ops/all_gather_matmul_hbm_buffer.py | 28 ++++++--- tests/ops/test_all_gather_matmul.py | 78 ++++++++++++++++++++++++ 3 files changed, 100 insertions(+), 8 deletions(-) diff --git a/iris/ops/all_gather_matmul.py b/iris/ops/all_gather_matmul.py index c99e83157..53805f935 100644 --- a/iris/ops/all_gather_matmul.py +++ b/iris/ops/all_gather_matmul.py @@ -154,7 +154,7 @@ def _fused_all_gather_matmul_kernel( # Store result using tritonblas Tile rm, rn = out_tile.indices() - C_ptr = C + rm[:, None] * stride_cm + rn[None, :] * stride_cn + C_ptr = C + rm.to(tl.int64)[:, None] * stride_cm + rn.to(tl.int64)[None, :] * stride_cn mask = (rm[:, None] < M) & (rn[None, :] < N) tl.store(C_ptr, c, mask=mask) diff --git a/iris/ops/all_gather_matmul_hbm_buffer.py b/iris/ops/all_gather_matmul_hbm_buffer.py index 4e18a4ea1..caf3bb09b 100644 --- a/iris/ops/all_gather_matmul_hbm_buffer.py +++ b/iris/ops/all_gather_matmul_hbm_buffer.py @@ -283,9 +283,12 @@ def _hbm_buffer_all_gather_matmul_kernel( for fg_idx in range(stage_pid, total_fg_stage, stage_fetch_sms): m_group = fg_idx // tiles_per_m_group within_group = fg_idx % tiles_per_m_group - k_flag_group = within_group // GROUP_SIZE_M - m_in_group = within_group % GROUP_SIZE_M - m_tile = stage_m_start + m_group * GROUP_SIZE_M + m_in_group + # The last group of a stage can hold fewer than GROUP_SIZE_M M-tiles. + group_m_start = m_group * GROUP_SIZE_M + group_sz = min(stage_m_count - group_m_start, GROUP_SIZE_M) + k_flag_group = within_group // group_sz + m_in_group = within_group % group_sz + m_tile = stage_m_start + group_m_start + m_in_group m_tile = min(m_tile, NUM_M_TILES - 1) k_block_start = k_flag_group * K_PER_FLAG @@ -304,7 +307,9 @@ def _hbm_buffer_all_gather_matmul_kernel( rk = k_block_global * BLOCK_SIZE_K + tl.arange(0, BLOCK_SIZE_K) rk = tl.max_contiguous(tl.multiple_of(rk, BLOCK_SIZE_K), BLOCK_SIZE_K) - staged_ptrs = staged_a + rm.to(tl.int64)[:, None] * stride_sa_m + rk[None, :] * stride_sa_k + staged_ptrs = ( + staged_a + rm.to(tl.int64)[:, None] * stride_sa_m + rk.to(tl.int64)[None, :] * stride_sa_k + ) for compile_rank in range(world_size): if src_rank_idx == compile_rank: @@ -329,7 +334,8 @@ def _hbm_buffer_all_gather_matmul_kernel( group_id = gemm_local_id // num_pid_in_group first_pid_m = stage_m_start + group_id * GROUP_SIZE_M first_pid_m = min(first_pid_m, NUM_M_TILES - 1) - group_sz = min(NUM_M_TILES - first_pid_m, GROUP_SIZE_M) + stage_m_end = min(stage_m_start + M_PER_STAGE, NUM_M_TILES) + group_sz = min(stage_m_end - first_pid_m, GROUP_SIZE_M) pid_m = first_pid_m + ((gemm_local_id % num_pid_in_group) % group_sz) pid_n = (gemm_local_id % num_pid_in_group) // group_sz pid_m = min(pid_m, NUM_M_TILES - 1) @@ -373,7 +379,7 @@ def _hbm_buffer_all_gather_matmul_kernel( rk = k_block * BLOCK_SIZE_K + tl.arange(0, BLOCK_SIZE_K) rk = tl.max_contiguous(tl.multiple_of(rk, BLOCK_SIZE_K), BLOCK_SIZE_K) - a_ptrs = staged_a + rm.to(tl.int64)[:, None] * stride_sa_m + rk[None, :] * stride_sa_k + a_ptrs = staged_a + rm.to(tl.int64)[:, None] * stride_sa_m + rk.to(tl.int64)[None, :] * stride_sa_k a = tl.load(a_ptrs) B_ptrs = B + rk[:, None] * stride_bk + rn[None, :] * stride_bn @@ -389,7 +395,9 @@ def _hbm_buffer_all_gather_matmul_kernel( acc = acc + bias_val[:, None] c = acc.to(C.type.element_ty) - C_ptrs = C + rm[:, None] * stride_cm + rn[None, :] * stride_cn + stride_cm_i64 = tl.cast(stride_cm, tl.int64) + stride_cn_i64 = tl.cast(stride_cn, tl.int64) + C_ptrs = C + rm.to(tl.int64)[:, None] * stride_cm_i64 + rn.to(tl.int64)[None, :] * stride_cn_i64 c_mask = (rm[:, None] < M) & (rn[None, :] < N) tl.store(C_ptrs, c, mask=c_mask, cache_modifier=".wt") @@ -436,6 +444,10 @@ def all_gather_matmul_hbm_buffer_preamble( num_m_tiles = M // config.block_size_m num_k_blocks = K // config.block_size_k + # Keep the default valid for small and single-rank problems where fewer + # than eight K-blocks are available. + while k_per_flag > 1 and num_k_blocks % k_per_flag != 0: + k_per_flag //= 2 assert num_k_blocks % k_per_flag == 0 num_flag_groups_k = num_k_blocks // k_per_flag @@ -587,6 +599,8 @@ def all_gather_matmul_hbm_buffer( fetch_block_k = config.block_size_k num_k_blocks = K // config.block_size_k + while k_per_flag > 1 and num_k_blocks % k_per_flag != 0: + k_per_flag //= 2 assert num_k_blocks % k_per_flag == 0 if workspace is None: diff --git a/tests/ops/test_all_gather_matmul.py b/tests/ops/test_all_gather_matmul.py index ce2b6e3ce..fd20c8679 100644 --- a/tests/ops/test_all_gather_matmul.py +++ b/tests/ops/test_all_gather_matmul.py @@ -415,6 +415,84 @@ def test_all_gather_matmul_hbm_buffer_trace(): ) +def test_all_gather_matmul_hbm_buffer_default_k_per_flag_small_problem(): + """The default k_per_flag is reduced to divide the K-blocks of small problems.""" + if not dist.is_initialized(): + pytest.skip("torch.distributed not initialized") + + ctx = iris.iris(2**33) + rank = ctx.get_rank() + world_size = ctx.get_num_ranks() + + dtype = torch.float16 + config = FusedConfig(block_size_m=64, block_size_n=64, block_size_k=32) + M, K_local, N = 128, 32, 64 # 1-2 K-blocks in total: fewer than the default k_per_flag of 8 + A_sharded, B, ref_output = _make_reference(rank, world_size, M, K_local, N, dtype) + A_sharded_shmem = ctx.zeros((M, K_local), dtype=dtype) + A_sharded_shmem.copy_(A_sharded) + B_shmem = ctx.zeros((K_local * world_size, N), dtype=dtype) + B_shmem.copy_(B) + output = ctx.zeros((M, N), dtype=dtype) + ctx.barrier() + + all_gather_matmul_hbm_buffer(ctx, output, A_sharded_shmem, B_shmem, config=config, trace=False) + torch.cuda.synchronize() + ctx.barrier() + + assert torch.allclose(output, ref_output, atol=1e-2, rtol=1e-2) + + +@pytest.mark.parametrize( + "m_tiles,group_size_m,num_fetch_stages", + [ + (5, 2, 1), # last group of the (only) stage is partial + (5, 8, 1), # group larger than the whole problem + (5, 4, 2), # stages of 3 + 2 M-tiles: both partial + (6, 4, 2), # stages of 3 + 3 M-tiles + (7, 3, 2), # stages of 4 + 3 M-tiles: first stage partial + ], +) +def test_all_gather_matmul_hbm_buffer_partial_m_groups(m_tiles, group_size_m, num_fetch_stages): + """M-tiles per stage need not be a multiple of group_size_m (partial last group).""" + if not dist.is_initialized(): + pytest.skip("torch.distributed not initialized") + + ctx = iris.iris(2**33) + rank = ctx.get_rank() + world_size = ctx.get_num_ranks() + + dtype = torch.float16 + config = FusedConfig(block_size_m=64, block_size_n=64, block_size_k=32, group_size_m=group_size_m) + M, K_local, N = m_tiles * config.block_size_m, 64, 128 + + A_sharded, B, ref_output = _make_reference(rank, world_size, M, K_local, N, dtype) + A_sharded_shmem = ctx.zeros((M, K_local), dtype=dtype) + A_sharded_shmem.copy_(A_sharded) + B_shmem = ctx.zeros((K_local * world_size, N), dtype=dtype) + B_shmem.copy_(B) + output = ctx.zeros((M, N), dtype=dtype) + ctx.barrier() + + # k_per_flag=1 gives several flag groups along K, which the partial-group decode must all cover. + all_gather_matmul_hbm_buffer( + ctx, + output, + A_sharded_shmem, + B_shmem, + config=config, + k_per_flag=1, + num_fetch_stages=num_fetch_stages, + trace=False, + ) + torch.cuda.synchronize() + ctx.barrier() + + max_diff = (output - ref_output).abs().max().item() + assert torch.allclose(output, ref_output, atol=1e-2, rtol=1e-2), ( + f"Rank {rank}: max diff {max_diff} (m_tiles={m_tiles}, group_size_m={group_size_m}, stages={num_fetch_stages})" + ) + + # ────────────────────────────────────────────────────────────────────── # Unit tests for _auto_config (no distributed context required) # ──────────────────────────────────────────────────────────────────────