From f7de9375ace71f46882d095bb3fc23f706bad660 Mon Sep 17 00:00:00 2001 From: Yuan Luo Date: Fri, 6 Mar 2026 21:42:44 +0800 Subject: [PATCH] [GDN][Qwen3-Next][Qwen3.5] Fuse fused_gdn_gating and fused_recurrent_gated_delta_rule_update in verify_target (#19775) --- .../tests/test_fused_verify_triton_gdn.py | 231 ++++++++++++++++++ .../fla/fused_sigmoid_gating_recurrent.py | 142 +++++++++-- .../layers/attention/linear/gdn_backend.py | 32 +-- .../linear/kernels/gdn_flashinfer.py | 20 +- .../attention/linear/kernels/gdn_triton.py | 21 +- .../linear/kernels/kernel_backend.py | 6 +- 6 files changed, 395 insertions(+), 57 deletions(-) create mode 100644 python/sglang/jit_kernel/tests/test_fused_verify_triton_gdn.py diff --git a/python/sglang/jit_kernel/tests/test_fused_verify_triton_gdn.py b/python/sglang/jit_kernel/tests/test_fused_verify_triton_gdn.py new file mode 100644 index 000000000..a6e048a40 --- /dev/null +++ b/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"]) diff --git a/python/sglang/srt/layers/attention/fla/fused_sigmoid_gating_recurrent.py b/python/sglang/srt/layers/attention/fla/fused_sigmoid_gating_recurrent.py index 486537b0e..f140ccae4 100644 --- a/python/sglang/srt/layers/attention/fla/fused_sigmoid_gating_recurrent.py +++ b/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,8 +151,12 @@ 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) - b_a = tl.load(p_a).to(tl.float32) - b_dt_bias = tl.load(p_dt_bias).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) # Compute g = -exp(A_log) * softplus(a + dt_bias) x = b_a + b_dt_bias @@ -144,26 +198,46 @@ 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_k += H * K + p_q += stride_q + p_k += stride_k + p_v += stride_v + p_b += stride_b p_o += HV * V - p_v += HV * V - p_b += HV - p_a += HV + if IS_KDA: + p_a += HV * K + else: + p_a += HV # Store final state back to h0_source with bounds checking - if USE_INITIAL_STATE: - idx = tl.load(h0_indices + i_n) - if idx >= 0: - p_h0 = ( - h0_source - + idx * HV * K * V - + i_hv * K * V - + o_k[:, None] * V - + o_v[None, :] - ) - tl.store(p_h0, b_h.to(p_h0.dtype.element_ty), mask=mask_h) + if not DISABLE_STATE_UPDATE: + if USE_INITIAL_STATE: + idx = tl.load(h0_indices + i_n) + if idx >= 0: + p_h0 = ( + h0_source + + idx * HV * K * V + + i_hv * K * V + + o_k[:, None] * V + + o_v[None, :] + ) + tl.store(p_h0, b_h.to(p_h0.dtype.element_ty), mask=mask_h) def fused_sigmoid_gating_delta_rule_update( @@ -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, ) diff --git a/python/sglang/srt/layers/attention/linear/gdn_backend.py b/python/sglang/srt/layers/attention/linear/gdn_backend.py index 3725456fe..8b7851dfa 100644 --- a/python/sglang/srt/layers/attention/linear/gdn_backend.py +++ b/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, - beta: torch.Tensor, + a: 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, - k, - v, - g, - beta, + A_log=A_log, + dt_bias=dt_bias, + q=q, + k=k, + 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, - beta=beta, + a=a, + 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, diff --git a/python/sglang/srt/layers/attention/linear/kernels/gdn_flashinfer.py b/python/sglang/srt/layers/attention/linear/kernels/gdn_flashinfer.py index 8ad0fc5a0..32ffd9664 100644 --- a/python/sglang/srt/layers/attention/linear/kernels/gdn_flashinfer.py +++ b/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, - beta: torch.Tensor, + a: 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. - # 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: + if a is None or b 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) - b_mtp = b_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.view(batch_size, draft_token_num, num_v_heads) output_fi, _ = self._mtp_fn( q=query_mtp, diff --git a/python/sglang/srt/layers/attention/linear/kernels/gdn_triton.py b/python/sglang/srt/layers/attention/linear/kernels/gdn_triton.py index a1cf16bba..106113cda 100644 --- a/python/sglang/srt/layers/attention/linear/kernels/gdn_triton.py +++ b/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, - beta: torch.Tensor, + a: 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, - beta=beta, + a=a, + 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, diff --git a/python/sglang/srt/layers/attention/linear/kernels/kernel_backend.py b/python/sglang/srt/layers/attention/linear/kernels/kernel_backend.py index 83a7fe0e7..5e8018699 100644 --- a/python/sglang/srt/layers/attention/linear/kernels/kernel_backend.py +++ b/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, - beta: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, *, ssm_states: torch.Tensor, cache_indices: torch.Tensor,