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,