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

[mem_cache] Move mamba state and retraction_backup into ReqKvInfo (#37164)

Browse Source
This commit is contained in:
Liangsheng Yin
2026-08-30 22:03:01 -07:00
committed by GitHub
parent 9a9e167179
commit 5d12ad4fd7
30 changed files with 375 additions and 410 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/disaggregation/decode_kvcache_offload_manager.py
+1 -1
View File
@@ -285,7 +285,7 @@ class DecodeKVCacheOffloadManager:
self.token_to_kv_pool_allocator.free(overalloc_indices)
self.req_to_token_pool.free(req)
req.kv.mark_released()
req.kv.mark_kv_released()
self.tree_cache.protected_size_ -= len(req.prefix_indices)
self.offloaded_state.pop(req, None)
python/sglang/srt/disaggregation/prefill.py
+1 -1
View File
@@ -1036,7 +1036,7 @@ class SchedulerDisaggregationPrefillMixin:
else:
logger.warning(error_message)
req.time_stats.trace_ctx.abort(abort_info={"reason": error_message})
if req.kv.is_held or req.mamba_pool_idx is not None:
if req.kv.holds_kv or req.kv.holds_mamba:
release_kv_cache(req, self.tree_cache)
maybe_release_metadata_buffer(req, self.req_to_metadata_buffer_idx_allocator)
req.pending_bootstrap = False
python/sglang/srt/hardware_backend/mlx/kv_cache/auxiliary_state.py
+28 -28
View File
@@ -222,7 +222,7 @@ class MlxAuxiliaryStateReqToTokenPool(ReqToTokenPool):
* Radix cache enabled: the ``MlxAuxiliaryStateComponent`` of the unified
radix cache owns release — on finish it either frees the slot or
transfers it to the tree, nulling ``req.mamba_pool_idx`` before the
transfers it to the tree, nulling ``req.kv.mamba_pool_idx`` before the
request row is freed. The pool must NOT free auxiliary slots itself.
* Radix cache disabled (``ChunkCache``): no tree component exists, and
``release_kv_cache``'s ``free_mamba_cache`` fallback is gated on
@@ -270,13 +270,13 @@ class MlxAuxiliaryStateReqToTokenPool(ReqToTokenPool):
auxiliary_state_indices = []
for req in reqs:
if getattr(req, "mamba_pool_idx", None) is not None:
mid = req.mamba_pool_idx
if req.kv.holds_mamba:
mid = req.kv.mamba_pool_idx
else:
allocated = self.auxiliary_state_pool.alloc(1)
assert allocated is not None, "Not enough MLX auxiliary state slots"
mid = allocated[0]
req.mamba_pool_idx = mid
req.kv.mamba_pool_idx = mid
auxiliary_state_indices.append(mid.to(dtype=torch.int32))
self.req_index_to_auxiliary_state_index_mapping[select_index] = torch.stack(
auxiliary_state_indices
@@ -293,16 +293,16 @@ class MlxAuxiliaryStateReqToTokenPool(ReqToTokenPool):
return 0
def free_mamba_cache(self, req, mamba_ping_pong_track_buffer_to_keep=None):
if getattr(req, "mamba_pool_idx", None) is not None:
self.auxiliary_state_pool.free(req.mamba_pool_idx.unsqueeze(0))
req.mamba_pool_idx = None
track_buffer = getattr(req, "mamba_ping_pong_track_buffer", None)
if req.kv.holds_mamba:
self.auxiliary_state_pool.free(req.kv.mamba_pool_idx.unsqueeze(0))
req.kv.mamba_pool_idx = None
track_buffer = req.kv.mamba_ping_pong_track_buffer
if track_buffer is not None:
if mamba_ping_pong_track_buffer_to_keep is None:
self.auxiliary_state_pool.free(track_buffer)
req.mamba_ping_pong_track_buffer = None
req.mamba_next_track_idx = None
req.mamba_last_track_idx = None
req.kv.mamba_ping_pong_track_buffer = None
req.kv.mamba_next_track_idx = None
req.kv.mamba_last_track_idx = None
def free_auxiliary_state_cache(self, req, track_buffer_to_keep=None):
self.free_mamba_cache(
@@ -314,7 +314,7 @@ class MlxAuxiliaryStateReqToTokenPool(ReqToTokenPool):
if self._owns_auxiliary_state_release:
# No-radix configuration: nothing else will ever release the
# auxiliary slot, so return it with the request row. Keyed on
# req.mamba_pool_idx (None-safe, nulled by free_mamba_cache), NOT
# req.kv.mamba_pool_idx (None-safe, nulled by free_mamba_cache), NOT
# on req_index_to_auxiliary_state_index_mapping, which may point
# at a slot the radix tree owns.
self.free_mamba_cache(req)
@@ -347,13 +347,13 @@ class MlxAuxiliaryStateComponent(MambaComponent):
@staticmethod
def _tracked_value(req) -> tuple[object | None, bool]:
track_buffer = getattr(req, "mamba_ping_pong_track_buffer", None)
track_len = getattr(req, "mamba_last_track_seqlen", None)
track_buffer = req.kv.mamba_ping_pong_track_buffer
track_len = req.kv.mamba_last_track_seqlen
if track_buffer is not None and track_len is not None:
return track_buffer[0].unsqueeze(-1).clone(), True
if getattr(req, "mamba_pool_idx", None) is None:
if not req.kv.holds_mamba:
return None, False
return req.mamba_pool_idx.unsqueeze(-1).clone(), False
return req.kv.mamba_pool_idx.unsqueeze(-1).clone(), False
def prepare_for_caching_req(
self,
@@ -362,7 +362,7 @@ class MlxAuxiliaryStateComponent(MambaComponent):
token_ids_len: int,
is_finished: bool,
) -> int | None:
cache_len = getattr(req, "mamba_last_track_seqlen", None)
cache_len = req.kv.mamba_last_track_seqlen
auxiliary_value, uses_track_slot = self._tracked_value(req)
setattr(insert_params, "mlx_auxiliary_state_uses_track_slot", uses_track_slot)
@@ -408,13 +408,13 @@ class MlxAuxiliaryStateComponent(MambaComponent):
if bool(
getattr(insert_params, "mlx_auxiliary_state_uses_track_slot", False)
):
track_buffer = getattr(req, "mamba_ping_pong_track_buffer", None)
track_buffer = req.kv.mamba_ping_pong_track_buffer
if track_buffer is not None:
self.cache.req_to_token_pool.auxiliary_state_pool.free(track_buffer)
req.mamba_ping_pong_track_buffer = None
req.mamba_next_track_idx = None
req.mamba_last_track_idx = None
req.mamba_last_track_seqlen = None
req.kv.mamba_ping_pong_track_buffer = None
req.kv.mamba_next_track_idx = None
req.kv.mamba_last_track_idx = None
req.kv.mamba_last_track_seqlen = None
return
auxiliary_value_exists = (
@@ -433,11 +433,11 @@ class MlxAuxiliaryStateComponent(MambaComponent):
self.cache.req_to_token_pool.free_auxiliary_state_cache(req)
else:
# The radix tree now owns the live auxiliary-state slot.
track_buffer = getattr(req, "mamba_ping_pong_track_buffer", None)
track_buffer = req.kv.mamba_ping_pong_track_buffer
if track_buffer is not None:
self.cache.req_to_token_pool.auxiliary_state_pool.free(track_buffer)
req.mamba_ping_pong_track_buffer = None
req.mamba_next_track_idx = None
req.mamba_last_track_idx = None
req.mamba_pool_idx = None
req.mamba_last_track_seqlen = None
req.kv.mamba_ping_pong_track_buffer = None
req.kv.mamba_next_track_idx = None
req.kv.mamba_last_track_idx = None
req.kv.mamba_pool_idx = None
req.kv.mamba_last_track_seqlen = None
python/sglang/srt/hardware_backend/mlx/model_runner.py
+7 -7
View File
@@ -379,7 +379,7 @@ class MlxModelRunner:
chunk_size = mamba_cache_chunk_size()
track_len = prefix_len + (new_token_count // chunk_size) * chunk_size
branching_len = getattr(req, "mamba_branching_seqlen", None)
branching_len = req.mamba_branching_seqlen
if (
branching_len is not None
and prefix_len < branching_len <= prefix_len + new_token_count
@@ -407,7 +407,7 @@ class MlxModelRunner:
if pool is None or not hasattr(pool, "store_cache"):
return
track_buffer = getattr(req, "mamba_ping_pong_track_buffer", None)
track_buffer = req.kv.mamba_ping_pong_track_buffer
if track_buffer is None:
track_buffer = pool.alloc(1)
if track_buffer is None:
@@ -416,16 +416,16 @@ class MlxModelRunner:
"falling back to leaf-only auxiliary-state radix caching."
)
return
req.mamba_ping_pong_track_buffer = track_buffer
req.mamba_next_track_idx = 0
req.mamba_last_track_idx = 0
req.kv.mamba_ping_pong_track_buffer = track_buffer
req.kv.mamba_next_track_idx = 0
req.kv.mamba_last_track_idx = 0
pool.store_cache(
track_buffer[0],
cache,
self._cache_layout.auxiliary_layer_indices,
)
req.mamba_last_track_seqlen = track_len
req.kv.mamba_last_track_seqlen = track_len
def _cache_with_pool_backed_attention(
self, prefix_slot_ids: list[int], prefix_len: int
@@ -905,7 +905,7 @@ class MlxModelRunner:
"""
prefix_len = len(prefix_slot_ids)
if req is not None:
req.mamba_last_track_seqlen = None
req.kv.mamba_last_track_seqlen = None
if self._enable_sampling:
self._req_sampling[req_id] = (
MlxSamplingParams.from_req(
python/sglang/srt/hardware_backend/mlx/tp_worker.py
+1 -1
View File
@@ -172,7 +172,7 @@ class MlxTpModelWorker(TpModelWorker):
self._mlx_runner.store_auxiliary_state_for_request(req.rid)
# Prefer the just-snapshotted live auxiliary state for the final
# insert. Any older tracked slot is released during component cleanup.
req.mamba_last_track_seqlen = None
req.kv.mamba_last_track_seqlen = None
def _route_extend_request(self, rid: str, decoding_rids: set[str]) -> str:
"""Classify a request within an extend / mixed batch.
python/sglang/srt/managers/schedule_batch.py
+79 -65
View File
@@ -816,11 +816,12 @@ class ReqLogprob:
@dataclasses.dataclass(slots=True, kw_only=True)
class ReqKvInfo:
# Device KV a request holds outside the prefix cache. Always present on the Req;
# whether any KV is held is `is_held` (a row is registered).
# whether any KV is held is `holds_kv` (a row is registered).
# Match observations and scheduling state stay on the Req itself.
req_pool_idx: Optional[int] = None # req_to_token row, the register for the slots
# The request's own KV is [cache_protected_len, kv_allocated_len).
cache_protected_len: int = 0 # tree cache owns [0, here) (matched or inserted)
cache_protected_len: int = 0 # Tree cache owns [0, here) (matched or inserted)
kv_committed_len: int = 0 # KV content committed up to here, <= kv_allocated_len
kv_allocated_len: int = 0
@@ -828,6 +829,21 @@ class ReqKvInfo:
swa_evict_floor: int = 0 # [0, here) never window-evicted (prefill-aware SWA)
swa_evicted_seqlen: int = 0 # SWA eviction cursor
# Host-side KV backup the request holds across a retraction (unified cache).
retraction_backup: Optional[RetractionBackup] = None
# Mamba state: an independent resource; whether it is held is `holds_mamba`.
mamba_pool_idx: Optional[torch.Tensor] = None # shape (1)
mamba_ping_pong_track_buffer: Optional[torch.Tensor] = None # shape (2)
mamba_next_track_idx: Optional[int] = None # 0 or 1
mamba_last_track_idx: Optional[int] = None # 0 or 1
# Seq len of the last cached mamba state
mamba_last_track_seqlen: Optional[int] = None
# Deferred COW: source mamba pool index from radix cache node (copy on forward stream)
mamba_cow_src_index: Optional[torch.Tensor] = None
# Deferred clear: newly allocated mamba slot needs zeroing on forward stream
mamba_needs_clear: bool = False
def swa_dead_lo(self, page_size: int) -> int:
# Lowest SWA position this request may free itself: above the tree-owned
# prefix and above the eviction shield, page-aligned upward.
@@ -837,14 +853,18 @@ class ReqKvInfo:
return lo
@property
def is_held(self) -> bool:
def holds_kv(self) -> bool:
return self.req_pool_idx is not None
@property
def is_released(self) -> bool:
def holds_mamba(self) -> bool:
return self.mamba_pool_idx is not None
@property
def is_kv_released(self) -> bool:
return self.kv_allocated_len == 0 and self.swa_evicted_seqlen == 0
def mark_released(self) -> None:
def mark_kv_released(self) -> None:
self.kv_allocated_len = 0
self.swa_evicted_seqlen = 0
@@ -929,7 +949,6 @@ class Req(ReqDllmMixin):
# For req-level memory management
self.kv = ReqKvInfo()
self.retraction_backup: Optional[RetractionBackup] = None
# for cross-encoder model
self.token_type_ids = token_type_ids
@@ -974,21 +993,6 @@ class Req(ReqDllmMixin):
self.lora_id = lora_id
self.routing_key = routing_key
# Memory pool info
self.mamba_pool_idx: Optional[torch.Tensor] = None # shape (1)
self.mamba_ping_pong_track_buffer: Optional[torch.Tensor] = None # shape (2)
self.mamba_next_track_idx: Optional[int] = None # 0 or 1
self.mamba_last_track_idx: Optional[int] = None # 0 or 1
self.mamba_last_track_seqlen: Optional[int] = (
None # seq len of the last cached mamba state
)
# the branching point seqlen to track mamba state. If set, given by prefix match,
# it will be the tracked seqlen in the ping pong buffer for the right prefill pass.
self.mamba_branching_seqlen: Optional[int] = None
# Deferred COW: source mamba pool index from radix cache node (copy on forward stream)
self.mamba_cow_src_index: Optional[torch.Tensor] = None
# Deferred clear: newly allocated mamba slot needs zeroing on forward stream
self.mamba_needs_clear: bool = False
# Lazy extra buffer: skip radix cache insert when prealloc failed at
# boundary — the forward overwrites the only slot, corrupting the state.
self.mamba_lazy_is_insert: bool = True
@@ -1041,6 +1045,10 @@ class Req(ReqDllmMixin):
self.host_hit_length = 0
self.swa_host_hit_length = 0
self.mamba_host_hit_length = 0
# The branching point seqlen to track mamba state. If set, given by prefix
# match, it will be the tracked seqlen in the ping pong buffer for the
# right prefill pass.
self.mamba_branching_seqlen: Optional[int] = None
# Total cached prefix length (on-device prefix_indices + host_hit_length),
# capped at the max allowed prefix. Set during prefix matching at schedule
# time and used to estimate uncached tokens / sort by longest prefix for
@@ -1744,16 +1752,16 @@ class Req(ReqDllmMixin):
self.temp_input_token_ids_logprobs_val = None
self.temp_input_token_ids_logprobs_idx = None
self.inflight_middle_chunks = 0
self.mamba_pool_idx = None
self.mamba_ping_pong_track_buffer = None
self.mamba_next_track_idx = None
self.mamba_last_track_idx = None
self.mamba_last_track_seqlen = None
self.kv.mamba_pool_idx = None
self.kv.mamba_ping_pong_track_buffer = None
self.kv.mamba_next_track_idx = None
self.kv.mamba_last_track_idx = None
self.kv.mamba_last_track_seqlen = None
self.mamba_branching_seqlen = None
self.mamba_cow_src_index = None
self.mamba_needs_clear = False
self.kv.mamba_cow_src_index = None
self.kv.mamba_needs_clear = False
self.already_computed = 0
assert not self.kv.is_held, "expect it is already released"
assert not self.kv.holds_kv, "expect it is already released"
self.kv.kv_committed_len = 0
self.extend_batch_idx = 0
self.decode_batch_idx = 0
@@ -1784,34 +1792,34 @@ class Req(ReqDllmMixin):
mamba_pool = self._mamba_pool_needing_backup(
req_to_token_pool, token_to_kv_pool_allocator
)
self.retraction_backup = RetractionBackup(
self.kv.retraction_backup = RetractionBackup(
cpu_tensors=token_to_kv_pool_allocator.get_cpu_copy(
token_indices, mamba_indices=self.mamba_pool_idx
token_indices, mamba_indices=self.kv.mamba_pool_idx
),
mamba_cpu=(
mamba_pool.get_cpu_copy(self.mamba_pool_idx.unsqueeze(0))
if mamba_pool is not None and self.mamba_pool_idx is not None
mamba_pool.get_cpu_copy(self.kv.mamba_pool_idx.unsqueeze(0))
if mamba_pool is not None and self.kv.holds_mamba
else None
),
)
def load_kv_cache(self, req_to_token_pool, token_to_kv_pool_allocator):
assert self.retraction_backup is not None
assert self.kv.retraction_backup is not None
token_indices = req_to_token_pool.req_to_token[
self.kv.req_pool_idx, : self.seqlen - 1
]
# Loads both the kv cache and mamba state if exists
mamba_cpu = self.retraction_backup.mamba_cpu
if mamba_cpu is not None and self.mamba_pool_idx is not None:
mamba_cpu = self.kv.retraction_backup.mamba_cpu
if mamba_cpu is not None and self.kv.holds_mamba:
req_to_token_pool.mamba_pool.load_cpu_copy(
mamba_cpu, self.mamba_pool_idx.unsqueeze(0)
mamba_cpu, self.kv.mamba_pool_idx.unsqueeze(0)
)
token_to_kv_pool_allocator.load_cpu_copy(
self.retraction_backup.cpu_tensors,
self.kv.retraction_backup.cpu_tensors,
token_indices,
mamba_indices=self.mamba_pool_idx,
mamba_indices=self.kv.mamba_pool_idx,
)
self.retraction_backup = None
self.kv.retraction_backup = None
def build_rebootstrap_payload(self) -> dict:
"""Build the prefill ``/generate`` payload that asks the original prefill
@@ -1964,7 +1972,11 @@ def set_mamba_track_indices_from_reqs(
# gone through _alloc_ping_pong_buffer yet (e.g., spec v2 verify path).
# Default to 0 (first ping-pong slot) to avoid TypeError.
track_positions = [
req.mamba_next_track_idx if req.mamba_next_track_idx is not None else 0
(
req.kv.mamba_next_track_idx
if req.kv.mamba_next_track_idx is not None
else 0
)
for req in batch.reqs
]
batch.mamba_track_buffer_indices = list(track_positions)
@@ -2171,8 +2183,8 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
# For hybrid GDN prefix cache
mamba_track_indices: torch.Tensor = None # shape: [b], int64
# Per-batch snapshot of the logical ping-pong positions selected for this
# forward (normally req.mamba_next_track_idx; spec may override it). Result
# processing uses it to update req.mamba_last_track_idx, since both req-level
# forward (normally req.kv.mamba_next_track_idx; spec may override it). Result
# processing uses it to update req.kv.mamba_last_track_idx, since both req-level
# indices may advance under overlap.
mamba_track_buffer_indices: Optional[List[int]] = None # shape: [b], 0 or 1
mamba_track_mask: torch.Tensor = None # shape: [b], bool
@@ -2694,7 +2706,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
)
# Collect mamba init info for deferred ops on forward stream
if any(req.mamba_pool_idx is not None for req in reqs):
if any(req.kv.holds_mamba for req in reqs):
self._collect_deferred_mamba_cow_and_clear(reqs)
if self.model_config.is_encoder_decoder:
@@ -2736,7 +2748,9 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
return i + 1
mask = req.extend_range.length >= checkpoint_grid
track_index = req.mamba_ping_pong_track_buffer[req.mamba_next_track_idx].item()
track_index = req.kv.mamba_ping_pong_track_buffer[
req.kv.mamba_next_track_idx
].item()
mamba_track_seqlen = -1
if mask:
# mamba_track_seqlen is used to calculate the indices to track in
@@ -2769,11 +2783,11 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
# In lazy mode, skip the swap — the second ping-pong slot is not
# allocated yet; it will be allocated on demand at the track boundary
# in mamba_lazy_prealloc_at_boundary during prepare_for_decode.
req.mamba_last_track_idx = req.mamba_next_track_idx
req.kv.mamba_last_track_idx = req.kv.mamba_next_track_idx
if not mamba_extra_buffer_lazy_enabled():
req.mamba_next_track_idx = (
req.kv.mamba_next_track_idx = (
self.req_to_token_pool.get_mamba_ping_pong_other_idx(
req.mamba_next_track_idx
req.kv.mamba_next_track_idx
)
)
if req.mamba_branching_seqlen is not None:
@@ -2792,7 +2806,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
# See _force_track_h() for more details.
mamba_track_seqlen = _force_track_h(req.mamba_branching_seqlen)
mamba_track_seqlen_aligned = req.mamba_branching_seqlen
req.mamba_last_track_seqlen = mamba_track_seqlen_aligned
req.kv.mamba_last_track_seqlen = mamba_track_seqlen_aligned
return _MambaRadixCacheV2TrackEntry(
track_mask=mask,
@@ -2806,14 +2820,14 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
cow_dst_tensors = []
clear_tensors = []
for req in reqs:
if req.mamba_cow_src_index is not None:
cow_src_tensors.append(req.mamba_cow_src_index)
cow_dst_tensors.append(req.mamba_pool_idx.unsqueeze(0))
req.mamba_cow_src_index = None
req.mamba_needs_clear = False
elif req.mamba_needs_clear:
clear_tensors.append(req.mamba_pool_idx.unsqueeze(0))
req.mamba_needs_clear = False
if req.kv.mamba_cow_src_index is not None:
cow_src_tensors.append(req.kv.mamba_cow_src_index)
cow_dst_tensors.append(req.kv.mamba_pool_idx.unsqueeze(0))
req.kv.mamba_cow_src_index = None
req.kv.mamba_needs_clear = False
elif req.kv.mamba_needs_clear:
clear_tensors.append(req.kv.mamba_pool_idx.unsqueeze(0))
req.kv.mamba_needs_clear = False
self.mamba_cow_src_indices = (
torch.cat(cow_src_tensors) if cow_src_tensors else None
)
@@ -3100,12 +3114,12 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
"""
pool = self.req_to_token_pool
for i, req in enumerate(self.reqs):
buf = req.mamba_ping_pong_track_buffer
buf = req.kv.mamba_ping_pong_track_buffer
assert buf is not None
# Skip reqs not at a track boundary
if self.seq_lens_cpu[i].item() % mamba_track_interval != 0:
continue
other_idx = 1 - req.mamba_next_track_idx
other_idx = 1 - req.kv.mamba_next_track_idx
if buf[other_idx].item() != -1:
# With overlap the previous forward's post-processing
# (which frees this slot) hasn't run yet. Skip.
@@ -3118,7 +3132,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
new_slot = pool.mamba_allocator.alloc(1)
if new_slot is not None:
pool.set_mamba_ping_pong_slot(req, other_idx, new_slot[0])
req.mamba_next_track_idx = other_idx
req.kv.mamba_next_track_idx = other_idx
def mamba_lazy_spec_prepare(self, mamba_track_interval: int, max_draft_tokens: int):
"""Lazy-mode spec counterpart of mamba_lazy_prealloc_at_boundary.
@@ -3134,16 +3148,16 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
pool = self.req_to_token_pool
track_positions: List[int] = []
for req in self.reqs:
buf = req.mamba_ping_pong_track_buffer
buf = req.kv.mamba_ping_pong_track_buffer
assert buf is not None
if not mamba_lazy_spec_in_window(
req, mamba_track_interval, max_draft_tokens
):
# No crossing reachable: the scatter mask stays -1, the
# position is never written.
track_positions.append(req.mamba_next_track_idx)
track_positions.append(req.kv.mamba_next_track_idx)
continue
other_idx = 1 - req.mamba_next_track_idx
other_idx = 1 - req.kv.mamba_next_track_idx
has_pending = buf[other_idx].item() != -1
if not has_pending:
if envs.SGLANG_TEST_MAMBA_LAZY_ALLOC_FAIL.get():
@@ -3157,7 +3171,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
has_pending = True
# On failure the verify scatters in place into the keep slot.
track_positions.append(
other_idx if has_pending else req.mamba_next_track_idx
other_idx if has_pending else req.kv.mamba_next_track_idx
)
self.mamba_lazy_spec_track_positions_cpu = track_positions
@@ -3497,7 +3511,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
# seqlen progress is monotonic per KV handle.
if (
req.decode_batch_idx >= 1
and req.kv.is_held
and req.kv.holds_kv
and req.seqlen - 1 - sliding_window_size
>= req.kv.swa_evicted_seqlen + eviction_interval
):
python/sglang/srt/managers/schedule_policy.py
+1 -1
View File
@@ -833,7 +833,7 @@ class PrefillAdder:
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:
if self._mamba_slot_cost and not req.kv.holds_mamba:
return self._mamba_slot_cost
return 0
python/sglang/srt/managers/scheduler.py
+7 -12
View File
@@ -3587,15 +3587,13 @@ class Scheduler(
if not added:
# init_next_round_input() may stage deferred Mamba COW/clear
# metadata before add_one_req() rejects the request.
req.mamba_cow_src_index = None
req.mamba_needs_clear = False
if req.mamba_pool_idx is not None and not getattr(
req, "session", None
):
req.kv.mamba_cow_src_index = None
req.kv.mamba_needs_clear = False
if req.kv.holds_mamba and not getattr(req, "session", None):
self.tree_cache.req_to_token_pool.mamba_allocator.free(
req.mamba_pool_idx.unsqueeze(-1)
req.kv.mamba_pool_idx.unsqueeze(-1)
)
req.mamba_pool_idx = None
req.kv.mamba_pool_idx = None
break
if mamba_allocator is not None:
@@ -4852,7 +4850,7 @@ class Scheduler(
# For mamba radix cache
if (
req.mamba_pool_idx is not None
req.kv.holds_mamba
and self.disaggregation_mode != DisaggregationMode.DECODE
):
release_kv_cache(req, self.tree_cache, is_insert=False)
@@ -4867,10 +4865,7 @@ class Scheduler(
self.ipc_channels.send_to_tokenizer.send_output(
_make_abort_req(req), req
)
if (
req.kv.req_pool_idx is not None
or getattr(req, "mamba_pool_idx", None) is not None
):
if req.kv.holds_kv or req.kv.holds_mamba:
release_kv_cache(req, self.tree_cache, is_insert=False)
logger.debug(f"Abort dLLM queued request. {req.rid=}")
python/sglang/srt/managers/scheduler_components/batch_result_processor.py
+21 -21
View File
@@ -1125,14 +1125,14 @@ class SchedulerBatchResultProcessor:
known_mamba_boundary = bool(batch.mamba_track_mask_next_cpu[i])
if completed_mamba_boundary and not lazy:
req.mamba_last_track_idx = batch.mamba_track_buffer_indices[i]
req.mamba_last_track_seqlen = req.kv.kv_committed_len - lookahead
req.kv.mamba_last_track_idx = batch.mamba_track_buffer_indices[i]
req.kv.mamba_last_track_seqlen = req.kv.kv_committed_len - lookahead
elif (
req.finished()
and lazy
and lookahead == 1
and known_mamba_boundary
and req.mamba_next_track_idx == req.mamba_last_track_idx
and req.kv.mamba_next_track_idx == req.kv.mamba_last_track_idx
):
req.mamba_lazy_is_insert = False
@@ -1216,7 +1216,7 @@ class SchedulerBatchResultProcessor:
Lazy: keep the same index (prealloc handles the swap) and run
post-decode cleanup to free the temporary second slot.
"""
if req.mamba_ping_pong_track_buffer is None:
if req.kv.mamba_ping_pong_track_buffer is None:
return
lazy = mamba_extra_buffer_lazy_enabled()
@@ -1238,17 +1238,17 @@ class SchedulerBatchResultProcessor:
if not at_boundary:
return
track_idx = req.mamba_next_track_idx
track_idx = req.kv.mamba_next_track_idx
if not known_boundary and batch.mamba_track_buffer_indices is not None:
track_idx = batch.mamba_track_buffer_indices[i]
if not known_boundary:
req.mamba_last_track_seqlen = track_seqlen
req.kv.mamba_last_track_seqlen = track_seqlen
if lazy:
self.mamba_lazy_post_decode_at_boundary(req, batch, track_idx)
else:
if not known_boundary:
req.mamba_last_track_idx = track_idx
req.mamba_next_track_idx = (
req.kv.mamba_last_track_idx = track_idx
req.kv.mamba_next_track_idx = (
batch.req_to_token_pool.get_mamba_ping_pong_other_idx(track_idx)
)
@@ -1274,12 +1274,12 @@ class SchedulerBatchResultProcessor:
if req.finished():
# Skip the donation if a scatter wrote or may still write the keep slot.
keep_written_by_this_step = (
crossed and planned_pos == req.mamba_next_track_idx
crossed and planned_pos == req.kv.mamba_next_track_idx
)
other_idx = 1 - req.mamba_next_track_idx
other_idx = 1 - req.kv.mamba_next_track_idx
# Recompute the in-flight verify's plan (kv_committed_len is
# frozen since its prepare, so the recompute is exact).
keep_may_be_written_in_flight = req.mamba_ping_pong_track_buffer[
keep_may_be_written_in_flight = req.kv.mamba_ping_pong_track_buffer[
other_idx
].item() == -1 and mamba_lazy_spec_in_window(
req,
@@ -1296,18 +1296,18 @@ class SchedulerBatchResultProcessor:
if not crossed or planned_pos is None:
return
if planned_pos != req.mamba_next_track_idx:
if planned_pos != req.kv.mamba_next_track_idx:
# Promote pending -> keep: free the old checkpoint, repoint.
pool = batch.req_to_token_pool
keep_idx = req.mamba_next_track_idx
keep_val = req.mamba_ping_pong_track_buffer[keep_idx]
keep_idx = req.kv.mamba_next_track_idx
keep_val = req.kv.mamba_ping_pong_track_buffer[keep_idx]
pool.mamba_allocator.free(keep_val.unsqueeze(0))
pool.set_mamba_ping_pong_slot(req, keep_idx, -1)
req.mamba_next_track_idx = planned_pos
req.kv.mamba_next_track_idx = planned_pos
# else: in-place fallback, or promoted by an earlier confirmation —
# keep holds the track_seqlen state either way.
req.mamba_last_track_idx = planned_pos
req.mamba_last_track_seqlen = track_seqlen
req.kv.mamba_last_track_idx = planned_pos
req.kv.mamba_last_track_seqlen = track_seqlen
@staticmethod
def _mamba_assert_committed_len_lookahead(req: Req) -> None:
@@ -1355,13 +1355,13 @@ class SchedulerBatchResultProcessor:
self, req: Req, batch: ScheduleBatch, track_idx: int
):
"""Commit a completed lazy-mode boundary and free its old slot."""
req.mamba_last_track_idx = track_idx
req.mamba_next_track_idx = track_idx
req.kv.mamba_last_track_idx = track_idx
req.kv.mamba_next_track_idx = track_idx
other_idx = 1 - track_idx
other_val = req.mamba_ping_pong_track_buffer[other_idx].item()
other_val = req.kv.mamba_ping_pong_track_buffer[other_idx].item()
if other_val != -1:
pool = batch.req_to_token_pool
pool.mamba_allocator.free(
req.mamba_ping_pong_track_buffer[other_idx].unsqueeze(0)
req.kv.mamba_ping_pong_track_buffer[other_idx].unsqueeze(0)
)
pool.set_mamba_ping_pong_slot(req, other_idx, -1)
python/sglang/srt/managers/scheduler_components/invariant_checker.py
+3 -3
View File
@@ -252,7 +252,7 @@ class SchedulerInvariantChecker:
swa_uncached = 0
for batch in batches:
for req in batch.reqs:
if not req.kv.is_held:
if not req.kv.holds_kv:
continue
allocated_len = req.kv.kv_allocated_len
@@ -324,7 +324,7 @@ class SchedulerInvariantChecker:
batch = self.get_last_batch()
if batch is not None:
for req in batch.reqs:
if not req.kv.is_held:
if not req.kv.holds_kv:
continue
_add_owner(
req,
@@ -336,7 +336,7 @@ class SchedulerInvariantChecker:
sess = getattr(self.tree_cache, "slots", None)
if sess:
for sid, slot in sess.items():
if slot.kv.is_held:
if slot.kv.holds_kv:
_add_owner(
slot,
f"slot {sid[:8]}",
python/sglang/srt/managers/scheduler_pp_mixin.py
+3 -3
View File
@@ -723,15 +723,15 @@ class SchedulerPPMixin:
latencies.append(latency_ms)
# Release KV and Mamba cache
if req.kv.is_held:
if req.kv.holds_kv:
kv_indices = self.req_to_token_pool.req_to_token[
req.kv.req_pool_idx, : req.extend_range.end
]
self.token_to_kv_pool_allocator.free(kv_indices)
if req.mamba_pool_idx is not None:
if req.kv.holds_mamba:
self.req_to_token_pool.free_mamba_cache(req)
self.req_to_token_pool.free(req)
req.kv.mark_released()
req.kv.mark_kv_released()
logger.info(
f"[PP Dynamic Chunk] [PP0] Profiled {len(seq_lens)} samples: "
python/sglang/srt/mem_cache/common.py
+19 -19
View File
@@ -62,7 +62,7 @@ def free_swa_out_of_window_slots(
is_chunk_cache: bool = False,
retain_floor: int | None = None,
) -> None:
if not req.kv.is_held:
if not req.kv.holds_kv:
return
# For swa radix cache, we need to evict the tokens that are not in the tree cache and also not in the sliding window
@@ -159,8 +159,8 @@ def retraction_backup(
return True
unified_cache = cast("UnifiedRadixCache", tree_cache)
req.retraction_backup = unified_cache.retraction_backup(req)
return req.retraction_backup is not None
req.kv.retraction_backup = unified_cache.retraction_backup(req)
return req.kv.retraction_backup is not None
def retraction_restore(
@@ -179,38 +179,38 @@ def retraction_restore(
return
unified_cache = cast("UnifiedRadixCache", tree_cache)
assert req.retraction_backup is not None
unified_cache.retraction_restore(req, req.retraction_backup)
req.retraction_backup = None
assert req.kv.retraction_backup is not None
unified_cache.retraction_restore(req, req.kv.retraction_backup)
req.kv.retraction_backup = None
def retraction_discard(req: Req, tree_cache: BasePrefixCache, backend: str) -> None:
if backend == "cpu_tensor":
req.retraction_backup = None
req.kv.retraction_backup = None
return
if backend != "host_pool":
raise ValueError(f"Unknown retraction backup backend: {backend}")
if req.retraction_backup is None:
if req.kv.retraction_backup is None:
return
unified_cache = cast("UnifiedRadixCache", tree_cache)
unified_cache.retraction_discard(req.retraction_backup)
req.retraction_backup = None
unified_cache.retraction_discard(req.kv.retraction_backup)
req.kv.retraction_backup = None
def release_kv_cache(req: Req, tree_cache: BasePrefixCache, is_insert: bool = True):
assert (not req.kv.is_held) == req.kv.is_released
assert (not req.kv.holds_kv) == req.kv.is_kv_released
# MambaRadixCache may alloc mamba state before alloc KV cache
if not req.kv.is_held:
if not req.kv.holds_kv:
assert (
tree_cache.supports_mamba()
), "Only MambaRadixCache allow freeing before alloc"
# TODO (csy, hanming): clean up this early allocation logic
if req.mamba_pool_idx is not None:
if req.kv.holds_mamba:
tree_cache.req_to_token_pool.mamba_allocator.free(
req.mamba_pool_idx.unsqueeze(-1)
req.kv.mamba_pool_idx.unsqueeze(-1)
)
req.mamba_pool_idx = None
req.kv.mamba_pool_idx = None
return
effective_kv_committed_len = req.effective_kv_committed_len()
@@ -222,8 +222,8 @@ def release_kv_cache(req: Req, tree_cache: BasePrefixCache, is_insert: bool = Tr
# StreamingSession.cache_finished_req handles speculative tail trim
# internally, then sets req_pool_idx = None.
assert (not req.kv.is_held) == req.kv.is_released
if not req.kv.is_held:
assert (not req.kv.holds_kv) == req.kv.is_kv_released
if not req.kv.holds_kv:
return
start_p, end_p = effective_kv_committed_len, req.kv.kv_allocated_len
@@ -234,13 +234,13 @@ def release_kv_cache(req: Req, tree_cache: BasePrefixCache, is_insert: bool = Tr
not tree_cache.supports_mamba()
):
assert (
req.mamba_pool_idx is not None
req.kv.holds_mamba
), "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.
tree_cache.req_to_token_pool.free(req)
req.kv.mark_released()
req.kv.mark_kv_released()
def _release_overallocated_kv_indices(
python/sglang/srt/mem_cache/mamba_radix_cache.py
+17 -17
View File
@@ -561,7 +561,7 @@ class MambaRadixCache(BasePrefixCache):
if is_insert:
if self.enable_mamba_extra_buffer:
cache_len = req.mamba_last_track_seqlen
cache_len = req.kv.mamba_last_track_seqlen
else:
cache_len = len(token_ids)
# ReplaySSM (no_buffer): `temporal[slot]` lags the live state by
@@ -571,8 +571,8 @@ class MambaRadixCache(BasePrefixCache):
# with its key length. page_size is asserted == 1, so no realign.
write_pos_buf = self.req_to_token_pool.mamba_pool.replayssm_write_pos
if write_pos_buf is not None:
cache_len -= int(write_pos_buf[req.mamba_pool_idx].item())
write_pos_buf[req.mamba_pool_idx] = 0
cache_len -= int(write_pos_buf[req.kv.mamba_pool_idx].item())
write_pos_buf[req.kv.mamba_pool_idx] = 0
if cache_len is None:
cache_len = 0
if cache_len != len(token_ids):
@@ -602,16 +602,16 @@ class MambaRadixCache(BasePrefixCache):
mamba_ping_pong_track_buffer_to_keep = (
self.req_to_token_pool.get_mamba_ping_pong_keep_idx(req)
)
src_active = req.mamba_ping_pong_track_buffer[
src_active = req.kv.mamba_ping_pong_track_buffer[
mamba_ping_pong_track_buffer_to_keep
].unsqueeze(-1)
if _MAMBA_DEBUG_ASSERTS:
# .item() forces a cudaStreamSynchronize; only pay it when debugging.
assert src_active.item() != -1, (
f"Cached mamba slot is -1: keep_idx={mamba_ping_pong_track_buffer_to_keep}, "
f"buf={req.mamba_ping_pong_track_buffer.tolist()}, "
f"next_track_idx={req.mamba_next_track_idx}, "
f"last_track_seqlen={req.mamba_last_track_seqlen}, "
f"buf={req.kv.mamba_ping_pong_track_buffer.tolist()}, "
f"next_track_idx={req.kv.mamba_next_track_idx}, "
f"last_track_seqlen={req.kv.mamba_last_track_seqlen}, "
f"rid={req.rid}"
)
if self.int8_ckpt_pool is not None:
@@ -623,10 +623,10 @@ class MambaRadixCache(BasePrefixCache):
else:
if self.int8_ckpt_pool is not None:
mamba_value = self._commit_int8_checkpoint(
req.mamba_pool_idx.unsqueeze(-1)
req.kv.mamba_pool_idx.unsqueeze(-1)
)
else:
mamba_value = req.mamba_pool_idx.unsqueeze(-1).clone()
mamba_value = req.kv.mamba_pool_idx.unsqueeze(-1).clone()
mamba_ping_pong_track_buffer_to_keep = None
result = self.insert(
@@ -685,7 +685,7 @@ class MambaRadixCache(BasePrefixCache):
token_ids = req.get_fill_ids()
cache_len = (
req.mamba_last_track_seqlen
req.kv.mamba_last_track_seqlen
if self.enable_mamba_extra_buffer
else len(token_ids)
)
@@ -726,7 +726,7 @@ class MambaRadixCache(BasePrefixCache):
self.req_to_token_pool.mamba_allocator.free(src_active)
else:
mamba_value_donated = self._commit_int8_checkpoint(
req.mamba_pool_idx.view(-1)
req.kv.mamba_pool_idx.view(-1)
)
elif self.enable_mamba_extra_buffer:
new_slot = self._alloc_mamba_slot()
@@ -739,7 +739,7 @@ class MambaRadixCache(BasePrefixCache):
# 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(
translate(req.mamba_pool_idx.unsqueeze(0)),
translate(req.kv.mamba_pool_idx.unsqueeze(0)),
translate(mamba_value_donated),
)
@@ -799,7 +799,7 @@ class MambaRadixCache(BasePrefixCache):
[new_indices, kv_indices_orig[len(new_indices) :]]
)
req.kv.cache_protected_len = len(new_indices)
req.mamba_last_track_seqlen = None
req.kv.mamba_last_track_seqlen = None
req.last_node = new_last_node
def pretty_print(self) -> None:
@@ -1179,7 +1179,7 @@ class MambaRadixCache(BasePrefixCache):
# Defer COW to forward stream: record source index, allocate destination
if cow_mamba and last_node.mamba_value is not None:
if req.mamba_pool_idx is None:
if not req.kv.holds_mamba:
dst_index = self.req_to_token_pool.mamba_allocator.alloc(1)
if dst_index is None:
self.inc_lock_ref(last_node)
@@ -1187,9 +1187,9 @@ class MambaRadixCache(BasePrefixCache):
dst_index = self.req_to_token_pool.mamba_allocator.alloc(1)
self.dec_lock_ref(last_node)
assert dst_index is not None, "Can not alloc mamba cache"
req.mamba_pool_idx = dst_index[0]
req.mamba_cow_src_index = last_node.mamba_value
req.mamba_needs_clear = False
req.kv.mamba_pool_idx = dst_index[0]
req.kv.mamba_cow_src_index = last_node.mamba_value
req.kv.mamba_needs_clear = False
value = value[:best_value_len]
if value:
python/sglang/srt/mem_cache/memory_pool.py
+26 -24
View File
@@ -1359,32 +1359,34 @@ class HybridReqToTokenPool(ReqToTokenPool):
mamba_indices: list[torch.Tensor] = []
mamba_ping_pong_track_buffers: list[torch.Tensor] = []
for req in reqs:
if req.mamba_pool_idx is not None: # for radix cache / continuing chunked
if req.kv.holds_mamba: # for radix cache / continuing chunked
pass
else:
mid = self.mamba_allocator.alloc(1)
assert (
mid is not None
), f"Not enough space for mamba cache, try to increase --mamba-full-memory-ratio or --max-mamba-cache-size. {mid=}, {self.mamba_pool.size=}, {self.mamba_allocator.available_size()=}, {len(reqs)=}"
req.mamba_pool_idx = mid[0]
req.mamba_needs_clear = True
req.kv.mamba_pool_idx = mid[0]
req.kv.mamba_needs_clear = True
# GDN ReplaySSM: a freshly (re)assigned slot starts an empty
# ring. write_pos=0 means "ring empty", so the decode kernel
# ignores ring contents and reads only the checkpoint state
# (the post-prefill state that prefill wrote into this slot).
if self.mamba_pool.replayssm_write_pos is not None:
self.mamba_pool.replayssm_write_pos[req.mamba_pool_idx] = 0
self.mamba_pool.replayssm_write_pos[req.kv.mamba_pool_idx] = 0
# ReplaySSM spec-verify ring: an empty ring also resets the
# circular origin + flush flag so the first verify step on this
# freshly-prefilled slot reconstructs from the checkpoint alone.
if self.mamba_pool.replayssm_cache_base is not None:
self.mamba_pool.replayssm_cache_base[req.mamba_pool_idx] = 0
self.mamba_pool.replayssm_is_flush[req.mamba_pool_idx] = 0
mamba_indices.append(req.mamba_pool_idx)
self.mamba_pool.replayssm_cache_base[req.kv.mamba_pool_idx] = 0
self.mamba_pool.replayssm_is_flush[req.kv.mamba_pool_idx] = 0
mamba_indices.append(req.kv.mamba_pool_idx)
if self.enable_mamba_extra_buffer:
if req.mamba_ping_pong_track_buffer is None:
if req.kv.mamba_ping_pong_track_buffer is None:
self._alloc_ping_pong_buffer(req)
mamba_ping_pong_track_buffers.append(req.mamba_ping_pong_track_buffer)
mamba_ping_pong_track_buffers.append(
req.kv.mamba_ping_pong_track_buffer
)
assert len(select_index) == len(
mamba_indices
), "Not enough space for mamba cache, try to increase --mamba-full-memory-ratio or --max-mamba-cache-size."
@@ -1449,7 +1451,7 @@ class HybridReqToTokenPool(ReqToTokenPool):
def get_mamba_ping_pong_keep_idx(self, req: Req) -> int:
"""Return the ping-pong index holding the most recent tracked state."""
return req.mamba_last_track_idx
return req.kv.mamba_last_track_idx
def _alloc_ping_pong_buffer(self, req: Req):
"""Allocate the ping-pong track buffer for a new request.
@@ -1474,9 +1476,9 @@ class HybridReqToTokenPool(ReqToTokenPool):
device=slots.device,
)
buf[:n] = slots
req.mamba_ping_pong_track_buffer = buf
req.mamba_next_track_idx = 0
req.mamba_last_track_idx = (
req.kv.mamba_ping_pong_track_buffer = buf
req.kv.mamba_next_track_idx = 0
req.kv.mamba_last_track_idx = (
0
if self.enable_mamba_extra_buffer_lazy
else self.get_mamba_ping_pong_other_idx(0)
@@ -1489,9 +1491,9 @@ class HybridReqToTokenPool(ReqToTokenPool):
req_index_to_mamba_ping_pong_track_buffer_mapping in sync so that
set_mamba_track_indices_from_reqs reads correct slot indices.
"""
req.mamba_ping_pong_track_buffer[idx] = value
req.kv.mamba_ping_pong_track_buffer[idx] = value
self.req_index_to_mamba_ping_pong_track_buffer_mapping[req.kv.req_pool_idx] = (
req.mamba_ping_pong_track_buffer
req.kv.mamba_ping_pong_track_buffer
)
def donate_mamba_ping_pong_slot(
@@ -1504,14 +1506,14 @@ class HybridReqToTokenPool(ReqToTokenPool):
"""
donate_idx = self.get_mamba_ping_pong_keep_idx(req)
mamba_value_donated = (
req.mamba_ping_pong_track_buffer[donate_idx].unsqueeze(-1).clone()
req.kv.mamba_ping_pong_track_buffer[donate_idx].unsqueeze(-1).clone()
)
if _MAMBA_DEBUG_ASSERTS:
# .item() forces a cudaStreamSynchronize; only pay it when debugging.
assert mamba_value_donated.item() != -1, (
f"Donated mamba slot is -1: donate_idx={donate_idx}, "
f"buf={req.mamba_ping_pong_track_buffer.tolist()}, "
f"next_track_idx={req.mamba_next_track_idx}, "
f"buf={req.kv.mamba_ping_pong_track_buffer.tolist()}, "
f"next_track_idx={req.kv.mamba_next_track_idx}, "
f"rid={req.rid}"
)
self.set_mamba_ping_pong_slot(req, donate_idx, new_slot[0])
@@ -1520,10 +1522,10 @@ class HybridReqToTokenPool(ReqToTokenPool):
def free_mamba_cache(
self, req: Req, mamba_ping_pong_track_buffer_to_keep: Optional[int] = None
):
mamba_index = req.mamba_pool_idx
mamba_index = req.kv.mamba_pool_idx
assert mamba_index is not None, "double free? mamba_index is None"
self.mamba_allocator.free(mamba_index.unsqueeze(0))
req.mamba_pool_idx = None
req.kv.mamba_pool_idx = None
if self.enable_mamba_extra_buffer:
mamba_ping_pong_track_buffer_to_free = (
@@ -1565,13 +1567,13 @@ class HybridReqToTokenPool(ReqToTokenPool):
]
)
self.mamba_allocator.free(mamba_ping_pong_track_buffer_to_free)
# Match the req.mamba_pool_idx=None clear above so the next
# Match the req.kv.mamba_pool_idx=None clear above so the next
# alloc() doesn't see a stale ping-pong reference on the req
# and skip allocation (which would silently reuse a freed
# tensor on the req side while the new pool slot leaks).
req.mamba_ping_pong_track_buffer = None
req.mamba_next_track_idx = None
req.mamba_last_track_idx = None
req.kv.mamba_ping_pong_track_buffer = None
req.kv.mamba_next_track_idx = None
req.kv.mamba_last_track_idx = None
def clear(self):
logger.info("Reset HybridReqToTokenPool")
python/sglang/srt/mem_cache/unified_cache/components/mamba_component.py
+15 -15
View File
@@ -197,7 +197,7 @@ class MambaComponent(TreeComponent):
return result
req = params.req
assert req is not None
if req.mamba_pool_idx is None:
if not req.kv.holds_mamba:
dst_index = self.cache.req_to_token_pool.mamba_allocator.alloc(1)
if dst_index is None:
# Pin the window via inc/dec_lock_ref so evict's SWA release
@@ -210,9 +210,9 @@ class MambaComponent(TreeComponent):
result.best_match_node, lock_result.to_dec_params()
)
assert dst_index is not None, "Can not alloc mamba cache"
req.mamba_pool_idx = dst_index[0]
req.mamba_cow_src_index = src_index
req.mamba_needs_clear = False
req.kv.mamba_pool_idx = dst_index[0]
req.kv.mamba_cow_src_index = src_index
req.kv.mamba_needs_clear = False
return result
def commit_insert_component_data(
@@ -534,7 +534,7 @@ class MambaComponent(TreeComponent):
is_finished: bool,
) -> Optional[int]:
if self.cache.enable_mamba_extra_buffer:
cache_len = req.mamba_last_track_seqlen
cache_len = req.kv.mamba_last_track_seqlen
else:
cache_len = token_ids_len
# ReplaySSM (no_buffer): `temporal[slot]` lags the live state by the
@@ -548,8 +548,8 @@ class MambaComponent(TreeComponent):
self.cache.req_to_token_pool.mamba_pool.replayssm_write_pos
)
if write_pos_buf is not None:
cache_len -= int(write_pos_buf[req.mamba_pool_idx].item())
write_pos_buf[req.mamba_pool_idx] = 0
cache_len -= int(write_pos_buf[req.kv.mamba_pool_idx].item())
write_pos_buf[req.kv.mamba_pool_idx] = 0
if is_finished:
if cache_len is None:
@@ -559,10 +559,10 @@ class MambaComponent(TreeComponent):
req
)
active_value = (
req.mamba_ping_pong_track_buffer[keep_idx].unsqueeze(-1).clone()
req.kv.mamba_ping_pong_track_buffer[keep_idx].unsqueeze(-1).clone()
)
else:
active_value = req.mamba_pool_idx.unsqueeze(-1).clone()
active_value = req.kv.mamba_pool_idx.unsqueeze(-1).clone()
if self.int8_ckpt_pool is not None:
insert_params.mamba_value = self._commit_int8_checkpoint(active_value)
else:
@@ -584,7 +584,7 @@ class MambaComponent(TreeComponent):
self.cache.req_to_token_pool.mamba_allocator.free(src_active)
else:
mamba_value_donated = self._commit_int8_checkpoint(
req.mamba_pool_idx.view(-1)
req.kv.mamba_pool_idx.view(-1)
)
elif self.cache.enable_mamba_extra_buffer:
new_slot = self._alloc_mamba_slot()
@@ -599,7 +599,7 @@ class MambaComponent(TreeComponent):
# 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(
translate(req.mamba_pool_idx.unsqueeze(0)),
translate(req.kv.mamba_pool_idx.unsqueeze(0)),
translate(mamba_value_donated),
)
insert_params.mamba_value = mamba_value_donated
@@ -647,7 +647,7 @@ class MambaComponent(TreeComponent):
insert_result is None or insert_result.mamba_exist
):
self._free_mamba_value(insert_params.mamba_value)
req.mamba_last_track_seqlen = None
req.kv.mamba_last_track_seqlen = None
def build_external_linker_transfer(
self,
@@ -669,7 +669,7 @@ class MambaComponent(TreeComponent):
) -> PrepareLoadBackResult:
if (
req is None
or req.mamba_pool_idx is not None
or req.kv.holds_mamba
or not self.tree_core.component_has_host_value_only(
node_id, self.component_type
)
@@ -680,7 +680,7 @@ class MambaComponent(TreeComponent):
self.cache.evict_for_alloc(EvictParams(num_tokens=0, mamba_num=1))
dst = self.cache.req_to_token_pool.mamba_allocator.alloc(1)
assert dst is not None, "Cannot alloc mamba for load_back"
req.mamba_pool_idx = dst[0]
req.kv.mamba_pool_idx = dst[0]
return PrepareLoadBackResult(allocated_mamba_slot=dst)
def finalize_load_back(
@@ -689,7 +689,7 @@ class MambaComponent(TreeComponent):
# A called-off load-back returns the slot prepare allocated and clears req (the H->D copy never ran).
if not success and prep.allocated_mamba_slot is not None:
self.cache.req_to_token_pool.mamba_allocator.free(prep.allocated_mamba_slot)
req.mamba_pool_idx = None
req.kv.mamba_pool_idx = None
def prepare_prefetch(
self,
python/sglang/srt/mem_cache/unified_cache/unified_tree_core.py
+1 -1
View File
@@ -1937,7 +1937,7 @@ class UnifiedTreeCore(UnifiedTreeCoreInterface):
) -> tuple[PoolTransfer, dict[ComponentType, list[PoolTransfer]]]:
"""Build the H->D load-back KV transfer plus per-component aux transfers."""
# Component hooks take primitives, not Req: extract its fields here.
mamba_pool_idx = req.mamba_pool_idx if req is not None else None
mamba_pool_idx = req.kv.mamba_pool_idx if req is not None else None
node = self.node_by_id(node_id)
kv_xfer = self.components_by_type[BASE_COMPONENT_TYPE].build_hicache_transfers(
node, CacheTransferPhase.LOAD_BACK
python/sglang/srt/mem_cache/unified_radix_cache.py
+1 -1
View File
@@ -2870,7 +2870,7 @@ class UnifiedRadixCache(BasePrefixCache):
def swa_retain_floor(self, req) -> int | None:
if not self.is_mamba_enabled or self._sliding_window_size is None:
return None
checkpoint = req.mamba_last_track_seqlen
checkpoint = req.kv.mamba_last_track_seqlen
if checkpoint is None:
return None
return checkpoint - self._sliding_window_size
python/sglang/srt/session/streaming_session.py
+23 -62
View File
@@ -53,14 +53,6 @@ class SessionSlot:
# releases only what it took (may share the node with another req).
skip_lock_node_ids: dict = field(default_factory=dict)
# Mamba states
mamba_pool_idx: Any = None
mamba_ping_pong_track_buffer: Any = None
mamba_next_track_idx: Any = None
mamba_last_track_idx: Any = None
mamba_last_track_seqlen: Any = None
mamba_branching_seqlen: Any = None
def save_from_req(self, req: Req, is_first: bool):
"""Save KV state from a finishing request into this slot."""
kv = req.detach_kv()
@@ -74,37 +66,14 @@ class SessionSlot:
# Later turns run on the slot's record (see restore_to_req).
assert kv is self.kv
self.mamba_pool_idx = req.mamba_pool_idx
self.mamba_ping_pong_track_buffer = req.mamba_ping_pong_track_buffer
self.mamba_next_track_idx = req.mamba_next_track_idx
self.mamba_last_track_idx = req.mamba_last_track_idx
self.mamba_last_track_seqlen = req.mamba_last_track_seqlen
self.mamba_branching_seqlen = req.mamba_branching_seqlen
# The mamba state moved to the slot too; clear the req's references so a
# later alloc/retract path cannot mistake slot-owned state for its own.
req.mamba_pool_idx = None
req.mamba_ping_pong_track_buffer = None
req.mamba_next_track_idx = None
req.mamba_last_track_idx = None
req.mamba_last_track_seqlen = None
req.mamba_branching_seqlen = None
def restore_to_req(self, req: Req):
"""Restore KV state from this slot into an incoming request."""
req.kv = self.kv
req.swa_uuid_for_lock = self.swa_uuid_for_lock
req.skip_lock_node_ids = self.skip_lock_node_ids
req.mamba_pool_idx = self.mamba_pool_idx
req.mamba_ping_pong_track_buffer = self.mamba_ping_pong_track_buffer
req.mamba_next_track_idx = self.mamba_next_track_idx
req.mamba_last_track_idx = self.mamba_last_track_idx
req.mamba_last_track_seqlen = self.mamba_last_track_seqlen
req.mamba_branching_seqlen = self.mamba_branching_seqlen
# NOTE: req_pool_idx and mamba_pool_idx are intentionally NOT cleared
# from the slot. During chunked prefill, a request may be rejected by
# NOTE: the slot keeps sharing the record it just handed out. During
# chunked prefill, a request may be rejected by
# the scheduler (e.g. budget exhausted) and retried in the next cycle.
# Each retry calls match_prefix -> restore_to_req again, so the slot
# must remain intact for idempotent restoration.
@@ -176,7 +145,7 @@ class StreamingSession(BasePrefixCache):
return session_id in self.slots
def any_holding_kv(self) -> bool:
return any(s.kv.is_held for s in self.slots.values())
return any(s.kv.holds_kv for s in self.slots.values())
# -- Try-handle entries for composition (see class docstring) --
@@ -204,7 +173,7 @@ class StreamingSession(BasePrefixCache):
if not _is_streaming(req):
return None
slot = self.slots.get(req.session.session_id)
if slot is None or not slot.kv.is_held:
if slot is None or not slot.kv.holds_kv:
return None
if req.to_finish is not None:
req.session.abort_req()
@@ -310,25 +279,17 @@ class StreamingSession(BasePrefixCache):
kv = req.detach_kv()
if slot is None:
# First-request mid-processing abort: create ephemeral
# slot from req state so release_session handles cleanup.
# Include last_node from the req so
# release_session calls dec_lock_ref on the tree lock.
# Also carry the mamba refs over so _free_slot_mamba can
# return the (possibly extra_buffer ping-pong) slots to
# the mamba pool; otherwise the abort orphans them.
# slot from req state so release_session handles cleanup;
# the detached record carries the mamba refs for
# _free_slot_mamba, and last_node lets release_session
# dec_lock_ref the tree lock.
slot = SessionSlot(
kv=kv,
last_node=req.last_node,
swa_uuid_for_lock=req.swa_uuid_for_lock,
skip_lock_node_ids=req.skip_lock_node_ids,
mamba_pool_idx=req.mamba_pool_idx,
mamba_ping_pong_track_buffer=req.mamba_ping_pong_track_buffer,
)
self.slots[session_id] = slot
# Slot now owns the mamba state — drop the req's refs so
# the abort fall-through doesn't double-free.
req.mamba_pool_idx = None
req.mamba_ping_pong_track_buffer = None
else:
assert kv is slot.kv
self.release_session(session_id)
@@ -424,7 +385,7 @@ class StreamingSession(BasePrefixCache):
protected_len = slot.kv.cache_protected_len
lock_node = slot.last_node
tokens_freed = (
max(0, slot.kv.kv_allocated_len - protected_len) if slot.kv.is_held else 0
max(0, slot.kv.kv_allocated_len - protected_len) if slot.kv.holds_kv else 0
)
logger.info(
"Session KV released: %s (%d tokens freed)", session_id, tokens_freed
@@ -439,7 +400,7 @@ class StreamingSession(BasePrefixCache):
),
)
if slot.kv.is_held:
if slot.kv.holds_kv:
start = protected_len
end = slot.kv.kv_allocated_len
if start < end:
@@ -467,7 +428,7 @@ class StreamingSession(BasePrefixCache):
active_pool_idxs is not None
and slot.kv.req_pool_idx in active_pool_idxs
)
if slot.kv.is_held and not in_batch:
if slot.kv.holds_kv and not in_batch:
allocated = ceil_align(slot.kv.kv_allocated_len, self.page_size)
total += allocated - slot.kv.cache_protected_len
return total
@@ -484,7 +445,7 @@ class StreamingSession(BasePrefixCache):
active_pool_idxs is not None
and slot.kv.req_pool_idx in active_pool_idxs
)
if slot.kv.is_held and not in_batch:
if slot.kv.holds_kv and not in_batch:
allocated = ceil_align(slot.kv.kv_allocated_len, self.page_size)
total += allocated - max(
slot.kv.cache_protected_len, slot.kv.swa_evicted_seqlen
@@ -498,7 +459,7 @@ class StreamingSession(BasePrefixCache):
in_batch = (
active_pool_idxs is not None and s.kv.req_pool_idx in active_pool_idxs
)
return s.kv.is_held and not in_batch
return s.kv.holds_kv and not in_batch
return sum(_owned(s) for s in self.slots.values())
@@ -517,10 +478,10 @@ class StreamingSession(BasePrefixCache):
)
if in_batch:
continue
if slot.mamba_pool_idx is not None:
total += slot.mamba_pool_idx.numel()
if slot.mamba_ping_pong_track_buffer is not None:
total += slot.mamba_ping_pong_track_buffer.numel()
if slot.kv.holds_mamba:
total += slot.kv.mamba_pool_idx.numel()
if slot.kv.mamba_ping_pong_track_buffer is not None:
total += slot.kv.mamba_ping_pong_track_buffer.numel()
return total
def _free_slot_mamba(self, slot: SessionSlot) -> None:
@@ -528,12 +489,12 @@ class StreamingSession(BasePrefixCache):
mamba_allocator = getattr(self.req_to_token_pool, "mamba_allocator", None)
if mamba_allocator is None:
return
if slot.mamba_pool_idx is not None:
mamba_allocator.free(slot.mamba_pool_idx.unsqueeze(0))
slot.mamba_pool_idx = None
if slot.mamba_ping_pong_track_buffer is not None:
mamba_allocator.free(slot.mamba_ping_pong_track_buffer)
slot.mamba_ping_pong_track_buffer = None
if slot.kv.holds_mamba:
mamba_allocator.free(slot.kv.mamba_pool_idx.unsqueeze(0))
slot.kv.mamba_pool_idx = None
if slot.kv.mamba_ping_pong_track_buffer is not None:
mamba_allocator.free(slot.kv.mamba_ping_pong_track_buffer)
slot.kv.mamba_ping_pong_track_buffer = None
# -- Internal helpers (streaming body bits) --
python/sglang/test/scripted_runtime/req_handle.py
+1 -1
View File
@@ -42,7 +42,7 @@ class ScriptedReqHandle:
@property
def kv_pages(self) -> int:
req = self.req
if req is None or not req.kv.is_held:
if req is None or not req.kv.holds_kv:
return 0
page_size = self.context.scheduler.page_size
return (req.kv.kv_allocated_len + page_size - 1) // page_size