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

[KDA] Fused-accept state advance for FlashInfer KDA MTP verify (#33722)

Browse Source 
Co-authored-by: luoyuan.luo <luoyuan.luo@antgroup.com>
This commit is contained in:
Yuan Luo
2026-08-31 16:11:49 +08:00
committed by GitHub
co-authored by luoyuan.luo
parent f61bb7b40a
commit 712a720c8a
5 changed files with 362 additions and 27 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py
+31
View File
@@ -10,6 +10,7 @@ from sglang.kernels.ops.mamba.mamba_state_indices_triton import (
fused_replay_state_indices,
)
from sglang.kernels.ops.mamba.mamba_state_scatter_triton import (
fused_conv_window_scatter_with_mask,
scatter_mamba_states_after_mtp_verify,
track_mamba_states_all_layers,
track_mamba_states_if_needed,
@@ -50,6 +51,11 @@ _validate_mamba_replay_state_indices = (
class MambaAttnBackendBase(AttentionBackend):
# Per-slot accept lengths for the KDA fused-accept spec path; allocated only
# by KDAAttnBackend where `_can_fuse_accept_state` holds. None everywhere
# else — update_mamba_state_after_mtp_verify keys the fused branch on it.
accept_lens_pool: Optional[torch.Tensor] = None
def __init__(self, model_runner: ModelRunner):
super().__init__()
self.pad_slot_id = PAD_SLOT_ID
@@ -1282,6 +1288,31 @@ class HybridLinearAttnBackend(AttentionBackend):
)
return
# KDA fused-accept: the next verify seeds itself in-kernel from the
# accepted checkpoint slot (recurrent_kda's num_accepted_tokens), so the
# SSM state never round-trips through `temporal` and only the conv
# windows still need the accept rollback. Recording this round's accept
# length is what selects that seed next round; chain layout only (see
# above), so accept_lens == last_correct_step_indices + 1. The pool
# exists only where KDAAttnBackend found the contract satisfied, which
# includes mamba radix tracking being off.
accept_lens_pool = self.linear_attn_backend.accept_lens_pool
if accept_lens_pool is not None:
assert mamba_track_indices is None, "fused-accept runs with radix off"
for conv_states, intermediate_conv_window in zip(
mamba_caches.conv, mamba_caches.intermediate_conv_window
):
fused_conv_window_scatter_with_mask(
conv_states,
intermediate_conv_window,
state_indices_tensor,
last_correct_step_indices,
)
accept_lens_pool[state_indices_tensor.to(torch.int64)] = (
last_correct_step_indices.to(torch.int32) + 1
)
return
scatter_mamba_states_after_mtp_verify(
mamba_caches,
state_indices_tensor,
python/sglang/srt/layers/attention/linear/kda_backend.py
+108
View File
@@ -10,6 +10,9 @@ from sglang.kernels.ops.mamba.causal_conv1d_triton import (
)
from sglang.srt.environ import envs
from sglang.srt.layers.attention.hybrid_linear_attn_backend import MambaAttnBackendBase
from sglang.srt.layers.attention.linear.kernels.kda_flashinfer import (
build_fused_accept_indices,
)
from sglang.srt.layers.attention.linear.kernels.kda_triton import TritonKDAKernel
from sglang.srt.layers.attention.linear.utils import (
LinearAttnKernelBackend,
@@ -33,6 +36,9 @@ elif is_cpu():
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
from sglang.srt.model_executor.model_runner import ModelRunner
from sglang.srt.runtime_context import (
get_disagg,
get_exec,
get_memory,
get_spec,
)
@@ -435,6 +441,80 @@ class KDAAttnBackend(MambaAttnBackendBase):
model_runner.device,
)
)
# Fused-accept spec path (flashinfer recurrent_kda): the next verify
# seeds itself in-kernel from the accepted checkpoint slot
# (num_accepted_tokens), so the per-round SSM commit scatter is skipped.
# accept_lens_pool holds last round's accept length per mamba slot;
# extend stages fresh requests with 1 (read slot 0). Its presence is the
# signal that switches the post-verify commit to conv-only.
if self._can_fuse_accept_state(verify_backend):
self.accept_lens_pool = torch.ones(
self.req_to_token_pool.size + 1,
dtype=torch.int32,
device=model_runner.device,
)
@staticmethod
def _can_fuse_accept_state(verify_backend) -> bool:
"""Whether the verify kernel can seed itself from the accepted checkpoint.
The seed comes from recurrent_kda's ``num_accepted_tokens``, which makes
the previous round's accepted state addressable in-kernel so the SSM
state never has to round-trip through the committed pool (`temporal`).
Hence the keying on the verify backend -- the one that selects the
target_verify kernel -- and not on decode, which is set separately.
`temporal` then goes stale between verifies, which is what the remaining
conditions rule out: each is a reader of the committed state that the
skipped scatter would starve. Falling short of the contract falls back to
the commit scatter rather than raising -- this is a capability, not a
mode. Every condition reads its namespace bag rather than the record:
these are fields resolution decides, so the record would answer with
what the operator typed instead of what was decided.
"""
if not verify_backend.is_flashinfer():
return False # only recurrent_kda takes num_accepted_tokens
if get_spec().speculative_algorithm is None:
return False # no verify round, and no intermediate scratch
if not get_memory().disable_radix_cache:
return False # mamba radix tracking snapshots `temporal`
if get_exec().mamba.enable_linear_replayssm_spec:
return False # the ring already owns the verify-round commitment
if get_disagg().disaggregation_mode != "null":
return False # the PD hand-off transfers `temporal`
return True
def _fused_accept_indices(
self,
*,
cache_indices: torch.Tensor,
query_start_loc: torch.Tensor,
intermediate_state_cache: torch.Tensor,
draft_token_num: int,
):
"""Slot-indexed verify rows + accept lengths for this forward.
Every KDA layer of a forward verifies the same requests over the same
draft window, so the build is hoisted onto the shared forward metadata:
layer 0 builds, the rest reuse. Under cuda graph the capture then holds
a single build reading the static slot buffer, instead of one per layer.
"""
metadata = self.forward_metadata
if metadata.fused_accept_state_indices is None:
batch_size = query_start_loc.shape[0] - 1
(
metadata.fused_accept_state_indices,
metadata.fused_accept_num_accepted,
) = build_fused_accept_indices(
slots=cache_indices[:batch_size],
scratch_steps=intermediate_state_cache.shape[1],
draft_token_num=draft_token_num,
accept_lens_pool=self.accept_lens_pool,
)
return (
metadata.fused_accept_state_indices,
metadata.fused_accept_num_accepted,
)
def init_forward_metadata(self, forward_batch: ForwardBatch):
super().init_forward_metadata(forward_batch)
@@ -732,6 +812,19 @@ class KDAAttnBackend(MambaAttnBackendBase):
forward_batch, h, ssm_states, self.forward_metadata
)
if (
self.accept_lens_pool is not None
and not forward_batch.forward_mode.is_draft_extend_v2()
):
# Fused-accept staging: the extend kernel just wrote this request's
# committed state; copy it into scratch slot 0 and reset the accept
# length to 1 so the first verify reads slot 0. Runs once per KDA
# layer (the nat write is idempotent; the scratch copy is per-layer).
slots = cache_indices.to(torch.int64)
intermediate_ssm = mamba_cache_params.intermediate_ssm
intermediate_ssm[slots, 0] = ssm_states[slots].to(intermediate_ssm.dtype)
self.accept_lens_pool[slots] = 1
return core_attn_out
def _forward_target_verify(
@@ -951,6 +1044,21 @@ class KDAAttnBackend(MambaAttnBackendBase):
retrieve_parent_token=retrieve_parent_token,
lower_bound=layer.lower_bound,
**ring_kwargs,
**(
dict(
zip(
("fused_accept_state_indices", "fused_accept_num_accepted"),
self._fused_accept_indices(
cache_indices=cache_indices,
query_start_loc=query_start_loc,
intermediate_state_cache=intermediate_state_cache,
draft_token_num=draft_token_num,
),
)
)
if self.accept_lens_pool is not None
else {}
),
)
if dense_token_indices is not None:
# Kernel output is empty-allocated and the capped qsl skips the
python/sglang/srt/layers/attention/linear/kernels/kda_flashinfer.py
+75 -27
View File
@@ -57,6 +57,33 @@ def _get_flashinfer_kda_kernel():
return _flashinfer_kda_available, _flashinfer_recurrent_kda
def build_fused_accept_indices(
*,
slots: torch.Tensor,
scratch_steps: int,
draft_token_num: int,
accept_lens_pool: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
"""Slot-indexed verify indices + accept lengths for fused-accept mode.
Row n of the returned ``[N, T]`` index tensor addresses the scratch slots of
the request holding mamba slot ``slots[n]``. A padded row (``slots[n] < 0``)
yields ONLY negative indices (``-scratch_steps + step`` with
``step < scratch_steps``), which recurrent_kda treats as inactive — the
padding contract must survive any refactor of this arithmetic. The nat
gather clamps padded slots to row 0 of the pool; their value is never
consumed (inactive rows). All ops are device-side and capture-safe.
"""
step = torch.arange(draft_token_num, device=slots.device, dtype=torch.int32)
ssm_state_indices = (
slots.to(torch.int32)[:, None] * scratch_steps + step[None, :]
).contiguous() # [N, T]
num_accepted_tokens = accept_lens_pool.index_select(
0, slots.clamp(min=0).to(torch.int64)
)
return ssm_state_indices, num_accepted_tokens
class FlashInferKDAKernel(LinearAttnKernelBase):
"""FlashInfer KDA kernel: SM100 decode + MTP (target_verify), topk=1.
@@ -235,6 +262,8 @@ class FlashInferKDAKernel(LinearAttnKernelBase):
cache_steps: int,
retrieve_parent_token: torch.Tensor,
lower_bound: Optional[float] = None,
fused_accept_state_indices: Optional[torch.Tensor] = None,
fused_accept_num_accepted: Optional[torch.Tensor] = None,
**kwargs,
) -> torch.Tensor:
if retrieve_parent_token is not None:
@@ -272,36 +301,54 @@ class FlashInferKDAKernel(LinearAttnKernelBase):
f"but intermediate_ssm only has {scratch_steps}."
)
base_rows = intermediate_state_indices[:batch_size]
cache_key = (
id(intermediate_state_indices),
batch_size,
draft_token_num,
scratch_steps,
)
ssm_state_indices = self._verify_idx_cache.get(cache_key)
if ssm_state_indices is None:
# The fast seed copy below assumes row n in scratch belongs to request n.
expected = torch.arange(
batch_size, device=base_rows.device, dtype=base_rows.dtype
if fused_accept_state_indices is not None:
# Fused-accept mode: rows are the requests' mamba SLOTS (stable for
# the request lifetime, unlike batch positions), so last round's
# checkpoints are addressable this round. The kernel seeds each row
# from slot[nat - 1] (nat = last round's accept length, gathered
# from accept_lens_pool; fresh requests were staged with nat = 1 at
# extend) and overwrites all T slots in place — no committed-pool
# seed copy here and no SSM commit scatter after verify. Padded
# graph rows carry slot -1: every derived index stays negative,
# which recurrent_kda treats as inactive. Both tensors are built
# once per forward by the backend (see KDAAttnBackend), so the 20
# KDA layers of a step share one build.
ssm_state_indices = fused_accept_state_indices
num_accepted_tokens = fused_accept_num_accepted
else:
num_accepted_tokens = None
base_rows = intermediate_state_indices[:batch_size]
cache_key = (
id(intermediate_state_indices),
batch_size,
draft_token_num,
scratch_steps,
)
if not torch.equal(base_rows, expected):
raise RuntimeError(
"FlashInfer KDA verify requires an identity intermediate row-map "
"(verify_intermediate_state_indices must be arange)."
ssm_state_indices = self._verify_idx_cache.get(cache_key)
if ssm_state_indices is None:
# The fast seed copy below assumes row n in scratch belongs to
# request n.
expected = torch.arange(
batch_size, device=base_rows.device, dtype=base_rows.dtype
)
step = torch.arange(draft_token_num, device=q.device, dtype=torch.int32)
ssm_state_indices = (
base_rows.to(torch.int32)[:, None] * scratch_steps + step[None, :]
).contiguous() # [N, T]
self._verify_idx_cache[cache_key] = ssm_state_indices
if not torch.equal(base_rows, expected):
raise RuntimeError(
"FlashInfer KDA verify requires an identity intermediate "
"row-map (verify_intermediate_state_indices must be arange)."
)
step = torch.arange(draft_token_num, device=q.device, dtype=torch.int32)
ssm_state_indices = (
base_rows.to(torch.int32)[:, None] * scratch_steps + step[None, :]
).contiguous() # [N, T]
self._verify_idx_cache[cache_key] = ssm_state_indices
# Seed step 0 from committed state, then recurrent_kda overwrites it with
# token-0 post-state. Padded graph rows clamp to slot 0; their output is ignored.
base_state = ssm_states.index_select(
0, cache_indices[:batch_size].clamp(min=0).to(torch.int64)
)
scratch[:batch_size, 0].copy_(base_state)
# Seed step 0 from committed state, then recurrent_kda overwrites it
# with token-0 post-state. Padded graph rows clamp to slot 0; their
# output is ignored.
base_state = ssm_states.index_select(
0, cache_indices[:batch_size].clamp(min=0).to(torch.int64)
)
scratch[:batch_size, 0].copy_(base_state)
# Same storage as scratch, flattened over the allocated step stride.
state_pool = scratch.view(
@@ -325,6 +372,7 @@ class FlashInferKDAKernel(LinearAttnKernelBase):
cu_seqlens=query_start_loc.to(torch.int32),
ssm_state_indices=ssm_state_indices,
num_spec_tokens=num_spec_tokens,
num_accepted_tokens=num_accepted_tokens,
)
return output_fi.view(1, seq_len, num_v_heads, head_v_dim)
python/sglang/srt/layers/attention/mamba/mamba2_metadata.py
+7
View File
@@ -62,6 +62,13 @@ class ForwardMetadata:
is_target_verify: bool = False
draft_token_num: int = 1
# KDA fused-accept: the [N, T] slot-indexed scratch rows and the per-request
# accept length that seed the verify kernel. Every KDA layer of a forward
# sees the same slots and draft window, so these are built once and shared:
# a cuda-graph capture then holds one build instead of one per layer.
fused_accept_state_indices: Optional[torch.Tensor] = None
fused_accept_num_accepted: Optional[torch.Tensor] = None
has_mamba_track_mask: bool = False
mamba_track_mask_indices: Optional[torch.Tensor] = None
conv_states_mask_indices: Optional[torch.Tensor] = None
test/registered/attention/unittests/hybrid_linear/test_kda_fused_accept_indices.py
+141
View File
@@ -0,0 +1,141 @@
"""Padding-safety invariants of the KDA fused-accept verify index builder.
``build_fused_accept_indices`` produces the slot-indexed ``[N, T]``
``ssm_state_indices`` and the per-row ``num_accepted_tokens`` gather that
flashinfer ``recurrent_kda`` consumes in fused-accept mode. The kernel's
padding contract is: a row is inactive iff its raw slot index is negative.
Padded sglang rows carry mamba slot ``-1``, so EVERY derived index
``-1 * scratch_steps + step`` must stay negative for all
``step < scratch_steps`` — an arithmetic reorder (e.g. adding the step before
the multiply) would silently activate padded rows and corrupt neighbor state.
The nat gather must clamp padded slots in-bounds (their value is never
consumed) and keep real slots' accept lengths intact.
CPU tensors only — the invariants are pure index arithmetic.
"""
import unittest
import torch
from sglang.srt.layers.attention.linear.kernels.kda_flashinfer import (
build_fused_accept_indices,
)
from sglang.test.ci.ci_register import register_cuda_ci
from sglang.test.test_utils import CustomTestCase
register_cuda_ci(est_time=2, stage="base-b", runner_config="1-gpu-large")
class TestBuildFusedAcceptIndices(CustomTestCase):
def test_real_slots_address_their_scratch_rows(self):
for scratch_steps, draft_token_num in ((4, 4), (8, 4), (33, 33)):
slots = torch.tensor([0, 3, 7], dtype=torch.int32)
pool = torch.arange(10, dtype=torch.int32) + 1 # accept len = slot + 1
indices, nat = build_fused_accept_indices(
slots=slots,
scratch_steps=scratch_steps,
draft_token_num=draft_token_num,
accept_lens_pool=pool,
)
self.assertEqual(indices.shape, (3, draft_token_num))
self.assertEqual(indices.dtype, torch.int32)
step = torch.arange(draft_token_num, dtype=torch.int32)
expected = slots[:, None] * scratch_steps + step[None, :]
self.assertTrue(torch.equal(indices, expected))
self.assertEqual(nat.dtype, torch.int32)
self.assertTrue(torch.equal(nat, slots + 1))
def test_padded_slot_rows_stay_fully_negative(self):
# T == scratch_steps is the tight case: the largest step must still
# land below zero for slot -1.
for scratch_steps, draft_token_num in ((4, 4), (8, 8), (8, 4), (33, 33)):
slots = torch.tensor([2, -1, 5, -1], dtype=torch.int32)
pool = torch.full((8,), 3, dtype=torch.int32)
indices, nat = build_fused_accept_indices(
slots=slots,
scratch_steps=scratch_steps,
draft_token_num=draft_token_num,
accept_lens_pool=pool,
)
padded_rows = indices[slots < 0]
self.assertTrue(
(padded_rows < 0).all(),
f"padded row leaked a non-negative index "
f"({scratch_steps=}, {draft_token_num=}): {padded_rows.tolist()}",
)
real_rows = indices[slots >= 0]
self.assertTrue((real_rows >= 0).all())
# nat gather clamps padded slots in-bounds (value unused).
self.assertEqual(nat.shape[0], 4)
def test_nat_gather_reads_pool_values(self):
slots = torch.tensor([1, 4, -1], dtype=torch.int32)
pool = torch.tensor([9, 2, 9, 9, 5, 9], dtype=torch.int32)
_, nat = build_fused_accept_indices(
slots=slots,
scratch_steps=4,
draft_token_num=4,
accept_lens_pool=pool,
)
self.assertEqual(nat[0].item(), 2)
self.assertEqual(nat[1].item(), 5)
# Padded row clamps to pool row 0; the value is never consumed but the
# gather itself must stay in-bounds.
self.assertEqual(nat[2].item(), 9)
class TestFusedAcceptPerForwardCache(CustomTestCase):
"""The verify indices are built once per forward and shared by every KDA
layer. That sharing is only sound while the cache dies with the forward: a
cache that outlived it would seed the next batch from the previous batch's
mamba slots, which is a silent wrong-state bug (no shape or index error).
"""
@staticmethod
def _build(slots, pool_values, draft_token_num=4, scratch_steps=4):
device = "cuda" if torch.cuda.is_available() else "cpu"
return build_fused_accept_indices(
slots=torch.tensor(slots, dtype=torch.int32, device=device),
scratch_steps=scratch_steps,
draft_token_num=draft_token_num,
accept_lens_pool=torch.tensor(
pool_values, dtype=torch.int32, device=device
),
)
def test_shared_build_matches_a_per_layer_build(self):
"""What every layer reuses must equal what it would have built itself."""
pool = [1] * 8
pool[3], pool[5] = 2, 4
first_idx, first_nat = self._build([3, 5], pool)
second_idx, second_nat = self._build([3, 5], pool)
self.assertTrue(torch.equal(first_idx, second_idx))
self.assertTrue(torch.equal(first_nat, second_nat))
def test_a_different_batch_builds_different_rows(self):
"""Guards the staleness mode: reusing a previous forward's tensor would
address the previous forward's slots, and the values must differ so the
cache-reset is observable rather than accidentally correct."""
pool = [1] * 8
pool[3], pool[5], pool[6] = 2, 4, 3
idx_a, nat_a = self._build([3, 5], pool)
idx_b, nat_b = self._build([6, 5], pool)
self.assertFalse(torch.equal(idx_a, idx_b))
self.assertFalse(torch.equal(nat_a, nat_b))
def test_metadata_starts_uncached(self):
"""A forward's metadata must arrive with no indices carried over: the
backend keys 'build once' on these being None."""
from sglang.srt.layers.attention.mamba.mamba2_metadata import ForwardMetadata
metadata = ForwardMetadata(
query_start_loc=torch.zeros(2, dtype=torch.int32),
mamba_cache_indices=torch.zeros(1, dtype=torch.int32),
)
self.assertIsNone(metadata.fused_accept_state_indices)
self.assertIsNone(metadata.fused_accept_num_accepted)
if __name__ == "__main__":
unittest.main()