[Fix] Don't write conv state from the fused KDA verify kernel (#39219)

Co-authored-by: mmangkad <mohammad.angkad@radixark.ai>

python/sglang/kernels/ops/attention/fla/fused_kda_conv_recurrent_verify.py

@@ -12,15 +12,16 @@ two transpose copies the unfused path needs to feed the conv kernel.
 
 Scope (v1): chain speculation only (``speculative_eagle_topk == 1``, i.e.
 ``retrieve_next_token is None``). The tree path keeps the unfused reference
-kernels. Requires ``T >= kernel_width - 1`` (the rolled conv state is then
+kernels. Requires ``T >= kernel_width - 1``.
-exactly the last ``kernel_width - 1`` input tokens, matching the reference
-kernel's store).
 
-Numerics: deliberately bit-aligned with the unfused pair. The conv output is
+State: conv_state and the SSM state are read-only. Verify is speculative, and
-rounded to the activation dtype (bf16) before entering the recurrence —
+the commit scatter advances them from the selected intermediate window.
-exactly what the unfused path does through its intermediate tensor — and all
+
-expressions mirror the reference kernels line by line, with the same
+Numerics: aligned with the unfused pair. The conv output is rounded to the
-num_warps so reduction order matches.
+activation dtype (bf16) before entering the recurrence — exactly what the
+unfused path does through its intermediate tensor — and all expressions mirror
+the reference kernels line by line. Reduction order still splits differently
+where many V heads share one Q/K head, worth ~1 ulp on the output.
 """
 
 from typing import Optional
@@ -318,19 +319,8 @@ def fused_kda_conv_gating_verify_kernel(
                 )
                 tl.store(cache_ptr, b_h.to(cache_ptr.dtype.element_ty), mask=mask_h)
 
-    # Rolled conv state after consuming T >= W-1 tokens is exactly the last
+    # No conv-state writeback: every V tile reads the same Q/K history, so a
-    # W-1 input tokens — which are the current window registers. The verify
+    # tile in a later wave would read what i_v == 0 had overwritten.
-    # pass never writes the ssm state back (rollback happens at commit).
-    if is_qk_owner:
-        tl.store(cs_base + q_ch + 0 * stride_cs_tok, q_c0, mask=mask_k)
-        tl.store(cs_base + q_ch + 1 * stride_cs_tok, q_c1, mask=mask_k)
-        tl.store(cs_base + q_ch + 2 * stride_cs_tok, q_c2, mask=mask_k)
-        tl.store(cs_base + k_ch + 0 * stride_cs_tok, k_c0, mask=mask_k)
-        tl.store(cs_base + k_ch + 1 * stride_cs_tok, k_c1, mask=mask_k)
-        tl.store(cs_base + k_ch + 2 * stride_cs_tok, k_c2, mask=mask_k)
-    tl.store(cs_base + v_ch + 0 * stride_cs_tok, v_c0, mask=mask_v)
-    tl.store(cs_base + v_ch + 1 * stride_cs_tok, v_c1, mask=mask_v)
-    tl.store(cs_base + v_ch + 2 * stride_cs_tok, v_c2, mask=mask_v)
 
 
 def fused_kda_conv_gating_verify(
@@ -358,14 +348,11 @@ def fused_kda_conv_gating_verify(
     softplus_beta: float = 1.0,
     softplus_threshold: float = 20.0,
     use_qk_l2norm_in_kernel: bool = True,
-    # num_warps=4 is ~1.3x faster than the unfused pair in-graph; the output,
+    # num_warps=4 is ~1.3x faster than the unfused pair in-graph; 1 restores the
-    # conv_state and conv-window caches stay bit-identical to the reference.
+    # reference reduction order but is ~2.4x slower, for numerics debugging only.
-    # Only the fp32 intermediate-ssm rollback cache differs: the tl.sum
+    # The fp32 intermediate-ssm rollback cache carries the reduction-order delta
-    # reduction-order delta (~1 ulp/step) compounds through the delta-rule
+    # furthest: ~6e-8 at T=4 standard gate (the production MTP shape), ~2e-3 at
-    # recurrence — measured ~6e-8 at T=4 standard gate (the production MTP
+    # T=8 safe gate. conv_state is not comparable to the reference at all.
-    # shape), ~1.5e-5 at T=4 safe gate, ~2e-3 at T=8 safe gate. num_warps=1
-    # reproduces the reference reduction order exactly (all buffers
-    # bit-identical) but is ~2.4x slower in-graph — numerics debugging only.
     num_warps: int = 4,
 ) -> torch.Tensor:
     """Chain-verify fast path. Returns ``o`` of shape [1, seq_len, HV, V],

test/registered/kernels/test_fused_kda_conv_recurrent_verify.py

@@ -166,12 +166,12 @@ def _compare_case(case, num_warps):
 
     idx_vals = inp["idx_vals"]
     valid_rows = [i for i, slot in enumerate(idx_vals) if slot >= 0]
-    touched_slots = [slot for slot in idx_vals if slot >= 0]
 
     o_ref_v = o_ref.reshape(B, T, HV, V)[valid_rows]
     o_fus_v = o_fus.reshape(B, T, HV, V)[valid_rows]
     assert torch.equal(o_ref_v, o_fus_v)
-    assert torch.equal(conv_ref[touched_slots], conv_fus[touched_slots])
+    # conv_state is read-only in verify; the commit scatter advances it.
+    assert torch.equal(inp["conv_pool"], conv_fus)
     assert torch.equal(win_ref[valid_rows], win_fus[valid_rows])
     torch.testing.assert_close(
         ic_ref[valid_rows], ic_fus[valid_rows], atol=4e-3, rtol=0
@@ -183,5 +183,27 @@ def test_matches_unfused_reference(case):
     _compare_case(case, num_warps=4)
 
 
+def test_output_does_not_depend_on_cta_scheduling():
+    """The verify output must not change with how the CTAs happen to be
+    scheduled. H=1 with HV=16 shares one Q/K history across 16 V tiles."""
+    if torch.cuda.get_device_capability()[0] < 9:
+        pytest.skip("green contexts need SM90 or newer")
+    from flashinfer.green_ctx import split_device_green_ctx_by_sm_count
+
+    case = (1, 6, 1, 16, 128, 128, 4, False, None, False, 1)
+    B, T, H, HV, K, V, W, has_bias, lower_bound, neg_slot, seed = case
+    inp = _make_inputs(B, T, H, HV, K, V, W, has_bias, neg_slot, seed)
+    full = _run_fused(inp, B, T, H, HV, K, V, lower_bound, num_warps=4)[0]
+
+    streams, _ = split_device_green_ctx_by_sm_count(torch.device("cuda:0"), [8])
+    # The green stream is non-blocking, so it must be told to wait for the
+    # inputs produced above; synchronize() afterwards only waits on the consumer.
+    streams[0].wait_stream(torch.cuda.current_stream())
+    with torch.cuda.stream(streams[0]):
+        squeezed = _run_fused(inp, B, T, H, HV, K, V, lower_bound, num_warps=4)[0]
+    streams[0].synchronize()
+    assert torch.equal(full, squeezed)
+
+
 if __name__ == "__main__":
     sys.exit(pytest.main([__file__]))