[None][fix] wait for K1 readers before reusing KDA prefill stages

tensorrt_llm/_torch/cute_dsl_kernels/blackwell/kimi_k3_kda/fused_k123.py

@@ -22,23 +22,25 @@
   Total work units = (NT/4) * H * B, distributed round-robin across SMs
   Block i processes work units i, i+NUM_SMS, i+2*NUM_SMS, ...
 Block: 992 threads (31 warps), warp-specialized with setmaxnreg:
-  Warps 0-15:  TMA+K1 fused (8×2, vec2, prefetch pipeline) – 4 WGs, 56 regs
-  Warps 16-26: K2 MMA compute (10 active + warp 26 as TMA producer) – 72 regs
+  Warps 0-15:  K1 compute (8×2, vec2) – 4 WGs, 56 regs
+  Warps 16-25: K2 MMA compute – 72 regs
+  Warp 26:     Dedicated TMA producer (TMA_WARP_ID) – 72 regs
   Warps 27-30: Store/Inversion warps – 24 regs
 
 Pipeline (single for_generate, warp groups separated by if-blocks):
   per work unit:
-    Warps 0-15:  prefetch chunk 0→stage 0 (warp 0), then loop:
-                    TMA next chunk (warp 0), wait cur chunk, K1 compute, arrive(k1_done)
-    Warps 16-26: wait(k1_done)+wait(store_done), MMA, arrive(mma_done+stage_reuse)
+    Warps 0-15:  wait(tma), K1 cumsum, arrive(k1_done),
+                    Pass 2b Q/K reads and scaling, arrive(stage_reuse)
+    Warps 16-25: wait(k1_done)+wait(store_done), MMA, arrive(mma_done+stage_reuse)
+    Warp 26:     wait(stage_reuse), TMA Q/K/G for each chunk, signal(tma)
     Warps 27-30: wait(mma_done), store sAqk/sAkk→GMEM, arrive(store_done)
   All warp-group invariants are computed inside each group's if-block (not hoisted)
   to eliminate cross-group register pressure — same budget as the _all version.
   Mbarrier phases self-reset after 4 iterations (2 stages × 2 phases).
 
 Mbarriers:
-  tma_mbars[2]:          count=1, warp 0 lane 0 → K1+MMA wait for TMA data
-  stage_reuse_mbars[2]:  count=320, MMA(10 warps) → warp 0 waits before TMA reuse
+  tma_mbars[2]:          count=1, TMA_WARP_ID (26) lane 0 → K1+MMA wait for TMA data
+  stage_reuse_mbars[2]:  count=832, K1(16)+MMA(10 warps) → TMA waits before reuse
   k1_done_mbars[2]:      count=512, K1(16 warps) → MMA waits for g_cumsum ready
   mma_done_mbars[2]:     count=320, MMA(10 warps) → Store waits for sAqk/sAkk ready
   store_done_mbars[2]:   count=128, Store(4 warps) → MMA waits for sAqk/sAkk stage free
@@ -1035,8 +1037,11 @@ def fused_kernel123(
     if tidx == 0:
         for s in range(NUM_STAGES):
             cute.arch.mbarrier_init(tma_mbars + s, 1)
-            # TMA warp is waiter (not arriver) on stage_reuse; and it skips mma_done arrive
-            cute.arch.mbarrier_init(stage_reuse_mbars + s, (NUM_MMA_WARPS - 1) * 32)
+            # Both K1 Pass 2b and MMA read Q/K before TMA can reuse a stage.
+            # The dedicated TMA warp does not arrive on either barrier.
+            cute.arch.mbarrier_init(
+                stage_reuse_mbars + s, (NUM_MMA_WARPS - 1 + NUM_K1_TMA_WARPS) * 32
+            )
             cute.arch.mbarrier_init(k1_done_mbars + s, NUM_K1_TMA_WARPS * 32)
             cute.arch.mbarrier_init(mma_done_mbars + s, (NUM_MMA_WARPS - 1) * 32)
             cute.arch.mbarrier_init(store_done_mbars + s, NUM_STORE_WARPS * 32)
@@ -1070,8 +1075,8 @@ def fused_kernel123(
     cute.arch.barrier()
 
     # =====================================================================
-    # Pre-arrive (MMA warps only)
-    # stage_reuse_mbars: warp 0 waits before MMA arrives → pre-arrive all 10 MMA warps
+    # Pre-arrive before the first TMA load and MMA store
+    # stage_reuse_mbars: pre-arrive all 16 K1 and 10 MMA reader warps
     # store_done_mbars:  MMA waits before Store arrives → pre-arrive first 4 MMA warps
     # =====================================================================
     if (
@@ -1085,6 +1090,10 @@ def fused_kernel123(
             if mma_warp_tmp < NUM_STORE_WARPS:
                 cute.arch.mbarrier_arrive(store_done_mbars + s)
 
+    if warp_idx < NUM_K1_TMA_WARPS:
+        for s in range(NUM_STAGES):
+            cute.arch.mbarrier_arrive(stage_reuse_mbars + s)
+
     # =================================================================
     # Persistent outer loop. Single for_generate at top level (required).
     # Opaque asm barrier on work_id prevents MLIR LICM from hoisting
@@ -1365,9 +1374,13 @@ def fused_kernel123(
                     else:
                         cute.autovec_copy(rGkOut, mGkLast[i_b, chunk_idx, i_h, col_vec_idx, None])
 
+                # K1 still reads Q/K in Pass 2b after signaling k1_done.
+                # TMA may reuse this stage only once those reads finish too.
+                cute.arch.mbarrier_arrive(stage_reuse_mbars + cur_stage)
+
         # =============================================================
         # Warp 26 (TMA_WARP_ID): dedicated TMA producer.
-        # Waits stage_reuse (gated by MMA arrives), issues TMA for Q/K/G,
+        # Waits stage_reuse (K1 and MMA readers), issues TMA for Q/K/G,
         # signals tma_mbar. Decouples MMA -> TMA dependency from K1 compute.
         # =============================================================
         if warp_idx == TMA_WARP_ID:

tests/unittest/_torch/modules/kimi_kda/test_kda_prefill_state_parity.py

@@ -397,3 +397,75 @@ def test_chunked_continuation_matches_fla(
     assert_kda_close("chunked/first_output", first_output, full_output[:, :split])
     assert_kda_close("chunked/second_output", second_output, full_output[:, split:])
     assert_kda_close("chunked/state", second_state, full_state)
+
+
+@pytest.mark.parametrize("use_cuda_graph", [False, True], ids=["eager", "cuda_graph"])
+@torch.no_grad()
+def test_prefill_stage_reuse_is_repeatable(use_cuda_graph: bool) -> None:
+    """TMA must wait for K1's Q/K reads as well as MMA before reusing a stage."""
+    generator = torch.Generator(device="cuda").manual_seed(20260921)
+    length = 1186
+    shape = (1, length, NUM_HEADS, HEAD_DIM)
+
+    def random_tensor(shape: tuple[int, ...]) -> torch.Tensor:
+        return torch.randn(shape, device="cuda", generator=generator)
+
+    q = torch.nn.functional.normalize(random_tensor(shape), dim=-1).bfloat16()
+    k = torch.nn.functional.normalize(random_tensor(shape), dim=-1).bfloat16()
+    v = random_tensor(shape).bfloat16()
+    g = random_tensor(shape).bfloat16()
+    beta = random_tensor(shape[:-1])
+    a_log = torch.zeros(NUM_HEADS, device="cuda")
+    dt_bias = torch.zeros(NUM_HEADS * HEAD_DIM, device="cuda")
+    state = torch.empty(1, NUM_HEADS, HEAD_DIM, HEAD_DIM, device="cuda")
+    indices = torch.zeros(1, dtype=torch.int64, device="cuda")
+    cu_seqlens = _make_cu_seqlens([length])
+    chunk_count = (length + 63) // 64
+    chunk_indices = torch.stack(
+        (
+            torch.zeros(chunk_count, device="cuda", dtype=torch.int64),
+            torch.arange(chunk_count, device="cuda"),
+        ),
+        dim=-1,
+    )
+    dispatch = KDAKernelDispatch(use_optimized_prefill=True, use_optimized_decode=False)
+
+    def run() -> torch.Tensor:
+        state.zero_()
+        output, _ = dispatch.prefill_chunk_kda(
+            q=q,
+            k=k,
+            v=v,
+            g=g,
+            beta=beta,
+            A_log=a_log,
+            dt_bias=dt_bias,
+            scale=HEAD_DIM**-0.5,
+            initial_state=None,
+            safe_gate=True,
+            lower_bound=-5.0,
+            cu_seqlens=cu_seqlens,
+            chunk_indices=chunk_indices,
+            state_pool=state,
+            state_indices=indices,
+            varlen_is_aligned=False,
+            single_sequence_length=length,
+        )
+        return output
+
+    for _ in range(3):
+        output = run()
+    expected = output.clone()
+    expected_state = state.clone()
+    graph = None
+    if use_cuda_graph:
+        graph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(graph):
+            output = run()
+    for _ in range(128):
+        if graph is None:
+            output = run()
+        else:
+            graph.replay()
+        torch.testing.assert_close(output, expected, rtol=0, atol=0)
+        torch.testing.assert_close(state, expected_state, rtol=0, atol=0)