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) # ──────────────────────────────────────────────────────────────────────