[GDN][Qwen3-Next][Qwen3.5] Fuse fused_gdn_gating and fused_recurrent_gated_delta_rule_update in verify_target (#19775)

python/sglang/jit_kernel/tests/test_fused_verify_triton_gdn.py

@@ -0,0 +1,231 @@
+"""Tests for fused sigmoid gating delta rule MTP kernel (GDN target_verify).
+
+Compares the fused kernel `fused_sigmoid_gating_delta_rule_update` against
+the reference two-step implementation:
+    1. g, beta = fused_gdn_gating(A_log, a, b, dt_bias)
+    2. o = fused_recurrent_gated_delta_rule_update(q, k, v, g, beta, ...)
+"""
+
+import pytest
+import torch
+
+try:
+    from sglang.srt.layers.attention.fla.fused_gdn_gating import fused_gdn_gating
+    from sglang.srt.layers.attention.fla.fused_recurrent import (
+        fused_recurrent_gated_delta_rule_update,
+    )
+    from sglang.srt.layers.attention.fla.fused_sigmoid_gating_recurrent import (
+        fused_sigmoid_gating_delta_rule_update,
+    )
+
+    KERNELS_AVAILABLE = True
+except ImportError:
+    KERNELS_AVAILABLE = False
+
+
+def _make_tensors(N, T, H, HV, K, V, device="cuda", seed=2025):
+    """Create input tensors for GDN target_verify."""
+    torch.manual_seed(seed)
+    A_log = torch.randn(HV, dtype=torch.float32, device=device)
+    dt_bias = torch.randn(HV, dtype=torch.bfloat16, device=device)
+    a = torch.randn(1, N * T, HV, dtype=torch.bfloat16, device=device)
+    b = torch.randn(1, N * T, HV, dtype=torch.bfloat16, device=device)
+    q = torch.randn(1, N * T, H, K, dtype=torch.bfloat16, device=device)
+    k = torch.randn(1, N * T, H, K, dtype=torch.bfloat16, device=device)
+    v = torch.randn(1, N * T, HV, V, dtype=torch.bfloat16, device=device)
+    indices = torch.arange(N, dtype=torch.int32, device=device)
+    initial_state = torch.randn(N, HV, K, V, dtype=torch.float, device=device)
+    cu_seqlens = torch.arange(0, N * T + 1, T, dtype=torch.int32, device=device)
+    return A_log, dt_bias, a, b, q, k, v, initial_state, indices, cu_seqlens
+
+
+def run_reference(
+    A_log,
+    dt_bias,
+    q,
+    k,
+    v,
+    a,
+    b,
+    initial_state_source,
+    initial_state_indices,
+    cu_seqlens,
+    disable_state_update=True,
+    intermediate_states_buffer=None,
+    intermediate_state_indices=None,
+    cache_steps=None,
+    retrieve_parent_token=None,
+):
+    """Reference: fused_gdn_gating + fused_recurrent_gated_delta_rule_update."""
+    # fused_gdn_gating expects 2D [seq_len, HV]
+    a_2d = a.view(-1, a.shape[-1])
+    b_2d = b.view(-1, b.shape[-1])
+    g, beta = fused_gdn_gating(A_log, a_2d, b_2d, dt_bias)
+    # fused_recurrent expects 3D [B, T, HV]
+    g = g.view(a.shape)
+    beta = beta.view(b.shape)
+
+    # fused_recurrent requires intermediate_state_indices when cu_seqlens is used
+    if cu_seqlens is not None and intermediate_state_indices is None:
+        N = len(cu_seqlens) - 1
+        intermediate_state_indices = torch.arange(N, dtype=torch.int32, device=q.device)
+
+    return fused_recurrent_gated_delta_rule_update(
+        q=q,
+        k=k,
+        v=v,
+        g=g,
+        beta=beta,
+        initial_state_source=initial_state_source,
+        initial_state_indices=initial_state_indices,
+        cu_seqlens=cu_seqlens,
+        use_qk_l2norm_in_kernel=True,
+        disable_state_update=disable_state_update,
+        intermediate_states_buffer=intermediate_states_buffer,
+        intermediate_state_indices=intermediate_state_indices,
+        cache_steps=cache_steps,
+        retrieve_parent_token=retrieve_parent_token,
+    )
+
+
+def run_fused_mtp(
+    A_log,
+    dt_bias,
+    q,
+    k,
+    v,
+    a,
+    b,
+    initial_state_source,
+    initial_state_indices,
+    cu_seqlens,
+    disable_state_update=True,
+    intermediate_states_buffer=None,
+    intermediate_state_indices=None,
+    cache_steps=None,
+    retrieve_parent_token=None,
+):
+    """Fused: fused_sigmoid_gating_delta_rule_update."""
+    return fused_sigmoid_gating_delta_rule_update(
+        A_log=A_log,
+        dt_bias=dt_bias,
+        q=q,
+        k=k,
+        v=v,
+        a=a,
+        b=b,
+        initial_state_source=initial_state_source,
+        initial_state_indices=initial_state_indices,
+        cu_seqlens=cu_seqlens,
+        use_qk_l2norm_in_kernel=True,
+        softplus_beta=1.0,
+        softplus_threshold=20.0,
+        is_kda=False,
+        disable_state_update=disable_state_update,
+        intermediate_states_buffer=intermediate_states_buffer,
+        intermediate_state_indices=intermediate_state_indices,
+        cache_steps=cache_steps,
+        retrieve_parent_token=retrieve_parent_token,
+    )
+
+
+@pytest.mark.skipif(not KERNELS_AVAILABLE, reason="Kernel not available")
+@pytest.mark.parametrize("N", [1, 8, 16])
+@pytest.mark.parametrize("T", [1, 4, 8])
+def test_fused_gdn_mtp_precision(N: int, T: int):
+    """Compare fused MTP output against reference."""
+    H, HV, K, V = 16, 32, 128, 128
+
+    A_log, dt_bias, a, b, q, k, v, state, indices, cu_seqlens = _make_tensors(
+        N, T, H, HV, K, V
+    )
+
+    state_ref = state.clone()
+    state_fused = state.clone()
+
+    out_ref = run_reference(
+        A_log,
+        dt_bias,
+        q,
+        k,
+        v,
+        a,
+        b,
+        state_ref,
+        indices,
+        cu_seqlens,
+        disable_state_update=True,
+    )
+    out_fused = run_fused_mtp(
+        A_log,
+        dt_bias,
+        q,
+        k,
+        v,
+        a,
+        b,
+        state_fused,
+        indices,
+        cu_seqlens,
+        disable_state_update=True,
+    )
+
+    torch.testing.assert_close(out_ref, out_fused, rtol=1e-2, atol=1e-2)
+
+
+@pytest.mark.skipif(not KERNELS_AVAILABLE, reason="Kernels not available")
+@pytest.mark.parametrize("N", [1, 16, 128])
+def test_mtp_single_step_decode(N: int):
+    """Verify MTP kernel matches reference for T=1 (decode scenario)."""
+    T = 1
+    H, HV, K, V = 16, 32, 128, 128
+
+    A_log, dt_bias, a, b, q, k, v, state, indices, cu_seqlens = _make_tensors(
+        N, T, H, HV, K, V
+    )
+
+    state_ref = state.clone()
+    state_fused = state.clone()
+
+    out_ref = run_reference(
+        A_log,
+        dt_bias,
+        q,
+        k,
+        v,
+        a,
+        b,
+        state_ref,
+        indices,
+        cu_seqlens,
+        disable_state_update=False,
+    )
+    out_fused = run_fused_mtp(
+        A_log,
+        dt_bias,
+        q,
+        k,
+        v,
+        a,
+        b,
+        state_fused,
+        indices,
+        cu_seqlens,
+        disable_state_update=False,
+    )
+
+    torch.testing.assert_close(out_ref, out_fused, rtol=1e-2, atol=1e-2)
+
+    # Also verify states match after update
+    state_diff = (state_ref.float() - state_fused.float()).abs()
+    state_max_diff = state_diff.max().item()
+    state_fail_rate = (state_diff > 0.1).float().mean().item() * 100
+    print(
+        f"  single_step state N={N}: max_diff={state_max_diff:.2e}, "
+        f"fail_rate={state_fail_rate:.2f}%"
+    )
+    assert state_fail_rate < 0.01, f"State mismatch: fail_rate={state_fail_rate:.2f}%"
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v", "-s"])

python/sglang/srt/layers/attention/fla/fused_sigmoid_gating_recurrent.py

@@ -20,12 +20,21 @@ def fused_sigmoid_gating_delta_rule_update_kernel(
     h0_source,
     h0_indices,
     cu_seqlens,
+    # Parameters for target_verify support (unused for decode)
+    intermediate_states_buffer,
+    intermediate_state_indices,
+    cache_steps,
+    retrieve_parent_token_ptr,
+    stride_retrieve_parent_token_seq: tl.constexpr,
+    stride_retrieve_parent_token_token: tl.constexpr,
+    # ================================================
     scale,
     T,
     stride_q,
     stride_k,
     stride_v,
     stride_b,
+    NP2_T: tl.constexpr,
     B: tl.constexpr,
     H: tl.constexpr,
     HV: tl.constexpr,
@@ -37,6 +46,10 @@ def fused_sigmoid_gating_delta_rule_update_kernel(
     USE_QK_L2NORM_IN_KERNEL: tl.constexpr,
     IS_VARLEN: tl.constexpr,
     IS_KDA: tl.constexpr,
+    # Optional flags for target_verify support (default False for decode)
+    DISABLE_STATE_UPDATE: tl.constexpr = False,
+    CACHE_INTERMEDIATE_STATES: tl.constexpr = False,
+    HAS_EAGLE_TREE_CUSTOM_ATTN_MASK: tl.constexpr = False,
 ):
     """
     Fused kernel that combines sigmoid gating computation with recurrent delta rule update.
@@ -91,7 +104,44 @@ def fused_sigmoid_gating_delta_rule_update_kernel(
             )
             b_h += tl.load(p_h0, mask=mask_h, other=0).to(tl.float32)
 
+    # Preload tree attention data if needed
+    if HAS_EAGLE_TREE_CUSTOM_ATTN_MASK:
+        token_indices = tl.arange(0, NP2_T)
+        mask_retrieve = token_indices < T
+        retrieve_parent_token_base = (
+            retrieve_parent_token_ptr
+            + (i_n * stride_retrieve_parent_token_seq)
+            + token_indices * stride_retrieve_parent_token_token
+        )
+        parent_idx_tokens = tl.load(
+            retrieve_parent_token_base, mask=mask_retrieve, other=0
+        )
+
+    # Prepare intermediate state cache index if enabled
+    cache_idx = -1
+    if CACHE_INTERMEDIATE_STATES:
+        cache_idx = tl.load(intermediate_state_indices + i_n)
+
+    step_idx = 0
     for _ in range(0, T):
+        # Tree attention: load parent's cached state
+        if HAS_EAGLE_TREE_CUSTOM_ATTN_MASK:
+            # step_idx == 0 uses b_h from USE_INITIAL_STATE
+            if step_idx != 0 and cache_idx >= 0:
+                parent_step_idx = tl.sum(
+                    tl.where(token_indices == step_idx, parent_idx_tokens, 0)
+                )
+                step_offset = parent_step_idx * HV * K * V
+                cache_ptr = (
+                    intermediate_states_buffer
+                    + cache_idx * cache_steps * HV * K * V
+                    + step_offset
+                    + i_hv * K * V
+                    + o_k[:, None] * V
+                    + o_v[None, :]
+                )
+                b_h = tl.load(cache_ptr, mask=mask_h, other=0).to(tl.float32)
+
         # Load inputs
         b_q = tl.load(p_q, mask=mask_k, other=0).to(tl.float32)
         b_k = tl.load(p_k, mask=mask_k, other=0).to(tl.float32)
@@ -101,6 +151,10 @@ def fused_sigmoid_gating_delta_rule_update_kernel(
         # Compute sigmoid gating
         # Load gating parameters
         b_A_log = tl.load(p_A_log).to(tl.float32)
+        if IS_KDA:
+            b_a = tl.load(p_a, mask=mask_k, other=0).to(tl.float32)
+            b_dt_bias = tl.load(p_dt_bias, mask=mask_k, other=0).to(tl.float32)
+        else:
             b_a = tl.load(p_a).to(tl.float32)
             b_dt_bias = tl.load(p_dt_bias).to(tl.float32)
 
@@ -144,15 +198,35 @@ def fused_sigmoid_gating_delta_rule_update_kernel(
         b_o = tl.sum(b_h * b_q[:, None], 0)
         tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=mask_v)
 
+        # Cache intermediate states if enabled
+        if CACHE_INTERMEDIATE_STATES:
+            if cache_idx >= 0:
+                step_offset = step_idx * HV * K * V
+                cache_ptr = (
+                    intermediate_states_buffer
+                    + cache_idx * cache_steps * HV * K * V
+                    + step_offset
+                    + i_hv * K * V
+                    + o_k[:, None] * V
+                    + o_v[None, :]
+                )
+                tl.store(cache_ptr, b_h.to(cache_ptr.dtype.element_ty), mask=mask_h)
+
+        step_idx += 1
+
         # Update pointers for next timestep
-        p_q += H * K
+        p_q += stride_q
-        p_k += H * K
+        p_k += stride_k
+        p_v += stride_v
+        p_b += stride_b
         p_o += HV * V
-        p_v += HV * V
+        if IS_KDA:
-        p_b += HV
+            p_a += HV * K
+        else:
             p_a += HV
 
     # Store final state back to h0_source with bounds checking
+    if not DISABLE_STATE_UPDATE:
         if USE_INITIAL_STATE:
             idx = tl.load(h0_indices + i_n)
             if idx >= 0:
@@ -182,11 +256,22 @@ def fused_sigmoid_gating_delta_rule_update(
     use_qk_l2norm_in_kernel: bool = False,
     cu_seqlens: Optional[torch.Tensor] = None,
     is_kda: bool = False,
+    # Optional parameters for target_verify support
+    disable_state_update: bool = False,
+    intermediate_states_buffer: Optional[torch.Tensor] = None,
+    intermediate_state_indices: Optional[torch.Tensor] = None,
+    cache_steps: Optional[int] = None,
+    retrieve_parent_token: Optional[torch.Tensor] = None,
 ):
     """
     Fused triton implementation of sigmoid gating delta rule update.
     This function uses a single fused kernel that combines both sigmoid gating computation
     and the recurrent delta rule update for better performance.
+
+    Supports both decode and target_verify modes:
+    - decode: standard single-step update with state write-back
+    - target_verify: multi-step with intermediate state caching, optional tree attention,
+                     and optional state update disable
     """
     B, T, H, K, V = *k.shape, v.shape[-1]
     stride_q = q.stride()[1]
@@ -207,6 +292,17 @@ def fused_sigmoid_gating_delta_rule_update(
         assert scale > 0, "scale must be positive"
 
     o = q.new_empty(NK, *v.shape)
+
+    # Prepare retrieve_parent_token strides
+    if retrieve_parent_token is not None:
+        stride_retrieve_parent_token_seq = retrieve_parent_token.stride(0)
+        stride_retrieve_parent_token_token = retrieve_parent_token.stride(1)
+    else:
+        stride_retrieve_parent_token_seq = 0
+        stride_retrieve_parent_token_token = 0
+
+    NP2_T = triton.next_power_of_2(T)
+
     grid = (NK, NV, N * HV)
 
     fused_sigmoid_gating_delta_rule_update_kernel[grid](
@@ -223,12 +319,19 @@ def fused_sigmoid_gating_delta_rule_update(
         h0_source=initial_state_source,
         h0_indices=initial_state_indices,
         cu_seqlens=cu_seqlens,
+        intermediate_states_buffer=intermediate_states_buffer,
+        intermediate_state_indices=intermediate_state_indices,
+        cache_steps=0 if cache_steps is None else cache_steps,
+        retrieve_parent_token_ptr=retrieve_parent_token,
+        stride_retrieve_parent_token_seq=stride_retrieve_parent_token_seq,
+        stride_retrieve_parent_token_token=stride_retrieve_parent_token_token,
         scale=scale,
         T=T,
         stride_q=stride_q,
         stride_k=stride_k,
         stride_v=stride_v,
         stride_b=stride_b,
+        NP2_T=NP2_T,
         B=B,
         H=H,
         HV=HV,
@@ -240,6 +343,9 @@ def fused_sigmoid_gating_delta_rule_update(
         USE_QK_L2NORM_IN_KERNEL=use_qk_l2norm_in_kernel,
         IS_VARLEN=cu_seqlens is not None,
         IS_KDA=is_kda,
+        DISABLE_STATE_UPDATE=disable_state_update,
+        CACHE_INTERMEDIATE_STATES=intermediate_states_buffer is not None,
+        HAS_EAGLE_TREE_CUSTOM_ATTN_MASK=retrieve_parent_token is not None,
         num_warps=num_warps,
         num_stages=num_stages,
     )

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

@@ -171,11 +171,13 @@ class GDNKernelDispatcher:
 
     def target_verify(
         self,
+        A_log: torch.Tensor,
+        dt_bias: torch.Tensor,
         q: torch.Tensor,
         k: torch.Tensor,
         v: torch.Tensor,
-        g: torch.Tensor,
+        a: torch.Tensor,
-        beta: torch.Tensor,
+        b: torch.Tensor,
         *,
         ssm_states: torch.Tensor,
         cache_indices: torch.Tensor,
@@ -183,11 +185,13 @@ class GDNKernelDispatcher:
         **kwargs,
     ) -> torch.Tensor:
         return self.verify_kernel.target_verify(
-            q,
+            A_log=A_log,
-            k,
+            dt_bias=dt_bias,
-            v,
+            q=q,
-            g,
+            k=k,
-            beta,
+            v=v,
+            a=a,
+            b=b,
             ssm_states=ssm_states,
             cache_indices=cache_indices,
             query_start_loc=query_start_loc,
@@ -364,15 +368,15 @@ class GDNAttnBackend(MambaAttnBackendBase):
         key = key.view(1, actual_seq_len, layer.num_k_heads, layer.head_k_dim)
         value = value.view(1, actual_seq_len, layer.num_v_heads, layer.head_v_dim)
 
-        g, beta = fused_gdn_gating(layer.A_log, a, b, layer.dt_bias)
-
         if is_target_verify:
             core_attn_out = self.kernel_dispatcher.target_verify(
+                A_log=layer.A_log,
+                dt_bias=layer.dt_bias,
                 q=query,
                 k=key,
                 v=value,
-                g=g,
+                a=a,
-                beta=beta,
+                b=b,
                 ssm_states=ssm_states,
                 cache_indices=cache_indices,
                 query_start_loc=query_start_loc,
@@ -380,13 +384,9 @@ class GDNAttnBackend(MambaAttnBackendBase):
                 intermediate_state_indices=intermediate_state_indices,
                 cache_steps=forward_batch.spec_info.draft_token_num,
                 retrieve_parent_token=retrieve_parent_token,
-                # Pass raw pre-gating values for FlashInfer MTP kernel
-                a_raw=a,
-                b_raw=b,
-                A_log=layer.A_log,
-                dt_bias=layer.dt_bias,
             )
         else:
+            g, beta = fused_gdn_gating(layer.A_log, a, b, layer.dt_bias)
             core_attn_out, last_recurrent_state, h = self.kernel_dispatcher.extend(
                 q=query,
                 k=key,

python/sglang/srt/layers/attention/linear/kernels/gdn_flashinfer.py

@@ -251,11 +251,13 @@ class FlashInferGDNKernel(LinearAttnKernelBase):
 
     def target_verify(
         self,
+        A_log: torch.Tensor,
+        dt_bias: torch.Tensor,
         q: torch.Tensor,
         k: torch.Tensor,
         v: torch.Tensor,
-        g: torch.Tensor,
+        a: torch.Tensor,
-        beta: torch.Tensor,
+        b: torch.Tensor,
         *,
         ssm_states: torch.Tensor,
         cache_indices: torch.Tensor,
@@ -293,22 +295,14 @@ class FlashInferGDNKernel(LinearAttnKernelBase):
         value_mtp = v.view(batch_size, draft_token_num, num_v_heads, head_v_dim)
 
         # a, b from g/beta: [1, seq, HV] -> [B, T, HV]
-        # But the MTP kernel expects raw a, b (pre-gating), not g, beta.
+        if a is None or b is None or A_log is None or dt_bias is None:
-        # We need to recover a and b from the gdn_backend caller.
-        # The caller passes them via **kwargs from the dispatcher.
-        a_raw = kwargs.get("a_raw")
-        b_raw = kwargs.get("b_raw")
-        A_log = kwargs.get("A_log")
-        dt_bias = kwargs.get("dt_bias")
-
-        if a_raw is None or b_raw is None or A_log is None or dt_bias is None:
             raise RuntimeError(
                 "FlashInfer GDN MTP kernel requires a_raw, b_raw, A_log, "
                 "dt_bias to be passed via kwargs."
             )
 
-        a_mtp = a_raw.view(batch_size, draft_token_num, num_v_heads)
+        a_mtp = a.view(batch_size, draft_token_num, num_v_heads)
-        b_mtp = b_raw.view(batch_size, draft_token_num, num_v_heads)
+        b_mtp = b.view(batch_size, draft_token_num, num_v_heads)
 
         output_fi, _ = self._mtp_fn(
             q=query_mtp,

python/sglang/srt/layers/attention/linear/kernels/gdn_triton.py

@@ -7,9 +7,6 @@ from sglang.srt.utils import is_cpu, is_npu
 
 if not is_cpu():
     from sglang.srt.layers.attention.fla.chunk import chunk_gated_delta_rule
-    from sglang.srt.layers.attention.fla.fused_recurrent import (
-        fused_recurrent_gated_delta_rule_update,
-    )
     from sglang.srt.layers.attention.fla.fused_sigmoid_gating_recurrent import (
         fused_sigmoid_gating_delta_rule_update,
     )
@@ -98,11 +95,13 @@ class TritonGDNKernel(LinearAttnKernelBase):
 
     def target_verify(
         self,
+        A_log: torch.Tensor,
+        dt_bias: torch.Tensor,
         q: torch.Tensor,
         k: torch.Tensor,
         v: torch.Tensor,
-        g: torch.Tensor,
+        a: torch.Tensor,
-        beta: torch.Tensor,
+        b: torch.Tensor,
         *,
         ssm_states: torch.Tensor,
         cache_indices: torch.Tensor,
@@ -113,16 +112,22 @@ class TritonGDNKernel(LinearAttnKernelBase):
         retrieve_parent_token: torch.Tensor,
         **kwargs,
     ) -> torch.Tensor:
-        return fused_recurrent_gated_delta_rule_update(
+        return fused_sigmoid_gating_delta_rule_update(
+            A_log=A_log,
+            dt_bias=dt_bias,
             q=q,
             k=k,
             v=v,
-            g=g,
+            a=a,
-            beta=beta,
+            b=b,
             initial_state_source=ssm_states,
             initial_state_indices=cache_indices,
             cu_seqlens=query_start_loc,
             use_qk_l2norm_in_kernel=True,
+            softplus_beta=1.0,
+            softplus_threshold=20.0,
+            is_kda=False,
+            # target_verify specific parameters
             disable_state_update=True,
             intermediate_states_buffer=intermediate_states_buffer,
             intermediate_state_indices=intermediate_state_indices,

python/sglang/srt/layers/attention/linear/kernels/kernel_backend.py

@@ -44,11 +44,13 @@ class LinearAttnKernelBase(ABC):
 
     def target_verify(
         self,
+        A_log: torch.Tensor,
+        dt_bias: torch.Tensor,
         q: torch.Tensor,
         k: torch.Tensor,
         v: torch.Tensor,
-        g: torch.Tensor,
+        a: torch.Tensor,
-        beta: torch.Tensor,
+        b: torch.Tensor,
         *,
         ssm_states: torch.Tensor,
         cache_indices: torch.Tensor,