From ef9e58fd6d0140f9d2bade6a31dbab779013d038 Mon Sep 17 00:00:00 2001 From: caihuali95 <42954765+caihuali95@users.noreply.github.com> Date: Mon, 31 Aug 2026 15:10:12 -0700 Subject: [PATCH] feat(unified-memory): three sub-pools for mamba + hybrid-SWA models (#35177) Co-authored-by: Caihua Li Co-authored-by: Claude Fable 5 Co-authored-by: Cheng Wan --- python/sglang/srt/managers/schedule_policy.py | 17 +- .../scheduler_components/invariant_checker.py | 18 +- .../pool_stats_observer.py | 17 +- python/sglang/srt/mem_cache/common.py | 16 +- .../srt/mem_cache/kv_cache_configurator.py | 140 +- .../srt/mem_cache/multi_ended_allocator.py | 1534 +++++++++++++++-- .../srt/mem_cache/unified_memory_pool.py | 312 +++- .../models_e2e/test_inkling_unified.py | 240 +++ .../test_unified_memory_move_gate.py | 13 +- .../mem_cache/test_multi_ended_allocator.py | 481 +++++- .../mem_cache/test_unified_capacity_memo.py | 299 ++++ .../test_unified_free_no_host_sync.py | 196 ++- .../mem_cache/test_unified_npool_sweep.py | 271 +++ .../unit/mem_cache/test_unified_tri_pool.py | 1404 +++++++++++++++ 14 files changed, 4778 insertions(+), 180 deletions(-) create mode 100644 test/registered/models_e2e/test_inkling_unified.py create mode 100644 test/registered/unit/mem_cache/test_unified_capacity_memo.py create mode 100644 test/registered/unit/mem_cache/test_unified_npool_sweep.py create mode 100644 test/registered/unit/mem_cache/test_unified_tri_pool.py diff --git a/python/sglang/srt/managers/schedule_policy.py b/python/sglang/srt/managers/schedule_policy.py index 5bf710ad0..ae9fa0dc9 100644 --- a/python/sglang/srt/managers/schedule_policy.py +++ b/python/sglang/srt/managers/schedule_policy.py @@ -61,6 +61,7 @@ from sglang.srt.mem_cache.base_prefix_cache import ( zero_match_result, ) from sglang.srt.mem_cache.multi_ended_allocator import ( + UnifiedMambaSWATokenToKVPoolAllocator, UnifiedMambaTokenToKVPoolAllocator, ) from sglang.srt.mem_cache.radix_cache import RadixCache, RadixKey, TreeNode @@ -547,15 +548,17 @@ class PrefillAdder: self.rem_swa_token_offset = 0 - # Unified-pool joint budget: a new mamba state consumes shared-gap bytes - # that `rem_total_tokens` (full KV) otherwise counts as free, so reserve - # the gap per new mamba slot or admission over-commits. Gate on the - # ALLOCATOR being the unified Mamba composite, NOT on `is_hybrid_ssm_cache` - # (False for `ChunkCache`, which would skip the reservation on the - # chunk-cache path): the gap coupling is a property of the byte buffer. + # 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`, + # which would skip the reservation on the chunk-cache path. self._mamba_slot_cost = 0 if isinstance( - self.token_to_kv_pool_allocator, UnifiedMambaTokenToKVPoolAllocator + self.token_to_kv_pool_allocator, + ( + UnifiedMambaTokenToKVPoolAllocator, + UnifiedMambaSWATokenToKVPoolAllocator, + ), ): self._mamba_slot_cost = ( self.token_to_kv_pool_allocator.mamba_slot_full_token_cost() diff --git a/python/sglang/srt/managers/scheduler_components/invariant_checker.py b/python/sglang/srt/managers/scheduler_components/invariant_checker.py index cdb4a8c70..e2af23cfb 100644 --- a/python/sglang/srt/managers/scheduler_components/invariant_checker.py +++ b/python/sglang/srt/managers/scheduler_components/invariant_checker.py @@ -23,6 +23,9 @@ from sglang.srt.managers.scheduler_components.pool_stats_observer import ( from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache from sglang.srt.mem_cache.memory_pool import ReqToTokenPool +from sglang.srt.mem_cache.multi_ended_allocator import ( + UnifiedMambaSWATokenToKVPoolAllocator, +) from sglang.srt.runtime_context import get_parallel from sglang.srt.utils.common import ( ceil_align, @@ -116,9 +119,14 @@ 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", - ps.full_available_size, + full_available, full_evictable_size, protected, session_held, @@ -134,9 +142,15 @@ class SchedulerInvariantChecker: return leak, msg def _check_swa_pool(self, ps: PoolStats, uncached: int = 0) -> Tuple[bool, str]: + allocator = self.token_to_kv_pool_allocator + 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", - ps.swa_available_size, + swa_available, ps.swa_evictable_size, self.tree_cache.swa_protected_size(), self.pool_stats_observer.session_held_swa_tokens(), diff --git a/python/sglang/srt/managers/scheduler_components/pool_stats_observer.py b/python/sglang/srt/managers/scheduler_components/pool_stats_observer.py index 5529b770f..7dc5f63b8 100644 --- a/python/sglang/srt/managers/scheduler_components/pool_stats_observer.py +++ b/python/sglang/srt/managers/scheduler_components/pool_stats_observer.py @@ -11,6 +11,10 @@ from typing import ( Tuple, ) +from sglang.srt.mem_cache.multi_ended_allocator import ( + UnifiedMambaSWATokenToKVPoolAllocator, +) + if TYPE_CHECKING: from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache @@ -284,9 +288,18 @@ class SchedulerPoolStatsObserver: ) def _get_swa_token_info(self) -> PoolStats: - full_available_size = self.token_to_kv_pool_allocator.full_available_size() + # `*_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_evictable_size = self.tree_cache.full_evictable_size() - swa_available_size = self.token_to_kv_pool_allocator.swa_available_size() swa_evictable_size = self.tree_cache.swa_evictable_size() full_num_used = self.full_tokens_per_layer - ( full_available_size + full_evictable_size diff --git a/python/sglang/srt/mem_cache/common.py b/python/sglang/srt/mem_cache/common.py index 1b646392e..0915b5bb3 100644 --- a/python/sglang/srt/mem_cache/common.py +++ b/python/sglang/srt/mem_cache/common.py @@ -101,7 +101,21 @@ def free_swa_out_of_window_slots( free_slots = req_to_token_pool.req_to_token[ req.kv.req_pool_idx, req.kv.swa_evicted_seqlen : new_swa_evicted_seqlen ] - token_to_kv_pool_allocator.free_swa(free_slots) + # Local import: multi_ended_allocator imports this module lazily for + # eviction; a module-level import here would be a cycle hazard. + from sglang.srt.mem_cache.multi_ended_allocator import ( + UnifiedSWATokenToKVPoolAllocator, + ) + + if isinstance(token_to_kv_pool_allocator, UnifiedSWATokenToKVPoolAllocator): + # Contiguous range with host-int bounds: hand the composite its + # start position so the free stays host-sync-free (`free_segment` + # derives page reps by stride math instead of `torch.unique`). + token_to_kv_pool_allocator.free_swa( + free_slots, start_pos=req.kv.swa_evicted_seqlen + ) + else: + token_to_kv_pool_allocator.free_swa(free_slots) req.kv.swa_evicted_seqlen = new_swa_evicted_seqlen diff --git a/python/sglang/srt/mem_cache/kv_cache_configurator.py b/python/sglang/srt/mem_cache/kv_cache_configurator.py index 17a157a8e..98634cc8d 100644 --- a/python/sglang/srt/mem_cache/kv_cache_configurator.py +++ b/python/sglang/srt/mem_cache/kv_cache_configurator.py @@ -409,7 +409,30 @@ class KVCacheConfigurator: # (req_to_token_pool is None); supports hybrid Mamba and hybrid SWA (not DSV4). if get_memory().enable_unified_memory and req_to_token_pool is None: pd_enabled = get_disagg().disaggregation_mode != "null" - if self.mambaish_config is not None: + is_dsv4 = is_deepseek_v4(self.model_config.hf_config) + # Order matters: an Inkling-class model is BOTH mambaish and + # hybrid-SWA, and the mamba pair would store every SWA layer's KV at + # FULL lifetime -- its branch reads the HF config's + # full_attention_layer_ids, which for Inkling is ALL layers. + if self.mambaish_config is not None and self.is_hybrid_swa and not is_dsv4: + if pd_enabled: + # Same limitation as the 2-pool SWA branch below: the + # tri-pool carries an SWA sub-pool, and there is no + # whole-envelope transfer scheme for it. + raise ValueError( + "--enable-unified-memory with PD disaggregation does " + "not support hybrid-SWA models yet (no whole-envelope " + "transfer scheme for the SWA sub-pool); this model " + "routes to the mamba+SWA tri-pool, which has one. Drop " + "--enable-unified-memory or run without PD." + ) + bundle = self._init_unified_mamba_swa_pools( + 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, + ) + elif self.mambaish_config is not None: if pd_enabled and not self.use_mla_backend: raise ValueError( "--enable-unified-memory with PD disaggregation " @@ -423,7 +446,7 @@ class KVCacheConfigurator: 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): + elif self.is_hybrid_swa and not is_dsv4: if pd_enabled: raise ValueError( "--enable-unified-memory with PD disaggregation does " @@ -654,6 +677,119 @@ class KVCacheConfigurator: ) return bundle + def _init_unified_mamba_swa_pools( + self, + *, + 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: + """TRI-pool stack for models that are BOTH mambaish and hybrid-SWA + (Inkling-class): full KV + SWA KV + mamba/conv state in one buffer, + chain [mamba(up) | swa(float) | full(down)].""" + from sglang.srt.mem_cache.unified_memory_pool import ( + init_unified_mamba_swa_pools, + ) + + config = self.mambaish_config + assert config is not None and self.is_hybrid_swa + assert self.page_size >= 1, f"page_size must be >= 1, got {self.page_size}" + assert ( + not self.use_mla_backend + ), "unified tri-pool does not support an MLA full side yet" + # Mirror the non-shared path's extra_max_context_len computation. + extra_max_context_len = 4 + if get_spec().speculative_num_draft_tokens is not None: + extra_max_context_len += get_spec().speculative_num_draft_tokens + + head_num = self.model_config.get_num_kv_heads( + get_parallel().attn_tp_size, get_parallel().attn_dcp_size + ) + head_dim = self.model_config.head_dim + if self.is_hybrid_swa_compress: + # Asymmetric full/SWA head geometry (Inkling): SWA dims from the + # hf text config, same as the 2-pool SWA wrapper. + v_head_dim = self.model_config.hf_text_config.v_head_dim + swa_head_num = max( + 1, + self.model_config.hf_text_config.swa_num_key_value_heads + // get_parallel().attn_tp_size, + ) + swa_head_dim = self.model_config.hf_text_config.swa_head_dim + swa_v_head_dim = self.model_config.hf_text_config.swa_v_head_dim + else: + v_head_dim = head_dim + swa_head_num = head_num + swa_head_dim = head_dim + swa_v_head_dim = head_dim + + # From the sglang ModelConfig WRAPPER, never the HF config's + # full_attention_layer_ids: that property feeds the conv/attention + # pairing, not the KV-lifetime split, and returns ALL layers. + swa_attention_layer_ids = [ + i + for i in self.model_config.swa_attention_layer_ids + if self.layer_info.start_layer <= i < self.layer_info.end_layer + ] + full_attention_layer_ids = [ + i + for i in self.model_config.full_attention_layer_ids + if self.layer_info.start_layer <= i < self.layer_info.end_layer + ] + n_local_layers = self.layer_info.end_layer - self.layer_info.start_layer + assert ( + len(full_attention_layer_ids) + len(swa_attention_layer_ids) + == n_local_layers + ), ( + "tri-pool KV split must cover every local attention layer exactly " + f"once: full={len(full_attention_layer_ids)} + " + f"swa={len(swa_attention_layer_ids)} != {n_local_layers} layers in " + f"[{self.layer_info.start_layer}, {self.layer_info.end_layer}) — " + "the ModelConfig full/swa split is wrong for this architecture" + ) + mamba_layer_ids = [ + i + for i in config.mamba2_cache_params.layers + if self.layer_info.start_layer <= i < self.layer_info.end_layer + ] + + return init_unified_mamba_swa_pools( + device=self.device, + kv_cache_dtype=self.kv_cache_dtype, + head_num=head_num, + head_dim=head_dim, + v_head_dim=v_head_dim, + swa_head_num=swa_head_num, + swa_head_dim=swa_head_dim, + swa_v_head_dim=swa_v_head_dim, + page_size=self.page_size, + start_layer=self.layer_info.start_layer, + end_layer=self.layer_info.end_layer, + swa_attention_layer_ids=swa_attention_layer_ids, + full_attention_layer_ids=full_attention_layer_ids, + mamba_layer_ids=mamba_layer_ids, + mamba2_cache_params=config.mamba2_cache_params, + full_max_total_num_tokens=full_max_total_num_tokens, + swa_max_total_num_tokens=swa_max_total_num_tokens, + max_mamba_cache_size=get_schedule().max_mamba_cache_size, + model_context_len=self.model_config.context_len, + extra_max_context_len=extra_max_context_len, + max_num_reqs=max_num_reqs, + enable_memory_saver=get_exec().features.enable_memory_saver, + enable_mamba_extra_buffer=self.server_args.enable_mamba_extra_buffer(), + disable_overlap_schedule=get_schedule().disable_overlap_schedule, + need_sort=get_disagg().disaggregation_mode in ("decode", "prefill"), + speculative_num_draft_tokens=get_spec().speculative_num_draft_tokens, + forward_stream=self.forward_stream, + 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 input (context len is already passed). + sliding_window_size=self.model_config.sliding_window_size, + ) + def _init_unified_swa_pools( self, *, diff --git a/python/sglang/srt/mem_cache/multi_ended_allocator.py b/python/sglang/srt/mem_cache/multi_ended_allocator.py index bad11c227..d8dbbd71a 100644 --- a/python/sglang/srt/mem_cache/multi_ended_allocator.py +++ b/python/sglang/srt/mem_cache/multi_ended_allocator.py @@ -25,7 +25,17 @@ from __future__ import annotations import inspect import logging import os -from typing import Callable, Dict, List, Optional, Sequence, Set, Tuple +from typing import ( + Callable, + Dict, + Generic, + List, + Optional, + Sequence, + Set, + Tuple, + TypeVar, +) import torch from torch.profiler import record_function @@ -95,9 +105,155 @@ def _install_signal_handlers_once() -> None: pass +_T = TypeVar("_T") + + +class _CapacityField(Generic[_T]): + """Data descriptor for a capacity-bearing allocator field. + + Every rebind bumps the owner's ``_capacity_epoch``, so the epoch-keyed + capacity memos (``available_size`` / ``schedulable_available_size`` on + every chain member plus the composite joint views) invalidate by + construction — mutation sites need no explicit hook, and future mutators + cannot forget one. Contract: these fields are REBOUND, never mutated in + place (all current writes are; ``_free_phys_pages`` slicing/cat/sort + always rebinds). + """ + + __slots__ = ("_name",) + + def __set_name__(self, owner, name: str) -> None: + self._name = name + + def __get__(self, obj, objtype=None) -> _T: + if obj is None: + return self # type: ignore[return-value] + try: + return obj.__dict__[self._name] + except KeyError: + raise AttributeError(self._name) from None + + def __set__(self, obj, value: _T) -> None: + obj.__dict__[self._name] = value + obj._capacity_epoch += 1 + + +def _float_open_short_side(flt, demand) -> None: + """THE float-relocate policy, driven by a DEMAND VECTOR -- one entry per + band, in PAGES of that band, zero for bands the operation does not touch + (e.g. mamba during a decode-token alloc). Any allocation event — a + band's own pages, a coupled token spanning several bands, or a future + combined admission vector — expresses itself the same way; nothing here + names a member or an operation. + + Each END band's unpayable remainder (demand − its drainable holes) lands + on the float band on ITS side (a grow-down end faces the float's HIGH + side, a grow-up end its LOW side); the float's own remainder F can + extend into either band. With surplus = band − end-demand per side: + + any demanded band's INDEX space too small -> skip (bytes cannot fix); + both sides short -> skip: relocation is ZERO-SUM between the bands + (opening one side closes the other) — the ladder falls through to + evict/retract; + one side short -> open exactly that side, folding F in after + crediting the far side's surplus; + only F short -> open the LARGER-surplus side by the remainder; + nothing short -> no relocation. + + `make_room`'s ``min_bytes`` is a TARGET for that side's whole band, so + the ask is demand + remainder + one page of slack (largest demanded + page) — never a delta, which under-asks whenever the band is partially + free. Best-effort: one relocation per ladder round, re-checked by the + caller; `make_room` leaves state untouched on an impossible ask. + """ + if flt is None or flt._is_frontier_transparent(): + return # no float involved / empty float never blocks + if not any(pages > 0 for pages in demand.values()): + return # nothing demanded — nothing to open (also keeps slack's max() total) + for band_alloc, pages in demand.items(): + if pages <= 0: + continue + index_room = ( + band_alloc.num_pages + - band_alloc.min_page_index + - band_alloc._allocated_pages() + ) + if pages > index_room: + return # index space binds; bytes cannot fix this + sides = {"low": 0, "high": 0} + for band_alloc, pages in demand.items(): + if band_alloc is flt or pages <= 0: + continue + holes = len(band_alloc._free_phys_pages) if band_alloc.lazy_compaction else 0 + ext = max(0, pages - holes) + side = "high" if band_alloc.grow_direction == "down" else "low" + sides[side] += ext * band_alloc.entry_bytes_per_page + band = { + "low": max( + 0, flt._byte_low_frontier() - flt._chain_high_frontier_below_bytes() + ), + "high": max( + 0, flt._chain_low_frontier_above_bytes() - flt._byte_high_frontier() + ), + } + surplus = {side: band[side] - sides[side] for side in ("low", "high")} + f_pages = demand.get(flt, 0) + f_bytes = max(0, f_pages - flt._hole_pages()) * flt.entry_bytes_per_page + slack = max(b.entry_bytes_per_page for b, pages in demand.items() if pages > 0) + if surplus["low"] < 0 and surplus["high"] < 0: + return # zero-sum: opening one side closes the other + if surplus["low"] < 0 or surplus["high"] < 0: + short, far = ("low", "high") if surplus["low"] < 0 else ("high", "low") + target = sides[short] + max(0, f_bytes - max(0, surplus[far])) + slack + if target > band[short]: + flt.make_room(side=short, min_bytes=target) + return + if f_bytes > max(surplus.values()): + short = "low" if surplus["low"] >= surplus["high"] else "high" + flt.make_room(side=short, min_bytes=sides[short] + f_bytes + slack) + + +def _relieve_for_alloc(short_pool, need_tokens: int) -> bool: + """THE shortfall ladder. Every allocation shortfall in the unified pool -- + a single band's own alloc, or a composite's coupled multi-band alloc — + runs exactly this, cheapest remedy first: + + 1. flush targets flush (absorb; ENDS also compact) + 2. enough? -> done + 3. the float, if one can help, slides (relocate) + 4. enough? -> done, else the caller evicts / retracts + + ``short_pool`` is the allocator that FAILED — a band when its own pages + ran out (e.g. mamba state slots), the composite when a coupled alloc + (one token = a page on EVERY member) missed its joint gate. It supplies + the two policies as methods, each documented where it is defined: + + _flush_targets() who can raise MY availability by flushing + _ask_float_for_room(N) how MY deficit maps to a float relocation + + `_flush` is called unconditionally: an eager END no-ops (it compacted at + free time) and a FLOAT always has boundary absorption to do — so the + ladder itself never branches on lazy mode, member kind, or layout. + """ + for m in short_pool._flush_targets(): + m._flush(urgent=True) + if need_tokens <= short_pool.available_size(): + return True + short_pool._ask_float_for_room(need_tokens) + return need_tokens <= short_pool.available_size() + + class MultiEndedAllocator(BaseTokenToKVPoolAllocator): """Allocator for one sub-pool over a `UnifiedKVPool`.""" + # Capacity-bearing state: any rebind bumps `_capacity_epoch`, invalidating + # the epoch-keyed capacity memos across the whole chain (see + # `_CapacityField` / `_chain_capacity_epoch`). + _capacity_epoch: int = 0 + watermark_physical: _CapacityField[int] = _CapacityField() + live_page_count: _CapacityField[int] = _CapacityField() + _free_phys_pages: _CapacityField[torch.Tensor] = _CapacityField() + def __init__( self, *, @@ -169,7 +325,10 @@ class MultiEndedAllocator(BaseTokenToKVPoolAllocator): # Back-compat alias (count of virtual PAGES) consulted by is_slot_allocated. self.num_virtual_ids = self.num_pages - self._peer: Optional[MultiEndedAllocator] = None + # Chain neighbours: `low_peer` toward byte 0, `high_peer` toward + # `total_bytes`. Ends have one (`bind_peer`), float middles have both. + self.low_peer: Optional[MultiEndedAllocator] = None + self.high_peer: Optional[MultiEndedAllocator] = None # Inverse history of relocations (spec rollback), at PAGE granularity. self._inverse_history: List[Tuple[torch.Tensor, torch.Tensor, torch.Tensor]] = ( @@ -233,6 +392,14 @@ class MultiEndedAllocator(BaseTokenToKVPoolAllocator): os.environ.get("SGLANG_LAZY_COMPACTION_MAX_MOVES_PER_CALL", "4096") ) + # Epoch-keyed memos for the capacity views -- pure functions of chain + # state between mutations, but schedulers read them O(queue) times per + # step (see `available_size` / `schedulable_available_size`). + self._avail_memo_epoch: Optional[int] = None + self._avail_memo_tokens: int = 0 + self._sched_avail_memo_epoch: Optional[int] = None + self._sched_avail_memo_tokens: int = 0 + self.clear() logger.info( @@ -254,23 +421,47 @@ class MultiEndedAllocator(BaseTokenToKVPoolAllocator): self.num_pages - self.min_page_index, ) - # -- peer binding -- + # -- chain-neighbor binding -- def bind_peer(self, peer: MultiEndedAllocator) -> None: - self._peer = peer + """2-pool END-pair compat: bind the OTHER end as this end's growth-side + neighbor (grow-up's neighbor sits above; grow-down's below). Float + middles must be wired explicitly — calling this on/with one raises. + """ + assert self.grow_direction in ("up", "down") and peer.grow_direction in ( + "up", + "down", + ), ( + f"bind_peer is END-pool-only; got {self.sub_pool_name!r} " + f"({self.grow_direction}) <-> {peer.sub_pool_name!r} " + f"({peer.grow_direction}); wire floats via bind_low_peer/bind_high_peer" + ) + if self.grow_direction == "up": + self.high_peer = peer + else: + self.low_peer = peer + self._capacity_epoch += 1 # rewiring changes what the chain walks see - @property - def peer(self) -> Optional[MultiEndedAllocator]: - return self._peer + def bind_low_peer(self, peer: MultiEndedAllocator) -> None: + self.low_peer = peer + self._capacity_epoch += 1 # rewiring changes what the chain walks see + + def bind_high_peer(self, peer: MultiEndedAllocator) -> None: + self.high_peer = peer + self._capacity_epoch += 1 # rewiring changes what the chain walks see # -- state -- - def clear(self) -> None: - """Reset to initial state. Pages in `[0, min_page_index)` are reserved.""" + def _reset_watermarks(self) -> None: + """Reset frontier state to empty (float middles override).""" if self.grow_direction == "up": self.watermark_physical = self.min_page_index else: self.watermark_physical = self.num_pages - 1 + + def clear(self) -> None: + """Reset to initial state. Pages in `[0, min_page_index)` are reserved.""" + self._reset_watermarks() self.virtual_to_physical.fill_(-1) # Virtual page 0 <-> physical page 0 (padding sink). self.virtual_to_physical[0] = 0 @@ -368,6 +559,32 @@ class MultiEndedAllocator(BaseTokenToKVPoolAllocator): f"[{self.sub_pool_name}] span {wm_span} != live " f"{self.live_page_count} + holes {holes} + pending {pending}" ) + out.extend(self._capacity_memo_violations()) + return out + + def _capacity_memo_violations(self) -> List[str]: + """Memo-coherence check (idle-time): a current-epoch capacity memo must + equal a fresh recompute; divergence means a mutation bypassed + `_CapacityField` (e.g. an in-place write). Empty == healthy.""" + out: List[str] = [] + epoch = self._chain_capacity_epoch() + if self._avail_memo_epoch == epoch: + actual = self._available_tokens() + if self._avail_memo_tokens != actual: + out.append( + f"[{self.sub_pool_name}] stale available_size memo: " + f"cached={self._avail_memo_tokens}, actual={actual}" + ) + if self._sched_avail_memo_epoch == epoch: + actual = self._available_tokens( + extra_gap_bytes=self._peer_drainable_hole_bytes() + ) + if self._sched_avail_memo_tokens != actual: + out.append( + f"[{self.sub_pool_name}] stale schedulable_available_size " + f"memo: cached={self._sched_avail_memo_tokens}, " + f"actual={actual}" + ) return out def _byte_low_frontier(self) -> int: @@ -376,19 +593,75 @@ class MultiEndedAllocator(BaseTokenToKVPoolAllocator): return self.min_page_index * self.entry_bytes_per_page return (self.watermark_physical + 1) * self.entry_bytes_per_page + # -- chain frontier walk -- + + def _is_frontier_transparent(self) -> bool: + """Whether neighbors' frontier walks may see THROUGH this pool. + + End pools are always opaque (an empty end's frontier already sits at + its buffer end, so opacity yields the correct gap). Float middles + override: an empty float occupies no bytes anywhere and must never + wall off free space. + """ + return False + + def _chain_low_frontier_above_bytes(self) -> int: + """Byte low-frontier of the nearest NON-transparent chain member above + this pool; the buffer top if none.""" + p = self.high_peer + while p is not None and p._is_frontier_transparent(): + p = p.high_peer + if p is None: + return self.unified_buffer.total_bytes + return p._byte_low_frontier() + + def _chain_high_frontier_below_bytes(self) -> int: + """Byte high-frontier of the nearest NON-transparent chain member below + this pool; 0 if none.""" + p = self.low_peer + while p is not None and p._is_frontier_transparent(): + p = p.low_peer + if p is None: + return 0 + return p._byte_high_frontier() + + def _chain_capacity_epoch(self) -> int: + """Sum of `_capacity_epoch` over the whole chain (self included). + + Capacity views read chain-neighbor frontiers (gap/transparency walks), + so a memo stays valid only while EVERY member is unmutated; the sum + moves whenever any member does (epochs only ever increment). + """ + total = self._capacity_epoch + p = self.low_peer + while p is not None: + total += p._capacity_epoch + p = p.low_peer + p = self.high_peer + while p is not None: + total += p._capacity_epoch + p = p.high_peer + return total + + def _growth_side_neighbor(self) -> Optional[MultiEndedAllocator]: + """Nearest NON-transparent chain member on this pool's GROWTH side -- + the one whose compaction/flush releases bytes reachable at this pool's + frontier.""" + p = self.high_peer if self.grow_direction == "up" else self.low_peer + while p is not None and p._is_frontier_transparent(): + p = p.high_peer if self.grow_direction == "up" else p.low_peer + return p + def _current_gap_bytes(self) -> int: - """Free byte band between this side's frontier and the peer's CURRENT frontier.""" + """Free byte band between this side's frontier and the nearest + non-transparent chain frontier (2-pool: the peer's, byte-identical).""" if self.grow_direction == "up": - my_high = self._byte_high_frontier() - peer_low = ( - self._peer._byte_low_frontier() - if self._peer is not None - else self.unified_buffer.total_bytes + return max( + 0, self._chain_low_frontier_above_bytes() - self._byte_high_frontier() ) - return max(0, peer_low - my_high) - my_low = self._byte_low_frontier() - peer_high = self._peer._byte_high_frontier() if self._peer is not None else 0 - return max(0, my_low - peer_high) + return max( + 0, self._byte_low_frontier() - self._chain_high_frontier_below_bytes() + ) def _available_tokens(self, extra_gap_bytes: int = 0) -> int: """Tokens allocatable given `extra_gap_bytes` of ADDED gap room @@ -412,38 +685,63 @@ class MultiEndedAllocator(BaseTokenToKVPoolAllocator): Alloc shortfall gates consult this to decide whether to peer-flush, so it MUST NOT fold in peer holes (use `schedulable_available_size()` for that). + Memoized on the chain capacity epoch (pure between mutations). """ - return self._available_tokens() + epoch = self._chain_capacity_epoch() + if self._avail_memo_epoch != epoch: + self._avail_memo_tokens = self._available_tokens() + self._avail_memo_epoch = epoch + return self._avail_memo_tokens def _peer_drainable_hole_bytes(self) -> int: - """Gap bytes a peer urgent flush would release. Only `_free_phys_pages` - count — NOT `_pending_reuse` (awaiting an event) — so the credit is realizable. + """Gap bytes an urgent flush of the growth-side chain neighbor would + release. Only `_free_phys_pages` count — NOT `_pending_reuse` (awaiting + an event) — so the credit is realizable. (2-pool: the peer's holes, + byte-identical.) """ - peer = self._peer - if peer is None or not peer.lazy_compaction: + neighbor = self._growth_side_neighbor() + if neighbor is None or not neighbor.lazy_compaction: return 0 - if peer.disagg_move_gate is not None and not peer.disagg_move_gate(): - # The peer cannot compact while a PD transfer is in flight, so these - # holes are not realizable. Crediting them would let the scheduler - # admit work that `_flush_peer_for_alloc` then cannot satisfy, and - # the caller treats a failed alloc as a memory-estimation bug. + if neighbor.disagg_move_gate is not None and not neighbor.disagg_move_gate(): + # Not realizable: a PD transfer blocks the neighbour's compaction. + # Crediting them admits work `_flush_peer_for_alloc` cannot satisfy, + # which the caller reads as a memory-estimation bug. return 0 - return len(peer._free_phys_pages) * peer.entry_bytes_per_page + return len(neighbor._free_phys_pages) * neighbor.entry_bytes_per_page def schedulable_available_size(self) -> int: - """Tokens allocatable AFTER a peer urgent-flush (realizable-with-compaction). - Used by composite views; alloc gates use `available_size()`. + """Tokens allocatable AFTER a neighbor urgent-flush (realizable-with- + compaction). Used by composite views; alloc gates use `available_size()`. + Memoized on the chain capacity epoch (pure between mutations). """ - return self._available_tokens(extra_gap_bytes=self._peer_drainable_hole_bytes()) + epoch = self._chain_capacity_epoch() + if self._sched_avail_memo_epoch != epoch: + self._sched_avail_memo_tokens = self._available_tokens( + extra_gap_bytes=self._peer_drainable_hole_bytes() + ) + self._sched_avail_memo_epoch = epoch + return self._sched_avail_memo_tokens - def _flush_peer_for_alloc(self, need_tokens: int) -> bool: - """One urgent peer-flush on alloc shortfall; returns whether THIS side now - has enough. Only PEER compaction releases gap bytes (own compaction is net 0). + def _flush_targets(self): + """A band short on its OWN alloc asks only its growth-side neighbour + to flush. Never itself: for its own allocation, holes and gap are + interchangeable (`take_physical_pages` drains holes first), so own + compaction trades one hole for one gap byte — net zero for self; only + a NEIGHBOUR's compaction releases bytes into the shared gap that own + extension consumes. """ - if not (self.lazy_compaction and self._peer is not None): - return False - self._peer._flush(urgent=True) - return need_tokens <= self.available_size() + neighbor = self._growth_side_neighbor() + return () if neighbor is None else (neighbor,) + + def _ask_float_for_room(self, need_tokens: int) -> None: + """A band short on its OWN pages: demand vector = {me: pages}; the + float, if the nearest non-transparent growth-side member is one, + opens the side facing me. Everything else — side derivation, index + guard, total-target ask -- is the shared policy.""" + blocker = self._growth_side_neighbor() + if not isinstance(blocker, FloatMultiEndedAllocator): + return + _float_open_short_side(blocker, {self: -(-need_tokens // self.page_size)}) # -- physical-slot / physical-page primitives -- @@ -536,32 +834,34 @@ class MultiEndedAllocator(BaseTokenToKVPoolAllocator): def _extend_watermark(self, num_pages: int) -> bool: """Advance the watermark by `num_pages` (lazy-path helper). Returns False - on index-space overflow OR crossing the PEER's byte frontier. + on index-space overflow OR crossing the nearest non-transparent chain + frontier. (Unbound chain side degenerates to the index-space check: the + walk returns the buffer end, whose page conversion equals `num_pages` / + 0 exactly — byte-identical to the old peerless branch.) """ if self.grow_direction == "up": new_wm = self.watermark_physical + num_pages if new_wm > self.num_pages: return False - # Peer (grow-down) sits ABOVE; don't extend past its low frontier. - if self._peer is not None: - peer_low_pages = ( - self._peer._byte_low_frontier() // self.entry_bytes_per_page - ) - if new_wm > peer_low_pages: - return False + # The chain above; don't extend past its low frontier. + chain_low_pages = ( + self._chain_low_frontier_above_bytes() // self.entry_bytes_per_page + ) + if new_wm > chain_low_pages: + return False self.watermark_physical = new_wm else: new_wm = self.watermark_physical - num_pages if new_wm < self.min_page_index - 1: return False - # Peer (grow-up) sits BELOW; `new_wm + 1` (our new lowest live page) - # must stay strictly above the peer's high frontier. - if self._peer is not None: - peer_high_pages = ( - self._peer._byte_high_frontier() // self.entry_bytes_per_page - ) - if new_wm + 1 < peer_high_pages: - return False + # `new_wm + 1` must stay strictly above the chain's high frontier + # below. Backstop only: callers gate on `available_size()`, whose + # floor'd gap already guarantees the extension fits. + chain_high_pages = ( + self._chain_high_frontier_below_bytes() // self.entry_bytes_per_page + ) + if new_wm + 1 < chain_high_pages: + return False self.watermark_physical = new_wm return True @@ -802,7 +1102,7 @@ class MultiEndedAllocator(BaseTokenToKVPoolAllocator): # Shortfall: flush the PEER, not own. Own compaction is net 0 # (each move trades 1 hole for +1 gap byte); only peer compaction # releases bytes into the shared gap that own extension consumes. - if not self._flush_peer_for_alloc(need_size): + if not _relieve_for_alloc(self, need_size): return None num_pages = need_size // self.page_size v_pages = self.free_virtual_ids[:num_pages] @@ -873,7 +1173,7 @@ class MultiEndedAllocator(BaseTokenToKVPoolAllocator): # compaction is internal — see `alloc`). need_tokens = num_new_pages * self.page_size if need_tokens > self.available_size(): - if not self._flush_peer_for_alloc(need_tokens): + if not _relieve_for_alloc(self, need_tokens): return None bs = len(prefix_lens) if self.need_sort and extend_num_tokens // self.page_size + bs + 1 > len( @@ -941,7 +1241,7 @@ class MultiEndedAllocator(BaseTokenToKVPoolAllocator): # Lazy: physical-capacity pre-check; on shortfall flush PEER. need_tokens = num_new_pages * self.page_size if need_tokens > self.available_size(): - if not self._flush_peer_for_alloc(need_tokens): + if not _relieve_for_alloc(self, need_tokens): return None if self.need_sort and bs > len(self.free_virtual_ids): self.merge_and_sort_free() @@ -1185,40 +1485,42 @@ class MultiEndedAllocator(BaseTokenToKVPoolAllocator): if src_list: src_pages = torch.tensor(src_list, dtype=torch.int64, device=self.device) dst_pages = torch.tensor(dst_list, dtype=torch.int64, device=self.device) - v_moved = self.physical_to_virtual[ - src_pages - ].clone() # read before clearing - - # Expand page ids to token ids for the token-granular move kernel. - if self.page_size == 1: - src_t, dst_t = src_pages, dst_pages - else: - offsets = torch.arange( - self.page_size, dtype=torch.int64, device=self.device - ) - src_t = (src_pages[:, None] * self.page_size + offsets).reshape(-1) - dst_t = (dst_pages[:, None] * self.page_size + offsets).reshape(-1) - - # Un-translated copy: the public copy_from translates virtual ids, - # which we must NOT do here. - move_fn = getattr(self._kvcache, "move_kv_cache", None) - if move_fn is not None: - move_fn(dst_t, src_t) - else: - copy_phys = getattr(self._kvcache, "_copy_from_physical", None) - assert copy_phys is not None, ( - f"sub-pool {self.sub_pool_name!r} supports neither move_kv_cache " - "nor _copy_from_physical" - ) - copy_phys(src_t, dst_t) - # Clear the vacated band, then re-bind the relocated dst pages. + # `dst` holes are outside the vacated band by construction, so + # rebinding them before the band wipe is order-equivalent. + self._move_pages_and_rebind(src_pages, dst_pages) self.physical_to_virtual[vacated_lo:vacated_hi] = -1 - self.virtual_to_physical[v_moved] = dst_pages - self.physical_to_virtual[dst_pages] = v_moved - self._inverse_history.append((src_pages, dst_pages, v_moved)) else: self.physical_to_virtual[vacated_lo:vacated_hi] = -1 + def _move_pages_and_rebind( + self, src_pages: torch.Tensor, dst_pages: torch.Tensor + ) -> torch.Tensor: + """Copy live pages src->dst (disjoint sets), rebind v2p/p2v for the + moved virtuals, and record inverse history. Does NOT clear p2v[src] — + callers own vacated-region clearing (end pools wipe the whole vacated + band; float middles clear exactly the src set). Returns the moved + virtual page ids. + """ + v_moved = self.physical_to_virtual[src_pages].clone() # read pre-wipe + + # Expand page ids to token ids for the token-granular move kernel. + if self.page_size == 1: + src_t, dst_t = src_pages, dst_pages + else: + offsets = torch.arange( + self.page_size, dtype=torch.int64, device=self.device + ) + src_t = (src_pages[:, None] * self.page_size + offsets).reshape(-1) + dst_t = (dst_pages[:, None] * self.page_size + offsets).reshape(-1) + + # Un-translated copy: the public copy_from translates virtual ids, + # which we must NOT do here. + self._kvcache.move_kv_cache(dst_t, src_t) + self.virtual_to_physical[v_moved] = dst_pages + self.physical_to_virtual[dst_pages] = v_moved + self._inverse_history.append((src_pages, dst_pages, v_moved)) + return v_moved + # -- lazy compaction primitives -- def set_latest_forward_done_event(self, event: Optional[torch.cuda.Event]) -> None: @@ -1708,16 +2010,7 @@ class MultiEndedAllocator(BaseTokenToKVPoolAllocator): ) src_t = (src_pages_t[:, None] * self.page_size + offsets).reshape(-1) dst_t = (dst_pages_t[:, None] * self.page_size + offsets).reshape(-1) - move_fn = getattr(self._kvcache, "move_kv_cache", None) - if move_fn is not None: - move_fn(dst_t, src_t) - else: - copy_phys = getattr(self._kvcache, "_copy_from_physical", None) - assert copy_phys is not None, ( - f"sub-pool {self.sub_pool_name!r} supports neither " - "move_kv_cache nor _copy_from_physical" - ) - copy_phys(src_t, dst_t) + self._kvcache.move_kv_cache(dst_t, src_t) # ONE bulk remap (single-writer on schedule_stream). self.virtual_to_physical[v_moveds_t] = dst_pages_t self.physical_to_virtual[dst_pages_t] = v_moveds_t @@ -1783,14 +2076,16 @@ def _chain_byte_accounting_violations( 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). + Transparent members (an empty/parked float occupies no bytes anywhere) + are skipped by the ordering walk — their per-pool conservation still runs. """ out: List[str] = [] for a in chain: out.extend(a._byte_accounting_violations()) frontier = 0 for a in chain: + if a._is_frontier_transparent(): + continue lo_b, hi_b = a._byte_low_frontier(), a._byte_high_frontier() if lo_b < frontier: out.append( @@ -1809,6 +2104,592 @@ def _end_pair_chain( return sorted((a, b), key=lambda x: x.grow_direction != "up") +class FloatMultiEndedAllocator(MultiEndedAllocator): + """Float MIDDLE cache pool: a span ``[low_wm_page, high_wm_page)`` between + two chain neighbors, with freed HOLES allowed inside the span. + + Holes-first model (a middle CACHE pool is not a band): + - ``free`` marks interior holes (zero copies) and absorbs boundary holes; + - alloc reuses holes first (zero copies — steady-state churn recycles in + place), then extends the boundary on the side with the LARGER free gap; + from empty it positions the span at the MIDPOINT of the inter-frontier + region, so free gap exists on both sides and neighbor growth does not + immediately force a data move; + - data moves happen only ON DEMAND: ``make_room(side, min_bytes)`` opens + contiguous space on ``side`` by relocating live boundary pages into + interior holes / the far gap (cost min(L_live, G): when the demand + exceeds the live bytes this degenerates into moving every live page — + the whole-pool leapfrog); ``compact_holes`` closes all holes, shrinking + the span from a chosen side. + - An EMPTY float (no live pages) resets its span and is + frontier-transparent: it occupies no bytes and must never wall off free + space (its parked position is irrelevant to neighbors). + + Floats skip the lazy event pipeline (`lazy_compaction` must be False): + frees/allocs are zero-copy by design, so only the on-demand moves need + write-set safety, which their scheduler-phase call sites provide. + """ + + # The span IS this pool's capacity state (it has no watermark): moving it + # changes its own availability and, through transparency, both neighbours' + # gaps. Same `_CapacityField` contract as the ends' `watermark_physical`. + low_wm_page: _CapacityField[int] = _CapacityField() + high_wm_page: _CapacityField[int] = _CapacityField() + + # Only `free` can make a boundary page a hole (alloc drains holes into live + # pages, extension adds live ones), so a clean flag proves both boundaries + # are live and the deferred absorb skips its D2H. Relocation re-arms it. + _holes_dirty: bool = False + + def __init__(self, **kwargs): + assert not kwargs.get("lazy_compaction", False), ( + "FloatMultiEndedAllocator is holes-first; the lazy event pipeline " + "is end-pool machinery and must stay off for float middles" + ) + # Base __init__ ends with self.clear(), which reads these via our + # _reset_watermarks override -- pre-seed so the override can run. + self.low_wm_page = 0 + self.high_wm_page = 0 + super().__init__(**kwargs) + assert self.grow_direction == "float", ( + f"FloatMultiEndedAllocator needs a 'float' sub-pool spec; got " + f"{self.grow_direction!r}" + ) + + # -- span / frontier state -- + + def _reset_watermarks(self) -> None: + # Park empty at the buffer top; empty-transparency makes the parked + # position irrelevant to neighbors. + self.low_wm_page = self.num_pages + self.high_wm_page = self.num_pages + self.watermark_physical = -1 # unused for float pools (logs only) + + def _span_pages(self) -> int: + return self.high_wm_page - self.low_wm_page + + def _hole_pages(self) -> int: + return int(self._free_phys_pages.numel()) + + def _live_pages(self) -> int: + return self._span_pages() - self._hole_pages() + + def _is_frontier_transparent(self) -> bool: + return self._live_pages() == 0 + + def _allocated_pages(self) -> int: + return self._live_pages() + + def _byte_low_frontier(self) -> int: + return self.low_wm_page * self.entry_bytes_per_page + + def _byte_high_frontier(self) -> int: + return self.high_wm_page * self.entry_bytes_per_page + + def _region_bounds_pages(self) -> Tuple[int, int]: + """Page bounds ``[lo, hi)`` of the inter-frontier region available to + this float (chain-transparent walk; clamped to the slot-0 sink + reservation). Rounded conservatively: ``lo`` up, ``hi`` down.""" + epp = self.entry_bytes_per_page + lo = (self._chain_high_frontier_below_bytes() + epp - 1) // epp + lo = max(lo, self.min_page_index) + hi = self._chain_low_frontier_above_bytes() // epp + hi = min(hi, self.num_pages) + return lo, hi + + def pages_in_band(self, *, low_byte: int, high_byte: int) -> int: + """Pages obtainable from ``[low_byte, high_byte)`` on this pool's OWN + page grid. A raw ``(high - low) // entry_bytes_per_page`` over-counts by + a page whenever ``low_byte`` is off the grid, which is the generic case: + the bounding frontier is a multiple of the NEIGHBOUR's entry size. + """ + epp = self.entry_bytes_per_page + lo = max((low_byte + epp - 1) // epp, self.min_page_index) + hi = min(high_byte // epp, self.num_pages) + return max(0, hi - lo) + + def _gap_pages(self) -> Tuple[int, int]: + """(gap_low, gap_high) in own page units; both == the whole region + when the span is empty/parked.""" + lo, hi = self._region_bounds_pages() + if self._is_frontier_transparent(): + room = max(0, hi - lo) + return room, room + return max(0, self.low_wm_page - lo), max(0, hi - self.high_wm_page) + + # -- availability -- + + def _side_drainable_hole_bytes(self, side: str) -> int: + """Realizable gap bytes an urgent flush of the neighbour on ``side`` + would release, walking past transparent members like the frontier walk. + """ + p = self.low_peer if side == "low" else self.high_peer + while p is not None and p._is_frontier_transparent(): + p = p.low_peer if side == "low" else p.high_peer + if p is None or not p.lazy_compaction: + return 0 + if p.disagg_move_gate is not None and not p.disagg_move_gate(): + return 0 + return len(p._free_phys_pages) * p.entry_bytes_per_page + + def _peer_drainable_hole_bytes(self) -> int: + """The better of the two sides. `_growth_side_neighbor()` is undefined + for a float -- its `grow_direction` is "float", so the base answers + `low_peer` and never sees the high neighbour. + """ + return max( + self._side_drainable_hole_bytes("low"), + self._side_drainable_hole_bytes("high"), + ) + + def _available_tokens(self, extra_gap_bytes: int = 0) -> int: + gap_low, gap_high = self._gap_pages() + if extra_gap_bytes > 0: + # Per side: the base hands down one undirected scalar because an + # END pool grows one way, but a float grows both. + epp = self.entry_bytes_per_page + gap_low += self._side_drainable_hole_bytes("low") // epp + gap_high += self._side_drainable_hole_bytes("high") // epp + gap_pages = max(gap_low, gap_high) # a single alloc extends ONE side + pages_by_index_space = self.num_pages - self.min_page_index - self._live_pages() + pages_extend = min(gap_pages, pages_by_index_space) + return (pages_extend + self._hole_pages()) * self.page_size + + # -- physical page primitives (holes-first) -- + + def take_physical_pages(self, num_pages: int) -> Optional[torch.Tensor]: + if num_pages <= 0: + return torch.empty(0, dtype=torch.int64, device=self.device) + n_drain = min(num_pages, self._hole_pages()) + need_more = num_pages - n_drain + + fresh: Optional[torch.Tensor] = None + if need_more > 0: + lo, hi = self._region_bounds_pages() + if self._is_frontier_transparent(): + # Reposition-on-alloc-from-empty: collapse to the midpoint so + # free gap remains on BOTH sides. + if need_more > hi - lo: + return None + start = lo + (hi - lo - need_more) // 2 + self.low_wm_page = start + self.high_wm_page = start + need_more + fresh = torch.arange( + start, start + need_more, dtype=torch.int64, device=self.device + ) + else: + gap_low = self.low_wm_page - lo + gap_high = hi - self.high_wm_page + # Extend toward the roomier gap; fall back to the other side. + sides = ("high", "low") if gap_high >= gap_low else ("low", "high") + for side in sides: + if side == "high" and need_more <= gap_high: + start = self.high_wm_page + self.high_wm_page += need_more + break + if side == "low" and need_more <= gap_low: + start = self.low_wm_page - need_more + self.low_wm_page = start + break + else: + return None # neither side fits; state untouched + fresh = torch.arange( + start, start + need_more, dtype=torch.int64, device=self.device + ) + + if n_drain > 0: + drained = self._free_phys_pages[:n_drain].clone() + self._free_phys_pages = self._free_phys_pages[n_drain:] + else: + drained = None + + if drained is None: + return fresh + if fresh is None: + return drained + return torch.cat([drained, fresh]) + + def take_physical(self, need_size: int) -> Optional[torch.Tensor]: + if need_size <= 0: + return torch.empty(0, dtype=torch.int64, device=self.device) + assert need_size % self.page_size == 0, ( + f"take_physical: need_size={need_size} must be a multiple of " + f"page_size={self.page_size}" + ) + return self.take_physical_pages(need_size // self.page_size) + + def _alloc_bind_fast_or_slow( + self, v_pages: torch.Tensor, N: int + ) -> Optional[torch.Tensor]: + # Holes-first always routes through take_physical_pages (no fused + # watermark fast path -- float alloc cadence doesn't need it). + if N == 0: + return torch.empty(0, dtype=torch.int64, device=self.device) + phys_pages = self.take_physical_pages(N) + if phys_pages is None: + return None + self.bind(v_pages, phys_pages) + return phys_pages + + # -- free: hole-marking, boundary absorption, park-on-empty -- + + def free( + self, free_index: torch.Tensor, *, _pages: Optional[torch.Tensor] = None + ) -> None: + """Mark the freed pages as interior HOLES / absorb boundary ones. + + `_pages` carries virtual PAGE ids the caller already derived (segment + frees from `start_pos` arithmetic; the SWA composite's page-rep + release) — same contract as the base allocator, and it must be + honoured here for the same reason: deriving them again via + `torch.unique` is a data-dependent-shape op, i.e. a HOST SYNC on the + per-step free path. + """ + with record_function("FloatMultiEndedAlloc.free"): + if free_index is None or free_index.numel() == 0: + return + if self.free_group is not None: + self.free_group.append(self._copy_for_free_group(free_index)) + return + # Page-derivation ladder, as `_free_lazy`: caller ids, the ps==1 + # identity, then dedup. No stale-slot assert -- callers must not + # double-free (a tombstoned page would join the hole list); the + # composite's filters uphold it and the byte verifier catches a miss. + free_v_pages_raw = free_index.detach().to(torch.int64) + if _pages is not None: + free_v_pages = _pages + elif self.page_size == 1: + free_v_pages = free_v_pages_raw + else: + free_v_pages = torch.unique(free_v_pages_raw // self.page_size) + freed_p_pages = self.virtual_to_physical[free_v_pages] + # `index_fill_`, never `t[idx] = -1`: see the END free path. + self.virtual_to_physical.index_fill_(0, free_v_pages, -1) + self.physical_to_virtual.index_fill_(0, freed_p_pages, -1) + if self.is_id_owner: + self.free_virtual_ids = torch.cat([self.free_virtual_ids, free_v_pages]) + self._free_phys_pages = torch.cat([self._free_phys_pages, freed_p_pages]) + # Park is sync-free (span/hole COUNTS only); boundary absorption is + # DEFERRED -- see `_absorb_span_boundary_holes`. + self._holes_dirty = True + self._park_if_empty() + + def _park_if_empty(self) -> bool: + """Reset the span and go frontier-transparent once no live page + remains. Sync-free: `_live_pages()` is span minus hole COUNT, both + host-side (`numel()` is tensor metadata). Returns whether it parked.""" + if self._live_pages() != 0: + return False + self._reset_watermarks() + self._free_phys_pages = torch.empty(0, dtype=torch.int64, device=self.device) + self._holes_dirty = False + return True + + def _absorb_span_boundary_holes(self) -> int: + """Shrink the span past any holes touching its boundaries (zero-copy), + returning the number of pages handed back to the neighbours. + + DEFERRED, not per-free: deciding how far to walk needs the hole set on + the HOST (the watermarks are host ints), so this is the float's one + D2H — exactly the base allocator's model, whose `_free_lazy` does "no + boundary absorb" and pays a single sync inside `_flush`. Doing it per + free put a host sync on the per-decode-step path. + + Called where a sync is already free or already warranted: the per-step + opportunistic flush (the scheduler runs it at the sync boundary with + the forward stream drained) and the head of the tri's shortfall ladder + (a stale-wide span would otherwise inflate the rebalance deficit and + buy data movement that this zero-copy shrink makes unnecessary). + + Skipping it is only ever CONSERVATIVE: the span reads wider than its + live content, so neighbours see less gap. `_live_pages()`, hence + transparency and the byte-conservation identity, stay exact either way. + """ + if self._park_if_empty(): + self._holes_dirty = False + return 0 + if not self._holes_dirty or self._free_phys_pages.numel() == 0: + # Nothing freed since the last absorb => both boundaries are still + # live => the walk provably finds nothing. Skip the D2H; steady + # churn with only INTERIOR holes then costs no sync at all. + self._holes_dirty = False + return 0 + self._holes_dirty = False + before = self._span_pages() + holes = set(int(x) for x in self._free_phys_pages.tolist()) + changed = False + while self.low_wm_page in holes: + holes.remove(self.low_wm_page) + self.low_wm_page += 1 + changed = True + while (self.high_wm_page - 1) in holes: + holes.remove(self.high_wm_page - 1) + self.high_wm_page -= 1 + changed = True + if changed: + self._free_phys_pages = torch.tensor( + sorted(holes), dtype=torch.int64, device=self.device + ) + return before - self._span_pages() + + # -- on-demand data movement -- + + def make_room(self, *, side: str, min_bytes: int) -> int: + """Open >= ``min_bytes`` of CONTIGUOUS free space between this pool's + ``side`` boundary and the region bound on that side, relocating the + minimum set of live boundary pages (holes-first destinations, then the + far gap). Returns the bytes now open on ``side`` (may exceed the ask; + < min_bytes iff impossible now — state is then unchanged). + + Cost model: moving k pages costs k page-copies; k <= min(L_live, G). + Scheduler-phase only. Stream safety is owned HERE, not by the caller: + the entry settles the in-flight forward before the first copy. + """ + assert side in ("low", "high"), f"side must be 'low'|'high'; got {side!r}" + # Order the copies after the in-flight forward, or they carry pre-write + # bytes and the rebind sends readers to a destination that never got + # them. One wait covers read AND write: the event is post-forward. + self._settle_inflight_forward() + epp = self.entry_bytes_per_page + lo, hi = self._region_bounds_pages() + gap_low, gap_high = self._gap_pages() + gap_side_bytes = (gap_low if side == "low" else gap_high) * epp + if gap_side_bytes >= min_bytes or self._is_frontier_transparent(): + return gap_side_bytes + + # Capacity: even packing every live page flush against the far side + # cannot open more than (region - live) bytes. + live = self._live_pages() + if (hi - lo - live) * epp < min_bytes: + return gap_side_bytes # impossible now; untouched + + need_pages = (min_bytes - gap_side_bytes + epp - 1) // epp + + holes = set(int(x) for x in self._free_phys_pages.tolist()) + span = range(self.low_wm_page, self.high_wm_page) + live_pages_sorted = [p for p in span if p not in holes] + + if need_pages >= live: + # Whole-pool LEAPFROG: pack every live page flush against the far + # region edge (cost L_live <= G); the capacity check above + # guarantees the resulting gap satisfies the ask. + if side == "high": + final = list(range(lo, lo + live)) + else: + final = list(range(hi - live, hi)) + self._relocate_to_positions(live_pages_sorted, final) + gap_low2, gap_high2 = self._gap_pages() + return (gap_low2 if side == "low" else gap_high2) * epp + + # Boundary relocation (G < L_live): + # Sources: live pages nearest the demanded side, retreating inward. + if side == "high": + srcs = list(reversed(live_pages_sorted))[: min(need_pages, live)] + else: + srcs = live_pages_sorted[: min(need_pages, live)] + src_set = set(srcs) + + # Strictly on the far side of EVERY source: keeps the batched move + # src/dst-disjoint and actually retreats the edge. + if side == "high": + usable_holes = sorted(h for h in holes if h < min(srcs)) + else: + usable_holes = sorted((h for h in holes if h > max(srcs)), reverse=True) + dsts: List[int] = list(usable_holes[: len(srcs)]) + n_fresh = len(srcs) - len(dsts) + if n_fresh > 0: + # Far-gap feasibility for the fresh destinations. + if side == "high": + if n_fresh > gap_low: + return gap_side_bytes + dsts += list( + range(self.low_wm_page - 1, self.low_wm_page - 1 - n_fresh, -1) + ) + else: + if n_fresh > gap_high: + return gap_side_bytes + dsts += list(range(self.high_wm_page, self.high_wm_page + n_fresh)) + + if srcs: + self._move_pages_and_rebind( + torch.tensor(srcs, dtype=torch.int64, device=self.device), + torch.tensor(dsts, dtype=torch.int64, device=self.device), + ) + self.physical_to_virtual.index_fill_( + 0, torch.tensor(srcs, dtype=torch.int64, device=self.device), -1 + ) + + # Reconstruct the span from final live positions: everything between + # the extremes is span; non-live pages inside are holes. + final_live = sorted((set(live_pages_sorted) - src_set) | set(dsts)) + self.low_wm_page = final_live[0] + self.high_wm_page = final_live[-1] + 1 + final_live_set = set(final_live) + new_holes = [ + p + for p in range(self.low_wm_page, self.high_wm_page) + if p not in final_live_set + ] + self._free_phys_pages = torch.tensor( + new_holes, dtype=torch.int64, device=self.device + ) + self._holes_dirty = True # the span moved; re-check its boundaries + self._absorb_span_boundary_holes() + + gap_low2, gap_high2 = self._gap_pages() + return (gap_low2 if side == "low" else gap_high2) * epp + + def _relocate_to_positions(self, live_sorted: List[int], final: List[int]) -> int: + """Order-preserving relocation of the live pages onto the ``final`` + positions (an ascending hole-free block). Batched disjoint move when + possible; otherwise ORDERED singleton moves (uniform shift direction: + each destination is a hole or an already-vacated source by induction). + Sets span to the final block, clears holes. Returns pages moved. + """ + assert len(live_sorted) == len(final) + pairs = [(s, d) for s, d in zip(live_sorted, final) if s != d] + if pairs: + src_set = {s for s, _ in pairs} + dst_set = {d for _, d in pairs} + if src_set.isdisjoint(dst_set): + src_t = torch.tensor( + [s for s, _ in pairs], dtype=torch.int64, device=self.device + ) + dst_t = torch.tensor( + [d for _, d in pairs], dtype=torch.int64, device=self.device + ) + self._move_pages_and_rebind(src_t, dst_t) + self.physical_to_virtual.index_fill_(0, src_t, -1) + else: + # Overlapping shift: process toward the move direction so each + # destination is free by the time it is written. + ordered = pairs if final[0] <= live_sorted[0] else list(reversed(pairs)) + # Built ONCE: `torch.tensor(..., device=cuda)` in the loop is a + # pageable H2D per page, and a shift can span the whole pool. + src_all = torch.tensor( + [s for s, _ in ordered], dtype=torch.int64, device=self.device + ) + dst_all = torch.tensor( + [d for _, d in ordered], dtype=torch.int64, device=self.device + ) + for i in range(len(ordered)): + s_t, d_t = src_all[i : i + 1], dst_all[i : i + 1] + self._move_pages_and_rebind(s_t, d_t) + self.physical_to_virtual.index_fill_(0, s_t, -1) + if final: + self.low_wm_page = final[0] + self.high_wm_page = final[-1] + 1 + else: + self._reset_watermarks() + self._free_phys_pages = torch.empty(0, dtype=torch.int64, device=self.device) + return len(pairs) + + def _byte_accounting_violations(self) -> List[str]: + out: List[str] = [] + total = self.unified_buffer.total_bytes + lo_b, hi_b = self._byte_low_frontier(), self._byte_high_frontier() + if not self._is_frontier_transparent() and not (0 <= lo_b <= hi_b <= total): + out.append( + f"[{self.sub_pool_name}] float span out of bounds: " + f"low={lo_b}, high={hi_b}, total={total}" + ) + # Independent live count from the p2v table (`_live_pages()` is + # DERIVED as span - holes, so checking against it would be circular): + # every span page must be either p2v-bound or an interior hole. + if self._span_pages() > 0: + bound = int( + (self.physical_to_virtual[self.low_wm_page : self.high_wm_page] != -1) + .sum() + .item() + ) + if self._span_pages() != bound + self._hole_pages(): + out.append( + f"[{self.sub_pool_name}] float span {self._span_pages()} != " + f"p2v-bound {bound} + holes {self._hole_pages()}" + ) + out.extend(self._capacity_memo_violations()) + return out + + def _flush(self, *, urgent: bool) -> int: + """Boundary absorption only -- never data movement. The base `_flush` + treats `_free_phys_pages` as a lazy compaction backlog to be drained, + but for a float those entries are INTERIOR HOLES, reusable assets by + design; relocation happens on demand via `make_room` / + `compact_holes`. What a float CAN do at a flush point is hand back + span it no longer needs, which is where its deferred D2H belongs — so + neighbours' urgent-flush ladders and the per-step opportunistic flush + both reclaim the boundary holes.""" + return self._absorb_span_boundary_holes() + + def flush_opportunistic(self) -> int: + """Public gated wrapper around `_flush(urgent=False)` -- the base's + exact shape. The ONLY reason for the override is the gate: the base + keys its fast path on `lazy_compaction`, which a float never has; a + float's flushable work is its deferred boundary absorption, so the + fast path keys on `_holes_dirty` instead. The scheduler calls this at + the sync boundary with the forward stream drained, so the D2H the + flush costs is the cheapest one available; the clean fast path keeps + the common step sync-free.""" + with record_function("FloatMultiEndedAlloc.flush_opportunistic"): + if not self._holes_dirty or self._free_phys_pages.numel() == 0: + return 0 + return self._flush(urgent=False) + + def backup_state(self): + # Span-aware snapshot (base backs up watermark_physical, meaningless + # here). Spec decode is asserted off under unified today; kept correct + # for when the gate lifts. + return ( + self.low_wm_page, + self.high_wm_page, + self._free_phys_pages.clone(), + (len(self.free_virtual_ids) if self.is_id_owner else None), + len(self._inverse_history), + ) + + def restore_state(self, state): + low_wm, high_wm, holes, _n_free_virtual, n_inverse = state + self.low_wm_page = low_wm + self.high_wm_page = high_wm + self._free_phys_pages = holes + new_entries = self._inverse_history[n_inverse:] + if new_entries: + logger.warning( + "FloatMultiEndedAllocator.restore_state: %d relocation(s) inside " + "a backup window (sub_pool=%s) — float moves are not reversible.", + len(new_entries), + self.sub_pool_name, + ) + del self._inverse_history[n_inverse:] + return new_entries + + def compact_holes(self, *, retreat_side: str) -> int: + """Close ALL interior holes by packing live pages toward the side + OPPOSITE ``retreat_side`` (order-preserving), shrinking the span on + ``retreat_side`` by the hole count. Returns pages moved.""" + assert retreat_side in ("low", "high") + if self._hole_pages() == 0: + return 0 + # Settle before the first copy -- see `make_room`. + self._settle_inflight_forward() + holes = set(int(x) for x in self._free_phys_pages.tolist()) + live_sorted = [ + p for p in range(self.low_wm_page, self.high_wm_page) if p not in holes + ] + if retreat_side == "high": + final = list(range(self.low_wm_page, self.low_wm_page + len(live_sorted))) + else: + final = list(range(self.high_wm_page - len(live_sorted), self.high_wm_page)) + return self._relocate_to_positions(live_sorted, final) + + # -- band-incompatible base APIs -- + + def bind_peer(self, peer: MultiEndedAllocator) -> None: # pragma: no cover + raise AssertionError( + "float middles must be wired via bind_low_peer/bind_high_peer" + ) + + class UnifiedMambaTokenToKVPoolAllocator(BaseTokenToKVPoolAllocator): """Composite allocator for the MHA (full-attn) + Mamba hybrid pair. @@ -2236,19 +3117,21 @@ class UnifiedSWATokenToKVPoolAllocator(SWATokenToKVPoolAllocator): forward_stream=forward_stream, lazy_compaction=lazy_compaction, ) - self.swa_attn_allocator = MultiEndedAllocator( + self.swa_attn_allocator = self._build_swa_attn_allocator( kvcache=kvcache.swa_kv_pool, unified_buffer=unified_buffer, - sub_pool_name="swa", device=device, - is_id_owner=False, # non-owner; consumes virtuals minted by full page_size=page_size, need_sort=need_sort, forward_stream=forward_stream, lazy_compaction=lazy_compaction, ) - self.full_attn_allocator.bind_peer(self.swa_attn_allocator) - self.swa_attn_allocator.bind_peer(self.full_attn_allocator) + self._wire_peers() + + # 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 + self._joint_avail_memo_tokens: int = 0 # The full/SWA KV pools need no allocator wiring (write locations resolved # in attention metadata); the composite keeps allocators for read-path translates. @@ -2279,12 +3162,40 @@ class UnifiedSWATokenToKVPoolAllocator(SWATokenToKVPoolAllocator): self.available_size(), ) + # -- construction hooks (the tri-pool subclass overrides both) -- + + def _build_swa_attn_allocator(self, **kwargs) -> MultiEndedAllocator: + """The swa sub-allocator: an END pool here (2-pool pair); the tri-pool + subclass overrides to build the swa FLOAT middle instead.""" + return MultiEndedAllocator( + sub_pool_name="swa", + is_id_owner=False, # non-owner; consumes virtuals minted by full + **kwargs, + ) + + def _wire_peers(self) -> None: + """2-pool end-pair wiring; the tri-pool subclass wires the full chain + (mamba end <-> swa float <-> full end) after its mamba end exists.""" + 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: """Tokens available for `alloc(N)` / `alloc_extend(N)` (TOKENS). - Joint byte-budget: each composite alloc(1) consumes one full-side AND one + Memoized on the chain capacity epoch (the compute walks every chain + frontier; see `_compute_available_size`, which the tri-pool subclass + overrides with its three-band variant). + """ + epoch = self.full_attn_allocator._chain_capacity_epoch() + if self._joint_avail_memo_epoch != epoch: + self._joint_avail_memo_tokens = self._compute_available_size() + self._joint_avail_memo_epoch = epoch + return self._joint_avail_memo_tokens + + def _compute_available_size(self) -> int: + """Joint byte-budget: each composite alloc(1) consumes one full-side AND one swa-side page (same virtual id). The 3-phase lazy formula consumes both sides' holes maximally before extending toward the gap (H_f/H_s = holes, e_f/e_s = bytes/page, R_f/R_s = extension room, G = byte gap): @@ -2360,6 +3271,16 @@ class UnifiedSWATokenToKVPoolAllocator(SWATokenToKVPoolAllocator): self.schedulable_swa_available_size(), ) + # Slot-conservation views for the LEAK INVARIANT only, which pairs the static + # per-layer total with (static cap - live). Schedulers keep the `min(...)` + # views above: under the floating boundary the byte term dips below the + # conserve cap, so bytes lent to a peer sub-pool would read as a leak. + def conserve_full_available_size(self) -> int: + return self._conserve_full_available_size() + + def conserve_swa_available_size(self) -> int: + return self._conserve_swa_available_size() + # Byte-coordinated, realizable-with-compaction views (peer drainable holes # credited — see `MultiEndedAllocator.schedulable_available_size`). def schedulable_full_available_size(self) -> int: @@ -2368,16 +3289,19 @@ class UnifiedSWATokenToKVPoolAllocator(SWATokenToKVPoolAllocator): def schedulable_swa_available_size(self) -> int: return self.swa_attn_allocator.schedulable_available_size() - def _flush_both_for_alloc(self, need_tokens: int) -> bool: - """SWA analogue of `_flush_peer_for_alloc`. Each composite alloc consumes a - full AND a swa page and either side's compaction opens gap for the other, - so flush BOTH (one urgent pass each). + def _flush_targets(self): + """A coupled alloc consumes a page on EVERY member under one virtual + id, so a hole on ONE side is unusable once the gap is dry — there is + nothing on the other side to pair it with. Each member's compaction + converts such dead one-sided holes into SHARED gap, which serves the + joint gate: flush ALL members, including ones that are themselves + short. """ - if not self.lazy_compaction: - return need_tokens <= self.available_size() - self.full_attn_allocator._flush(urgent=True) - self.swa_attn_allocator._flush(urgent=True) - return need_tokens <= self.available_size() + 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 # `size_full` / `size_swa` are inherited; they read `_size_full`/`_size_swa` # (set to the static caps). We do NOT report `max_slots - 1`: under unified @@ -2453,7 +3377,7 @@ class UnifiedSWATokenToKVPoolAllocator(SWATokenToKVPoolAllocator): # 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 self._flush_both_for_alloc(need_size): + if not _relieve_for_alloc(self, need_size): return None # Snapshot the virtual PAGES full will consume, to bind them on swa too. num_pages = need_size // self.page_size @@ -2491,7 +3415,7 @@ class UnifiedSWATokenToKVPoolAllocator(SWATokenToKVPoolAllocator): ) need_tokens = num_new_pages * self.page_size if need_tokens > self.available_size(): - if not self._flush_both_for_alloc(need_tokens): + if not _relieve_for_alloc(self, need_tokens): return None # Snapshot the virtual PAGES the kernel will consume; clone so swa keeps @@ -2530,7 +3454,7 @@ class UnifiedSWATokenToKVPoolAllocator(SWATokenToKVPoolAllocator): ) need_tokens = num_new_pages * self.page_size if need_tokens > self.available_size(): - if not self._flush_both_for_alloc(need_tokens): + if not _relieve_for_alloc(self, need_tokens): return None fa = self.full_attn_allocator @@ -2579,24 +3503,52 @@ class UnifiedSWATokenToKVPoolAllocator(SWATokenToKVPoolAllocator): self.full_attn_allocator.clear_inverse_history() self.swa_attn_allocator.clear_inverse_history() - def free_swa(self, free_index: torch.Tensor) -> None: + def free_swa( + self, free_index: torch.Tensor, *, start_pos: Optional[int] = None + ) -> None: """SWA tombstone path: release swa-physical, leave virtual id and - full-physical live. Called by `SWARadixCache._evict_swa_only` when a node - ages past the sliding-window horizon. `swa.v2p_page[v_page] = -1` IS the - tombstone. + full-physical live. Called by the per-step window ratchet and by radix + SWA eviction when a node ages past the sliding-window horizon. + `swa.v2p_page[v_page] = -1` IS the tombstone. + + ``start_pos`` is the `free_segment` contract: when the caller frees a + CONTIGUOUS ascending range whose first token sits at prefix position + `start_pos` (the window ratchet does — host-int, page-aligned bounds), + page representatives come from stride arithmetic and the swa side is + freed with caller-supplied page ids — no `torch.unique`, keeping the + per-decode-step free host-sync-free. Without it (radix eviction hands + arbitrary node values) the swa side falls back to its own dedup. """ if free_index is None or free_index.numel() == 0: return - # Keep only tokens whose virtual PAGE is still bound on swa (calling - # `swa.free` on an already-tombstoned one would assert). v = free_index.detach().to(torch.int64) - v_pages = v // self.page_size + ps = self.page_size + if start_pos is not None and ps > 1: + pieces = self.swa_attn_allocator._page_reps_pieces(v, start_pos) + reps = pieces[0] if len(pieces) == 1 else torch.cat(pieces) + # Keep only pages still bound on swa (freeing a tombstoned one + # would corrupt the hole list). `> 0` strict: -1 = tombstoned, + # page 0 = padding sink (never freeable). + rep_pages = reps // ps + swa_v2p_pages = self.swa_attn_allocator.virtual_to_physical[rep_pages] + live_reps = reps[swa_v2p_pages > 0] + if live_reps.numel() == 0: + return + self.swa_attn_allocator.free(live_reps, _pages=live_reps // ps) + self.swa_attn_allocator.clear_inverse_history() + return + v_pages = v // ps # `> 0` strict: -1 = tombstoned, page 0 = padding sink (never freeable). swa_v2p_pages = self.swa_attn_allocator.virtual_to_physical[v_pages] live = v[swa_v2p_pages > 0] if live.numel() == 0: return - self.swa_attn_allocator.free(live) + if ps == 1: + # token == page and the live filter just deduped against the v2p + # table, so these ARE unique page ids -- same skip as `_free_lazy`. + self.swa_attn_allocator.free(live, _pages=live) + else: + self.swa_attn_allocator.free(live) self.swa_attn_allocator.clear_inverse_history() def free_full(self, free_index: torch.Tensor) -> None: @@ -2667,10 +3619,29 @@ class UnifiedSWATokenToKVPoolAllocator(SWATokenToKVPoolAllocator): 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) + 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.""" + if ( + self._joint_avail_memo_epoch + != self.full_attn_allocator._chain_capacity_epoch() + ): + return [] + actual = self._compute_available_size() + if self._joint_avail_memo_tokens == actual: + return [] + return [ + f"[joint] stale available_size memo: " + f"cached={self._joint_avail_memo_tokens}, actual={actual}" + ] + def clear(self) -> None: self.full_attn_allocator.clear() self.swa_attn_allocator.clear() @@ -2717,3 +3688,308 @@ class UnifiedSWATokenToKVPoolAllocator(SWATokenToKVPoolAllocator): ): return 0 return fa.flush_opportunistic() + sa.flush_opportunistic() + + +class UnifiedMambaSWATokenToKVPoolAllocator(UnifiedSWATokenToKVPoolAllocator): + """Tri-pool composite for models with full KV + SWA KV + mamba/conv state + (Inkling-class: both `mambaish_config` and `is_hybrid_swa`). + + Chain (low byte -> high byte): + + [ mamba/conv (grow-up END) | swa (FLOAT middle) | full (grow-down END) ] + + Placement rationale: end pools never relocate — the request-granular, + fat-slot state pool and the unbounded per-step grower (full) take the + ends; SWA is window-capped (steady-state span ~= sum(min(seq, window))) + with the cheapest slots to move, so it floats. Out-of-window `free_swa` + tombstones become the float's interior HOLES, recycled in place by the + next per-step allocs — steady-state SWA churn costs zero copies. + + Token surface: inherited from the SWA composite (full = id-owner of the + per-token virtual ids; swa binds the same ids via `alloc_with_virtual`, + now on a `FloatMultiEndedAllocator`). Per-request state surface: the + `mamba_allocator` end MEA, wrapped by `UnifiedMambaSlotAllocator` exactly + like the 2-pool mamba composite. + """ + + def __init__( + self, + *, + unified_buffer: UnifiedKVPool, + kvcache, # UnifiedSWAKVPool + mamba_kvcache, # UnifiedMambaPool (req_to_token_pool.mamba_pool) + device: str, + full_max_total_num_tokens: int, + swa_max_total_num_tokens: int, + page_size: int = 1, + need_sort: bool = False, + forward_stream: Optional[torch.cuda.Stream] = None, + lazy_compaction: bool = False, + ): + super().__init__( + unified_buffer=unified_buffer, + kvcache=kvcache, + 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, + ) + # Per-request state END pool (grow-up; page_size=1 -- state is + # per-request, orthogonal to KV paging). + self.mamba_allocator = MultiEndedAllocator( + kvcache=mamba_kvcache, + unified_buffer=unified_buffer, + sub_pool_name="mamba", + device=device, + is_id_owner=True, + page_size=1, + need_sort=need_sort, + forward_stream=forward_stream, + lazy_compaction=lazy_compaction, + ) + # Chain wiring: mamba <-> swa(float) <-> full. + self.mamba_allocator.bind_high_peer(self.swa_attn_allocator) + self.swa_attn_allocator.bind_low_peer(self.mamba_allocator) + self.swa_attn_allocator.bind_high_peer(self.full_attn_allocator) + self.full_attn_allocator.bind_low_peer(self.swa_attn_allocator) + + # None, not empty: the checker's mamba census mixes physical free-lists + # with tree-held VIRTUAL ids, meaningless here. `free_pages is None` is + # its documented skip contract. + self.free_pages = None + self.release_pages = None + + logger.info( + "[unified-memory-pool] UnifiedMambaSWATokenToKVPoolAllocator ready: " + "chain=[mamba(up) | swa(float) | full(down)], " + "mamba max_slots=%d (entry_bytes=%d), joint available=%d", + self.mamba_allocator.max_slots, + self.mamba_allocator.entry_bytes, + self.available_size(), + ) + + # -- construction hooks -- + + def _build_swa_attn_allocator(self, **kwargs) -> MultiEndedAllocator: + # The swa side is the FLOAT middle. Holes-first: the float never runs + # the lazy event pipeline regardless of the composite's flag (frees + # mark holes; allocs recycle them in place). + kwargs["lazy_compaction"] = False + return FloatMultiEndedAllocator( + sub_pool_name="swa", + is_id_owner=False, # non-owner; consumes virtuals minted by full + **kwargs, + ) + + def _wire_peers(self) -> None: + # Chain wired in __init__ once the mamba end exists. + return + + # -- capacity -- + + def _compute_available_size(self) -> int: + """Joint TOKENS for `alloc(N)`: N costs N full pages AND N swa pages. + + (Memoized by the inherited `available_size` wrapper — the chain epoch + covers the mamba end via the frontier walks below.) + + The two sides draw on DIFFERENT free bands: full extends only downward + into the HIGH band (between the float's high frontier — or the mamba + end's when the float is empty/transparent — and full's low frontier); + the swa float extends either side but a single batch alloc extends ONE + side. Monotone feasibility predicate, solved by binary search: + + ext_f = max(0, N - H_f) must fit: ext_f*e_f <= B_high + ext_s = max(0, N - H_s) must fit: ext_s*e_s <= max(B_low, + B_high - ext_f*e_f) + N <= H_f + R_f, N <= H_s + R_s (index-space caps) + + where H_* are drainable holes (full: lazy only; swa: always — holes + are the float's design), B_low is the band between the mamba end and + the float's low frontier (0 when the float is transparent — the whole + region is already in B_high), and R_* are index rooms. Order matches + the alloc path: full takes from B_high first, then the float extends. + """ + fa, sa = self.full_attn_allocator, self.swa_attn_allocator + e_f, e_s = fa.entry_bytes_per_page, sa.entry_bytes_per_page + # full is grow-down: its chain gap IS the high band. + b_high = fa._current_gap_bytes() + if sa._is_frontier_transparent(): + b_low = 0 + else: + b_low = max( + 0, + sa._byte_low_frontier() - sa._chain_high_frontier_below_bytes(), + ) + h_f = len(fa._free_phys_pages) if fa.lazy_compaction else 0 + h_s = sa._hole_pages() + r_f = fa.num_pages - fa.min_page_index - fa._allocated_pages() + r_s = sa.num_pages - sa.min_page_index - sa._allocated_pages() + + def feasible(n: int) -> bool: + if n > h_f + r_f or n > h_s + r_s: + return False + ext_f = max(0, n - h_f) + if ext_f * e_f > b_high: + return False + ext_s = max(0, n - h_s) + # On the float's page grid, never in raw bytes: a byte budget + # credits a page `take_physical_pages` cannot yield. + full_low_after = fa._byte_low_frontier() - ext_f * e_f + if sa._is_frontier_transparent(): + room = sa.pages_in_band( + low_byte=sa._chain_high_frontier_below_bytes(), + high_byte=full_low_after, + ) + return ext_s <= room + p_low = sa.pages_in_band( + low_byte=sa._chain_high_frontier_below_bytes(), + high_byte=sa._byte_low_frontier(), + ) + p_high = sa.pages_in_band( + low_byte=sa._byte_high_frontier(), + high_byte=full_low_after, + ) + return ext_s <= max(p_low, p_high) + + lo_n, hi_n = 0, min(h_f + r_f, h_s + r_s) + while lo_n < hi_n: + mid = (lo_n + hi_n + 1) // 2 + if feasible(mid): + lo_n = mid + else: + hi_n = mid - 1 + return lo_n * self.page_size + + def _flush_targets(self): + """All three members, same reasoning as the 2-pool pair with one + addition each way: the FLOAT's `_flush` is zero-copy boundary + absorption, and running it before `_ask_float_for_room` keeps the + deficit math from pricing a span that still claims absorbed holes + (which would buy a relocation the free shrink already covered); the + MAMBA end's compaction feeds the low band, which the float's own + extension for the same tokens can draw on. + """ + return ( + self.swa_attn_allocator, + self.full_attn_allocator, + self.mamba_allocator, + ) + + def _alloc_demand(self, need_tokens: int): + """Demand VECTOR for one composite allocation, in pages per band -- + zero for bands the operation does not touch. A composite token + (prefill extend and decode alike) needs a full page AND a swa page; + it never draws a state slot — those are per-REQUEST allocations that + run the band-level ladder with their own {mamba: k} vector, so mamba + is an explicit 0 here, not an omission. A future 3-pool composite + (e.g. C128 | swa-float | C4) overrides just this vector and inherits + the whole relocation policy. + """ + need_n = -(-need_tokens // self.page_size) + return { + self.full_attn_allocator: need_n, + self.swa_attn_allocator: need_n, + self.mamba_allocator: 0, + } + + def _ask_float_for_room(self, need_tokens: int) -> None: + """Composite shortfall: hand the demand vector to the shared policy; + the float is whichever demanded band floats.""" + demand = self._alloc_demand(need_tokens) + flt = None + for b in demand: + if isinstance(b, FloatMultiEndedAllocator): + flt = b + _float_open_short_side(flt, demand) + + def mamba_slot_full_token_cost(self) -> int: + """Full-token-equivalents one mamba/conv slot removes from the shared + buffer. A tri-pool token costs e_f + e_s bytes, so: + ceil(mamba_entry_bytes / (e_f + e_s)). Conservative (rounded up).""" + e_tok = ( + self.full_attn_allocator.entry_bytes + self.swa_attn_allocator.entry_bytes + ) + return -(-self.mamba_allocator.entry_bytes_per_page // e_tok) + + def debug_print(self) -> str: + sa = self.swa_attn_allocator + return ( + super().debug_print() + + f", #mamba-available={self.mamba_allocator.available_size()}" + + f", swa-float span=[{sa.low_wm_page},{sa.high_wm_page}) " + + f"holes={sa._hole_pages()}" + ) + + # -- lifecycle fanout (adds the mamba end) -- + + def clear(self) -> None: + super().clear() + self.mamba_allocator.clear() + + def set_latest_forward_done_event(self, event: Optional[torch.cuda.Event]) -> None: + super().set_latest_forward_done_event(event) + self.mamba_allocator.set_latest_forward_done_event(event) + + def set_inflight_forward( + self, + forward_done: torch.cuda.Event, + out_cache_loc_virtual: Optional[torch.Tensor], + ) -> None: + # full + swa are written per new token via set_kv_buffer; the mamba + # state is written by the conv kernels, not out_cache_loc -- pass None + # (the 2-pool mamba composite's convention). + super().set_inflight_forward(forward_done, out_cache_loc_virtual) + self.mamba_allocator.set_inflight_forward(forward_done, None) + + def evict_to_free_tokens(self, tree_cache, num_tokens: int) -> None: + """Joint-aware eviction: evicting one tri-lifetime tree node frees + bytes on several sides at once, and the default single pass's per-side + shortfall math can leave the JOINT gate short. Bounded re-check loop: + evict until the joint availability covers the ask or a pass stops + making progress (then the capacity gate reports the shortfall).""" + from sglang.srt.mem_cache.common import evict_from_tree_cache + + for _ in range(4): + before = self.available_size() + if before >= num_tokens: + return + evict_from_tree_cache(tree_cache, num_tokens) + if self.available_size() <= before: + return # no progress + + def verify_byte_accounting(self) -> List[str]: + return ( + _chain_byte_accounting_violations( + [ + self.mamba_allocator, + self.swa_attn_allocator, + self.full_attn_allocator, + ] + ) + + self._joint_capacity_memo_violations() + ) + + def flush_opportunistic(self) -> int: + """Per-step reclaim across the whole chain. The float participates: + its holes are not flushable BACKLOG (never moved here), but its + deferred boundary absorption is exactly the work this quiescent point + exists for -- and it is where the float's single D2H is paid.""" + fa, ma = self.full_attn_allocator, self.mamba_allocator + sa = self.swa_attn_allocator + if ( + fa._free_phys_pages.numel() == 0 + and not fa._pending_reuse + and ma._free_phys_pages.numel() == 0 + and not ma._pending_reuse + and sa._free_phys_pages.numel() == 0 + ): + return 0 + return ( + fa.flush_opportunistic() + + ma.flush_opportunistic() + + sa.flush_opportunistic() + ) diff --git a/python/sglang/srt/mem_cache/unified_memory_pool.py b/python/sglang/srt/mem_cache/unified_memory_pool.py index 9c69045af..f674ba81b 100644 --- a/python/sglang/srt/mem_cache/unified_memory_pool.py +++ b/python/sglang/srt/mem_cache/unified_memory_pool.py @@ -11,13 +11,15 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -"""UnifiedKVPool — one physical `uint8` byte buffer shared by 2 sub-pools. +"""UnifiedKVPool -- one physical `uint8` byte buffer shared by N sub-pools. -Two `MultiEndedAllocator`s grow from opposite ends; eager-compacting `free` -keeps each pool's byte range hole-free. Layout is envelope-major (a slot's data -for all its layers in one contiguous byte envelope) so a freed slot vacates a -region the peer can grow into. Everything above the allocator stores virtual -slot IDs; the allocator owns the per-sub-pool virtual<->physical tables and +Two END `MultiEndedAllocator`s grow inward from opposite ends; optional +"float" MIDDLE pools live between their frontiers (chain order +`[up end, floats..., down end]`). Eager- or lazy-compacting `free` keeps each +pool's byte range reclaimable. Layout is envelope-major (a slot's data for all +its layers in one contiguous byte envelope) so a freed slot vacates a region a +neighbor can grow into. Everything above the allocator stores virtual slot +IDs; the allocator owns the per-sub-pool virtual<->physical tables and compaction only mutates those (no reference rewriting). """ @@ -26,7 +28,7 @@ from __future__ import annotations import logging from abc import ABC, abstractmethod from dataclasses import dataclass -from typing import Dict, List, NamedTuple, Optional, Tuple +from typing import ClassVar, Dict, List, NamedTuple, Optional, Tuple import torch from torch.profiler import record_function @@ -70,17 +72,28 @@ def _store_dtype_for(kv_cache_dtype: torch.dtype) -> torch.dtype: @dataclass(frozen=True, kw_only=True) class SubPoolSpec(ABC): - """Abstract per-slot layout of one sub-pool in a `UnifiedKVPool`.""" + """Abstract per-slot layout of one sub-pool in a `UnifiedKVPool`. + + ``grow_direction`` places the sub-pool in the buffer's chain: the two + ``"up"``/``"down"`` END pools own the buffer's two ends and grow inward; + ``"float"`` middles live between the ends' frontiers (a float's position is + carried entirely by its allocator's two watermarks — every view spans the + whole buffer, so relocation never rebuilds views). + """ + + # Grow directions this subclass accepts. Scratch-class specs (e.g. the + # spec-decode band) narrow this to ("float",). + _allowed_grow_directions: ClassVar[Tuple[str, ...]] = ("up", "down", "float") name: str layer_num: int - grow_direction: str # "up" | "down" + grow_direction: str # "up" | "down" | "float" def __post_init__(self): - assert self.grow_direction in ( - "up", - "down", - ), f"grow_direction must be 'up' or 'down'; got {self.grow_direction!r}" + assert self.grow_direction in self._allowed_grow_directions, ( + f"{type(self).__name__}.grow_direction must be one of " + f"{self._allowed_grow_directions}; got {self.grow_direction!r}" + ) assert self.layer_num > 0, f"layer_num must be positive; got {self.layer_num}" @abstractmethod @@ -278,9 +291,11 @@ def _reserved_floor_bytes(sub_pool_specs: List[SubPoolSpec], page_size: int) -> class UnifiedKVPool: - """One physical `uint8` byte buffer shared by 2 sub-pools, each exposing + """One physical `uint8` byte buffer shared by N sub-pools, each exposing per-layer views over its own byte range (contiguous per layer for KV, - strided for the Mamba state). Allocators keep byte ranges disjoint; no usage tracking here. + strided for the Mamba state). Two END pools (one grow-up, one grow-down) + own the buffer's ends; optional "float" MIDDLE pools live between their + frontiers. Allocators keep byte ranges disjoint; no usage tracking here. """ def __init__( @@ -293,21 +308,33 @@ class UnifiedKVPool: page_size: int = 1, ): assert page_size >= 1, f"page_size must be >= 1; got {page_size}" - assert len(sub_pool_specs) == 2, ( - f"UnifiedKVPool currently supports exactly 2 sub-pools; got " - f"{len(sub_pool_specs)} (N>2 is not yet implemented)" - ) + assert ( + len(sub_pool_specs) >= 2 + ), f"UnifiedKVPool needs >= 2 sub-pools; got {len(sub_pool_specs)}" names = [s.name for s in sub_pool_specs] - assert len(set(names)) == 2, f"sub-pool names must be unique; got {names}" - directions = sorted(s.grow_direction for s in sub_pool_specs) - assert directions == ["down", "up"], ( - f"UnifiedKVPool needs one grow-up and one grow-down sub-pool; " - f"got {directions}" + assert len(set(names)) == len( + names + ), f"sub-pool names must be unique; got {names}" + # Per-spec direction validity already ran in each spec's __post_init__. + up_specs = [s for s in sub_pool_specs if s.grow_direction == "up"] + down_specs = [s for s in sub_pool_specs if s.grow_direction == "down"] + float_specs = [s for s in sub_pool_specs if s.grow_direction == "float"] + assert len(up_specs) == 1 and len(down_specs) == 1, ( + f"UnifiedKVPool needs exactly one grow-up and one grow-down END " + f"sub-pool; got directions " + f"{[s.grow_direction for s in sub_pool_specs]}" ) self.device = device self.total_bytes = total_bytes - self.sub_pool_specs = sub_pool_specs + # Canonical chain order, low byte end -> high: grow-up end, float + # middles (input order preserved), grow-down end. The allocators' + # neighbour wiring follows it; all other access is by-name. + self.sub_pool_specs: List[SubPoolSpec] = [ + up_specs[0], + *float_specs, + down_specs[0], + ] self._page_size = page_size self._specs_by_name: Dict[str, SubPoolSpec] = { s.name: s for s in sub_pool_specs @@ -348,9 +375,9 @@ 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. - reserved_floor = _reserved_floor_bytes(sub_pool_specs, page_size) + reserved_floor = _reserved_floor_bytes(self.sub_pool_specs, page_size) - for spec in sub_pool_specs: + for spec in self.sub_pool_specs: entry_bytes = spec.entry_bytes() max_slots = total_bytes // entry_bytes min_slot_index = (reserved_floor + entry_bytes - 1) // entry_bytes # ceil @@ -390,9 +417,9 @@ class UnifiedKVPool: "%d sub-pool(s)", total_bytes / GB, total_bytes, - len(sub_pool_specs), + len(self.sub_pool_specs), ) - for s in sub_pool_specs: + for s in self.sub_pool_specs: logger.info( "[unified-memory-pool] sub-pool %r: kind=%s, layer_num=%d, grow=%s, " "entry_bytes=%d, max_slots=%d, min_slot_index=%d (slots [0,%d) reserved)", @@ -843,9 +870,11 @@ class UnifiedMambaPool(MambaPool): # Inherited MambaPool state ops (copy_from/clear_slots/get_cpu_copy/load_cpu_copy) # take PHYSICAL slot ids; callers translate via the slot allocator first. - def _copy_from_physical(self, src_index: torch.Tensor, dst_index: torch.Tensor): - # Physical-slot copy used by the allocator's `_compact_pending`. - MambaPool.copy_from(self, src_index, dst_index) + def move_kv_cache(self, tgt_loc: torch.Tensor, src_loc: torch.Tensor): + # Cross-pool physical-move contract, implemented by every pool the + # MultiEndedAllocator wraps. Ids are PHYSICAL slots; `MambaPool.copy_from` + # takes (src, dst), hence the swap. + MambaPool.copy_from(self, src_loc, tgt_loc) # -- PD state transfer (StateType.MAMBA) -- # The transfer item is the whole per-slot envelope, addressed as @@ -1776,3 +1805,222 @@ def init_unified_swa_pools( token_to_kv_pool=token_to_kv_pool, token_to_kv_pool_allocator=allocator, ) + + +def init_unified_mamba_swa_pools( + *, + device: str, + kv_cache_dtype: torch.dtype, + head_num: int, + head_dim: int, + v_head_dim: int, + swa_head_num: int, + swa_head_dim: int, + swa_v_head_dim: int, + page_size: int, + start_layer: int, + end_layer: int, + swa_attention_layer_ids: List[int], + full_attention_layer_ids: List[int], + mamba_layer_ids: List[int], + mamba2_cache_params, + full_max_total_num_tokens: int, + swa_max_total_num_tokens: int, + max_mamba_cache_size: int, + model_context_len: int, + extra_max_context_len: int, + max_num_reqs: int, + enable_memory_saver: bool, + enable_mamba_extra_buffer: bool, + disable_overlap_schedule: bool, + need_sort: bool, + speculative_num_draft_tokens: Optional[int] = None, + forward_stream: Optional[torch.cuda.Stream] = None, + lazy_compaction: bool = False, + unified_total_bytes: Optional[int] = None, + sliding_window_size: Optional[int] = None, +) -> UnifiedPoolBundle: + """Build the TRI-pool unified-memory-pool stack for models with full KV + + SWA KV + mamba/conv state (Inkling-class: `mambaish_config` AND + `is_hybrid_swa` simultaneously — Inkling's SConv state is conv-only but + rides the mamba machinery, so "mamba" here == the conv state pool). + + Chain: ``[mamba (up END) | swa (FLOAT) | full (down END)]``. The KV side + is a `UnifiedSWAKVPool` (per-layer full/swa routing, asymmetric head + geometry supported); the state side is a `UnifiedHybridReqToTokenPool` + whose `mamba_pool` the model reads directly (sconv: + `req_to_token_pool.mamba2_layer_cache(layer).conv[...]` with + `translate_mamba_indices` for v->p). + + Sizing inputs are the same token counts the 2-pool factories take (ratio- + fed until the byte configurator lands); the buffer budget is their byte + sum and the runtime split floats. + """ + from sglang.srt.mem_cache.multi_ended_allocator import ( + UnifiedMambaSWATokenToKVPoolAllocator, + ) + + assert page_size >= 1, f"page_size must be >= 1, got {page_size}" + assert ( + len(full_attention_layer_ids) > 0 + ), "tri-pool with zero full-attention layers is degenerate" + assert ( + len(swa_attention_layer_ids) > 0 + ), "tri-pool with zero SWA-attention layers is degenerate" + assert len(mamba_layer_ids) > 0, "tri-pool with zero state layers is degenerate" + + store_dtype = _store_dtype_for(kv_cache_dtype) + # mamba/conv at the LOWEST bytes, full KV at the HIGHEST, SWA floating + # between: ends never relocate, and SWA's window-capped span is cheapest to move. + full_spec = MHASubPoolSpec( + name="full", + layer_num=len(full_attention_layer_ids), + head_num=head_num, + head_dim=head_dim, + v_head_dim=v_head_dim, + store_dtype=store_dtype, + grow_direction="down", + ) + swa_spec = MHASubPoolSpec( + name="swa", + layer_num=len(swa_attention_layer_ids), + head_num=swa_head_num, + head_dim=swa_head_dim, + v_head_dim=swa_v_head_dim, + store_dtype=store_dtype, + grow_direction="float", + ) + cp = mamba2_cache_params + mamba_spec = MambaSubPoolSpec( + name="mamba", + layer_num=len(mamba_layer_ids), + conv_state_shapes=tuple(tuple(int(x) for x in s) for s in cp.shape.conv), + conv_dtype=cp.dtype.conv, + temporal_state_shape=tuple(int(x) for x in cp.shape.temporal), + temporal_dtype=cp.dtype.temporal, + grow_direction="up", + ) + 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 = ( + full_max_total_num_tokens * full_spec.entry_bytes() + + swa_max_total_num_tokens * swa_spec.entry_bytes() + + max_mamba_cache_size * mamba_spec.entry_bytes() + ) + # bs=1 floor: ONE sliding window of swa KV (+ a page of slack, clamped to + # the context) + the state slots one running request locks (1 active + 2 + # radix checkpoints, a FLOOR not headroom) + the slot-0 sink. The + # full-attention side is not charged: `max_req_len` already clamps to the pool. + 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()), + ("bs1_state_slots", 3 * mamba_spec.entry_bytes()), + ( + "sink", + _reserved_floor_bytes([full_spec, swa_spec, mamba_spec], page_size), + ), + ], + factory="init_unified_mamba_swa_pools", + ) + shared_pool = UnifiedKVPool( + total_bytes=total_bytes, + sub_pool_specs=[full_spec, swa_spec, mamba_spec], + device=device, + enable_memory_saver=enable_memory_saver, + page_size=page_size, + ) + token_to_kv_pool = UnifiedSWAKVPool( + unified_buffer=shared_pool, + swa_attention_layer_ids=swa_attention_layer_ids, + full_attention_layer_ids=full_attention_layer_ids, + page_size=page_size, + start_layer=start_layer, + end_layer=end_layer, + enable_memory_saver=enable_memory_saver, + ) + req_to_token_pool = UnifiedHybridReqToTokenPool( + unified_buffer=shared_pool, + mamba_sub_pool_name="mamba", + size=max_num_reqs, + mamba_spec_state_size=max_num_reqs, + max_context_len=model_context_len + extra_max_context_len, + device=device, + enable_memory_saver=enable_memory_saver, + cache_params=mamba2_cache_params, + mamba_layer_ids=mamba_layer_ids, + enable_mamba_extra_buffer=enable_mamba_extra_buffer, + speculative_num_draft_tokens=speculative_num_draft_tokens, + enable_overlap_schedule=not disable_overlap_schedule, + start_layer=start_layer, + ) + allocator = UnifiedMambaSWATokenToKVPoolAllocator( + unified_buffer=shared_pool, + kvcache=token_to_kv_pool, + mamba_kvcache=req_to_token_pool.mamba_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, + ) + # Wrap the composite's mamba end in the slot allocator (PHYSICAL view) the + # radix MambaComponent / model-side sconv reads consume. + mamba_slot_allocator = UnifiedMambaSlotAllocator( + allocator.mamba_allocator, + max_size=req_to_token_pool._shared_mamba_size, + device=device, + ) + req_to_token_pool.mamba_allocator = mamba_slot_allocator + + logger.info( + "[unified-memory-pool] ============================================================" + ) + logger.info( + "[unified-memory-pool] UNIFIED MEMORY POOL ENABLED -- path=SWA+Mamba tri-pool" + ) + logger.info( + "[unified-memory-pool] full_layers=%d, swa_layers=%d, state_layers=%d, " + "head_num=%d/%d, head_dim=%d/%d, page_size=%d", + len(full_attention_layer_ids), + len(swa_attention_layer_ids), + len(mamba_layer_ids), + head_num, + swa_head_num, + head_dim, + swa_head_dim, + page_size, + ) + logger.info( + "[unified-memory-pool] total_bytes=%d (=%.2f GB), full_max=%d, swa_max=%d, " + "max_mamba_cache_size=%d, max_num_reqs=%d, joint_available=%d", + total_bytes, + total_bytes / GB, + full_max_total_num_tokens, + swa_max_total_num_tokens, + max_mamba_cache_size, + max_num_reqs, + allocator.available_size(), + ) + logger.info( + "[unified-memory-pool] ============================================================" + ) + return UnifiedPoolBundle( + unified_memory_pool=shared_pool, + token_to_kv_pool=token_to_kv_pool, + token_to_kv_pool_allocator=allocator, + req_to_token_pool=req_to_token_pool, + ) diff --git a/test/registered/models_e2e/test_inkling_unified.py b/test/registered/models_e2e/test_inkling_unified.py new file mode 100644 index 000000000..4e089d980 --- /dev/null +++ b/test/registered/models_e2e/test_inkling_unified.py @@ -0,0 +1,240 @@ +"""Inkling under ``--enable-unified-memory`` -- the first TRI-pool model. + +Boots the shrunken ``thinkingmachines/Inkling`` checkpoint (``test`` revision) +with the unified memory pool: one byte buffer, chain +``[mamba/conv (up END) | swa (FLOAT) | full (down END)]``. Inkling is the only +in-tree model that is BOTH mambaish (conv-only SConv state riding the mamba +machinery) and hybrid-SWA, so booting AT ALL proves the tri routing branch +(the 2-pool branches would either mis-store SWA KV at full lifetime — the +pre-tri hazard — or fail loud). + +Guards (undertrained checkpoint — code-path correctness, not answer quality): + - tri boot + generation through the Triton backend (unified forces triton; + Inkling's fa4 default is NOT unified-compatible); + - decode/prefill KV consistency via the input-vs-output logprobs match + (catches wrong-slot reads through the v2p translate on any of the three + pools); + - multi-turn prefix reuse: a repeated prefix must reproduce identical + logprobs (radix reuse + SWA tombstone recycling + conv COW); + - a long-generation turn that slides past the SWA window (exercises + out-of-window free_swa -> float holes -> in-place reuse during decode). + +An optional env-gated parity class re-runs fixed prompts on the STATIC pools +and compares logprobs (``INKLING_UNIFIED_PARITY=1``) — the eval-host lane; +kept out of per-commit CI to bound cost. + + python -m pytest test/registered/models/test_inkling_unified.py -v +""" + +import os +import unittest + +import requests + +from sglang.srt.utils import kill_process_tree +from sglang.test.ci.ci_register import register_cuda_ci + +# Aliased so pytest does not collect the imported `test_`-prefixed helper. +from sglang.test.kl_test_utils import ( + test_input_output_logprobs_match_helper as assert_logprobs_match, +) +from sglang.test.test_utils import ( + DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + DEFAULT_URL_FOR_TEST, + CustomTestCase, + popen_launch_server, +) + +register_cuda_ci(est_time=600, stage="base-b", runner_config="1-gpu-large") + +_MODEL_PATH = os.environ.get("INKLING_TEST_MODEL_PATH", "thinkingmachines/Inkling") +_MODEL_REVISION = os.environ.get("INKLING_TEST_MODEL_REVISION", "test") + + +def _unified_args(): + """Server args for the tri-pool boot. Mirrors test_inkling.py's fixture + minus the multimodal/parser surface (KV-path focus), plus the unified + flags. The ratios still feed boot sizing until the byte configurator + lands; the runtime split floats regardless.""" + args = [ + "--trust-remote-code", + "--enable-unified-memory", + # Unified requires the Triton strided page-major read/write paths. + "--attention-backend", + "triton", + "--page-size", + "128", + "--mamba-radix-cache-strategy", + "extra_buffer", + # Inkling defaults to a FULL prefill graph, which unified rejects at + # boot: the prefill graph runner bypasses the virtual->physical rebind. + "--cuda-graph-backend-prefill", + "disabled", + "--swa-full-tokens-ratio", + "0.1", + "--mamba-full-memory-ratio", + "0.1", + "--mem-fraction-static", + "0.5", + ] + if _MODEL_REVISION: + args += ["--revision", _MODEL_REVISION] + return args + + +def _static_args(): + args = [ + "--trust-remote-code", + "--attention-backend", + "triton", + "--page-size", + "128", + "--mamba-radix-cache-strategy", + "extra_buffer", + # Inkling defaults to a FULL prefill graph, which unified rejects at + # boot: the prefill graph runner bypasses the virtual->physical rebind. + "--cuda-graph-backend-prefill", + "disabled", + "--swa-full-tokens-ratio", + "0.1", + "--mamba-full-memory-ratio", + "0.1", + "--mem-fraction-static", + "0.5", + ] + if _MODEL_REVISION: + args += ["--revision", _MODEL_REVISION] + return args + + +_PARITY_PROMPTS = [ + "The capital of France is", + "1 + 2 + 3 + 4 + 5 =", + "List three prime numbers:", +] + + +def _greedy_generate(base_url, text, max_new_tokens=32, logprobs=False): + payload = { + "text": text, + "sampling_params": {"temperature": 0.0, "max_new_tokens": max_new_tokens}, + } + if logprobs: + payload["return_logprob"] = True + payload["logprob_start_len"] = 0 + resp = requests.post(f"{base_url}/generate", json=payload, timeout=120) + assert resp.status_code == 200, resp.text + return resp.json() + + +class TestInklingUnifiedTriPool(CustomTestCase): + @classmethod + def setUpClass(cls): + cls.model = _MODEL_PATH + cls.base_url = DEFAULT_URL_FOR_TEST + cls.process = popen_launch_server( + cls.model, + cls.base_url, + timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + other_args=_unified_args(), + env={**os.environ, "SGLANG_ENABLE_UNIFIED_RADIX_TREE": "1"}, + ) + + @classmethod + def tearDownClass(cls): + if getattr(cls, "process", None) is not None: + kill_process_tree(cls.process.pid) + + def test_generation_basic(self): + """Booting IS the tri-routing gate; every prompt must complete.""" + for prompt in _PARITY_PROMPTS: + data = _greedy_generate(self.base_url, prompt, max_new_tokens=16) + self.assertIn("text", data, data) + self.assertGreater(len(data["text"].strip()), 0, data) + + def test_input_output_logprobs_match(self): + """Prefill-vs-decode KV consistency through all three sub-pools' + translates (wrong-slot reads surface as logprob mismatches).""" + assert_logprobs_match( + self.base_url, + {self.model: {"kl_div": 1e-2}}, + self.model, + max_samples=4, + max_new_tokens=256, + trust_remote_code=True, + ) + + def test_repeated_prefix_reproduces_logprobs(self): + """Multi-turn prefix reuse: radix hit + conv COW + swa recycling must + not change the numerics of a greedy re-run.""" + prompt = ( + "In a quiet village by the sea, a clockmaker kept a ledger of " + "every tide. One morning the ledger read:" + ) + first = _greedy_generate( + self.base_url, prompt, max_new_tokens=24, logprobs=True + ) + second = _greedy_generate( + self.base_url, prompt, max_new_tokens=24, logprobs=True + ) + self.assertEqual(first["text"], second["text"]) + lp1 = [t[0] for t in first["meta_info"]["output_token_logprobs"]] + lp2 = [t[0] for t in second["meta_info"]["output_token_logprobs"]] + for a, b in zip(lp1, lp2): + self.assertAlmostEqual(a, b, places=3) + + def test_long_decode_slides_past_swa_window(self): + """A generation long enough to age tokens out of the SWA window + exercises free_swa -> float holes -> in-place reuse mid-decode.""" + data = _greedy_generate( + self.base_url, + "Write an unbroken story about a lighthouse: ", + max_new_tokens=512, + ) + self.assertGreater(len(data["text"].strip()), 0, data) + + +@unittest.skipUnless( + os.environ.get("INKLING_UNIFIED_PARITY") == "1", + "eval-host lane: set INKLING_UNIFIED_PARITY=1 (two sequential server boots)", +) +class TestInklingUnifiedVsStaticParity(CustomTestCase): + """Greedy logprob parity: unified tri-pool vs static pools, same prompts. + Two sequential boots -- the strongest wrong-slot tripwire short of GSM8K.""" + + @classmethod + def _collect(cls, other_args): + proc = popen_launch_server( + _MODEL_PATH, + DEFAULT_URL_FOR_TEST, + timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + other_args=other_args, + env={**os.environ, "SGLANG_ENABLE_UNIFIED_RADIX_TREE": "1"}, + ) + try: + out = [] + for p in _PARITY_PROMPTS: + data = _greedy_generate( + DEFAULT_URL_FOR_TEST, p, max_new_tokens=32, logprobs=True + ) + out.append( + ( + data["text"], + [t[0] for t in data["meta_info"]["output_token_logprobs"]], + ) + ) + return out + finally: + kill_process_tree(proc.pid) + + def test_parity(self): + static = self._collect(_static_args()) + unified = self._collect(_unified_args()) + for (s_text, s_lp), (u_text, u_lp) in zip(static, unified): + self.assertEqual(s_text, u_text) + for a, b in zip(s_lp, u_lp): + self.assertAlmostEqual(a, b, places=2) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/disaggregation/test_unified_memory_move_gate.py b/test/registered/unit/disaggregation/test_unified_memory_move_gate.py index fc93bf7d4..34ae984e1 100644 --- a/test/registered/unit/disaggregation/test_unified_memory_move_gate.py +++ b/test/registered/unit/disaggregation/test_unified_memory_move_gate.py @@ -148,9 +148,20 @@ class TestGatedPeerHolesAreNotSchedulable(CustomTestCase): self.entry_bytes_per_page = 512 self.disagg_move_gate = gate + def _is_frontier_transparent(self): + return False + class _Owner: + """Stands in for a grow-up END pool: the credit walks the chain from + `_growth_side_neighbor()`, so the stub must expose what that walk reads, + not the pre-chain `_peer` slot it used to.""" + def __init__(self, peer): - self._peer = peer + self.grow_direction = "up" + self.high_peer = peer + self.low_peer = None + + _growth_side_neighbor = MultiEndedAllocator._growth_side_neighbor def _credit(self, gate): peer = self._Peer(gate) diff --git a/test/registered/unit/mem_cache/test_multi_ended_allocator.py b/test/registered/unit/mem_cache/test_multi_ended_allocator.py index 263222b24..3876deb96 100644 --- a/test/registered/unit/mem_cache/test_multi_ended_allocator.py +++ b/test/registered/unit/mem_cache/test_multi_ended_allocator.py @@ -30,6 +30,7 @@ import unittest import torch from sglang.srt.mem_cache.multi_ended_allocator import ( + FloatMultiEndedAllocator, MultiEndedAllocator, UnifiedMambaTokenToKVPoolAllocator, UnifiedSWATokenToKVPoolAllocator, @@ -1825,8 +1826,8 @@ class TestPagedMultiEndedAllocator(unittest.TestCase): full_kv.attach_allocator = lambda allocator: None mamba_kv = _FakeKVCache(pool.max_slots("mamba")) mamba_kv.attach_allocator = lambda allocator: None - # _copy_from_physical for the mamba sub-pool (kept un-translated). - mamba_kv._copy_from_physical = lambda src, dst: None + # The physical-move contract for the mamba sub-pool (un-translated); + # _FakeKVCache already provides move_kv_cache, so nothing extra. class _FakeHybridLinearKVPool: full_kv_pool = full_kv @@ -2822,5 +2823,481 @@ class TestPs64MLACompositeFeasibility(unittest.TestCase): self.assertTrue(bool((got < 2**31).all().item())) +class _ChainStub: + """Duck-typed chain member for frontier-walk tests: only the attributes the + walk itself touches. Lets the tests position an (opaque|transparent) middle + at exact byte coordinates without the full allocator machinery (the real + float allocator lands in a later commit).""" + + def __init__(self, *, low_byte: int, high_byte: int, transparent: bool): + self._low_byte = low_byte + self._high_byte = high_byte + self.transparent = transparent + self.low_peer = None + self.high_peer = None + self.lazy_compaction = False + self._free_phys_pages = torch.empty(0, dtype=torch.int64) + self.entry_bytes_per_page = 1 + self.sub_pool_name = "stub" + self.grow_direction = "float" + + def _is_frontier_transparent(self): + return self.transparent + + def _byte_low_frontier(self): + return self._low_byte + + def _byte_high_frontier(self): + return self._high_byte + + +class TestChainFrontierWalk(unittest.TestCase): + """N-pool chain walk: 2-pool byte-identity to the old single-peer formulas, + transparent-middle skipping, and growth-side-neighbor credit routing. + + Guarded failure modes: (a) a rewrite of the walk silently changes the + 2-pool gap math (golden identity); (b) an empty/parked middle walls off + free space it does not occupy; (c) drainable-hole credit reads the wrong + chain member. + """ + + def _build_pair(self): + full = _make_mha_spec("full", "up", layer_num=2) + mamba = _make_mamba_spec("mamba", "down", layer_num=2) + total = full.entry_bytes() * 64 + mamba.entry_bytes() * 16 + pool = UnifiedKVPool( + total_bytes=total, + sub_pool_specs=[full, mamba], + device=_DEV, + enable_memory_saver=False, + ) + fa = MultiEndedAllocator( + kvcache=_FakeKVCache(pool.max_slots("full")), + unified_buffer=pool, + sub_pool_name="full", + device=_DEV, + is_id_owner=True, + ) + ma = MultiEndedAllocator( + kvcache=_FakeKVCache(pool.max_slots("mamba")), + unified_buffer=pool, + sub_pool_name="mamba", + device=_DEV, + is_id_owner=True, + ) + fa.bind_peer(ma) + ma.bind_peer(fa) + return pool, fa, ma + + def test_bind_peer_mirrors_growth_side(self): + _, fa, ma = self._build_pair() + self.assertIs(fa.high_peer, ma) # grow-up's neighbor sits above + self.assertIsNone(fa.low_peer) + self.assertIs(ma.low_peer, fa) # grow-down's neighbor sits below + self.assertIsNone(ma.high_peer) + + def test_bind_peer_rejects_float_members(self): + _, fa, ma = self._build_pair() + stub = _ChainStub(low_byte=0, high_byte=0, transparent=True) + with self.assertRaisesRegex(AssertionError, "END-pool-only"): + fa.bind_peer(stub) + fa.grow_direction = "float" + try: + with self.assertRaisesRegex(AssertionError, "END-pool-only"): + fa.bind_peer(ma) + finally: + fa.grow_direction = "up" + + def test_two_pool_gap_equals_old_single_peer_formula(self): + # Golden identity: with no middles, the chain walk must reproduce the + # pre-chain closed form gap_up = peer_low - my_high ; + # gap_down = my_low - peer_high at every allocation state. + _, fa, ma = self._build_pair() + for n_full, n_mamba in ((0, 0), (8, 0), (8, 4), (32, 16)): + fa.clear() + ma.clear() + if n_full: + self.assertIsNotNone(fa.alloc(n_full)) + if n_mamba: + self.assertIsNotNone(ma.alloc(n_mamba)) + self.assertEqual( + fa._current_gap_bytes(), + max(0, ma._byte_low_frontier() - fa._byte_high_frontier()), + ) + # Old down-side closed form: my_low - peer_high (symmetric band). + self.assertEqual( + ma._current_gap_bytes(), + max(0, ma._byte_low_frontier() - fa._byte_high_frontier()), + ) + + def test_transparent_middle_is_skipped(self): + pool, fa, ma = self._build_pair() + mid_lo = fa.entry_bytes_per_page * 8 + mid_hi = fa.entry_bytes_per_page * 12 + stub = _ChainStub(low_byte=mid_lo, high_byte=mid_hi, transparent=True) + fa.bind_high_peer(stub) + stub.low_peer = fa + stub.high_peer = ma + ma.bind_low_peer(stub) + + # Transparent: both ends see straight through to each other. + self.assertEqual( + fa._current_gap_bytes(), + ma._byte_low_frontier() - fa._byte_high_frontier(), + ) + self.assertEqual( + ma._current_gap_bytes(), + ma._byte_low_frontier() - fa._byte_high_frontier(), + ) + + # Opaque: each end's gap stops at the middle's near frontier. + stub.transparent = False + self.assertEqual(fa._current_gap_bytes(), mid_lo - fa._byte_high_frontier()) + self.assertEqual(ma._current_gap_bytes(), ma._byte_low_frontier() - mid_hi) + + def test_multi_hop_walk_stops_at_first_opaque(self): + _, fa, ma = self._build_pair() + t1 = _ChainStub(low_byte=100, high_byte=100, transparent=True) + t2 = _ChainStub(low_byte=200, high_byte=260, transparent=False) + fa.bind_high_peer(t1) + t1.low_peer = fa + t1.high_peer = t2 + t2.low_peer = t1 + t2.high_peer = ma + ma.bind_low_peer(t2) + self.assertEqual(fa._current_gap_bytes(), 200 - fa._byte_high_frontier()) + self.assertIs(fa._growth_side_neighbor(), t2) + t2.transparent = True + self.assertIs(fa._growth_side_neighbor(), ma) + self.assertEqual( + fa._current_gap_bytes(), + ma._byte_low_frontier() - fa._byte_high_frontier(), + ) + + def test_drainable_credit_reads_walked_neighbor(self): + _, fa, ma = self._build_pair() + stub = _ChainStub(low_byte=64, high_byte=128, transparent=True) + fa.bind_high_peer(stub) + stub.low_peer = fa + stub.high_peer = ma + ma.bind_low_peer(stub) + + # Walked-through to the far end: its holes are credited iff lazy. + self.assertEqual(fa._peer_drainable_hole_bytes(), 0) # ma not lazy + ma.lazy_compaction = True + ma._free_phys_pages = torch.arange(3, dtype=torch.int64) + self.assertEqual(fa._peer_drainable_hole_bytes(), 3 * ma.entry_bytes_per_page) + # Opaque non-lazy middle blocks the far end's credit. + stub.transparent = False + self.assertEqual(fa._peer_drainable_hole_bytes(), 0) + + +class TestFloatMultiEndedAllocator(unittest.TestCase): + """Holes-first float middle: midpoint placement, in-place hole recycling, + larger-gap boundary extension, boundary absorption + park-on-empty + transparency, and the on-demand data movers (`make_room` boundary + relocation / leapfrog, `compact_holes` ordered pack) with data-move + verification through the fake KV marker. + + Guarded failure modes: a float that copies on steady-state churn, walls + off free space when empty, extends toward the tighter gap, moves more + than min(L_live, G) pages, loses data across relocation, or corrupts + v2p/p2v span/hole bookkeeping. + """ + + def _build_tri(self, n_state=8, n_float=32, n_full=32): + state = _make_mamba_spec("state", "up", layer_num=2) + fl = _make_mha_spec("swa", "float", layer_num=2) + full = _make_mha_spec("full", "down", layer_num=2) + total = ( + state.entry_bytes() * n_state + + fl.entry_bytes() * n_float + + full.entry_bytes() * n_full + ) + pool = UnifiedKVPool( + total_bytes=total, + sub_pool_specs=[full, fl, state], + device=_DEV, + enable_memory_saver=False, + ) + sa = MultiEndedAllocator( + kvcache=_FakeKVCache(pool.max_slots("state")), + unified_buffer=pool, + sub_pool_name="state", + device=_DEV, + is_id_owner=True, + ) + fkv = _FakeKVCache(pool.max_slots("swa")) + fla = FloatMultiEndedAllocator( + kvcache=fkv, + unified_buffer=pool, + sub_pool_name="swa", + device=_DEV, + is_id_owner=True, + ) + dkv = _FakeKVCache(pool.max_slots("full")) + da = MultiEndedAllocator( + kvcache=dkv, + unified_buffer=pool, + sub_pool_name="full", + device=_DEV, + is_id_owner=True, + ) + # Chain wiring state <-> float <-> full. + sa.bind_high_peer(fla) + fla.bind_low_peer(sa) + fla.bind_high_peer(da) + da.bind_low_peer(fla) + return pool, sa, fla, da, fkv + + def _stamp(self, alloc, kv, v): + kv.buf[alloc.virtual_to_physical[v]] = v + + def _interior_block(self, fla, blocks): + """The allocated block whose physical pages touch neither span + boundary (robust to the extend-direction policy).""" + for v in blocks: + pages = set(int(x) for x in fla.virtual_to_physical[v].tolist()) + if fla.low_wm_page not in pages and (fla.high_wm_page - 1) not in pages: + return v + raise AssertionError("no interior block in layout") + + def _check_float_state(self, fla, kv): + holes = set(int(x) for x in fla._free_phys_pages.tolist()) + span = range(fla.low_wm_page, fla.high_wm_page) + live = [p for p in span if p not in holes] + self.assertEqual(fla._live_pages(), len(live)) + for h in holes: + self.assertTrue(fla.low_wm_page < h < fla.high_wm_page - 1) + for p in live: + v = int(fla.physical_to_virtual[p].item()) + self.assertNotEqual(v, -1, f"live page {p} unbound") + self.assertEqual(int(fla.virtual_to_physical[v].item()), p) + self.assertEqual(int(kv.buf[p].item()), v, f"data lost at {p}") + + def test_midpoint_initial_placement(self): + _, _, fla, _, _ = self._build_tri() + self.assertTrue(fla._is_frontier_transparent()) + v = fla.alloc(4) + self.assertIsNotNone(v) + lo, hi = fla._region_bounds_pages() + self.assertEqual(fla.low_wm_page, lo + (hi - lo - 4) // 2) + self.assertEqual(fla._span_pages(), 4) + # Gap on BOTH sides. + gap_low, gap_high = fla._gap_pages() + self.assertGreater(gap_low, 0) + self.assertGreater(gap_high, 0) + + def test_holes_first_reuse_is_zero_copy(self): + _, _, fla, _, kv = self._build_tri() + va = fla.alloc(2) + vb = fla.alloc(2) + vc = fla.alloc(2) + for v in (va, vb, vc): + self._stamp(fla, kv, v) + span_before = (fla.low_wm_page, fla.high_wm_page) + fla.free(self._interior_block(fla, (va, vb, vc))) # interior -> holes + self.assertEqual(fla._hole_pages(), 2) + self.assertEqual((fla.low_wm_page, fla.high_wm_page), span_before) + vd = fla.alloc(2) # must recycle the holes in place + self._stamp(fla, kv, vd) + self.assertEqual(fla._hole_pages(), 0) + self.assertEqual((fla.low_wm_page, fla.high_wm_page), span_before) + self.assertEqual(len(fla._inverse_history), 0) # zero copies + self._check_float_state(fla, kv) + + def test_boundary_free_absorbed_at_the_deferred_point(self): + """Boundary holes shrink the span ZERO-COPY -- but the shrink is + DEFERRED out of `free`, which must stay host-sync-free (deciding how + far to walk needs the hole set on the host). `free` records the holes; + `_absorb_span_boundary_holes` (per-step flush / shortfall ladder) reclaims + the span. Skipping it is only ever conservative.""" + _, _, fla, _, kv = self._build_tri() + va = fla.alloc(2) + vb = fla.alloc(2) # extends one side; frees at that edge absorb + self._stamp(fla, kv, va) + self._stamp(fla, kv, vb) + span = fla._span_pages() + fla.free(vb) + # Deferred: span still claims the freed edge, the pages are holes. + self.assertEqual(fla._hole_pages(), 2) + self.assertEqual(fla._span_pages(), span) + self.assertEqual(fla._live_pages(), 2) # exact regardless + absorbed = fla._flush(urgent=False) + self.assertEqual(absorbed, 2) + self.assertEqual(fla._hole_pages(), 0) + self.assertEqual(fla._span_pages(), span - 2) + self.assertEqual(len(fla._inverse_history), 0) # zero copies + self._check_float_state(fla, kv) + + def test_free_is_host_sync_free(self): + """The per-decode-step property: no D2H anywhere in the float's free + (the base's lazy free earns this by deferring absorption to `_flush`; + the float now follows the same model).""" + from unittest import mock + + _, _, fla, _, kv = self._build_tri() + v = fla.alloc(4) + self._stamp(fla, kv, v) + with mock.patch.object( + torch.Tensor, "tolist", side_effect=AssertionError("tolist = D2H") + ), mock.patch.object( + torch.Tensor, "item", side_effect=AssertionError("item = D2H") + ), mock.patch.object( + torch, "unique", side_effect=AssertionError("unique = host sync") + ): + fla.free(v[:2], _pages=v[:2]) + + def test_deferred_absorption_reaches_the_same_state_as_eager(self): + """Derived property: deferring must not change WHERE the span lands, + only when. Two floats, identical ops; one absorbs after every free, + one only at the end.""" + _, _, f1, _, kv1 = self._build_tri() + _, _, f2, _, kv2 = self._build_tri() + for f, kv, absorb_each in ((f1, kv1, True), (f2, kv2, False)): + blocks = [f.alloc(2) for _ in range(3)] + for v in blocks: + self._stamp(f, kv, v) + for v in (blocks[2], blocks[0]): + f.free(v) + if absorb_each: + f._flush(urgent=False) + f2._flush(urgent=False) + self.assertEqual(f1.low_wm_page, f2.low_wm_page) + self.assertEqual(f1.high_wm_page, f2.high_wm_page) + self.assertEqual(f1._hole_pages(), f2._hole_pages()) + self.assertEqual(f1.available_size(), f2.available_size()) + + def test_park_on_empty_restores_transparency(self): + _, sa, fla, da, _ = self._build_tri() + base_gap = da._current_gap_bytes() + v = fla.alloc(4) + self.assertLess(da._current_gap_bytes(), base_gap) # float blocks + fla.free(v) + self.assertTrue(fla._is_frontier_transparent()) + self.assertEqual(fla._hole_pages(), 0) + self.assertEqual(da._current_gap_bytes(), base_gap) # sees through again + self.assertEqual(sa._current_gap_bytes(), base_gap) + + def test_extends_toward_larger_gap(self): + _, _, fla, da, kv = self._build_tri() + v = fla.alloc(4) + self._stamp(fla, kv, v) + # Consume most of the high gap with the full end; low gap now larger. + self.assertIsNotNone(da.alloc(24)) + gap_low, gap_high = fla._gap_pages() + self.assertGreater(gap_low, gap_high) + lo_before = fla.low_wm_page + hi_before = fla.high_wm_page + v2 = fla.alloc(2) + self.assertIsNotNone(v2) + self._stamp(fla, kv, v2) + self.assertEqual(fla.low_wm_page, lo_before - 2) # grew low side + self.assertEqual(fla.high_wm_page, hi_before) + self._check_float_state(fla, kv) + + def test_available_is_max_gap_plus_holes(self): + _, _, fla, da, _ = self._build_tri() + va = fla.alloc(2) + vb = fla.alloc(2) + vc = fla.alloc(2) + fla.free(self._interior_block(fla, (va, vb, vc))) + self.assertEqual(fla._hole_pages(), 2) + gap_low, gap_high = fla._gap_pages() + self.assertEqual(fla.available_size(), max(gap_low, gap_high) + 2) + del da + + def test_make_room_boundary_relocation(self): + _, _, fla, da, kv = self._build_tri() + v = fla.alloc(6) + self._stamp(fla, kv, v) + epp = fla.entry_bytes_per_page + _, gap_high = fla._gap_pages() + ask = (gap_high + 3) * epp # 3 pages beyond the current high gap + opened = fla.make_room(side="high", min_bytes=ask) + self.assertGreaterEqual(opened, ask) + # Cost min(L_live, G): moved exactly the 3 boundary pages, not 6. + moved = sum(int(s.numel()) for s, _, _ in fla._inverse_history) + self.assertEqual(moved, 3) + self._check_float_state(fla, kv) + # The opened space is real: the full end can now take it. + self.assertGreaterEqual(da.available_size(), 3) + + def test_make_room_leapfrog_cost_bounded_by_live(self): + _, _, fla, _, kv = self._build_tri() + v = fla.alloc(2) # tiny live mass + self._stamp(fla, kv, v) + epp = fla.entry_bytes_per_page + _, gap_high = fla._gap_pages() + ask = (gap_high + 10) * epp # demand >> live + opened = fla.make_room(side="high", min_bytes=ask) + self.assertGreaterEqual(opened, ask) + moved = sum(int(s.numel()) for s, _, _ in fla._inverse_history) + self.assertEqual(moved, 2) # min(L_live, G) == L_live + self._check_float_state(fla, kv) + + def test_make_room_impossible_leaves_state_unchanged(self): + _, _, fla, _, kv = self._build_tri(n_float=8) + v = fla.alloc(6) + self._stamp(fla, kv, v) + lo, hi = fla._region_bounds_pages() + epp = fla.entry_bytes_per_page + snapshot = ( + fla.low_wm_page, + fla.high_wm_page, + fla._hole_pages(), + len(fla._inverse_history), + ) + opened = fla.make_room(side="high", min_bytes=(hi - lo) * epp) + self.assertLess(opened, (hi - lo) * epp) + self.assertEqual( + snapshot, + ( + fla.low_wm_page, + fla.high_wm_page, + fla._hole_pages(), + len(fla._inverse_history), + ), + ) + self._check_float_state(fla, kv) + + def test_make_room_uses_far_holes_before_far_gap(self): + _, _, fla, _, kv = self._build_tri() + va = fla.alloc(2) + vb = fla.alloc(2) + vc = fla.alloc(2) + for v in (va, vb, vc): + self._stamp(fla, kv, v) + lo_before = fla.low_wm_page + fla.free(self._interior_block(fla, (va, vb, vc))) # 2 interior holes + self.assertEqual(fla._hole_pages(), 2) + epp = fla.entry_bytes_per_page + _, gap_high = fla._gap_pages() + opened = fla.make_room(side="high", min_bytes=(gap_high + 2) * epp) + self.assertGreaterEqual(opened, (gap_high + 2) * epp) + # The two holes absorbed the two moved pages: low side untouched. + self.assertEqual(fla.low_wm_page, lo_before) + self.assertEqual(fla._hole_pages(), 0) + self._check_float_state(fla, kv) + + def test_compact_holes_ordered_pack(self): + _, _, fla, _, kv = self._build_tri() + vs = [fla.alloc(2) for _ in range(4)] + for v in vs: + self._stamp(fla, kv, v) + fla.free(vs[1]) # interleaved holes + span_before = fla._span_pages() + moved = fla.compact_holes(retreat_side="high") + self.assertEqual(fla._hole_pages(), 0) + self.assertEqual(fla._span_pages(), span_before - 2) + self.assertGreater(moved, 0) + self._check_float_state(fla, kv) + + def test_bind_peer_raises_on_float(self): + _, sa, fla, _, _ = self._build_tri() + with self.assertRaises(AssertionError): + fla.bind_peer(sa) + + if __name__ == "__main__": unittest.main() diff --git a/test/registered/unit/mem_cache/test_unified_capacity_memo.py b/test/registered/unit/mem_cache/test_unified_capacity_memo.py new file mode 100644 index 000000000..c3611f5c9 --- /dev/null +++ b/test/registered/unit/mem_cache/test_unified_capacity_memo.py @@ -0,0 +1,299 @@ +# 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. +# ============================================================================== +"""Epoch-memoized capacity views on the allocator chain (2-pool subset). + +The capacity views (`available_size` / `schedulable_available_size` per band, +plus the composite joint view) are pure functions of a handful of +CPU-resident fields across the chain; schedulers read them O(queue) times +between mutations. `_CapacityField` descriptors bump `_capacity_epoch` on +every rebind, so the memos invalidate by construction. + +The failure mode being guarded: a memo serving a STALE value after a mutation +the epoch machinery missed — either a new mutation site writing a field the +descriptors don't cover, or an in-place write that bypasses `__set__`. Stale +capacity is silent over-/under-admission, not a crash. Hence: + + * every mutation kind is followed by memo == fresh-recompute assertions; + * a randomized op-sequence property test (seeded) catches interactions no + hand-written sequence covers; + * a deliberate descriptor-bypassing write must be caught by the IDLE check + (`verify_byte_accounting`) — readers cannot detect it, the battery must. + +The tri-pool cases (float span fields, three-band joint view) join at the +tri phase; this file pins the 2-pool chain form they build on. + + python -m pytest test/registered/unit/mem_cache/test_unified_capacity_memo.py -v +""" + +import random +import unittest + +from sglang.test.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=20, suite="base-a-test-cpu") + + +def _build(lazy: bool): + # Function-scope import: the fixture is a TestCase subclass, and a + # module-scope binding would make pytest collect its tests AGAIN here. + from test_multi_ended_allocator import ( + TestUnifiedSWATokenToKVPoolAllocator as _SwaFixture, + ) + + inst = _SwaFixture([m for m in dir(_SwaFixture) if m.startswith("test_")][0]) + pool, allocator, kvcache = inst._build() + allocator.full_attn_allocator.lazy_compaction = lazy + allocator.swa_attn_allocator.lazy_compaction = lazy + allocator.lazy_compaction = lazy + return inst, allocator, kvcache + + +class TestCapacityMemoCoherence(unittest.TestCase): + def _assert_memos_fresh(self, allocator): + """Every memoized capacity view must equal a fresh recompute.""" + self.assertEqual( + allocator.available_size(), allocator._compute_available_size() + ) + for band in ( + allocator.full_attn_allocator, + allocator.swa_attn_allocator, + ): + self.assertEqual( + band.available_size(), + band._available_tokens(), + msg=f"stale available_size memo on {band.sub_pool_name!r}", + ) + self.assertEqual( + band.schedulable_available_size(), + band._available_tokens( + extra_gap_bytes=band._peer_drainable_hole_bytes() + ), + msg=f"stale schedulable memo on {band.sub_pool_name!r}", + ) + self.assertEqual(allocator.verify_byte_accounting(), []) + + def test_memos_track_every_mutation_kind(self): + for lazy in (False, True): + with self.subTest(lazy_compaction=lazy): + inst, allocator, kvcache = _build(lazy) + self._assert_memos_fresh(allocator) + + v1 = inst._alloc(allocator, kvcache, 8) # both-side bind + self.assertIsNotNone(v1) + self._assert_memos_fresh(allocator) + + allocator.free_swa(v1[2:6]) # swa-side tombstones + self._assert_memos_fresh(allocator) + + inst._free(allocator, kvcache, v1) # both-side free + self._assert_memos_fresh(allocator) + + v2 = inst._alloc(allocator, kvcache, 4) + self.assertIsNotNone(v2) + allocator.free_group_begin() # grouped free path + allocator.free(v2) + allocator.free_group_end() + self._assert_memos_fresh(allocator) + + if lazy: + allocator.full_attn_allocator._flush(urgent=True) + self._assert_memos_fresh(allocator) + + allocator.clear() + self._assert_memos_fresh(allocator) + + def test_random_op_sequence_value_identity(self): + """Property: after ANY mutation sequence, memoized views equal fresh + recomputes. Seeded (deterministic) — the sequences cover interleavings + (alloc / partial swa free / grouped free / flush / clear) that no + hand-written case enumerates.""" + rng = random.Random(0xC0FFEE) + for lazy in (False, True): + with self.subTest(lazy_compaction=lazy): + inst, allocator, kvcache = _build(lazy) + live = [] + for step in range(60): + op = rng.choice(("alloc", "free", "free_swa", "flush", "clear")) + if op == "alloc": + v = inst._alloc(allocator, kvcache, rng.choice((1, 2, 4))) + if v is not None: + live.append(v) + elif op == "free" and live: + inst._free(allocator, kvcache, live.pop()) + elif op == "free_swa" and live: + v = live[-1] + if v.numel() > 1: + allocator.free_swa(v[: v.numel() // 2]) + elif op == "flush": + allocator.full_attn_allocator._flush(urgent=True) + allocator.swa_attn_allocator._flush(urgent=True) + elif op == "clear": + allocator.clear() + live.clear() + self._assert_memos_fresh(allocator) + + def test_bypassing_write_is_caught_by_the_idle_check(self): + inst, allocator, kvcache = _build(lazy=False) + v = inst._alloc(allocator, kvcache, 8) + self.assertIsNotNone(v) + fa = allocator.full_attn_allocator + # Prime every memo at the current epoch. + allocator.available_size() + fa.available_size() + fa.schedulable_available_size() + # Mutate capacity state BYPASSING the _CapacityField descriptor -- the + # epoch does not move, so the memos go stale undetectably for readers... + fa.__dict__["watermark_physical"] = fa.watermark_physical + 2 + # ...but the idle-time coherence check must flag it. + violations = allocator.verify_byte_accounting() + self.assertTrue( + any("stale" in msg for msg in violations), + msg=f"bypassing write not caught: {violations}", + ) + + def test_joint_memo_invalidates_on_swa_only_mutation(self): + """The joint view depends on the swa end's frontier through the chain + walk; an swa-side-only mutation (tombstoning) must invalidate the + composite memo even though the full side never moved.""" + inst, allocator, kvcache = _build(lazy=True) + v = inst._alloc(allocator, kvcache, 8) + self.assertIsNotNone(v) + before = allocator.available_size() + allocator.free_swa(v[:4]) # swa band only + after = allocator.available_size() + self.assertEqual(after, allocator._compute_available_size()) + self.assertGreaterEqual(after, before) # holes only ever add room + + def test_float_only_span_move_invalidates_every_memo(self): + """A hole-free float alloc rebinds NO free-list and has no watermark -- + the span fields are its ONLY capacity state. If they are not + `_CapacityField` descriptors, the float's own memo AND both + neighbours' (the span flips transparency, walling off their gaps) + keep serving pre-move values. + + Driven on a hand-wired end+float+end chain (the composite arrives + with the tri phase); the float is exercised alone so no end-pool + descriptor write can mask a missing span bump.""" + from test_multi_ended_allocator import TestFloatMultiEndedAllocator + + inst = TestFloatMultiEndedAllocator( + [m for m in dir(TestFloatMultiEndedAllocator) if m.startswith("test_")][0] + ) + _pool, sa, fla, da, _kv = inst._build_tri() + self.assertEqual(fla._hole_pages(), 0) # hole-free extension path + self.assertTrue(fla._is_frontier_transparent()) + + # Prime every memo while the float is empty/transparent. + float_cached = fla.available_size() + low_end_cached = sa.available_size() + high_end_cached = da.available_size() + + v = fla.alloc(4) # float-only mutation: span move, no end-pool write + self.assertIsNotNone(v) + self.assertFalse(fla._is_frontier_transparent()) # span now opaque + + self.assertEqual(fla.available_size(), fla._available_tokens()) + self.assertEqual(sa.available_size(), sa._available_tokens()) + self.assertEqual(da.available_size(), da._available_tokens()) + # The opaque midpoint span must actually reduce what the neighbours + # see, i.e. the memos above were not merely re-serving primed values. + self.assertLess(sa.available_size(), low_end_cached) + self.assertLess(da.available_size(), high_end_cached) + self.assertLessEqual(fla.available_size(), float_cached) + + def test_bind_rewiring_bumps_the_epoch(self): + """Rewiring changes what the chain walks see; a memo primed before a + re-bind must not survive it.""" + inst, allocator, kvcache = _build(lazy=False) + fa = allocator.full_attn_allocator + e0 = fa._chain_capacity_epoch() + fa.bind_peer(allocator.swa_attn_allocator) # re-bind (same peer) + self.assertGreater(fa._chain_capacity_epoch(), e0) + + +class TestTriCapacityMemoCoherence(unittest.TestCase): + """Tri-composite twins of the 2-pool cases: the joint view walks THREE + bands (mamba end, swa float, full end), so a mutation on ANY of them must + invalidate the composite memo — including the two mutations only the tri + has: a mamba-end state draw and a float span move behind the composite.""" + + def _build_tri(self, lazy=False): + from test_unified_tri_pool import TestUnifiedTriPool + + inst = TestUnifiedTriPool( + [m for m in dir(TestUnifiedTriPool) if m.startswith("test_")][0] + ) + pool, allocator, kvcache, mamba_kv = inst._build(lazy_compaction=lazy) + return inst, allocator + + def _assert_memos_fresh(self, allocator): + self.assertEqual( + allocator.available_size(), allocator._compute_available_size() + ) + for band in ( + allocator.full_attn_allocator, + allocator.swa_attn_allocator, + allocator.mamba_allocator, + ): + self.assertEqual( + band.available_size(), + band._available_tokens(), + msg=f"stale available_size memo on {band.sub_pool_name!r}", + ) + self.assertEqual(allocator.verify_byte_accounting(), []) + + def test_memos_track_tri_mutation_kinds(self): + for lazy in (False, True): + with self.subTest(lazy_compaction=lazy): + inst, allocator = self._build_tri(lazy) + ma = allocator.mamba_allocator + self._assert_memos_fresh(allocator) + + v1 = allocator.alloc(8) # composite alloc (full + swa bind) + self.assertIsNotNone(v1) + self._assert_memos_fresh(allocator) + + s1 = ma.alloc(2) # mamba end alloc + self.assertIsNotNone(s1) + self._assert_memos_fresh(allocator) + + allocator.free_swa(v1[2:6]) # interior float holes + self._assert_memos_fresh(allocator) + + allocator.free(v1) # both-side free + self._assert_memos_fresh(allocator) + + ma.free(s1) # mamba free + self._assert_memos_fresh(allocator) + + allocator.clear() + ma.clear() + self._assert_memos_fresh(allocator) + + def test_joint_memo_invalidates_on_mamba_only_mutation(self): + """The joint view depends on the mamba end's frontier through the + chain walk; a mamba-only mutation must invalidate the composite memo + even though neither KV side moved.""" + inst, allocator = self._build_tri() + ma = allocator.mamba_allocator + before = allocator.available_size() + slots = ma.alloc(4) + self.assertIsNotNone(slots) + after = allocator.available_size() + self.assertEqual(after, allocator._compute_available_size()) + self.assertLessEqual(after, before) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/mem_cache/test_unified_free_no_host_sync.py b/test/registered/unit/mem_cache/test_unified_free_no_host_sync.py index e9d2011c2..9e6245706 100644 --- a/test/registered/unit/mem_cache/test_unified_free_no_host_sync.py +++ b/test/registered/unit/mem_cache/test_unified_free_no_host_sync.py @@ -65,12 +65,55 @@ def _paged_allocator(lazy: bool): # 1. tombstone scatters # -------------------------------------------------------------------------- +_TABLES = {"virtual_to_physical", "physical_to_virtual"} + +# Methods that MUST tombstone through index_fill_. Explicit, because "this +# method writes a tombstone" is a design fact per method, not something a scan +# can infer -- but `test_every_allocator_free_path_is_listed` below fails if a +# new allocator arrives with its own free path and is not added here. _TOMBSTONE_METHODS = [ (mea.MultiEndedAllocator, "_free_lazy"), (mea.MultiEndedAllocator, "free"), (mea.MultiEndedAllocator, "_commit_move_batch"), + (mea.FloatMultiEndedAllocator, "free"), + (mea.FloatMultiEndedAllocator, "make_room"), + (mea.FloatMultiEndedAllocator, "_relocate_to_positions"), ] -_TABLES = {"virtual_to_physical", "physical_to_virtual"} + + +def _allocators_in_module(): + """Every allocator class DEFINED in multi_ended_allocator (not imported).""" + return sorted( + ( + c + for c in vars(mea).values() + if isinstance(c, type) + and c.__module__ == mea.__name__ + and "Allocator" in c.__name__ + ), + key=lambda c: c.__name__, + ) + + +def _table_touching_methods(): + """Every own method of every allocator whose source names a page table. + + DISCOVERY, not a list: a hardcoded list stops guarding the moment a new + allocator class arrives with its own free path -- which is what happened + when FloatMultiEndedAllocator was added and inherited no coverage. + """ + out = [] + for cls in _allocators_in_module(): + for name, fn in vars(cls).items(): + if not inspect.isfunction(fn): + continue + try: + src = inspect.getsource(fn) + except OSError: + continue + if any(f".{t}[" in src for t in _TABLES): + out.append((cls, name)) + return sorted(out, key=lambda pair: (pair[0].__name__, pair[1])) def _scalar_index_assignments(fn): @@ -102,13 +145,22 @@ def _scalar_index_assignments(fn): continue if isinstance(tgt.slice, ast.Slice): continue + # A CONSTANT integer index (`t[0] = 0`, `t[-1] = -1`) is a + # single-element sentinel write, not the tensor-index tombstone this + # guard exists to find, and the only such writes are in `clear()`. + if _is_scalar_literal(tgt.slice): + continue bad.append(ast.unparse(node)) return bad class TestTombstonesDoNotCrossTheBus(unittest.TestCase): def test_no_scalar_index_assignment(self): - for cls, name in _TOMBSTONE_METHODS: + discovered = _table_touching_methods() + self.assertGreaterEqual( + len(discovered), len(_TOMBSTONE_METHODS), "discovery scan went blind" + ) + for cls, name in discovered: with self.subTest(method=f"{cls.__name__}.{name}"): bad = _scalar_index_assignments(getattr(cls, name)) self.assertEqual( @@ -132,6 +184,41 @@ class TestTombstonesDoNotCrossTheBus(unittest.TestCase): self.assertEqual(len(_scalar_index_assignments(_offender)), 1) + def test_every_allocator_free_path_is_listed(self): + """The positive list must name every allocator that owns a free path. + + REGRESSION: the list used to hold three MultiEndedAllocator methods, so + adding FloatMultiEndedAllocator with its own `free` silently dropped that + free path out of coverage -- and it shipped a scalar tombstone. Fail here + instead, loudly, the next time an allocator arrives. + """ + listed = {(cls.__name__, name) for cls, name in _TOMBSTONE_METHODS} + for cls in _allocators_in_module(): + for name, fn in vars(cls).items(): + if not inspect.isfunction(fn): + continue + try: + src = inspect.getsource(fn) + except OSError: + continue + # WRITES a page table -- either correctly (index_fill_) or in the + # banned scalar form the scan below catches. A method that only + # READS a table has nothing to tombstone. + if not ( + any(f"{t}.index_fill_" in src for t in _TABLES) + or _scalar_index_assignments(fn) + ): + continue + self.assertIn( + (cls.__name__, name), + listed, + msg=( + f"{cls.__name__}.{name} writes a page table but is not in " + f"_TOMBSTONE_METHODS, so the index_fill_ guard does not " + f"cover it. Add it." + ), + ) + def test_free_paths_actually_use_index_fill(self): """Positive form, so deleting the scatter entirely cannot pass.""" for cls, name in _TOMBSTONE_METHODS: @@ -302,5 +389,110 @@ class TestEveryUnifiedAllocatorOverridesFreeSegment(unittest.TestCase): self.assertIn("free_page_reps_group", inspect.getsource(cls)) +class TestFreeSwaWindowRatchetNoHostSync(unittest.TestCase): + """The per-decode-step SWA window ratchet frees a CONTIGUOUS row slice + with host-int, page-aligned bounds — the same shape `free_segment` was + built for. `free_swa(..., start_pos=)` must therefore reach the swa side + with caller-derived page ids: no `torch.unique` (data-dependent shape = + host sync) and no stale-slot `.item()` on the per-step path. + + Poisoning the ops is the decisive form (a textual guard can be fooled). + """ + + PS = 4 + + def _swa_composite(self, lazy=True): + from test_multi_ended_allocator import _FakeKVCache, _make_mha_spec + + from sglang.srt.mem_cache.unified_memory_pool import UnifiedKVPool + + full = _make_mha_spec("full", "up", layer_num=4) + swa = _make_mha_spec("swa", "down", layer_num=2) + total = 64 * full.entry_bytes() + 64 * swa.entry_bytes() + pool = UnifiedKVPool( + total_bytes=total, + sub_pool_specs=[full, swa], + device="cpu", + enable_memory_saver=False, + page_size=self.PS, + ) + + class _KV: + def __init__(self, p): + self.full_kv_pool = _FakeKVCache(p.max_slots("full")) + self.swa_kv_pool = _FakeKVCache(p.max_slots("swa")) + + def attach_allocators(self, **kwargs): + pass + + return mea.UnifiedSWATokenToKVPoolAllocator( + unified_buffer=pool, + kvcache=_KV(pool), + device="cpu", + full_max_total_num_tokens=64, + swa_max_total_num_tokens=64, + page_size=self.PS, + need_sort=False, + forward_stream=None, + lazy_compaction=lazy, + ) + + def test_ratchet_shape_free_swa_never_syncs(self): + """Aligned bounds (the ratchet guarantees them at ps>1): no unique, + no item -- on the lazy production config.""" + alloc = self._swa_composite(lazy=True) + v = alloc.alloc(8 * self.PS) + self.assertIsNotNone(v) + with mock.patch.object( + torch, "unique", side_effect=AssertionError("unique = host sync") + ), mock.patch.object( + torch.Tensor, "item", side_effect=AssertionError("item = host sync") + ): + alloc.free_swa(v[: 4 * self.PS], start_pos=0) + alloc.free_swa(v[4 * self.PS :], start_pos=4 * self.PS) + + def test_unaligned_start_pos_still_no_sync(self): + """`_page_reps_pieces` covers a misaligned start with a second piece; + the sync-free property must not depend on alignment.""" + alloc = self._swa_composite(lazy=True) + v = alloc.alloc(8 * self.PS) + with mock.patch.object( + torch, "unique", side_effect=AssertionError("unique = host sync") + ): + alloc.free_swa(v[1 : 5 * self.PS], start_pos=1) + + def test_start_pos_path_matches_the_fallback_end_state(self): + """Derived property: the stride-rep path and the dedup fallback must + leave IDENTICAL allocator state (v2p tombstones, capacity).""" + for lazy in (True, False): + with self.subTest(lazy=lazy): + a1 = self._swa_composite(lazy=lazy) + a2 = self._swa_composite(lazy=lazy) + v1 = a1.alloc(6 * self.PS) + v2 = a2.alloc(6 * self.PS) + self.assertTrue(torch.equal(v1, v2)) + a1.free_swa(v1[: 4 * self.PS], start_pos=0) + a2.free_swa(v2[: 4 * self.PS]) # fallback (radix shape) + self.assertTrue( + torch.equal( + a1.swa_attn_allocator.virtual_to_physical, + a2.swa_attn_allocator.virtual_to_physical, + ) + ) + self.assertEqual(a1.available_size(), a2.available_size()) + self.assertEqual( + a1.swa_attn_allocator.schedulable_available_size(), + a2.swa_attn_allocator.schedulable_available_size(), + ) + + def test_double_ratchet_is_filtered_not_crashed(self): + """Freeing an already-tombstoned range again must no-op through the + liveness filter (radix eviction and the ratchet can overlap).""" + alloc = self._swa_composite(lazy=True) + v = alloc.alloc(4 * self.PS) + alloc.free_swa(v, start_pos=0) + alloc.free_swa(v, start_pos=0) # all tombstoned -> filtered to empty + + if __name__ == "__main__": unittest.main() diff --git a/test/registered/unit/mem_cache/test_unified_npool_sweep.py b/test/registered/unit/mem_cache/test_unified_npool_sweep.py new file mode 100644 index 000000000..4e81da222 --- /dev/null +++ b/test/registered/unit/mem_cache/test_unified_npool_sweep.py @@ -0,0 +1,271 @@ +# 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. +# ============================================================================== +"""N-sub-pool construction sweep for ``UnifiedKVPool``. + +The pool accepts N sub-pool specs: exactly one grow-up END, exactly one +grow-down END, and >= 0 "float" MIDDLE pools between their frontiers. These +tests pin the constructor contract the N-pool chain machinery builds on: + + - canonical chain order ``[up end, floats (input order), down end]`` — + input list order is irrelevant (2-pool configs stay byte-identical); + - by-name geometry (``max_slots = total_bytes // entry_bytes``, + ``min_slot_index`` past the shared reserved floor) independent of N; + - the reserved slot-0 sink covers EVERY sub-pool's page-0 dummy-write + envelope, floats included (mamba stays page_size=1); + - validation: unique names, exactly one up + one down, >= 2 specs, and + per-spec ``_allowed_grow_directions`` narrowing; + - float sub-pool views build and round-trip like end-pool views (all views + span the whole buffer at anchor 0; keeping the bands disjoint is the + allocators' job). + +Pure CPU geometry — no allocator, no GPU. + + python -m pytest test/registered/unit/mem_cache/test_unified_npool_sweep.py -v +""" + +import unittest + +import torch + +from sglang.srt.mem_cache.unified_memory_pool import ( + MambaSubPoolSpec, + MHASubPoolSpec, + MLASubPoolSpec, + UnifiedKVPool, +) +from sglang.test.ci.ci_register import register_cpu_ci + +# Plain unittest.TestCase, importing only ci_register -- the deliberate +# hermetic convention of the pool-geometry tests in this directory (see +# test_multi_ended_allocator.py): no heavy sglang.test.test_utils import +# chain, so the suite runs in a lean torch-only environment. +register_cpu_ci(est_time=30, suite="base-a-test-cpu") + +_DEV = "cpu" + + +def _mha( + name: str, + grow_direction: str, + *, + layer_num: int = 2, + head_num: int = 2, + head_dim: int = 8, +) -> MHASubPoolSpec: + return MHASubPoolSpec( + name=name, + layer_num=layer_num, + grow_direction=grow_direction, + head_num=head_num, + head_dim=head_dim, + store_dtype=torch.bfloat16, + ) + + +def _mla(name: str, grow_direction: str, *, layer_num: int = 2) -> MLASubPoolSpec: + return MLASubPoolSpec( + name=name, + layer_num=layer_num, + grow_direction=grow_direction, + kv_lora_rank=16, + qk_rope_head_dim=8, + store_dtype=torch.bfloat16, + ) + + +def _mamba(name: str, grow_direction: str, *, layer_num: int = 2) -> MambaSubPoolSpec: + return MambaSubPoolSpec( + name=name, + layer_num=layer_num, + grow_direction=grow_direction, + conv_state_shapes=((4, 6),), + conv_dtype=torch.bfloat16, + temporal_state_shape=(2, 4, 4), + temporal_dtype=torch.float32, + ) + + +def _make_pool(specs, *, total_bytes: int = 1 << 20, page_size: int = 1): + return UnifiedKVPool( + total_bytes=total_bytes, + sub_pool_specs=specs, + device=_DEV, + enable_memory_saver=False, + page_size=page_size, + ) + + +def _chain_names(pool: UnifiedKVPool): + return [s.name for s in pool.sub_pool_specs] + + +class TestNPoolCanonicalOrder(unittest.TestCase): + def test_two_pool_input_order_irrelevant(self): + for specs in ( + [_mha("full", "down"), _mamba("mamba", "up")], + [_mamba("mamba", "up"), _mha("full", "down")], + ): + pool = _make_pool(specs) + self.assertEqual(_chain_names(pool), ["mamba", "full"]) + + def test_three_pool_float_in_the_middle(self): + for specs in ( + [_mha("full", "down"), _mha("swa", "float"), _mamba("conv", "up")], + [_mha("swa", "float"), _mamba("conv", "up"), _mha("full", "down")], + [_mamba("conv", "up"), _mha("full", "down"), _mha("swa", "float")], + ): + pool = _make_pool(specs) + self.assertEqual(_chain_names(pool), ["conv", "swa", "full"]) + + def test_four_pool_float_input_order_preserved(self): + pool = _make_pool( + [ + _mha("full", "down"), + _mha("f1", "float"), + _mamba("state", "up"), + _mha("f0", "float", layer_num=1), + ] + ) + # Ends canonical; floats keep INPUT order between them. + self.assertEqual(_chain_names(pool), ["state", "f1", "f0", "full"]) + + def test_by_name_geometry_independent_of_n(self): + two = _make_pool([_mha("full", "down"), _mamba("mamba", "up")]) + three = _make_pool( + [_mha("full", "down"), _mha("swa", "float"), _mamba("mamba", "up")] + ) + for name in ("full", "mamba"): + self.assertEqual( + two.max_slots(name), + two.total_bytes // two.spec(name).entry_bytes(), + ) + self.assertEqual(two.max_slots(name), three.max_slots(name)) + for pool in (two, three): + for s in pool.sub_pool_specs: + self.assertEqual(pool.anchor_bytes(s.name), 0) + + +class TestNPoolValidation(unittest.TestCase): + def test_duplicate_names_rejected(self): + with self.assertRaisesRegex(AssertionError, "unique"): + _make_pool([_mha("x", "down"), _mamba("x", "up")]) + + def test_fewer_than_two_specs_rejected(self): + with self.assertRaisesRegex(AssertionError, ">= 2 sub-pools"): + _make_pool([_mha("full", "down")]) + + def test_two_ups_rejected(self): + with self.assertRaisesRegex(AssertionError, "exactly one grow-up"): + _make_pool([_mha("a", "up"), _mamba("b", "up")]) + + def test_missing_down_end_rejected(self): + with self.assertRaisesRegex(AssertionError, "exactly one grow-up"): + _make_pool([_mha("a", "up"), _mha("b", "float")]) + + def test_missing_up_end_rejected(self): + with self.assertRaisesRegex(AssertionError, "exactly one grow-up"): + _make_pool([_mha("a", "down"), _mha("b", "float"), _mha("c", "float")]) + + def test_bogus_direction_rejected_at_spec_level(self): + with self.assertRaisesRegex(AssertionError, "grow_direction"): + _mha("a", "sideways") + + def test_float_accepted_on_all_cache_spec_kinds(self): + # Every cache-class spec kind may float (the chain decides placement). + pool = _make_pool( + [ + _mamba("state", "up"), + _mha("f_mha", "float"), + _mla("f_mla", "float", layer_num=1), + _mamba("f_mamba", "float", layer_num=1), + _mha("full", "down"), + ] + ) + self.assertEqual( + _chain_names(pool), ["state", "f_mha", "f_mla", "f_mamba", "full"] + ) + + +class TestReservedFloorWithFloats(unittest.TestCase): + def test_float_page_envelope_extends_the_sink(self): + # The float MHA has the largest page-0 envelope; every pool's + # min_slot_index must clear it (mamba is page_size=1 and excluded from + # the page-aware term, but still must clear the byte floor). + page_size = 4 + big_float = _mha("swa", "float", layer_num=8, head_num=4, head_dim=32) + specs = [_mamba("state", "up"), big_float, _mha("full", "down")] + pool = _make_pool(specs, total_bytes=1 << 22, page_size=page_size) + floor = max( + max(s.entry_bytes() for s in specs), + page_size * big_float.entry_bytes(), + page_size * specs[2].entry_bytes(), + ) + for s in specs: + e = s.entry_bytes() + self.assertEqual(pool.min_slot_index(s.name), (floor + e - 1) // e) + + def test_too_small_buffer_fails_loud(self): + # 2048 B with page_size=16 and 128 B/entry MHA specs: the page-0 sink + # (16*128 = 2048 B) consumes the whole buffer -> min_slot_index == + # max_slots for the MHA pools -> no allocatable slot -> loud error. + with self.assertRaisesRegex(RuntimeError, "no room"): + _make_pool( + [_mamba("state", "up"), _mha("swa", "float"), _mha("full", "down")], + total_bytes=2048, + page_size=16, + ) + + +class TestFloatViews(unittest.TestCase): + def test_float_mha_views_shape_and_roundtrip(self): + page_size = 2 + spec = _mha("swa", "float", layer_num=3, head_num=2, head_dim=8) + pool = _make_pool( + [_mamba("state", "up"), spec, _mha("full", "down")], + total_bytes=1 << 20, + page_size=page_size, + ) + k_views, v_views = pool.mha_views_for("swa") + self.assertEqual(len(k_views), spec.layer_num) + self.assertEqual(len(v_views), spec.layer_num) + num_pages = pool.max_slots("swa") // page_size + blocks = 2 * spec.layer_num # K at block 2l, V at 2l+1 + n_rows = num_pages * blocks * page_size + for k in (*k_views, *v_views): + # Stock 3-D per-layer MHA signature; the row index is the + # kernel-facing id, each view's storage_offset folding in its block + # origin (see `build_mha_views`). + self.assertEqual(tuple(k.shape), (n_rows, spec.head_num, spec.head_dim)) + # Round-trip: a float view is a real strided window into _raw. + slot = pool.min_slot_index("swa") + row = (slot // page_size) * (page_size * blocks) + slot % page_size + pattern = ( + torch.arange(spec.head_num * spec.head_dim, dtype=torch.float32) + .reshape(spec.head_num, spec.head_dim) + .to(torch.bfloat16) + ) + k_views[1][row] = pattern + torch.testing.assert_close(k_views[1][row], pattern) + + def test_float_mamba_views_zero_visible(self): + pool = _make_pool( + [_mamba("state", "up"), _mamba("fstate", "float"), _mha("full", "down")] + ) + conv_views, temporal = pool.mamba_views_for("fstate") + self.assertTrue(all(v.eq(0).all() for v in conv_views)) + self.assertTrue(temporal.eq(0).all()) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/mem_cache/test_unified_tri_pool.py b/test/registered/unit/mem_cache/test_unified_tri_pool.py new file mode 100644 index 000000000..5aae3bdb0 --- /dev/null +++ b/test/registered/unit/mem_cache/test_unified_tri_pool.py @@ -0,0 +1,1404 @@ +# 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. +# ============================================================================== +"""Tri-pool composite (`UnifiedMambaSWATokenToKVPoolAllocator`) -- full KV + +SWA KV + mamba/conv state in ONE unified byte buffer, chain +``[mamba (up END) | swa (FLOAT) | full (down END)]``. + +Pinned contracts (each guards a distinct failure mode): + - chain wiring + the swa side being a `FloatMultiEndedAllocator`; + - the JOINT `available_size()` feasibility contract: `alloc(N)` for + N == available_size() must succeed (full extends into the high band + first, then the float extends a single side — the predicate models + exactly that order; over-promising here is the fail-loud + `alloc_with_virtual` assert, i.e. a crash in production); + - `free_swa` tombstones become float HOLES recycled IN PLACE by later + allocs (steady-state SWA churn == zero copies), while the full side + keeps the token; + - the per-request state surface (`UnifiedMambaSlotAllocator` over the + mamba END) and its independence from the token surface; + - urgent flushes drain the two ENDS but never touch the float's holes. + +Pure CPU; fakes stand in for the KV pools (data markers verify moves). + + python -m pytest test/registered/unit/mem_cache/test_unified_tri_pool.py -v +""" + +import inspect +import unittest + +import torch + +import sglang.srt.mem_cache.multi_ended_allocator as mea +from sglang.srt.mem_cache.multi_ended_allocator import ( + FloatMultiEndedAllocator, + UnifiedMambaSWATokenToKVPoolAllocator, +) +from sglang.srt.mem_cache.unified_memory_pool import ( + MambaSubPoolSpec, + MHASubPoolSpec, + UnifiedKVPool, + UnifiedMambaSlotAllocator, + init_unified_mamba_swa_pools, +) +from sglang.test.ci.ci_register import register_cpu_ci + +# Hermetic convention of this directory's pool tests: plain unittest.TestCase, +# only ci_register imported (no heavy sglang.test.test_utils chain). +register_cpu_ci(est_time=30, suite="base-a-test-cpu") + +_DEV = "cpu" + + +class _FakeKVCache: + """buf[p] == virtual id stored at physical slot p (-1 free); moves copy it.""" + + def __init__(self, max_slots: int): + self.buf = torch.full((max_slots,), -1, dtype=torch.int64) + + def move_kv_cache(self, dst_loc: torch.Tensor, src_loc: torch.Tensor): + self.buf[dst_loc] = self.buf[src_loc].clone() + + +class _FakeUnifiedSWAKVPool: + class _SubKV(_FakeKVCache): + def __init__(self, max_slots): + super().__init__(max_slots) + self.allocator = None + + def attach_allocator(self, allocator): + self.allocator = allocator + + def __init__(self, shared_pool: UnifiedKVPool): + self.full_kv_pool = self._SubKV(shared_pool.max_slots("full")) + self.swa_kv_pool = self._SubKV(shared_pool.max_slots("swa")) + self._full_allocator = None + self._swa_allocator = None + + def attach_allocators(self, *, full_allocator, swa_allocator): + self._full_allocator = full_allocator + self._swa_allocator = swa_allocator + + +def _tri_specs( + full_layer_num=4, swa_layer_num=2, state_layer_num=2, head_num=2, head_dim=4 +): + full = MHASubPoolSpec( + name="full", + layer_num=full_layer_num, + head_num=head_num, + head_dim=head_dim, + store_dtype=torch.float16, + grow_direction="down", + ) + swa = MHASubPoolSpec( + name="swa", + layer_num=swa_layer_num, + head_num=head_num, + head_dim=head_dim, + store_dtype=torch.float16, + grow_direction="float", + ) + mamba = MambaSubPoolSpec( + name="mamba", + layer_num=state_layer_num, + conv_state_shapes=((3, 8),), + conv_dtype=torch.bfloat16, + temporal_state_shape=(0, 0, 0), # Inkling: conv-only, no SSM state + temporal_dtype=torch.float32, + grow_direction="up", + ) + return full, swa, mamba + + +class TestUnifiedTriPool(unittest.TestCase): + def _build( + self, + n_full=32, + n_swa=16, + n_state=8, + lazy_compaction=False, + ): + full, swa, mamba = _tri_specs() + total = ( + n_full * full.entry_bytes() + + n_swa * swa.entry_bytes() + + n_state * mamba.entry_bytes() + ) + pool = UnifiedKVPool( + total_bytes=total, + sub_pool_specs=[full, swa, mamba], + device=_DEV, + enable_memory_saver=False, + ) + kvcache = _FakeUnifiedSWAKVPool(pool) + mamba_kv = _FakeKVCache(pool.max_slots("mamba")) + allocator = UnifiedMambaSWATokenToKVPoolAllocator( + unified_buffer=pool, + kvcache=kvcache, + mamba_kvcache=mamba_kv, + device=_DEV, + full_max_total_num_tokens=n_full, + swa_max_total_num_tokens=n_swa, + need_sort=False, + forward_stream=None, + lazy_compaction=lazy_compaction, + ) + return pool, allocator, kvcache, mamba_kv + + def _stamp(self, allocator, kvcache, v): + fa = allocator.full_attn_allocator + sa = allocator.swa_attn_allocator + kvcache.full_kv_pool.buf[fa.virtual_to_physical[v]] = v + kvcache.swa_kv_pool.buf[sa.virtual_to_physical[v]] = v + + # -- construction -- + + def test_chain_wiring_and_float_swa(self): + pool, allocator, kvcache, _ = self._build() + fa = allocator.full_attn_allocator + sa = allocator.swa_attn_allocator + ma = allocator.mamba_allocator + self.assertIsInstance(sa, FloatMultiEndedAllocator) + self.assertIs(ma.high_peer, sa) + self.assertIs(sa.low_peer, ma) + self.assertIs(sa.high_peer, fa) + self.assertIs(fa.low_peer, sa) + # Canonical chain order in the pool. + self.assertEqual( + [s.name for s in pool.sub_pool_specs], ["mamba", "swa", "full"] + ) + # KV pool got the allocators. + self.assertIs(kvcache._full_allocator, fa) + self.assertIs(kvcache._swa_allocator, sa) + + def test_empty_float_is_transparent_to_the_ends(self): + _, allocator, _, _ = self._build() + fa = allocator.full_attn_allocator + ma = allocator.mamba_allocator + self.assertTrue(allocator.swa_attn_allocator._is_frontier_transparent()) + # full's chain gap reaches the mamba end's frontier straight through. + self.assertEqual( + fa._current_gap_bytes(), + fa._byte_low_frontier() - ma._byte_high_frontier(), + ) + + # -- the joint availability contract -- + + def test_available_size_alloc_contract(self): + for lazy in (False, True): + _, allocator, kvcache, _ = self._build(lazy_compaction=lazy) + avail = allocator.available_size() + self.assertGreater(avail, 0) + v = allocator.alloc(avail) + self.assertIsNotNone( + v, f"alloc(available_size()={avail}) must succeed (lazy={lazy})" + ) + self.assertEqual(int(v.numel()), avail) + # Both sides bound for every allocated virtual id. + fa = allocator.full_attn_allocator + sa = allocator.swa_attn_allocator + self.assertTrue(bool((fa.virtual_to_physical[v] >= 0).all())) + self.assertTrue(bool((sa.virtual_to_physical[v] >= 0).all())) + + def test_available_shrinks_as_state_slots_grow(self): + _, allocator, _, _ = self._build() + before = allocator.available_size() + slots = allocator.mamba_allocator.alloc(4) + self.assertIsNotNone(slots) + after = allocator.available_size() + self.assertLess(after, before) + allocator.mamba_allocator.free(slots) + self.assertEqual(allocator.available_size(), before) + + # -- steady-state SWA churn: tombstones -> holes -> in-place reuse -- + + def _swa_interior_block(self, allocator, blocks): + """The block whose SWA-physical pages touch neither float boundary -- + `free_swa` on it must create interior holes (a boundary block would be + absorbed instead; both are zero-copy, different mechanisms).""" + sa = allocator.swa_attn_allocator + for v in blocks: + pages = set(int(x) for x in sa.virtual_to_physical[v].tolist()) + if sa.low_wm_page not in pages and (sa.high_wm_page - 1) not in pages: + return v + raise AssertionError("no interior block in layout") + + def test_free_swa_holes_recycled_in_place_zero_copy(self): + _, allocator, kvcache, _ = self._build() + blocks = [allocator.alloc(4) for _ in range(3)] + for v in blocks: + self.assertIsNotNone(v) + self._stamp(allocator, kvcache, v) + sa = allocator.swa_attn_allocator + fa = allocator.full_attn_allocator + span_before = (sa.low_wm_page, sa.high_wm_page) + + # Window slide: an INTERIOR block ages out of the SWA window. + v_mid = self._swa_interior_block(allocator, blocks) + allocator.free_swa(v_mid) + # The full side keeps the token; the swa side tombstoned it. + self.assertTrue(bool((fa.virtual_to_physical[v_mid] >= 0).all())) + self.assertTrue(bool((sa.virtual_to_physical[v_mid] == -1).all())) + self.assertEqual(sa._hole_pages(), 4) + self.assertEqual((sa.low_wm_page, sa.high_wm_page), span_before) + + # The next alloc recycles the holes IN PLACE: no span growth, no moves. + vd = allocator.alloc(4) + self.assertIsNotNone(vd) + self._stamp(allocator, kvcache, vd) + self.assertEqual(sa._hole_pages(), 0) + self.assertEqual((sa.low_wm_page, sa.high_wm_page), span_before) + self.assertEqual(len(sa._inverse_history), 0) # zero copies + + def test_free_swa_boundary_block_absorbed_zero_copy(self): + # The OTHER zero-copy mechanism: a boundary block's tombstones shrink + # the span, handing bytes back to the neighbours. The shrink is + # DEFERRED out of the per-step free (it needs the hole set on the + # host); the per-step opportunistic flush is where it lands. + _, allocator, kvcache, _ = self._build() + blocks = [allocator.alloc(4) for _ in range(2)] + for v in blocks: + self._stamp(allocator, kvcache, v) + sa = allocator.swa_attn_allocator + span_pages = sa._span_pages() + # Pick a block holding a span-boundary page. + boundary = None + for v in blocks: + pages = set(int(x) for x in sa.virtual_to_physical[v].tolist()) + if sa.low_wm_page in pages or (sa.high_wm_page - 1) in pages: + boundary = v + break + self.assertIsNotNone(boundary) + allocator.free_swa(boundary) + allocator.flush_opportunistic() # the deferred reclaim point + self.assertEqual(sa._hole_pages(), 0) # absorbed, not holed + self.assertEqual(sa._span_pages(), span_pages - 4) + self.assertEqual(len(sa._inverse_history), 0) # zero copies + + def test_free_releases_both_sides_and_filters_tombstones(self): + _, allocator, kvcache, _ = self._build() + va = allocator.alloc(4) + self._stamp(allocator, kvcache, va) + allocator.free_swa(va) # tombstone first (aged out of window) + allocator.free(va) # then the request finishes + fa = allocator.full_attn_allocator + sa = allocator.swa_attn_allocator + self.assertTrue(bool((fa.virtual_to_physical[va] == -1).all())) + self.assertTrue(bool((sa.virtual_to_physical[va] == -1).all())) + # Fully-freed float parks and is transparent again. + self.assertTrue(sa._is_frontier_transparent()) + + # -- per-request state surface -- + + def test_mamba_slot_allocator_surface(self): + pool, allocator, _, mamba_kv = self._build() + slot_alloc = UnifiedMambaSlotAllocator( + allocator.mamba_allocator, + max_size=pool.max_slots("mamba") - 1, + device=_DEV, + ) + v = slot_alloc.alloc(3) + self.assertIsNotNone(v) + p = slot_alloc.translate(v) + self.assertTrue(bool((p >= 0).all())) + mamba_kv.buf[p] = v + self.assertEqual( + slot_alloc.available_size(), + (pool.max_slots("mamba") - 1) - 3, + ) + slot_alloc.free(v) + self.assertEqual(slot_alloc.available_size(), pool.max_slots("mamba") - 1) + # Group prefetch draw-down + surplus return. + slot_alloc.alloc_group_begin(4) + s1 = slot_alloc.alloc(1) + self.assertIsNotNone(s1) + slot_alloc.alloc_group_end() + self.assertEqual( + slot_alloc.available_size(), + (pool.max_slots("mamba") - 1) - 1, + ) + + # -- cost + flush semantics -- + + def test_mamba_slot_full_token_cost_formula(self): + _, allocator, _, _ = self._build() + e_tok = ( + allocator.full_attn_allocator.entry_bytes + + allocator.swa_attn_allocator.entry_bytes + ) + m = allocator.mamba_allocator.entry_bytes_per_page + self.assertEqual(allocator.mamba_slot_full_token_cost(), -(-m // e_tok)) + + def test_urgent_flush_preserves_float_holes(self): + _, allocator, kvcache, _ = self._build(lazy_compaction=True) + blocks = [allocator.alloc(4) for _ in range(3)] + for v in blocks: + self._stamp(allocator, kvcache, v) + allocator.free_swa(self._swa_interior_block(allocator, blocks)) + sa = allocator.swa_attn_allocator + holes = sa._hole_pages() + self.assertGreater(holes, 0) + from sglang.srt.mem_cache.multi_ended_allocator import _relieve_for_alloc + + _relieve_for_alloc(allocator, 1) + self.assertEqual(sa._hole_pages(), holes) # holes are assets, not backlog + + +class TestTriPagedFreeGroup(unittest.TestCase): + """The tri composite at PAGE SIZE > 1, driven through the production free + path: free_group_begin -> free_segment -> free_group_end. + + Regression (GPU eval_440, Inkling ps=128 boot crash): every other tri test + runs at page_size=1, where the page-REPRESENTATIVE machinery + (`free_page_reps_group` / `_release_page_reps`, added when the free path + was made host-sync-free) is entirely dead code — `free_segment` frees + directly. At ps>1 the composite releases reps by calling + `swa_attn_allocator.free(..., _pages=...)`, and the float allocator was + ported from a base that predates that keyword, so the first real decode + batch died with `TypeError: unexpected keyword argument '_pages'`. + + `_pages` is not cosmetic: honouring it is what keeps the free path free of + the data-dependent `torch.unique` host sync, so this also pins that the + float takes the caller's page ids rather than re-deriving them. + """ + + def _build_paged(self, page_size=4, n_full=64, n_swa=32, n_state=8): + full, swa, mamba = _tri_specs() + total = ( + n_full * full.entry_bytes() + + n_swa * swa.entry_bytes() + + n_state * mamba.entry_bytes() + ) + pool = UnifiedKVPool( + total_bytes=total, + sub_pool_specs=[full, swa, mamba], + device=_DEV, + enable_memory_saver=False, + page_size=page_size, + ) + kvcache = _FakeUnifiedSWAKVPool(pool) + mamba_kv = _FakeKVCache(pool.max_slots("mamba")) + allocator = UnifiedMambaSWATokenToKVPoolAllocator( + unified_buffer=pool, + kvcache=kvcache, + mamba_kvcache=mamba_kv, + device=_DEV, + full_max_total_num_tokens=n_full, + swa_max_total_num_tokens=n_swa, + page_size=page_size, + need_sort=False, + forward_stream=None, + ) + return pool, allocator + + def test_free_group_segment_release_reaches_the_float(self): + """The exact production sequence the scheduler runs per decode batch.""" + pool, allocator = self._build_paged() + v = allocator.alloc(8) + self.assertIsNotNone(v) + before = allocator.available_size() + + allocator.free_group_begin() + allocator.free_segment(v, start_pos=0) + allocator.free_group_end() # -> _release_page_reps -> float.free(_pages=) + + self.assertEqual(allocator.verify_byte_accounting(), []) + self.assertGreaterEqual(allocator.available_size(), before) + # Capacity fully recovered: the float parked, both ends rewound. + self.assertTrue(allocator.swa_attn_allocator._is_frontier_transparent()) + + def test_float_free_honours_caller_supplied_pages(self): + """`_pages` must be USED, not merely accepted -- re-deriving it is the + host sync the paged free path exists to avoid.""" + pool, allocator = self._build_paged() + v = allocator.alloc(8) + self.assertIsNotNone(v) + sa = allocator.swa_attn_allocator + ps = allocator.page_size + pages = (v[::ps] // ps).clone() + live_before = sa._live_pages() + sa.free(v[::ps] * 0 + v[::ps], _pages=pages) + self.assertEqual(sa._live_pages(), live_before - pages.numel()) + + def test_ungrouped_segment_free_also_reaches_the_float(self): + """`free_segment` outside a free group releases reps immediately -- + the same float call, one frame shallower.""" + pool, allocator = self._build_paged() + v = allocator.alloc(8) + self.assertIsNotNone(v) + allocator.free_segment(v, start_pos=0) + self.assertEqual(allocator.verify_byte_accounting(), []) + + +class TestTriFreeSwaNoHostSync(unittest.TestCase): + """The tri's swa side is the FLOAT, and the float can never run the lazy + event pipeline — so unless the per-step frees carry caller-derived page + ids, the tri silently reintroduces the host syncs the sync-free free + path removed. Poison the ops to pin the property. + + (Fixtures at page_size > 1 on purpose: ps==1 short-circuits the whole + page machinery and hides exactly this class of bug.) + """ + + PS = 4 + + def _tri(self): + inst = TestTriPagedFreeGroup( + [m for m in dir(TestTriPagedFreeGroup) if m.startswith("test_")][0] + ) + return inst._build_paged(page_size=self.PS)[1] + + def test_ratchet_shape_free_swa_never_syncs_on_the_float(self): + alloc = self._tri() + v = alloc.alloc(8 * self.PS) + self.assertIsNotNone(v) + from unittest import mock + + with mock.patch.object( + torch, "unique", side_effect=AssertionError("unique = host sync") + ), mock.patch.object( + torch.Tensor, "item", side_effect=AssertionError("item = host sync") + ): + alloc.free_swa(v[: 4 * self.PS], start_pos=0) + self.assertEqual(alloc.verify_byte_accounting(), []) + + def test_float_free_has_no_stale_slot_item_sync(self): + """The float's free must not `.item()`-assert per free (the lazy-path + contract: callers must not double-free; the idle span == p2v-bound + + holes conservation catches violations without a per-free sync).""" + alloc = self._tri() + v = alloc.alloc(4 * self.PS) + sa = alloc.swa_attn_allocator + from unittest import mock + + with mock.patch.object( + torch.Tensor, "item", side_effect=AssertionError("item = host sync") + ): + sa.free(v[:: self.PS], _pages=v[:: self.PS] // self.PS) + + def test_fallback_free_swa_still_correct_for_radix_shapes(self): + """Radix eviction hands arbitrary node values (no start_pos): the + dedup fallback must keep working and end in the same state as the + stride path.""" + a1, a2 = self._tri(), self._tri() + v1, v2 = a1.alloc(6 * self.PS), a2.alloc(6 * self.PS) + self.assertTrue(torch.equal(v1, v2)) + a1.free_swa(v1[: 4 * self.PS], start_pos=0) + a2.free_swa(v2[: 4 * self.PS]) + self.assertTrue( + torch.equal( + a1.swa_attn_allocator.virtual_to_physical, + a2.swa_attn_allocator.virtual_to_physical, + ) + ) + self.assertEqual(a1.available_size(), a2.available_size()) + self.assertEqual(a1.verify_byte_accounting(), []) + self.assertEqual(a2.verify_byte_accounting(), []) + + +class TestGeneralizedRebalance(unittest.TestCase): + """The float must yield to WHICHEVER end is short, with the direction + computed from the layout — not only to the token path's hard-coded side. + + The mechanism (`make_room`) was always side-agnostic; these pin the + POLICY: any end pool's own-alloc shortfall reaches + `_ask_float_for_room`, which derives the side from the caller's + growth direction.""" + + PS = 4 + + def _tri(self): + inst = TestTriPagedFreeGroup( + [m for m in dir(TestTriPagedFreeGroup) if m.startswith("test_")][0] + ) + return inst._build_paged(page_size=self.PS)[1] + + def test_state_end_shortfall_slides_the_float_low(self): + """The previously-missing direction: mamba (grow-up END) starved while + free bytes idle ABOVE the float. The remedy must slide the float up + (open its LOW side) and let the state alloc succeed.""" + alloc = self._tri() + v = alloc.alloc(4 * self.PS) # places the float mid-region + self.assertIsNotNone(v) + ma = alloc.mamba_allocator + sa = alloc.swa_attn_allocator + self.assertFalse(sa._is_frontier_transparent()) + # Fill the LOW band exactly: as many state slots as fit below the + # float's low frontier. + e_m = ma.entry_bytes_per_page + fit = (sa._byte_low_frontier() - ma._byte_high_frontier()) // e_m + self.assertGreater(fit, 0) + got = ma.alloc(int(fit) * ma.page_size) + self.assertIsNotNone(got) + low_before = sa.low_wm_page + # One more slot does NOT fit below the float -- only a rebalance helps. + more = ma.alloc(ma.page_size) + self.assertIsNotNone(more, "state alloc must succeed via float rebalance") + self.assertGreater(sa.low_wm_page, low_before) # float slid UP + self.assertEqual(alloc.verify_byte_accounting(), []) + + def test_direction_is_derived_from_growth_on_both_ends(self): + """Raw end+float+end chain, BOTH orientations in one fixture: the + up-growing end opens the float's LOW side; the down-growing end opens + its HIGH side. No layout assumption survives.""" + from test_multi_ended_allocator import TestFloatMultiEndedAllocator + + inst = TestFloatMultiEndedAllocator( + [m for m in dir(TestFloatMultiEndedAllocator) if m.startswith("test_")][0] + ) + _pool, up_end, fla, down_end, _kv = inst._build_tri() + v = fla.alloc(8) # opaque float mid-region + self.assertIsNotNone(v) + + # UP end: exhaust its band below the float, then ask for more. + e_up = up_end.entry_bytes_per_page + fit = int((fla._byte_low_frontier() - up_end._byte_high_frontier()) // e_up) + if fit > 0: + self.assertIsNotNone(up_end.alloc(fit * up_end.page_size)) + low_before = fla.low_wm_page + self.assertIsNotNone(up_end.alloc(up_end.page_size)) + self.assertGreater(fla.low_wm_page, low_before) # opened LOW side + + # DOWN end: exhaust its band above the float, then ask for more. + e_dn = down_end.entry_bytes_per_page + fit = int((down_end._byte_low_frontier() - fla._byte_high_frontier()) // e_dn) + if fit > 0: + self.assertIsNotNone(down_end.alloc(fit * down_end.page_size)) + high_before = fla.high_wm_page + self.assertIsNotNone(down_end.alloc(down_end.page_size)) + self.assertLess(fla.high_wm_page, high_before) # opened HIGH side + + def test_two_pool_chain_rebalance_is_a_noop(self): + """No float in the chain => the remedy must change nothing (the + 2-pool composites keep their exact pre-existing behavior).""" + from test_multi_ended_allocator import ( + TestPagedMultiEndedAllocator as _PagedFixture, + ) + + inst = _PagedFixture( + [m for m in dir(_PagedFixture) if m.startswith("test_")][0] + ) + _pool, full, swa, _fkv, _skv = inst._build() + v = full.alloc(full.page_size * 2) + self.assertIsNotNone(v) + wm = full.watermark_physical + full._ask_float_for_room(full.page_size * 1000) # absurd ask + self.assertEqual(full.watermark_physical, wm) # untouched + + def test_index_cap_guard_never_moves_data_uselessly(self): + """When the caller's own INDEX space binds, no amount of float + movement helps -- make_room must not be called (poisoned).""" + from unittest import mock + + alloc = self._tri() + alloc.alloc(4 * self.PS) + ma = alloc.mamba_allocator + sa = alloc.swa_attn_allocator + huge = (ma.num_pages + 10) * ma.page_size # beyond index space + with mock.patch.object( + sa, "make_room", side_effect=AssertionError("useless make_room") + ): + ma._ask_float_for_room(huge) + + +class TestComputedShortSide(unittest.TestCase): + """`_ask_float_for_room` must open the side that MEASURES short -- never + "the side facing full". These pin the per-side computation, including + the coupled-ends-on-both-sides shape a DSV4-style composite + (C128 | swa-float | C4) will need. + """ + + PS = 4 + + def _tri(self): + inst = TestTriPagedFreeGroup( + [m for m in dir(TestTriPagedFreeGroup) if m.startswith("test_")][0] + ) + return inst._build_paged(page_size=self.PS)[1] + + def _sides(self, alloc): + sa = alloc.swa_attn_allocator + low = max(0, sa._byte_low_frontier() - sa._chain_high_frontier_below_bytes()) + high = max(0, sa._chain_low_frontier_above_bytes() - sa._byte_high_frontier()) + return low, high + + def test_float_share_short_opens_the_state_side(self): + """RED-LINE: full's demand fits its band, the float's own share fits + NEITHER band, and the state side has the larger surplus — the policy + must open the STATE side (the float slides toward full during a + TOKEN alloc), which the old "side facing full" policy could never do. + """ + from unittest import mock + + alloc = self._tri() + v = alloc.alloc(6 * self.PS) # places the float mid-region + self.assertIsNotNone(v) + sa = alloc.swa_attn_allocator + fa = alloc.full_attn_allocator + e_f, e_s = fa.entry_bytes_per_page, sa.entry_bytes_per_page + + # Position: slide the float LOW (setup uses the mechanism directly), + # so the low band is small and the geometry below is expressible. + b_low0, b_high0 = self._sides(alloc) + # Two positioning moves: pack the float low (leapfrog over-opens by + # design), then open the LOW side back to ~2 full-pages -- small + # enough that F outgrows it, wide enough that the integer need_n + # window below is non-empty. + sa.make_room(side="high", min_bytes=b_low0 + b_high0 - 2 * e_f) + sa.make_room(side="low", min_bytes=2 * e_f) + + # Find a need_n where: D_high = need_n*e_f fits band_high, F = + # need_n*e_s exceeds BOTH surpluses, and low has the larger surplus. + chosen = None + for need_n in range(1, 64): + b_low, b_high = self._sides(alloc) + s_low = b_low # no coupled end on the low side + s_high = b_high - need_n * e_f + if s_high < 0: + break + F = need_n * e_s + if F > s_low and F > s_high and s_low >= s_high: + chosen = need_n + break + self.assertIsNotNone(chosen, "fixture cannot express the geometry") + + calls = [] + real = sa.make_room + with mock.patch.object( + sa, "make_room", side_effect=lambda **kw: calls.append(kw) or real(**kw) + ): + alloc._ask_float_for_room(chosen * self.PS) + self.assertEqual(len(calls), 1, calls) + self.assertEqual(calls[0]["side"], "low") # the STATE side + + def test_full_side_short_target_matches_the_closed_form(self): + """Equivalence: when the full side is the short one (today's only + reachable end-shortage), the ask must equal the documented formula + demand + max(0, F - far_surplus) + slack — i.e. the historical + behavior is the special case, preserved.""" + from unittest import mock + + alloc = self._tri() + v = alloc.alloc(4 * self.PS) + self.assertIsNotNone(v) + sa, fa = alloc.swa_attn_allocator, alloc.full_attn_allocator + e_f, e_s = fa.entry_bytes_per_page, sa.entry_bytes_per_page + chosen = None + for need_n in range(1, 256): + b_low, b_high = self._sides(alloc) + if b_high - need_n * e_f < 0 and b_low >= 0: + chosen = need_n + break + self.assertIsNotNone(chosen) + b_low, b_high = self._sides(alloc) + F = max(0, chosen - sa._hole_pages()) * e_s + want = chosen * e_f + max(0, F - b_low) + max(e_f, e_s) + calls = [] + with mock.patch.object( + sa, "make_room", side_effect=lambda **kw: calls.append(kw) + ): + alloc._ask_float_for_room(chosen * self.PS) + self.assertEqual(len(calls), 1) + self.assertEqual(calls[0]["side"], "high") + self.assertEqual(calls[0]["min_bytes"], want) + + def test_two_coupled_ends_lands_demand_on_both_sides(self): + """DSV4 shape (C128 | float | C4): a coupled set with ends on BOTH + sides. One-side-short must open that side; BOTH-sides-short must not + move at all (relocation is zero-sum between the bands).""" + from unittest import mock + + alloc = self._tri() + v = alloc.alloc(6 * self.PS) + self.assertIsNotNone(v) + sa, fa, ma = ( + alloc.swa_attn_allocator, + alloc.full_attn_allocator, + alloc.mamba_allocator, + ) + # Synthetic coupling: the state end joins the demand vector, exactly + # the override a DSV4-style composite would ship. + need = lambda self, t: { + fa: -(-t // self.page_size), + sa: -(-t // self.page_size), + ma: -(-t // self.page_size), + } + with mock.patch.object(type(alloc), "_alloc_demand", need): + # (a) both sides short: absurd need -> both demands exceed their + # bands -> make_room must NOT be called. + with mock.patch.object( + sa, "make_room", side_effect=AssertionError("zero-sum move") + ): + alloc._ask_float_for_room(10_000 * self.PS) + + # (b) one side short: find a need where the HIGH side (full) is + # short while the LOW side (mamba demand) still fits. + e_f, e_m = fa.entry_bytes_per_page, ma.entry_bytes_per_page + chosen = None + for need_n in range(1, 256): + b_low, b_high = self._sides(alloc) + if b_high - need_n * e_f < 0 and b_low - need_n * e_m >= 0: + chosen = need_n + break + if chosen is not None: + calls = [] + with mock.patch.object( + sa, "make_room", side_effect=lambda **kw: calls.append(kw) + ): + alloc._ask_float_for_room(chosen * self.PS) + self.assertEqual(len(calls), 1) + self.assertEqual(calls[0]["side"], "high") + + def test_nothing_short_means_no_relocation(self): + """Everything fits -> the policy must not move a single page.""" + from unittest import mock + + alloc = self._tri() + alloc.alloc(4 * self.PS) + sa = alloc.swa_attn_allocator + with mock.patch.object( + sa, "make_room", side_effect=AssertionError("needless move") + ): + alloc._ask_float_for_room(1) + + +class TestFloatPolicyTotalTarget(unittest.TestCase): + """`make_room`'s min_bytes is a TARGET for the whole band, not a delta. + + Regression: the band-level policy passed `deficit + one page` — with a + PARTIALLY free band that is below the current gap, so `make_room` + no-oped and the allocation failed even though the float had room to + slide. (Its own test missed this because it filled the band exactly, + making deficit ≈ the whole need.) The demand-vector policy computes the + total target, so a partial gap under-asks never. + """ + + PS = 4 + + def _tri(self): + inst = TestTriPagedFreeGroup( + [m for m in dir(TestTriPagedFreeGroup) if m.startswith("test_")][0] + ) + return inst._build_paged(page_size=self.PS)[1] + + def test_partial_gap_state_alloc_still_succeeds(self): + alloc = self._tri() + v = alloc.alloc(6 * self.PS) + self.assertIsNotNone(v) + ma, sa = alloc.mamba_allocator, alloc.swa_attn_allocator + e_m = ma.entry_bytes_per_page + gap_slots = int((sa._byte_low_frontier() - ma._byte_high_frontier()) // e_m) + self.assertGreater(gap_slots, 2) + low_before = sa.low_wm_page + # Need = partial-gap + 3: the old delta-ask was BELOW the current + # gap, so nothing moved and this returned None. + got = ma.alloc((gap_slots + 3) * ma.page_size) + self.assertIsNotNone(got, "partial-gap shortfall must relocate, not fail") + self.assertGreater(sa.low_wm_page, low_before) + self.assertEqual(alloc.verify_byte_accounting(), []) + + def test_zero_demand_bands_are_inert(self): + """The tri's token vector carries {mamba: 0}: a zero entry must + neither move the float for mamba's sake nor trip the index guard.""" + from unittest import mock + + alloc = self._tri() + alloc.alloc(4 * self.PS) + demand = alloc._alloc_demand(2 * self.PS) + self.assertEqual(demand[alloc.mamba_allocator], 0) + sa = alloc.swa_attn_allocator + with mock.patch.object( + sa, "make_room", side_effect=AssertionError("needless move") + ): + alloc._ask_float_for_room(1) # nothing short -> no relocation + + +class TestTriDeferredAbsorption(unittest.TestCase): + """Boundary absorption is deferred out of the per-step free and paid once + at a quiescent point — the base allocator's model (its lazy free does "no + boundary absorb" and `_flush` pays a single D2H). These pin WHERE it is + now paid, and that skipping it stays merely conservative.""" + + PS = 4 + + def _tri(self): + inst = TestTriPagedFreeGroup( + [m for m in dir(TestTriPagedFreeGroup) if m.startswith("test_")][0] + ) + return inst._build_paged(page_size=self.PS)[1] + + def test_per_step_flush_reclaims_the_span(self): + alloc = self._tri() + v = alloc.alloc(8 * self.PS) + sa = alloc.swa_attn_allocator + span = sa._span_pages() + alloc.free_swa(v[6 * self.PS :], start_pos=6 * self.PS) # high edge + self.assertGreater(sa._hole_pages(), 0) # deferred + self.assertEqual(sa._span_pages(), span) + moved = alloc.flush_opportunistic() + self.assertGreater(moved, 0) + self.assertLess(sa._span_pages(), span) + self.assertEqual(alloc.verify_byte_accounting(), []) + + def test_shortfall_ladder_absorbs_before_the_deficit_math(self): + """The zero-copy rung must run FIRST: a stale-wide span would inflate + the rebalance deficit and buy a `make_room` relocation the shrink + already covers.""" + alloc = self._tri() + v = alloc.alloc(8 * self.PS) + sa = alloc.swa_attn_allocator + alloc.free_swa(v[6 * self.PS :], start_pos=6 * self.PS) + self.assertGreater(sa._hole_pages(), 0) + moves_before = len(sa._inverse_history) + from sglang.srt.mem_cache.multi_ended_allocator import _relieve_for_alloc + + _relieve_for_alloc(alloc, 1) # the ladder + self.assertEqual(sa._hole_pages(), 0) # rung 0 ran + self.assertEqual(len(sa._inverse_history), moves_before) # zero copies + + def test_deferral_is_conservative_never_over_reports(self): + """Availability with a stale-wide span must never EXCEED the absorbed + value -- under-reporting is safe, over-reporting would over-admit.""" + alloc = self._tri() + v = alloc.alloc(8 * self.PS) + alloc.free_swa(v[6 * self.PS :], start_pos=6 * self.PS) + deferred = alloc.available_size() + alloc.swa_attn_allocator._flush(urgent=False) + absorbed = alloc.available_size() + self.assertLessEqual(deferred, absorbed) + self.assertEqual(alloc.verify_byte_accounting(), []) + + def test_clean_flush_skips_the_d2h_entirely(self): + """Only `free` can put a hole ON a boundary (alloc DRAINS holes into + live pages; extension adds live pages), so with nothing freed since + the last absorb the walk provably finds nothing — and must not pay + the D2H. Steady churn with only interior holes then costs no sync.""" + from unittest import mock + + alloc = self._tri() + v = alloc.alloc(8 * self.PS) + alloc.free_swa(v[2 * self.PS : 4 * self.PS], start_pos=2 * self.PS) + alloc.flush_opportunistic() # consumes the dirty flag + sa = alloc.swa_attn_allocator + self.assertGreater(sa._hole_pages(), 0) # interior holes remain + with mock.patch.object( + torch.Tensor, "tolist", side_effect=AssertionError("tolist = D2H") + ): + self.assertEqual(alloc.flush_opportunistic(), 0) + self.assertEqual(sa._flush(urgent=False), 0) + + def test_alloc_between_frees_cannot_hide_a_boundary_hole(self): + """Soundness of the skip: an alloc drains holes and can change the + hole COUNT back to a previously-seen value, so the flag must be armed + by `free`, not inferred from `numel()`.""" + alloc = self._tri() + v = alloc.alloc(8 * self.PS) + sa = alloc.swa_attn_allocator + alloc.free_swa(v[: 2 * self.PS], start_pos=0) # low-edge holes + n_after_free = sa._hole_pages() + alloc.alloc(2 * self.PS) # drains them back to live + alloc.free_swa(v[6 * self.PS :], start_pos=6 * self.PS) # high edge + self.assertEqual(sa._hole_pages(), n_after_free) # same COUNT as before + span = sa._span_pages() + self.assertGreater(alloc.flush_opportunistic(), 0) # still absorbed + self.assertLess(sa._span_pages(), span) + self.assertEqual(alloc.verify_byte_accounting(), []) + + def test_transparency_still_exact_without_absorption(self): + """Park-on-empty stays in `free` because it is sync-free -- a float + that empties must go transparent immediately, with no flush needed.""" + alloc = self._tri() + v = alloc.alloc(4 * self.PS) + sa = alloc.swa_attn_allocator + self.assertFalse(sa._is_frontier_transparent()) + from unittest import mock + + with mock.patch.object( + torch.Tensor, "tolist", side_effect=AssertionError("tolist = D2H") + ): + alloc.free_swa(v, start_pos=0) + self.assertTrue(sa._is_frontier_transparent()) + self.assertEqual(sa._hole_pages(), 0) + + +class TestTriFactorySizing(unittest.TestCase): + """Factory-level contracts: byte-budget sizing, the bs=1 feasibility + floor, and the boot signature the GPU harness greps for (3 sub-pools, + swa grow=float).""" + + def _factory_kwargs(self, **over): + import types + + cp = types.SimpleNamespace( + shape=types.SimpleNamespace(conv=[(3, 8)], temporal=(0, 0, 0)), + dtype=types.SimpleNamespace(conv=torch.bfloat16, temporal=torch.float32), + layers=[0, 1], + ) + kw = dict( + device=_DEV, + kv_cache_dtype=torch.float16, + head_num=2, + head_dim=4, + v_head_dim=4, + swa_head_num=2, + swa_head_dim=4, + swa_v_head_dim=4, + page_size=1, + start_layer=0, + end_layer=2, + swa_attention_layer_ids=[1], + full_attention_layer_ids=[0], + mamba_layer_ids=[0, 1], + mamba2_cache_params=cp, + full_max_total_num_tokens=64, + swa_max_total_num_tokens=32, + max_mamba_cache_size=4, + model_context_len=16, + extra_max_context_len=4, + max_num_reqs=4, + enable_memory_saver=False, + enable_mamba_extra_buffer=False, + disable_overlap_schedule=True, + need_sort=False, + ) + kw.update(over) + return kw + + def test_budget_sizing_and_boot_signature(self): + budget = 1 << 20 + bundle = init_unified_mamba_swa_pools( + **self._factory_kwargs(unified_total_bytes=budget) + ) + pool = bundle.unified_memory_pool + # Buffer = budget + the state pool's bytes (budget captured AFTER the + # state carve-out), never the token-count re-sum. + state_bytes = 4 * pool.spec("mamba").entry_bytes() + self.assertEqual(pool.total_bytes, budget + state_bytes) + # Boot signature: 3 sub-pools in chain order, swa is the float. + self.assertEqual(len(pool.sub_pool_specs), 3) + self.assertEqual( + [(sp.name, sp.grow_direction) for sp in pool.sub_pool_specs], + [("mamba", "up"), ("swa", "float"), ("full", "down")], + ) + + def test_fallback_is_the_token_count_resum(self): + bundle = init_unified_mamba_swa_pools(**self._factory_kwargs()) + pool = bundle.unified_memory_pool + want = ( + 64 * pool.spec("full").entry_bytes() + + 32 * pool.spec("swa").entry_bytes() + + 4 * pool.spec("mamba").entry_bytes() + ) + self.assertEqual(pool.total_bytes, want) + + def test_bs1_floor_fails_loud_before_construction(self): + """A budget far below one worst-case request must raise BEFORE any + pool construction -- under-sizing is a retract LIVELOCK at runtime.""" + with self.assertRaisesRegex(RuntimeError, "bs=1 floor"): + init_unified_mamba_swa_pools( + **self._factory_kwargs( + unified_total_bytes=1024, # << ctx * e_f alone + model_context_len=100_000, + sliding_window_size=64, + ) + ) + + +class TestTriPoolHardening(unittest.TestCase): + """C1.7 pressure lanes: the planned-rebalance remedy in the alloc path + (a mis-positioned float must not fail an alloc that fits in total bytes), + retract-loop convergence through check_decode_capacity, and bounded copy + traffic under alternating end pressure. + """ + + def _build(self, **kw): + return TestUnifiedTriPool._build(self, **kw) + + def test_alloc_rebalances_a_blocking_float(self): + # Fill much of the high band so the float (midpoint-placed) walls off + # the low band's free bytes from `full`; the next alloc must succeed + # by SLIDING the float, not fail while total bytes suffice. + _, allocator, kvcache, _ = self._build(n_full=32, n_swa=24, n_state=8) + sa = allocator.swa_attn_allocator + v0 = allocator.alloc(4) # places the float at the region midpoint + self.assertIsNotNone(v0) + TestUnifiedTriPool._stamp(self, allocator, kvcache, v0) + # Exhaust the high band directly on the full end (full-only growth, + # e.g. long decode of already-admitted requests). + fa = allocator.full_attn_allocator + b_high_pages = fa._current_gap_bytes() // fa.entry_bytes_per_page + grab = fa.alloc(max(0, (b_high_pages - 2))) + self.assertIsNotNone(grab) + # The honest gate under-reports (no slide credit) -- asking BEYOND it + # is what fires the rebalance remedy; the ask still fits total free + # bytes because the LOW band holds them behind the float. + avail = allocator.available_size() + need = avail + 4 + live_before = sa._live_pages() + moves_before = len(sa._inverse_history) + v1 = allocator.alloc(need) + self.assertIsNotNone( + v1, "alloc must rebalance the blocking float instead of failing" + ) + self.assertEqual(int(v1.numel()), need) + moved = sum(int(s.numel()) for s, _, _ in sa._inverse_history[moves_before:]) + self.assertGreater(moved, 0, "the rebalance path must have fired") + # Cost bound min(L_live, G): never more than the live pages present + # when the slide ran (the leapfrog cap). + self.assertLessEqual(moved, live_before) + self.assertEqual(allocator.verify_byte_accounting(), []) + + def test_check_decode_capacity_retract_convergence(self): + # Simulated retract loop: requests' token blocks freed one at a time + # until the next-step allocation fits; must converge before bs=1 and + # never report capacity while the gate is short. + _, allocator, _, _ = self._build(n_full=32, n_swa=24, n_state=8) + reqs = [] + while True: + v = allocator.alloc(4) + if v is None or allocator.available_size() < 4: + if v is not None: + reqs.append(v) + break + reqs.append(v) + self.assertGreater(len(reqs), 2) + # Pool saturated: a large decode step does not fit. + need = 16 + while not allocator.check_decode_capacity(num_tokens=need, tree_cache=None): + self.assertGreater(len(reqs), 1, "retract must converge before bs=1") + allocator.free(reqs.pop()) + self.assertGreaterEqual(allocator.available_size(), need) + self.assertEqual(allocator.verify_byte_accounting(), []) + + def test_alternating_pressure_copy_traffic_bounded(self): + # Alternating full-grow / swa-churn cycles: total float moves stay + # bounded (hole recycling + absorption do the steady-state work; the + # rebalance fires only on real positional deficits). + _, allocator, kvcache, _ = self._build(n_full=48, n_swa=32, n_state=8) + sa = allocator.swa_attn_allocator + fa = allocator.full_attn_allocator + total_alloc_pages = 0 + for _ in range(6): + v = allocator.alloc(8) + self.assertIsNotNone(v) + total_alloc_pages += 8 + TestUnifiedTriPool._stamp(self, allocator, kvcache, v) + allocator.free_swa(v) # window slide: tombstones -> holes/absorb + g = fa.alloc(4) # full-side decode growth + self.assertIsNotNone(g) + total_alloc_pages += 4 + fa.free(g) + moved = sum(int(s.numel()) for s, _, _ in sa._inverse_history) + self.assertLessEqual( + moved, + total_alloc_pages // 2, + "steady-state churn must be predominantly zero-copy", + ) + self.assertEqual(allocator.verify_byte_accounting(), []) + + def test_joint_eviction_loop_stops_without_progress(self): + # tree_cache=None: the default helper no-ops; the bounded loop must + # return promptly (no infinite re-check) and the gate reports honestly. + _, allocator, _, _ = self._build() + big = allocator.available_size() + 64 + self.assertFalse( + allocator.check_decode_capacity(num_tokens=big, tree_cache=None) + ) + + +class TestJointCapacityIsHonoured(unittest.TestCase): + """`alloc(available_size())` must never fail. + + REGRESSION: the joint predicate priced the swa float's extension in RAW + BYTES while `take_physical_pages` can only use whole pages on the float's + OWN grid -- `_region_bounds_pages` rounds the band's low edge UP. The + bounding frontier is a multiple of the NEIGHBOUR's entry size, which is + unrelated to the float's, so the byte budget credited a page the grid could + not yield and the very first alloc tripped `alloc_with_virtual`'s backstop + assert. Swept over geometries rather than pinned to one, so a symmetric + mistake on the FULL side would surface here too. + """ + + def _build(self, *, page_size, n_full, n_swa, n_state, lazy, specs): + full, swa, mamba = specs + total = ( + n_full * full.entry_bytes() + + n_swa * swa.entry_bytes() + + n_state * mamba.entry_bytes() + ) + pool = UnifiedKVPool( + total_bytes=total, + sub_pool_specs=[full, swa, mamba], + device=_DEV, + enable_memory_saver=False, + page_size=page_size, + ) + kvcache = _FakeUnifiedSWAKVPool(pool) + return pool, UnifiedMambaSWATokenToKVPoolAllocator( + unified_buffer=pool, + kvcache=kvcache, + mamba_kvcache=_FakeKVCache(pool.max_slots("mamba")), + device=_DEV, + full_max_total_num_tokens=n_full, + swa_max_total_num_tokens=n_swa, + need_sort=False, + forward_stream=None, + lazy_compaction=lazy, + ) + + def test_fresh_boot_alloc_of_available_size_succeeds(self): + # Geometries chosen so the mamba end's frontier (a multiple of the + # STATE entry size) lands off the swa float's page grid -- the + # misalignment the byte budget used to ignore. + for page_size in (1, 2, 4): + for fl, sl, ml in ((4, 3, 1), (4, 2, 2), (6, 3, 1), (3, 5, 2)): + for n_full, n_swa, n_state in ((24, 16, 4), (32, 16, 8), (20, 12, 6)): + for lazy in (False, True): + specs = _tri_specs( + full_layer_num=fl, + swa_layer_num=sl, + state_layer_num=ml, + head_num=1, + head_dim=8, + ) + with self.subTest( + ps=page_size, + layers=(fl, sl, ml), + n=(n_full, n_swa, n_state), + lazy=lazy, + ): + _pool, alloc = self._build( + page_size=page_size, + n_full=n_full * page_size, + n_swa=n_swa * page_size, + n_state=n_state, + lazy=lazy, + specs=specs, + ) + n = alloc.available_size() + if n <= 0: + continue + # The whole point: the number the scheduler reads + # must be allocatable, with no backstop assert. + out = alloc.alloc(n) + self.assertIsNotNone( + out, + f"alloc(available_size()={n}) returned None", + ) + self.assertEqual(out.numel(), n) + + def test_available_size_never_exceeds_the_float_page_grid(self): + """Direct form: the joint answer, converted to float pages, must fit + inside what `_region_bounds_pages` actually offers.""" + for page_size in (1, 4): + specs = _tri_specs( + full_layer_num=4, + swa_layer_num=3, + state_layer_num=1, + head_num=1, + head_dim=8, + ) + _pool, alloc = self._build( + page_size=page_size, + n_full=24 * page_size, + n_swa=16 * page_size, + n_state=4, + lazy=False, + specs=specs, + ) + sa = alloc.swa_attn_allocator + n_pages = alloc.available_size() // page_size + lo, hi = sa._region_bounds_pages() + with self.subTest(ps=page_size): + self.assertLessEqual( + n_pages - sa._hole_pages(), + max(0, hi - lo), + "joint available_size() promises more float pages than the " + "float's own page grid can yield", + ) + + +class TestFloatRelocationIsOrderedAgainstTheForward(unittest.TestCase): + """Float relocation must settle the in-flight forward BEFORE its first copy. + + REGRESSION: `make_room` / `compact_holes` issued `move_kv_cache` and rebound + `virtual_to_physical` with no ordering against the running forward, so the + copy could carry pre-write bytes and the rebind then pointed every later + reader at a destination that never received those writes -- silently wrong + KV, no crash. The END pools guard exactly this hazard in + `_flush(urgent=True)` via `_settle_inflight_forward`; the float had no + `forward_stream` / `wait_event` / settle call anywhere in its body. + """ + + def _tri(self, lazy=True): + full, swa, mamba = _tri_specs(head_num=1, head_dim=8) + total = ( + 48 * full.entry_bytes() + 32 * swa.entry_bytes() + 8 * mamba.entry_bytes() + ) + pool = UnifiedKVPool( + total_bytes=total, + sub_pool_specs=[full, swa, mamba], + device=_DEV, + enable_memory_saver=False, + ) + kvcache = _FakeUnifiedSWAKVPool(pool) + alloc = UnifiedMambaSWATokenToKVPoolAllocator( + unified_buffer=pool, + kvcache=kvcache, + mamba_kvcache=_FakeKVCache(pool.max_slots("mamba")), + device=_DEV, + full_max_total_num_tokens=48, + swa_max_total_num_tokens=32, + need_sort=False, + forward_stream=None, + lazy_compaction=lazy, + ) + return pool, alloc, kvcache + + def _trace(self, flt): + """Record the order of (settle, move) on the float.""" + order = [] + real_settle = flt._settle_inflight_forward + real_move = flt._move_pages_and_rebind + + def settle(): + order.append("settle") + return real_settle() + + def move(src, dst): + order.append("move") + return real_move(src, dst) + + flt._settle_inflight_forward = settle + flt._move_pages_and_rebind = move + return order + + def test_make_room_settles_before_the_first_move(self): + _pool, alloc, _kv = self._tri() + flt = alloc.swa_attn_allocator + # Occupy the float, then free an interior page so a relocation has + # something to move and somewhere to move it. + v = alloc.alloc(12) + self.assertIsNotNone(v) + alloc.free(v[:4]) + order = self._trace(flt) + flt.make_room(side="low", min_bytes=flt.entry_bytes_per_page) + self.assertIn("settle", order, "make_room never settled the forward") + if "move" in order: + self.assertLess( + order.index("settle"), + order.index("move"), + f"a copy was issued before the settle: {order}", + ) + + def test_compact_holes_settles_before_the_first_move(self): + _pool, alloc, _kv = self._tri() + flt = alloc.swa_attn_allocator + v = alloc.alloc(12) + self.assertIsNotNone(v) + alloc.free(v[2:6]) # interior holes, so compact_holes has work + order = self._trace(flt) + flt.compact_holes(retreat_side="high") + if not order: + self.skipTest("no holes reached compact_holes in this geometry") + self.assertEqual(order[0], "settle", f"first action was not a settle: {order}") + + def test_the_settle_is_a_stream_wait_not_a_host_sync(self): + """Pin the mechanism: `_settle_inflight_forward` must stream-wait, so + the fix costs no host sync on the shortfall path.""" + src = inspect.getsource( + mea.MultiEndedAllocator._settle_inflight_forward # noqa: SLF001 + ) + self.assertIn("wait_event", src) + self.assertNotIn(".item()", src) + self.assertNotIn("synchronize()", src) + + +class TestFloatHoleCreditIsPerSide(unittest.TestCase): + """A float's schedulable credit must follow the side the holes are on. + + REGRESSION: the base `_peer_drainable_hole_bytes` asks + `_growth_side_neighbor()`, which reads `grow_direction`. A float's is + "float", so the base fell through to `low_peer` -- it never saw the HIGH + neighbour, and the single scalar it returned was then added to + `max(gap_low, gap_high)`, landing a LOW neighbour's holes on the HIGH gap. + Over-reporting `schedulable_available_size` makes the scheduler admit work + the shortfall ladder cannot satisfy, which the caller treats as a + memory-estimation bug. + """ + + def _float(self): + full, swa, mamba = _tri_specs(head_num=1, head_dim=8) + total = ( + 48 * full.entry_bytes() + 32 * swa.entry_bytes() + 8 * mamba.entry_bytes() + ) + pool = UnifiedKVPool( + total_bytes=total, + sub_pool_specs=[full, swa, mamba], + device=_DEV, + enable_memory_saver=False, + ) + alloc = UnifiedMambaSWATokenToKVPoolAllocator( + unified_buffer=pool, + kvcache=_FakeUnifiedSWAKVPool(pool), + mamba_kvcache=_FakeKVCache(pool.max_slots("mamba")), + device=_DEV, + full_max_total_num_tokens=48, + swa_max_total_num_tokens=32, + need_sort=False, + forward_stream=None, + lazy_compaction=True, + ) + return alloc, alloc.swa_attn_allocator + + def test_credit_sees_both_neighbours(self): + _alloc, flt = self._float() + self.assertIsInstance(flt, FloatMultiEndedAllocator) + low = flt._side_drainable_hole_bytes("low") + high = flt._side_drainable_hole_bytes("high") + self.assertEqual(flt._peer_drainable_hole_bytes(), max(low, high)) + # The base would have answered with the LOW side alone. + self.assertGreaterEqual(flt._peer_drainable_hole_bytes(), high) + + def test_schedulable_never_exceeds_the_sum_of_the_two_sides(self): + """Upper bound that the undirected scalar could violate: no side may be + credited with the other side's holes on top of its own gap.""" + alloc, flt = self._float() + v = alloc.alloc(10) + self.assertIsNotNone(v) + alloc.free(v[:3]) + epp = flt.entry_bytes_per_page + gap_low, gap_high = flt._gap_pages() + c_low = flt._side_drainable_hole_bytes("low") // epp + c_high = flt._side_drainable_hole_bytes("high") // epp + bound = ( + min( + max(gap_low + c_low, gap_high + c_high), + flt.num_pages - flt.min_page_index - flt._live_pages(), + ) + + flt._hole_pages() + ) * flt.page_size + self.assertLessEqual(flt.schedulable_available_size(), bound) + + def test_memo_verifier_agrees_with_the_per_side_formula(self): + """The staleness verifier recomputes through the same entry, so the + override must not make the memo look stale.""" + _alloc, flt = self._float() + flt.available_size() + flt.schedulable_available_size() + self.assertEqual(flt._byte_accounting_violations(), []) + + +if __name__ == "__main__": + unittest.main()