feat(unified-memory): byte-budget sizing, feasibility floor, and a conservation verifier (#35158)

Co-authored-by: Caihua Li <caihua.li@bytedance.com>
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
Co-authored-by: Cheng Wan <cheng.wan@radixark.ai>

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."""

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()

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 ""
 

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
             )

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()

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],

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."