diff --git a/python/sglang/srt/layers/attention/linear/inkling_sconv_backend.py b/python/sglang/srt/layers/attention/linear/inkling_sconv_backend.py index 2130ff563..c59c05305 100644 --- a/python/sglang/srt/layers/attention/linear/inkling_sconv_backend.py +++ b/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): - """Per-(layer, step) conv-state handle handed to Inkling's conv kernels. - - ``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. +class InklingShortConvMetadata(msgspec.Struct): + """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: Any - cache_indices: torch.Tensor # per-request slot ids, int32 + cache_indices: Optional[torch.Tensor] = None # 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._precomputed: Optional[SconvExtendMetadata | SconvDecodeMetadata] = None - self._track_conv_indices: Optional[torch.Tensor] = None + self.sconv_metadata = InklingShortConvMetadata() 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._precomputed = None - self._track_conv_indices = None + self.sconv_metadata = InklingShortConvMetadata() @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, - self._has_initial_state, - self._precomputed, + md.query_start_loc, + md.has_initial_state, + 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 - self._has_initial_state = has_initial_state - self._precomputed = precomputed + md = self.sconv_metadata + md.query_start_loc = query_start_loc + 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 - shares one gather, one fused launch and one track-index build.""" - del forward_batch - return InklingShortConvMetadata( - layer_cache=self.req_to_token_pool.mamba2_layer_cache(layer_id), - cache_indices=self._cache_indices, - query_start_loc=self._query_start_loc, - has_initial_state=self._has_initial_state, - precomputed=self._precomputed, - track_conv_indices=self._track_conv_indices, - ) + """The step's metadata: resolved once during prep, so this is a pure read.""" + del layer_id, forward_batch + return self.sconv_metadata + + def sconv_state(self, *, layer_id: int, stream: int) -> torch.Tensor: + """``layer_id``'s conv state for one ``SconvType`` stream.""" + pool_layer = self.req_to_token_pool.mamba2_layer_index(layer_id) + return self._mamba_cache.conv[stream][pool_layer] + + 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 - :meth:`conv_state_metadata`. Four departures from + full-attention child and the sidecar is reached only through its metadata and + 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 diff --git a/python/sglang/srt/mem_cache/memory_pool.py b/python/sglang/srt/mem_cache/memory_pool.py index 24b5a8058..c06a72889 100644 --- a/python/sglang/srt/mem_cache/memory_pool.py +++ b/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 - # 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 {} + return self.mamba_cache.at_layer_idx(layer_id) + # 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() diff --git a/python/sglang/srt/models/inkling_common/sconv.py b/python/sglang/srt/models/inkling_common/sconv.py index ae9a2cc62..33e726658 100644 --- a/python/sglang/srt/models/inkling_common/sconv.py +++ b/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. - - ``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. - """ + """The step's conv-state metadata, resolved once by the attention backend.""" return get_attn_backend().conv_state_metadata(self.layer_id, forward_batch) - def _sconv_cache(self, meta) -> torch.Tensor: - return meta.layer_cache.conv[self.sconv_type.value] + def _sconv_cache(self) -> torch.Tensor: + """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(meta), + sconv_cache=self._sconv_cache(), 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, diff --git a/test/registered/unit/models/test_inkling_sconv_metadata_once.py b/test/registered/unit/models/test_inkling_sconv_metadata_once.py deleted file mode 100644 index 1becff702..000000000 --- a/test/registered/unit/models/test_inkling_sconv_metadata_once.py +++ /dev/null @@ -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()