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

[KDA] Support ReplaySSM ring-write in the fused chain-verify kernel (#36821)

Browse Source 
Co-authored-by: luoyuan.luo <luoyuan.luo@antgroup.com>
This commit is contained in:
Yuan Luo
2026-09-15 10:22:41 +08:00
committed by GitHub
co-authored by luoyuan.luo
parent a23fd557ed
commit 99060191e7
5 changed files with 935 additions and 50 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/kernels/ops/attention/fla/fused_kda_conv_recurrent_verify.py
+127 -13
View File
@@ -16,11 +16,17 @@ kernels. Requires ``T >= kernel_width - 1`` (the rolled conv state is then
exactly the last ``kernel_width - 1`` input tokens, matching the reference
kernel's store).
Numerics: deliberately bit-aligned with the unfused pair. The conv output is
rounded to the activation dtype (bf16) before entering the recurrence —
exactly what the unfused path does through its intermediate tensor — and all
expressions mirror the reference kernels line by line, with the same
num_warps so reduction order matches.
ReplaySSM (``cache_ring``): instead of per-step [HV, V, K] fp32 state
snapshots, stash each step's raw inputs (pre-l2norm k, pre-delta v, gate,
beta) into the per-slot rings the commit-time exact fold replays
(kda_replayssm_spec_decode.py) -- same CACHE_RING contract as the unfused
fused_sigmoid_gating_delta_rule_update, so the two paths fill identical rings.
Numerics: aligned with the unfused pair. The conv output is rounded to the
activation dtype (bf16) before entering the recurrence — exactly what the
unfused path does through its intermediate tensor — and all expressions mirror
the reference kernels line by line. Reduction order still splits differently
where many V heads share one Q/K head, worth ~1 ulp on the output.
"""
from typing import Optional
@@ -83,6 +89,18 @@ def fused_kda_conv_gating_verify_kernel(
SAVE_INTERMEDIATE_WINDOW: tl.constexpr,
CACHE_INTERMEDIATE_STATES: tl.constexpr,
USE_GDC: tl.constexpr = False,
# ReplaySSM fused ring-write (spec verify): per-slot rings consumed by the
# commit-time exact fold (kda_replayssm_spec_decode.py). Off -> dead code.
replayssm_rawv=None, # [slots, HV, L, V] activation dtype
replayssm_rawk=None, # [slots, H, L, K] activation dtype
replayssm_g=None, # [slots, HV, L, K] fp32
replayssm_beta=None, # [slots, HV, L] fp32
stride_rawv_slot: tl.constexpr = 0,
stride_rawk_slot: tl.constexpr = 0,
stride_g_slot: tl.constexpr = 0,
stride_beta_slot: tl.constexpr = 0,
MAX_CACHE_LEN: tl.constexpr = 0,
CACHE_RING: tl.constexpr = False,
):
# PDL: overlap prologue with the tail of the producer qkv-projection GEMM;
# every global load (conv_state_indices, mixed_qkv, weights) happens after
@@ -289,6 +307,53 @@ def fused_kda_conv_gating_verify_kernel(
b_beta = 1.0 / (1.0 + tl.exp(-b_b))
# ReplaySSM ring-write. Must sit here: b_k still pre-l2norm, b_v still
# pre-delta, b_g/b_beta formed -- so the commit fold's replay is
# bit-identical to the update below (mirrors the CACHE_RING block in
# fused_sigmoid_gating_recurrent.py). rawk dedups via is_qk_owner
# (per k-head); g/beta write once per v-head at i_v == 0. The
# t < MAX_CACHE_LEN guard drops absorb-overflow steps instead of
# smashing the next slot's ring.
if CACHE_RING:
if h0_idx >= 0 and t < MAX_CACHE_LEN:
ring_slot = h0_idx.to(tl.int64)
tl.store(
replayssm_rawv
+ ring_slot * stride_rawv_slot
+ i_hv * MAX_CACHE_LEN * V
+ t * V
+ o_v,
b_v.to(replayssm_rawv.dtype.element_ty),
mask=mask_v,
)
if is_qk_owner:
tl.store(
replayssm_rawk
+ ring_slot * stride_rawk_slot
+ i_h * MAX_CACHE_LEN * K
+ t * K
+ o_k,
b_k.to(replayssm_rawk.dtype.element_ty),
mask=mask_k,
)
if i_v == 0:
tl.store(
replayssm_g
+ ring_slot * stride_g_slot
+ i_hv * MAX_CACHE_LEN * K
+ t * K
+ o_k,
b_g,
mask=mask_k,
)
tl.store(
replayssm_beta
+ ring_slot * stride_beta_slot
+ i_hv * MAX_CACHE_LEN
+ t,
b_beta,
)
if USE_QK_L2NORM_IN_KERNEL:
b_q = b_q / (tl.sqrt(tl.sum(b_q * b_q) + 1e-6))
b_k = b_k / (tl.sqrt(tl.sum(b_k * b_k) + 1e-6))
@@ -358,15 +423,23 @@ def fused_kda_conv_gating_verify(
softplus_beta: float = 1.0,
softplus_threshold: float = 20.0,
use_qk_l2norm_in_kernel: bool = True,
# num_warps=4 is ~1.3x faster than the unfused pair in-graph; the output,
# conv_state and conv-window caches stay bit-identical to the reference.
# Only the fp32 intermediate-ssm rollback cache differs: the tl.sum
# reduction-order delta (~1 ulp/step) compounds through the delta-rule
# recurrence — measured ~6e-8 at T=4 standard gate (the production MTP
# shape), ~1.5e-5 at T=4 safe gate, ~2e-3 at T=8 safe gate. num_warps=1
# reproduces the reference reduction order exactly (all buffers
# bit-identical) but is ~2.4x slower in-graph — numerics debugging only.
# num_warps=4 is ~1.3x faster than the unfused pair in-graph; conv_state
# and the conv-window cache stay bit-identical to the reference, the bf16
# output within one ulp (the BV=4 tile reduces K in a different order).
# The fp32 intermediate-ssm rollback cache carries that ~1 ulp/step delta
# through the delta-rule recurrence — measured ~6e-8 at T=4 standard gate
# (the production MTP shape), ~1.5e-5 at T=4 safe gate, ~2e-3 at T=8 safe
# gate. num_warps=1 is ~2.4x slower in-graph — numerics debugging only.
# The ReplaySSM ring values are bit-exact at any num_warps: they are
# elementwise (conv FMA chain, gate, sigmoid), upstream of every tl.sum.
num_warps: int = 4,
# ReplaySSM fused ring-write; same parameter names as the unfused
# fused_sigmoid_gating_delta_rule_update so ring_kwargs pass through both.
cache_ring: bool = False,
replayssm_rawv: Optional[torch.Tensor] = None,
replayssm_rawk: Optional[torch.Tensor] = None,
replayssm_g: Optional[torch.Tensor] = None,
replayssm_beta: Optional[torch.Tensor] = None,
) -> torch.Tensor:
"""Chain-verify fast path. Returns ``o`` of shape [1, seq_len, HV, V],
matching the unfused ``target_verify`` output layout."""
@@ -420,6 +493,37 @@ def fused_kda_conv_gating_verify(
if intermediate_states_buffer is not None:
assert intermediate_states_buffer.is_contiguous()
if cache_ring:
# Per-layer ring views (memory_pool.py KDA spec rings). The kernel uses
# stride(0) as the slot pitch and packs within a slot from
# MAX_CACHE_LEN and the head/dim extents, so inner dims must be packed.
assert (
replayssm_rawv is not None
and replayssm_rawk is not None
and replayssm_g is not None
and replayssm_beta is not None
), "cache_ring requires all four replayssm_* rings"
max_cache_len = replayssm_rawv.shape[-2]
assert tuple(replayssm_rawv.shape[1:]) == (HV, max_cache_len, V)
assert tuple(replayssm_rawk.shape[1:]) == (H, max_cache_len, K)
assert tuple(replayssm_g.shape[1:]) == (HV, max_cache_len, K)
assert tuple(replayssm_beta.shape[1:]) == (HV, max_cache_len)
assert replayssm_rawv.stride()[1:] == (max_cache_len * V, V, 1)
assert replayssm_rawk.stride()[1:] == (max_cache_len * K, K, 1)
assert replayssm_g.stride()[1:] == (max_cache_len * K, K, 1)
assert replayssm_beta.stride()[1:] == (max_cache_len, 1)
assert replayssm_rawv.dtype == mixed_qkv.dtype
assert replayssm_rawk.dtype == mixed_qkv.dtype
assert replayssm_g.dtype == torch.float32
assert replayssm_beta.dtype == torch.float32
stride_rawv_slot = replayssm_rawv.stride(0)
stride_rawk_slot = replayssm_rawk.stride(0)
stride_g_slot = replayssm_g.stride(0)
stride_beta_slot = replayssm_beta.stride(0)
else:
max_cache_len = 0
stride_rawv_slot = stride_rawk_slot = stride_g_slot = stride_beta_slot = 0
grid = (NV, B * HV)
# PDL (sm90+): chain behind the producer qkv-projection GEMM and signal the
# downstream o_norm / o_proj. Scheduling only — bit-exactness unaffected.
@@ -480,6 +584,16 @@ def fused_kda_conv_gating_verify(
USE_LOWER_BOUND=lower_bound is not None,
SAVE_INTERMEDIATE_WINDOW=intermediate_conv_window is not None,
CACHE_INTERMEDIATE_STATES=intermediate_states_buffer is not None,
replayssm_rawv=replayssm_rawv,
replayssm_rawk=replayssm_rawk,
replayssm_g=replayssm_g,
replayssm_beta=replayssm_beta,
stride_rawv_slot=stride_rawv_slot,
stride_rawk_slot=stride_rawk_slot,
stride_g_slot=stride_g_slot,
stride_beta_slot=stride_beta_slot,
MAX_CACHE_LEN=max_cache_len,
CACHE_RING=cache_ring,
# num_warps=1 matches the reference kernels' reduction order exactly;
# higher values must be re-validated for bit-exactness before use.
num_warps=num_warps,
python/sglang/srt/layers/attention/linear/kda_backend.py
+127 -24
View File
@@ -39,6 +39,7 @@ from sglang.srt.runtime_context import (
get_disagg,
get_exec,
get_memory,
get_platform,
get_spec,
)
@@ -988,6 +989,27 @@ class KDAAttnBackend(MambaAttnBackendBase):
"KDA target_verify requires a speculative mamba cache "
"(MambaPool.SpeculativeState); none found."
)
# ReplaySSM: the ring-write is fused into the verify kernel
# (CACHE_RING) on both the fused chain-verify and the unfused triton
# paths; commit replays the ring instead of reading per-step
# snapshots. ring_kwargs stays empty for non-triton verify kernels,
# which never see replayssm. Ragged layouts work natively on the
# unfused path -- step_idx is the within-row step under varlen, so
# row i writes ring[slot][0..verify_lens[i]) and commit folds at most
# commit_lens of them (absorb overflow is bounded in-kernel).
replayssm_rawk = replayssm_g = replayssm_beta = None
ring_kwargs = {}
if replayssm_on:
replayssm_rawk = mamba_cache_params.replayssm_rawk
replayssm_g = mamba_cache_params.replayssm_g
replayssm_beta = mamba_cache_params.replayssm_beta
ring_kwargs = dict(
cache_ring=True,
replayssm_rawv=replayssm_rawv,
replayssm_rawk=replayssm_rawk,
replayssm_g=replayssm_g,
replayssm_beta=replayssm_beta,
)
intermediate_conv_window_cache = mamba_cache_params.intermediate_conv_window[0]
intermediate_state_indices = self.verify_intermediate_state_indices
@@ -1047,6 +1069,9 @@ class KDAAttnBackend(MambaAttnBackendBase):
retrieve_next_sibling=retrieve_next_sibling,
retrieve_parent_token=retrieve_parent_token,
replayssm_rawv=replayssm_rawv,
replayssm_rawk=replayssm_rawk,
replayssm_g=replayssm_g,
replayssm_beta=replayssm_beta,
):
return self._fused_chain_verify_fn(
mixed_qkv=mixed_qkv,
@@ -1077,6 +1102,7 @@ class KDAAttnBackend(MambaAttnBackendBase):
head_k_dim=layer.head_k_dim,
head_v_dim=layer.head_v_dim,
lower_bound=layer.lower_bound,
**ring_kwargs,
)
dense_token_indices = None
mixed_qkv_dense = mixed_qkv.view(batch_size, draft_token_num, -1)
@@ -1135,22 +1161,6 @@ class KDAAttnBackend(MambaAttnBackendBase):
k = k.unflatten(-1, (-1, layer.head_k_dim)).unsqueeze(0)
v = v.unflatten(-1, (-1, layer.head_v_dim)).unsqueeze(0)
# ReplaySSM: the ring-write is fused into the triton verify kernel
# (CACHE_RING). Ragged layouts work natively -- step_idx is the
# within-row step under varlen, so row i writes
# ring[slot][0..verify_lens[i]) and commit folds at most commit_lens
# of them (absorb overflow is bounded in-kernel). ring_kwargs stays
# empty for non-triton verify kernels, which never see replayssm.
ring_kwargs = {}
if replayssm_rawv is not None:
ring_kwargs = dict(
cache_ring=True,
replayssm_rawv=replayssm_rawv,
replayssm_rawk=mamba_cache_params.replayssm_rawk,
replayssm_g=mamba_cache_params.replayssm_g,
replayssm_beta=mamba_cache_params.replayssm_beta,
)
core_attn_out = self.kernel_dispatcher.target_verify(
A_log=layer.A_log,
dt_bias=layer.dt_bias,
@@ -1210,10 +1220,13 @@ class KDAAttnBackend(MambaAttnBackendBase):
retrieve_next_sibling: Optional[torch.Tensor],
retrieve_parent_token: Optional[torch.Tensor],
replayssm_rawv: Optional[torch.Tensor],
replayssm_rawk: Optional[torch.Tensor],
replayssm_g: Optional[torch.Tensor],
replayssm_beta: Optional[torch.Tensor],
) -> bool:
if self._fused_chain_verify_fn is None or not mixed_qkv.is_cuda:
return False
if replayssm_rawv is not None or any(
if any(
value is not None
for value in (
retrieve_next_token,
@@ -1222,6 +1235,7 @@ class KDAAttnBackend(MambaAttnBackendBase):
)
):
return False
replayssm_on = replayssm_rawv is not None
if draft_token_num < 3 or mixed_qkv.shape[0] % draft_token_num != 0:
return False
if (
@@ -1240,6 +1254,15 @@ class KDAAttnBackend(MambaAttnBackendBase):
seq_len, dim = mixed_qkv.shape
batch_size = seq_len // draft_token_num
if replayssm_on and (
batch_size != 1 or not (get_platform().is_sm90 or get_platform().is_sm100)
):
# The runtime still uses BV=4, not the benchmark's best-BV sweep:
# fused+ring wins at B=1 but regresses from B=4 (B=2 at T=8) on
# both enabled architectures. Keep the ring path conservative until
# other batch/architecture combinations are measured. The snapshot
# path and the separate CuTe path are unchanged.
return False
expected_dim = (
2 * layer.num_q_heads * layer.head_k_dim
+ layer.num_v_heads * layer.head_v_dim
@@ -1273,7 +1296,21 @@ class KDAAttnBackend(MambaAttnBackendBase):
layer.A_log.dtype != torch.float32
or layer.dt_bias.dtype != torch.float32
or ssm_states.dtype != torch.float32
or intermediate_state_cache is None
):
return False
if replayssm_on:
if not self._replayssm_ring_ok(
layer=layer,
draft_token_num=draft_token_num,
mixed_qkv=mixed_qkv,
replayssm_rawv=replayssm_rawv,
replayssm_rawk=replayssm_rawk,
replayssm_g=replayssm_g,
replayssm_beta=replayssm_beta,
):
return False
elif (
intermediate_state_cache is None
or intermediate_state_cache.dtype != torch.float32
):
return False
@@ -1283,7 +1320,7 @@ class KDAAttnBackend(MambaAttnBackendBase):
or b.stride(-1) != 1
or not conv_states.is_contiguous()
or not ssm_states.is_contiguous()
or not intermediate_state_cache.is_contiguous()
or (not replayssm_on and not intermediate_state_cache.is_contiguous())
):
return False
if (
@@ -1301,10 +1338,15 @@ class KDAAttnBackend(MambaAttnBackendBase):
or ssm_states.ndim != 4
or tuple(ssm_states.shape[-3:])
!= (layer.num_v_heads, layer.head_v_dim, layer.head_k_dim)
or intermediate_state_cache.ndim != 5
or intermediate_state_cache.shape[1] < draft_token_num
or tuple(intermediate_state_cache.shape[-3:])
!= (layer.num_v_heads, layer.head_v_dim, layer.head_k_dim)
or (
not replayssm_on
and (
intermediate_state_cache.ndim != 5
or intermediate_state_cache.shape[1] < draft_token_num
or tuple(intermediate_state_cache.shape[-3:])
!= (layer.num_v_heads, layer.head_v_dim, layer.head_k_dim)
)
)
):
return False
if (
@@ -1324,15 +1366,76 @@ class KDAAttnBackend(MambaAttnBackendBase):
b,
conv_states,
ssm_states,
intermediate_state_cache,
intermediate_conv_window_cache,
cache_indices,
intermediate_state_indices,
)
if layer.bias is not None:
tensors += (layer.bias,)
# Ring devices are validated in _replayssm_ring_ok.
if not replayssm_on:
tensors += (intermediate_state_cache,)
return all(tensor.device == mixed_qkv.device for tensor in tensors)
@staticmethod
def _replayssm_ring_ok(
*,
layer: RadixLinearAttention,
draft_token_num: int,
mixed_qkv: torch.Tensor,
replayssm_rawv: torch.Tensor,
replayssm_rawk: Optional[torch.Tensor],
replayssm_g: Optional[torch.Tensor],
replayssm_beta: Optional[torch.Tensor],
) -> bool:
"""Whether the per-layer ReplaySSM rings fit the fused ring-write.
Layouts follow memory_pool.py's KDA spec rings: rawv [slots, HV, L, V]
and rawk [slots, H, L, K] in the activation dtype, g [slots, HV, L, K]
fp32 (per-K KDA gate), beta [slots, HV, L] fp32. The kernel uses
stride(0) as the slot pitch and assumes packed inner dims; anything
else falls back to the unfused path, which handles it.
"""
if replayssm_rawk is None or replayssm_g is None or replayssm_beta is None:
return False
if (
replayssm_rawv.ndim != 4
or replayssm_rawk.ndim != 4
or replayssm_g.ndim != 4
or replayssm_beta.ndim != 3
):
return False
H, HV = layer.num_q_heads, layer.num_v_heads
K, V = layer.head_k_dim, layer.head_v_dim
ring_len = replayssm_rawv.shape[-2]
if ring_len < draft_token_num:
return False
if (
tuple(replayssm_rawv.shape[1:]) != (HV, ring_len, V)
or tuple(replayssm_rawk.shape[1:]) != (H, ring_len, K)
or tuple(replayssm_g.shape[1:]) != (HV, ring_len, K)
or tuple(replayssm_beta.shape[1:]) != (HV, ring_len)
):
return False
if (
replayssm_rawv.dtype != mixed_qkv.dtype
or replayssm_rawk.dtype != mixed_qkv.dtype
or replayssm_g.dtype != torch.float32
or replayssm_beta.dtype != torch.float32
):
return False
if (
replayssm_rawv.stride()[1:] != (ring_len * V, V, 1)
or replayssm_rawk.stride()[1:] != (ring_len * K, K, 1)
or replayssm_g.stride()[1:] != (ring_len * K, K, 1)
or replayssm_beta.stride()[1:] != (ring_len, 1)
):
return False
return all(
ring.device == mixed_qkv.device
for ring in (replayssm_rawv, replayssm_rawk, replayssm_g, replayssm_beta)
)
def _can_run_dspark_cutedsl_mtp(
self,
*,