chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Yuan Luo
2026-03-06 21:42:44 +08:00
committed by GitHub
parent e3b581ce6b
commit f7de9375ac
6 changed files with 395 additions and 57 deletions
Download Patch File Download Diff File Expand all files Collapse all files

231
python/sglang/jit_kernel/tests/test_fused_verify_triton_gdn.py Normal file
View File

@@ -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"])

142
python/sglang/srt/layers/attention/fla/fused_sigmoid_gating_recurrent.py
View File

@@ -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,
)

32
python/sglang/srt/layers/attention/linear/gdn_backend.py
View File

@@ -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,

20
python/sglang/srt/layers/attention/linear/kernels/gdn_flashinfer.py
View File

@@ -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,

21
python/sglang/srt/layers/attention/linear/kernels/gdn_triton.py
View File

@@ -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,

6
python/sglang/srt/layers/attention/linear/kernels/kernel_backend.py
View File

@@ -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,