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

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,