[gdn] fused replayssm ring write into flashinfer gdn mtp verify kernel (#33102)

python/sglang/srt/layers/attention/linear/gdn_backend.py

@@ -183,8 +183,10 @@ class GDNKernelDispatcher:
             decode_backend.is_flashinfer() or prefill_backend.is_flashinfer()
         ) and flashinfer_kernel.supports_target_verify:
             self.verify_kernel = flashinfer_kernel
+            self.verify_kernel_is_flashinfer = True
         else:
             self.verify_kernel = triton_kernel
+            self.verify_kernel_is_flashinfer = False
 
         self.supports_packed_decode = getattr(
             self.decode_kernel, "supports_packed_decode", False
@@ -709,9 +711,10 @@ class GDNAttnBackend(MambaAttnBackendBase):
         query_start_loc: torch.Tensor,
         retrieve_parent_token: Optional[torch.Tensor],
     ) -> torch.Tensor:
-        """Recurrent verify + fused ring-write; the commit fold replays the
-        accepted prefix into ``temporal``. Called directly, not via the kernel
-        dispatcher: the ring-write exists only in the Triton kernel."""
+        """Ring-writing verify; the commit fold replays the accepted prefix
+        into ``temporal``. Uses the vendored CuTe DSL MTP kernel when the
+        dispatcher selected the FlashInfer bf16-state verify, else the Triton
+        recurrent kernel (both store the same raw window)."""
         from sglang.kernels.ops.attention.fla.fused_sigmoid_gating_recurrent import (
             fused_sigmoid_gating_delta_rule_update,
         )
@@ -720,6 +723,39 @@ class GDNAttnBackend(MambaAttnBackendBase):
             "ReplaySSM fold-every-commit supports a linear draft chain only "
             "(topk <= 1); EAGLE tree verify must use the recurrent verify."
         )
+        seq_len = query.shape[1]
+        batch_size = query_start_loc.shape[0] - 1
+        draft_token_num = seq_len // batch_size
+        if (
+            self.kernel_dispatcher.verify_kernel_is_flashinfer
+            and ssm_states.dtype == torch.bfloat16
+            and draft_token_num >= 3
+        ):
+            from sglang.kernels.ops.attention.cutedsl_gdn_mtp_ring import (
+                gated_delta_rule_mtp,
+            )
+
+            num_v_heads = value.shape[2]
+            head_v_dim = value.shape[3]
+            out = gated_delta_rule_mtp(
+                A_log=layer.A_log.detach(),
+                a=a.view(batch_size, draft_token_num, num_v_heads),
+                dt_bias=layer.dt_bias.detach(),
+                q=query.view(batch_size, draft_token_num, *query.shape[2:]),
+                k=key.view(batch_size, draft_token_num, *key.shape[2:]),
+                v=value.view(batch_size, draft_token_num, num_v_heads, head_v_dim),
+                b=b.view(batch_size, draft_token_num, num_v_heads),
+                initial_state_source=ssm_states,
+                initial_state_indices=cache_indices,
+                use_qk_l2norm_in_kernel=True,
+                disable_state_update=True,
+                cache_ring=True,
+                replayssm_rawv=layer_cache.replayssm_rawv,
+                replayssm_rawk=layer_cache.replayssm_rawk,
+                replayssm_g=layer_cache.replayssm_g,
+                replayssm_beta=layer_cache.replayssm_beta,
+            )
+            return out.view(1, seq_len, num_v_heads, head_v_dim)
         return fused_sigmoid_gating_delta_rule_update(
             A_log=layer.A_log,
             dt_bias=layer.dt_bias,

test/registered/attention/unittests/gdn/test_gdn_cutedsl_ring_verify.py

@@ -0,0 +1,150 @@
+"""Vendored CuTe DSL GDN MTP verify with the fused ReplaySSM ring-write.
+
+The ring-write must be a pure side channel and the ring must feed the fold a
+faithful raw window, so the anchor is:
+
+  * the verify OUTPUT is bitwise unchanged by cache_ring (both the ilp4 and
+    the wide_vec kernels);
+  * rawv/rawk are bitwise copies of the kernel inputs; g matches the Triton
+    gating and beta matches the fp32 sigmoid to fastmath tolerance;
+  * folding the ring (Triton fold kernel) reproduces the CuTe DSL kernel's
+    OWN committed state (disable_state_update=False run) to bf16-ulp
+    tolerance -- the mixed-numerics bound replacing the triton-vs-triton
+    bitwise anchor.
+"""
+
+import unittest
+
+import torch
+
+from sglang.kernels.ops.attention.cutedsl_gdn_mtp_ring import gated_delta_rule_mtp
+from sglang.kernels.ops.attention.fla.fused_gdn_gating import fused_gdn_gating
+from sglang.kernels.ops.attention.fla.gdn_replayssm_spec_fold import (
+    commit_gdn_replayssm_fold_all_layers,
+)
+from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cuda_ci(est_time=120, stage="base-b", runner_config="1-gpu-large")
+
+T, H, HV, K, V, SLOTS = 4, 4, 16, 128, 128, 16
+DEVICE = "cuda"
+
+
+def _case(B):
+    gen = torch.Generator(device=DEVICE).manual_seed(7)
+
+    def rnd(*shape, dtype=torch.bfloat16):
+        return torch.randn(*shape, device=DEVICE, dtype=dtype, generator=gen)
+
+    gating = {
+        "A_log": (torch.randn(HV, device=DEVICE, generator=gen) * 0.1).float(),
+        "dt_bias": (torch.randn(HV, device=DEVICE, generator=gen) * 0.1).float(),
+    }
+    inputs = {
+        "q": rnd(B, T, H, K),
+        "k": rnd(B, T, H, K),
+        "v": rnd(B, T, HV, V),
+        "a": rnd(B, T, HV),
+        "b": rnd(B, T, HV),
+    }
+    state0 = rnd(SLOTS, HV, V, K)
+    slots = torch.arange(B, device=DEVICE, dtype=torch.int32) + 2
+    rings = {
+        "rawv": torch.zeros(1, SLOTS, HV, T, V, device=DEVICE, dtype=torch.bfloat16),
+        "rawk": torch.zeros(1, SLOTS, H, T, K, device=DEVICE, dtype=torch.bfloat16),
+        "g": torch.zeros(1, SLOTS, HV, T, device=DEVICE, dtype=torch.float32),
+        "beta": torch.zeros(1, SLOTS, HV, T, device=DEVICE, dtype=torch.float32),
+    }
+    return gating, inputs, state0, slots, rings
+
+
+def _verify(gating, inputs, state, slots, rings=None, disable_state_update=True):
+    kwargs = {}
+    if rings is not None:
+        kwargs.update(
+            cache_ring=True,
+            replayssm_rawv=rings["rawv"][0],
+            replayssm_rawk=rings["rawk"][0],
+            replayssm_g=rings["g"][0],
+            replayssm_beta=rings["beta"][0],
+        )
+    return gated_delta_rule_mtp(
+        gating["A_log"],
+        inputs["a"],
+        gating["dt_bias"],
+        q=inputs["q"],
+        k=inputs["k"],
+        v=inputs["v"],
+        b=inputs["b"],
+        initial_state_source=state,
+        initial_state_indices=slots,
+        use_qk_l2norm_in_kernel=True,
+        disable_state_update=disable_state_update,
+        **kwargs,
+    )
+
+
+class TestGdnCuteDSLRingVerify(CustomTestCase):
+    def _run(self, B):
+        gating, inputs, state0, slots, rings = _case(B)
+
+        out_ref = _verify(gating, inputs, state0.clone(), slots)
+        out_ring = _verify(gating, inputs, state0.clone(), slots, rings=rings)
+        self.assertTrue(torch.equal(out_ref, out_ring), f"{B=}")
+
+        for i, s in enumerate(slots.tolist()):
+            self.assertTrue(
+                torch.equal(rings["rawv"][0, s], inputs["v"][i].transpose(0, 1))
+            )
+            self.assertTrue(
+                torch.equal(rings["rawk"][0, s], inputs["k"][i].transpose(0, 1))
+            )
+
+        g_ref, _ = fused_gdn_gating(
+            gating["A_log"],
+            inputs["a"].view(B * T, HV),
+            inputs["b"].view(B * T, HV),
+            gating["dt_bias"],
+        )
+        g_ref = g_ref.view(B, T, HV).transpose(1, 2).float()
+        beta_ref = torch.sigmoid(inputs["b"].float()).transpose(1, 2)
+        self.assertLess((rings["g"][0, slots.long()] - g_ref).abs().max().item(), 5e-5)
+        self.assertLess(
+            (rings["beta"][0, slots.long()] - beta_ref).abs().max().item(), 5e-5
+        )
+
+        state_ref = state0.clone()
+        _verify(gating, inputs, state_ref, slots, disable_state_update=False)
+        fold_state = state0.clone().unsqueeze(0)
+        commit_gdn_replayssm_fold_all_layers(
+            checkpoint_state=fold_state,
+            rawv_cache=rings["rawv"],
+            rawk_cache=rings["rawk"],
+            g_cache=rings["g"],
+            beta_cache=rings["beta"],
+            ssm_state_indices=slots,
+            accept_lens=torch.full((B,), T, device=DEVICE, dtype=torch.int32),
+            max_cache_len=T,
+            num_k_heads=H,
+        )
+        touched = slots.long()
+        err = (
+            (fold_state[0, touched].float() - state_ref[touched].float())
+            .abs()
+            .max()
+            .item()
+        )
+        self.assertLess(err, 3e-2, f"{B=} fold vs own update: {err}")
+        untouched = [s for s in range(SLOTS) if s not in slots.tolist()]
+        self.assertTrue(torch.equal(fold_state[0, untouched], state0[untouched]))
+
+    def test_ilp4_small_batch(self):
+        self._run(1)
+
+    def test_wide_vec_batch(self):
+        self._run(8)
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/attention/unittests/gdn/test_gdn_replayssm_spec_fold.py

@@ -9,14 +9,10 @@ verify->commit chain stays bitwise equal at every step (no accumulation channel
 on the state path -- the long-decode drift failure mode).
 """
 
-import sys
 import unittest
-from pathlib import Path
 
 import torch
 
-sys.path.insert(0, str(Path(__file__).resolve().parents[4]))
-
 from sglang.kernels.ops.attention.fla.fused_sigmoid_gating_recurrent import (
     fused_sigmoid_gating_delta_rule_update,
 )