Parallelize partial output-tile stores across the store warp

benchmarks/bench_fwd.py

@@ -141,14 +141,20 @@ def run_chunk_gated_delta_rule():
     [1024] * 8,
 ]
 
+# Isolates the partial output-tile path: one tail tile per sequence and head.
+# Compare the same command before and after a kernel change.
+TAIL_STORE_CASES = [
+    [15],
+]
+
 
 def main():
     import argparse
     p = argparse.ArgumentParser()
     p.add_argument("--warmup", type=int, default=30)
     p.add_argument("--iters", type=int, default=200)
     p.add_argument("--repeats", type=int, default=5)
-    p.add_argument("--mode", choices=["fixed", "varlen", "all"], default="all")
+    p.add_argument("--mode", choices=["fixed", "varlen", "tail-store", "all"], default="all")
     p.add_argument("--H", type=int, default=96)
     p.add_argument("--D", type=int, default=128)
     args = p.parse_args()
@@ -158,6 +164,8 @@ def main():
         cases.extend(FIXED_CASES)
     if args.mode in ("varlen", "all"):
         cases.extend(VARLEN_CASES)
+    if args.mode == "tail-store":
+        cases.extend(TAIL_STORE_CASES)
 
     for seq_lens in cases:
         run_case(seq_lens, args.H, args.D, args.warmup, args.iters, args.repeats)

csrc/smxx/fwd_kernel2.cuh

@@ -739,28 +739,32 @@ __global__ void __launch_bounds__(NumThreads) _flash_kda_fwd_recurrence(
     }
 
 #ifndef TMA_DISABLE_ALL
-    if (warp_role == WarpRole::STORE && lane_predicate) {
+    if (warp_role == WarpRole::STORE) {
         Tensor g_out = tma_store_out.get_tma_tensor(make_shape(H, T_total, D));
         auto cta_tma_store = tma_store_out.get_slice(Int<0>{});
         StorePipelineState out_read;
+        int lane_idx = threadIdx.x % kWarpSize;
         for (int t = 0; t < t_tiles; ++t) {
-            store_pipeline.consumer_wait(out_read);
+            if (lane_predicate) {
+                store_pipeline.consumer_wait(out_read);
+            }
+            __syncwarp();
+
             int stage = out_read.index();
             int actual_len = min(CHUNK, seq_len - t * CHUNK);
 
             BF16* out_stage_ptr = shared_storage.output[stage].out.begin();
 
             if (actual_len < CHUNK) {
                 // Manual store for tail tile to avoid overwriting next sequence
-                // Only one thread (lane_predicate) runs here, so loop over all D
                 Tensor s_out = make_tensor(make_smem_ptr(out_stage_ptr), VOLayout{});
-                for (int row = 0; row < actual_len; ++row) {
+                for (int idx = lane_idx; idx < actual_len * D; idx += kWarpSize) {
+                    int row = idx / D;
+                    int col = idx % D;
                     int64_t global_base = (bos + t * CHUNK + row) * H * D + head_idx * D;
-                    for (int col = 0; col < D; ++col) {
-                        out_raw_ptr[global_base + col] = s_out(row, col);
-                    }
+                    out_raw_ptr[global_base + col] = s_out(row, col);
                 }
-            } else {
+            } else if (lane_predicate) {
                 // TMA store for full tiles
                 auto out_off = g_out.layout()(head_idx, int(bos) + t * CHUNK, 0);
                 Tensor g_out_tile = make_tensor(g_out.data() + out_off,
@@ -774,26 +778,34 @@ __global__ void __launch_bounds__(NumThreads) _flash_kda_fwd_recurrence(
                 tma_store_arrive();
             }
 
-            tma_store_wait<0>();
-            store_pipeline.consumer_release(out_read);
+            if (lane_predicate) {
+                tma_store_wait<0>();
+            }
+            // Every lane must finish reading the stage before its leader releases it.
+            __syncwarp();
+            if (lane_predicate) {
+                store_pipeline.consumer_release(out_read);
+            }
             ++out_read;
         }
 
         if constexpr (HasStateOut && !StateFP32) {
-            // BF16 state: TMA store directly from state_acc
-            Tensor g_final = tma_store_final_state.get_tma_tensor(make_shape(N * H, D, D));
-            auto state_off = g_final.layout()(seq_idx * H + head_idx, 0, 0);
-            Tensor g_final_tile = make_tensor(g_final.data() + state_off,
-                make_layout(make_shape(Int<1>{}, Int<D>{}, Int<D>{}), stride(g_final.layout())));
-            Tensor s_state = make_tensor(make_smem_ptr(shared_storage.state_acc.begin()), TMAStateSmemLayout{});
-
-            auto cta_tma_store_state = tma_store_final_state.get_slice(Int<0>{});
-            cute::copy(
-                tma_store_final_state,
-                cta_tma_store_state.partition_S(s_state),
-                cta_tma_store_state.partition_D(g_final_tile)
-            );
-            tma_store_arrive();
+            if (lane_predicate) {
+                // BF16 state: TMA store directly from state_acc
+                Tensor g_final = tma_store_final_state.get_tma_tensor(make_shape(N * H, D, D));
+                auto state_off = g_final.layout()(seq_idx * H + head_idx, 0, 0);
+                Tensor g_final_tile = make_tensor(g_final.data() + state_off,
+                    make_layout(make_shape(Int<1>{}, Int<D>{}, Int<D>{}), stride(g_final.layout())));
+                Tensor s_state = make_tensor(make_smem_ptr(shared_storage.state_acc.begin()), TMAStateSmemLayout{});
+
+                auto cta_tma_store_state = tma_store_final_state.get_slice(Int<0>{});
+                cute::copy(
+                    tma_store_final_state,
+                    cta_tma_store_state.partition_S(s_state),
+                    cta_tma_store_state.partition_D(g_final_tile)
+                );
+                tma_store_arrive();
+            }
         }
     }
 

tests/test_fwd_full.py

@@ -150,6 +150,35 @@ def test_fwd_varlen(seq_lens, H, state_dtype, has_in, has_out):
             f"final_state mismatch: seqs={seq_lens} H={H} dtype={state_dtype} in={has_in} out={has_out}"
 
 
+def test_fwd_varlen_all_tail_lengths():
+    """Exercise every partial 16-row output tile at adjacent sequence boundaries."""
+    seq_lens = list(range(1, 16))
+    T_total = sum(seq_lens)
+    H = 4
+    cu_seqlens = torch.tensor(
+        [0] + list(torch.cumsum(torch.tensor(seq_lens), dim=0).tolist()),
+        dtype=torch.long, device="cuda",
+    )
+
+    q, k, v, g, beta, A_log, dt_bias, scale = _make_inputs(T_total, H)
+    out_kernel = torch.full_like(q, float("nan"))
+    flash_kda.fwd(
+        q, k, v, g, beta, scale, out_kernel,
+        A_log=A_log, dt_bias=dt_bias, lower_bound=LOWER_BOUND,
+        cu_seqlens=cu_seqlens,
+    )
+    torch.cuda.synchronize()
+
+    out_ref = torch.empty_like(q)
+    torch_ref(
+        q, k, v, g, beta, scale, out_ref,
+        A_log=A_log, dt_bias=dt_bias, lower_bound=LOWER_BOUND,
+        cu_seqlens=cu_seqlens,
+    )
+
+    assert torch.equal(out_kernel, out_ref)
+
+
 # ---------------------------------------------------------------------------
 # Batched tests (B > 1, equal-length sequences)
 # ---------------------------------------------------------------------------