Fix all_gather_matmul_hbm_buffer: partial M groups, small K, int64 offsets

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)
 

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:

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