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:
caihuali95
co-authored by
Caihua Li
Claude Fable 5
Cheng Wan
Caihua Li
Claude Fable 5
Cheng Wan
parent
98cb3535b7
commit
ef9e58fd6d
@@ -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