[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,

test/registered/attention/unittests/hybrid_linear/test_verify_buffer_fixup_hook.py

@@ -0,0 +1,131 @@
+"""The mamba plan-stream verify fixup hook must refresh draft-produced tree links.
+
+Under ``SGLANG_ENABLE_OVERLAP_PLAN_STREAM``, ``_replay_metadata`` copies
+``spec_info.retrieve_next_token`` / ``retrieve_next_sibling`` into the captured
+per-bs buffers on the plan stream, racing the draft's ``build_tree`` on the
+compute stream. ``update_verify_buffers_to_fill_after_draft`` is the post-join
+fixup that re-copies them into the ``cuda_graph_bs`` buffers; chain mode
+(``topk == 1``), the eager path (``cuda_graph_bs is None``), and dummy/capture
+runs (``retrieve_next_token is None``) must stay no-ops.
+
+Pure host-side buffer logic — CPU tensors, no CUDA required.
+"""
+
+import unittest
+
+import torch
+
+from sglang.srt.layers.attention.hybrid_linear_attn_backend import (
+    MambaAttnBackendBase,
+)
+from sglang.srt.speculative.eagle_info import EagleVerifyInput
+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")
+
+_STALE = -555
+_MAX_BS = 4
+_DRAFT_TOKEN_NUM = 8
+
+
+def _make_backend(topk: int) -> MambaAttnBackendBase:
+    """Bare backend with only the fields the hook reads (no ModelRunner)."""
+    backend = object.__new__(MambaAttnBackendBase)
+    backend.topk = topk
+    # Mirror init_cuda_graph_state's per-bs buffer shapes: (bs, draft_token_num).
+    backend.retrieve_next_token_list = [
+        torch.full((bs, _DRAFT_TOKEN_NUM), _STALE, dtype=torch.int32)
+        for bs in range(1, _MAX_BS + 1)
+    ]
+    backend.retrieve_next_sibling_list = [
+        torch.full((bs, _DRAFT_TOKEN_NUM), _STALE, dtype=torch.int32)
+        for bs in range(1, _MAX_BS + 1)
+    ]
+    return backend
+
+
+def _make_verify_input(bs_without_pad: int, base: int) -> EagleVerifyInput:
+    """EagleVerifyInput carrying just the tree-link tensors the hook consumes."""
+    spec_info = object.__new__(EagleVerifyInput)
+    numel = bs_without_pad * _DRAFT_TOKEN_NUM
+    spec_info.retrieve_next_token = (
+        torch.arange(base, base + numel, dtype=torch.int32)
+    ).reshape(bs_without_pad, _DRAFT_TOKEN_NUM)
+    spec_info.retrieve_next_sibling = (
+        torch.arange(base + numel, base + 2 * numel, dtype=torch.int32)
+    ).reshape(bs_without_pad, _DRAFT_TOKEN_NUM)
+    return spec_info
+
+
+class TestVerifyBufferFixupHook(CustomTestCase):
+    def test_refresh_overwrites_stale_links(self):
+        backend = _make_backend(topk=2)
+        cuda_graph_bs = _MAX_BS
+        bs_without_pad = _MAX_BS - 1  # one padded row
+        spec_info = _make_verify_input(bs_without_pad=bs_without_pad, base=100)
+
+        backend.update_verify_buffers_to_fill_after_draft(
+            spec_info=spec_info, cuda_graph_bs=cuda_graph_bs
+        )
+
+        for buf_list, fresh in (
+            (backend.retrieve_next_token_list, spec_info.retrieve_next_token),
+            (backend.retrieve_next_sibling_list, spec_info.retrieve_next_sibling),
+        ):
+            buf = buf_list[cuda_graph_bs - 1]
+            self.assertTrue(
+                torch.equal(buf[:bs_without_pad], fresh),
+                "fresh tree links not copied into the captured buffer",
+            )
+            # The padded tail row is not covered by the copy.
+            self.assertTrue(
+                (buf[bs_without_pad:] == _STALE).all(),
+                "rows beyond bs_without_pad must not be written",
+            )
+            # Buffers of other captured batch sizes stay untouched.
+            for other_bs in range(1, _MAX_BS):
+                self.assertTrue(
+                    (buf_list[other_bs - 1] == _STALE).all(),
+                    f"buffer for bs={other_bs} must stay untouched",
+                )
+
+    def test_chain_topk1_is_noop(self):
+        backend = _make_backend(topk=1)
+        spec_info = _make_verify_input(bs_without_pad=2, base=100)
+        backend.update_verify_buffers_to_fill_after_draft(
+            spec_info=spec_info, cuda_graph_bs=2
+        )
+        self.assertTrue((backend.retrieve_next_token_list[1] == _STALE).all())
+        self.assertTrue((backend.retrieve_next_sibling_list[1] == _STALE).all())
+
+    def test_eager_path_is_noop(self):
+        backend = _make_backend(topk=2)
+        spec_info = _make_verify_input(bs_without_pad=2, base=100)
+        backend.update_verify_buffers_to_fill_after_draft(
+            spec_info=spec_info, cuda_graph_bs=None
+        )
+        self.assertTrue((backend.retrieve_next_token_list[1] == _STALE).all())
+        self.assertTrue((backend.retrieve_next_sibling_list[1] == _STALE).all())
+
+    def test_dummy_run_none_links_is_noop(self):
+        backend = _make_backend(topk=2)
+        spec_info = object.__new__(EagleVerifyInput)
+        spec_info.retrieve_next_token = None  # dummy / capture run
+        spec_info.retrieve_next_sibling = None
+        backend.update_verify_buffers_to_fill_after_draft(
+            spec_info=spec_info, cuda_graph_bs=2
+        )
+        self.assertTrue((backend.retrieve_next_token_list[1] == _STALE).all())
+        self.assertTrue((backend.retrieve_next_sibling_list[1] == _STALE).all())
+
+    def test_non_eagle_spec_input_is_noop(self):
+        backend = _make_backend(topk=2)
+        backend.update_verify_buffers_to_fill_after_draft(
+            spec_info=None, cuda_graph_bs=2
+        )
+        self.assertTrue((backend.retrieve_next_token_list[1] == _STALE).all())
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/kernels/ops/mamba/test_fused_replay_state_indices.py

@@ -0,0 +1,192 @@
+"""fused_replay_state_indices must be bit-identical to the unfused prep.
+
+The unfused reference is the exact op sequence ``_replay_metadata`` used to
+launch for the static hybrid pool:
+
+    req_pool_indices[valid_bs:total_bs] = 0        # zero padded rows (side effect)
+    mamba_indices = mapping[req_pool_indices]      # get_mamba_indices gather
+    # identity v2p translate (static pool)
+    mamba_indices[valid_bs:] = -1                  # padding sentinel
+    state_indices[:total_bs].copy_(mamba_indices)
+
+The two paths must agree bit-for-bit, INCLUDING the side effect of zeroing the
+padded rows of the static ``req_pool_indices`` replay buffer — captured kernels
+gather with that buffer, so a non-zeroed padded row is a delayed illegal memory
+access, not a visible diff. Both paths run on guard-padded buffers across a
+bs x num_padding matrix (non-power-of-two sizes exercise the BS_UPPER masking):
+
+1. the produced state indices are identical over the whole ``[0, total_bs)``
+   range (padding sentinel rows included);
+2. the ``req_pool_indices`` buffer ends up identical (padded rows zeroed);
+3. neither buffer is written beyond ``total_bs`` (guard tails stay intact).
+"""
+
+import unittest
+
+import torch
+
+from sglang.kernels.ops.mamba.mamba_state_indices_triton import (
+    fused_replay_state_indices,
+)
+from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cuda_ci(est_time=5, stage="base-b-kernel-unit", runner_config="1-gpu-large")
+
+# Guard tail appended to every buffer; must stay untouched by both paths.
+_GUARD = 8
+_GUARD_SENTINEL = -7777
+# Poison for the out buffer so unwritten cells inside [0, total_bs) are caught.
+_OUT_POISON = -12345
+
+_REQ_POOL_SIZE = 160
+_MAMBA_POOL_SIZE = 4096
+
+
+def _reference_chain(
+    req_pool_indices: torch.Tensor,
+    mapping: torch.Tensor,
+    out_buf: torch.Tensor,
+    valid_bs: int,
+    total_bs: int,
+) -> None:
+    """Replicates the _replay_metadata reference ops, in order, in place."""
+    req_pool_indices[valid_bs:total_bs] = 0
+    mamba_indices = mapping[req_pool_indices[:total_bs]]
+    # static pool: _translate_mamba_indices is the identity
+    mamba_indices[valid_bs:] = -1
+    out_buf[: len(mamba_indices)].copy_(mamba_indices)
+
+
+@unittest.skipUnless(torch.cuda.is_available(), "requires CUDA (triton kernel)")
+class TestFusedReplayStateIndices(CustomTestCase):
+    def _run_case(self, total_bs: int, num_padding: int, seed: int) -> None:
+        device = torch.device("cuda")
+        gen = torch.Generator(device="cpu").manual_seed(seed)
+        valid_bs = total_bs - num_padding
+
+        # Production dtypes: req_pool_indices int64 (static replay buffer),
+        # req_index_to_mamba_index_mapping int32, state_indices_list int32.
+        req_pool = torch.randint(
+            0, _REQ_POOL_SIZE, (total_bs + _GUARD,), generator=gen, dtype=torch.int64
+        )
+        req_pool[total_bs:] = _GUARD_SENTINEL
+        mapping = torch.randint(
+            0, _MAMBA_POOL_SIZE, (_REQ_POOL_SIZE,), generator=gen, dtype=torch.int32
+        )
+        out = torch.full((total_bs + _GUARD,), _OUT_POISON, dtype=torch.int32)
+
+        req_pool_ref = req_pool.to(device)
+        req_pool_fused = req_pool.to(device)
+        mapping_dev = mapping.to(device)
+        out_ref = out.to(device)
+        out_fused = out.to(device)
+
+        _reference_chain(
+            req_pool_indices=req_pool_ref,
+            mapping=mapping_dev,
+            out_buf=out_ref,
+            valid_bs=valid_bs,
+            total_bs=total_bs,
+        )
+        returned = fused_replay_state_indices(
+            req_pool_indices=req_pool_fused,
+            mamba_index_mapping=mapping_dev,
+            out_state_indices=out_fused,
+            valid_bs=valid_bs,
+            total_bs=total_bs,
+        )
+        torch.cuda.synchronize()
+
+        case = f"{total_bs=} {num_padding=} {seed=}"
+        # 1. state indices bit-identical over [0, total_bs), sentinels included
+        self.assertTrue(
+            torch.equal(out_ref[:total_bs], out_fused[:total_bs]),
+            f"state indices mismatch ({case}):\n"
+            f"  ref   {out_ref[:total_bs].tolist()}\n"
+            f"  fused {out_fused[:total_bs].tolist()}",
+        )
+        # The returned view is what _replay_metadata forwards downstream.
+        self.assertTrue(
+            torch.equal(returned, out_fused[:total_bs]),
+            f"returned view is not the filled buffer ({case})",
+        )
+        # 2. req_pool_indices side effect bit-identical (padded rows zeroed)
+        self.assertTrue(
+            torch.equal(req_pool_ref[:total_bs], req_pool_fused[:total_bs]),
+            f"req_pool_indices mismatch ({case}):\n"
+            f"  ref   {req_pool_ref[:total_bs].tolist()}\n"
+            f"  fused {req_pool_fused[:total_bs].tolist()}",
+        )
+        # Explicit re-statement of the contract, independent of the reference:
+        self.assertTrue(
+            (req_pool_fused[valid_bs:total_bs] == 0).all(),
+            f"padded req_pool_indices rows not zeroed ({case})",
+        )
+        self.assertTrue(
+            (out_fused[valid_bs:total_bs] == -1).all(),
+            f"padding sentinel rows not -1 ({case})",
+        )
+        self.assertFalse(
+            (out_fused[:total_bs] == _OUT_POISON).any(),
+            f"unwritten cells inside [0, total_bs) ({case})",
+        )
+        # 3. no out-of-range writes past total_bs (BS_UPPER > total_bs masking)
+        for name, buf in (("req_pool", req_pool_fused), ("out", out_fused)):
+            expected = _GUARD_SENTINEL if name == "req_pool" else _OUT_POISON
+            self.assertTrue(
+                (buf[total_bs:] == expected).all(),
+                f"{name} guard tail clobbered ({case}): {buf[total_bs:].tolist()}",
+            )
+
+    def test_matrix(self):
+        # Non-power-of-two sizes (7, 33) exercise the BS_UPPER in_range mask;
+        # num_padding sweeps none / one / half / all-but-one padded rows.
+        for total_bs in (1, 2, 7, 32, 33):
+            paddings = sorted(
+                {0, 1, total_bs // 2, total_bs - 1} & set(range(total_bs))
+            )
+            for num_padding in paddings:
+                for seed in (0, 1, 2):
+                    with self.subTest(
+                        total_bs=total_bs, num_padding=num_padding, seed=seed
+                    ):
+                        self._run_case(
+                            total_bs=total_bs, num_padding=num_padding, seed=seed
+                        )
+
+    def test_shared_mamba_slots(self):
+        # Multiple requests mapping to the same mamba slot (mapping is not
+        # injective in general) must gather identically on both paths.
+        mapping_const = torch.full((_REQ_POOL_SIZE,), 3, dtype=torch.int32)
+        device = torch.device("cuda")
+        total_bs, num_padding = 7, 2
+        valid_bs = total_bs - num_padding
+        req_pool = torch.arange(total_bs + _GUARD, dtype=torch.int64)
+        out = torch.full((total_bs + _GUARD,), _OUT_POISON, dtype=torch.int32)
+
+        req_ref, req_fused = req_pool.to(device), req_pool.to(device)
+        out_ref, out_fused = out.to(device), out.to(device)
+        mapping_dev = mapping_const.to(device)
+
+        _reference_chain(
+            req_pool_indices=req_ref,
+            mapping=mapping_dev,
+            out_buf=out_ref,
+            valid_bs=valid_bs,
+            total_bs=total_bs,
+        )
+        fused_replay_state_indices(
+            req_pool_indices=req_fused,
+            mamba_index_mapping=mapping_dev,
+            out_state_indices=out_fused,
+            valid_bs=valid_bs,
+            total_bs=total_bs,
+        )
+        torch.cuda.synchronize()
+        self.assertTrue(torch.equal(out_ref, out_fused))
+        self.assertTrue(torch.equal(req_ref, req_fused))
+
+
+if __name__ == "__main__":
+    unittest.main()