feat(unified-memory): three sub-pools for mamba + hybrid-SWA models (#35177)
Co-authored-by: Caihua Li <caihua.li@bytedance.com> Co-authored-by: Claude Fable 5 <noreply@anthropic.com> Co-authored-by: Cheng Wan <cheng.wan@radixark.ai>
This commit is contained in:
co-authored by
Caihua Li
Claude Fable 5
Cheng Wan
parent
98cb3535b7
commit
ef9e58fd6d
@@ -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()
|
||||
|
||||
@@ -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(),
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -101,6 +101,20 @@ 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
|
||||
]
|
||||
# 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
|
||||
|
||||
|
||||
@@ -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,
|
||||
*,
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -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,
|
||||
)
|
||||
|
||||
@@ -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()
|
||||
@@ -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)
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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()
|
||||
@@ -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()
|
||||
|
||||
@@ -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()
|
||||
File diff suppressed because it is too large
Load Diff
Reference in New Issue
Block a user