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

[Mamba][GDN] Deduplicate spec conv-window intermediate cache via sliding window layout (#28302)

Browse Source 
Co-authored-by: luoyuan.luo <luoyuan.luo@antgroup.com>
This commit is contained in:
Yuan Luo
2026-06-17 20:42:31 -07:00
committed by GitHub
co-authored by luoyuan.luo
parent 3340f4e3da
commit 3b9db3a1f0
5 changed files with 283 additions and 20 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/disaggregation/decode.py
+2
View File
@@ -195,6 +195,7 @@ class HybridMambaDecodeReqToTokenPool(HybridReqToTokenPool):
enable_overlap_schedule: bool,
mamba_size: int = None,
start_layer: int = None,
speculative_eagle_topk: Optional[int] = None,
):
DecodeReqToTokenPool.__init__(
self,
@@ -238,6 +239,7 @@ class HybridMambaDecodeReqToTokenPool(HybridReqToTokenPool):
device=device,
enable_mamba_extra_buffer=self.enable_mamba_extra_buffer,
speculative_num_draft_tokens=speculative_num_draft_tokens,
speculative_eagle_topk=speculative_eagle_topk,
)
def clear(self):
python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py
+5 -2
View File
@@ -11,6 +11,7 @@ from sglang.srt.layers.attention.mamba.mamba2_metadata import (
Mamba2Metadata,
)
from sglang.srt.layers.attention.mamba.mamba_state_scatter_triton import (
fused_conv_window_scatter_with_mask,
fused_mamba_state_scatter_with_mask,
track_mamba_states_if_needed,
)
@@ -922,7 +923,9 @@ class HybridLinearAttnBackend(AttentionBackend):
state_indices_tensor,
last_correct_step_indices,
)
fused_mamba_state_scatter_with_mask(
# conv intermediate uses the deduplicated sliding-window (overlapping)
# layout, so it needs the strided-read scatter variant.
fused_conv_window_scatter_with_mask(
conv_states,
intermediate_conv_window_cache,
state_indices_tensor,
@@ -939,7 +942,7 @@ class HybridLinearAttnBackend(AttentionBackend):
mamba_track_indices,
mamba_steps_to_track,
)
fused_mamba_state_scatter_with_mask(
fused_conv_window_scatter_with_mask(
conv_states,
intermediate_conv_window_cache,
mamba_track_indices,
python/sglang/srt/layers/attention/mamba/mamba_state_scatter_triton.py
+153
View File
@@ -286,3 +286,156 @@ def fused_mamba_state_scatter_with_mask(
dst_req_size,
BLOCK_SIZE=BLOCK_SIZE,
)
@triton.jit
def _fused_conv_window_scatter_with_mask_kernel(
src_ptr,
dst_ptr,
dst_indices_raw_ptr, # [total_requests]
step_indices_raw_ptr, # [total_requests], entry >= 0 means valid
elem_per_entry: tl.constexpr, # dim * (K-1)
KM1: tl.constexpr, # K-1 (conv window width)
src_layer_stride,
src_req_stride,
src_step_stride,
src_dim_stride,
src_win_stride,
dst_layer_stride,
dst_req_stride,
src_req_size,
src_step_size,
dst_req_size,
BLOCK_SIZE: tl.constexpr,
):
"""Scatter accepted conv windows from the deduplicated sliding-window source.
Unlike ``_fused_mamba_state_scatter_with_mask_kernel`` (which flat-copies a
contiguous per-step state row), the source here is an *overlapping* view: the
deduplicated layout keeps one shared ``[dim, D+K-2]`` buffer per (layer, slot)
and step ``t``'s window is the slice ``shared[:, t:t+K-1]``. That window is
non-contiguous, so we index every ``(dim, win)`` element through the view's
strides (``src_step_stride`` / ``src_dim_stride`` / ``src_win_stride``). The
destination conv-state row stays contiguous in ``(dim, K-1)`` order.
"""
pid_req = tl.program_id(0)
pid_layer = tl.program_id(1).to(tl.int64)
pid_block = tl.program_id(2).to(tl.int64)
step_idx = tl.load(step_indices_raw_ptr + pid_req).to(tl.int64)
if step_idx < 0:
return
dst_idx = tl.load(dst_indices_raw_ptr + pid_req).to(tl.int64)
src_idx = pid_req
if not (
(dst_idx >= 0)
& (dst_idx < dst_req_size)
& (src_idx < src_req_size)
& (step_idx < src_step_size)
):
return
start = pid_block * BLOCK_SIZE
e = start + tl.arange(0, BLOCK_SIZE)
mask = e < elem_per_entry
# Decode the flat (dim, K-1)-row element index into (dim, win) coordinates.
d = e // KM1
w = e % KM1
src_off = (
pid_layer * src_layer_stride
+ src_idx * src_req_stride
+ step_idx * src_step_stride
+ d * src_dim_stride
+ w * src_win_stride
)
# dst window is contiguous in (dim, K-1) order -> flat element index `e`.
dst_off = pid_layer * dst_layer_stride + dst_idx * dst_req_stride + e
data = tl.load(src_ptr + src_off, mask=mask, other=0.0)
tl.store(dst_ptr + dst_off, data, mask=mask)
def fused_conv_window_scatter_with_mask(
dst: torch.Tensor, # conv_states [num_layers, cache_size, dim, K-1] (contiguous)
src: torch.Tensor, # deduped conv-window view [num_layers, spec_size, draft_tokens, dim, K-1]
dst_indices_raw: torch.Tensor, # [total_requests]
step_indices_raw: torch.Tensor, # [total_requests], entry >= 0 means valid
):
"""Conv-window variant of :func:`fused_mamba_state_scatter_with_mask`.
``src`` is the deduplicated sliding-window conv-intermediate cache: an
overlapping ``as_strided`` view over a shared ``[..., dim, D+K-2]`` buffer,
so its per-step windows are intentionally non-contiguous. This kernel indexes
``(dim, win)`` elements through the view's strides instead of flat-copying.
``dst`` (the real conv-state pool) is the usual contiguous
``[layers, cache, dim, K-1]``.
"""
total_requests = step_indices_raw.shape[0]
if total_requests == 0:
return
if not (dst.is_cuda and src.is_cuda and dst.device == src.device):
raise ValueError(
"fused_conv_window_scatter_with_mask requires dst and src to be CUDA "
f"tensors on the same device ({dst.device=}, {src.device=})."
)
if dst.ndim != 4 or src.ndim != 5:
raise ValueError(f"Unexpected ranks: {dst.ndim=} (want 4) {src.ndim=} (want 5)")
if dst.shape[0] != src.shape[0]:
raise ValueError(f"Layer dim mismatch: {dst.shape[0]=} vs {src.shape[0]=}")
if dst.shape[2:] != src.shape[3:]:
raise ValueError(f"Window dims mismatch: {dst.shape[2:]=} vs {src.shape[3:]=}")
if dst_indices_raw.ndim != 1 or step_indices_raw.ndim != 1:
raise ValueError(
f"indices must be 1D: {dst_indices_raw.shape=} {step_indices_raw.shape=}"
)
if dst_indices_raw.shape[0] != step_indices_raw.shape[0]:
raise ValueError(
f"indices length mismatch: {dst_indices_raw.shape[0]=} vs {step_indices_raw.shape[0]=}"
)
num_layers = dst.shape[0]
dim = dst.shape[2]
km1 = dst.shape[3]
elem_per_entry = dim * km1
src_req_size = src.shape[1]
src_step_size = src.shape[2]
dst_req_size = dst.shape[1]
# `dst` stays contiguous; `src` is an intentionally non-contiguous (overlapping)
# view, so we do NOT assert src contiguity here (unlike the dense scatter).
if not dst.is_contiguous():
raise ValueError(
"dst tensor in fused_conv_window_scatter_with_mask must be contiguous"
)
dst_indices_raw = dst_indices_raw.to(torch.int32).contiguous()
step_indices_raw = step_indices_raw.to(torch.int32).contiguous()
BLOCK_SIZE = 1024
grid = (total_requests, num_layers, triton.cdiv(elem_per_entry, BLOCK_SIZE))
_fused_conv_window_scatter_with_mask_kernel[grid](
src,
dst,
dst_indices_raw,
step_indices_raw,
elem_per_entry,
km1,
src.stride(0),
src.stride(1),
src.stride(2),
src.stride(3),
src.stride(4),
dst.stride(0),
dst.stride(1),
src_req_size,
src_step_size,
dst_req_size,
BLOCK_SIZE=BLOCK_SIZE,
)
python/sglang/srt/mem_cache/memory_pool.py
+121 -18
View File
@@ -90,6 +90,26 @@ _is_fp8_fnuz = is_fp8_fnuz()
_use_aiter = bool(envs.SGLANG_USE_AITER.get()) and _is_hip
def conv_window_dedup_enabled(
is_npu: bool, is_cpu: bool, speculative_eagle_topk: Optional[int]
) -> bool:
"""Whether the deduplicated sliding-window conv-intermediate layout is safe.
It is only correct for a *linear* draft chain (``speculative_eagle_topk <= 1``,
i.e. NEXTN / MTP): consecutive draft tokens then form a true sliding window, so
the overlapping physical columns hold identical values. Under EAGLE *tree*
verify (``topk > 1``) the conv kernel walks per-token tree ancestors, so aliased
columns can need different values from different parent chains -> fall back to
the dense layout. NPU/CPU also keep the dense layout (their kernels assume
contiguous per-step windows). See ``MambaPool.__init__``.
"""
return (
not is_npu
and not is_cpu
and (speculative_eagle_topk is None or speculative_eagle_topk <= 1)
)
def get_tensor_size_bytes(t: Union[torch.Tensor, List[torch.Tensor]]):
if isinstance(t, list):
return sum(get_tensor_size_bytes(x) for x in t)
@@ -316,6 +336,7 @@ class MambaPool:
device: str,
enable_memory_saver: bool = False,
speculative_num_draft_tokens: Optional[int] = None,
speculative_eagle_topk: Optional[int] = None,
):
conv_state_shape = cache_params.shape.conv
temporal_state_shape = cache_params.shape.temporal
@@ -393,22 +414,86 @@ class MambaPool:
dtype=ssm_dtype,
device="cuda",
)
# Cache intermediate conv windows (last K-1 inputs) per draft token during target verify
# Shape: [num_layers, size + 1, speculative_num_draft_tokens, dim, K-1]
intermediate_conv_window_cache = [
torch.zeros(
size=(
num_mamba_layers,
spec_state_size + 1,
speculative_num_draft_tokens,
conv_shape[0],
conv_shape[1],
),
dtype=conv_dtype,
device="cuda",
)
for conv_shape in conv_state_shape
]
# Cache intermediate conv windows (last K-1 inputs) per draft token
# during target verify.
#
# On CUDA (Triton conv kernel + Triton scatter) we use a
# *deduplicated sliding-window* layout: consecutive draft tokens'
# (K-1)-wide windows overlap by (K-2), so instead of D separate
# [dim, K-1] windows we store one shared [dim, D+K-2] buffer per
# (layer, slot) and expose an overlapping `as_strided` view of
# logical shape [num_layers, size+1, draft_tokens, dim, K-1] where
# step `t`'s window is the slice shared[..., :, t:t+K-1]. This
# halves the conv-intermediate footprint (D*(K-1) -> D+K-2 columns)
# with no numerical change: both the conv kernel write (idempotent
# overlapping stores) and `fused_conv_window_scatter_with_mask`
# consume the view through its strides.
#
# Dedup the sliding-window conv-intermediate only when it is safe:
# CUDA + a linear draft chain (topk <= 1). NPU/CPU and EAGLE tree
# verify (topk > 1) keep the dense layout -- see
# `conv_window_dedup_enabled` for the full rationale. The
# `fused_conv_window_scatter_with_mask` scatter is layout-agnostic,
# so the dense fallback reads correctly through the same code path.
dedup_conv_window = conv_window_dedup_enabled(
_is_npu, _is_cpu, speculative_eagle_topk
)
self._intermediate_conv_window_phys = []
if dedup_conv_window:
intermediate_conv_window_cache = []
for conv_shape in conv_state_shape:
conv_dim, win = conv_shape # win == conv_kernel - 1 == K-1
shared_win = (
speculative_num_draft_tokens + win - 1
) # D + (K-1) - 1
phys = torch.zeros(
size=(
num_mamba_layers,
spec_state_size + 1,
conv_dim,
shared_win,
),
dtype=conv_dtype,
device="cuda",
)
# view[l, s, step, d, w] = phys[l, s, d, step + w]
view = phys.as_strided(
(
phys.shape[0],
phys.shape[1],
speculative_num_draft_tokens,
conv_dim,
win,
),
(
phys.stride(0),
phys.stride(1),
phys.stride(3), # step -> shared-win axis (stride 1)
phys.stride(2), # dim
phys.stride(3), # win -> shared-win axis (stride 1)
),
)
self._intermediate_conv_window_phys.append(phys)
intermediate_conv_window_cache.append(view)
else:
# Original dense layout (NPU/CPU, or EAGLE tree verify): one
# [dim, K-1] window per draft token.
# Shape: [num_layers, size+1, draft_tokens, dim, K-1]
intermediate_conv_window_cache = [
torch.zeros(
size=(
num_mamba_layers,
spec_state_size + 1,
speculative_num_draft_tokens,
conv_shape[0],
conv_shape[1],
),
dtype=conv_dtype,
device="cuda",
)
for conv_shape in conv_state_shape
]
self._intermediate_conv_window_phys = intermediate_conv_window_cache
self.mamba_cache = self.SpeculativeState(
conv=conv_state,
temporal=temporal_state,
@@ -421,7 +506,9 @@ class MambaPool:
f"conv_state size: {get_tensor_size_bytes(conv_state) / GB:.2f}GB, "
f"ssm_state size: {get_tensor_size_bytes(temporal_state) / GB:.2f}GB "
f"intermediate_ssm_state_cache size: {get_tensor_size_bytes(intermediate_ssm_state_cache) / GB:.2f}GB "
f"intermediate_conv_window_cache size: {get_tensor_size_bytes(intermediate_conv_window_cache) / GB:.2f}GB "
# Report the deduplicated PHYSICAL conv-window buffers (the view
# over-reports its logical, un-deduplicated size).
f"intermediate_conv_window_cache size: {get_tensor_size_bytes(self._intermediate_conv_window_phys) / GB:.2f}GB "
)
else:
self.mamba_cache = self.State(conv=conv_state, temporal=temporal_state)
@@ -431,7 +518,19 @@ class MambaPool:
f"conv_state size: {get_tensor_size_bytes(conv_state) / GB:.2f}GB, "
f"ssm_state size: {get_tensor_size_bytes(temporal_state) / GB:.2f}GB "
)
self.mem_usage = self.mamba_cache.mem_usage_bytes() / GB
mem_usage_bytes = self.mamba_cache.mem_usage_bytes()
if isinstance(self.mamba_cache, self.SpeculativeState):
# `intermediate_conv_window` is an as_strided view whose logical
# shape over-reports its real footprint; charge the physical buffers
# instead. No-op for the dense layout, where the view and the
# physical tensors coincide.
mem_usage_bytes -= get_tensor_size_bytes(
self.mamba_cache.intermediate_conv_window
)
mem_usage_bytes += get_tensor_size_bytes(
self._intermediate_conv_window_phys
)
self.mem_usage = mem_usage_bytes / GB
self.num_mamba_layers = num_mamba_layers
def get_speculative_mamba2_params_all_layers(self) -> SpeculativeState:
@@ -561,6 +660,7 @@ class HybridReqToTokenPool(ReqToTokenPool):
enable_mamba_extra_buffer: bool,
enable_mamba_extra_buffer_lazy: bool = False,
speculative_num_draft_tokens: int = None,
speculative_eagle_topk: Optional[int] = None,
enable_overlap_schedule: bool = True,
start_layer: Optional[int] = None,
):
@@ -585,6 +685,7 @@ class HybridReqToTokenPool(ReqToTokenPool):
device=device,
enable_mamba_extra_buffer=enable_mamba_extra_buffer,
speculative_num_draft_tokens=speculative_num_draft_tokens,
speculative_eagle_topk=speculative_eagle_topk,
)
def _init_mamba_pool(
@@ -596,6 +697,7 @@ class HybridReqToTokenPool(ReqToTokenPool):
device: str,
enable_mamba_extra_buffer: bool,
speculative_num_draft_tokens: int = None,
speculative_eagle_topk: Optional[int] = None,
):
self.mamba_pool = MambaPool(
size=mamba_size,
@@ -605,6 +707,7 @@ class HybridReqToTokenPool(ReqToTokenPool):
device=device,
enable_memory_saver=self.enable_memory_saver,
speculative_num_draft_tokens=speculative_num_draft_tokens,
speculative_eagle_topk=speculative_eagle_topk,
)
self.mamba_allocator = MambaSlotAllocator(
size=mamba_size,
python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py
+2
View File
@@ -350,6 +350,7 @@ class ModelRunnerKVCacheMixin:
]
),
speculative_num_draft_tokens=max_spec_draft_tokens,
speculative_eagle_topk=self.server_args.speculative_eagle_topk,
enable_mamba_extra_buffer=self.server_args.enable_mamba_extra_buffer(),
pre_alloc_size=pre_alloc_size,
enable_overlap_schedule=not self.server_args.disable_overlap_schedule,
@@ -385,6 +386,7 @@ class ModelRunnerKVCacheMixin:
enable_mamba_extra_buffer=self.server_args.enable_mamba_extra_buffer(),
enable_mamba_extra_buffer_lazy=self.server_args.enable_mamba_extra_buffer_lazy(),
speculative_num_draft_tokens=max_spec_draft_tokens,
speculative_eagle_topk=self.server_args.speculative_eagle_topk,
enable_overlap_schedule=not self.server_args.disable_overlap_schedule,
start_layer=self.start_layer,
)