diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index 40ab8dc18..bb9038d8c 100755 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -2979,11 +2979,15 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): return total def check_decode_mem(self, selected_indices: Optional[List[int]] = None): - """Reclaim evictable tree-cache entries (shortfall only), then report - whether the next decode step fits in the KV pool.""" + """Whether the next decode step fits in the KV pool. The ALLOCATOR owns + the capacity gate (eviction + any per-step reservations of its own) — + the retract loop converges on this same check, so allocator-side + shortfalls retract gracefully instead of tripping fail-loud alloc + errors.""" num_tokens = self.new_tokens_required_next_decode(selected_indices) - evict_from_tree_cache(self.tree_cache, num_tokens) - return self.token_to_kv_pool_allocator.available_size() >= num_tokens + return self.token_to_kv_pool_allocator.check_decode_capacity( + num_tokens=num_tokens, tree_cache=self.tree_cache + ) def retract_decode(self) -> Tuple[List[Req], float, List[Req]]: """Retract the decoding requests when there is not enough memory.""" diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index ecbc6dd6d..c4b1eab24 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -4383,6 +4383,13 @@ class Scheduler( if has_leak: self.invariant_checker._report_leak("pool", "\n".join(messages)) self.invariant_checker._check_req_pool() + # Byte-conservation diagnostic (allocator-owned; static pools + # return [] — the token identity above can't see byte leaks). + byte_violations = self.token_to_kv_pool_allocator.verify_byte_accounting() + if byte_violations: + self.invariant_checker._report_leak( + "pool-bytes", "\n".join(byte_violations) + ) # tree cache sanity check self.invariant_checker._check_tree_cache() diff --git a/python/sglang/srt/mem_cache/allocator/base.py b/python/sglang/srt/mem_cache/allocator/base.py index 851d32db6..c9fb67c42 100644 --- a/python/sglang/srt/mem_cache/allocator/base.py +++ b/python/sglang/srt/mem_cache/allocator/base.py @@ -51,6 +51,39 @@ class BaseTokenToKVPoolAllocator(abc.ABC): def size_full(self): return self.size + # -- scheduler-facing capacity hooks -- + # The scheduler calls these UNCONDITIONALLY (zero feature branches on its + # side); the defaults reproduce the historical token behavior exactly, and + # unified composites override them with byte-denominated logic. + + def evict_to_free_tokens(self, tree_cache, num_tokens: int) -> None: + """Ask the prefix cache to evict unlocked entries until this allocator + can serve ``num_tokens`` (or nothing evictable remains). Default = the + shared token-count eviction; joint-byte composites override (evicting + one multi-lifetime tree node frees bytes on several sides at once). + """ + from sglang.srt.mem_cache.common import evict_from_tree_cache + + evict_from_tree_cache(tree_cache, num_tokens) + + def check_decode_capacity(self, *, num_tokens: int, tree_cache) -> bool: + """Whether the NEXT decode step's ``num_tokens`` allocation fits, + evicting reclaimable cache first. The retract loop converges on this + same check, so allocator-side shortfalls retract gracefully instead of + tripping fail-loud alloc errors. Default reproduces the historical + ``ScheduleBatch.check_decode_mem`` body; unified composites override + with byte gates + per-step reservations of their own. + """ + self.evict_to_free_tokens(tree_cache, num_tokens) + return self.available_size() >= num_tokens + + def verify_byte_accounting(self) -> list: + """Idle-time conservation diagnostic: recompute this allocator's + byte/slot accounting and return human-readable violation strings + (empty == healthy). Default: static pools have no byte model. + """ + return [] + def debug_print(self) -> str: return "" diff --git a/python/sglang/srt/mem_cache/kv_cache_configurator.py b/python/sglang/srt/mem_cache/kv_cache_configurator.py index dfbbd2ada..17a157a8e 100644 --- a/python/sglang/srt/mem_cache/kv_cache_configurator.py +++ b/python/sglang/srt/mem_cache/kv_cache_configurator.py @@ -230,6 +230,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_total_bytes: Optional[int] = None @dataclass(slots=True, kw_only=True) @@ -390,6 +391,7 @@ class KVCacheConfigurator: c128_state_pool_size=c128_state_pool_size, c4_state_dtype=c4_state_dtype, c128_state_dtype=c128_state_dtype, + unified_total_bytes=config.unified_total_bytes, ) def _init_pools( @@ -419,6 +421,7 @@ class KVCacheConfigurator: bundle = self._init_unified_mamba_pools( max_num_reqs=sizes.max_running_requests, max_total_num_tokens=sizes.max_total_num_tokens, + unified_total_bytes=sizes.unified_total_bytes, ) elif self.is_hybrid_swa and not is_deepseek_v4(self.model_config.hf_config): if pd_enabled: @@ -432,6 +435,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_total_bytes=sizes.unified_total_bytes, ) else: # Fail loud, not silently fall through to the normal pools (which would @@ -568,7 +572,11 @@ class KVCacheConfigurator: ) def _init_unified_mamba_pools( - self, *, max_num_reqs: int, max_total_num_tokens: int + self, + *, + max_num_reqs: int, + max_total_num_tokens: int, + unified_total_bytes: Optional[int] = None, ) -> UnifiedPoolBundle: """Build the shared-KV-pool stack for a hybrid-Mamba model: one byte buffer split between the full-attn MHA KV pool and the @@ -640,6 +648,9 @@ class KVCacheConfigurator: forward_stream=self.forward_stream, # Lazy compaction: default ON, 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), ) return bundle @@ -649,6 +660,7 @@ class KVCacheConfigurator: max_num_reqs: int, full_max_total_num_tokens: Optional[int], swa_max_total_num_tokens: Optional[int], + unified_total_bytes: Optional[int] = None, ) -> UnifiedPoolBundle: """Build the unified-pool stack for a hybrid-SWA model (Triton): one byte buffer split between the full-attention and SWA KV pools.""" @@ -731,6 +743,14 @@ 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`. + model_context_len=self.model_config.context_len, + sliding_window_size=self.model_config.sliding_window_size, ) return UnifiedPoolBundle( unified_memory_pool=bundle.unified_memory_pool, @@ -2034,10 +2054,18 @@ class KVCacheConfigurator: config = configurator.calculate_pool_sizes( budget_bytes, get_schedule().page_size ) + if get_memory().enable_unified_memory: + # Floor-align to 4096 B: the factories `.view()` the whole uint8 + # buffer as the KV/state dtype, so the total must be a dtype-size + # multiple and a profiled budget is not. Flooring never overcommits. + config.unified_total_bytes = budget_bytes - (budget_bytes % 4096) max_tokens = self._apply_token_constraints(config.max_total_num_tokens) if cap_tokens is not None: max_tokens = min(max_tokens, cap_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. config = configurator.calculate_pool_sizes_from_max_tokens( max_tokens, get_schedule().page_size ) diff --git a/python/sglang/srt/mem_cache/multi_ended_allocator.py b/python/sglang/srt/mem_cache/multi_ended_allocator.py index 7b89ea1be..bad11c227 100644 --- a/python/sglang/srt/mem_cache/multi_ended_allocator.py +++ b/python/sglang/srt/mem_cache/multi_ended_allocator.py @@ -345,6 +345,31 @@ class MultiEndedAllocator(BaseTokenToKVPoolAllocator): return self.watermark_physical * self.entry_bytes_per_page return self.num_pages * self.entry_bytes_per_page + def _byte_accounting_violations(self) -> List[str]: + """Per-sub-pool conservation strings (empty == healthy): the watermark + span must equal live + holes + pending pages, and frontiers must lie + inside the buffer. Idle-time diagnostic — pure host arithmetic.""" + out: List[str] = [] + total = self.unified_buffer.total_bytes + lo_b, hi_b = self._byte_low_frontier(), self._byte_high_frontier() + if not (0 <= lo_b <= hi_b <= total): + out.append( + f"[{self.sub_pool_name}] frontier out of bounds: " + f"low={lo_b}, high={hi_b}, total={total}" + ) + if self.lazy_compaction: + # Lazy end: the watermark span contains live + holes + pending + # (eager has no holes/pending — span == live by construction). + holes = int(self._free_phys_pages.numel()) + pending = len(self._pending_reuse_pages_cpu) + wm_span = self._allocated_pages() + if wm_span != self.live_page_count + holes + pending: + out.append( + f"[{self.sub_pool_name}] span {wm_span} != live " + f"{self.live_page_count} + holes {holes} + pending {pending}" + ) + return out + def _byte_low_frontier(self) -> int: """Byte starting this side's allocatable range (grow-up) / just below its lowest live page (grow-down).""" if self.grow_direction == "up": @@ -1750,6 +1775,40 @@ class MultiEndedAllocator(BaseTokenToKVPoolAllocator): self.free(reps, _pages=reps // self.page_size) +def _chain_byte_accounting_violations( + chain: List[MultiEndedAllocator], +) -> List[str]: + """Conservation for an ordered low→high chain of band allocators: each + member's own accounting, plus the frontier total order — a member's low + frontier must clear the previous member's high frontier, or the bands + overlap in the shared byte buffer. + + Today's chains are the 2-pool end pairs; the N-pool track inserts float + middles here (and teaches the walk to skip empty/parked ones). + """ + out: List[str] = [] + for a in chain: + out.extend(a._byte_accounting_violations()) + frontier = 0 + for a in chain: + lo_b, hi_b = a._byte_low_frontier(), a._byte_high_frontier() + if lo_b < frontier: + out.append( + f"[chain] {a.sub_pool_name} low frontier {lo_b} overlaps the " + f"previous pool's high frontier {frontier}" + ) + frontier = max(frontier, hi_b) + return out + + +def _end_pair_chain( + a: MultiEndedAllocator, b: MultiEndedAllocator +) -> List[MultiEndedAllocator]: + """Order an end pair low→high by grow direction (the factories and the + unit fixtures orient the pair differently; the chain check must not care).""" + return sorted((a, b), key=lambda x: x.grow_direction != "up") + + class UnifiedMambaTokenToKVPoolAllocator(BaseTokenToKVPoolAllocator): """Composite allocator for the MHA (full-attn) + Mamba hybrid pair. @@ -2042,6 +2101,11 @@ class UnifiedMambaTokenToKVPoolAllocator(BaseTokenToKVPoolAllocator): self.full_attn_allocator.clear_inverse_history() self.mamba_allocator.clear_inverse_history() + def verify_byte_accounting(self) -> List[str]: + return _chain_byte_accounting_violations( + _end_pair_chain(self.mamba_allocator, self.full_attn_allocator) + ) + def free_group_begin(self) -> None: super().free_group_begin() self.free_page_reps_group = [] @@ -2602,6 +2666,11 @@ 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) + ) + def clear(self) -> None: self.full_attn_allocator.clear() self.swa_attn_allocator.clear() diff --git a/python/sglang/srt/mem_cache/unified_memory_pool.py b/python/sglang/srt/mem_cache/unified_memory_pool.py index afd4d3c7d..9c69045af 100644 --- a/python/sglang/srt/mem_cache/unified_memory_pool.py +++ b/python/sglang/srt/mem_cache/unified_memory_pool.py @@ -253,6 +253,30 @@ def _assert_kernel_id_bound(*, sub_pool_name: str, n_rows: int) -> None: ) +def _reserved_floor_bytes(sub_pool_specs: List[SubPoolSpec], page_size: int) -> int: + """Bytes at the bottom of the buffer reserved as the slot-0 padding sink. + + Slot-0 dummy writes for every sub-pool land here; each sub-pool's first + allocatable slot is chosen so real data starts past it. For a PAGE-AWARE + sub-pool the slot-0 write touches layer blocks spread across the whole + page-0 envelope (page_size * entry_bytes), not just one slot envelope -- + but a mamba sub-pool is page_size=1, so its entry is charged ONCE. Charging + a mamba entry per page would reserve page_size * ~100 MB of buffer that the + sink never touches. + + Single source of truth: `UnifiedKVPool` reserves exactly this, and the + factories' bs=1 feasibility floors charge exactly this. + """ + return max( + [max(s.entry_bytes() for s in sub_pool_specs)] + + [ + page_size * s.entry_bytes() + for s in sub_pool_specs + if not isinstance(s, MambaSubPoolSpec) # mamba is page_size=1 + ] + ) + + class UnifiedKVPool: """One physical `uint8` byte buffer shared by 2 sub-pools, each exposing per-layer views over its own byte range (contiguous per layer for KV, @@ -324,15 +348,7 @@ class UnifiedKVPool: # For a page-aware sub-pool the slot-0 write touches layer blocks spread # across the WHOLE page-0 envelope (up to page_size * entry_bytes), not # just one slot envelope — reserve the max of both. - entry_max = max(s.entry_bytes() for s in sub_pool_specs) - reserved_floor = max( - [entry_max] - + [ - page_size * s.entry_bytes() - for s in sub_pool_specs - if not isinstance(s, MambaSubPoolSpec) # mamba is page_size=1 - ] - ) + reserved_floor = _reserved_floor_bytes(sub_pool_specs, page_size) for spec in sub_pool_specs: entry_bytes = spec.entry_bytes() @@ -1086,6 +1102,31 @@ class UnifiedPoolBundle(NamedTuple): req_to_token_pool: object # UnifiedHybridReqToTokenPool +def _check_bs1_feasibility_floor( + *, + total_bytes: int, + floor_terms: List[Tuple[str, int]], + factory: str, +) -> None: + """bs=1 feasibility FLOOR — the retract loop's terminal guarantee. + + The scheduler retracts requests until the LAST one fits; if one worst-case + request running ALONE does not fit in the buffer, under-sizing is a retract + LIVELOCK at runtime, not a perf bug. Fail loud at boot, before any pool + construction, with the itemized requirement. + """ + floor = sum(b for _, b in floor_terms) + if total_bytes >= floor: + return + detail = " + ".join(f"{name}={b}" for name, b in floor_terms) + raise RuntimeError( + f"[unified-memory-pool] {factory}: byte budget {total_bytes} cannot fit " + f"ONE worst-case request (bs=1 floor {floor} = {detail}). A pool this " + f"size retract-livelocks at runtime. Raise --mem-fraction-static, lower " + f"the model context length, or reduce reserved memory." + ) + + def init_unified_mamba_pools( *, device: str, @@ -1116,6 +1157,7 @@ def init_unified_mamba_pools( forward_stream: Optional[torch.cuda.Stream] = None, lazy_compaction: bool = False, decode_pre_alloc_size: int = 0, + unified_total_bytes: Optional[int] = None, ) -> UnifiedPoolBundle: """Build the Mamba-hybrid unified-memory-pool stack.""" from sglang.srt.mem_cache.multi_ended_allocator import ( @@ -1164,9 +1206,29 @@ def init_unified_mamba_pools( conv_slice_axis=getattr(cp.shape, "conv_slice_axis", 0), grow_direction="up", ) - total_bytes = ( - max_total_num_tokens * full_spec.entry_bytes() - + max_mamba_cache_size * mamba_spec.entry_bytes() + if unified_total_bytes is not None: + # PROFILED byte budget for the token side (captured pre-ratio-floor); + # the state pool's bytes ride on top. The token counts stay boot + # labels / conserve caps -- the runtime split floats. + total_bytes = ( + unified_total_bytes + max_mamba_cache_size * mamba_spec.entry_bytes() + ) + else: + total_bytes = ( + max_total_num_tokens * full_spec.entry_bytes() + + max_mamba_cache_size * mamba_spec.entry_bytes() + ) + # bs=1 floor: the state slots one running request locks (1 active + 2 radix + # checkpoints, a FLOOR not headroom) + the slot-0 sink. The token side is + # not charged -- `TpModelWorker.get_worker_info` already clamps max_req_len + # to the pool, so a too-long request is refused at admission, not livelocked. + _check_bs1_feasibility_floor( + total_bytes=total_bytes, + floor_terms=[ + ("bs1_state_slots", 3 * mamba_spec.entry_bytes()), + ("sink", _reserved_floor_bytes([full_spec, mamba_spec], page_size)), + ], + factory="init_unified_mamba_pools", ) shared_pool = UnifiedKVPool( total_bytes=total_bytes, @@ -1576,6 +1638,9 @@ def init_unified_swa_pools( 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: """Build the SWA-hybrid unified-memory-pool stack.""" from sglang.srt.mem_cache.multi_ended_allocator import ( @@ -1612,10 +1677,33 @@ def init_unified_swa_pools( store_dtype=store_dtype, grow_direction="up", ) - total_bytes = ( - full_max_total_num_tokens * full_spec.entry_bytes() - + swa_max_total_num_tokens * swa_spec.entry_bytes() - ) + 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: + total_bytes = ( + full_max_total_num_tokens * full_spec.entry_bytes() + + swa_max_total_num_tokens * swa_spec.entry_bytes() + ) + 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 + # (max_req_len clamps it); the swa sub-pool is sized independently of + # that clamp, which is why the window term stays. + swa_bs1_tokens = ( + min(model_context_len, sliding_window_size + page_size) + if sliding_window_size is not None + else model_context_len + ) + _check_bs1_feasibility_floor( + total_bytes=total_bytes, + floor_terms=[ + ("swa_window_kv", swa_bs1_tokens * swa_spec.entry_bytes()), + ("sink", _reserved_floor_bytes([full_spec, swa_spec], page_size)), + ], + factory="init_unified_swa_pools", + ) shared_pool = UnifiedKVPool( total_bytes=total_bytes, sub_pool_specs=[full_spec, swa_spec], diff --git a/python/sglang/srt/model_executor/pool_configurator.py b/python/sglang/srt/model_executor/pool_configurator.py index 4f6dc72ae..38b7fccd4 100644 --- a/python/sglang/srt/model_executor/pool_configurator.py +++ b/python/sglang/srt/model_executor/pool_configurator.py @@ -71,6 +71,13 @@ class MemoryPoolConfig: mem_fraction_static: Optional[float] = None + # Unified pool only: the PROFILED byte budget for the token-granular + # sub-pools. Set, the factories size the buffer from it directly instead of + # re-summing ratio-derived token counts, which keeps the re-sum's floor + # losses out of the buffer; the token counts stay boot labels / conserve + # caps. None on the token-capped path -- a user token cap IS the budget. + unified_total_bytes: Optional[int] = None + def __post_init__(self): if self.max_total_num_tokens <= 0: msg = "Not enough memory. Please try to increase --mem-fraction-static." diff --git a/test/registered/unit/mem_cache/test_unified_byte_accounting.py b/test/registered/unit/mem_cache/test_unified_byte_accounting.py new file mode 100644 index 000000000..683a94d20 --- /dev/null +++ b/test/registered/unit/mem_cache/test_unified_byte_accounting.py @@ -0,0 +1,157 @@ +# Copyright 2023-2026 SGLang Team +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Byte-conservation verifier for the unified 2-pool composites. + +`verify_byte_accounting` is the idle-time tripwire the token-identity leak +check cannot provide: the unified pool's correctness rests on BYTE bookkeeping +(watermark spans, holes, pending compaction, frontier ordering inside one +shared buffer), and a drifted counter admits requests into memory that is not +actually free — silent corruption territory, not a crash. + +Derived properties pinned here: + + * Conservation: on a lazy end pool the watermark span must equal + live + holes + pending pages at EVERY point of a healthy lifecycle + (alloc, partial free, group free, flush) — not just at rest. + * The check is not vacuous: drifting any single term (live count, watermark, + a leaked hole) reports loudly, naming the sub-pool. + * Chain order: one member's low frontier clearing the other's high frontier + is what "two pools share one buffer without overlap" MEANS; the pair check + must hold regardless of which member grows up. + * The strict escalation env defaults OFF: promoting the diagnostic to a + RuntimeError is a validation posture, not the production one. + + python -m pytest test/registered/unit/mem_cache/test_unified_byte_accounting.py -v +""" + +import unittest + +from test_multi_ended_allocator import TestPagedMultiEndedAllocator as _PagedFixture +from test_multi_ended_allocator import ( + TestUnifiedSWATokenToKVPoolAllocator as _SwaFixture, +) + +from sglang.srt.mem_cache import multi_ended_allocator as mea +from sglang.test.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=15, suite="base-a-test-cpu") + + +def _swa_composite(): + inst = _SwaFixture([m for m in dir(_SwaFixture) if m.startswith("test_")][0]) + pool, allocator, kvcache = inst._build() + return inst, allocator, kvcache + + +def _paged_pair(lazy: bool): + inst = _PagedFixture([m for m in dir(_PagedFixture) if m.startswith("test_")][0]) + _pool, full, swa, _fkv, _skv = inst._build() + full.lazy_compaction = lazy + return full, swa + + +class TestHealthyLifecycleReportsClean(unittest.TestCase): + def test_swa_composite_clean_at_every_step(self): + inst, allocator, kvcache = _swa_composite() + self.assertEqual(allocator.verify_byte_accounting(), []) + v = inst._alloc(allocator, kvcache, 8) + self.assertEqual(allocator.verify_byte_accounting(), []) + allocator.free_swa(v[:4]) # tombstone half the swa side + self.assertEqual(allocator.verify_byte_accounting(), []) + inst._free(allocator, kvcache, v) + self.assertEqual(allocator.verify_byte_accounting(), []) + allocator.clear() + self.assertEqual(allocator.verify_byte_accounting(), []) + + def test_lazy_end_pool_clean_through_free_and_flush(self): + full, _swa = _paged_pair(lazy=True) + self.assertEqual(full._byte_accounting_violations(), []) + v = full.alloc(full.page_size * 4) + self.assertEqual(full._byte_accounting_violations(), []) + full.free(v[: full.page_size * 2]) # lazy: holes, no compaction yet + self.assertEqual(full._byte_accounting_violations(), []) + full._flush(urgent=True) + self.assertEqual(full._byte_accounting_violations(), []) + + +class TestDriftReportsLoudly(unittest.TestCase): + """Each mutation below models a distinct bookkeeping bug; the verifier + must name the drifted sub-pool. Without these, a regression in any single + counter passes every other test (the pool still 'works' — it just lies + about capacity).""" + + def _lazy_full(self): + full, _swa = _paged_pair(lazy=True) + v = full.alloc(full.page_size * 4) + full.free(v[: full.page_size]) # one hole so all three terms are live + self.assertEqual(full._byte_accounting_violations(), []) + return full + + def test_drifted_live_count(self): + full = self._lazy_full() + full.live_page_count += 1 + self.assertTrue(any("span" in s for s in full._byte_accounting_violations())) + + def test_leaked_hole(self): + full = self._lazy_full() + full._free_phys_pages = full._free_phys_pages[:-1] # hole vanished + self.assertTrue(any("span" in s for s in full._byte_accounting_violations())) + + def test_drifted_watermark(self): + full = self._lazy_full() + full.watermark_physical += 1 + self.assertTrue(any("span" in s for s in full._byte_accounting_violations())) + + def test_composite_report_names_the_sub_pool(self): + """Frontier-bounds drift (checked in BOTH lazy and eager modes): push + the swa band's watermark outside the buffer.""" + inst, allocator, kvcache = _swa_composite() + inst._alloc(allocator, kvcache, 8) + swa = allocator.swa_attn_allocator + # grow-down member: low frontier = (wm+1)*bytes; wm == num_pages puts + # it past the buffer top. + self.assertEqual(swa.grow_direction, "down") + swa.watermark_physical = swa.num_pages + out = allocator.verify_byte_accounting() + self.assertTrue(out and any("[swa]" in s for s in out), out) + + +class TestChainFrontierOrder(unittest.TestCase): + def test_overlapping_frontiers_report(self): + """Both bands hold pages, then the up member's watermark is pushed past + the down member's LIVE low frontier: the two bands now claim the same + bytes of one buffer. (An empty down band cannot overlap — its low + frontier IS the buffer top — so both sides must be populated for the + scenario to be a real corruption.)""" + full, swa = _paged_pair(lazy=False) + chain = mea._end_pair_chain(full, swa) + up, down = chain + self.assertEqual(up.grow_direction, "up") + self.assertIsNotNone(down.alloc(down.page_size * 2)) # down side live + self.assertLess(down._byte_low_frontier(), up.unified_buffer.total_bytes) + up.watermark_physical = up.num_pages # up band swallows the buffer + out = mea._chain_byte_accounting_violations(chain) + self.assertTrue(any("overlap" in s for s in out), out) + + def test_pair_order_is_direction_agnostic(self): + """The factories and the unit fixtures orient the pair differently; + the check must order by grow direction, not by argument position.""" + full, swa = _paged_pair(lazy=False) + a = mea._end_pair_chain(full, swa) + b = mea._end_pair_chain(swa, full) + self.assertEqual([x.sub_pool_name for x in a], [x.sub_pool_name for x in b]) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/mem_cache/test_unified_byte_budget_sizing.py b/test/registered/unit/mem_cache/test_unified_byte_budget_sizing.py new file mode 100644 index 000000000..d17c99280 --- /dev/null +++ b/test/registered/unit/mem_cache/test_unified_byte_budget_sizing.py @@ -0,0 +1,244 @@ +# Copyright 2023-2026 SGLang Team +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Byte-budget buffer sizing for the unified 2-pool factories. + +Derived properties pinned here: + + * Budget honored EXACTLY: with ``unified_total_bytes`` set, the swa pair's + buffer is that many bytes (the mamba pair adds the state pool's bytes on + top — the budget is captured AFTER the state carve-out). Sizing from the + ratio-derived token counts instead re-introduces the configurator's + rounding: the swa split floors the budget by the cell size and then + page-aligns EACH side's token count, so the re-sum reconstructs less + than the profiled budget by up to about one page of tokens per side. + * Fallback: without the budget, sizing is the historical token-count re-sum, + bit-for-bit. + * bs=1 feasibility floor: a budget that cannot fit ONE worst-case request + (full KV at max context, plus one SWA window / the state slots a single + running request locks) raises at BOOT, before any pool construction — + under-sizing is a retract LIVELOCK at runtime, not a perf bug. + * The 4096-byte alignment exists because the factories ``.view()`` the whole + uint8 buffer as the KV dtype; an unaligned budget must be floored, never + rounded up (rounding up overcommits profiled memory). + + python -m pytest test/registered/unit/mem_cache/test_unified_byte_budget_sizing.py -v +""" + +import unittest + +import torch + +from sglang.srt.mem_cache.unified_memory_pool import ( + MambaSubPoolSpec, + MHASubPoolSpec, + UnifiedKVPool, + _check_bs1_feasibility_floor, + _reserved_floor_bytes, + init_unified_swa_pools, +) +from sglang.test.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=15, suite="base-a-test-cpu") + +_DEV = "cpu" + + +def _swa_factory(**over): + kw = dict( + device=_DEV, + kv_cache_dtype=torch.float16, + head_num=2, + head_dim=8, + v_head_dim=8, + swa_head_num=2, + swa_head_dim=8, + swa_v_head_dim=8, + page_size=1, + start_layer=0, + end_layer=4, + swa_attention_layer_ids=[1, 3], + full_attention_layer_ids=[0, 2], + full_max_total_num_tokens=64, + swa_max_total_num_tokens=32, + enable_memory_saver=False, + need_sort=False, + ) + kw.update(over) + return init_unified_swa_pools(**kw) + + +def _entry_bytes(): + full = MHASubPoolSpec( + name="full", + layer_num=2, + head_num=2, + head_dim=8, + store_dtype=torch.float16, + grow_direction="up", + ) + return full.entry_bytes() + + +class TestBudgetSizing(unittest.TestCase): + def test_swa_factory_honors_the_budget_exactly(self): + e = _entry_bytes() + budget = 96 * e + 512 # deliberately NOT a token-count multiple + bundle = _swa_factory(unified_total_bytes=budget) + self.assertEqual(bundle.unified_memory_pool.total_bytes, budget) + + def test_fallback_is_the_token_count_resum(self): + e = _entry_bytes() + bundle = _swa_factory() + self.assertEqual(bundle.unified_memory_pool.total_bytes, (64 + 32) * e) + + def test_budget_beats_resum_on_rounding(self): + """The property that motivates the whole phase: the re-sum cannot + represent a budget that is not a whole-token multiple per side, so it + strands bytes the buffer could have held.""" + e = _entry_bytes() + budget = (64 + 32) * e + (e - 2) # almost one more entry + bundle = _swa_factory(unified_total_bytes=budget) + self.assertEqual(bundle.unified_memory_pool.total_bytes, budget) + self.assertGreater(budget, (64 + 32) * e) + + +class TestReservedFloorIsOneSourceOfTruth(unittest.TestCase): + """The bs=1 floor charges the slot-0 sink, and MUST charge exactly what + `UnifiedKVPool` actually reserves. + + Regression (GPU eval_434/436, Falcon-H1 boot): the floor hand-copied the + formula as `page_size * max(entry_bytes)`, applying the page multiplier to + the MAMBA spec. The pool deliberately excludes mamba (it is page_size=1), + so with page_size=256 and a ~139 MB state entry the floor over-charged the + sink by 256x — ~33 GiB of phantom requirement — and a healthy config + failed to boot with 25 GiB of real headroom. + """ + + def _specs(self, page_size): + full = MHASubPoolSpec( + name="full", + layer_num=2, + head_num=2, + head_dim=8, + store_dtype=torch.float16, + grow_direction="down", + ) + # A state entry vastly larger than a KV token entry — the real ratio + # (~139 MB vs ~45 KB) is what made the over-charge fatal. + mamba = MambaSubPoolSpec( + name="mamba", + layer_num=2, + conv_state_shapes=((4, 256),), + conv_dtype=torch.float16, + temporal_state_shape=(4, 256, 64), + temporal_dtype=torch.float16, + grow_direction="up", + ) + return full, mamba + + def test_mamba_entry_is_not_multiplied_by_page_size(self): + full, mamba = self._specs(page_size=256) + got = _reserved_floor_bytes([full, mamba], 256) + self.assertEqual(got, max(mamba.entry_bytes(), 256 * full.entry_bytes())) + self.assertLess(got, 256 * mamba.entry_bytes()) # the bug's value + + def test_floor_sink_equals_what_the_pool_reserves(self): + """Pin the two against each other so the formula cannot drift again.""" + for page_size in (1, 4, 256): + with self.subTest(page_size=page_size): + full, mamba = self._specs(page_size) + floor = _reserved_floor_bytes([full, mamba], page_size) + pool = UnifiedKVPool( + total_bytes=floor + 64 * mamba.entry_bytes(), + sub_pool_specs=[full, mamba], + device=_DEV, + enable_memory_saver=False, + page_size=page_size, + ) + # min_slot_index is ceil(reserved_floor / entry_bytes) per side. + for spec in (full, mamba): + self.assertEqual( + pool.min_slot_index(spec.name), + -(-floor // spec.entry_bytes()), + ) + + +class TestBs1FeasibilityFloor(unittest.TestCase): + def test_infeasible_budget_raises_before_construction(self): + """The buffer cannot hold one sliding window plus the sink, so boot + must fail loud instead of livelocking later.""" + with self.assertRaises(RuntimeError) as ctx: + _swa_factory( + unified_total_bytes=8 * _entry_bytes(), + model_context_len=4096, + sliding_window_size=4096, + ) + self.assertIn("bs=1 floor", str(ctx.exception)) + self.assertIn("swa_window_kv", str(ctx.exception)) + + def test_context_longer_than_the_pool_is_not_rejected(self): + """REGRESSION: the floor must NOT charge the full-attention token side. + `TpModelWorker.get_worker_info` clamps max_req_len to the pool, so a + context far larger than the buffer is refused at admission, not a + livelock -- and it is an ordinary way to serve a long-context model on + one GPU. Charging it here made such configs fail at boot.""" + e = _entry_bytes() + bundle = _swa_factory( + unified_total_bytes=200 * e, + model_context_len=1_000_000, # far beyond what the buffer holds + sliding_window_size=16, + ) + self.assertEqual(bundle.unified_memory_pool.total_bytes, 200 * e) + + def test_feasible_config_boots_with_floor_inputs_present(self): + e = _entry_bytes() + bundle = _swa_factory( + unified_total_bytes=200 * e, + model_context_len=64, + sliding_window_size=16, + ) + self.assertEqual(bundle.unified_memory_pool.total_bytes, 200 * e) + + def test_window_term_is_clamped_to_context(self): + """A window larger than the context must charge at most the context — + otherwise short-context models over-raise.""" + e = _entry_bytes() + bundle = _swa_factory( + unified_total_bytes=200 * e, + model_context_len=64, + sliding_window_size=10_000, # window >> context + ) + self.assertIsNotNone(bundle) + + def test_floor_message_itemizes_terms(self): + with self.assertRaises(RuntimeError) as ctx: + _check_bs1_feasibility_floor( + total_bytes=10, + floor_terms=[("a", 8), ("b", 8)], + factory="test", + ) + msg = str(ctx.exception) + self.assertIn("a=8", msg) + self.assertIn("b=8", msg) + self.assertIn("16", msg) + + def test_exact_floor_passes(self): + """Boundary: total == floor must NOT raise (>= is the contract).""" + _check_bs1_feasibility_floor( + total_bytes=16, floor_terms=[("a", 8), ("b", 8)], factory="test" + ) + + +if __name__ == "__main__": + unittest.main()