minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Cheng Wan
2026-07-31 18:11:35 -07:00
committed by GitHub
parent 58974ca16c
commit 934a13ce3e
4 changed files with 80 additions and 424 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/attention/linear/inkling_sconv_backend.py
+47 -37
View File
@@ -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
python/sglang/srt/mem_cache/memory_pool.py
+14 -16
View File
@@ -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()
python/sglang/srt/models/inkling_common/sconv.py
+19 -19
View File
@@ -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,
test/registered/unit/models/test_inkling_sconv_metadata_once.py
-352
View File
@@ -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()