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

Co-authored-by: luoyuan.luo <luoyuan.luo@antgroup.com>

python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py

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

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

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

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