[Inkling] Hold the short-conv per-step state on one metadata struct (#33116)

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

@@ -15,7 +15,8 @@
 
 A :mod:`~sglang.srt.layers.attention.linear.short_conv_backend` sidecar. Four short
 convs per decoder layer keep per-request conv state in the centralized
-``MambaPool``; the model reaches this via :meth:`conv_state_metadata`, never
+``MambaPool``; the model reaches this via :meth:`conv_state_metadata` for the
+step's metadata and :meth:`sconv_state` for a layer's own conv stream, never
 through ``forward_decode`` / ``forward_extend``.
 
 On top of what :class:`ShortConvAttnBackend` owns, Inkling's kernels take a
@@ -34,8 +35,9 @@ tensor a captured kernel reads lives in a graph-static buffer refilled in place.
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Any, NamedTuple, Optional
+from typing import TYPE_CHECKING, Optional
 
+import msgspec
 import torch
 
 from sglang.kernels.ops.mamba.mamba_state_scatter_triton import (
@@ -67,15 +69,12 @@ if TYPE_CHECKING:
     from sglang.srt.model_executor.model_runner import ModelRunner
 
 
-class InklingShortConvMetadata(NamedTuple):
+class InklingShortConvMetadata(msgspec.Struct):
-    """Per-(layer, step) conv-state handle handed to Inkling's conv kernels.
+    """The step's conv-state metadata, filled during metadata prep. On the graph
-
+    path every tensor here is a static buffer refilled in place.
-    ``layer_cache`` holds this layer's pool views indexed by ``SconvType``; the
-    rest is step-global, and on the graph path is a static buffer refilled in place.
     """
 
-    layer_cache: Any
+    cache_indices: Optional[torch.Tensor] = None  # per-request slot ids, int32
-    cache_indices: torch.Tensor  # per-request slot ids, int32
     query_start_loc: Optional[torch.Tensor] = None  # cu-seqlens, int32
     has_initial_state: Optional[torch.Tensor] = None  # "resumes a cached prefix"
     precomputed: Optional[SconvExtendMetadata | SconvDecodeMetadata] = None
@@ -96,7 +95,9 @@ class InklingShortConvAttnBackend(ShortConvAttnBackend):
 
     def __init__(self, model_runner: ModelRunner):
         super().__init__(model_runner)
-        # conv[i] is [n_layers, n_slots, conv_kernel - 1, conv_dim].
+        # Pool-wide, bound at pool construction: conv[stream] is
+        # [n_layers, n_slots, conv_kernel - 1, conv_dim].
+        self._mamba_cache = self.req_to_token_pool.mamba_pool.mamba_cache
         self.conv_state_len: int = self.conv_states_shape[2]
         self.mamba_cache_chunk_size = get_server_args().mamba_cache_chunk_size
         # A plain table lookup is recordable; the unified pool's translate is an
@@ -106,9 +107,7 @@ class InklingShortConvAttnBackend(ShortConvAttnBackend):
             is HybridReqToTokenPool.translate_mamba_indices
         )
 
-        self._query_start_loc: Optional[torch.Tensor] = None
+        self.sconv_metadata = InklingShortConvMetadata()
-        self._precomputed: Optional[SconvExtendMetadata | SconvDecodeMetadata] = None
-        self._track_conv_indices: Optional[torch.Tensor] = None
 
         self._alloc_graph_buffers()
 
@@ -186,9 +185,7 @@ class InklingShortConvAttnBackend(ShortConvAttnBackend):
 
     def _reset_step_state(self):
         super()._reset_step_state()
-        self._query_start_loc = None
+        self.sconv_metadata = InklingShortConvMetadata()
-        self._precomputed = None
-        self._track_conv_indices = None
 
     @staticmethod
     def _phase_records_metadata(forward_batch: ForwardBatch) -> bool:
@@ -264,6 +261,7 @@ class InklingShortConvAttnBackend(ShortConvAttnBackend):
     ):
         if self._cache_indices is None:
             return
+        self.sconv_metadata.cache_indices = self._cache_indices
         mode = forward_batch.forward_mode
         if mode.is_decode_or_idle():
             self._refresh_decode_metadata(forward_batch, on_graph_path)
@@ -279,10 +277,11 @@ class InklingShortConvAttnBackend(ShortConvAttnBackend):
         self, forward_batch: ForwardBatch, on_graph_path: bool
     ):
         B = forward_batch.batch_size
+        md = self.sconv_metadata
         (
-            self._query_start_loc,
+            md.query_start_loc,
-            self._has_initial_state,
+            md.has_initial_state,
-            self._precomputed,
+            md.precomputed,
         ) = fused_decode_sconv_metadata(
             B=B,
             cache_indices=self._cache_indices,
@@ -365,9 +364,10 @@ class InklingShortConvAttnBackend(ShortConvAttnBackend):
                 cu=precomputed["cu"],
                 si=precomputed["si"][:T],
             )
-        self._query_start_loc = query_start_loc
+        md = self.sconv_metadata
-        self._has_initial_state = has_initial_state
+        md.query_start_loc = query_start_loc
-        self._precomputed = precomputed
+        md.has_initial_state = has_initial_state
+        md.precomputed = precomputed
 
     def _unfused_extend_metadata(self, forward_batch: ForwardBatch):
         """Unfused query_start_loc / has_initial_state prep; fallback only."""
@@ -422,7 +422,7 @@ class InklingShortConvAttnBackend(ShortConvAttnBackend):
         if forward_batch.mamba_track_mask is None:
             return
         rows = forward_batch.batch_size
-        query_start_loc = self._query_start_loc
+        query_start_loc = self.sconv_metadata.query_start_loc
         live = min(
             rows,
             forward_batch.mamba_track_seqlens.shape[0],
@@ -464,7 +464,7 @@ class InklingShortConvAttnBackend(ShortConvAttnBackend):
         )
         if live < rows:
             out[live:].zero_()
-        self._track_conv_indices = out
+        self.sconv_metadata.track_conv_indices = out
 
     def commit_conv_state_after_mtp_verify(
         self,
@@ -492,25 +492,27 @@ class InklingShortConvAttnBackend(ShortConvAttnBackend):
     def conv_state_metadata(
         self, layer_id: int, forward_batch: ForwardBatch
     ) -> InklingShortConvMetadata:
-        """``layer_id``'s handle for this step: a pure read, so every conv layer
+        """The step's metadata: resolved once during prep, so this is a pure read."""
-        shares one gather, one fused launch and one track-index build."""
+        del layer_id, forward_batch
-        del forward_batch
+        return self.sconv_metadata
-        return InklingShortConvMetadata(
+
-            layer_cache=self.req_to_token_pool.mamba2_layer_cache(layer_id),
+    def sconv_state(self, *, layer_id: int, stream: int) -> torch.Tensor:
-            cache_indices=self._cache_indices,
+        """``layer_id``'s conv state for one ``SconvType`` stream."""
-            query_start_loc=self._query_start_loc,
+        pool_layer = self.req_to_token_pool.mamba2_layer_index(layer_id)
-            has_initial_state=self._has_initial_state,
+        return self._mamba_cache.conv[stream][pool_layer]
-            precomputed=self._precomputed,
+
-            track_conv_indices=self._track_conv_indices,
+    def sconv_intermediate_window(self, *, layer_id: int, stream: int) -> torch.Tensor:
-        )
+        """One stream's per-draft-token conv windows. TARGET_VERIFY only."""
+        pool_layer = self.req_to_token_pool.mamba2_layer_index(layer_id)
+        return self._mamba_cache.intermediate_conv_window[stream][pool_layer]
 
 
 class InklingShortConvHybridAttnBackend(ShortConvHybridAttnBackend):
     """Full-attention backend plus Inkling's conv-state sidecar.
 
     Inkling has NO linear-attention layers, so every layer routes to the
-    full-attention child and the sidecar is reached only via
+    full-attention child and the sidecar is reached only through its metadata and
-    :meth:`conv_state_metadata`. Four departures from
+    conv-state accessors. Four departures from
     :class:`ShortConvHybridAttnBackend`: every layer is full attention (including
     the draft's, so the base's ``full_attn_layers = [0]`` does not hold);
     DRAFT_EXTEND_V2 still inits the sidecar (the draft runs its own convs, unlike
@@ -519,6 +521,14 @@ class InklingShortConvHybridAttnBackend(ShortConvHybridAttnBackend):
     is Inkling's own, not the generic mamba scatter.
     """
 
+    def sconv_state(self, *, layer_id: int, stream: int) -> torch.Tensor:
+        return self.short_conv_backend.sconv_state(layer_id=layer_id, stream=stream)
+
+    def sconv_intermediate_window(self, *, layer_id: int, stream: int) -> torch.Tensor:
+        return self.short_conv_backend.sconv_intermediate_window(
+            layer_id=layer_id, stream=stream
+        )
+
     def _is_full_attn(self, layer=None, layer_id: Optional[int] = None) -> bool:
         del layer, layer_id
         return True

python/sglang/srt/mem_cache/memory_pool.py

@@ -850,21 +850,11 @@ class MambaPool:
         return self.mamba_cache
 
     def mamba2_layer_cache(self, layer_id: int):
-        # The per-layer views are pool-stable (mamba_cache is only bound at
+        return self.mamba_cache.at_layer_idx(layer_id)
-        # construction), so each layer's State is built once.
-        cached = self._layer_cache_by_id.get(layer_id)
-        if cached is None:
-            cached = self.mamba_cache.at_layer_idx(layer_id)
-            self._layer_cache_by_id[layer_id] = cached
-        return cached
-
-    # These properties are pool-stable (conv tensors don't move after allocation)
-    # so they're cached per instance on first use. Defined as cached_property
-    # rather than set in __init__ because UnifiedMambaPool skips super().__init__.
-    @cached_property
-    def _layer_cache_by_id(self) -> dict:
-        return {}
 
+    # Pool-stable (conv tensors don't move after allocation) so cached per instance
+    # on first use. A cached_property rather than set in __init__ because
+    # UnifiedMambaPool skips super().__init__.
     @cached_property
     def _conv_fuse_ok(self) -> bool:
         """Whether clear/copy may use the fused kernel: CUDA bf16 contiguous conv.
@@ -1334,11 +1324,19 @@ class HybridReqToTokenPool(ReqToTokenPool):
         / get_cpu_copy / load_cpu_copy)."""
         return mamba_indices
 
-    def mamba2_layer_cache(self, layer_id: int):
+    def mamba2_layer_index(self, layer_id: int) -> int:
+        """Pool-side index of ``layer_id``'s state, gated on its HiCache transfer.
+
+        For a caller that wants one specific state tensor: it indexes the pool
+        tensor itself instead of taking a ``State`` sliced over every field.
+        """
         assert layer_id in self.mamba_map
         if self.layer_transfer_counter is not None:
             self.layer_transfer_counter.wait_until(layer_id - self.start_layer)
-        return self.mamba_pool.mamba2_layer_cache(self.mamba_map[layer_id])
+        return self.mamba_map[layer_id]
+
+    def mamba2_layer_cache(self, layer_id: int):
+        return self.mamba_pool.mamba2_layer_cache(self.mamba2_layer_index(layer_id))
 
     def get_speculative_mamba2_params_all_layers(self) -> MambaPool.SpeculativeState:
         return self.mamba_pool.get_speculative_mamba2_params_all_layers()

python/sglang/srt/models/inkling_common/sconv.py

@@ -7,7 +7,6 @@ import triton.language as tl
 from einops import rearrange
 from torch.nn.parameter import Parameter
 
-from sglang.srt.mem_cache.memory_pool import MambaPool
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch
 from sglang.srt.model_executor.forward_context import get_attn_backend
 from sglang.srt.models.inkling_common.kernels.sconv import (
@@ -123,16 +122,20 @@ class ShortConvolution(nn.Module):
         param_data.copy_(loaded_weight)
 
     def _conv_state(self, forward_batch: ForwardBatch):
-        """This layer's conv-state handle for the current step.
+        """The step's conv-state metadata, resolved once by the attention backend."""
-
-        ``InklingShortConvAttnBackend`` resolved the whole step-global metadata set
-        once during metadata prep, so this is a pure read shared by every conv
-        module in the step.
-        """
         return get_attn_backend().conv_state_metadata(self.layer_id, forward_batch)
 
-    def _sconv_cache(self, meta) -> torch.Tensor:
+    def _sconv_cache(self) -> torch.Tensor:
-        return meta.layer_cache.conv[self.sconv_type.value]
+        """This module's own conv-state stream for this layer."""
+        return get_attn_backend().sconv_state(
+            layer_id=self.layer_id, stream=self.sconv_type.value
+        )
+
+    def _intermediate_window(self) -> torch.Tensor:
+        """This module's per-draft-token conv windows. TARGET_VERIFY only."""
+        return get_attn_backend().sconv_intermediate_window(
+            layer_id=self.layer_id, stream=self.sconv_type.value
+        )
 
     def _weight_2d(self) -> torch.Tensor:
         return rearrange(self.weight, "d 1 w -> d w")
@@ -187,7 +190,6 @@ class ShortConvolution(nn.Module):
     def _save_intermediate_conv_windows(
         self,
         forward_batch: ForwardBatch,
-        cache: MambaPool.SpeculativeState,
         sconv_cache: torch.Tensor,
         cache_indices: torch.Tensor,
         hidden_states: torch.Tensor,
@@ -204,7 +206,7 @@ class ShortConvolution(nn.Module):
             sconv_cache=sconv_cache,
             hidden_states=hidden_states,
             cache_indices=cache_indices,
-            intermediate_out=cache.intermediate_conv_window[self.sconv_type.value],
+            intermediate_out=self._intermediate_window(),
             batch_size=forward_batch.batch_size,
             draft_token_num=forward_batch.spec_info.draft_token_num,
         )
@@ -275,7 +277,7 @@ class ShortConvolution(nn.Module):
         ``(sconv_cache, cache_indices, cache_mask, weight_2d)``."""
         meta = self._conv_state(forward_batch)
         return (
-            self._sconv_cache(meta),
+            self._sconv_cache(),
             meta.cache_indices,
             meta.precomputed["cache_mask"],
             self._weight_2d(),
@@ -289,11 +291,11 @@ class ShortConvolution(nn.Module):
         meta = self._conv_state(forward_batch)
         b = forward_batch.batch_size
         return (
-            self._sconv_cache(meta),
+            self._sconv_cache(),
             meta.cache_indices[:b],
             meta.has_initial_state,
             self._weight_2d(),
-            meta.layer_cache.intermediate_conv_window[self.sconv_type.value],
+            self._intermediate_window(),
         )
 
     def extend_fused_ar_inputs(self, forward_batch: ForwardBatch):
@@ -322,7 +324,7 @@ class ShortConvolution(nn.Module):
             track_mask = torch.empty((0,), dtype=torch.bool, device=dev)
             track_dst = torch.empty((0,), dtype=torch.int64, device=dev)
         return (
-            self._sconv_cache(meta),
+            self._sconv_cache(),
             precomputed["safe_idx"],
             precomputed["cache_mask"].view(-1),
             precomputed["cu"],
@@ -349,8 +351,7 @@ class ShortConvolution(nn.Module):
         meta = self._conv_state(forward_batch)
         self._save_intermediate_conv_windows(
             forward_batch=forward_batch,
-            cache=meta.layer_cache,
+            sconv_cache=self._sconv_cache(),
-            sconv_cache=self._sconv_cache(meta),
             cache_indices=cache_indices,
             hidden_states=x_scratch,
         )
@@ -375,7 +376,7 @@ class ShortConvolution(nn.Module):
 
         meta = self._conv_state(forward_batch)
         cache_indices = meta.cache_indices
-        sconv_cache = self._sconv_cache(meta)
+        sconv_cache = self._sconv_cache()
         precomputed = meta.precomputed
         weight = self._weight_2d()
 
@@ -391,7 +392,6 @@ class ShortConvolution(nn.Module):
             )
             self._save_intermediate_conv_windows(
                 forward_batch=forward_batch,
-                cache=meta.layer_cache,
                 sconv_cache=sconv_cache,
                 cache_indices=cache_indices,
                 hidden_states=hidden_states,

test/registered/unit/models/test_inkling_sconv_metadata_once.py

@@ -1,352 +0,0 @@
-"""Inkling's short-conv metadata must be resolved exactly ONCE per forward step.
-
-A decoder layer holds FOUR ``ShortConvolution`` modules, and per-layer ownership
-would recompute the whole set once per module. Pinned here: one resolution per step
-however many modules ask, every module gets the *same* tensors, and the graph-path
-destinations stay address-stable across steps -- including across a later
-``init_cuda_graph_state``, where reallocating would move an address an
-already-captured prefill graph reads.
-"""
-
-import unittest
-from types import SimpleNamespace
-
-import torch
-
-from sglang.srt.model_executor.forward_batch_info import ForwardMode
-from sglang.srt.server_args import ServerArgs, set_global_server_args_for_scheduler
-from sglang.test.ci.ci_register import register_cuda_ci
-from sglang.test.test_utils import CustomTestCase
-
-register_cuda_ci(est_time=20, stage="base-b", runner_config="1-gpu-small")
-
-NUM_LAYERS = 4
-NUM_SCONV_STREAMS = 6  # pool-wide streams: k/v full, k/v local, attn, mlp
-NUM_MODULES_PER_LAYER = 4  # k_sconv, v_sconv, attn_sconv, mlp_sconv
-POOL_SLOTS = 32
-CONV_KERNEL = 4
-CONV_DIM = 8
-
-
-class _MockMambaPool:
-    enable_linear_replayssm = False
-
-    def __init__(self):
-        conv = [
-            torch.zeros(
-                (NUM_LAYERS, POOL_SLOTS + 1, CONV_KERNEL - 1, CONV_DIM),
-                dtype=torch.bfloat16,
-                device="cuda",
-            )
-            for _ in range(NUM_SCONV_STREAMS)
-        ]
-        self.mamba_cache = SimpleNamespace(conv=conv, temporal=None)
-
-    def mamba2_layer_cache(self, layer_id: int):
-        return SimpleNamespace(
-            conv=[c[layer_id] for c in self.mamba_cache.conv],
-            intermediate_conv_window=None,
-        )
-
-
-class _MockReqToTokenPool:
-    """The four methods the backend calls, plus ``size`` (its max-bs bound)."""
-
-    def __init__(self):
-        self.size = POOL_SLOTS
-        self.mamba_pool = _MockMambaPool()
-        self.req_index_to_mamba_index_mapping = torch.arange(
-            POOL_SLOTS + 1, dtype=torch.int32, device="cuda"
-        )
-        self.gather_calls = 0
-
-    def get_mamba_indices(self, req_indices: torch.Tensor) -> torch.Tensor:
-        self.gather_calls += 1
-        return self.req_index_to_mamba_index_mapping[req_indices]
-
-    def translate_mamba_indices(self, mamba_indices: torch.Tensor) -> torch.Tensor:
-        return mamba_indices
-
-    def mamba2_layer_cache(self, layer_id: int):
-        return self.mamba_pool.mamba2_layer_cache(layer_id)
-
-    def get_speculative_mamba2_params_all_layers(self):
-        return self.mamba_pool.mamba_cache
-
-
-def _decode_batch(bs: int):
-    return SimpleNamespace(
-        forward_mode=ForwardMode.DECODE,
-        batch_size=bs,
-        req_pool_indices=torch.arange(bs, dtype=torch.int64, device="cuda"),
-        seq_lens=torch.full((bs,), 64, dtype=torch.int64, device="cuda"),
-        spec_info=None,
-        mamba_track_mask=None,
-        mamba_track_seqlens=None,
-        mamba_track_indices=None,
-    )
-
-
-def _extend_batch(seq_lens):
-    bs = len(seq_lens)
-    lens = torch.tensor(seq_lens, dtype=torch.int64, device="cuda")
-    return SimpleNamespace(
-        forward_mode=ForwardMode.EXTEND,
-        batch_size=bs,
-        req_pool_indices=torch.arange(bs, dtype=torch.int64, device="cuda"),
-        seq_lens=lens,
-        extend_seq_lens=lens,
-        extend_prefix_lens=torch.zeros(bs, dtype=torch.int64, device="cuda"),
-        extend_num_tokens=int(sum(seq_lens)),
-        spec_info=None,
-        mamba_track_mask=torch.ones(bs, dtype=torch.bool, device="cuda"),
-        mamba_track_seqlens=lens,
-        mamba_track_indices=torch.arange(bs, dtype=torch.int64, device="cuda"),
-    )
-
-
-class TestInklingSconvMetadataOnce(CustomTestCase):
-    @classmethod
-    def setUpClass(cls):
-        if not torch.cuda.is_available():
-            raise unittest.SkipTest("Inkling's conv metadata kernels are CUDA-only.")
-        server_args = ServerArgs(
-            model_path="dummy",
-            page_size=1,
-            # Skips the model-config load in the Inkling prefill-graph default.
-            disable_prefill_cuda_graph=True,
-            disable_cuda_graph=True,
-        )
-        # Pre-seed the cached property so it does not reach for a real HF config.
-        server_args._mamba_cache_chunk_size = 64
-        set_global_server_args_for_scheduler(server_args)
-
-    def _build_backend(self):
-        from sglang.srt.layers.attention.linear.inkling_sconv_backend import (
-            InklingShortConvAttnBackend,
-        )
-
-        pool = _MockReqToTokenPool()
-        from sglang.srt.runtime_context import get_server_args
-
-        runner = SimpleNamespace(
-            device="cuda",
-            server_args=get_server_args(),
-            is_draft_worker=False,
-            req_to_token_pool=pool,
-            token_to_kv_pool=None,
-        )
-        return InklingShortConvAttnBackend(runner), pool
-
-    def _count_fused_calls(self, backend):
-        """Wrap the two fused metadata entry points with counters."""
-        import sglang.srt.layers.attention.linear.inkling_sconv_backend as mod
-
-        counts = {"decode": 0, "extend": 0}
-        real_decode = mod.fused_decode_sconv_metadata
-        real_extend = mod.fused_extend_sconv_metadata
-
-        def decode(*a, **kw):
-            counts["decode"] += 1
-            return real_decode(*a, **kw)
-
-        def extend(*a, **kw):
-            counts["extend"] += 1
-            return real_extend(*a, **kw)
-
-        mod.fused_decode_sconv_metadata = decode
-        mod.fused_extend_sconv_metadata = extend
-        self.addCleanup(setattr, mod, "fused_decode_sconv_metadata", real_decode)
-        self.addCleanup(setattr, mod, "fused_extend_sconv_metadata", real_extend)
-        return counts
-
-    def _drain_all_conv_modules(self, backend, forward_batch):
-        """Mimic every ShortConvolution in the model asking for its handle."""
-        handles = []
-        for layer_id in range(NUM_LAYERS):
-            for _module in range(NUM_MODULES_PER_LAYER):
-                handles.append(backend.conv_state_metadata(layer_id, forward_batch))
-        return handles
-
-    def test_decode_resolves_once_per_step(self):
-        backend, pool = self._build_backend()
-        counts = self._count_fused_calls(backend)
-        fb = _decode_batch(bs=3)
-
-        backend.init_forward_metadata(fb)
-        handles = self._drain_all_conv_modules(backend, fb)
-
-        self.assertEqual(counts["decode"], 1)
-        self.assertEqual(pool.gather_calls, 1)
-        self.assertEqual(len(handles), NUM_LAYERS * NUM_MODULES_PER_LAYER)
-        first = handles[0]
-        for h in handles[1:]:
-            self.assertIs(h.cache_indices, first.cache_indices)
-            self.assertIs(h.precomputed, first.precomputed)
-            self.assertIs(h.query_start_loc, first.query_start_loc)
-            self.assertIs(h.has_initial_state, first.has_initial_state)
-
-    def test_extend_resolves_once_per_step(self):
-        backend, pool = self._build_backend()
-        counts = self._count_fused_calls(backend)
-        fb = _extend_batch([7, 5, 3])
-
-        backend.init_forward_metadata(fb)
-        handles = self._drain_all_conv_modules(backend, fb)
-
-        self.assertEqual(counts["extend"], 1)
-        self.assertEqual(pool.gather_calls, 1)
-        first = handles[0]
-        self.assertIsNotNone(first.track_conv_indices)
-        self.assertEqual(tuple(first.track_conv_indices.shape), (3, CONV_KERNEL - 1))
-        for h in handles[1:]:
-            self.assertIs(h.track_conv_indices, first.track_conv_indices)
-            self.assertIs(h.precomputed, first.precomputed)
-
-    def test_each_step_re_resolves(self):
-        """A second forward must recompute; nothing may leak across steps."""
-        backend, pool = self._build_backend()
-        counts = self._count_fused_calls(backend)
-        fb = _decode_batch(bs=2)
-
-        for _ in range(3):
-            backend.init_forward_metadata(fb)
-            self._drain_all_conv_modules(backend, fb)
-
-        self.assertEqual(counts["decode"], 3)
-        self.assertEqual(pool.gather_calls, 3)
-
-    def test_graph_destinations_are_address_stable(self):
-        for slots_in_graph in (False, True):
-            with self.subTest(slots_in_graph=slots_in_graph):
-                self._check_address_stable(slots_in_graph)
-
-    def _check_address_stable(self, slots_in_graph: bool):
-        """A captured graph holds each metadata tensor's address, so steps refill in
-        place and a later ``init_cuda_graph_state`` must not reallocate."""
-        backend, _pool = self._build_backend()
-        # Cover both halves of the slot split (the mock's translate is not the base
-        # one, so slots would otherwise always stay eager).
-        backend._slot_gather_recordable = slots_in_graph
-        fb = _decode_batch(bs=2)
-
-        # Mirrors the decode runner: out-of-graph prep, then the recorded hook.
-        backend.init_forward_metadata_out_graph(fb, in_capture=True)
-        backend.init_forward_metadata_in_graph(fb)
-        h0 = backend.conv_state_metadata(0, fb)
-        ptrs = (
-            h0.cache_indices.data_ptr(),
-            h0.query_start_loc.data_ptr(),
-            h0.has_initial_state.data_ptr(),
-            h0.precomputed["cache_mask"].data_ptr(),
-            h0.precomputed["safe_idx"].data_ptr(),
-            h0.precomputed["cu"].data_ptr(),
-            h0.precomputed["si"].data_ptr(),
-        )
-
-        backend.init_cuda_graph_state(max_bs=8, max_num_tokens=8)
-        backend.init_forward_metadata_out_graph(fb)
-        backend.init_forward_metadata_in_graph(fb)
-        h1 = backend.conv_state_metadata(0, fb)
-        self.assertEqual(
-            ptrs,
-            (
-                h1.cache_indices.data_ptr(),
-                h1.query_start_loc.data_ptr(),
-                h1.has_initial_state.data_ptr(),
-                h1.precomputed["cache_mask"].data_ptr(),
-                h1.precomputed["safe_idx"].data_ptr(),
-                h1.precomputed["cu"].data_ptr(),
-                h1.precomputed["si"].data_ptr(),
-            ),
-        )
-
-
-class TestInklingMtpVerifyCommit(CustomTestCase):
-    """The commit runs after the forward context exits, so the per-step slot buffer
-    may already belong to a later forward. Sourcing slot ids from
-    ``forward_metadata`` (as the generic mamba path does) therefore mismatches the
-    verify batch; they must come from the passed ``req_pool_indices``.
-    """
-
-    @classmethod
-    def setUpClass(cls):
-        if not torch.cuda.is_available():
-            raise unittest.SkipTest("Inkling's conv-state kernels are CUDA-only.")
-        TestInklingSconvMetadataOnce.setUpClass()
-
-    def _build_wrapper(self):
-        from sglang.srt.layers.attention.linear.inkling_sconv_backend import (
-            InklingShortConvAttnBackend,
-            InklingShortConvHybridAttnBackend,
-        )
-        from sglang.srt.runtime_context import get_server_args
-
-        pool = _MockReqToTokenPool()
-        runner = SimpleNamespace(
-            device="cuda",
-            server_args=get_server_args(),
-            is_draft_worker=False,
-            req_to_token_pool=pool,
-            token_to_kv_pool=None,
-        )
-        sidecar = InklingShortConvAttnBackend(runner)
-        full = SimpleNamespace(
-            token_to_kv_pool=None,
-            req_to_token_pool=pool,
-            needs_cpu_seq_lens=True,
-        )
-        wrapper = InklingShortConvHybridAttnBackend(
-            full, sidecar, list(range(NUM_LAYERS))
-        )
-        return wrapper, sidecar, pool
-
-    def test_commit_uses_passed_req_pool_indices_not_step_metadata(self):
-        wrapper, sidecar, pool = self._build_wrapper()
-
-        # The hazard: a later forward left a SHORTER slot buffer than the verify
-        # batch this commit is for.
-        sidecar.init_forward_metadata(_decode_batch(bs=3))
-        self.assertEqual(sidecar._cache_indices.shape[0], 3)
-
-        seen = {}
-
-        def fake_scatter(caches, state_indices, last_correct, track, steps):
-            seen["state_indices"] = state_indices
-
-        import sglang.srt.layers.attention.linear.inkling_sconv_backend as mod
-
-        real = mod.scatter_mamba_states_after_mtp_verify
-        mod.scatter_mamba_states_after_mtp_verify = fake_scatter
-        self.addCleanup(setattr, mod, "scatter_mamba_states_after_mtp_verify", real)
-
-        req_pool_indices = torch.arange(5, dtype=torch.int64, device="cuda")
-        wrapper.update_mamba_state_after_mtp_verify(
-            last_correct_step_indices=torch.zeros(5, dtype=torch.int64, device="cuda"),
-            mamba_track_indices=None,
-            mamba_steps_to_track=None,
-            model=None,
-            req_pool_indices=req_pool_indices,
-        )
-        # 5 rows from req_pool_indices, not the 3 on the step buffer.
-        self.assertEqual(seen["state_indices"].shape[0], 5)
-        self.assertTrue(
-            torch.equal(seen["state_indices"], pool.get_mamba_indices(req_pool_indices))
-        )
-
-    def test_commit_requires_req_pool_indices(self):
-        """The generic caller signature makes it optional; Inkling cannot guess it."""
-        wrapper, _sidecar, _pool = self._build_wrapper()
-        with self.assertRaises(AssertionError):
-            wrapper.update_mamba_state_after_mtp_verify(
-                last_correct_step_indices=torch.zeros(
-                    2, dtype=torch.int64, device="cuda"
-                ),
-                mamba_track_indices=None,
-                mamba_steps_to_track=None,
-                model=None,
-            )
-
-
-if __name__ == "__main__":
-    unittest.main()