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

feat(mem_cache): unified memory pool for hybrid Mamba / SWA models (#29678)

Browse Source 
Co-authored-by: lch1475369 <lch1475369@gmail.com>
This commit is contained in:
Cheng Wan
2026-07-01 13:21:59 -07:00
committed by GitHub
co-authored by lch1475369
parent 3adfd0f34b
commit 4a8e76805c
28 changed files with 8378 additions and 551 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/environ.py
+9
View File
@@ -301,6 +301,15 @@ class Envs:
SGLANG_OTLP_EXPORTER_SCHEDULE_DELAY_MILLIS = EnvInt(500)
SGLANG_OTLP_EXPORTER_MAX_EXPORT_BATCH_SIZE = EnvInt(64)
SGLANG_NATIVE_MOVE_KV_CACHE = EnvBool(False)
# Disable lazy compaction in the unified memory pool allocator and
# fall back to the per-free eager compaction. Used for production
# A/B and quick rollback. Default False (lazy compaction on).
SGLANG_DISABLE_LAZY_COMPACTION = EnvBool(False)
# Sort the multi-ended allocator's free list after a merge (perf A/B knob).
SGLANG_SORT_FREE_LIST_AFTER_MERGE = EnvBool(False)
# Periodically log lazy-compaction stats per sub-pool (observability only).
SGLANG_LOG_LAZY_COMPACTION_STATS = EnvBool(False)
SGLANG_LOG_LAZY_COMPACTION_STATS_INTERVAL_SEC = EnvInt(30)
SGLANG_ENABLE_TP_MEMORY_INBALANCE_CHECK = EnvBool(True)
SGLANG_TEST_DISAGG_FAILURE_PROB = EnvFloat(0.0)
python/sglang/srt/hardware_backend/npu/memory_pool_npu.py
+2 -2
View File
@@ -179,7 +179,7 @@ class NPUMHATokenToKVPool(MHATokenToKVPool):
layer_id_override: Optional[int] = None,
dcp_kv_mask: Optional[torch.Tensor] = None,
):
loc, _ = unwrap_write_loc(loc_info)
loc, _, _ = unwrap_write_loc(loc_info)
if layer_id_override is not None:
layer_id = layer_id_override
else:
@@ -441,7 +441,7 @@ class NPUMLATokenToKVPool(MLATokenToKVPool):
cache_k: torch.Tensor,
cache_v: torch.Tensor,
):
loc, _ = unwrap_write_loc(loc_info)
loc, _, _ = unwrap_write_loc(loc_info)
layer_id = layer.layer_id
if cache_k.dtype != self.dtype:
cache_k = cache_k.to(self.dtype)
python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py
+125 -238
View File
@@ -35,15 +35,15 @@ class MambaAttnBackendBase(AttentionBackend):
self.is_draft_worker = model_runner.is_draft_worker
self.req_to_token_pool: HybridReqToTokenPool = model_runner.req_to_token_pool
self.token_to_kv_pool = model_runner.token_to_kv_pool
self.enable_unified_memory = model_runner.server_args.enable_unified_memory
self.forward_metadata: ForwardMetadata = None
self.state_indices_list = []
# GDN ReplaySSM (slice 1b): per-bs STATIC per-row write-cursor buffers
# for cuda-graph. Allocated lazily in init_cuda_graph_state only when
# --enable-linear-replayssm is set; stays None otherwise.
# Static (max_bs,) track-dest buffer captured by pointer, refreshed in-place
# each replay; the captured track-save reads this, not the InputBuffer slot.
self.mamba_track_indices_buf = None
# Per-bs static write-cursor / force-flush buffers for cuda-graph; None
# unless --enable-linear-replayssm is set.
self.replayssm_write_pos_list = None
# GDN ReplaySSM (slice 2b): per-bs STATIC per-row force-flush buffers
# for cuda-graph, parallel to replayssm_write_pos_list. Same lifetime
# (None unless the flag is on).
self.replayssm_force_flush_list = None
self.query_start_loc_list = []
self.retrieve_next_token_list = []
@@ -66,8 +66,10 @@ class MambaAttnBackendBase(AttentionBackend):
forward_batch.mamba_clear_indices is not None
and len(forward_batch.mamba_clear_indices) > 0
):
# mamba_pool is a pure PHYSICAL store; translate before zeroing or
# clear_slots zeroes the wrong physical slots.
self.req_to_token_pool.mamba_pool.clear_slots(
forward_batch.mamba_clear_indices
self._translate_mamba_indices(forward_batch.mamba_clear_indices)
)
if (
forward_batch.mamba_cow_src_indices is not None
@@ -75,22 +77,28 @@ class MambaAttnBackendBase(AttentionBackend):
):
ckpt_pool = getattr(self.req_to_token_pool, "mamba_ckpt_pool", None)
if ckpt_pool is not None:
# int8 checkpoints: dequantize the cached state (src = int8 ckpt slot)
# into the request's active bf16 slot (dst).
# int8 checkpoints: dequantize src int8 ckpt slot into the active bf16 dst.
ckpt_pool.load_to_active(
self.req_to_token_pool.mamba_pool,
forward_batch.mamba_cow_src_indices,
forward_batch.mamba_cow_dst_indices,
)
else:
# mamba_pool is a pure PHYSICAL store; translate both COW slot ids.
self.req_to_token_pool.mamba_pool.copy_from(
forward_batch.mamba_cow_src_indices,
forward_batch.mamba_cow_dst_indices,
self._translate_mamba_indices(forward_batch.mamba_cow_src_indices),
self._translate_mamba_indices(forward_batch.mamba_cow_dst_indices),
)
forward_batch.mamba_clear_indices = None
forward_batch.mamba_cow_src_indices = None
forward_batch.mamba_cow_dst_indices = None
def _translate_mamba_indices(self, mamba_indices: torch.Tensor) -> torch.Tensor:
"""Virtual->physical mamba slot-id translate (identity for the non-unified
pool). Must run everywhere mamba ids feed the SSM/conv kernels or mamba-pool
state ops, incl. the cuda-graph replay-prep copy into ``state_indices_list``."""
return self.req_to_token_pool.translate_mamba_indices(mamba_indices)
def _forward_metadata(self, forward_batch: ForwardBatch):
bs = forward_batch.batch_size
@@ -106,7 +114,14 @@ class MambaAttnBackendBase(AttentionBackend):
mamba_cache_indices = self.req_to_token_pool.get_mamba_indices(
forward_batch.req_pool_indices
)
_real_bs = getattr(forward_batch, "_original_batch_size", None)
# Translate virtual->physical BEFORE the padding sentinel below, so the
# gather reads only real ids; padded rows are then poisoned to -1 (skipped).
mamba_cache_indices = self._translate_mamba_indices(mamba_cache_indices)
if forward_batch.mamba_track_indices is not None:
forward_batch.mamba_track_indices = self._translate_mamba_indices(
forward_batch.mamba_track_indices
)
_real_bs = forward_batch._original_batch_size
if _real_bs is not None and _real_bs < mamba_cache_indices.shape[0]:
mamba_cache_indices = mamba_cache_indices.clone()
mamba_cache_indices[_real_bs:] = -1
@@ -117,11 +132,8 @@ class MambaAttnBackendBase(AttentionBackend):
query_start_loc = torch.arange(
0, bs + 1, dtype=torch.int32, device=self.device
)
# GDN ReplaySSM (slice 1a): the ring cursor is a per-slot
# decode-position counter shared by ALL GDN layers in this forward.
# Manage it exactly ONCE here (not per-layer): snapshot this step's
# value for the batch's slots, hand it to the layers, then advance
# the persistent buffer mod L for the next step.
# The ring cursor is a per-slot decode counter shared by all GDN layers;
# manage it once here (snapshot, hand to layers, advance mod L), not per-layer.
mamba_pool = getattr(self.req_to_token_pool, "mamba_pool", None)
write_pos_buf = (
getattr(mamba_pool, "replayssm_write_pos", None)
@@ -130,27 +142,16 @@ class MambaAttnBackendBase(AttentionBackend):
)
if write_pos_buf is not None:
slots = mamba_cache_indices.to(torch.long)
# Padded rows carry slot == -1; clamp so the per-row gather stays
# in-bounds (the kernel zeroes padded rows via state_idx < 0).
# Padded rows carry slot == -1; clamp the gather in-bounds (kernel
# zeroes padded rows via state_idx < 0).
safe_slots = slots.clamp(min=0)
replayssm_write_pos = write_pos_buf[safe_slots].clone()
L = mamba_pool.linear_replayssm_cache_len
# KDA (per-K gate) ships without radix coordination for now: no
# track-boundary force-flush, so the ring flushes only at the
# natural write_pos == L-1 wrap. GDN keeps the radix-aligned
# force-flush (slice 2b). Gate on the pool's recorded gate type.
# KDA has no radix coordination: flush only on the natural write_pos
# == L-1 wrap. GDN adds the radix-aligned force-flush below.
is_kda = getattr(mamba_pool, "replayssm_is_kda", False)
# GDN ReplaySSM (slice 2b): per-row force-flush at the radix
# track boundary. THE alignment: the radix mamba track snapshots
# temporal[slot] when seq_lens_cpu % mamba_track_interval == 0
# (extra_buffer: schedule_batch.prepare_for_decode builds
# `mamba_track_mask = (seq_lens_cpu % mamba_track_interval == 0)`
# off the SAME post-increment seq_lens_cpu used here). We source
# the flush from the identical seq_lens + condition so the kernel
# folds the ring into temporal[slot] on EXACTLY the steps the
# snapshot reads it. seq_lens_cpu is the committed length AFTER
# this decode token (incremented in prepare_for_decode before the
# forward), matching the track. int32, one entry per batch row.
# Force-flush on the radix track's seq_lens % mamba_track_interval
# == 0 boundary so the ring folds into temporal[slot] when read.
if not is_kda:
force_flush_bool = self._replayssm_track_flush_mask(
forward_batch.seq_lens_cpu, bs
@@ -158,16 +159,11 @@ class MambaAttnBackendBase(AttentionBackend):
replayssm_force_flush = force_flush_bool.to(
device=self.device, dtype=torch.int32
)
# Advance only the VALID (non-padded) slots. Scatter over the
# unique valid slots to avoid duplicate-index races (padded rows
# all clamp to slot 0, which a real row may also occupy). A
# forced flush empties the ring -> next write_pos is 0 (same as
# the natural wrap at write_pos == L-1).
# Advance only valid slots, scattered over unique slots (dup-index
# race; padded rows clamp to 0); a forced flush -> next write_pos 0.
valid_mask = slots >= 0
valid_slots = slots[valid_mask]
if valid_slots.numel() > 0:
# Per-row "did this step flush?": natural wrap OR forced.
# (KDA has no forced flush -> force_flush is None -> pure wrap.)
flushed = replayssm_write_pos == (L - 1)
if replayssm_force_flush is not None:
flushed = flushed | (replayssm_force_flush != 0)
@@ -176,9 +172,8 @@ class MambaAttnBackendBase(AttentionBackend):
torch.zeros_like(replayssm_write_pos),
(replayssm_write_pos + 1) % L,
)
# Dedup valid slots; for duplicates a scatter picks one
# arbitrary row, but all rows of a given slot share the same
# write_pos/flush, so the value is identical regardless.
# Dedup: rows sharing a slot share write_pos/flush, so the
# scattered value is identical regardless of which row wins.
uniq_slots, inv = torch.unique(valid_slots, return_inverse=True)
next_for_valid = next_pos[valid_mask]
new_vals = torch.empty(
@@ -190,9 +185,8 @@ class MambaAttnBackendBase(AttentionBackend):
write_pos_buf[uniq_slots] = new_vals
elif forward_batch.forward_mode.is_extend(include_draft_extend_v2=True):
if forward_batch.forward_mode.is_draft_extend_v2():
# HybridLinearAttnBackend.init_forward_metadata calls all sub-backends
# unconditionally, but DRAFT_EXTEND_V2 only runs full-attn layers in
# the draft model, so mamba metadata can be skipped.
# DRAFT_EXTEND_V2 runs only full-attn layers in the draft model;
# skip mamba metadata.
query_start_loc = None
elif forward_batch.forward_mode.is_target_verify():
query_start_loc = torch.arange(
@@ -208,7 +202,7 @@ class MambaAttnBackendBase(AttentionBackend):
retrieve_next_sibling = (
forward_batch.spec_info.retrieve_next_sibling
)
# retrieve_next_token is None during dummy run so skip tensor creation
# None during dummy run
if retrieve_next_token is not None:
retrieve_parent_token = torch.empty_like(retrieve_next_token)
else:
@@ -245,6 +239,9 @@ class MambaAttnBackendBase(AttentionBackend):
return ForwardMetadata(
query_start_loc=query_start_loc,
mamba_cache_indices=mamba_cache_indices,
# Physical track destinations (None when tracking off); cuda-graph
# supplies this via the static backend buffer in _replay_metadata.
mamba_track_indices=getattr(forward_batch, "mamba_track_indices", None),
retrieve_next_token=retrieve_next_token,
retrieve_next_sibling=retrieve_next_sibling,
retrieve_parent_token=retrieve_parent_token,
@@ -263,8 +260,6 @@ class MambaAttnBackendBase(AttentionBackend):
forward_batch: ForwardBatch,
in_capture: bool = False,
):
# seq_lens_cpu is unused by _replay_metadata for the non-target-verify
# case but kept in the contract for compatibility.
self.forward_metadata = self._replay_metadata(
forward_batch.batch_size,
forward_batch.req_pool_indices,
@@ -275,6 +270,7 @@ class MambaAttnBackendBase(AttentionBackend):
0 if in_capture else getattr(forward_batch, "num_padding", None)
),
in_capture=in_capture,
mamba_track_indices=getattr(forward_batch, "mamba_track_indices", None),
)
def init_forward_metadata(self, forward_batch: ForwardBatch):
@@ -284,28 +280,10 @@ class MambaAttnBackendBase(AttentionBackend):
def _init_track_conv_indices(
self, query_start_loc: torch.Tensor, forward_batch: ForwardBatch
):
"""
Compute indices for extracting conv states from the input sequence during extend.
In Mamba models, the conv layer maintains a sliding window of recent inputs.
After processing a prefill chunk, we need to save the last `conv_state_len` tokens
of the processed region for prefix caching.
The key insight is that FLA (Flash Linear Attention) and Mamba2 processes sequences in chunks
of the chunk size (FLA_CHUNK_SIZE=64 for FLA, mamba_chunk_size for Mamba2).
We only track the conv state up to the last complete chunk boundary (aligned_len).
start_indices is the starting token index of the conv state to track in this extend batch.
indices include all pos to track in this extend batch, conv_state_len for each req that
needs to be tracked (i.e. mamba_track_mask is True)
Returns:
indices: Tensor of shape [num_tracked_requests, conv_state_len] containing
flattened positions into the packed input tensor.
"""
"""Flattened input positions of conv states to track during extend (up to
the last complete chunk boundary, mamba_track_mask rows only)."""
conv_state_len = self.conv_states_shape[-1]
# Calculate the end position of the last aligned chunk
lens_to_track = (
forward_batch.mamba_track_seqlens - forward_batch.extend_prefix_lens
)
@@ -314,7 +292,6 @@ class MambaAttnBackendBase(AttentionBackend):
start_indices = query_start_loc[:-1] + aligned_len - conv_state_len
start_indices = start_indices[forward_batch.mamba_track_mask]
# Create indices: [batch_size, conv_state_len]
indices = start_indices.unsqueeze(-1) + torch.arange(
conv_state_len,
device=self.device,
@@ -326,46 +303,11 @@ class MambaAttnBackendBase(AttentionBackend):
def _init_track_ssm_indices(
self, mamba_cache_indices: torch.Tensor, forward_batch: ForwardBatch
):
"""
Compute source and destination indices for tracking SSM states for prefix caching.
After processing a prefill, we need to save the SSM recurrent state for prefix caching.
The kernel outputs intermediate hidden states `h` at each chunk boundary,
plus a `last_recurrent_state` at the end of the chunked prefill size.
The chunk size varies by model type:
- FLA models: FLA_CHUNK_SIZE (64)
- Mamba2 models: mamba_chunk_size (256)
The challenge is that sequences may or may not end on a chunk boundary:
- Aligned case (len % chunk_size == 0): The to-cache state is stored in
the last_recurrent_state.
- Unaligned case (len % chunk_size != 0): The last_recurrent_state includes the
unaligned position, but we only want state up to the last chunk boundary.
We must extract from the intermediate `h` tensor at the appropriate chunk index.
We compute the src and dst indices for all requests that need to be cached
(i.e. mamba_track_mask is True) based on the rule above.
For example (assuming chunk_size=64):
1. If chunked prefill length is < chunk_size, then only final state has value.
In this case we cache `final` state.
2. If chunked prefill length == chunk_size, then only final state has value.
In this case we cache pos chunk_size, from `final` state.
3. If chunked prefill length > chunk_size and < 2 * chunk_size, then both h and
final state have value. We cache pos chunk_size from `h` state.
4. If chunked prefill length == 2 * chunk_size, then both h and final state have
value. We cache pos 2 * chunk_size from `final` state. Note `h` doesn't include
the final position.
Returns:
track_ssm_h_src: Source indices into the packed `h` tensor (for unaligned seqs)
track_ssm_h_dst: Destination cache slot indices (for unaligned seqs)
track_ssm_final_src: Source indices into last_recurrent_state buffer (for aligned seqs)
track_ssm_final_dst: Destination cache slot indices (for aligned seqs)
"""
"""src/dst indices to track SSM states for prefix caching: aligned seqs
cache last_recurrent_state, unaligned cache intermediate `h` at the last
chunk boundary."""
mamba_cache_chunk_size = get_global_server_args().mamba_cache_chunk_size
# Move to CPU to avoid kernel launches for masking operations
# CPU to avoid kernel launches for the masking ops
mamba_track_mask = forward_batch.mamba_track_mask.cpu()
extend_seq_lens = forward_batch.extend_seq_lens.cpu()
mamba_track_indices = forward_batch.mamba_track_indices.cpu()
@@ -373,38 +315,33 @@ class MambaAttnBackendBase(AttentionBackend):
mamba_track_seqlens = forward_batch.mamba_track_seqlens.cpu()
prefix_lens = forward_batch.extend_prefix_lens.cpu()
# Calculate the number of hidden states per request
if isinstance(self, Mamba2AttnBackend):
num_h_states = extend_seq_lens // mamba_cache_chunk_size
else:
num_h_states = (extend_seq_lens - 1) // mamba_cache_chunk_size + 1
# Calculate the starting offset for each sequence in the packed batch
track_ssm_src_offset = torch.zeros_like(num_h_states)
track_ssm_src_offset[1:] = torch.cumsum(num_h_states[:-1], dim=0)
# Filter variables by track mask
lens_to_track = mamba_track_seqlens - prefix_lens
lens_masked = lens_to_track[mamba_track_mask]
offset_masked = track_ssm_src_offset[mamba_track_mask]
dst_masked = mamba_track_indices[mamba_track_mask]
# Determine if the sequence ends at a chunk boundary
is_aligned = (lens_masked % mamba_cache_chunk_size) == 0
# Case 1: Aligned. Use last_recurrent_state from ssm_states.
# Aligned: last_recurrent_state from ssm_states.
track_ssm_final_src = mamba_cache_indices[mamba_track_mask][is_aligned]
track_ssm_final_dst = dst_masked[is_aligned]
# Case 2: Unaligned. Use intermediate state from h.
# TODO: if support mamba_cache_chunk_size % page size != 0, then need to modify this
# Unaligned: intermediate state from h.
# TODO: handle mamba_cache_chunk_size % page size != 0
not_aligned = ~is_aligned
track_ssm_h_src = offset_masked[not_aligned] + (
lens_masked[not_aligned] // mamba_cache_chunk_size
)
track_ssm_h_dst = dst_masked[not_aligned]
# Move back to GPU
return (
track_ssm_h_src.to(self.device, non_blocking=True),
track_ssm_h_dst.to(self.device, non_blocking=True),
@@ -427,12 +364,8 @@ class MambaAttnBackendBase(AttentionBackend):
)
def _replayssm_enabled(self) -> bool:
"""True iff --enable-linear-replayssm allocated the persistent ring cursor.
The per-slot ``replayssm_write_pos`` buffer on MambaPool is None unless
the flag is set, so it doubles as the on/off gate (same signal that
``_forward_metadata`` / ``GDNAttnBackend.forward_decode`` already use).
"""
"""True iff --enable-linear-replayssm allocated the ring cursor
(MambaPool.replayssm_write_pos doubles as the on/off gate)."""
mamba_pool = getattr(self.req_to_token_pool, "mamba_pool", None)
if mamba_pool is None:
return False
@@ -441,21 +374,12 @@ class MambaAttnBackendBase(AttentionBackend):
def _replayssm_track_flush_mask(
self, seq_lens_cpu: torch.Tensor, bs: int
) -> torch.Tensor:
"""Per-row bool flush mask == the radix mamba-track snapshot condition.
THE alignment (slice 2b): the radix mamba track snapshots temporal[slot]
exactly when ``seq_lens_cpu % mamba_track_interval == 0`` (the same mask
``schedule_batch.prepare_for_decode`` builds for extra_buffer at
``mamba_track_mask = (seq_lens_cpu % mamba_track_interval == 0)``). Both
read the SAME post-increment ``seq_lens_cpu`` (committed length AFTER
this decode token), so the kernel force-flush fires on EXACTLY the steps
the snapshot reads the checkpoint -- no off-by-one. Returns a CPU bool
tensor of length ``bs`` (caller moves it to device as int32).
"""
"""Per-row (length bs) bool flush mask = the radix track's seq_lens_cpu %
mamba_track_interval == 0, so force-flush and snapshot fire on the same
steps (no off-by-one)."""
interval = get_global_server_args().mamba_track_interval
if seq_lens_cpu is None:
# Decode without a CPU seq-len mirror should not happen for the
# supported (no_buffer, radix-on) config, but stay safe: never flush.
# Should not happen for the supported config; stay safe and never flush.
return torch.zeros((bs,), dtype=torch.bool)
mask = (seq_lens_cpu[:bs].to(torch.int64) % interval) == 0
if mask.shape[0] < bs:
@@ -468,16 +392,15 @@ class MambaAttnBackendBase(AttentionBackend):
max_num_tokens % max_bs == 0
), f"max_num_tokens={max_num_tokens} must be divisible by max_bs={max_bs}"
draft_token_num = max_num_tokens // max_bs
# GDN ReplaySSM (slice 1b): per-batch-size STATIC per-row write-cursor
# buffers the kernel reads. Captured into the graph by pointer, so they
# must be the SAME tensor objects refreshed in-place each replay. Sized
# and indexed like state_indices_list ((i+1,), indexed [bs - 1]). Left
# None when the flag is off so the dispatch falls through unchanged.
# Per-bs static write-cursor / force-flush buffers, captured by pointer +
# refreshed in-place each replay; sized like state_indices_list. None when off.
self.replayssm_write_pos_list = [] if self._replayssm_enabled() else None
# GDN ReplaySSM (slice 2b): static per-bs force-flush buffers, captured
# by pointer and refreshed in-place per replay just like the write-pos
# buffers. None when the flag is off.
self.replayssm_force_flush_list = [] if self._replayssm_enabled() else None
# int64 to match DecodeInputBuffers.mamba_track_indices + the track-save
# kernel's int64 index load. Refreshed in-place by _replay_metadata.
self.mamba_track_indices_buf = torch.zeros(
(max_bs,), dtype=torch.int64, device=self.device
)
for i in range(max_bs):
self.state_indices_list.append(
torch.full(
@@ -556,32 +479,26 @@ class MambaAttnBackendBase(AttentionBackend):
else:
raise ValueError(f"Invalid forward mode: {forward_mode=}")
mamba_indices = self.req_to_token_pool.get_mamba_indices(req_pool_indices)
# Captured Mamba kernels read state_indices_list as PHYSICAL ids; translate
# before copying (no-op for non-unified pool).
mamba_indices = self._translate_mamba_indices(mamba_indices)
self.state_indices_list[bs - 1][: len(mamba_indices)].copy_(mamba_indices)
# GDN ReplaySSM (slice 1b): point at the STATIC per-bs write-cursor
# buffer (no advance, no snapshot — capture records the pointer; its
# zeros are overwritten in-place by _replay_metadata before each
# replay). None when the flag is off. Same per-bs tensor object that
# _replay_metadata refreshes, so the captured pointer stays valid.
# Capture records the pointer to the static per-bs buffers; their zeros are
# overwritten in-place by _replay_metadata before each replay. None when off.
replayssm_write_pos = (
self.replayssm_write_pos_list[bs - 1]
if self.replayssm_write_pos_list is not None
else None
)
# GDN ReplaySSM (slice 2b): point at the STATIC per-bs force-flush
# buffer (same capture-by-pointer contract as write_pos; refreshed
# in-place by _replay_metadata before each replay). None when off.
replayssm_force_flush = (
self.replayssm_force_flush_list[bs - 1]
if self.replayssm_force_flush_list is not None
else None
)
# If topk > 1, we need to use retrieve_next_token and retrieve_next_sibling to handle the eagle tree custom attention mask
if forward_mode.is_target_verify() and self.topk > 1:
# They are None during cuda graph capture so skip the copy_...
# self.retrieve_next_token_list[bs - 1].copy_(spec_info.retrieve_next_token)
# self.retrieve_next_sibling_list[bs - 1].copy_(spec_info.retrieve_next_sibling)
# retrieve_* are None during capture, so skip the copy.
return ForwardMetadata(
query_start_loc=self.query_start_loc_list[bs - 1],
mamba_cache_indices=self.state_indices_list[bs - 1],
@@ -608,6 +525,7 @@ class MambaAttnBackendBase(AttentionBackend):
seq_lens_cpu: Optional[torch.Tensor],
num_padding: Optional[int] = None,
in_capture: bool = False,
mamba_track_indices: Optional[torch.Tensor] = None,
):
if num_padding is None:
if seq_lens_cpu is None:
@@ -619,16 +537,22 @@ class MambaAttnBackendBase(AttentionBackend):
# Make sure forward metadata is correctly handled for padding reqs
req_pool_indices[bs - num_padding :] = 0
mamba_indices = self.req_to_token_pool.get_mamba_indices(req_pool_indices)
# Translate using the LIVE v2p table BEFORE the padding sentinel below;
# captured Mamba kernels read state_indices_list as PHYSICAL ids.
mamba_indices = self._translate_mamba_indices(mamba_indices)
mamba_indices[bs - num_padding :] = -1
self.state_indices_list[bs - 1][: len(mamba_indices)].copy_(mamba_indices)
# GDN ReplaySSM (slice 1b): refresh the STATIC per-row write cursor the
# kernel reads, mirroring the eager snapshot-then-advance in
# _forward_metadata but writing in-place into the captured per-bs buffer
# so the graph's recorded pointer stays valid across replays. Done once
# per forward here (out_graph host op), not per layer. Skipped during
# capture (in_capture): capture runs on dummy slots, so advancing the
# persistent counter then would corrupt real per-slot ring positions;
# the captured buffer's contents are irrelevant at capture time anyway.
# Refresh the static track-dest buffer in-place (translated); the captured
# track-save reads it, leaving the handed-in InputBuffer slot read-only.
track_buf = None
if mamba_track_indices is not None:
track_buf = self.mamba_track_indices_buf
track_buf[: len(mamba_track_indices)].copy_(
self._translate_mamba_indices(mamba_track_indices)
)
# Refresh the static write cursor in-place (mirrors the eager
# snapshot-then-advance). Skip the advance during capture: dummy slots
# would corrupt real ring positions.
replayssm_write_pos = None
replayssm_force_flush = None
if self.replayssm_write_pos_list is not None:
@@ -636,29 +560,20 @@ class MambaAttnBackendBase(AttentionBackend):
write_pos_buf = mamba_pool.replayssm_write_pos
static_wp = self.replayssm_write_pos_list[bs - 1]
static_ff = self.replayssm_force_flush_list[bs - 1]
# Hand the full captured per-bs buffers to the kernel, mirroring how
# mamba_cache_indices = self.state_indices_list[bs - 1] is the full
# (bs,) tensor; the kernel indexes them per decode row.
# Hand the full captured per-bs buffers to the kernel; it indexes per row.
replayssm_write_pos = static_wp
replayssm_force_flush = static_ff
if write_pos_buf is not None:
# mamba_indices: this replay's per-row physical slots (padded
# rows == -1, same tensor fed to state_indices_list above).
# this replay's per-row physical slots (padded rows == -1)
slots = mamba_indices.to(torch.long)
safe_slots = slots.clamp(min=0)
# Snapshot THIS step's per-slot cursor into the captured buffer
# the kernel reads (in-place copy_, never reassign the object).
# Snapshot this step's cursor into the captured buffer in-place
# (copy_, never reassign the object).
static_wp[: len(mamba_indices)].copy_(write_pos_buf[safe_slots])
# GDN ReplaySSM (slice 2b): refresh the captured force-flush
# buffer in-place from THIS step's seq_lens. THE alignment: same
# `seq_lens_cpu % mamba_track_interval == 0` the radix track uses
# (see _replayssm_track_flush_mask / schedule_batch). During
# capture (seq_lens_cpu is None) leave it zeroed: capture content
# is irrelevant and decode replays overwrite it below.
# Refresh the force-flush buffer in-place from this step's seq_lens
# (same condition as the radix track). Zeroed during capture.
force_flush_dev = None
# KDA: no radix coordination -> leave static_ff zeroed and
# force_flush_dev None so the advance below is a pure wrap,
# matching the kernel (a zeroed force_flush flushes nothing).
# KDA: no radix coordination -> leave zeroed so the advance is a pure wrap.
is_kda = getattr(mamba_pool, "replayssm_is_kda", False)
if (
not is_kda
@@ -672,10 +587,8 @@ class MambaAttnBackendBase(AttentionBackend):
static_ff.zero_()
if not in_capture:
L = mamba_pool.linear_replayssm_cache_len
# Advance only VALID (non-padded) slots. A forced flush
# empties the ring -> next write_pos is 0 (same as the
# natural wrap at write_pos == L-1). Use this step's snapshot
# cursor (write_pos_buf[safe_slots]) + the flush flag.
# Advance only valid (non-padded) slots; a forced flush empties
# the ring -> next write_pos 0, like the natural L-1 wrap.
valid_mask = slots >= 0
valid_slots = slots[valid_mask]
if valid_slots.numel() > 0:
@@ -688,8 +601,7 @@ class MambaAttnBackendBase(AttentionBackend):
torch.zeros_like(cur_pos),
(cur_pos + 1) % L,
)
# Dedup; rows sharing a slot share write_pos+flush, so
# the scattered value is identical for either row.
# Dedup; rows sharing a slot share write_pos+flush.
uniq_slots, inv = torch.unique(valid_slots, return_inverse=True)
next_for_valid = next_pos[valid_mask]
new_vals = torch.empty(
@@ -726,7 +638,6 @@ class MambaAttnBackendBase(AttentionBackend):
else:
raise ValueError(f"Invalid forward mode: {forward_mode=}")
# If topk > 1, we need to use retrieve_next_token and retrieve_next_sibling to handle the eagle tree custom attention mask
if forward_mode.is_target_verify() and self.topk > 1:
if (
spec_info is not None
@@ -742,6 +653,7 @@ class MambaAttnBackendBase(AttentionBackend):
return ForwardMetadata(
query_start_loc=self.query_start_loc_list[bs - 1],
mamba_cache_indices=self.state_indices_list[bs - 1],
mamba_track_indices=track_buf,
retrieve_next_token=self.retrieve_next_token_list[bs - 1],
retrieve_next_sibling=self.retrieve_next_sibling_list[bs - 1],
retrieve_parent_token=self.retrieve_parent_token_list[bs - 1],
@@ -752,6 +664,7 @@ class MambaAttnBackendBase(AttentionBackend):
return ForwardMetadata(
query_start_loc=self.query_start_loc_list[bs - 1],
mamba_cache_indices=self.state_indices_list[bs - 1],
mamba_track_indices=track_buf,
replayssm_write_pos=replayssm_write_pos,
replayssm_force_flush=replayssm_force_flush,
)
@@ -769,27 +682,18 @@ class MambaAttnBackendBase(AttentionBackend):
ssm_states: torch.Tensor,
cache_indices: torch.Tensor,
):
"""
Track and copy Mamba conv/SSM states during decode for prefix caching.
During decode, each token update modifies conv_states and ssm_states in-place
at positions indexed by cache_indices (the working slots). For prefix caching,
we need to copy these updated states to persistent cache slots (mamba_track_indices)
so they can be prefix cached.
This delegates to `track_mamba_states_if_needed`, which performs:
conv_states[mamba_track_indices[i]] = conv_states[cache_indices[i]]
ssm_states[mamba_track_indices[i]] = ssm_states[cache_indices[i]]
for all requests where mamba_track_mask[i] is True.
"""
"""Copy decode conv/SSM states to track slots for prefix caching. Track
dests come from the metadata (under cuda-graph: the static buffer), so the
InputBuffer registry slot is never mutated."""
if forward_batch.mamba_track_mask is not None:
track_mamba_states_if_needed(
conv_states,
ssm_states,
cache_indices,
forward_batch.mamba_track_mask,
forward_batch.mamba_track_indices,
self.forward_metadata.mamba_track_indices,
forward_batch.batch_size,
check_freed_slots=self.enable_unified_memory,
)
def _track_mamba_state_extend(
@@ -799,18 +703,8 @@ class MambaAttnBackendBase(AttentionBackend):
ssm_states: torch.Tensor,
forward_metadata: ForwardMetadata,
):
"""
Track and copy SSM states during extend for prefix caching.
After the chunked prefill kernel runs, we need to save the SSM recurrent
state at the last chunk boundary so it can be reused for prefix caching.
The source of the state depends on whether the sequence length is aligned
to the chunk size. See `_init_track_ssm_indices` for more details on how
the source and destination indices are computed.
Note: Conv state tracking for extend is handled separately via gather operations
using indices computed by `_init_track_conv_indices`.
"""
"""Copy extend SSM state at the last chunk boundary to track slots (source
depends on chunk alignment; see `_init_track_ssm_indices`)."""
if forward_metadata.has_mamba_track_mask:
h = h.squeeze(0)
@@ -861,6 +755,7 @@ class Mamba2AttnBackend(MambaAttnBackendBase):
0 if in_capture else getattr(forward_batch, "num_padding", None)
),
in_capture=in_capture,
mamba_track_indices=getattr(forward_batch, "mamba_track_indices", None),
)
spec_info = forward_batch.spec_info
draft_token_num = spec_info.draft_token_num if spec_info is not None else 1
@@ -891,6 +786,12 @@ class Mamba2AttnBackend(MambaAttnBackendBase):
use_triton_causal_conv: bool = False,
):
assert isinstance(self.forward_metadata, Mamba2Metadata)
# Page-major stores state strided; only the stride-aware Triton causal-conv
# reads it (CUDA causal_conv1d garbles it). A model may also force Triton.
use_triton_causal_conv = (
use_triton_causal_conv
or get_global_server_args().enable_page_major_kv_layout
)
layer_cache = self.req_to_token_pool.mamba2_layer_cache(layer_id)
mixer_out, intermediate_states = mixer.forward(
hidden_states=hidden_states,
@@ -922,8 +823,9 @@ class Mamba2AttnBackend(MambaAttnBackendBase):
layer_cache.temporal,
self.forward_metadata.mamba_cache_indices[-num_decodes:],
forward_batch.mamba_track_mask[-num_decodes:],
forward_batch.mamba_track_indices[-num_decodes:],
self.forward_metadata.mamba_track_indices[-num_decodes:],
num_decodes,
check_freed_slots=self.enable_unified_memory,
)
return mixer_out
@@ -952,7 +854,6 @@ class HybridLinearAttnBackend(AttentionBackend):
self.full_attn_backend = full_attn_backend
self.linear_attn_backend = linear_attn_backend
self.attn_backend_list = [full_attn_backend, linear_attn_backend]
# Dispatcher aliases the full-attn backend's pool refs.
self.token_to_kv_pool = full_attn_backend.token_to_kv_pool
self.req_to_token_pool = full_attn_backend.req_to_token_pool
self.max_context_len = getattr(full_attn_backend, "max_context_len", None)
@@ -981,8 +882,8 @@ class HybridLinearAttnBackend(AttentionBackend):
def init_forward_metadata(self, forward_batch: ForwardBatch):
if forward_batch.forward_mode.is_draft_extend_v2():
# DRAFT_EXTEND_V2 only runs full-attn layers in the draft model,
# so skip linear/mamba backend metadata which requires query_start_loc.
# DRAFT_EXTEND_V2 runs only full-attn layers in the draft model; skip
# linear/mamba metadata (it requires query_start_loc).
self.full_attn_backend.init_forward_metadata(forward_batch)
return
for attn_backend in self.attn_backend_list:
@@ -991,9 +892,8 @@ class HybridLinearAttnBackend(AttentionBackend):
def init_mha_chunk_metadata(
self, forward_batch: ForwardBatch, disable_flashinfer_ragged: bool = False
):
# Hybrid MLA models (Ring/Ling, Kimi-Linear) resolve this via
# get_attn_backend(), which returns this wrapper; delegate to the
# full-attn backend so its chunked/one-shot prefill metadata is planned.
# Hybrid MLA models resolve get_attn_backend() to this wrapper; delegate
# so the full-attn backend plans its chunked-prefill metadata.
init = getattr(self.full_attn_backend, "init_mha_chunk_metadata", None)
if init is not None:
init(forward_batch, disable_flashinfer_ragged)
@@ -1050,7 +950,6 @@ class HybridLinearAttnBackend(AttentionBackend):
return self.full_attn_backend.forward_decode(
q, k, v, layer, forward_batch, save_kv_cache, **kwargs
)
# Linear attention backend
return self.linear_attn_backend.forward_decode(
q=q,
k=k,
@@ -1081,7 +980,6 @@ class HybridLinearAttnBackend(AttentionBackend):
return self.full_attn_backend.forward_extend(
q, k, v, layer, forward_batch, save_kv_cache, **kwargs
)
# Linear attention backend
return self.linear_attn_backend.forward_extend(
q=q,
k=k,
@@ -1150,15 +1048,7 @@ class HybridLinearAttnBackend(AttentionBackend):
mamba_steps_to_track: Optional[torch.Tensor],
model,
):
"""
Update mamba states after MTP verify using fully fused Triton kernel.
This replaces the original advanced indexing operations with a single fused
gather-scatter kernel that also handles masking internally, avoiding:
- index_elementwise_kernel from tensor[bool_mask]
- index_select kernel launches
- nonzero kernel launches
"""
"""Update mamba states after MTP verify via a fused gather-scatter kernel."""
request_number = last_correct_step_indices.shape[0]
state_indices_tensor = (
@@ -1176,16 +1066,14 @@ class HybridLinearAttnBackend(AttentionBackend):
intermediate_state_cache = mamba_caches.intermediate_ssm
intermediate_conv_window_cache = mamba_caches.intermediate_conv_window[0]
# Use fully fused kernel that handles masking internally
# This avoids separate nonzero() and index_select() calls
fused_mamba_state_scatter_with_mask(
ssm_states,
intermediate_state_cache,
state_indices_tensor,
last_correct_step_indices,
)
# conv intermediate uses the deduplicated sliding-window (overlapping)
# layout, so it needs the strided-read scatter variant.
# conv intermediate uses the deduplicated sliding-window layout, so it
# needs the strided-read scatter variant.
fused_conv_window_scatter_with_mask(
conv_states,
intermediate_conv_window_cache,
@@ -1193,10 +1081,9 @@ class HybridLinearAttnBackend(AttentionBackend):
last_correct_step_indices,
)
# Track indices used for tracking mamba states for prefix cache
# Track indices for prefix cache
if mamba_track_indices is not None:
assert mamba_steps_to_track is not None
# Use fully fused kernel for track scatter operations
fused_mamba_state_scatter_with_mask(
ssm_states,
intermediate_state_cache,
python/sglang/srt/layers/attention/mamba/mamba2_metadata.py
+7
View File
@@ -30,6 +30,11 @@ class ForwardMetadata:
query_start_loc: torch.Tensor
mamba_cache_indices: torch.Tensor
mamba_cache_indices_gdn: Optional[torch.Tensor] = None
# Mamba track DESTINATION slots (PHYSICAL, length == batch). Like
# mamba_cache_indices: a backend-owned static buffer under cuda-graph (translated
# in-place each replay), eager sets the translated decode tensor. The decode
# track-save reads THIS, never forward_batch.mamba_track_indices.
mamba_track_indices: Optional[torch.Tensor] = None
# GDN ReplaySSM (slice 1a): per-decode-row snapshot of the ring write
# cursor for THIS decode step (gathered from the persistent per-slot
# buffer, then advanced once for the next step). int32, length == batch.
@@ -179,6 +184,7 @@ class Mamba2Metadata(ForwardMetadata):
return Mamba2Metadata(
query_start_loc=forward_metadata.query_start_loc,
mamba_cache_indices=forward_metadata.mamba_cache_indices,
mamba_track_indices=forward_metadata.mamba_track_indices,
retrieve_next_token=forward_metadata.retrieve_next_token,
retrieve_next_sibling=forward_metadata.retrieve_next_sibling,
retrieve_parent_token=forward_metadata.retrieve_parent_token,
@@ -274,6 +280,7 @@ class Mamba2Metadata(ForwardMetadata):
return Mamba2Metadata(
query_start_loc=query_start_loc,
mamba_cache_indices=forward_metadata.mamba_cache_indices,
mamba_track_indices=forward_metadata.mamba_track_indices,
retrieve_next_token=forward_metadata.retrieve_next_token,
retrieve_next_sibling=forward_metadata.retrieve_next_sibling,
retrieve_parent_token=forward_metadata.retrieve_parent_token,
python/sglang/srt/layers/attention/mamba/mamba_state_scatter_triton.py
+9
View File
@@ -23,6 +23,7 @@ def track_mamba_state_if_needed_kernel(
conv_state_numel_per_row: tl.constexpr, # total elements per row
ssm_state_numel_per_row: tl.constexpr, # total elements per row
BLOCK_SIZE: tl.constexpr,
check_freed_slots: tl.constexpr, # only the shared/unified KV pool emits -1
):
"""
Track conv_states and ssm_states rows based on track mask.
@@ -51,6 +52,12 @@ def track_mamba_state_if_needed_kernel(
src_idx = tl.load(cache_indices_ptr + batch_idx).to(tl.int64)
dst_idx = tl.load(mamba_track_indices_ptr + batch_idx).to(tl.int64)
# Skip freed slots (-1): `state_ptr + (-1)*stride` would fault. Only the unified
# pool emits -1 tombstones (from the v2p translate); compiled out for static.
if check_freed_slots:
if src_idx < 0 or dst_idx < 0:
return
# Copy conv_states
# Each thread handles BLOCK_SIZE elements
for offset in range(0, conv_state_numel_per_row, BLOCK_SIZE):
@@ -82,6 +89,7 @@ def track_mamba_states_if_needed(
mamba_track_mask: torch.Tensor,
mamba_track_indices: torch.Tensor,
batch_size: int,
check_freed_slots: bool = False,
):
"""
Track mamba states using Triton kernel for better performance.
@@ -113,6 +121,7 @@ def track_mamba_states_if_needed(
conv_state_numel_per_row,
ssm_state_numel_per_row,
BLOCK_SIZE,
check_freed_slots,
)
python/sglang/srt/layers/attention/triton_backend.py
+149 -102
View File
@@ -98,6 +98,9 @@ class ForwardMetadata:
swa_attn_logits: Optional[torch.Tensor] = None
# full->SWA translated out_cache_loc (SWA KV-store write target)
swa_out_cache_loc: Optional[torch.Tensor] = None
# PHYSICAL full-attn write target for the unified pool (eager: translated tensor;
# cuda-graph: capture-stable buffer view). None for non-unified pools.
out_cache_loc_full_physical: Optional[torch.Tensor] = None
class TritonAttnBackend(AttentionBackend):
@@ -132,39 +135,31 @@ class TritonAttnBackend(AttentionBackend):
extend_attention_fwd_unified
)
self.build_unified_kv_indices = torch.compiler.disable(build_unified_kv_indices)
# Split-KV EAGLE-verify kernel (ROCm/Triton). Registered here; enabled
# below once topk is known (the path is only valid at topk == 1).
# Split-KV EAGLE-verify kernel; enabled below once topk is known (valid only at topk == 1).
self.verify_splitkv_fwd = torch.compiler.disable(verify_splitkv_fwd)
# Parse args
self.skip_prefill = skip_prefill
max_bs = model_runner.req_to_token_pool.size
self.sliding_window_size = model_runner.sliding_window_size
# Pool refs — captured at construction so they survive deletion of the
# corresponding ForwardBatch fields.
self.req_to_token_pool = model_runner.req_to_token_pool
self.token_to_kv_pool = model_runner.token_to_kv_pool
self.req_to_token = model_runner.req_to_token_pool.req_to_token
self.token_to_kv_pool_allocator = model_runner.token_to_kv_pool_allocator
self.use_sliding_window_kv_pool = isinstance(self.token_to_kv_pool, SWAKVPool)
# Pass-through to the Triton attention wrappers so they can extract the
# KV view strides and specialize on the PAGE_SIZE constexpr. At
# page_size=1 the kernel path matches the slot-based envelope addresses.
# `model_runner.page_size` defaults to 1 when `server_args.page_size` is
# None, avoiding the Optional case here.
# Lets the Triton wrappers specialize on PAGE_SIZE; page_size=1 is
# byte-identical to the slot-based envelope.
self.page_size = getattr(model_runner, "page_size", 1) or 1
# Unified pool v2p hook (None = no-op): req_to_token holds VIRTUAL ids but
# kernels need PHYSICAL. Applied eagerly so the captured graph has no translate.
self._translate_kv_loc = getattr(
self.token_to_kv_pool_allocator, "translate_kv_loc", None
)
self.num_draft_tokens = model_runner.server_args.speculative_num_draft_tokens
self.speculative_num_steps = model_runner.server_args.speculative_num_steps
self.topk = model_runner.server_args.speculative_eagle_topk or 0
# Split-KV verify matches extend_attention_fwd only when the EAGLE tree
# reduces to a pure-causal chain, i.e. topk == 1 (the same condition the
# aiter backend's unified-verify uses). For topk > 1 the tree custom_mask
# is not causal, so leave the path off and fall back to the baseline.
# gfx95-only (MI350X/CDNA4): the kernel uses ROCm/CDNA Triton launch hints
# (waves_per_eu, matrix_instr_nonkdim) and its block config is tuned and
# validated only on gfx950. NVIDIA's Triton rejects those kwargs, and the
# path is unvalidated on NV and on other AMD archs, so restrict it to gfx95
# and fall back to extend_attention_fwd everywhere else.
# Split-KV verify is bit-equivalent only for a pure-causal chain (topk==1)
# and is gfx95-only; else fall back to extend_attention_fwd.
self.use_verify_splitkv = (
is_gfx95_supported()
and envs.SGLANG_ENABLE_SPLITKV_VERIFY.get()
@@ -179,12 +174,9 @@ class TritonAttnBackend(AttentionBackend):
self.num_kv_head = model_runner.model_config.get_num_kv_heads(
get_parallel().attn_tp_size
)
# The decode triton kernel derives attn_lse offsets from attn_logits
# strides via integer division by v_head_dim (the "// Lv" trick in
# _fwd_kernel_stage1/stage2), so attn_logits.shape[-1] must exactly
# match the layer's v_head_dim. For hybrid SWA models where SWA and
# full-attention layers use different v_head_dim (e.g. Gemma 4:
# swa=256, full=512), we allocate a second buffer for SWA layers.
# The decode kernel's "// Lv" stride trick requires attn_logits.shape[-1]
# to exactly match the layer's v_head_dim, so hybrid SWA models with
# differing SWA/full v_head_dim need a second buffer for SWA layers.
full_v_head_dim = model_runner.model_config.v_head_dim
swa_v_head_dim = model_runner.model_config.swa_v_head_dim
if self.sliding_window_size is not None and swa_v_head_dim != full_v_head_dim:
@@ -217,13 +209,9 @@ class TritonAttnBackend(AttentionBackend):
self.max_context_len,
)
if _is_gfx942:
# gfx942 (MI300X / MI325X) has 304 CUs, so #20479's next_power_of_2(sm_count)
# rounds up to 512 — twice MI355X's natural cap of 256 — and the persistent
# cuda_graph_attn_logits fp32 buffer hits ~4 GiB on Kimi-K2.6 (v_head_dim=512),
# faulting in ROCm CUDA graph replay
# (https://github.com/sgl-project/sglang/actions/runs/25513282022/job/74877480809).
# Pin the cap at 256 so gfx942 matches the gfx950 (MI355X) behavior that we
# already validated end-to-end.
# gfx942's 304 CUs round up to 512 splits, doubling the persistent
# fp32 attn_logits buffer to ~4 GiB on Kimi-K2.6 and faulting in
# ROCm graph replay; pin to 256 to match validated gfx950 behavior.
self.max_kv_splits = min(self.max_kv_splits, 256)
if _is_cuda:
self.use_pdl = is_arch_support_pdl()
@@ -235,18 +223,15 @@ class TritonAttnBackend(AttentionBackend):
and model_runner.server_args.chunked_prefill_size == -1
)
# Decide whether enable deterministic inference with batch-invariant operations
self.enable_deterministic = (
model_runner.server_args.enable_deterministic_inference
)
# Configure deterministic inference settings
if self.enable_deterministic:
# Use fixed split tile size for batch invariance
# Fixed split tile size for batch invariance
self.split_tile_size = get_int_env_var(
"SGLANG_TRITON_DECODE_SPLIT_TILE_SIZE", 256
)
# Set static_kv_splits to False to use deterministic logic instead
self.static_kv_splits = False
else:
self.split_tile_size = (
@@ -258,14 +243,12 @@ class TritonAttnBackend(AttentionBackend):
self.max_context_len + self.split_tile_size - 1
) // self.split_tile_size
# Check arguments
assert not (
model_runner.sliding_window_size is not None
and model_runner.model_config.is_encoder_decoder
), "Sliding window and cross attention are not supported together"
# Initialize buffers
# TODO(Jianan Ji): Make sure it behaves as expected when kv_indptr_buf is provided and sliding window is enabled
# TODO(Jianan Ji): verify behavior when kv_indptr_buf is provided and sliding window is enabled
if kv_indptr_buf is None:
self.kv_indptr = torch.zeros(
(max_bs + 1,), dtype=torch.int32, device=model_runner.device
@@ -273,8 +256,7 @@ class TritonAttnBackend(AttentionBackend):
else:
self.kv_indptr = kv_indptr_buf
# If sliding window is enabled, we might need two sets of buffers
# because of interleaved attention types (e.g. for Gemma3)
# Sliding window may need a second buffer for interleaved attention types
self.window_kv_indptr = None
if self.sliding_window_size is not None and self.sliding_window_size > 0:
if kv_indptr_buf is None:
@@ -282,7 +264,6 @@ class TritonAttnBackend(AttentionBackend):
(max_bs + 1,), dtype=torch.int32, device=model_runner.device
)
else:
# When provided a buffer, create a clone for the second buffer
self.window_kv_indptr = torch.zeros_like(kv_indptr_buf)
if not self.skip_prefill:
@@ -294,7 +275,6 @@ class TritonAttnBackend(AttentionBackend):
(max_bs + 1,), dtype=torch.int64, device=model_runner.device
)
# Initialize forward metadata
self.forward_metadata: ForwardMetadata = None
self.cuda_graph_custom_mask = None
@@ -314,16 +294,13 @@ class TritonAttnBackend(AttentionBackend):
num_group * num_seq == num_token
), f"num_seq({num_seq}), num_token({num_token}), something goes wrong!"
# Legacy dynamic splitting logic (non-deterministic)
if (
self.static_kv_splits or self.device_core_count <= 0
) and not self.enable_deterministic:
num_kv_splits.fill_(self.max_kv_splits)
return
# deterministic
if self.split_tile_size is not None and self.enable_deterministic:
# expand seq_lens to match num_token
if num_group > 1:
expanded_seq_lens = seq_lens.repeat_interleave(num_group)
else:
@@ -362,9 +339,8 @@ class TritonAttnBackend(AttentionBackend):
kv_indices: Optional[torch.Tensor] = None,
kv_start_idx: Optional[torch.Tensor] = None,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
# Build per-DCP-rank sharded KV indptr/indices. eager passes
# kv_indices=None (allocate a fresh tensor); the cuda-graph path passes
# a fixed address-stable buffer to fill in place.
# Build per-DCP-rank sharded KV indptr/indices. eager passes kv_indices=None
# (fresh tensor); cuda-graph passes an address-stable buffer to fill in place.
dcp_lens = self._dcp_lens(lens, kv_start_idx)
kv_indptr[1 : len(req_pool_indices) + 1] = torch.cumsum(dcp_lens, dim=0)
kv_indptr = kv_indptr[: len(req_pool_indices) + 1]
@@ -422,10 +398,8 @@ class TritonAttnBackend(AttentionBackend):
seq_lens = seq_lens[:bs]
req_pool_indices = req_pool_indices[:bs]
if self.dcp_size > 1:
# DCP: kv_indptr cumsum and kv_indices are per-rank sharded. Write
# them into the same cuda-graph buffers that
# _build_cuda_graph_forward_metadata reads back
# (self.kv_indptr / self.cuda_graph_kv_indices).
# DCP: per-rank sharded; write into the same cuda-graph buffers
# _build_cuda_graph_forward_metadata reads back.
_, _, dcp_seq_lens = self._dcp_kv_indices(
req_pool_indices,
seq_lens,
@@ -439,10 +413,15 @@ class TritonAttnBackend(AttentionBackend):
kv_indptr = self._fill_kv_indptr_and_indices(
bs, seq_lens, req_pool_indices, self.cuda_graph_kv_indices
)
# Unified pool: VIRTUAL ids written here are translated to PHYSICAL in
# init_forward_metadata_out_graph (replay-prep) so the captured graph
# carries zero translate nodes.
num_kv_splits_lens = seq_lens
window_kv_indptr = self.window_kv_indptr
window_kv_lens = None
if self.sliding_window_size is not None and self.sliding_window_size > 0:
# Unified pool: leave the window VIRTUAL too (translated alongside the
# full kv_indices later); baseline SWA keeps the eager window translate.
window_kv_indptr, _, window_kv_lens, _ = update_sliding_window_buffer(
self.window_kv_indptr,
self.req_to_token,
@@ -452,6 +431,7 @@ class TritonAttnBackend(AttentionBackend):
bs,
token_to_kv_pool=self.token_to_kv_pool,
window_kv_indices=self.cuda_graph_window_kv_indices,
skip_full_to_swa_translation=(self._translate_kv_loc is not None),
)
return kv_indptr, window_kv_indptr, window_kv_lens, num_kv_splits_lens
@@ -462,12 +442,7 @@ class TritonAttnBackend(AttentionBackend):
req_pool_indices: torch.Tensor,
spec_info,
):
"""Fill all cuda-graph buffers for target_verify mode.
Returns the ForwardMetadata components:
``(qo_indptr, kv_indptr, custom_mask, mask_indptr,
window_kv_indptr, window_kv_indices, window_num_kv_splits, window_kv_offsets)``
"""
"""Fill all cuda-graph buffers for target_verify mode."""
qo_indptr = self.qo_indptr[: bs + 1]
qo_indptr[: bs + 1] = torch.arange(
0,
@@ -529,13 +504,10 @@ class TritonAttnBackend(AttentionBackend):
forward_mode: ForwardMode,
spec_info: Optional[SpecInput],
):
"""Fill QO + KV cuda-graph buffers for draft_extend mode.
Returns (qo_indptr, kv_indptr, num_tokens_per_bs).
"""
"""Fill QO + KV cuda-graph buffers for draft_extend mode."""
seq_lens = seq_lens[:bs]
# V2 draft-extend fills num_draft_tokens per req (the cuda-graph runner's
# token layout); num_steps+1 only equals that when topk == 1.
# V2 draft-extend fills num_draft_tokens per req; num_steps+1 only equals
# that when topk == 1.
num_tokens_per_bs = (
self.num_draft_tokens
if forward_mode.is_draft_extend_v2()
@@ -549,17 +521,9 @@ class TritonAttnBackend(AttentionBackend):
dtype=torch.int32,
device=self.device,
)
# DRAFT_EXTEND_V2: seq_lens = prefix + extend (bumped on the draft-extend path).
# Triton extend kernel receives extend K/V as separate tensors, so
# kv_indptr/kv_indices must cover only the prefix portion.
# extend_seq_lens_tensor is only attached to spec_info at real
# replay (eagle_draft_extend_cuda_graph_runner.replay); during the
# capture-time warmup it's absent, so fall back to zeros (matches
# the pre-unification capture path in #26651). Clamp at 0 because
# padded rows (raw_bs..bs) leave seq_lens at the fill value (1)
# while extend_seq_lens stays at num_tokens_per_bs, which would
# otherwise produce negative kv_lens; padded rows reference
# reserved req-pool slot 0 and their output is discarded.
# DRAFT_EXTEND_V2: kv_indptr/kv_indices cover only the prefix (extend K/V go
# separately). Capture warmup lacks extend_seq_lens_tensor -> fall back to
# zeros; clamp at 0 so padded rows (seq_lens==fill 1) don't go negative.
if (
spec_info is not None
and getattr(spec_info, "extend_seq_lens_tensor", None) is not None
@@ -586,8 +550,8 @@ class TritonAttnBackend(AttentionBackend):
if in_capture:
assert forward_batch.encoder_lens is None, "Not supported"
# Multi-step speculative decode: kv buffers come from spec_info
# rather than the cuda-graph pool, so replay is not involved.
# Multi-step spec decode: kv buffers come from spec_info, not the
# cuda-graph pool, so replay is not involved.
if forward_mode.is_decode_or_idle() and spec_info is not None:
self.forward_metadata = ForwardMetadata(
attn_logits=self.cuda_graph_attn_logits,
@@ -614,9 +578,16 @@ class TritonAttnBackend(AttentionBackend):
forward_mode=forward_mode,
spec_info=spec_info,
)
out_cache_loc_full_physical = self._translate_cuda_graph_shared_pool_locs(
forward_batch, bs
)
swa_out_cache_loc = self._fill_cuda_graph_swa_out_cache_loc(forward_batch)
self.forward_metadata = self._build_cuda_graph_forward_metadata(
bs, forward_mode, spec_info, swa_out_cache_loc
bs,
forward_mode,
spec_info,
swa_out_cache_loc,
out_cache_loc_full_physical,
)
else:
self._apply_cuda_graph_metadata(
@@ -626,15 +597,16 @@ class TritonAttnBackend(AttentionBackend):
forward_mode=forward_mode,
spec_info=spec_info,
)
# Metadata view is reused from capture; just refill the buffer.
# Metadata view is reused from capture; just refill the buffers.
self._translate_cuda_graph_shared_pool_locs(forward_batch, bs)
self._fill_cuda_graph_swa_out_cache_loc(forward_batch)
def _fill_cuda_graph_swa_out_cache_loc(
self, forward_batch: ForwardBatch
) -> Optional[torch.Tensor]:
"""Refill the SWA write-target buffer from the live out_cache_loc and
return the [:n] view (None for non-SWA / multi-step draft), so the
captured store reads fresh slots on replay."""
"""Refill the SWA write-target buffer from live out_cache_loc, returning the
[:n] view (None for non-SWA / multi-step draft) so the captured store reads
fresh slots on replay."""
if not self.use_sliding_window_kv_pool:
return None
out_cache_loc = forward_batch.out_cache_loc
@@ -650,6 +622,47 @@ class TritonAttnBackend(AttentionBackend):
)
return self.cuda_graph_swa_out_cache_loc[:n]
def _translate_cuda_graph_shared_pool_locs(
self, forward_batch: ForwardBatch, bs: int
) -> Optional[torch.Tensor]:
"""Unified pool: eager v2p translate of the cuda-graph read+write LOC buffers,
run BEFORE graph.replay() reading the live post-compaction v2p, so the
captured graph carries zero translate nodes. No-op for non-unified pools.
Read buffers (full kv_indices, SWA window) are translated IN PLACE; the
full-attn WRITE loc is RETURNED as the [:n] view of the backend-owned
out_cache_loc_full_physical buffer. Eager .item() bounds are fine here
(out-of-graph), so no in-graph translate variant is needed.
"""
if self._translate_kv_loc is None:
return None
# Full-attention read path.
n_kv = int(self.kv_indptr[bs].item())
if n_kv > 0:
self.cuda_graph_kv_indices[:n_kv] = self._translate_kv_loc(
self.cuda_graph_kv_indices[:n_kv]
)
# SWA window read path (hybrid-SWA unified pools only).
if self.sliding_window_size is not None and self.sliding_window_size > 0:
n_win = int(self.window_kv_indptr[bs].item())
if n_win > 0:
self.cuda_graph_window_kv_indices[:n_win] = (
self.token_to_kv_pool.translate_loc_from_full_to_swa(
self.cuda_graph_window_kv_indices[:n_win]
)
)
# Full-attention write path: translate out_cache_loc -> physical into the
# capture-stable buffer and RETURN the [:n] view.
out_cache_loc = forward_batch.out_cache_loc
n = out_cache_loc.shape[0]
# Zero the padded tail first: a smaller replay batch leaves [n:] holding
# stale ids that the captured store would write; send them to slot 0 (sink).
self.cuda_graph_out_cache_loc_full_physical[n:].zero_()
self.cuda_graph_out_cache_loc_full_physical[:n].copy_(
self._translate_kv_loc(out_cache_loc)
)
return self.cuda_graph_out_cache_loc_full_physical[:n]
def init_forward_metadata(self, forward_batch: ForwardBatch):
"""Init auxiliary variables for triton attention backend."""
@@ -663,20 +676,17 @@ class TritonAttnBackend(AttentionBackend):
if forward_batch.forward_mode.is_decode_or_idle():
if spec_info is None or spec_info.kv_indptr is None:
# kv_indptr is None for draft-extend's idle batch (no tree
# indices); build plain metadata from seq_lens.
# kv_indptr is None for draft-extend's idle batch; build from seq_lens.
if self.dcp_size > 1:
# DCP: per-rank sharded KV indices (shares _dcp_kv_indices
# with the cuda-graph path). Building full contiguous
# indices here would make each rank read the whole KV
# instead of its owner shard.
# DCP: per-rank sharded KV indices, else each rank reads the
# whole KV instead of its owner shard.
kv_indptr, kv_indices, _ = self._dcp_kv_indices(
forward_batch.req_pool_indices,
forward_batch.seq_lens,
self.kv_indptr,
)
else:
# gpu_only: seq_lens_sum may be None; ub-allocate is safe (ragged write).
# gpu_only: seq_lens_sum may be None; over-allocate is safe (ragged write).
seq_lens_sum = forward_batch.seq_lens_sum
if seq_lens_sum is None:
seq_lens_sum = bs * self.max_context_len
@@ -689,7 +699,8 @@ class TritonAttnBackend(AttentionBackend):
forward_batch.req_pool_indices,
kv_indices,
)
# Sliding window
if self._translate_kv_loc is not None:
kv_indices = self._translate_kv_loc(kv_indices)
if (
self.sliding_window_size is not None
and self.sliding_window_size > 0
@@ -755,8 +766,7 @@ class TritonAttnBackend(AttentionBackend):
dtype=torch.int32,
device=self.device,
)
# Different with flashinfer kv_indptr and kv_indices construction.
# gpu_only: seq_lens_sum may be None; ub-allocate is safe (ragged write).
# gpu_only: seq_lens_sum may be None; over-allocate is safe (ragged write).
seq_lens_sum = forward_batch.seq_lens_sum
if seq_lens_sum is None:
seq_lens_sum = bs * self.max_context_len
@@ -771,7 +781,7 @@ class TritonAttnBackend(AttentionBackend):
)
if self.sliding_window_size is not None and self.sliding_window_size > 0:
# window_kv_offsets is used to calculate the start position in custom mask
# window_kv_offsets gives the start position in custom mask
(
window_kv_indptr,
window_kv_indices,
@@ -808,8 +818,7 @@ class TritonAttnBackend(AttentionBackend):
self.kv_indptr,
)
else:
# gpu_only leaves _cpu unset; ub-allocate is safe (ragged write
# from GPU tensor, extra tail unused).
# gpu_only leaves _cpu unset; over-allocate is safe (ragged write).
if forward_batch.extend_prefix_lens_cpu is not None:
kv_indices_len = sum(forward_batch.extend_prefix_lens_cpu)
else:
@@ -825,7 +834,8 @@ class TritonAttnBackend(AttentionBackend):
forward_batch.req_pool_indices,
kv_indices,
)
# Sliding window
if self._translate_kv_loc is not None:
kv_indices = self._translate_kv_loc(kv_indices)
if self.sliding_window_size is not None and self.sliding_window_size > 0:
(
window_kv_indptr,
@@ -850,8 +860,7 @@ class TritonAttnBackend(AttentionBackend):
mask_indptr = None
attn_logits = None
attn_lse = None
# Caller usually supplies extend_seq_lens_cpu (eagle_info gpu_only
# sets host-constant mirror); defensive GPU-max fallback if not.
# Defensive GPU-max fallback when extend_seq_lens_cpu is absent.
if forward_batch.extend_seq_lens_cpu is not None:
max_extend_len = max(forward_batch.extend_seq_lens_cpu)
else:
@@ -864,6 +873,17 @@ class TritonAttnBackend(AttentionBackend):
forward_batch.out_cache_loc
)
# Unified pool full-attention WRITE loc (virtual out_cache_loc -> physical),
# carried in the metadata (-> KVWriteLoc.full_loc). None for non-unified pools.
out_cache_loc_full_physical = None
if (
self._translate_kv_loc is not None
and forward_batch.out_cache_loc is not None
):
out_cache_loc_full_physical = self._translate_kv_loc(
forward_batch.out_cache_loc
)
self.forward_metadata = ForwardMetadata(
attn_logits,
attn_lse,
@@ -880,6 +900,7 @@ class TritonAttnBackend(AttentionBackend):
window_kv_offsets,
swa_attn_logits=swa_attn_logits,
swa_out_cache_loc=swa_out_cache_loc,
out_cache_loc_full_physical=out_cache_loc_full_physical,
)
def init_cuda_graph_state(
@@ -970,12 +991,22 @@ class TritonAttnBackend(AttentionBackend):
device=self.device,
)
if self._translate_kv_loc is not None:
# Unified pool full-attention write-target buffer, refilled at replay
# (-> KVWriteLoc.full_loc). Capture-stable, mirrors cuda_graph_swa_out_cache_loc.
self.cuda_graph_out_cache_loc_full_physical = torch.zeros(
(max_num_tokens,),
dtype=torch.int64,
device=self.device,
)
def _build_cuda_graph_forward_metadata(
self,
bs: int,
forward_mode: ForwardMode,
spec_info: Optional[SpecInput],
swa_out_cache_loc: Optional[torch.Tensor] = None,
out_cache_loc_full_physical: Optional[torch.Tensor] = None,
) -> ForwardMetadata:
"""Construct ForwardMetadata from the current cuda-graph buffer state.
@@ -1006,6 +1037,7 @@ class TritonAttnBackend(AttentionBackend):
window_kv_offsets=None,
swa_attn_logits=self.cuda_graph_swa_attn_logits,
swa_out_cache_loc=swa_out_cache_loc,
out_cache_loc_full_physical=out_cache_loc_full_physical,
)
elif forward_mode.is_target_verify():
custom_mask = (
@@ -1031,6 +1063,7 @@ class TritonAttnBackend(AttentionBackend):
),
window_kv_offsets=self.cuda_graph_window_kv_offsets if swa else None,
swa_out_cache_loc=swa_out_cache_loc,
out_cache_loc_full_physical=out_cache_loc_full_physical,
)
elif forward_mode.is_draft_extend_v2():
return ForwardMetadata(
@@ -1056,6 +1089,7 @@ class TritonAttnBackend(AttentionBackend):
window_num_kv_splits=None,
window_kv_offsets=None,
swa_out_cache_loc=swa_out_cache_loc,
out_cache_loc_full_physical=out_cache_loc_full_physical,
)
else:
raise ValueError(f"Invalid forward mode: {forward_mode=} for CUDA Graph.")
@@ -1183,6 +1217,7 @@ class TritonAttnBackend(AttentionBackend):
loc_info = KVWriteLoc(
forward_batch.out_cache_loc,
self.forward_metadata.swa_out_cache_loc,
full_loc=self.forward_metadata.out_cache_loc_full_physical,
)
if layer.k_scale is None:
self._set_kv_buffer(forward_batch, layer, loc_info, k, v)
@@ -1628,6 +1663,7 @@ class TritonAttnBackend(AttentionBackend):
KVWriteLoc(
forward_batch.out_cache_loc,
self.forward_metadata.swa_out_cache_loc,
full_loc=self.forward_metadata.out_cache_loc_full_physical,
),
k,
v,
@@ -1908,12 +1944,21 @@ def update_sliding_window_buffer(
device=None,
token_to_kv_pool=None,
window_kv_indices=None,
skip_full_to_swa_translation=False,
):
"""Fill window KV buffers for sliding-window attention.
Pass window_kv_indices to write into a pre-allocated buffer (CUDA-graph
path); omit it (or pass None) to allocate a fresh tensor (eager path,
requires device).
Pass ``window_kv_indices`` to write into a pre-allocated buffer (CUDA-graph
path); omit it (or pass ``None``) to allocate a fresh tensor (eager path,
requires ``device``).
``skip_full_to_swa_translation=True`` leaves ``window_kv_indices`` as VIRTUAL
full-token ids (no eager full->swa translate). The unified-memory-pool cuda-graph
builder passes this so the window translate is deferred to
``TritonAttnBackend._translate_cuda_graph_shared_pool_locs`` (run in
``init_forward_metadata_out_graph``, BEFORE ``graph.replay()``), which reads
the live v2p and rewrites the static window buffer to swa-physical in place;
baseline SWA leaves it False (eager).
"""
window_kv_lens = torch.minimum(
seq_lens,
@@ -1935,7 +1980,9 @@ def update_sliding_window_buffer(
window_kv_indices,
req_to_token.stride(0),
)
if hasattr(token_to_kv_pool, "translate_loc_from_full_to_swa"):
if not skip_full_to_swa_translation and hasattr(
token_to_kv_pool, "translate_loc_from_full_to_swa"
):
kv_last_index = window_kv_indptr[-1]
window_kv_indices[:kv_last_index] = (
token_to_kv_pool.translate_loc_from_full_to_swa(
python/sglang/srt/managers/schedule_policy.py
+101 -11
View File
@@ -52,6 +52,9 @@ from sglang.srt.mem_cache.base_prefix_cache import (
MatchPrefixParams,
zero_match_result,
)
from sglang.srt.mem_cache.multi_ended_allocator import (
UnifiedMambaTokenToKVPoolAllocator,
)
from sglang.srt.mem_cache.radix_cache import RadixCache, RadixKey, TreeNode
from sglang.srt.server_args import ServerArgs, get_global_server_args
@@ -495,6 +498,34 @@ class PrefillAdder:
self.rem_swa_token_offset = 0
# Unified-pool joint budget: a new mamba state consumes shared-gap bytes
# that `rem_total_tokens` (full KV) otherwise counts as free, so reserve
# the gap per new mamba slot or admission over-commits. Gate on the
# ALLOCATOR being the unified Mamba composite, NOT on `is_hybrid_ssm_cache`
# (False for `ChunkCache`, which would skip the reservation on the
# chunk-cache path): the gap coupling is a property of the byte buffer.
self._mamba_slot_cost = 0
if isinstance(
self.token_to_kv_pool_allocator, UnifiedMambaTokenToKVPoolAllocator
):
self._mamba_slot_cost = (
self.token_to_kv_pool_allocator.mamba_slot_full_token_cost()
)
# `mamba_gap_reserve` is charged to `rem_total_tokens`, which INCLUDES
# `full_evictable_size()` — but `alloc_req_slots` can only recover
# MAMBA-recoverable bytes for a mamba slot (shared gap + peer holes +
# mamba-evictable radix), NOT full-evictable. Gate new mamba slots on
# that mamba-recoverable budget separately or an over-admit hits the
# fail-loud `RuntimeError`. `None` outside the unified Mamba pool.
self.rem_mamba_slots = None
if self._mamba_slot_cost:
self.rem_mamba_slots = (
self.token_to_kv_pool_allocator.mamba_allocator.schedulable_available_size()
)
if self.is_hybrid_ssm_cache:
self.rem_mamba_slots += self.tree_cache.mamba_evictable_size()
self.priority_scheduling_preemption_threshold = (
priority_scheduling_preemption_threshold
)
@@ -602,6 +633,21 @@ class PrefillAdder:
budget += self.ceil_paged_tokens(swa_host_hit_length)
return budget
def _mamba_gap_budget_for_req(self, req: Req) -> int:
"""Shared-gap reservation (full-token-equivalents) for a request's new
mamba state. Charged only on the SHARED Mamba pool (`_mamba_slot_cost > 0`)
and only when the req has no state yet (`mamba_pool_idx is None`, mirroring
`HybridReqToTokenPool.alloc`); 0 keeps baseline / SWA / non-Mamba unchanged.
Conservative by design (`_mamba_slot_cost` rounds UP). Does NOT reserve
radix COW headroom or locked-but-evictable bytes — that residual is
backstopped by the fail-loud RuntimeError in `alloc_req_slots`. FIXME: if
over-admission crashes under pressure, make this more conservative (e.g.
multiply by `MAMBA_STATE_PER_REQ_PREFIX_CACHE`)."""
if self._mamba_slot_cost and req.mamba_pool_idx is None:
return self._mamba_slot_cost
return 0
def ceil_paged_tokens(self, tokens: int) -> int:
return -(-tokens // self.page_size) * self.page_size
@@ -609,6 +655,10 @@ class PrefillAdder:
no_token = self.rem_total_tokens <= 0 or self.cur_rem_tokens <= 0
if not no_token and self.is_hybrid_swa:
no_token = self.rem_swa_tokens <= 0
# Gate new mamba slots separately: rem_total_tokens' full_evictable can't
# cover a mamba slot, which needs mamba-recoverable bytes (see __init__).
if not no_token and self.rem_mamba_slots is not None:
no_token = self.rem_mamba_slots <= 0
if no_token:
return AddReqResult.NO_TOKEN
@@ -630,14 +680,27 @@ class PrefillAdder:
extend_input_len: int,
max_new_tokens: int,
retracted_stain: bool,
mamba_gap_reserve: int = 0,
):
# TODO(lsyin): check this workaround logic, which only ensures the prefill will not out of memory, and may be too conservative
extend_input_len = self.ceil_paged_tokens(extend_input_len)
# alloc_extend reserves an extra page_size per request to make sure the budget doesn't over-commit
page_overhead = self.page_size
self.rem_total_token_offset += extend_input_len + max_new_tokens + page_overhead
self.cur_rem_token_offset += extend_input_len + page_overhead
# `mamba_gap_reserve` (shared Mamba pool only; 0 otherwise) charges the new
# mamba state's shared-gap cost to BOTH full budgets: the slot is allocated
# immediately (counts against `cur_rem`) and held for the request lifetime
# (counts against `rem_total`). See `_mamba_gap_budget_for_req`.
self.rem_total_token_offset += (
extend_input_len + max_new_tokens + page_overhead + mamba_gap_reserve
)
self.cur_rem_token_offset += (
extend_input_len + page_overhead + mamba_gap_reserve
)
# The new mamba slot also consumes one mamba-recoverable slot (gated
# separately so full_evictable can't cover it — see __init__).
if mamba_gap_reserve and self.rem_mamba_slots is not None:
self.rem_mamba_slots -= 1
self.rem_input_tokens -= extend_input_len
if self.is_hybrid_swa:
@@ -681,7 +744,13 @@ class PrefillAdder:
self.can_run_list.append(req)
self._update_prefill_budget(prefix_len, trunc_len, 0, req.retracted_stain)
self._update_prefill_budget(
prefix_len,
trunc_len,
0,
req.retracted_stain,
mamba_gap_reserve=self._mamba_gap_budget_for_req(req),
)
def _req_inc_lock_ref(self, req: Req):
result = self.tree_cache.inc_lock_ref(req.last_node)
@@ -711,7 +780,11 @@ class PrefillAdder:
else 0
)
self._update_prefill_budget(
0, req.extend_range.length, max_new_tokens, req.retracted_stain
0,
req.extend_range.length,
max_new_tokens,
req.retracted_stain,
mamba_gap_reserve=self._mamba_gap_budget_for_req(req),
)
# Return based on remaining token availability
@@ -755,6 +828,7 @@ class PrefillAdder:
else 0
),
req.retracted_stain,
mamba_gap_reserve=self._mamba_gap_budget_for_req(req),
)
# Return if chunked prefill not finished
@@ -782,6 +856,9 @@ class PrefillAdder:
req.prefix_indices
)
paged_input = self.ceil_paged_tokens(cand_extend_input_len)
# Shared Mamba pool: fold the new mamba state's shared-gap cost into the
# budget gate so admission can't over-commit (0 for baseline / non-Mamba).
paged_input += self._mamba_gap_budget_for_req(req)
if paged_input > min(self.cur_rem_tokens, self.rem_total_tokens):
return AddReqResult.NO_TOKEN
if self.is_hybrid_swa:
@@ -863,6 +940,7 @@ class PrefillAdder:
req.extend_range.length,
min(req.sampling_params.max_new_tokens, CLIP_MAX_NEW_TOKENS),
req.retracted_stain,
mamba_gap_reserve=self._mamba_gap_budget_for_req(req),
)
else:
if self.rem_chunk_tokens <= 0:
@@ -877,7 +955,13 @@ class PrefillAdder:
)
self.can_run_list.append(req)
self.new_chunked_req = req
self._update_prefill_budget(0, trunc_len, 0, req.retracted_stain)
self._update_prefill_budget(
0,
trunc_len,
0,
req.retracted_stain,
mamba_gap_reserve=self._mamba_gap_budget_for_req(req),
)
return self.budget_state()
@@ -906,11 +990,9 @@ class PrefillAdder:
if req.sampling_params.ignore_eos and getattr(self.tree_cache, "disable", True):
return self.add_one_req_ignore_eos(req)
# Reserve page_size for page-alignment overhead. The paged allocator
# may consume up to one extra page per request (see alloc_extend), and
# _update_prefill_budget already accounts for this in the deduction.
# Without this, admission is more optimistic than the actual budget
# deduction, allowing over-admission when the pool is nearly full.
# Reserve page_size for page-alignment overhead: the paged allocator may
# consume one extra page per request (see alloc_extend), which
# _update_prefill_budget also deducts.
max_new = min(
max(req.sampling_params.max_new_tokens - len(req.output_ids), 0),
CLIP_MAX_NEW_TOKENS,
@@ -919,6 +1001,9 @@ class PrefillAdder:
req.prefix_indices
)
total_tokens = cand_extend_input_len + max_new + self.page_size
# Shared Mamba pool: fold the new mamba state's shared-gap cost into
# `total_tokens` so both `rem_total_tokens` gates reflect the joint budget.
total_tokens += self._mamba_gap_budget_for_req(req)
# adjusting the input_tokens based on host_hit_length and page_size
real_input_tokens = cand_extend_input_len - req.host_hit_length
@@ -1009,6 +1094,7 @@ class PrefillAdder:
CLIP_MAX_NEW_TOKENS,
),
req.retracted_stain,
mamba_gap_reserve=self._mamba_gap_budget_for_req(req),
)
else:
# Make sure at least one page is available
@@ -1045,7 +1131,11 @@ class PrefillAdder:
self._req_inc_lock_ref(req)
self._update_prefill_budget(
prefix_len, trunc_len, 0, req.retracted_stain
prefix_len,
trunc_len,
0,
req.retracted_stain,
mamba_gap_reserve=self._mamba_gap_budget_for_req(req),
)
return self.budget_state()
python/sglang/srt/managers/scheduler.py
+29 -2
View File
@@ -1496,7 +1496,7 @@ class Scheduler(
if self.device == "cpu":
self.schedule_stream.synchronize = lambda: None # No-op for CPU
# The global WAR barrier fences the scheduler's next shared-buffer write
# on the previous forward's read of the shared pool.
# on the previous forward's read of the unified memory pool.
self._war_barrier_enabled = is_cuda() or envs.SGLANG_ENABLE_WAR_BARRIER.get()
with self.device_module.StreamContext(self.schedule_stream):
dispatch_event_loop(self)
@@ -1580,6 +1580,14 @@ class Scheduler(
# we can process the last batch immediately.
if disable_overlap_for_batch:
pop_and_process()
# Opportunistic flush at the disable_overlap sync boundary:
# forward_stream is idle (prev forward drained, next not launched),
# so `_flush`'s non-urgent guard compacts freely. Sync-free, best-effort.
if self.server_args.enable_unified_memory:
try:
self.token_to_kv_pool_allocator.flush_opportunistic()
except Exception:
pass
# Launch the current batch
if batch:
@@ -2527,7 +2535,6 @@ class Scheduler(
release_kv_cache(req, self.tree_cache, is_insert=False)
self.chunked_req = None
self._chunked_req_scheduled_last_iter = False
self._pending_chunked_abort_req = None
self.ipc_channels.send_to_tokenizer.send_output(AbortReq(rid=req.rid), req)
logger.debug(f"Abort chunked prefill request. {req.rid=}")
@@ -3237,6 +3244,20 @@ class Scheduler(
self.batch_record_buf[self.batch_record_ct].extend(
batch_result.extra_keep_alive_refs
)
if self.server_args.enable_unified_memory:
# Record a `forward_done` event after the forward (before
# copy_to_cpu); lazy-compaction `_flush` gates src reuse on
# it. Only the unified pool's allocator exposes these hooks.
allocator = self.token_to_kv_pool_allocator
forward_done = self.device_module.Event()
forward_done.record(stream=self.forward_stream)
allocator.set_latest_forward_done_event(forward_done)
# Write-set classification: hand the allocator this
# forward's virtual out_cache_loc as a tensor ref (no GPU work).
allocator.set_inflight_forward(
forward_done,
batch.out_cache_loc,
)
# FIXME(lsyin): maybe move this to forward_batch_generation
batch_result.copy_done = self.device_module.Event()
if batch_result.delay_sample_func is None:
@@ -3497,6 +3518,12 @@ class Scheduler(
if not self.is_fully_idle():
return
if self.server_args.enable_unified_memory:
try:
self.token_to_kv_pool_allocator.flush_opportunistic()
except Exception:
pass
# memory leak check (skipped for hisparse — pool counters intentionally
# diverge during host-backup, see _get_swa_token_info clamp).
if not self.enable_hisparse:
python/sglang/srt/managers/scheduler_components/invariant_checker.py
+10 -2
View File
@@ -88,8 +88,16 @@ class SchedulerInvariantChecker:
protected = self.tree_cache.full_protected_size()
session_held = self.pool_stats_observer.session_held_full_tokens()
total = self.full_tokens_per_layer
elif self.is_hybrid_ssm and self.tree_cache.supports_mamba():
protected = self.tree_cache.full_protected_size()
elif self.is_hybrid_ssm:
# Branch on cache type for the protected accessor (MambaRadixCache
# splits full/mamba; ChunkCache only has the single protected_size).
# Use the allocator's `.size` for `total`: static max_total_num_tokens for
# non-unified pools, the dynamic byte-coordinated cap (matching
# `available_size`) for the unified pool.
if self.tree_cache.supports_mamba():
protected = self.tree_cache.full_protected_size()
else:
protected = self.tree_cache.protected_size()
session_held = self.pool_stats_observer.session_held_tokens()
total = self.token_to_kv_pool_allocator.size
else:
python/sglang/srt/mem_cache/allocator/mamba.py
+7
View File
@@ -48,6 +48,13 @@ class MambaSlotAllocator:
def available_size(self) -> int:
return len(self.free_slots)
def schedulable_available_size(self) -> int:
"""Planner-facing free count. Identity to ``available_size`` for the
static pool (slot-count and byte-coordinated views coincide); the shared
``UnifiedMambaSlotAllocator`` overrides it with the byte-coordinated view.
Lets ``alloc_req_slots`` call it uniformly without a getattr fallback."""
return self.available_size()
def alloc_group_begin(self, num_reqs: int):
"""Pre-allocate a batch of slots for match_prefix to amortize overhead."""
self._alloc_iter = None
python/sglang/srt/mem_cache/allocator/swa.py
+9
View File
@@ -112,6 +112,15 @@ class SWATokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
def swa_available_size(self):
return self.swa_attn_allocator.available_size()
# Slot-conservation views for the leak invariant. On the non-shared allocator
# the static budget IS physical (conserve == physical); the shared composite
# overrides these with the static-cap view.
def _conserve_full_available_size(self):
return self.full_available_size()
def _conserve_swa_available_size(self):
return self.swa_available_size()
@property
def size(self):
return min(self._size_full, self._size_swa)
python/sglang/srt/mem_cache/common.py
+36 -52
View File
@@ -87,11 +87,8 @@ def free_swa_out_of_window_slots(
req.swa_evicted_seqlen = max(req.swa_evicted_seqlen, evict_floor)
# Subtract an extra page_size so the eviction frontier never reaches the
# radix tree insert boundary (page_floor(seq_len)). This keeps at least one
# page of non-evicted SWA KV for the tree to store as a non-tombstone node,
# preserving cache reuse in multi-turn scenarios. Without this, leaf nodes
# may become tombstoned, causing SWA memory leak.
# See also: _insert_helper case 3 in swa_radix_cache.py (defensive counterpart).
# radix tree insert boundary, keeping >=1 page of non-evicted SWA KV for the
# tree to store as a non-tombstone node (else leaf nodes tombstone -> SWA leak).
if drop_page_margin or envs.SGLANG_OPT_SWA_EVICT_DROP_PAGE_MARGIN.get():
evict_threshold = pre_len - sliding_window_size
else:
@@ -220,9 +217,7 @@ def get_alloc_len_per_decode(server_args: Optional[ServerArgs] = None) -> int:
if server_args.speculative_algorithm is None:
return 1
# Spec decoding allocates max(topk * num_steps, num_draft_tokens) per
# decode step (draft chain and verify block share the reservation).
# Spec decoding allocates max(topk * num_steps, num_draft_tokens) per decode step.
spec_steps = server_args.speculative_num_steps or 1
spec_topk = server_args.speculative_eagle_topk or 1
spec_tokens = server_args.max_speculative_num_draft_tokens
@@ -234,10 +229,9 @@ def get_alloc_len_per_decode(server_args: Optional[ServerArgs] = None) -> int:
if page_size == 1 or spec_topk == 1 or not spec_algo.has_draft_kv():
return max(spec_steps * spec_topk, spec_tokens)
else:
# page_size > 1 + topk > 1 (spec v2 tree): worst-case page-aligned tree
# footprint. Per topk branch needs ceil((last_page_len + num_steps) / page)
# pages; the partial tail page can be up to page_size - 1, and each branch
# gets its own (duplicated) copy -- so reserve for all topk branches.
# spec v2 tree (page>1, topk>1): worst-case page-aligned footprint per
# topk branch is ceil((page_size-1 + num_steps) / page) pages, each branch
# duplicated -- reserve for all topk branches.
num_new_pages_per_topk = (
(page_size - 1) + spec_steps + page_size - 1
) // page_size
@@ -256,12 +250,10 @@ def get_alloc_reserve_per_decode(server_args: Optional[ServerArgs] = None) -> in
def get_req_to_token_extra_context_len(server_args: ServerArgs) -> int:
"""req_to_token row headroom beyond the model context length.
Sized to hold the decode over-allocation (kv_committed_len +
get_alloc_reserve_per_decode). The spec v2 page>1 topk>1 holey draft footprint
can outgrow the default num_draft_tokens headroom (PR #26972).
Sized to hold the decode over-allocation; the spec v2 page>1 topk>1 holey
draft footprint can outgrow the default num_draft_tokens headroom.
"""
# FIXME(lsyin): this is the temporary fix for the context length issue when
# using speculative decoding
# FIXME(lsyin): temporary fix for the context length issue under spec decoding
extra = 4 + (server_args.max_speculative_num_draft_tokens or 0)
if (
server_args.speculative_algorithm is not None
@@ -327,13 +319,8 @@ def evict_from_tree_cache(tree_cache: BasePrefixCache | None, num_tokens: int):
def _compute_dsv4_state_lens(batch, *, is_decode: bool):
"""Per-req c{4,128}_state pool alloc lens (a ``DSV4StateLens``) for this
alloc step. The DSV4-NPU allocator owns the computation (it also mutates the
per-req cumulative state on each ``Req``); we just trigger it here, right
before the paged alloc that consumes the result.
None on CUDA / non-V4 paths (allocator has no ``compute_dsv4_state_lens_*``)
so the ``alloc_paged_token_slots_*`` forwarding stays a no-op.
"""Per-req c{4,128}_state pool alloc lens (``DSV4StateLens``) for this step.
None on CUDA / non-V4 paths (allocator has no ``compute_dsv4_state_lens_*``).
"""
allocator = batch.token_to_kv_pool_allocator
if not hasattr(allocator, "compute_dsv4_state_lens_extend"):
@@ -371,8 +358,7 @@ def alloc_paged_token_slots_extend(
extra_alloc_kwargs = {}
if is_dsv4:
extra_alloc_kwargs["req_pool_indices"] = req_pool_indices
# Pass the per-req tables in per call for the c-pool / state last_loc
# lookup; the allocator holds no reference to the pool.
# Per-call per-req tables for the c-pool / state last_loc lookup.
if batch is not None:
extra_alloc_kwargs["req_to_token_pool"] = batch.req_to_token_pool
if dsv4_state_lens is not None:
@@ -415,10 +401,20 @@ def alloc_req_slots(
reqs: list[Req],
tree_cache: BasePrefixCache | None,
) -> list[int]:
"""Allocate request slots from the pool."""
"""Allocate request slots from the pool.
Fail-loud: raises ``RuntimeError`` if the pool can't satisfy the batch. An
alloc failure here means the admission budget (``PrefillAdder``) was wrong
and should surface rather than be masked.
"""
num_reqs = len(reqs)
if isinstance(req_to_token_pool, HybridReqToTokenPool):
mamba_available_size = req_to_token_pool.mamba_allocator.available_size()
# Byte-coordinated for the shared allocator (accounts for the peer full
# sub-pool's bytes); plain slot free count for the non-shared one.
mamba_available_size = (
req_to_token_pool.mamba_allocator.schedulable_available_size()
)
# Eviction headroom factor: 3x (or lazy variant) for radix COW, 1x for chunk.
if tree_cache.supports_mamba():
factor = (
MAMBA_STATE_PER_REQ_PREFIX_CACHE_LAZY
@@ -433,24 +429,19 @@ def alloc_req_slots(
mamba_num = max(0, mamba_state_needed - mamba_available_size)
tree_cache.evict(EvictParams(num_tokens=0, mamba_num=mamba_num))
req_pool_indices = req_to_token_pool.alloc(reqs)
if req_pool_indices is None:
raise RuntimeError(
"alloc_req_slots runs out of memory. "
"Please set a smaller number for `--max-running-requests`. "
f"{req_to_token_pool.available_size()=}, "
f"{num_reqs=}, "
f"{req_to_token_pool.available_size()=}, {num_reqs=}, "
)
return req_pool_indices
def _alloc_page_size(batch: ScheduleBatch) -> int:
# DCP (HIP & CUDA only) swaps in a PagedTokenToKVPoolAllocator whose
# page_size is server_args.page_size * dcp_size, so it can be > 1 even when
# tree_cache.page_size (== server_args.page_size) is 1. Only on the HIP DCP
# path do we branch on the real allocator's page_size so the paged path is
# taken; everywhere else tree_cache.page_size is authoritative and the two
# are equal (dcp_size == 1), so behavior is unchanged.
# DCP swaps in an allocator whose page_size is server_args.page_size *
# dcp_size, so it can be > 1 even when tree_cache.page_size is 1; branch on
# the real allocator's page_size there. Elsewhere the two are equal.
if (_is_hip or _is_cuda) and get_global_server_args().dcp_size > 1:
return batch.tree_cache.token_to_kv_pool_allocator.page_size
return batch.tree_cache.page_size
@@ -462,10 +453,9 @@ def alloc_for_extend(
"""
Allocate KV cache for extend batch and write to req_to_token_pool.
Returns:
out_cache_loc: allocated cache locations
req_pool_indices_device: request pool indices as a device tensor
req_pool_indices_cpu: request pool indices as a CPU tensor (host mirror)
Returns ``(out_cache_loc, req_pool_indices_device, req_pool_indices_cpu)``
(the last is the host/CPU mirror). ``alloc_req_slots`` raises ``RuntimeError``
if the pool can't satisfy the batch (fail-loud — see its docstring).
"""
# free out-of-window swa tokens
batch.maybe_evict_swa()
@@ -478,7 +468,7 @@ def alloc_for_extend(
prefix_lens_device = prefix_lens_cpu.to(batch.device, non_blocking=True)
extend_lens_device = extend_lens_cpu.to(batch.device, non_blocking=True)
# Allocate req slots
# Allocate req slots (raises RuntimeError if the pool is exhausted)
req_pool_indices = alloc_req_slots(
batch.req_to_token_pool, batch.reqs, batch.tree_cache
)
@@ -489,8 +479,6 @@ def alloc_for_extend(
if _alloc_page_size(batch) == 1:
out_cache_loc = alloc_token_slots(batch.tree_cache, batch.extend_num_tokens)
else:
# Since tree_cache.page_size is (page_size * dcp_world_size), for dcp
# on cuda platform, always use alloc_paged_token_slots_extend
# Paged allocation - build last_loc
last_loc = [
(t[-1:] if len(t) > 0 else torch.tensor([-1], device=batch.device))
@@ -524,8 +512,7 @@ def alloc_for_extend(
batch.req_to_token_pool,
)
# DSV4-NPU hook: write c4/c128/swa per-req tables from the stashed bundle.
# No-op on non-DSV4 paths (out_cache_loc_dsv4 stays None there).
# DSV4-NPU hook: no-op on non-DSV4 paths.
if _is_npu:
maybe_write_dsv4_extend(
batch,
@@ -559,8 +546,7 @@ def alloc_paged_token_slots_decode(
extra_alloc_kwargs = {}
if is_dsv4:
extra_alloc_kwargs["req_pool_indices"] = req_pool_indices
# Per-call per-req tables for the last_loc lookup; the allocator holds
# no reference to the pool.
# Per-call per-req tables for the last_loc lookup.
if batch is not None:
extra_alloc_kwargs["req_to_token_pool"] = batch.req_to_token_pool
if dsv4_state_lens is not None:
@@ -633,8 +619,7 @@ def alloc_for_decode(batch: ScheduleBatch, token_per_req: int) -> torch.Tensor:
(batch.req_pool_indices, locs), out_cache_loc.to(torch.int32)
)
# DSV4-NPU hook: post-decode write of c4/c128/swa per-req tables from the
# stashed bundle. No-op on non-DSV4 paths (out_cache_loc_dsv4 stays None).
# DSV4-NPU hook: no-op on non-DSV4 paths.
if _is_npu:
maybe_write_dsv4_decode(
batch,
@@ -698,8 +683,7 @@ def release_kv_cache(req: Req, tree_cache: BasePrefixCache, is_insert: bool = Tr
req.mamba_pool_idx is not None
), "mamba state is freed while the tree cache does not manage mamba states"
tree_cache.req_to_token_pool.free_mamba_cache(req)
# The DSV4-NPU ReqToTokenPool subclass's free() additionally releases the
# c4/c128 state pages; other ReqToTokenPool subclasses are a no-op here.
# DSV4-NPU's free() also releases c4/c128 state pages; no-op for others.
tree_cache.req_to_token_pool.free(req)
python/sglang/srt/mem_cache/mamba_radix_cache.py
+17 -4
View File
@@ -46,6 +46,9 @@ from sglang.srt.mem_cache.base_prefix_cache import (
)
from sglang.srt.mem_cache.events import KVCacheEventMixin
from sglang.srt.mem_cache.memory_pool import HybridReqToTokenPool
from sglang.srt.mem_cache.multi_ended_allocator import (
UnifiedMambaTokenToKVPoolAllocator,
)
from sglang.srt.mem_cache.radix_cache import RadixKey
from sglang.srt.mem_cache.utils import split_node_hash_value
from sglang.srt.server_args import get_global_server_args
@@ -420,9 +423,15 @@ class LRUList:
class MambaRadixCache(KVCacheEventMixin, BasePrefixCache):
def __init__(self, params: CacheInitParams):
assert isinstance(
params.token_to_kv_pool_allocator, TokenToKVPoolAllocator
) or isinstance(params.token_to_kv_pool_allocator, PagedTokenToKVPoolAllocator)
assert (
isinstance(params.token_to_kv_pool_allocator, TokenToKVPoolAllocator)
or isinstance(
params.token_to_kv_pool_allocator, PagedTokenToKVPoolAllocator
)
or isinstance(
params.token_to_kv_pool_allocator, UnifiedMambaTokenToKVPoolAllocator
)
)
self.req_to_token_pool: HybridReqToTokenPool = params.req_to_token_pool
self.token_to_kv_pool_allocator = params.token_to_kv_pool_allocator
self.mamba_cache_chunk_size = get_global_server_args().mamba_cache_chunk_size
@@ -694,8 +703,12 @@ class MambaRadixCache(KVCacheEventMixin, BasePrefixCache):
)
else:
mamba_value_donated = self._alloc_mamba_slot()
# mamba_pool is a pure PHYSICAL store; translate both slot ids
# virtual->physical (identity for the non-unified memory pool) before the copy.
translate = self.req_to_token_pool.translate_mamba_indices
self.req_to_token_pool.mamba_pool.copy_from(
req.mamba_pool_idx.unsqueeze(0), mamba_value_donated
translate(req.mamba_pool_idx.unsqueeze(0)),
translate(mamba_value_donated),
)
result = self.insert(
python/sglang/srt/mem_cache/memory_pool.py
+59 -19
View File
@@ -970,6 +970,14 @@ class HybridReqToTokenPool(ReqToTokenPool):
def get_mamba_indices(self, req_indices: torch.Tensor) -> torch.Tensor:
return self.req_index_to_mamba_index_mapping[req_indices]
def translate_mamba_indices(self, mamba_indices: torch.Tensor) -> torch.Tensor:
"""Virtual->physical mamba-slot translate. Identity for a static pool
(slots are physical); UnifiedHybridReqToTokenPool overrides it for the
unified memory pool, where mamba slot ids are virtual. Callers translate
before calling the pool's physical-id state ops (copy_from / clear_slots
/ get_cpu_copy / load_cpu_copy)."""
return mamba_indices
def mamba2_layer_cache(self, layer_id: int):
assert layer_id in self.mamba_map
if self.layer_transfer_counter is not None:
@@ -1134,29 +1142,43 @@ class HybridReqToTokenPool(ReqToTokenPool):
class KVWriteLoc:
"""Write target(s) for ``KVCache.set_kv_buffer``.
``loc`` is the full-pool write location; ``swa_loc`` is the pre-translated
full->SWA location for hybrid SWA pools (``None`` otherwise). Bundling them
lets a backend issue one ``set_kv_buffer`` call regardless of pool type.
All location info lives here (in the attention metadata), NOT in the pool:
- ``loc``: the generic per-token write location (the allocated
``out_cache_loc``). VIRTUAL under the unified memory pool (it indexes the
virtual slot space); already physical for a non-unified memory pool.
- ``swa_loc``: the pre-translated SWA-sub-pool PHYSICAL location for hybrid
SWA pools (``None`` otherwise).
- ``full_loc``: the pre-translated full-attention-sub-pool PHYSICAL location
for the unified memory pool (``None`` otherwise), computed once per forward in
attention metadata (``ForwardMetadata.out_cache_loc_full_physical``). The
shared full pool writes it directly; the pool never translates (replacing
the former per-layer v2p gather / ``set_full_loc`` pin).
``swa_loc`` and ``full_loc`` are the parallel pair (each a pre-resolved
PHYSICAL loc into its sub-pool, mirroring ``swa_kv_pool`` / ``full_kv_pool``);
``loc`` is the generic, possibly-virtual fallback. Bundling them lets a
backend issue one ``set_kv_buffer`` call regardless of pool type.
"""
loc: torch.Tensor
swa_loc: Optional[torch.Tensor] = None
full_loc: Optional[torch.Tensor] = None
def __post_init__(self):
# swa_out_cache_loc is computed once at metadata-init time from the
# full (possibly padded) out_cache_loc. Piecewise CUDA graphs later
# narrow out_cache_loc to real_num_tokens per layer, so swa_loc can
# be longer than loc. Slice to match since both are in the same
# per-token order.
# swa_loc / full_loc are resolved once at metadata-init from the full
# (padded) out_cache_loc; piecewise/DP-padded paths later narrow loc per
# layer, so slice these pre-resolved locs to match (same per-token order).
if self.swa_loc is not None and self.swa_loc.shape[0] != self.loc.shape[0]:
self.swa_loc = self.swa_loc[: self.loc.shape[0]]
if self.full_loc is not None and self.full_loc.shape[0] != self.loc.shape[0]:
self.full_loc = self.full_loc[: self.loc.shape[0]]
def unwrap_write_loc(loc_info):
"""Return ``(loc, swa_loc)`` from a ``KVWriteLoc`` or a bare loc tensor."""
"""Return ``(loc, swa_loc, full_loc)`` from a ``KVWriteLoc`` or a bare loc."""
if isinstance(loc_info, KVWriteLoc):
return loc_info.loc, loc_info.swa_loc
return loc_info, None
return loc_info.loc, loc_info.swa_loc, loc_info.full_loc
return loc_info, None, None
class KVCache(abc.ABC):
@@ -1648,7 +1670,7 @@ class MHATokenToKVPool(KVCache):
layer_id_override: Optional[int] = None,
dcp_kv_mask: Optional[torch.Tensor] = None,
):
loc, _ = unwrap_write_loc(loc_info)
loc, _, _ = unwrap_write_loc(loc_info)
# Catch stale slot ids here instead of as illegal-addr / silent KV
# corruption in the store_kvcache write (gated on SGLANG_ENABLE_ASYNC_ASSERT).
maybe_detect_oob(loc, 0, self.size + self.page_size, "set_kv_buffer (MHA)")
@@ -2123,7 +2145,7 @@ class MHATokenToKVPoolFP4(MHATokenToKVPool):
v_scale: Optional[float] = None,
layer_id_override: Optional[int] = None,
):
loc, _ = unwrap_write_loc(loc_info)
loc, _, _ = unwrap_write_loc(loc_info)
maybe_detect_oob(loc, 0, self.size + self.page_size, "set_kv_buffer (MHA-FP4)")
from sglang.srt.model_executor.runner import get_is_capture_mode
@@ -2346,6 +2368,9 @@ class HybridLinearKVPool(KVCache):
qk_rope_head_dim: int = None,
start_layer: Optional[int] = None,
full_kv_pool_class: Optional[type] = None,
# When provided (shared-KV-pool path), use this pool for the
# full-attention layers instead of constructing one internally.
full_kv_pool: Optional[KVCache] = None,
):
self.size = size
self.dtype = dtype
@@ -2357,8 +2382,15 @@ class HybridLinearKVPool(KVCache):
self.head_num = head_num
self.head_dim = head_dim
self.mamba_pool = mamba_pool
# virtual->physical mamba-slot translate for the HiCache offload path;
# identity for a static pool, the allocator's `translate` for the unified pool.
self._mamba_translate = lambda ids: ids
self.use_mla = use_mla
if not use_mla:
if full_kv_pool is not None:
# Shared-KV-pool path: the caller built a UnifiedMHATokenToKVPool
# aliasing the shared byte buffer.
self.full_kv_pool = full_kv_pool
elif not use_mla:
TokenToKVPoolClass = MHATokenToKVPool
if current_platform.is_out_of_tree():
@@ -2492,11 +2524,16 @@ class HybridLinearKVPool(KVCache):
v_scale: float = 1.0,
dcp_kv_mask: Optional[torch.Tensor] = None,
):
# Write-location info lives in the metadata (`KVWriteLoc`). `full_loc` is the
# unified pool's pre-translated PHYSICAL loc (None for a static pool, where
# `loc` is already physical) — either way the pool writes a PHYSICAL loc.
loc, _, full_loc = unwrap_write_loc(loc)
layer_id = self._transfer_full_attention_id(layer.layer_id)
if not self.use_mla:
write_loc = full_loc if full_loc is not None else loc
self.full_kv_pool.set_kv_buffer(
None,
loc,
write_loc,
cache_k,
cache_v,
k_scale,
@@ -2518,8 +2555,9 @@ class HybridLinearKVPool(KVCache):
def get_cpu_copy(self, indices, mamba_indices=None):
kv_cpu = self.full_kv_pool.get_cpu_copy(indices)
# mamba_pool stores PHYSICAL ids; translate the (unified-pool virtual) ids first.
mamba_cpu = (
self.mamba_pool.get_cpu_copy(mamba_indices)
self.mamba_pool.get_cpu_copy(self._mamba_translate(mamba_indices))
if mamba_indices is not None
else None
)
@@ -2529,7 +2567,9 @@ class HybridLinearKVPool(KVCache):
kv_cpu, mamba_cpu = cache_cpu
self.full_kv_pool.load_cpu_copy(kv_cpu, indices)
if mamba_cpu is not None and mamba_indices is not None:
self.mamba_pool.load_cpu_copy(mamba_cpu, mamba_indices)
self.mamba_pool.load_cpu_copy(
mamba_cpu, self._mamba_translate(mamba_indices)
)
def get_v_head_dim(self):
return self.full_kv_pool.get_value_buffer(0).shape[-1]
@@ -2676,7 +2716,7 @@ class MLATokenToKVPool(KVCache):
cache_k: torch.Tensor,
cache_v: torch.Tensor,
):
loc, _ = unwrap_write_loc(loc_info)
loc, _, _ = unwrap_write_loc(loc_info)
maybe_detect_oob(loc, 0, self.size + self.page_size, "set_kv_buffer (MLA)")
layer_id = layer.layer_id
assert not self.dsa_kv_cache_store_fp8
@@ -2882,7 +2922,7 @@ class MLATokenToKVPoolFP4(MLATokenToKVPool):
cache_v: torch.Tensor,
):
# loc_info may be a KVWriteLoc; MLA pools have no SWA target.
loc, _ = unwrap_write_loc(loc_info)
loc, _, _ = unwrap_write_loc(loc_info)
maybe_detect_oob(loc, 0, self.size + self.page_size, "set_kv_buffer (MLA-FP4)")
layer_id = layer.layer_id
assert not self.dsa_kv_cache_store_fp8
python/sglang/srt/mem_cache/multi_ended_allocator.py
+2474
View File
File diff suppressed because it is too large Load Diff
python/sglang/srt/mem_cache/swa_memory_pool.py
+1 -1
View File
@@ -160,7 +160,7 @@ class SWAKVPool(BaseSWAKVPool):
v_scale: float = 1.0,
):
# loc_info bundles the full loc and the pre-translated SWA loc.
loc, swa_loc = unwrap_write_loc(loc_info)
loc, swa_loc, _ = unwrap_write_loc(loc_info)
layer_id = layer.layer_id
layer_id_pool, is_swa_layer = self.layers_mapping[layer_id]
if is_swa_layer:
python/sglang/srt/mem_cache/triton_ops/virtual_slot.py
+96
View File
@@ -0,0 +1,96 @@
# Copyright 2023-2026 SGLang Team
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Virtual<->physical slot Triton kernels for the unified memory pool."""
from __future__ import annotations
import torch
import triton
import triton.language as tl
# Fused take-physical-pages + bind for the alloc fast path. Invoked ONLY when
# `_hole_count == 0`; otherwise the slow path drains holes first (Invariant B,
# greedy hole reuse). Caller advances `watermark_physical` and checks overflow
# BEFORE launch, passing the PRE-extension watermark. Cuda-graph safe (no
# `.item()`, no tensor branching); runs on the scheduler thread.
@triton.jit
def alloc_bind_inplace_kernel(
v_pages_ptr, # in: [N] int64 — virtual page ids
v2p_ptr, # in/out: int64 — virtual_to_physical table
p2v_ptr, # in/out: int64 — physical_to_virtual table
out_phys_ptr, # out: [N] int64 — physical page ids
N, # runtime: number of pages to allocate
start_phys, # runtime: lowest physical page id in the new range
BLOCK: tl.constexpr,
):
"""Fused: ascending arange + out_phys/v2p/p2v scatter.
Caller pre-adjusts `start_phys` per direction so the range is always
ascending (grow-up: start_wm; grow-down: start_wm - N + 1), making the
v->p mapping byte-identical to the `torch.arange` slow path.
"""
pid = tl.program_id(0)
offs = pid * BLOCK + tl.arange(0, BLOCK)
mask = offs < N
phys = (start_phys + offs).to(tl.int64)
v = tl.load(v_pages_ptr + offs, mask=mask, other=0).to(tl.int64)
# Masked stores skip out-of-range lanes, and `other=0` keeps us off the
# v2p[0]/p2v[0] padding-sink slot.
tl.store(out_phys_ptr + offs, phys, mask=mask)
tl.store(v2p_ptr + v, phys, mask=mask)
tl.store(p2v_ptr + phys, v, mask=mask)
ALLOC_BIND_BLOCK = 128
def alloc_bind_inplace(
v_pages: torch.Tensor,
v2p: torch.Tensor,
p2v: torch.Tensor,
start_phys: int,
) -> torch.Tensor:
"""Allocate N ascending physical pages from `start_phys` and bind to `v_pages`.
Caller must advance `watermark_physical` by N and verify overflow BEFORE
calling; this launcher does neither.
"""
N = int(v_pages.numel())
if N == 0:
return torch.empty(0, dtype=torch.int64, device=v_pages.device)
if not v_pages.is_cuda:
# Pure-torch CPU reference for the CUDA-only kernel.
phys_pages = torch.arange(
start_phys, start_phys + N, dtype=torch.int64, device=v_pages.device
)
v = v_pages.to(torch.int64)
v2p[v] = phys_pages
p2v[phys_pages] = v
return phys_pages
phys_pages = torch.empty(N, dtype=torch.int64, device=v_pages.device)
grid = (triton.cdiv(N, ALLOC_BIND_BLOCK),)
alloc_bind_inplace_kernel[grid](
v_pages,
v2p,
p2v,
phys_pages,
N,
start_phys,
BLOCK=ALLOC_BIND_BLOCK,
)
return phys_pages
python/sglang/srt/mem_cache/unified_cache_components/mamba_component.py
+5 -1
View File
@@ -333,8 +333,12 @@ class MambaComponent(TreeComponent):
)
else:
mamba_value_donated = self._alloc_mamba_slot()
# mamba_pool is a pure PHYSICAL store; translate both slot ids
# virtual->physical (identity for the non-unified memory pool) first.
translate = self.cache.req_to_token_pool.translate_mamba_indices
self.cache.req_to_token_pool.mamba_pool.copy_from(
req.mamba_pool_idx.unsqueeze(0), mamba_value_donated
translate(req.mamba_pool_idx.unsqueeze(0)),
translate(mamba_value_donated),
)
insert_params.mamba_value = mamba_value_donated
return cache_len
python/sglang/srt/mem_cache/unified_memory_pool.py
+1369
View File
File diff suppressed because it is too large Load Diff
python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py
+183
View File
@@ -62,6 +62,14 @@ if TYPE_CHECKING:
from sglang.srt.model_executor.pool_configurator import MemoryPoolConfig
def _should_enable_lazy_compaction() -> bool:
"""Lazy compaction default — ON unless
`SGLANG_DISABLE_LAZY_COMPACTION=1` (escape hatch for A/B / rollback).
Centralized here so both unified-memory-pool factory call sites stay in sync.
"""
return not envs.SGLANG_DISABLE_LAZY_COMPACTION.get()
# the ratio of mamba cache pool size to max_running_requests
MAMBA_CACHE_SIZE_MAX_RUNNING_REQUESTS_RATIO = 3
MAMBA_CACHE_V2_ADDITIONAL_RATIO_OVERLAP = 2
@@ -341,10 +349,185 @@ class ModelRunnerKVCacheMixin:
"attention, no HiSparse, and --kv-cache-dtype != fp4_e2m1."
)
def _init_unified_mamba_pools(self: ModelRunner, max_num_reqs: int):
"""Build the shared-KV-pool stack for a hybrid-Mamba model:
one byte buffer split between the full-attn MHA KV pool and the
per-request Mamba state pool, with virtual slot ids above the
allocator."""
from sglang.srt.mem_cache.unified_memory_pool import init_unified_mamba_pools
config = self.mambaish_config
assert config is not None
assert (
not self.use_mla_backend
), "unified memory pool does not support MLA-hybrid-Mamba yet"
# The full sub-pool is page-aware (via `MultiEndedAllocator(page_size=...)`);
# the mamba sub-pool stays page=1.
assert self.page_size >= 1, f"page_size must be >= 1, got {self.page_size}"
# Mirror the non-shared path's extra_max_context_len computation.
extra_max_context_len = 4
if self.server_args.speculative_num_draft_tokens is not None:
extra_max_context_len += self.server_args.speculative_num_draft_tokens
mamba_layer_ids = [
i
for i in config.mamba2_cache_params.layers
if self.start_layer <= i < self.end_layer
]
full_attention_layer_ids = [
i
for i in config.full_attention_layer_ids
if self.start_layer <= i < self.end_layer
]
bundle = init_unified_mamba_pools(
device=self.device,
kv_cache_dtype=self.kv_cache_dtype,
head_num=self.model_config.get_num_kv_heads(get_attention_tp_size()),
head_dim=self.model_config.head_dim,
page_size=self.page_size,
start_layer=self.start_layer,
end_layer=self.end_layer,
is_draft_worker=self.is_draft_worker,
use_mla_backend=self.use_mla_backend,
mamba_layer_ids=mamba_layer_ids,
full_attention_layer_ids=full_attention_layer_ids,
mamba2_cache_params=config.mamba2_cache_params,
model_context_len=self.model_config.context_len,
extra_max_context_len=extra_max_context_len,
max_total_num_tokens=self.max_total_num_tokens,
max_mamba_cache_size=self.server_args.max_mamba_cache_size,
max_num_reqs=max_num_reqs,
enable_memory_saver=self.server_args.enable_memory_saver,
enable_mamba_extra_buffer=self.server_args.enable_mamba_extra_buffer(),
speculative_num_draft_tokens=self.server_args.speculative_num_draft_tokens,
disable_overlap_schedule=self.server_args.disable_overlap_schedule,
need_sort=self.server_args.disaggregation_mode in ("decode", "prefill"),
mamba_full_memory_ratio=self.server_args.mamba_full_memory_ratio,
# Overlap mode: the allocator's `free` drops a wait_stream(forward_stream)
# barrier so eager compaction serializes after the in-flight forward's
# v2p/KV reads. Near-no-op in normal mode.
forward_stream=self.forward_stream,
# Lazy compaction: default ON, env-var escape hatch for rollback / A/B.
lazy_compaction=_should_enable_lazy_compaction(),
)
self.req_to_token_pool = bundle.req_to_token_pool
self.token_to_kv_pool = bundle.token_to_kv_pool
self.token_to_kv_pool_allocator = bundle.token_to_kv_pool_allocator
# Keep a reference so the shared byte buffer is not GC'd.
self._unified_memory_pool = bundle.unified_memory_pool
def _init_unified_swa_pools(self: ModelRunner, max_num_reqs: int):
"""Build the unified-pool stack for a hybrid-SWA model (Triton): one byte
buffer split between the full-attention and SWA KV pools."""
from sglang.srt.mem_cache.unified_memory_pool import init_unified_swa_pools
assert self.is_hybrid_swa, "_init_unified_swa_pools called on a non-SWA model"
# Both sub-pools are page-aware; the SWA composite runs alloc_extend_kernel
# once in virtual space and binds the new pages on both sub-allocators.
assert self.page_size >= 1, f"page_size must be >= 1, got {self.page_size}"
assert (
not self.use_mla_backend
), "unified memory pool does not support MLA-SWA hybrid yet"
# Mirror the non-shared path's extra_max_context_len computation.
extra_max_context_len = 4
if self.server_args.speculative_num_draft_tokens is not None:
extra_max_context_len += self.server_args.speculative_num_draft_tokens
self.req_to_token_pool = ReqToTokenPool(
size=max_num_reqs,
max_context_len=self.model_config.context_len + extra_max_context_len,
device=self.device,
enable_memory_saver=self.server_args.enable_memory_saver,
)
head_num = self.model_config.get_num_kv_heads(get_attention_tp_size())
head_dim = self.model_config.head_dim
if self.is_hybrid_swa_compress:
# Asymmetric head dims between full and SWA (NPU compress path):
# pull SWA-specific dims from the hf text config.
v_head_dim = self.model_config.hf_text_config.v_head_dim
swa_head_num = max(
1,
self.model_config.hf_text_config.swa_num_key_value_heads
// get_attention_tp_size(),
)
swa_head_dim = self.model_config.hf_text_config.swa_head_dim
swa_v_head_dim = self.model_config.hf_text_config.swa_v_head_dim
else:
v_head_dim = head_dim
swa_head_num = head_num
swa_head_dim = head_dim
swa_v_head_dim = head_dim
# Filter layer ids to this worker's [start_layer, end_layer) range.
swa_attention_layer_ids = [
i
for i in self.model_config.swa_attention_layer_ids
if self.start_layer <= i < self.end_layer
]
full_attention_layer_ids = [
i
for i in self.model_config.full_attention_layer_ids
if self.start_layer <= i < self.end_layer
]
bundle = init_unified_swa_pools(
device=self.device,
kv_cache_dtype=self.kv_cache_dtype,
head_num=head_num,
head_dim=head_dim,
v_head_dim=v_head_dim,
swa_head_num=swa_head_num,
swa_head_dim=swa_head_dim,
swa_v_head_dim=swa_v_head_dim,
page_size=self.page_size,
start_layer=self.start_layer,
end_layer=self.end_layer,
swa_attention_layer_ids=swa_attention_layer_ids,
full_attention_layer_ids=full_attention_layer_ids,
full_max_total_num_tokens=self.full_max_total_num_tokens,
swa_max_total_num_tokens=self.swa_max_total_num_tokens,
enable_memory_saver=self.server_args.enable_memory_saver,
need_sort=self.server_args.disaggregation_mode in ("decode", "prefill"),
# Overlap mode: same wait_stream(forward_stream) rationale as
# `_init_unified_mamba_pools`.
forward_stream=self.forward_stream,
# Lazy compaction: default ON, with env var escape hatch for rollback / A/B.
lazy_compaction=_should_enable_lazy_compaction(),
)
self.token_to_kv_pool = bundle.token_to_kv_pool
self.token_to_kv_pool_allocator = bundle.token_to_kv_pool_allocator
# Keep a reference so the shared byte buffer is not GC'd.
self._unified_memory_pool = bundle.unified_memory_pool
def _init_pools(self: ModelRunner):
"""Initialize the memory pools."""
max_num_reqs = self.max_running_requests
# Unified-pool fast path: build req_to_token + token_to_kv pool + allocator
# from one byte buffer, then return. Gated to the target worker
# (req_to_token_pool is None); supports hybrid Mamba and hybrid SWA (not DSV4).
if (
self.server_args.enable_unified_memory
and self.server_args.disaggregation_mode == "null"
and self.req_to_token_pool is None
):
if self.mambaish_config is not None:
self._init_unified_mamba_pools(max_num_reqs)
return
if self.is_hybrid_swa and not is_deepseek_v4(self.model_config.hf_config):
self._init_unified_swa_pools(max_num_reqs)
return
# Fail loud, not silently fall through to the normal pools (which would
# leave the flag a no-op). The feature replaces the HYBRID pools only.
raise ValueError(
"--enable-unified-memory only supports hybrid Mamba and "
"hybrid sliding-window-attention models (DeepSeek-V4 excluded); "
f"the current model ({self.model_config.hf_config.architectures}) "
"is neither, so the unified memory pool cannot be built. Drop "
"--enable-unified-memory for this model."
)
# Initialize req_to_token_pool
if self.req_to_token_pool is None:
max_spec_draft_tokens = self.server_args.max_speculative_num_draft_tokens
python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py
+19 -2
View File
@@ -157,6 +157,11 @@ def build_replay_fb_view(
encoder_lens=buffers.encoder_lens[:bs] if is_encoder_decoder else None,
out_cache_loc=getattr(forward_batch, "out_cache_loc", None),
out_cache_loc_dsv4=getattr(forward_batch, "out_cache_loc_dsv4", None),
# The mamba-track registry slot (VIRTUAL ids) is the v2p translate SOURCE
# for the backend, which copies the result into its own static buffer and
# reads THAT in the decode track-save — this slot is never mutated. None
# when mamba-track is disabled.
mamba_track_indices=getattr(buffers, "mamba_track_indices", None),
spec_info=forward_batch.spec_info,
)
@@ -705,6 +710,9 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner):
if self.enable_profile_cuda_graph:
self._post_process_after_profile(prof)
# No pool-side pin to clear: the captured full-physical write loc rides the
# backend's `ForwardMetadata.out_cache_loc_full_physical` (-> KVWriteLoc.full_loc).
def _capture_one_stream(self, stream_idx: Optional[int] = None) -> None:
avail_mem = get_available_gpu_memory(
self.model_runner.device,
@@ -775,10 +783,15 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner):
attn_backend.init_forward_metadata_out_graph(forward_batch, in_capture=True)
def run_once():
# Must run inside the capture block: warmup mutations here are
# undone by on_after_cuda_graph_warmup so capture starts clean.
# Graph-recordable metadata-prep hook. The unified memory pool
# records ZERO translate nodes here: all its read/write translates
# run eagerly in `init_forward_metadata_out_graph` (replay-prep), so
# the captured graph reads already-physical locs. Base no-op for triton.
attn_backend.init_forward_metadata_in_graph(forward_batch)
# No invalidate_loc_cache() here: the unified pool translates its
# locs in `init_forward_metadata_out_graph`, so no cache to invalidate.
forward_batch.dp_local_start_pos = forward_batch.dp_local_num_tokens = (
None
)
@@ -834,6 +847,10 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner):
if (c := self.model_runner.canary_manager) is not None
else contextlib.nullcontext()
)
# Full-physical write loc lives in the attention metadata (the backend's
# `out_cache_loc_full_physical` -> KVWriteLoc.full_loc), so the runner
# wires no buffer here. (SWA write loc rides the `swa_out_cache_loc` rail.)
with canary_ctx:
shape_key = self._make_graph_key(bs, stream_idx, variant_label)
self.backend.capture_one(
python/sglang/srt/server_args.py
+52
View File
@@ -783,6 +783,14 @@ class ServerArgs:
"(layer-major) layout. Requires the Triton attention / linear-attn / "
"Mamba backends.",
] = False
enable_unified_memory: A[
bool,
"Replace the statically-partitioned hybrid-model pools (full-attn KV + "
"SWA/Mamba state) with one byte buffer split dynamically between "
"sub-pools. Requires the Triton attention / linear-attn / Mamba "
"backends; not yet compatible with PD disaggregation or speculative "
"decoding.",
] = False
disable_chunked_prefix_cache: A[
bool,
"Disable chunked prefix cache feature for deepseek, which should save overhead for short sequences.",
@@ -2717,6 +2725,8 @@ class ServerArgs:
self._handle_page_major_kv_layout()
self._handle_unified_memory_pool()
# Handle diffusion LLM inference.
self._handle_dllm_inference()
@@ -6336,7 +6346,49 @@ class ServerArgs:
"NCCL_ALGO is set to 'allreduce:tree' and custom all reduce is disabled for deterministic inference when TP size > 1."
)
def _handle_unified_memory_pool(self):
if not self.enable_unified_memory:
return
assert self.disaggregation_mode == "null", (
"--enable-unified-memory is not yet compatible with PD " "disaggregation."
)
assert self.speculative_algorithm is None, (
"--enable-unified-memory is not yet compatible with speculative "
"decoding."
)
assert not (self.enable_hierarchical_cache or self.enable_lmcache), (
"--enable-unified-memory is not yet compatible with hierarchical / "
"host-tiered KV cache (--enable-hierarchical-cache / --enable-lmcache): "
"the unified-memory-pool init wires up no host pools, and its device mamba / "
"full-attention slots are VIRTUAL — the host-offload path does not "
"translate them to physical."
)
assert self.dcp_size == 1, (
"--enable-unified-memory is not yet compatible with decode context "
"parallelism (--dcp-size > 1): the pool has no DCP-aware masked write "
"path (UnifiedMHATokenToKVPool.set_kv_buffer asserts dcp_kv_mask is None), "
"so a DCP run would boot and then fail on the first KV write."
)
# Only monolithic decode cuda-graph capture is wired; piecewise prefill
# capture is not. Guard when the user opts into it.
_cg_cfg = self.cuda_graph_config
if _cg_cfg is not None and _cg_cfg.prefill.backend == Backend.TC_PIECEWISE:
raise ValueError(
"--enable-unified-memory supports monolithic (decode) "
"cuda-graph capture only; disable piecewise prefill capture "
"(e.g. --cuda-graph-backend-prefill=disabled)."
)
# The strided-layout Triton requirement is enforced via
# --enable-page-major-kv-layout (implied by the unified pool in
# _handle_page_major_kv_layout); the model-family gate is enforced at pool
# construction in model_runner_kv_cache_mixin._init_pools.
def _handle_page_major_kv_layout(self):
# The unified pool stores state in the page-major envelope-strided layout, so
# enabling it implies --enable-page-major-kv-layout — routing it through the
# single page-major path + stride-aware Triton asserts (set before the guard).
if self.enable_unified_memory:
self.enable_page_major_kv_layout = True
if not self.enable_page_major_kv_layout:
return
# Only the Triton attention kernels read the strided 4-D envelope K/V