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

Use a shared byte budget for unified hybrid-SWA memory (#36729)

Browse Source 
Co-authored-by: yhzhuang <yhzhuang@fb.com>
Co-authored-by: Cheng Wan <cheng.wan@radixark.ai>
This commit is contained in:
Yonghao Zhuang
2026-09-15 15:27:00 -07:00
committed by GitHub
co-authored by yhzhuang Cheng Wan
parent 4da5599e93
commit 2929a39927
32 changed files with 2220 additions and 588 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/arg_groups/overrides.py
+3
View File
@@ -1808,6 +1808,9 @@ def post_capture_kv_sizing_planned(server_args: Any) -> bool:
mla_enabled = use_mla_backend(server_args)
if not envs.SGLANG_ENABLE_POST_CAPTURE_KV_SIZING.get():
return False
# Unified arenas are fully backed before capture and cannot resize afterward.
if cfg.enable_unified_memory:
return False
if cfg.device != "cuda":
return False
if cfg.dcp_size != 1:
python/sglang/srt/managers/schedule_policy.py
+52 -246
View File
@@ -50,9 +50,7 @@ from sglang.srt.mem_cache.allocator.hisparse import (
DeepSeekV4HiSparseTokenToKVPoolAllocator,
)
from sglang.srt.mem_cache.allocator.swa import (
PureSWATokenToKVPoolAllocator,
SWATokenToKVPoolAllocator,
is_swa_req_ring,
)
from sglang.srt.mem_cache.allocator.unified_hybrid_swa import (
UnifiedMambaSWATokenToKVPoolAllocator,
@@ -581,8 +579,6 @@ class PrefillAdder:
self.prefill_tile_block_m = prefill_tile_block_m
self.tree_cache = tree_cache
self.token_to_kv_pool_allocator = token_to_kv_pool_allocator
# Per-request SWA ring: one fixed slot per request, not a token budget.
self._swa_req_ring = is_swa_req_ring(token_to_kv_pool_allocator)
self.running_batch = running_batch
self.new_token_ratio = new_token_ratio
self.rem_input_tokens = rem_input_tokens - num_mixed_decode_tokens
@@ -595,8 +591,9 @@ class PrefillAdder:
if self.rem_chunk_tokens is not None:
self.rem_chunk_tokens -= num_mixed_decode_tokens
self.rem_total_token_offset = num_mixed_decode_tokens
self.cur_rem_token_offset = num_mixed_decode_tokens
self.memory_budget = token_to_kv_pool_allocator.create_prefill_budget(
tree_cache, num_mixed_decode_tokens=num_mixed_decode_tokens
)
self.req_states = None
self.can_run_list = []
@@ -612,7 +609,7 @@ class PrefillAdder:
if running_batch is not None:
# Estimate the offset in the remaining token space
self.rem_total_token_offset += sum(
self.memory_budget.total_offset += sum(
[
self._get_running_request_total_token_offset(r)
for r in running_batch.reqs
@@ -625,13 +622,7 @@ class PrefillAdder:
self.token_to_kv_pool_allocator,
(SWATokenToKVPoolAllocator, DeepSeekV4HiSparseTokenToKVPoolAllocator),
)
self.is_all_swa = isinstance(
self.token_to_kv_pool_allocator, PureSWATokenToKVPoolAllocator
)
self.is_hybrid_ssm_cache = self.tree_cache.supports_mamba()
self.rem_swa_token_offset = 0
# A new state slot eats shared-gap bytes that `rem_total_tokens` counts
# as free, so reserve per slot or admission over-commits. Gate on the
# ALLOCATOR, not `is_hybrid_ssm_cache`: that is False for `ChunkCache`,
@@ -718,114 +709,11 @@ class PrefillAdder:
@property
def rem_total_tokens(self):
if self.is_all_swa:
available_and_evictable = (
self.token_to_kv_pool_allocator.swa_available_size()
+ self.tree_cache.swa_evictable_size()
)
elif self.is_hybrid_swa:
available_and_evictable = (
self.token_to_kv_pool_allocator.full_available_size()
+ self.tree_cache.full_evictable_size()
)
elif self.is_hybrid_ssm_cache:
available_and_evictable = (
self.token_to_kv_pool_allocator.available_size()
+ self.tree_cache.full_evictable_size()
)
else:
available_and_evictable = (
self.token_to_kv_pool_allocator.available_size()
+ self.tree_cache.evictable_size()
)
return available_and_evictable - self.rem_total_token_offset
@property
def rem_swa_tokens(self):
allocator = self.token_to_kv_pool_allocator
if self._swa_req_ring:
# swa_available_size() already reports ring capacity; tree
# swa_evictable is in linear token units and frees no ring space.
return allocator.swa_available_size() - self.rem_swa_token_offset
return (
allocator.swa_available_size()
+ self.tree_cache.swa_evictable_size()
- self.rem_swa_token_offset
)
return self.memory_budget.remaining_total
@property
def cur_rem_tokens(self):
if self.is_all_swa:
available_and_evictable = (
self.token_to_kv_pool_allocator.swa_available_size()
+ self.tree_cache.swa_evictable_size()
)
elif self.is_hybrid_swa:
available_and_evictable = (
self.token_to_kv_pool_allocator.full_available_size()
+ self.tree_cache.full_evictable_size()
)
elif self.is_hybrid_ssm_cache:
available_and_evictable = (
self.token_to_kv_pool_allocator.available_size()
+ self.tree_cache.full_evictable_size()
)
else:
available_and_evictable = (
self.token_to_kv_pool_allocator.available_size()
+ self.tree_cache.evictable_size()
)
return available_and_evictable - self.cur_rem_token_offset
def _swa_budget_for_req(
self, extend_input_len: int, max_new_tokens: int, swa_host_hit_length: int = 0
) -> int:
"""SWA pool budget per request. Only valid when is_hybrid_swa is True.
With chunked prefill + overlap scheduler, the peak SWA occupancy is:
chunk N (running, not yet in tree) + sliding window (locked in tree)
+ chunk N+1 (new allocation)
Since chunk N and locked tokens are already excluded from
swa_available + swa_evictable, the budget only needs to cover the
chunk N+1 allocation plus decode headroom:
budget = max(alloc - window, 0) + min(extend + max_new_tokens, window) + page
where alloc = min(extend, rem_chunk); the min() cap keeps the two terms
from double-counting extend, so budget <= extend + max_new_tokens + page.
"""
allocator = self.token_to_kv_pool_allocator
if self._swa_req_ring:
# One ring slot per request, in the same unit as swa_available_size.
return allocator.swa_ring_cost_tokens
if self.rem_chunk_tokens is not None:
alloc = min(extend_input_len, self.rem_chunk_tokens)
else:
alloc = extend_input_len
window = self.tree_cache.sliding_window_size
return max(alloc - window, 0) + self._swa_reserved_tokens(
extend_input_len, max_new_tokens, swa_host_hit_length
)
def _swa_reserved_tokens(
self, extend_input_len: int, max_new_tokens: int, swa_host_hit_length: int = 0
) -> int:
"""SWA slots a request adds to its own sliding window + page slack + the
load-back charge. Shared floor of _swa_budget_for_req and _swa_chunk_cap.
The headroom is min(extend + decode, window), not a constant window: a
request contributes only extend + decode fresh tokens to its window and
a cached SWA prefix funds the rest. Charging a full window double-counted
a short cached-prefix resume and livelocked admission at a ~2-window
pool; keeping extend in the min() holds the reservation >= the prefill
allocation so an admitted request cannot OOM."""
window = self.tree_cache.sliding_window_size
headroom = min(extend_input_len + max_new_tokens, window)
reserved = headroom + self.page_size
if swa_host_hit_length > 0:
reserved += self.ceil_paged_tokens(swa_host_hit_length)
return reserved
return self.memory_budget.remaining_current
def _swa_new_tokens(self, req: Req) -> int:
"""Tokens a request may still decode, for SWA headroom sizing. Mirrors
@@ -837,80 +725,22 @@ class PrefillAdder:
CLIP_MAX_NEW_TOKENS,
)
def _swa_chunk_cap(self, max_new_tokens: int, swa_host_hit_length: int = 0) -> int:
"""Largest page-aligned extend chunk the SWA pool can admit right now,
keeping a sliding window of headroom below rem_swa_tokens; 0 if not
even one page fits. Only valid when is_hybrid_swa is True.
Escape hatch for a request whose budget can never pass the
_swa_budget_for_req gate (extend near/above the pool size, or a large
load-back charge): without shrinking its chunk it would be rejected
forever (head-of-line livelock). Shrinking is sound because past a
chunk boundary only the sliding window stays locked — the rest turns
evictable — so each pass's transient footprint fits the pool."""
# extend_input_len=0: this solves for the extend chunk itself, so the
# reserved headroom is the post-chunk decode window only.
cap = int(self.rem_swa_tokens) - self._swa_reserved_tokens(
0, max_new_tokens, swa_host_hit_length
)
if cap <= 0:
return 0
return cap // self.page_size * self.page_size
def _swa_req_never_fits(
self, extend_input_len: int, max_new_tokens: int, swa_host_hit_length: int = 0
) -> bool:
"""True when a request's SWA budget exceeds the *entire* SWA pool, so it
can never be admitted whole no matter how far the pool drains.
This is the head-of-line livelock the _swa_chunk_cap escape hatch exists
for; the hatch must fire only in this case. A request that merely
exceeds *current* rem_swa (transient pressure) would fit once running
decodes free their windows, so it must wait — admitting it into the
decode headroom collapses the SWA evictable cushion and forces running
requests to retract (observed as a severe retraction/re-prefill storm on
hybrid-SWA models at high concurrency)."""
capacity = self.token_to_kv_pool_allocator.size_swa
return (
self._swa_budget_for_req(
extend_input_len, max_new_tokens, swa_host_hit_length
)
>= capacity
)
def _swa_admission_gate(
def _check_prefill_budget(
self,
req: Req,
*,
extend_input_len: int,
total_tokens: int,
swa_host_hit_length: int,
chunk_tokens_limit: Optional[int],
) -> tuple[Optional[AddReqResult], Optional[int]]:
"""SWA-pool gate: a non-None verdict rejects; otherwise the returned chunk
limit stands, tightened to the pool cap when never-fits fires."""
max_new_tokens = self._swa_new_tokens(req)
swa_needed = self._swa_budget_for_req(
extend_input_len, max_new_tokens, swa_host_hit_length=swa_host_hit_length
) -> tuple[bool, Optional[int]]:
return self.memory_budget.check_prefill(
extend_input_len=extend_input_len,
total_tokens=total_tokens,
max_new_tokens=self._swa_new_tokens(req),
input_tokens=len(req.full_untruncated_fill_ids),
swa_host_hit_length=swa_host_hit_length,
chunk_limit=self.rem_chunk_tokens,
)
# Ring-slot capacity is exact, so needing exactly what is left still
# fits; the legacy SWA-token path keeps its conservative `>=`.
fits = (
swa_needed <= self.rem_swa_tokens
if self._swa_req_ring
else swa_needed < self.rem_swa_tokens
)
if fits:
return None, chunk_tokens_limit
if not self._swa_req_never_fits(
extend_input_len, max_new_tokens, swa_host_hit_length
):
return AddReqResult.NO_TOKEN, chunk_tokens_limit
swa_cap = self._swa_chunk_cap(max_new_tokens, swa_host_hit_length)
if self.rem_chunk_tokens is None or swa_cap <= 0:
return AddReqResult.NO_TOKEN, chunk_tokens_limit
current = (
self.rem_chunk_tokens if chunk_tokens_limit is None else chunk_tokens_limit
)
return None, min(current, swa_cap)
def _mamba_gap_budget_for_req(self, req: Req) -> int:
"""Shared-gap reservation (full-token-equivalents) for a request's new
@@ -931,9 +761,7 @@ class PrefillAdder:
return -(-tokens // self.page_size) * self.page_size
def budget_state(self):
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
no_token = not self.memory_budget.has_capacity()
# 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:
@@ -980,17 +808,12 @@ class PrefillAdder:
if compute_charge is None:
compute_charge = 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
# `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
self.memory_budget.reserve(
extend_input_len,
max_new_tokens,
extra_tokens=mamba_gap_reserve,
chunk_limit=self.rem_chunk_tokens,
is_chunked_continuation=is_chunked_continuation,
)
# The new mamba slot also consumes one mamba-recoverable slot (gated
# separately so full_evictable can't cover it — see __init__).
@@ -998,14 +821,6 @@ class PrefillAdder:
self.rem_mamba_slots -= 1
self.rem_input_tokens -= compute_charge
if self.is_hybrid_swa:
# The ring slot is reserved once at first admission; charging it
# again on a continuation would double-count and over-throttle.
if not (self._swa_req_ring and is_chunked_continuation):
self.rem_swa_token_offset += self._swa_budget_for_req(
extend_input_len, max_new_tokens
)
if self.dllm_config is not None:
self.rem_dllm_tokens -= compute_charge
elif self.rem_chunk_tokens is not None:
@@ -1136,20 +951,11 @@ class PrefillAdder:
if self.dllm_config is not None:
_rem_tokens = self._get_dllm_remain_tokens()
else:
_rem_tokens = min(self.rem_chunk_tokens, int(self.rem_total_tokens))
if self.is_hybrid_swa and not self._swa_req_ring:
# alloc_extend needs extend_num_tokens + page_size per request,
# so reserve one page here to avoid OOM.
# Ring mode skips it: rem_swa_tokens counts slots, not chunk tokens.
_rem_tokens = min(
_rem_tokens, int(self.rem_swa_tokens) - self.page_size
)
# The chunked_req must be added to the list; otherwise, it will cause a memory leak.
# Therefore, in certain cases where _rem_tokens <= 0, it should be replaced with rem_chunk_tokens.
if _rem_tokens <= 0:
if self.is_hybrid_swa:
return req
_rem_tokens = self.rem_chunk_tokens
_rem_tokens = self.memory_budget.available_chunk_tokens(
self.rem_chunk_tokens
)
if _rem_tokens is None:
return req
# A mid-chunk rank prefills this pass regardless of the delayer
# verdict, so report prefillable=True and ignore the result.
@@ -1165,6 +971,13 @@ class PrefillAdder:
cand_extend_input_len = len(req.full_untruncated_fill_ids) - len(
req.prefix_indices
)
_rem_tokens = self.memory_budget.fit_chunk(
extend_input_len=cand_extend_input_len,
max_new_tokens=self._swa_new_tokens(req),
chunk_limit=_rem_tokens,
)
if _rem_tokens is None:
return req
truncated = cand_extend_input_len > _rem_tokens
new_len = min(cand_extend_input_len, _rem_tokens)
req.set_extend_range(len(req.prefix_indices), len(req.prefix_indices) + new_len)
@@ -1210,16 +1023,14 @@ class PrefillAdder:
# 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):
fits = self.memory_budget.can_allocate_prefill(
paged_input=paged_input,
extend_input_len=cand_extend_input_len,
max_new_tokens=self._swa_new_tokens(req),
chunk_limit=self.rem_chunk_tokens,
)
if not fits:
return AddReqResult.NO_TOKEN
if self.is_hybrid_swa:
if (
self._swa_budget_for_req(
cand_extend_input_len, self._swa_new_tokens(req)
)
> self.rem_swa_tokens
):
return AddReqResult.NO_TOKEN
def add_req_state(r, insert_sort=False):
new_token_ratio = (
@@ -1369,9 +1180,6 @@ class PrefillAdder:
mamba_gap_reserve = self._mamba_gap_budget_for_req(req)
total_tokens += mamba_gap_reserve
if total_tokens >= self.rem_total_tokens:
return AddReqResult.NO_TOKEN
# The temporary pin excludes this prefix from the evictable budget.
# Selection itself neither allocates slots nor materializes host hits.
with self._lock_node(req.last_node):
@@ -1466,9 +1274,6 @@ class PrefillAdder:
truncation_align_size: Optional[int],
) -> _PrefillAdmission | AddReqResult:
"""Select a prefill shape without allocating or publishing cached KV."""
if total_tokens >= self.rem_total_tokens:
return AddReqResult.NO_TOKEN
prefix_len = len(req.prefix_indices) + host_hit_length
extend_len = len(req.full_untruncated_fill_ids) - prefix_len
input_tokens = self.ceil_paged_tokens(extend_len)
@@ -1476,13 +1281,14 @@ class PrefillAdder:
# exact-chunk-fill, so a request whose ceiled length would spill is
# not needlessly split into a second chunk.
chunk_fit_tokens = extend_len if self.exact_chunk_fill else input_tokens
chunk_tokens_limit = self.rem_chunk_tokens
if self.is_hybrid_swa:
verdict, chunk_tokens_limit = self._swa_admission_gate(
req, input_tokens, swa_host_hit_length, chunk_tokens_limit
)
if verdict is not None:
return verdict
can_admit, chunk_tokens_limit = self._check_prefill_budget(
req,
extend_input_len=extend_len,
total_tokens=total_tokens,
swa_host_hit_length=swa_host_hit_length,
)
if not can_admit:
return AddReqResult.NO_TOKEN
# Without chunking, allow the first request even above the input cap.
if (
@@ -1619,7 +1425,7 @@ class PrefillAdder:
release_counter = 0
for i, running_req in enumerate(self.running_batch.reqs):
if running_req in preemptible_reqs:
self.rem_total_token_offset -= (
self.memory_budget.total_offset -= (
self._get_running_request_total_token_offset(running_req)
)
release_counter += 1
python/sglang/srt/managers/scheduler.py
+16 -11
View File
@@ -2523,23 +2523,28 @@ class Scheduler(
)
max_new_tokens = min(max_new_tokens, self.max_new_tokens_limit)
# Keep this bound consistent with PrefillAdder's admission budget:
# ceil_page(input_len) + max_new_tokens + page_size must be strictly
# smaller than max_total_num_tokens. Otherwise a request can be accepted
# into the waiting queue but can never be scheduled, blocking the queue
# and eventually making health checks fail.
paged_input_len = -(-input_len // self.page_size) * self.page_size
req.sampling_params.max_new_tokens = max(
# Keep this bound consistent with PrefillAdder's admission budget.
max_new_tokens = max(
0,
min(
max_new_tokens,
self.max_req_len - input_len - 1,
self.max_total_num_tokens * get_parallel().attn_dcp_size
- paged_input_len
- self.page_size
- 1,
),
)
max_new_tokens = self.token_to_kv_pool_allocator.max_new_tokens_for_memory(
input_len,
max_new_tokens,
token_capacity=self.max_total_num_tokens * get_parallel().attn_dcp_size,
sliding_window_size=self.sliding_window_size,
chunk_size=self.chunked_prefill_size,
)
if max_new_tokens is None:
req.set_finish_with_abort(
f"Request prompt exceeds the KV memory budget: input_len={input_len}."
)
max_new_tokens = 0
req.sampling_params.max_new_tokens = max(0, max_new_tokens)
# Clipping above can push max_new_tokens below min_new_tokens, which
# would suppress EOS for the whole generation. Restore the invariant.
if req.sampling_params.min_new_tokens > req.sampling_params.max_new_tokens:
python/sglang/srt/managers/scheduler_components/invariant_checker.py
+6 -19
View File
@@ -22,9 +22,6 @@ from sglang.srt.managers.scheduler_components.pool_stats_observer import (
)
from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator
from sglang.srt.mem_cache.allocator.swa import is_swa_req_ring
from sglang.srt.mem_cache.allocator.unified_hybrid_swa import (
UnifiedMambaSWATokenToKVPoolAllocator,
)
from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache
from sglang.srt.mem_cache.memory_pool import ReqToTokenPool
from sglang.srt.observability.scheduler_stage_metrics import (
@@ -94,12 +91,13 @@ class SchedulerInvariantChecker:
return leak, msg
def _check_full_pool(self, ps: PoolStats, uncached: int = 0) -> Tuple[bool, str]:
if self.is_hybrid_swa and not self.full_tokens_per_layer:
allocator = self.token_to_kv_pool_allocator
if self.is_hybrid_swa and not ps.full_capacity:
return False, ""
if self.is_hybrid_swa:
protected = self.tree_cache.full_protected_size()
session_held = self.pool_stats_observer.session_held_full_tokens()
total = self.full_tokens_per_layer
total = ps.full_capacity
elif self.is_hybrid_ssm:
# Branch on cache type for the protected accessor (MambaRadixCache
# splits full/mamba; ChunkCache only has the single protected_size).
@@ -119,7 +117,6 @@ class SchedulerInvariantChecker:
session_held = self.pool_stats_observer.session_held_tokens()
total = self.max_total_num_tokens
full_evictable_size = ps.full_evictable_size
allocator = self.token_to_kv_pool_allocator
if get_parallel().dcp_enabled and allocator.page_size > 1:
# DCP stores logical tokens in widened physical pages. Prefix cache
# counters are logical-token based, while the allocator frees whole
@@ -129,14 +126,9 @@ class SchedulerInvariantChecker:
// allocator.page_size
* allocator.page_size
)
full_available = ps.full_available_size
if isinstance(allocator, UnifiedMambaSWATokenToKVPoolAllocator):
# Pair the static per-layer total with the conserve view, never the
# byte-coordinated one -- see `conserve_full_available_size`.
full_available = allocator.conserve_full_available_size()
leak, msg = self._check_pool_invariant(
"full",
full_available,
ps.full_available_size,
full_evictable_size,
protected,
session_held,
@@ -162,18 +154,13 @@ class SchedulerInvariantChecker:
f"evictable={ps.swa_evictable_size}, "
f"total={self.swa_tokens_per_layer}"
)
swa_available = ps.swa_available_size
if isinstance(allocator, UnifiedMambaSWATokenToKVPoolAllocator):
# Tri-pool: same floating-boundary phantom as the full pool -- use the
# slot-conservation view, not the byte-coordinated min (see _check_full_pool).
swa_available = allocator.conserve_swa_available_size()
return self._check_pool_invariant(
"swa",
swa_available,
ps.swa_available_size,
ps.swa_evictable_size,
self.tree_cache.swa_protected_size(),
self.pool_stats_observer.session_held_swa_tokens(),
self.swa_tokens_per_layer,
ps.swa_capacity,
uncached,
)
python/sglang/srt/managers/scheduler_components/pool_stats_observer.py
+20 -23
View File
@@ -12,9 +12,6 @@ from typing import (
)
from sglang.srt.mem_cache.allocator.swa import is_swa_req_ring
from sglang.srt.mem_cache.allocator.unified_hybrid_swa import (
UnifiedMambaSWATokenToKVPoolAllocator,
)
if TYPE_CHECKING:
from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator
@@ -38,6 +35,8 @@ class PoolStats:
is_hisparse: bool = False
# For hybrid-swa pools
full_capacity: Optional[int] = None
swa_capacity: Optional[int] = None
swa_num_used: Optional[int] = None
swa_token_usage: Optional[float] = None
swa_available_size: Optional[int] = None
@@ -289,29 +288,20 @@ class SchedulerPoolStatsObserver:
)
def _get_swa_token_info(self) -> PoolStats:
# `*_num_used` is `static_cap - (available + evictable)`, so the
# available term must match the static cap's denomination: the conserve
# view, never the byte-coordinated one (see
# `conserve_full_available_size`). Measured ~25-90x inflated otherwise.
allocator = self.token_to_kv_pool_allocator
if isinstance(allocator, UnifiedMambaSWATokenToKVPoolAllocator):
full_available_size = allocator.conserve_full_available_size()
swa_available_size = allocator.conserve_swa_available_size()
else:
full_available_size = allocator.full_available_size()
swa_available_size = allocator.swa_available_size()
(full_capacity, full_available_size), (swa_capacity, swa_available_size) = (
self.token_to_kv_pool_allocator.swa_capacity_and_available(
full_capacity=self.full_tokens_per_layer,
swa_capacity=self.swa_tokens_per_layer,
)
)
full_evictable_size = self.tree_cache.full_evictable_size()
swa_evictable_size = self.tree_cache.swa_evictable_size()
# Per-request SWA ring: released with the req slot, yet cached radix
# prefixes still report swa_evictable; counting it drives usage negative.
if is_swa_req_ring(self.token_to_kv_pool_allocator):
swa_evictable_size = 0
full_num_used = self.full_tokens_per_layer - (
full_available_size + full_evictable_size
)
swa_num_used = self.swa_tokens_per_layer - (
swa_available_size + swa_evictable_size
)
full_num_used = full_capacity - (full_available_size + full_evictable_size)
swa_num_used = swa_capacity - (swa_available_size + swa_evictable_size)
# FIXME(hisparse): host-backup transiently over-releases the device pool
# counter, producing negative full_num_used / swa_num_used. We clamp to 0
# to keep token_usage / leak checks sane, but the underlying accounting
@@ -319,16 +309,23 @@ class SchedulerPoolStatsObserver:
if self.enable_hisparse:
full_num_used = max(0, full_num_used)
swa_num_used = max(0, swa_num_used)
if not self.full_tokens_per_layer:
if not full_capacity:
full_num_used = 0
full_available_size = 0
full_token_usage = 0.0
else:
full_token_usage = full_num_used / self.full_tokens_per_layer
swa_token_usage = swa_num_used / self.swa_tokens_per_layer
full_token_usage = full_num_used / full_capacity
if not swa_capacity:
swa_num_used = 0
swa_available_size = 0
swa_token_usage = 0.0
else:
swa_token_usage = swa_num_used / swa_capacity
return PoolStats(
is_hybrid_swa=True,
full_capacity=full_capacity,
swa_capacity=swa_capacity,
full_num_used=full_num_used,
full_token_usage=full_token_usage,
full_available_size=full_available_size,
python/sglang/srt/mem_cache/allocator/base.py
+40 -5
View File
@@ -69,12 +69,47 @@ class BaseTokenToKVPoolAllocator(abc.ABC):
# The scheduler calls these unconditionally, with no allocator-type branches
# on its side; byte-accounted composites override the token-count defaults.
def evict_to_free_tokens(self, tree_cache, num_tokens: int) -> None:
"""Evict unlocked prefix-cache entries until this allocator can serve
``num_tokens`` or nothing evictable remains."""
from sglang.srt.mem_cache.common import evict_from_tree_cache
def create_prefill_budget(self, tree_cache, *, num_mixed_decode_tokens=0):
from sglang.srt.mem_cache.prefill_budget import PrefillBudget
evict_from_tree_cache(tree_cache, num_tokens)
return PrefillBudget(
self, tree_cache, num_mixed_decode_tokens=num_mixed_decode_tokens
)
def max_new_tokens_for_memory(
self,
input_tokens: int,
max_new_tokens: int,
*,
token_capacity: int,
sliding_window_size: int | None,
chunk_size: int | None,
) -> int | None:
"""Clip generation to the empty-pool budget; None means prompt cannot fit.
token_capacity is the scheduler's configured capacity, including its
distributed token scaling. Shared pools use their physical byte layout.
"""
paged_input = -(-input_tokens // self.page_size) * self.page_size
return max(
0, min(max_new_tokens, token_capacity - paged_input - self.page_size - 1)
)
def evict_to_free_tokens(self, tree_cache, num_tokens: int) -> bool | None:
"""Evict unlocked prefix-cache entries until this allocator can serve
``num_tokens`` or nothing evictable remains.
Return whether capacity was realized, or None if it still needs checking.
"""
from sglang.srt.mem_cache.base_prefix_cache import EvictParams
from sglang.srt.mem_cache.common import _evict_until_allocatable
if tree_cache is None or tree_cache.is_chunk_cache():
return
shortfall = num_tokens - self.available_size()
if shortfall > 0:
tree_cache.evict_for_alloc(EvictParams(num_tokens=shortfall))
_evict_until_allocatable(tree_cache, self, num_tokens)
def check_decode_capacity(
self,
python/sglang/srt/mem_cache/allocator/hisparse.py
+14
View File
@@ -357,6 +357,20 @@ class DeepSeekV4HiSparseTokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
kv_indices
)
def create_prefill_budget(self, tree_cache, *, num_mixed_decode_tokens=0):
# Use this wrapper's capacity, which also includes the compressed pool.
from sglang.srt.mem_cache.prefill_budget import SWAPrefillBudget
return SWAPrefillBudget(
self, tree_cache, num_mixed_decode_tokens=num_mixed_decode_tokens
)
def swa_capacity_and_available(self, *, full_capacity, swa_capacity):
return (
(full_capacity, self.full_available_size()),
(swa_capacity, self.swa_available_size()),
)
def full_available_size(self):
return min(
self.logical_attn_allocator.full_available_size(),
python/sglang/srt/mem_cache/allocator/swa.py
+35
View File
@@ -158,6 +158,31 @@ class SWATokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
self.swa_attn_allocator.available_size(),
)
def create_prefill_budget(self, tree_cache, *, num_mixed_decode_tokens=0):
from sglang.srt.mem_cache.prefill_budget import SWAPrefillBudget
return SWAPrefillBudget(
self, tree_cache, num_mixed_decode_tokens=num_mixed_decode_tokens
)
def swa_capacity_and_available(self, *, full_capacity, swa_capacity):
return (
(full_capacity, self.full_available_size()),
(swa_capacity, self.swa_available_size()),
)
def evict_to_free_tokens(self, tree_cache, num_tokens: int) -> None:
from sglang.srt.mem_cache.base_prefix_cache import EvictParams
if tree_cache is None or tree_cache.is_chunk_cache():
return
full_shortfall = max(0, num_tokens - self.full_available_size())
swa_shortfall = max(0, num_tokens - self.swa_available_size())
if full_shortfall or swa_shortfall:
tree_cache.evict_for_alloc(
EvictParams(num_tokens=full_shortfall, swa_num_tokens=swa_shortfall)
)
def full_available_size(self):
return self.full_attn_allocator.available_size()
@@ -687,6 +712,16 @@ class PureSWATokenToKVPoolAllocator(SWATokenToKVPoolAllocator):
def available_size(self):
return self.swa_attn_allocator.available_size()
def create_prefill_budget(self, tree_cache, *, num_mixed_decode_tokens=0):
from sglang.srt.mem_cache.prefill_budget import SWAPrefillBudget
return SWAPrefillBudget(
self,
tree_cache,
num_mixed_decode_tokens=num_mixed_decode_tokens,
all_swa=True,
)
def full_available_size(self):
return self.swa_attn_allocator.available_size()
python/sglang/srt/mem_cache/allocator/unified_hybrid_swa.py
+479 -98
View File
@@ -17,6 +17,8 @@ sub-pools of one `UnifiedKVPool`, and the tri-pool variant that adds mamba state
from __future__ import annotations
import logging
import math
from abc import abstractmethod
from typing import Callable, List, Optional, Sequence, Tuple
import torch
@@ -42,11 +44,10 @@ from sglang.srt.utils.common import get_num_new_pages
logger = logging.getLogger(__name__)
class UnifiedSWATokenToKVPoolAllocator(SWATokenToKVPoolAllocator):
"""Composite allocator for the hybrid SWA pair (full + swa MHA sub-pools).
class UnifiedSWAAllocatorBase(SWATokenToKVPoolAllocator):
"""Shared FULL/SWA virtual IDs, allocation lifecycle, and index translation.
One alloc(N) binds N pages on BOTH sides under the same virtual id, so
`available_size()` (joint bytes, in TOKENS) is the only safe alloc pre-check.
Concrete allocators define the two-ended or Mamba/SWA/FULL capacity policy.
"""
# Parent's `size` property has no setter but base init does `self.size = size`;
@@ -65,26 +66,36 @@ class UnifiedSWATokenToKVPoolAllocator(SWATokenToKVPoolAllocator):
unified_buffer: UnifiedKVPool,
kvcache, # UnifiedSWAKVPool
device: str,
full_max_total_num_tokens: int,
swa_max_total_num_tokens: int,
full_max_total_num_tokens: Optional[int] = None,
swa_max_total_num_tokens: Optional[int] = None,
page_size: int = 1,
need_sort: bool = False,
forward_stream: Optional[torch.cuda.Stream] = None,
lazy_compaction: bool = False,
):
# Set _size_full / _size_swa BEFORE base init (read during it). STATIC
# partition caps -- the slot-conservation value the leak invariant expects.
self._size_full = full_max_total_num_tokens
self._size_swa = swa_max_total_num_tokens
self._full_max_total_num_tokens = full_max_total_num_tokens
self._swa_max_total_num_tokens = swa_max_total_num_tokens
if (full_max_total_num_tokens is None) != (swa_max_total_num_tokens is None):
raise ValueError(
"full_max_total_num_tokens and swa_max_total_num_tokens must "
"either both be set or both be omitted"
)
legacy_capacities = full_max_total_num_tokens is not None
self._size_full = (
int(full_max_total_num_tokens)
if legacy_capacities
else unified_buffer.max_slots("full") - 1
)
self._size_swa = (
int(swa_max_total_num_tokens)
if legacy_capacities
else unified_buffer.max_slots("swa") - 1
)
self.page_size = page_size
# The parent is inherited only for the isinstance contract: skip its
# static-partition sub-pool allocation, which the unified pool replaces.
BaseTokenToKVPoolAllocator.__init__(
self,
size=full_max_total_num_tokens,
size=self._size_full,
page_size=page_size,
dtype=unified_buffer.mha_spec("full").store_dtype,
device=device,
@@ -120,6 +131,16 @@ class UnifiedSWATokenToKVPoolAllocator(SWATokenToKVPoolAllocator):
)
self._wire_peers()
self._empty_shared_gap_bytes = self.full_attn_allocator._current_gap_bytes()
if not legacy_capacities:
self._size_full = self.full_attn_allocator.available_size()
self._size_swa = min(
self.swa_attn_allocator.available_size(),
len(self.full_attn_allocator.free_virtual_ids) * page_size,
)
self._full_max_total_num_tokens = self._size_full
self._swa_max_total_num_tokens = self._size_swa
# Epoch-keyed memo for the joint capacity view (any chain member's
# mutation invalidates -- see `MultiEndedAllocator._chain_capacity_epoch`).
self._joint_avail_memo_epoch: Optional[int] = None
@@ -143,7 +164,7 @@ class UnifiedSWATokenToKVPoolAllocator(SWATokenToKVPoolAllocator):
"[unified-memory-pool] UnifiedSWATokenToKVPoolAllocator ready: "
"full max_slots=%d (min_slot_index=%d, entry_bytes=%d), "
"swa max_slots=%d (min_slot_index=%d, entry_bytes=%d), "
"static caps full=%d swa=%d, joint available=%d",
"max capacity full=%d swa=%d, joint available=%d",
self.full_attn_allocator.max_slots,
self.full_attn_allocator.min_slot_index,
self.full_attn_allocator.entry_bytes,
@@ -157,18 +178,6 @@ class UnifiedSWATokenToKVPoolAllocator(SWATokenToKVPoolAllocator):
# -- construction hooks (the tri-pool subclass overrides both) --
def _build_swa_attn_allocator(self, **kwargs) -> MultiEndedAllocator:
"""The swa sub-allocator: an END pool in the 2-pool pair."""
return MultiEndedAllocator(
sub_pool_name="swa",
is_id_owner=False, # non-owner; consumes virtuals minted by full
**kwargs,
)
def _wire_peers(self) -> None:
self.full_attn_allocator.bind_peer(self.swa_attn_allocator)
self.swa_attn_allocator.bind_peer(self.full_attn_allocator)
# -- capacity reporting (three-way split) --
def available_size(self) -> int:
@@ -179,46 +188,6 @@ class UnifiedSWATokenToKVPoolAllocator(SWATokenToKVPoolAllocator):
self._joint_avail_memo_epoch = epoch
return self._joint_avail_memo_tokens
def _compute_available_size(self) -> int:
"""Joint byte budget in TOKENS: each composite alloc(1) consumes one
full-side AND one swa-side page under the same virtual id."""
fa, sa = self.full_attn_allocator, self.swa_attn_allocator
e_f = fa.entry_bytes_per_page
e_s = sa.entry_bytes_per_page
# Direction-agnostic shared gap: the free byte band between the two pools.
if fa.grow_direction == "up":
gap_bytes = max(0, sa._byte_low_frontier() - fa._byte_high_frontier())
else:
gap_bytes = max(0, fa._byte_low_frontier() - sa._byte_high_frontier())
R_f = fa.num_pages - fa.min_page_index - fa._allocated_pages()
R_s = sa.num_pages - sa.min_page_index - sa._allocated_pages()
if not self.lazy_compaction:
pages_by_bytes = gap_bytes // (e_f + e_s)
return min(pages_by_bytes, R_f, R_s) * self.page_size
H_f = len(fa._free_phys_pages)
H_s = len(sa._free_phys_pages)
K1 = min(H_f, H_s) # Phase 1: both drain
# Phase 2: fewer-holes side extends; more-holes side keeps draining.
if H_f <= H_s:
e_phase2 = e_f
K_phase2_max = H_s
else:
e_phase2 = e_s
K_phase2_max = H_f
K2_room = K_phase2_max - K1
K2 = min(K2_room, gap_bytes // e_phase2) if e_phase2 > 0 else K2_room
gap_bytes -= K2 * e_phase2
K3 = gap_bytes // (e_f + e_s) # Phase 3: both extend
K_total = K1 + K2 + K3
K_total = min(K_total, H_f + R_f, H_s + R_s) # index-space caps
return K_total * self.page_size
# Slot-conservation views for the leak invariant only; the byte-coordinated
# value would flag spurious leaks. `allocated_count()` is in TOKENS.
def _conserve_full_available_size(self) -> int:
@@ -261,17 +230,16 @@ class UnifiedSWATokenToKVPoolAllocator(SWATokenToKVPoolAllocator):
def schedulable_swa_available_size(self) -> int:
return self.swa_attn_allocator.schedulable_available_size()
def _flush_targets(self):
"""Flush ALL members, including ones that are not short themselves: a
one-sided hole is unusable, and compacting it yields SHARED gap."""
return (self.full_attn_allocator, self.swa_attn_allocator)
# `size_full` / `size_swa` bound each side independently; current capacities
# also account for the peer's live byte usage.
def _ask_float_for_room(self, need_tokens: int) -> None:
"""No float in a two-END chain -- nothing can slide."""
return None
@property
def current_full_capacity(self) -> int:
return self.full_available_size() + self.full_attn_allocator.allocated_count()
# `size_full` / `size_swa` are inherited and read the static caps; reporting
# `max_slots - 1` here would be ~= full_max + swa_max and over-promise.
@property
def current_swa_capacity(self) -> int:
return self.swa_available_size() + self.swa_attn_allocator.allocated_count()
@property
def draft_virtual_id_space(self) -> int:
@@ -396,11 +364,8 @@ class UnifiedSWATokenToKVPoolAllocator(SWATokenToKVPoolAllocator):
def alloc(self, need_size: int) -> Optional[torch.Tensor]:
with record_function("UnifiedSWAAlloc.alloc"):
# Joint pre-check. Both sides are mutual peers (each side's compaction
# opens gap for the other), so flush BOTH on shortfall.
if need_size > self.available_size():
if not _relieve_for_alloc(self, need_size):
return None
if not self.ensure_capacity(need_size, need_size):
return None
# Snapshot the virtual PAGES full will consume, to bind them on swa too.
num_pages = need_size // self.page_size
fa = self.full_attn_allocator
@@ -437,9 +402,8 @@ class UnifiedSWATokenToKVPoolAllocator(SWATokenToKVPoolAllocator):
prefix_lens=prefix_lens_cpu,
)
need_tokens = num_new_pages * self.page_size
if need_tokens > self.available_size():
if not _relieve_for_alloc(self, need_tokens):
return None
if not self.ensure_capacity(need_tokens, need_tokens):
return None
# Snapshot the virtual PAGES the kernel will consume; clone so swa keeps
# its view after the slice is consumed.
@@ -555,9 +519,8 @@ class UnifiedSWATokenToKVPoolAllocator(SWATokenToKVPoolAllocator):
seq_lens=seq_lens_cpu, page_size=self.page_size, decode=True
)
need_tokens = num_new_pages * self.page_size
if need_tokens > self.available_size():
if not _relieve_for_alloc(self, need_tokens):
return None
if not self.ensure_capacity(need_tokens, need_tokens):
return None
fa = self.full_attn_allocator
new_virtual_pages = fa.free_virtual_ids[:num_new_pages].clone()
@@ -727,14 +690,6 @@ class UnifiedSWATokenToKVPoolAllocator(SWATokenToKVPoolAllocator):
self.full_attn_allocator.clear_inverse_history()
self.swa_attn_allocator.clear_inverse_history()
def verify_byte_accounting(self) -> List[str]:
return (
_chain_byte_accounting_violations(
_end_pair_chain(self.full_attn_allocator, self.swa_attn_allocator)
)
+ self._joint_capacity_memo_violations()
)
def _joint_capacity_memo_violations(self) -> List[str]:
"""Idle-time twin of `MultiEndedAllocator._capacity_memo_violations`
for the composite joint view. Empty == healthy."""
@@ -781,6 +736,385 @@ class UnifiedSWATokenToKVPoolAllocator(SWATokenToKVPoolAllocator):
forward_done, out_cache_loc_virtual
)
@abstractmethod
def _build_swa_attn_allocator(self, **kwargs) -> MultiEndedAllocator: ...
@abstractmethod
def _wire_peers(self) -> None: ...
@abstractmethod
def _compute_available_size(self) -> int: ...
@abstractmethod
def ensure_capacity(self, full_tokens: int, swa_tokens: int) -> bool: ...
@abstractmethod
def _flush_targets(self): ...
@abstractmethod
def _ask_float_for_room(self, need_tokens: int) -> None: ...
class UnifiedSWATokenToKVPoolAllocator(UnifiedSWAAllocatorBase):
"""Two-ended FULL/SWA allocator with asymmetric shared-byte reservations."""
def _build_swa_attn_allocator(self, **kwargs) -> MultiEndedAllocator:
"""The swa sub-allocator: an END pool in the 2-pool pair."""
return MultiEndedAllocator(
sub_pool_name="swa",
is_id_owner=False, # non-owner; consumes virtuals minted by full
**kwargs,
)
def _wire_peers(self) -> None:
self.full_attn_allocator.bind_peer(self.swa_attn_allocator)
self.swa_attn_allocator.bind_peer(self.full_attn_allocator)
def _compute_available_size(self) -> int:
"""Joint byte budget in TOKENS: each composite alloc(1) consumes one
full-side AND one swa-side page under the same virtual id."""
fa, sa = self.full_attn_allocator, self.swa_attn_allocator
e_f = fa.entry_bytes_per_page
e_s = sa.entry_bytes_per_page
# Direction-agnostic shared gap: the free byte band between the two pools.
if fa.grow_direction == "up":
gap_bytes = max(0, sa._byte_low_frontier() - fa._byte_high_frontier())
else:
gap_bytes = max(0, fa._byte_low_frontier() - sa._byte_high_frontier())
R_f = fa.num_pages - fa.min_page_index - fa._allocated_pages()
R_s = sa.num_pages - sa.min_page_index - sa._allocated_pages()
if not self.lazy_compaction:
pages_by_bytes = gap_bytes // (e_f + e_s)
return min(pages_by_bytes, R_f, R_s) * self.page_size
H_f = len(fa._free_phys_pages)
H_s = len(sa._free_phys_pages)
K1 = min(H_f, H_s) # Phase 1: both drain
# Phase 2: fewer-holes side extends; more-holes side keeps draining.
if H_f <= H_s:
e_phase2 = e_f
K_phase2_max = H_s
else:
e_phase2 = e_s
K_phase2_max = H_f
K2_room = K_phase2_max - K1
K2 = min(K2_room, gap_bytes // e_phase2) if e_phase2 > 0 else K2_room
gap_bytes -= K2 * e_phase2
K3 = gap_bytes // (e_f + e_s) # Phase 3: both extend
K_total = K1 + K2 + K3
K_total = min(K_total, H_f + R_f, H_s + R_s) # index-space caps
return K_total * self.page_size
def _flush_targets(self):
"""Flush ALL members, including ones that are not short themselves: a
one-sided hole is unusable, and compacting it yields SHARED gap."""
return (self.full_attn_allocator, self.swa_attn_allocator)
def _ask_float_for_room(self, need_tokens: int) -> None:
"""No float in a two-END chain -- nothing can slide."""
return None
def reclaim_plan(
self,
full_tokens: int | float,
swa_tokens: int | float,
*,
full_evictable_tokens: int = 0,
swa_evictable_tokens: int = 0,
empty_pool: bool = False,
) -> Optional[Tuple[int, int]]:
"""Return cumulative FULL/SWA eviction targets, or None if impossible.
The tree evicts FULL before SWA. Minimize required SWA reclaim with all
evictable FULL available, then trim excess FULL reclaim. FULL eviction's
actual SWA cascade is counted by the tree's shared eviction tracker.
"""
if full_tokens < 0 or swa_tokens < 0:
return None
page_size = self.page_size
full_pages = (math.ceil(full_tokens) + page_size - 1) // page_size
swa_pages = (math.ceil(swa_tokens) + page_size - 1) // page_size
# Restoring host SWA for device-resident FULL can require more new
# SWA pages than new FULL pages; only the shared budget constrains it.
compacted = empty_pool or not self.lazy_compaction or self._compaction_allowed()
def fits(full_reclaim_pages: int, swa_reclaim_pages: int) -> bool:
return self._fits_page_demand(
full_pages,
swa_pages,
full_reclaim_pages=full_reclaim_pages,
swa_reclaim_pages=swa_reclaim_pages,
compacted=compacted,
empty_pool=empty_pool,
)
if empty_pool:
return (0, 0) if fits(0, 0) else None
if fits(0, 0):
return (0, 0)
max_full_pages = min(
self.full_attn_allocator.allocated_count() // page_size,
max(0, int(full_evictable_tokens)) // page_size,
)
max_swa_pages = min(
self.swa_attn_allocator.allocated_count() // page_size,
max(0, int(swa_evictable_tokens)) // page_size,
)
if not fits(max_full_pages, max_swa_pages):
return None
def first_fit(high: int, predicate: Callable[[int], bool]) -> int:
low = 0
while low < high:
mid = (low + high) // 2
if predicate(mid):
high = mid
else:
low = mid + 1
return low
swa_reclaim_pages = first_fit(
max_swa_pages, lambda value: fits(max_full_pages, value)
)
full_reclaim_pages = first_fit(
max_full_pages, lambda value: fits(value, swa_reclaim_pages)
)
return (
full_reclaim_pages * page_size,
swa_reclaim_pages * page_size,
)
def can_reserve(
self,
full_tokens: int | float,
swa_tokens: int | float,
*,
full_evictable_tokens: int = 0,
swa_evictable_tokens: int = 0,
empty_pool: bool = False,
require_token_slack: bool = False,
) -> bool:
"""Check pending FULL/SWA demand against the shared byte envelope.
Scheduler admission keeps the historical one-token strict slack at an
empty-pool boundary. Live admission checks the state reachable after
reclaiming the currently evictable pages from both sides.
"""
if full_tokens < 0 or swa_tokens < 0:
return False
if require_token_slack and (
full_tokens >= self.size_full or swa_tokens >= self.size_swa
):
return False
return (
self.reclaim_plan(
full_tokens,
swa_tokens,
full_evictable_tokens=full_evictable_tokens,
swa_evictable_tokens=swa_evictable_tokens,
empty_pool=empty_pool,
)
is not None
)
def _compaction_allowed(self) -> bool:
return all(
allocator.disagg_move_gate is None or allocator.disagg_move_gate()
for allocator in (self.full_attn_allocator, self.swa_attn_allocator)
)
def _fits_page_demand(
self,
num_full_pages: int,
num_swa_pages: int,
*,
full_reclaim_pages: int = 0,
swa_reclaim_pages: int = 0,
compacted: bool,
empty_pool: bool = False,
) -> bool:
"""Check one FULL/SWA page demand against a single allocator snapshot."""
if min(num_full_pages, num_swa_pages) < 0:
return False
fa, sa = self.full_attn_allocator, self.swa_attn_allocator
if empty_pool:
full_live_pages = swa_live_pages = 0
full_reclaim_pages = swa_reclaim_pages = 0
else:
full_live_pages = fa.allocated_count() // self.page_size
swa_live_pages = sa.allocated_count() // self.page_size
full_reclaim_pages = min(full_live_pages, max(0, int(full_reclaim_pages)))
swa_reclaim_pages = min(swa_live_pages, max(0, int(swa_reclaim_pages)))
full_live_pages -= full_reclaim_pages
swa_live_pages -= swa_reclaim_pages
full_total_pages = full_live_pages + num_full_pages
swa_total_pages = swa_live_pages + num_swa_pages
virtual_page_capacity = fa.num_virtual_ids - fa.min_page_index
full_page_capacity = min(
virtual_page_capacity,
fa.num_pages - fa.min_page_index,
)
swa_page_capacity = min(
virtual_page_capacity,
sa.num_pages - sa.min_page_index,
)
if full_total_pages > full_page_capacity or swa_total_pages > swa_page_capacity:
return False
if compacted:
full_virtual_room = virtual_page_capacity - full_live_pages
full_holes = swa_holes = 0
full_index_room = full_page_capacity - full_live_pages
swa_index_room = swa_page_capacity - swa_live_pages
gap_bytes = (
self._empty_shared_gap_bytes
- full_live_pages * fa.entry_bytes_per_page
- swa_live_pages * sa.entry_bytes_per_page
)
else:
full_virtual_room = len(fa.free_virtual_ids) + full_reclaim_pages
full_holes = len(fa._free_phys_pages) + full_reclaim_pages
swa_holes = len(sa._free_phys_pages) + swa_reclaim_pages
full_index_room = fa.num_pages - fa.min_page_index - fa._allocated_pages()
swa_index_room = sa.num_pages - sa.min_page_index - sa._allocated_pages()
gap_bytes = fa._current_gap_bytes()
if num_full_pages > full_virtual_room:
return False
if num_full_pages > full_holes + full_index_room:
return False
if num_swa_pages > swa_holes + swa_index_room:
return False
full_extensions = max(0, num_full_pages - full_holes)
swa_extensions = max(0, num_swa_pages - swa_holes)
return (
full_extensions * fa.entry_bytes_per_page
+ swa_extensions * sa.entry_bytes_per_page
<= max(0, gap_bytes)
)
def ensure_capacity(self, full_tokens: int, swa_tokens: int) -> bool:
"""Gate one allocation and compact both sides on shortfall."""
if full_tokens < 0 or swa_tokens < 0:
return False
if full_tokens == 0 and swa_tokens == 0:
return True
page_size = self.page_size
num_full_pages = (int(full_tokens) + page_size - 1) // page_size
num_swa_pages = (int(swa_tokens) + page_size - 1) // page_size
if self._fits_page_demand(
num_full_pages,
num_swa_pages,
compacted=False,
):
return True
if not self.lazy_compaction or not self._compaction_allowed():
return False
if not self._fits_page_demand(
num_full_pages,
num_swa_pages,
compacted=True,
):
return False
self.full_attn_allocator.flush_for_allocation()
self.swa_attn_allocator.flush_for_allocation()
return self._fits_page_demand(
num_full_pages,
num_swa_pages,
compacted=False,
)
def create_prefill_budget(self, tree_cache, *, num_mixed_decode_tokens=0):
from sglang.srt.mem_cache.prefill_budget import SharedSWAPrefillBudget
return SharedSWAPrefillBudget(
self, tree_cache, num_mixed_decode_tokens=num_mixed_decode_tokens
)
def swa_capacity_and_available(self, *, full_capacity, swa_capacity):
return (
(self.current_full_capacity, self.full_available_size()),
(self.current_swa_capacity, self.swa_available_size()),
)
def max_new_tokens_for_memory(
self,
input_tokens: int,
max_new_tokens: int,
*,
token_capacity: int,
sliding_window_size: int | None,
chunk_size: int | None,
) -> int | None:
from sglang.srt.mem_cache.prefill_budget import estimate_swa_kv_tokens
def fits(candidate):
return self.can_reserve(
input_tokens + candidate + self.page_size,
estimate_swa_kv_tokens(
input_tokens,
candidate,
sliding_window_size=sliding_window_size,
page_size=self.page_size,
allocation_limit=chunk_size,
),
empty_pool=True,
require_token_slack=True,
)
if not fits(0):
return None
if fits(max_new_tokens):
return max_new_tokens
lo, hi = 0, max_new_tokens
while lo < hi:
mid = (lo + hi + 1) // 2
if fits(mid):
lo = mid
else:
hi = mid - 1
return lo
def evict_to_free_tokens(self, tree_cache, num_tokens: int) -> bool | None:
from sglang.srt.mem_cache.base_prefix_cache import EvictParams
if tree_cache is None or tree_cache.is_chunk_cache():
return
reclaim_plan = self.reclaim_plan(
num_tokens,
num_tokens,
full_evictable_tokens=tree_cache.full_evictable_size(),
swa_evictable_tokens=tree_cache.swa_evictable_size(),
)
if reclaim_plan is None:
return
full_reclaim, swa_reclaim = reclaim_plan
if full_reclaim or swa_reclaim:
tree_cache.evict_for_alloc(
EvictParams(num_tokens=full_reclaim, swa_num_tokens=swa_reclaim)
)
# A zero-reclaim plan can still depend on compaction before allocation.
return self.ensure_capacity(num_tokens, num_tokens)
def verify_byte_accounting(self) -> List[str]:
return (
_chain_byte_accounting_violations(
_end_pair_chain(self.full_attn_allocator, self.swa_attn_allocator)
)
+ self._joint_capacity_memo_violations()
)
def flush_opportunistic(self) -> int:
"""Non-urgent flush of BOTH sub-allocators; sync-free."""
with record_function("UnifiedSWAAlloc.flush_opportunistic"):
@@ -796,7 +1130,7 @@ class UnifiedSWATokenToKVPoolAllocator(SWATokenToKVPoolAllocator):
return fa.flush_opportunistic() + sa.flush_opportunistic()
class UnifiedMambaSWATokenToKVPoolAllocator(UnifiedSWATokenToKVPoolAllocator):
class UnifiedMambaSWATokenToKVPoolAllocator(UnifiedSWAAllocatorBase):
"""Tri-pool composite for models with full KV + SWA KV + mamba/conv state
(both `mambaish_config` and `is_hybrid_swa`).
@@ -887,6 +1221,37 @@ class UnifiedMambaSWATokenToKVPoolAllocator(UnifiedSWATokenToKVPoolAllocator):
# -- capacity --
def can_reserve(
self,
full_tokens: int | float,
swa_tokens: int | float,
*,
full_evictable_tokens: int = 0,
swa_evictable_tokens: int = 0,
empty_pool: bool = False,
require_token_slack: bool = False,
) -> bool:
if (
full_tokens < 0
or swa_tokens < 0
or full_tokens != swa_tokens
or full_evictable_tokens
or swa_evictable_tokens
or empty_pool
):
return False
return full_tokens <= self.available_size()
def ensure_capacity(self, full_tokens: int, swa_tokens: int) -> bool:
if full_tokens < 0 or swa_tokens < 0 or full_tokens != swa_tokens:
return False
if full_tokens == 0:
return True
need_tokens = int(full_tokens)
if need_tokens <= self.available_size():
return True
return _relieve_for_alloc(self, need_tokens)
def _compute_available_size(self) -> int:
"""Joint TOKENS for `alloc(N)`: N costs N full pages AND N swa pages, drawn
from DIFFERENT bands -- full extends only into the high band, the float into
@@ -1036,17 +1401,33 @@ class UnifiedMambaSWATokenToKVPoolAllocator(UnifiedSWATokenToKVPoolAllocator):
super().set_inflight_forward(forward_done, out_cache_loc_virtual)
self.mamba_allocator.set_inflight_forward(forward_done, None)
def create_prefill_budget(self, tree_cache, *, num_mixed_decode_tokens=0):
# Mamba competes for the shared gap too; retain the tri-pool's existing
# token and state-slot admission until it has a three-way reservation.
return SWATokenToKVPoolAllocator.create_prefill_budget(
self, tree_cache, num_mixed_decode_tokens=num_mixed_decode_tokens
)
def max_new_tokens_for_memory(self, *args, **kwargs):
return BaseTokenToKVPoolAllocator.max_new_tokens_for_memory(
self, *args, **kwargs
)
def swa_capacity_and_available(self, *, full_capacity, swa_capacity):
return (
(full_capacity, self.conserve_full_available_size()),
(swa_capacity, self.conserve_swa_available_size()),
)
def evict_to_free_tokens(self, tree_cache, num_tokens: int) -> None:
"""Joint-aware eviction: one tri-lifetime node frees bytes on several sides
at once, so re-check the JOINT gate instead of the per-side shortfall."""
from sglang.srt.mem_cache.common import evict_from_tree_cache
# Arbitrary retry bound; a round that frees nothing ends the loop anyway.
for _ in range(4):
before = self.available_size()
if before >= num_tokens:
return
evict_from_tree_cache(tree_cache, num_tokens)
SWATokenToKVPoolAllocator.evict_to_free_tokens(self, tree_cache, num_tokens)
if self.available_size() <= before:
return # no progress
python/sglang/srt/mem_cache/allocator/unified_sub_pool.py
+5
View File
@@ -345,9 +345,13 @@ class MultiEndedAllocator(BaseTokenToKVPoolAllocator):
# v2p is indexed by VIRTUAL page id, p2v by PHYSICAL page id. A non-owner
# consumes the owner's ids, so the two counts are unrelated.
assert virtual_num_pages is None or not is_id_owner, (
"only a non-owner allocator may use another pool's virtual-id space"
)
self.num_virtual_ids = (
self.num_pages if virtual_num_pages is None else virtual_num_pages
)
assert self.num_virtual_ids > 0, "virtual page count must be positive"
# Page 0 is the padding anchor; the trailing row is the -1 sentinel.
self.virtual_to_physical = torch.full(
(self.num_virtual_ids + 1,),
@@ -554,6 +558,7 @@ class MultiEndedAllocator(BaseTokenToKVPoolAllocator):
f"is_id_owner={self.is_id_owner}, page_size={self.page_size}, "
f"min_page_index={self.min_page_index}, "
f"num_pages={self.num_pages}, "
f"num_virtual_ids={self.num_virtual_ids}, "
f"watermark_physical={self.watermark_physical}, "
f"allocated_pages={self._allocated_pages()}"
)
python/sglang/srt/mem_cache/common.py
+7 -29
View File
@@ -11,7 +11,6 @@ from sglang.kernels.ops.memory.common import (
)
from sglang.kernels.ops.memory.common import get_last_loc_kernel as get_last_loc_kernel
from sglang.srt.mem_cache.allocator.page_interleave import page_interleave_shard_size
from sglang.srt.mem_cache.allocator.swa import SWATokenToKVPoolAllocator
from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache, EvictParams
from sglang.srt.mem_cache.hicache_storage import PoolTransfer
from sglang.srt.mem_cache.memory_pool import HybridReqToTokenPool, ReqToTokenPool
@@ -161,34 +160,13 @@ def maybe_cache_unfinished_req(req: Req, tree_cache: BasePrefixCache, **kwargs):
tree_cache.cache_unfinished_req(req, **kwargs)
def evict_from_tree_cache(tree_cache: BasePrefixCache | None, num_tokens: int):
if tree_cache is None:
return
if tree_cache.is_chunk_cache():
return
allocator = tree_cache.token_to_kv_pool_allocator
if isinstance(allocator, SWATokenToKVPoolAllocator):
# Hybrid allocator
full_available_size = allocator.full_available_size()
swa_available_size = allocator.swa_available_size()
if full_available_size < num_tokens or swa_available_size < num_tokens:
full_num_tokens = max(0, num_tokens - full_available_size)
swa_num_tokens = max(0, num_tokens - swa_available_size)
tree_cache.evict_for_alloc(
EvictParams(num_tokens=full_num_tokens, swa_num_tokens=swa_num_tokens)
)
else:
# Standard allocator: evict only the shortfall (mirrors the SWA arm)
available_size = allocator.available_size()
if available_size < num_tokens:
tree_cache.evict_for_alloc(
EvictParams(num_tokens=num_tokens - available_size)
)
_evict_until_allocatable(tree_cache, allocator, num_tokens)
def evict_from_tree_cache(
tree_cache: BasePrefixCache | None, num_tokens: int
) -> bool | None:
if tree_cache is not None and not tree_cache.is_chunk_cache():
return tree_cache.token_to_kv_pool_allocator.evict_to_free_tokens(
tree_cache, num_tokens
)
def _evict_until_allocatable(
python/sglang/srt/mem_cache/kv_cache_configurator.py
+39 -12
View File
@@ -51,7 +51,7 @@ from sglang.srt.mem_cache.allocator.swa import (
is_swa_req_ring,
)
from sglang.srt.mem_cache.allocator.unified_hybrid_swa import (
UnifiedSWATokenToKVPoolAllocator,
UnifiedSWAAllocatorBase,
)
from sglang.srt.mem_cache.allocator.unified_mamba import (
UnifiedMambaTokenToKVPoolAllocator,
@@ -240,6 +240,7 @@ class _PoolSizes(msgspec.Struct, frozen=True, kw_only=True):
c128_state_pool_size: int
c4_state_dtype: Optional[torch.dtype]
c128_state_dtype: Optional[torch.dtype]
unified_memory_pool_bytes: Optional[int] = None
unified_total_bytes: Optional[int] = None
@@ -289,6 +290,23 @@ class KVCacheConfigurator:
self.draft_model_idx in self.model_config.swa_attention_layer_ids
)
def hybrid_swa_token_capacity(
self,
*,
allocator: BaseTokenToKVPoolAllocator,
full_capacity: Optional[int],
swa_capacity: Optional[int],
) -> int:
if get_memory().enable_unified_memory:
capacity = allocator.size_full
max_total_tokens = get_schedule().max_total_tokens
return (
min(capacity, max_total_tokens)
if max_total_tokens is not None
else capacity
)
return full_capacity or swa_capacity
def _build_fp4_quant_method(self, *, num_layers: int):
if not is_float4_e2m1fn_x2(self.kv_cache_dtype):
return None
@@ -423,6 +441,10 @@ class KVCacheConfigurator:
max_running_requests=max_running_requests,
full_max_total_num_tokens=full_max_total_num_tokens,
swa_max_total_num_tokens=swa_max_total_num_tokens,
# The target's byte envelope excludes the separate draft allocation.
unified_memory_pool_bytes=(
None if self.is_draft_worker else config.unified_memory_pool_bytes
),
c4_max_total_num_tokens=c4_max_total_num_tokens,
c128_max_total_num_tokens=c128_max_total_num_tokens,
c4_state_pool_size=c4_state_pool_size,
@@ -469,6 +491,7 @@ class KVCacheConfigurator:
max_num_reqs=sizes.max_running_requests,
full_max_total_num_tokens=sizes.full_max_total_num_tokens,
swa_max_total_num_tokens=sizes.swa_max_total_num_tokens,
unified_memory_pool_bytes=sizes.unified_memory_pool_bytes,
unified_total_bytes=sizes.unified_total_bytes,
)
else:
@@ -497,7 +520,7 @@ class KVCacheConfigurator:
token_to_kv_pool_allocator,
(
UnifiedMambaTokenToKVPoolAllocator,
UnifiedSWATokenToKVPoolAllocator,
UnifiedSWAAllocatorBase,
),
):
draft_virtual_id_space = (
@@ -520,7 +543,7 @@ class KVCacheConfigurator:
if (
isinstance(
token_to_kv_pool_allocator,
UnifiedSWATokenToKVPoolAllocator,
UnifiedSWAAllocatorBase,
)
and self.is_hybrid_swa
):
@@ -821,8 +844,9 @@ class KVCacheConfigurator:
self,
*,
max_num_reqs: int,
full_max_total_num_tokens: Optional[int],
swa_max_total_num_tokens: Optional[int],
full_max_total_num_tokens: Optional[int] = None,
swa_max_total_num_tokens: Optional[int] = None,
unified_memory_pool_bytes: Optional[int] = None,
unified_total_bytes: Optional[int] = None,
) -> UnifiedPoolBundle:
"""Build the unified-pool stack for a hybrid-SWA model (Triton): one byte
@@ -899,6 +923,12 @@ class KVCacheConfigurator:
if self.layer_info.start_layer <= i < self.layer_info.end_layer
]
total_bytes = unified_memory_pool_bytes
# An uncapped, draft-free pool owns the profiled budget, including bytes
# left over after rounding the FULL/SWA boot capacities to pages.
if unified_total_bytes is not None and self.spec_algorithm.is_none():
total_bytes = unified_total_bytes
bundle = init_unified_swa_pools(
device=self.device,
kv_cache_dtype=self.kv_cache_dtype,
@@ -915,6 +945,7 @@ class KVCacheConfigurator:
full_attention_layer_ids=full_attention_layer_ids,
full_max_total_num_tokens=full_max_total_num_tokens,
swa_max_total_num_tokens=swa_max_total_num_tokens,
total_bytes=total_bytes,
enable_memory_saver=get_exec().features.enable_memory_saver,
need_sort=get_disagg().disaggregation_mode in ("decode", "prefill"),
# Overlap mode: same wait_stream(forward_stream) rationale as
@@ -922,9 +953,6 @@ class KVCacheConfigurator:
forward_stream=self.forward_stream,
# Lazy compaction: default ON, with env var escape hatch for rollback / A/B.
lazy_compaction=_should_enable_lazy_compaction(),
# Draft workers keep the token-count byte sum (spec is asserted
# off under unified; belt only).
unified_total_bytes=(None if self.is_draft_worker else unified_total_bytes),
# bs=1 feasibility floor inputs. `model_context_len` bounds the
# sliding window term only -- the full-attention side is not
# charged, see `_check_bs1_feasibility_floor`.
@@ -2120,7 +2148,7 @@ class KVCacheConfigurator:
else:
swa_allocator = token_to_kv_pool_allocator
uses_unified_virtual_ids = isinstance(
swa_allocator, UnifiedSWATokenToKVPoolAllocator
swa_allocator, UnifiedSWAAllocatorBase
)
has_draft_swa_layers = (
not self.is_hybrid_swa_mtp_draft or self.draft_swa_full_capacity
@@ -2363,9 +2391,8 @@ class KVCacheConfigurator:
f"{config.max_total_num_tokens}"
)
if max_tokens != config.max_total_num_tokens:
# Token-capped re-derivation: the profiled budget no longer
# applies; the recalced config's unified_total_bytes stays None
# and the factories fall back to the token-count byte sum.
# Re-derive the capped budget: SWA carries unified_memory_pool_bytes;
# Mamba factories fall back to token-count sizing without unified_total_bytes.
config = configurator.calculate_pool_sizes_from_max_tokens(
max_tokens, get_schedule().page_size
)
python/sglang/srt/mem_cache/kv_index_translator.py
+14 -3
View File
@@ -65,7 +65,7 @@ from sglang.kernels.ops.kvcache.kv_read_table import (
build_kv_read_table_packed,
)
from sglang.srt.mem_cache.allocator.unified_hybrid_swa import (
UnifiedSWATokenToKVPoolAllocator,
UnifiedSWAAllocatorBase,
)
from sglang.srt.mem_cache.allocator.unified_mamba import (
UnifiedMambaTokenToKVPoolAllocator,
@@ -121,12 +121,13 @@ class KVIndexTranslator:
self.is_translating = (
isinstance(
token_to_kv_pool_allocator,
(UnifiedMambaTokenToKVPoolAllocator, UnifiedSWATokenToKVPoolAllocator),
(UnifiedMambaTokenToKVPoolAllocator, UnifiedSWAAllocatorBase),
)
and token_to_kv_pool_allocator.get_kvcache() is token_to_kv_pool
)
if self.is_translating:
alloc = token_to_kv_pool_allocator
self._capture_page_size = alloc.page_size
self._full_v2p_table = alloc.full_v2p_page_table
self._full_p2v_table = alloc.full_p2v_page_table
self._full_page_multiplier = alloc.kernel_page_multiplier
@@ -139,7 +140,7 @@ class KVIndexTranslator:
# DCP read ids stay WIDENED to the consumer: selecting this rank's
# share changes the length, so only the production site can do it.
self.defer_read_translate = get_parallel().attn_dcp_size > 1
if isinstance(alloc, UnifiedSWATokenToKVPoolAllocator):
if isinstance(alloc, UnifiedSWAAllocatorBase):
self._swa_v2p_table = alloc.swa_v2p_page_table
self._swa_page_multiplier = alloc.swa_kernel_page_multiplier
self._swa_write_loc_from_full = self._swa_write_loc_unified
@@ -171,6 +172,16 @@ class KVIndexTranslator:
)
self._index_table_memo: Optional[Tuple[weakref.ref, KVIndexTable]] = None
def capture_token_capacity(self, max_token_pool_size: int) -> int:
"""Host capture rows are indexed by request-token IDs, not kernel IDs.
Unified IDs span the whole virtual table even when admission is capped.
DCP widens allocator pages; the runner's page size stays physical.
"""
if self.is_translating:
return self._full_v2p_table.numel() * self._capture_page_size
return max_token_pool_size + self.page_size
# -- per-batch view --------------------------------------------------------
@property
python/sglang/srt/mem_cache/prefill_budget.py
+413
View File
@@ -0,0 +1,413 @@
# Copyright 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.
"""Memory reservations for one prefill pass.
The scheduler supplies token demand and its chunk/decode limits. These objects
account for admitted but not yet allocated work and query live cache capacity:
locking a prefix or preempting a request must affect the next admission check.
They neither select requests nor mutate the prefix cache or allocator.
"""
from typing import Optional
from sglang.srt.mem_cache.allocator.swa import is_swa_req_ring
def estimate_swa_kv_tokens(
extend_input_len: int,
max_new_tokens: int,
*,
sliding_window_size: Optional[int],
page_size: int,
allocation_limit: Optional[int] = None,
host_hit_length: int = 0,
) -> int:
"""Peak SWA reservation for one prefill/decode request."""
if sliding_window_size is None or sliding_window_size <= 0:
reserved = extend_input_len + max_new_tokens + page_size
else:
allocated = (
extend_input_len
if allocation_limit is None
else min(extend_input_len, allocation_limit)
)
allocated_tail = max(allocated - sliding_window_size, 0)
# With a roughly two-window SWA pool, a cached prefix can already lock
# one window. Charging another whole window for a short resume can then
# block admission forever on an idle pool. Reserve only its uncached
# tail plus decode headroom, capped at the window.
# Including the extension also keeps the reservation large enough for
# this pass's prefill allocation.
reserved = (
allocated_tail
+ min(extend_input_len + max_new_tokens, sliding_window_size)
+ page_size
)
if host_hit_length > 0:
reserved += -(-host_hit_length // page_size) * page_size
return reserved
class PrefillBudget:
"""Fixed token pool. Offsets include pending allocations and decode headroom."""
def __init__(self, allocator, tree_cache, *, num_mixed_decode_tokens: int = 0):
self.allocator = allocator
self.tree_cache = tree_cache
self.page_size = allocator.page_size
self.total_offset = num_mixed_decode_tokens
self.current_offset = num_mixed_decode_tokens
self.swa_offset = 0
def ceil_paged_tokens(self, tokens: int) -> int:
return -(-tokens // self.page_size) * self.page_size
def _available_and_evictable(self):
evictable = (
self.tree_cache.full_evictable_size()
if self.tree_cache.supports_mamba()
else self.tree_cache.evictable_size()
)
return self.allocator.available_size() + evictable
@property
def remaining_total(self):
return self._available_and_evictable() - self.total_offset
@property
def remaining_current(self):
return self._available_and_evictable() - self.current_offset
@property
def remaining_swa(self):
return 0
def has_capacity(self) -> bool:
return self.remaining_total > 0 and self.remaining_current > 0
def check_prefill(
self,
*,
extend_input_len: int,
total_tokens: int,
max_new_tokens: int,
input_tokens: int,
swa_host_hit_length: int,
chunk_limit: int | None,
) -> tuple[bool, int | None]:
"""Return admission feasibility and a memory bound on the requested chunk."""
return (
(True, chunk_limit)
if total_tokens < self.remaining_total
else (False, None)
)
def can_allocate_prefill(
self,
*,
paged_input: int,
extend_input_len: int,
max_new_tokens: int,
chunk_limit: int | None,
) -> bool:
return paged_input <= min(self.remaining_current, self.remaining_total)
def available_chunk_tokens(self, chunk_limit: int) -> int | None:
available = min(chunk_limit, int(self.remaining_total))
# Single-pool continuation must make progress to release its KV.
return available if available > 0 else chunk_limit
def fit_chunk(
self,
*,
extend_input_len: int,
max_new_tokens: int,
chunk_limit: int,
) -> int | None:
return chunk_limit
def reserve(
self,
extend_input_len: int,
max_new_tokens: int,
*,
extra_tokens: int = 0,
chunk_limit: int | None = None,
is_chunked_continuation: bool = False,
) -> None:
extend_input_len = self.ceil_paged_tokens(extend_input_len)
immediate = extend_input_len + self.page_size + extra_tokens
self.total_offset += immediate + max_new_tokens
self.current_offset += immediate
class SWAPrefillBudget(PrefillBudget):
"""Separate FULL/SWA partitions, including per-request SWA rings."""
def __init__(self, *args, all_swa=False, **kwargs):
super().__init__(*args, **kwargs)
self.all_swa = all_swa
self.req_ring = is_swa_req_ring(self.allocator)
def _available_and_evictable(self):
if self.all_swa:
return (
self.allocator.swa_available_size()
+ self.tree_cache.swa_evictable_size()
)
return (
self.allocator.full_available_size() + self.tree_cache.full_evictable_size()
)
@property
def remaining_swa(self):
evictable = 0 if self.req_ring else self.tree_cache.swa_evictable_size()
return self.allocator.swa_available_size() + evictable - self.swa_offset
def swa_tokens(
self,
extend_input_len,
max_new_tokens,
*,
chunk_limit=None,
swa_host_hit_length=0,
):
if self.req_ring:
return self.allocator.swa_ring_cost_tokens
return estimate_swa_kv_tokens(
extend_input_len,
max_new_tokens,
sliding_window_size=self.tree_cache.sliding_window_size,
page_size=self.page_size,
allocation_limit=chunk_limit,
host_hit_length=swa_host_hit_length,
)
def swa_never_fits(self, extend_input_len, max_new_tokens, **kwargs):
needed = self.swa_tokens(extend_input_len, max_new_tokens, **kwargs)
return (
needed > self.allocator.size_swa
if self.req_ring
else needed >= self.allocator.size_swa
)
def _chunk_cap(self, max_new_tokens, swa_host_hit_length=0):
# Only the sliding window stays locked between chunks, so a smaller
# chunk can bound the transient SWA footprint of a longer prompt.
headroom = self.swa_tokens(
0, max_new_tokens, swa_host_hit_length=swa_host_hit_length
)
cap = int(self.remaining_swa) - headroom
return max(0, cap // self.page_size * self.page_size)
def check_prefill(
self,
*,
extend_input_len: int,
total_tokens: int,
max_new_tokens: int,
input_tokens: int,
swa_host_hit_length: int,
chunk_limit: int | None,
) -> tuple[bool, int | None]:
if total_tokens >= self.remaining_total:
return False, None
extend_input_len = self.ceil_paged_tokens(extend_input_len)
needed = self.swa_tokens(
extend_input_len,
max_new_tokens,
chunk_limit=chunk_limit,
swa_host_hit_length=swa_host_hit_length,
)
fits = (
needed <= self.remaining_swa
if self.req_ring
else needed < self.remaining_swa
)
if fits:
return True, chunk_limit
# Only permanent shortfalls may shrink a chunk. Transient pressure waits
# so a new prefill does not consume running decodes' window headroom.
cap = 0
if self.swa_never_fits(
extend_input_len,
max_new_tokens,
chunk_limit=chunk_limit,
swa_host_hit_length=swa_host_hit_length,
):
cap = self._chunk_cap(max_new_tokens, swa_host_hit_length)
if chunk_limit is None or cap <= 0:
return False, None
return True, min(chunk_limit, cap)
def has_capacity(self) -> bool:
return super().has_capacity() and self.remaining_swa > 0
def available_chunk_tokens(self, chunk_limit: int) -> int | None:
available = min(chunk_limit, int(self.remaining_total))
if not self.req_ring:
available = min(available, int(self.remaining_swa) - self.page_size)
return available if available > 0 else None
def can_allocate_prefill(
self,
*,
paged_input: int,
extend_input_len: int,
max_new_tokens: int,
chunk_limit: int | None,
) -> bool:
return (
super().can_allocate_prefill(
paged_input=paged_input,
extend_input_len=extend_input_len,
max_new_tokens=max_new_tokens,
chunk_limit=chunk_limit,
)
and self.swa_tokens(
extend_input_len, max_new_tokens, chunk_limit=chunk_limit
)
<= self.remaining_swa
)
def reserve(
self,
extend_input_len: int,
max_new_tokens: int,
*,
extra_tokens: int = 0,
chunk_limit: int | None = None,
is_chunked_continuation: bool = False,
) -> None:
super().reserve(extend_input_len, max_new_tokens, extra_tokens=extra_tokens)
# A continuation already owns its ring slot.
if not (self.req_ring and is_chunked_continuation):
self.swa_offset += self.swa_tokens(
self.ceil_paged_tokens(extend_input_len),
max_new_tokens,
chunk_limit=chunk_limit,
)
class SharedSWAPrefillBudget(SWAPrefillBudget):
"""FULL and SWA reservations compete for the same physical byte budget."""
def __init__(self, *args, num_mixed_decode_tokens=0, **kwargs):
super().__init__(
*args, num_mixed_decode_tokens=num_mixed_decode_tokens, **kwargs
)
self.swa_offset = num_mixed_decode_tokens
def _fits(self, full_tokens, swa_tokens, *, empty_pool=False):
return self.allocator.can_reserve(
full_tokens + (0 if empty_pool else self.total_offset),
swa_tokens + (0 if empty_pool else self.swa_offset),
full_evictable_tokens=0
if empty_pool
else self.tree_cache.full_evictable_size(),
swa_evictable_tokens=0
if empty_pool
else self.tree_cache.swa_evictable_size(),
empty_pool=empty_pool,
require_token_slack=empty_pool,
)
def _joint_chunk_cap(self, *, max_chunk_tokens, chunk_limit, swa_host_hit_length=0):
lo, hi = 0, max(0, max_chunk_tokens) // self.page_size
while lo < hi:
mid = (lo + hi + 1) // 2
tokens = mid * self.page_size
if self._fits(
tokens + self.page_size,
self.swa_tokens(
tokens,
0,
chunk_limit=chunk_limit,
swa_host_hit_length=swa_host_hit_length,
),
):
lo = mid
else:
hi = mid - 1
return lo * self.page_size
def check_prefill(
self,
*,
extend_input_len: int,
total_tokens: int,
max_new_tokens: int,
input_tokens: int,
swa_host_hit_length: int,
chunk_limit: int | None,
) -> tuple[bool, int | None]:
needed = self.swa_tokens(
extend_input_len,
max_new_tokens,
chunk_limit=chunk_limit,
swa_host_hit_length=swa_host_hit_length,
)
if self._fits(total_tokens, needed):
return True, chunk_limit
if chunk_limit is None or self._fits(
input_tokens + max_new_tokens + self.page_size,
self.swa_tokens(input_tokens, max_new_tokens, chunk_limit=chunk_limit),
empty_pool=True,
):
return False, None
cap = self._joint_chunk_cap(
max_chunk_tokens=min(chunk_limit, max(0, extend_input_len - 1)),
chunk_limit=chunk_limit,
swa_host_hit_length=swa_host_hit_length,
)
return (True, min(chunk_limit, cap)) if cap > 0 else (False, None)
def has_capacity(self) -> bool:
return self._fits(0, 0)
def available_chunk_tokens(self, chunk_limit: int) -> int | None:
return chunk_limit if chunk_limit > 0 else None
def fit_chunk(
self,
*,
extend_input_len: int,
max_new_tokens: int,
chunk_limit: int,
) -> int | None:
candidate = min(extend_input_len, chunk_limit)
finishes = candidate >= extend_input_len
headroom = max_new_tokens if finishes else 0
if self._fits(
candidate + headroom + self.page_size,
self.swa_tokens(candidate, headroom, chunk_limit=chunk_limit),
):
return chunk_limit
cap = self._joint_chunk_cap(
max_chunk_tokens=max(0, candidate - 1) if finishes else candidate,
chunk_limit=chunk_limit,
)
return min(chunk_limit, cap) if cap > 0 else None
def can_allocate_prefill(
self,
*,
paged_input: int,
extend_input_len: int,
max_new_tokens: int,
chunk_limit: int | None,
) -> bool:
return self._fits(
extend_input_len + max_new_tokens + self.page_size,
self.swa_tokens(extend_input_len, max_new_tokens, chunk_limit=chunk_limit),
)
python/sglang/srt/mem_cache/unified_cache/components/swa.py
+2 -2
View File
@@ -226,11 +226,11 @@ class SWAComponent(TreeComponent):
def _unified_allocator(self):
"""The unified SWA composite, or None when running on the static pool."""
from sglang.srt.mem_cache.allocator.unified_hybrid_swa import (
UnifiedSWATokenToKVPoolAllocator,
UnifiedSWAAllocatorBase,
)
allocator = self.cache.token_to_kv_pool_allocator
if isinstance(allocator, UnifiedSWATokenToKVPoolAllocator):
if isinstance(allocator, UnifiedSWAAllocatorBase):
return allocator
return None
python/sglang/srt/mem_cache/unified_memory_pool.py
+20 -14
View File
@@ -1713,13 +1713,13 @@ def init_unified_swa_pools(
end_layer: int,
swa_attention_layer_ids: List[int],
full_attention_layer_ids: List[int],
full_max_total_num_tokens: int,
swa_max_total_num_tokens: int,
full_max_total_num_tokens: Optional[int] = None,
swa_max_total_num_tokens: Optional[int] = None,
total_bytes: Optional[int] = None,
enable_memory_saver: bool,
need_sort: bool,
forward_stream: Optional[torch.cuda.Stream] = None,
lazy_compaction: bool = False,
unified_total_bytes: Optional[int] = None,
model_context_len: Optional[int] = None,
sliding_window_size: Optional[int] = None,
) -> UnifiedSWAPoolBundle:
@@ -1758,15 +1758,20 @@ def init_unified_swa_pools(
store_dtype=store_dtype,
grow_direction="up",
)
if unified_total_bytes is not None:
# PROFILED byte budget, sized from directly: the re-sum's floor losses
# stay out of the buffer, and the token counts remain boot labels.
total_bytes = unified_total_bytes
else:
legacy_allocator_capacities = {}
if total_bytes is None:
if full_max_total_num_tokens is None or swa_max_total_num_tokens is None:
raise ValueError(
"total_bytes or both legacy full/SWA capacities must be provided"
)
total_bytes = (
full_max_total_num_tokens * full_spec.entry_bytes()
+ swa_max_total_num_tokens * swa_spec.entry_bytes()
)
legacy_allocator_capacities = {
"full_max_total_num_tokens": full_max_total_num_tokens,
"swa_max_total_num_tokens": swa_max_total_num_tokens,
}
if model_context_len is not None:
# bs=1 floor: ONE sliding window of swa KV (+ a page of slack for the
# page-granular walk) + the slot-0 sink. The full side is not charged
@@ -1785,6 +1790,8 @@ def init_unified_swa_pools(
],
factory="init_unified_swa_pools",
)
if total_bytes <= 0:
raise ValueError(f"total_bytes must be positive, got {total_bytes}")
shared_pool = UnifiedKVPool(
total_bytes=total_bytes,
sub_pool_specs=[full_spec, swa_spec],
@@ -1805,12 +1812,11 @@ def init_unified_swa_pools(
unified_buffer=shared_pool,
kvcache=token_to_kv_pool,
device=device,
full_max_total_num_tokens=full_max_total_num_tokens,
swa_max_total_num_tokens=swa_max_total_num_tokens,
page_size=page_size,
need_sort=need_sort,
forward_stream=forward_stream,
lazy_compaction=lazy_compaction,
**legacy_allocator_capacities,
)
logger.info(
@@ -1841,12 +1847,12 @@ def init_unified_swa_pools(
page_size,
)
logger.info(
"[unified-memory-pool] total_bytes=%d (=%.2f GB), full_max_total_num_tokens=%d, "
"swa_max_total_num_tokens=%d, joint_available=%d slots",
"[unified-memory-pool] total_bytes=%d (=%.2f GB), "
"full_capacity=%d, swa_capacity=%d, joint_available=%d slots",
total_bytes,
total_bytes / GB,
full_max_total_num_tokens,
swa_max_total_num_tokens,
allocator.size_full,
allocator.size_swa,
allocator.available_size(),
)
logger.info(
python/sglang/srt/model_executor/model_runner.py
+11 -3
View File
@@ -1117,7 +1117,9 @@ class ModelRunner:
RoutedExpertsCapturer.create(
model=self.model,
model_config=self.model_config,
num_tokens=self.max_token_pool_size + self.page_size,
num_tokens=self.kv_index_translator.capture_token_capacity(
self.max_token_pool_size
),
max_running_requests=self.max_running_requests,
device=self.device,
)
@@ -1127,7 +1129,9 @@ class ModelRunner:
set_global_indexer_capturer(
create_indexer_capturer(
model_config=self.model_config,
num_tokens=self.max_token_pool_size + self.page_size,
num_tokens=self.kv_index_translator.capture_token_capacity(
self.max_token_pool_size
),
max_running_requests=self.max_running_requests,
device=self.device,
)
@@ -1373,7 +1377,11 @@ class ModelRunner:
def effective_max_total_num_tokens(self):
"""Return the max token pool size considering hybrid swa settings."""
if self.is_hybrid_swa:
capacity = self.full_max_total_num_tokens or self.swa_max_total_num_tokens
capacity = self.kv_cache_configurator.hybrid_swa_token_capacity(
allocator=self.token_to_kv_pool_allocator,
full_capacity=self.full_max_total_num_tokens,
swa_capacity=self.swa_max_total_num_tokens,
)
else:
capacity = self.max_total_num_tokens
if (req_to_token_pool := getattr(self, "req_to_token_pool", None)) is not None:
python/sglang/srt/model_executor/pool_configurator.py
+25 -23
View File
@@ -71,6 +71,7 @@ class MemoryPoolConfig:
max_running_requests: Optional[int] = None
full_max_total_num_tokens: Optional[int] = None
swa_max_total_num_tokens: Optional[int] = None
unified_memory_pool_bytes: Optional[int] = None
# DSV4 compressed-attention pool sizes (target only; draft workers leave at 0).
c4_max_total_num_tokens: int = 0
@@ -698,6 +699,12 @@ class HybridSWAPoolConfigurator(MemoryPoolConfigurator):
+ self._draft_cell_size
)
def _unified_pool_bytes(self, full_tokens: int, swa_tokens: int) -> int:
return (
full_tokens * self._full_per_token * self._full_layers_num
+ swa_tokens * self._swa_per_token * self._swa_layers_num
)
def _max_unified_full_tokens(
self,
available_bytes: int,
@@ -707,7 +714,6 @@ class HybridSWAPoolConfigurator(MemoryPoolConfigurator):
"""Find the largest page-aligned full capacity whose allocations fit."""
draft_bytes_per_token = self._draft_pool_bytes_per_token()
target_full_bytes_per_token = self._full_per_token * self._full_layers_num
target_swa_bytes_per_token = self._swa_per_token * self._swa_layers_num
assert target_full_bytes_per_token > 0
def allocation_bytes(full_pages: int) -> int:
@@ -719,10 +725,7 @@ class HybridSWAPoolConfigurator(MemoryPoolConfigurator):
// page_size
* page_size
)
target_bytes = (
full_tokens * target_full_bytes_per_token
+ swa_tokens * target_swa_bytes_per_token
)
target_bytes = self._unified_pool_bytes(full_tokens, swa_tokens)
virtual_span = max(target_bytes // target_full_bytes_per_token - 1, 0)
draft_tokens = ceil_align(virtual_span, page_size) + page_size
return target_bytes + draft_tokens * draft_bytes_per_token
@@ -760,29 +763,34 @@ class HybridSWAPoolConfigurator(MemoryPoolConfigurator):
full_tokens = align_page_size(max_total_num_tokens)
swa_tokens = align_page_size(int(full_tokens * self._swa_full_tokens_ratio))
self.validate_swa_pool_size(
swa_tokens, self._sliding_window_size, self._page_size
)
if not self._enable_unified_memory:
self.validate_swa_pool_size(
swa_tokens, self._sliding_window_size, self._page_size
)
logger.info(
f"Use sliding window memory pool. "
f"full_layer_tokens={full_tokens}, swa_layer_tokens={swa_tokens}"
)
return self._make_pool_config(full_tokens, swa_tokens)
def _make_pool_config(self, full_tokens: int, swa_tokens: int) -> MemoryPoolConfig:
return MemoryPoolConfig(
max_total_num_tokens=full_tokens,
full_max_total_num_tokens=full_tokens,
swa_max_total_num_tokens=swa_tokens,
unified_memory_pool_bytes=(
self._unified_pool_bytes(full_tokens, swa_tokens)
if self._enable_unified_memory
else None
),
)
def calculate_pool_sizes(
self, available_bytes: int, page_size: int
) -> MemoryPoolConfig:
if (
self._enable_unified_memory
and self._full_layers_num > 0
and self._draft_pool_bytes_per_token() > 0
):
if self._enable_unified_memory and self._full_layers_num > 0:
max_total_num_tokens = self._max_unified_full_tokens(
available_bytes, page_size
)
@@ -873,7 +881,7 @@ class SWAChunkCapPoolConfigurator(HybridSWAPoolConfigurator):
* self._swa_per_token
* (self._swa_layers_num + self._draft_swa_layers_num)
)
if self._enable_unified_memory and self._draft_pool_bytes_per_token() > 0:
if self._enable_unified_memory:
full_tokens = self._max_unified_full_tokens(
available_bytes, page_size, fixed_swa_tokens=swa_tokens
)
@@ -894,11 +902,7 @@ class SWAChunkCapPoolConfigurator(HybridSWAPoolConfigurator):
f"Reduce --max-running-requests, lower SGLANG_SWA_EVICTION_INTERVAL, "
f"or increase --mem-fraction-static."
)
return MemoryPoolConfig(
max_total_num_tokens=full_tokens,
full_max_total_num_tokens=full_tokens,
swa_max_total_num_tokens=swa_tokens,
)
return self._make_pool_config(full_tokens, swa_tokens)
def calculate_pool_sizes_from_max_tokens(
self, max_total_num_tokens: int, page_size: int
@@ -906,10 +910,8 @@ class SWAChunkCapPoolConfigurator(HybridSWAPoolConfigurator):
# Constrained max_total goes to the full pool; SWA stays at its cap.
swa_tokens = ceil_align(self._swa_cap, page_size)
full_tokens = (max_total_num_tokens // page_size) * page_size
return MemoryPoolConfig(
max_total_num_tokens=full_tokens,
full_max_total_num_tokens=full_tokens,
swa_max_total_num_tokens=min(swa_tokens, max_total_num_tokens),
return self._make_pool_config(
full_tokens, min(swa_tokens, max_total_num_tokens)
)
python/sglang/srt/state_capturer/base.py
+10 -2
View File
@@ -201,13 +201,21 @@ class BaseTopkCapturer:
slice_gpu = self._get_local_slice(
forward_batch, can_run_graph, cuda_graph_batch
)
# get_topk reads req_to_token IDs; attention may have rebound the write
# loc to kernel-facing IDs, which are neither stable nor the same space.
out_cache_loc = forward_batch.out_cache_loc_virtual
if out_cache_loc is None:
out_cache_loc = forward_batch.out_cache_loc
else:
# Kernel batches may be padded; only real request tokens are stored.
slice_gpu = slice_gpu[: out_cache_loc.shape[0]]
if no_copy_to_cpu:
# Clone before the next overlapping forward reuses these buffers.
return TopkCaptureOutput(
out_cache_loc=forward_batch.out_cache_loc.clone(),
out_cache_loc=out_cache_loc.clone(),
topk=slice_gpu.clone(),
host_cache=self.host_cache,
)
out_cache_loc_cpu = forward_batch.out_cache_loc.cpu()
out_cache_loc_cpu = out_cache_loc.cpu()
self.host_cache.buffer[out_cache_loc_cpu] = slice_gpu.cpu()
return None