[MTP] Cut spec-v2 host-seam overhead in hybrid-linear MTP decode (#32219)

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

python/sglang/kernels/ops/mamba/mamba_state_indices_triton.py

@@ -0,0 +1,75 @@
+"""Fused replay-prep state-indices kernel for the mamba cuda-graph path.
+
+``MambaAttnBackendBase._replay_metadata`` refreshes the captured per-bs
+``state_indices_list`` buffer before every cuda-graph replay. The reference
+form is a chain of dispatched aten ops whose host cost shows up in the bs=1
+MTP inter-phase seam:
+
+    req_pool_indices[valid_bs:] = 0             # zero padded rows (side effect)
+    mamba_indices = mapping[req_pool_indices]   # get_mamba_indices gather
+    mamba_indices = translate(mamba_indices)    # identity for the static pool
+    mamba_indices[valid_bs:] = -1               # padding sentinel
+    state_indices[:total_bs].copy_(mamba_indices)
+
+This module fuses that chain into a single launch.
+"""
+
+import torch
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def _fused_replay_state_indices_kernel(
+    req_pool_indices_ptr,  # (total_bs,) int64 — static replay buffer
+    mamba_map_ptr,  # (req_pool_size,) int32 — req_index_to_mamba_index_mapping
+    out_ptr,  # (total_bs,) int32 — state_indices_list[bs - 1]
+    valid_bs,
+    total_bs,
+    BS_UPPER: tl.constexpr,
+):
+    offs = tl.arange(0, BS_UPPER)
+    in_range = offs < total_bs
+    valid = offs < valid_bs
+    req = tl.load(req_pool_indices_ptr + offs, mask=valid, other=0)
+    idx = tl.load(mamba_map_ptr + req, mask=valid, other=0)
+    out_val = tl.where(valid, idx.to(tl.int32), -1)
+    tl.store(out_ptr + offs, out_val, mask=in_range)
+    # Preserve the reference chain's side effect: padded rows of the static
+    # req_pool_indices buffer are zeroed so captured kernels that gather
+    # with them stay in-bounds.
+    zeros = tl.zeros([BS_UPPER], dtype=req.dtype)
+    tl.store(req_pool_indices_ptr + offs, zeros, mask=in_range & (~valid))
+
+
+def fused_replay_state_indices(
+    *,
+    req_pool_indices: torch.Tensor,
+    mamba_index_mapping: torch.Tensor,
+    out_state_indices: torch.Tensor,
+    valid_bs: int,
+    total_bs: int,
+) -> torch.Tensor:
+    """Fill the captured replay state-indices buffer in one launch.
+
+    Mapping gather + padding sentinel + store into ``out_state_indices``, plus
+    the reference chain's side effect of zeroing the padded rows of
+    ``req_pool_indices``. Rows ``[valid_bs, total_bs)`` are padding: they get
+    the ``-1`` sentinel (mamba kernels skip ``state_idx < 0``) and their
+    ``req_pool_indices`` entries are zeroed.
+
+    Callers must supply an identity virtual->physical mapping (the static
+    hybrid pool); the unified pool's allocator translate is not a flat table
+    gather and has to take the reference chain.
+
+    Returns the filled ``out_state_indices[:total_bs]`` view.
+    """
+    _fused_replay_state_indices_kernel[(1,)](
+        req_pool_indices,
+        mamba_index_mapping,
+        out_state_indices,
+        valid_bs,
+        total_bs,
+        BS_UPPER=triton.next_power_of_2(total_bs),
+    )
+    return out_state_indices[:total_bs]

python/sglang/kernels/ops/speculative/cache_locs.py

@@ -361,6 +361,74 @@ def assign_extend_cache_locs(
         save_offset += BLOCK_SIZE
 
 
+@triton.jit
+def assign_extend_cache_locs_uniform(
+    req_pool_indices,
+    req_to_token,
+    start_offset,
+    out_cache_loc,
+    pool_len: tl.constexpr,
+    draft_token_num: tl.constexpr,
+):
+    """Uniform-length variant of assign_extend_cache_locs: every row extends
+    exactly draft_token_num tokens, so the end offset is start +
+    draft_token_num (computed here, no end_offset tensor) and the output
+    offset is pid * draft_token_num (no cross-row prefix-sum loads)."""
+    BLOCK_SIZE: tl.constexpr = 64
+    pid = tl.program_id(axis=0)
+    kv_start = tl.load(start_offset + pid)
+    token_pool = req_to_token + tl.load(req_pool_indices + pid) * pool_len
+    out_cache_ptr = out_cache_loc + pid * draft_token_num
+
+    offs = tl.arange(0, BLOCK_SIZE)
+    num_loop = tl.cdiv(draft_token_num, BLOCK_SIZE)
+    for i in range(num_loop):
+        o = offs + i * BLOCK_SIZE
+        mask = o < draft_token_num
+        data = tl.load(token_pool + kv_start + o, mask=mask)
+        tl.store(out_cache_ptr + o, data, mask=mask)
+
+
+def assign_extend_cache_locs_uniform_func(
+    req_pool_indices: torch.Tensor,
+    req_to_token: torch.Tensor,
+    start_offset: torch.Tensor,
+    batch_size: int,
+    draft_token_num: int,
+    device,
+) -> torch.Tensor:
+    """assign_extend_cache_locs for the uniform case (all rows extend exactly
+    draft_token_num tokens, e.g. spec target-verify prep). Computes end
+    offsets inside the kernel, removing the eager `seq_lens + draft_token_num`
+    add from the host critical path."""
+    if _is_cuda or _is_hip or _is_musa or _is_xpu:
+        out_cache_loc = torch.empty(
+            (batch_size * draft_token_num,),
+            dtype=torch.int64,
+            device=device,
+        )
+        assign_extend_cache_locs_uniform[(batch_size,)](
+            req_pool_indices,
+            req_to_token,
+            start_offset,
+            out_cache_loc,
+            req_to_token.shape[1],
+            draft_token_num,
+        )
+        return out_cache_loc
+
+    # NPU / CPU platforms: fall back to the end_offset-tensor path.
+    return assign_extend_cache_locs_func(
+        req_pool_indices=req_pool_indices,
+        req_to_token=req_to_token,
+        start_offset=start_offset,
+        end_offset=start_offset + draft_token_num,
+        batch_size=batch_size,
+        draft_token_num=draft_token_num,
+        device=device,
+    )
+
+
 def assign_extend_cache_locs_func(
     req_pool_indices: torch.Tensor,
     req_to_token: torch.Tensor,

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

@@ -4,6 +4,9 @@ from typing import Optional, Union
 import torch
 
 from sglang.kernels.ops.mamba.causal_conv1d_triton import PAD_SLOT_ID
+from sglang.kernels.ops.mamba.mamba_state_indices_triton import (
+    fused_replay_state_indices,
+)
 from sglang.kernels.ops.mamba.mamba_state_scatter_triton import (
     scatter_mamba_states_after_mtp_verify,
     track_mamba_states_if_needed,
@@ -36,6 +39,16 @@ class MambaAttnBackendBase(AttentionBackend):
         self.req_to_token_pool: HybridReqToTokenPool = model_runner.req_to_token_pool
         self.token_to_kv_pool = model_runner.token_to_kv_pool
         self.enable_unified_memory = model_runner.server_args.enable_unified_memory
+        # Fused replay-prep state-indices fast path (fused_replay_state_indices):
+        # requires the static hybrid pool whose v2p translate is the identity —
+        # the unified pool overrides translate_mamba_indices with an allocator
+        # lookup that is not a flat table gather.
+        self._fused_state_indices_ok = (
+            str(self.device).startswith("cuda")
+            and isinstance(self.req_to_token_pool, HybridReqToTokenPool)
+            and type(self.req_to_token_pool).translate_mamba_indices
+            is HybridReqToTokenPool.translate_mamba_indices
+        )
         self.forward_metadata: ForwardMetadata = None
         self.state_indices_list = []
         # Static (max_bs,) track-dest buffer captured by pointer, refreshed in-place
@@ -242,6 +255,30 @@ class MambaAttnBackendBase(AttentionBackend):
     def init_forward_metadata(self, forward_batch: ForwardBatch):
         self.forward_metadata = self._forward_metadata(forward_batch)
 
+    def update_verify_buffers_to_fill_after_draft(
+        self, spec_info: SpecInput, cuda_graph_bs: Optional[int]
+    ):
+        # Plan-stream fixup: slot indices / static query_start_loc are
+        # draft-independent, but tree verify (topk > 1) copies the
+        # draft-produced tree links into the captured buffers on the plan
+        # stream, racing the draft — re-copy after the stream join. Eager
+        # verify reads the spec_info tensors directly; parent links are
+        # derived from these buffers at execution time.
+        if self.topk <= 1 or cuda_graph_bs is None:
+            return
+        if (
+            not isinstance(spec_info, EagleVerifyInput)
+            or spec_info.retrieve_next_token is None  # dummy / capture runs
+        ):
+            return
+        bs_without_pad = spec_info.retrieve_next_token.shape[0]
+        self.retrieve_next_token_list[cuda_graph_bs - 1][:bs_without_pad].copy_(
+            spec_info.retrieve_next_token
+        )
+        self.retrieve_next_sibling_list[cuda_graph_bs - 1][:bs_without_pad].copy_(
+            spec_info.retrieve_next_sibling
+        )
+
     def _init_track_conv_indices(
         self, query_start_loc: torch.Tensor, forward_batch: ForwardBatch
     ):
@@ -507,14 +544,26 @@ class MambaAttnBackendBase(AttentionBackend):
                 num_padding = torch.count_nonzero(
                     seq_lens_cpu == self.get_cuda_graph_seq_len_fill_value()
                 )
-        # Make sure forward metadata is correctly handled for padding reqs
-        req_pool_indices[bs - num_padding :] = 0
-        mamba_indices = self.req_to_token_pool.get_mamba_indices(req_pool_indices)
-        # Translate using the LIVE v2p table BEFORE the padding sentinel below;
-        # captured Mamba kernels read state_indices_list as PHYSICAL ids.
-        mamba_indices = self._translate_mamba_indices(mamba_indices)
-        mamba_indices[bs - num_padding :] = -1
-        self.state_indices_list[bs - 1][: len(mamba_indices)].copy_(mamba_indices)
+        if self._fused_state_indices_ok and self.replayssm_write_pos_list is None:
+            # Single-launch fast path: mapping gather + padding sentinel + store
+            # into the static buffer, plus zeroing padded req_pool_indices rows —
+            # bit-identical to the reference chain below.
+            mamba_indices = fused_replay_state_indices(
+                req_pool_indices=req_pool_indices,
+                mamba_index_mapping=self.req_to_token_pool.req_index_to_mamba_index_mapping,
+                out_state_indices=self.state_indices_list[bs - 1],
+                valid_bs=bs - int(num_padding),
+                total_bs=bs,
+            )
+        else:
+            # Make sure forward metadata is correctly handled for padding reqs
+            req_pool_indices[bs - num_padding :] = 0
+            mamba_indices = self.req_to_token_pool.get_mamba_indices(req_pool_indices)
+            # Translate using the LIVE v2p table BEFORE the padding sentinel below;
+            # captured Mamba kernels read state_indices_list as PHYSICAL ids.
+            mamba_indices = self._translate_mamba_indices(mamba_indices)
+            mamba_indices[bs - num_padding :] = -1
+            self.state_indices_list[bs - 1][: len(mamba_indices)].copy_(mamba_indices)
         # Refresh the static track-dest buffer in-place (translated); the captured
         # track-save reads it, leaving the handed-in InputBuffer slot read-only.
         track_buf = None
@@ -866,6 +915,24 @@ class HybridLinearAttnBackend(AttentionBackend):
         for attn_backend in self.attn_backend_list:
             attn_backend.on_after_cuda_graph_warmup()
 
+    def get_verify_buffers_to_fill_after_draft(self):
+        # Verify tree-mask / position buffers live on the full-attn child (the
+        # linear side consumes no mask). Handing them out lets the draft stage
+        # write straight into the captured verify buffers instead of allocating
+        # a fresh mask every step.
+        return self.full_attn_backend.get_verify_buffers_to_fill_after_draft()
+
+    def update_verify_buffers_to_fill_after_draft(
+        self, spec_info: SpecInput, cuda_graph_bs: Optional[int]
+    ):
+        # Plan-stream fixup after draft completes: forward to both children.
+        # Sub-backends that cannot run under the plan stream keep the fail-loud
+        # NotImplementedError base behavior.
+        for attn_backend in self.attn_backend_list:
+            attn_backend.update_verify_buffers_to_fill_after_draft(
+                spec_info=spec_info, cuda_graph_bs=cuda_graph_bs
+            )
+
     def init_forward_metadata(self, forward_batch: ForwardBatch):
         if forward_batch.forward_mode.is_draft_extend_v2():
             # DRAFT_EXTEND_V2 runs only full-attn layers in the draft model; skip

python/sglang/srt/layers/attention/linear/kda_backend.py

@@ -252,6 +252,13 @@ class KDAKernelDispatcher:
 class KDAAttnBackend(MambaAttnBackendBase):
     """Attention backend for KDA (Kimi Delta Attention) linear attention."""
 
+    # Same GPU-only contract as GDNAttnBackend / Mamba2AttnBackend: KDA metadata
+    # never reads the spec-v2 seq_lens_cpu mirror (replay padding comes from
+    # forward_batch.num_padding, and the replayssm track-flush mask paths are
+    # gated `not is_kda`), so don't force FutureMap's blocking per-step
+    # seq_lens D2H (~0.5 ms/step host stall in bs=1 MTP decode).
+    needs_cpu_seq_lens: bool = False
+
     def __init__(self, model_runner: ModelRunner):
         super().__init__(model_runner)
         # mamba_cache.conv is [..., kernel-1, dim] while conv_states_shape expects the window length (kernel-1) at shape[-1], hence the transpose.

python/sglang/srt/speculative/eagle_utils.py

@@ -491,7 +491,7 @@ def eagle_prepare_for_verify(
     target_worker: TpModelWorker,
 ):
     from sglang.kernels.ops.speculative.cache_locs import (
-        assign_extend_cache_locs_func,
+        assign_extend_cache_locs_uniform_func,
     )
     from sglang.srt.model_executor.forward_batch_info import (
         CaptureHiddenMode,
@@ -511,11 +511,13 @@ def eagle_prepare_for_verify(
             "v2 prepare_for_verify input_ids",
         )
         device = batch.device
-        batch.out_cache_loc = assign_extend_cache_locs_func(
+        # Uniform variant: end offsets (= start + draft_token_num) are computed
+        # inside the kernel, keeping the eager `seq_lens + N` add off the host
+        # critical path (bs=1 MTP inter-phase seam).
+        batch.out_cache_loc = assign_extend_cache_locs_uniform_func(
             req_pool_indices=batch.req_pool_indices,
             req_to_token=req_to_token_pool.req_to_token,
             start_offset=batch.seq_lens,
-            end_offset=batch.seq_lens + verify_input.draft_token_num,
             batch_size=bs,
             draft_token_num=verify_input.draft_token_num,
             device=device,