diff --git a/python/sglang/kernels/ops/attention/fla/fused_kda_conv_recurrent_verify.py b/python/sglang/kernels/ops/attention/fla/fused_kda_conv_recurrent_verify.py index dfb5bc0ef..2d4a56fde 100644 --- a/python/sglang/kernels/ops/attention/fla/fused_kda_conv_recurrent_verify.py +++ b/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 -exactly the last ``kernel_width - 1`` input tokens, matching the reference -kernel's store). +kernels. Requires ``T >= kernel_width - 1``. -Numerics: deliberately bit-aligned with the unfused pair. The conv output is -rounded to the 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, with the same -num_warps so reduction order matches. +State: conv_state and the SSM state are read-only. Verify is speculative, and +the commit scatter advances them from the selected intermediate window. + +Numerics: aligned with the unfused pair. The conv output is rounded to the +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 - # W-1 input tokens — which are the current window registers. The verify - # 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) + # No conv-state writeback: every V tile reads the same Q/K history, so a + # tile in a later wave would read what i_v == 0 had overwritten. 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, - # conv_state and conv-window caches stay bit-identical to the reference. - # Only the fp32 intermediate-ssm rollback cache differs: the tl.sum - # reduction-order delta (~1 ulp/step) compounds through the delta-rule - # recurrence — measured ~6e-8 at T=4 standard gate (the production MTP - # 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=4 is ~1.3x faster than the unfused pair in-graph; 1 restores the + # reference reduction order but is ~2.4x slower, for numerics debugging only. + # The fp32 intermediate-ssm rollback cache carries the reduction-order delta + # furthest: ~6e-8 at T=4 standard gate (the production MTP shape), ~2e-3 at + # T=8 safe gate. conv_state is not comparable to the reference at all. num_warps: int = 4, ) -> torch.Tensor: """Chain-verify fast path. Returns ``o`` of shape [1, seq_len, HV, V], diff --git a/test/registered/kernels/test_fused_kda_conv_recurrent_verify.py b/test/registered/kernels/test_fused_kda_conv_recurrent_verify.py index f67fe4634..3645cc1db 100644 --- a/test/registered/kernels/test_fused_kda_conv_recurrent_verify.py +++ b/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__]))