minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 10 Packages Projects Releases Wiki Activity

[Test] Prune redundant unified-memory allocator and pool tests (#38093)

Browse Source
This commit is contained in:
Liangsheng Yin
2026-09-04 20:47:01 -07:00
committed by GitHub
parent 92a4d8b5ee
commit d50e9a9756
21 changed files with 920 additions and 2487 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/models_e2e/test_kimi_linear_unified_memory.py
+16 -43
View File
@@ -1,28 +1,11 @@
"""Kimi-Linear (MLA full attention + KDA linear attention) served from the
unified memory pool.
Under `--enable-unified-memory` the MLA full side is exposed as per-layer
views and every loc the kernels see is a translated virtual id, so the whole
read/write path differs from the static pool. The unit tests pin that pool in
isolation; this is the end-to-end guard. `test_prefix_cache_branching` carries
most of the weight: a radix hit replays virtual locs whose physical pages may
have moved under compaction.
No `--attention-backend` is pinned on purpose -- the test runs whatever the host
resolves to (`fa3` on this suite's H100 runner, also the H200 default). Both
defects found in review on #32972 were reachable only under a resolved default,
which a pinned test hides by construction.
Reference GSM8K, all with `--enable-unified-memory`:
- 2x H200 TP2, resolved default (fa3): 0.917 @400, vs 0.915 static (1 sigma
~= 0.015). This file as written scores 0.920 @200.
- 2x H200 TP2, `--attention-backend triton`: 0.900 @400.
- 1x B300 TP1: 0.915 flashinfer, 0.900 trtllm_mla, @200.
Nightly-only: a second full 48B server launch is too much for per-PR CI on top
of the existing Kimi-Linear e2e coverage.
python -m pytest test/registered/models_e2e/test_kimi_linear_unified_memory.py -v
Under `--enable-unified-memory` the MLA full side is exposed as per-layer views
and every loc the kernels see is a translated virtual id, so the whole
read/write path differs from the static pool. `test_prefix_cache_branching`
carries most of the weight: a radix hit replays virtual locs whose physical
pages may have moved under compaction.
"""
import unittest
@@ -32,18 +15,23 @@ from sglang.test.kits.eval_accuracy_kit import GSM8KMixin
from sglang.test.kits.prefix_cache_branching_kit import PrefixCacheBranchingMixin
from sglang.test.server_fixtures.default_fixture import DefaultServerBase
register_cuda_ci(est_time=1200, stage="nightly", runner_config="4-gpu-h100")
register_cuda_ci(est_time=800, stage="nightly", runner_config="4-gpu-h100")
KIMI_LINEAR_MODEL = "moonshotai/Kimi-Linear-48B-A3B-Instruct"
class TestKimiLinearUnifiedMemory(
class TestKimiLinearUnifiedMemoryFlashMLA(
GSM8KMixin, PrefixCacheBranchingMixin, DefaultServerBase
):
"""flashmla at its ps=64 snap: the block-table route
(KVIndexTranslator.fill_read_table into flashmla's padded tables) plus the
ps=64 sub-pool sizing (64-token sink floor, per-layer-view tail pad).
Hopper-only, like the rest of this nightly suite."""
model = KIMI_LINEAR_MODEL
cache_chunk_size = 64
# Same bar as the static-pool Kimi-Linear e2e test: unified memory must not
# cost accuracy (measured 0.917 vs 0.915 static, see the module docstring).
# cost accuracy.
gsm8k_score_threshold = 0.88
other_args = [
"--trust-remote-code",
@@ -52,17 +40,6 @@ class TestKimiLinearUnifiedMemory(
"--chunked-prefill-size",
"2048",
"--enable-unified-memory",
]
class TestKimiLinearUnifiedMemoryFlashMLA(TestKimiLinearUnifiedMemory):
"""flashmla at its ps=64 snap: the canonical block-table route
(KVIndexTranslator.fill_read_table into flashmla's padded tables) plus the
ps=64 sub-pool sizing (64-token sink floor, per-layer-view tail pad) end to
end.
Hopper-only, like the rest of this nightly suite."""
other_args = TestKimiLinearUnifiedMemory.other_args + [
"--attention-backend",
"flashmla",
"--page-size",
@@ -73,13 +50,9 @@ class TestKimiLinearUnifiedMemoryFlashMLA(TestKimiLinearUnifiedMemory):
class TestKimiLinearUnifiedMemoryDCP(
GSM8KMixin, PrefixCacheBranchingMixin, DefaultServerBase
):
"""Unified memory + decode context parallelism.
`test_prefix_cache_branching` is the sharp one here: a radix hit replays
widened virtual locs whose pages may have moved under compaction, and each
rank must recover the same physical page from them while keeping a
different row inside it.
"""
"""Unified memory + decode context parallelism on flashinfer: on a radix hit
each rank must recover the same physical page from widened virtual locs
while keeping a different row inside it."""
model = KIMI_LINEAR_MODEL
cache_chunk_size = 64
test/registered/page_major/test_page_major_gpt_oss.py
+3 -13
View File
@@ -2,11 +2,8 @@
gpt-oss-20b is uniform-row hybrid-SWA, so its MHA and SWA sub-pools are
per-layer views and the fa3 cell reads them through the translator's read
tables. The resolved-default cell pins the no-pin path, since a pinned
backend hides default-resolution breakage by construction. flashinfer is
absent on purpose: gpt-oss uses attention sinks, which it does not support.
Registered to the label-gated ``run-ci-extra`` suite (opt-in, not per-commit).
tables. flashinfer is absent on purpose: gpt-oss uses attention sinks, which
it does not support.
"""
import unittest
@@ -17,7 +14,7 @@ from sglang.test.ci.ci_register import register_cuda_ci
from sglang.test.server_fixtures.default_fixture import DefaultServerBase
from sglang.test.test_utils import DEFAULT_MODEL_NAME_FOR_TEST_MXFP4_WITH_MOE
register_cuda_ci(est_time=1500, stage="extra-a", runner_config="1-gpu-large")
register_cuda_ci(est_time=1000, stage="extra-a", runner_config="1-gpu-large")
_UNIFIED_COMMON_ARGS = [
"--enable-unified-memory",
@@ -68,12 +65,5 @@ class TestUnifiedGptOssFa3(TestUnifiedGptOssTriton):
other_args = _UNIFIED_COMMON_ARGS + ["--attention-backend", "fa3"]
class TestUnifiedGptOssResolvedDefault(TestUnifiedGptOssTriton):
"""No backend pin: whatever the host resolves must be in the allow-list,
or the server fails to boot under its own defaults."""
other_args = _UNIFIED_COMMON_ARGS
if __name__ == "__main__":
unittest.main()
test/registered/page_major/test_page_major_qwen_hybrid.py
+4 -16
View File
@@ -4,11 +4,7 @@ Qwen3.5-4B is a gated-delta-net / linear-attention hybrid, which exercises the
path most prone to subtle bugs: the Mamba conv/SSM state stays a strided
envelope view (its kernels are stride-aware by design) while the
full-attention KV is per-layer views, which the fa3 / flashinfer cells read
through the translator's read tables. The resolved-default cell pins the
no-pin path, since a pinned backend hides default-resolution breakage by
construction.
Registered to the label-gated ``run-ci-extra`` suite (opt-in, not per-commit).
through the translator's read tables.
"""
import unittest
@@ -19,7 +15,7 @@ from sglang.test.ci.ci_register import register_cuda_ci
from sglang.test.server_fixtures.default_fixture import DefaultServerBase
from sglang.test.test_utils import DEFAULT_HYBRID_GDN_SMALL_MODEL_NAME_FOR_TEST
register_cuda_ci(est_time=1600, stage="extra-a", runner_config="1-gpu-large")
register_cuda_ci(est_time=1200, stage="extra-a", runner_config="1-gpu-large")
_UNIFIED_COMMON_ARGS = [
"--trust-remote-code",
@@ -40,9 +36,8 @@ class TestUnifiedQwenHybridTriton(DefaultServerBase):
model = DEFAULT_HYBRID_GDN_SMALL_MODEL_NAME_FOR_TEST
# Measured ~0.86 in this harness on both the static pools and the envelope
# layout; 0.80 leaves noise margin and still catches a corrupted prefill
# state, which reads ~0.61.
# Measured ~0.86 on both the static pools and the envelope layout; 0.80
# leaves noise margin and still catches a corrupted prefill state (~0.61).
gsm8k_threshold = 0.80
num_gsm8k_questions = 200
num_shots = 5
@@ -84,12 +79,5 @@ class TestUnifiedQwenHybridFlashinfer(TestUnifiedQwenHybridTriton):
other_args = _UNIFIED_COMMON_ARGS + ["--attention-backend", "flashinfer"]
class TestUnifiedQwenHybridResolvedDefault(TestUnifiedQwenHybridTriton):
"""No backend pin: whatever the host resolves must be in the allow-list,
or the server fails to boot under its own defaults."""
other_args = _UNIFIED_COMMON_ARGS
if __name__ == "__main__":
unittest.main()
test/registered/unit/layers/attention/test_kv_translate_ownership.py
+26 -41
View File
@@ -1,20 +1,16 @@
"""Nothing under layers/attention may translate KV ids for itself.
Ownership is exactly two places: `KVIndexTranslator` for READS (indices are
born kernel-facing, backends consume its tables) and the ForwardBatch rebind
(`rebind_write_loc`) for WRITES. Virtual and physical ids share a value range,
so a backend that forgets a translate -- or does one twice -- reads the wrong
rows and nothing crashes. This scan makes both unrepresentable.
Ownership is exactly two places: `KVIndexTranslator` for READS and the
ForwardBatch rebind (`rebind_write_loc`) for WRITES. Virtual and physical ids
share a value range, so a backend that forgets a translate -- or does one
twice -- reads the wrong rows and nothing crashes.
Out of scope, deliberately: the allocator-internal implementations
(`allocator/unified_*` / `unified_memory_pool`), which ARE the mechanism the
Deliberately out of scope: the allocator-internal implementations
(`allocator/unified_*`, `unified_memory_pool`), which ARE the mechanism the
translator calls; the PD transfer plane's `translate_kv_indices_for_transfer`,
which stages for RDMA outside the forward path; and the STATIC SWA pool's
legacy full->swa slot map, a different mapping kind with no virtual/physical
ambiguity -- its call sites are count-pinned below so new ones are added
consciously.
python3 -m pytest test/registered/unit/layers/attention/test_kv_translate_ownership.py -v
ambiguity.
"""
import ast
@@ -60,30 +56,21 @@ def _iter_sources():
class TestUnifiedTranslateBanned(CustomTestCase):
def test_no_unified_translate_calls(self):
"""No backend calls the unified translate surfaces. A hit here means
a backend re-grew its own id-space transition -- the design whose two
failure modes (forgotten translate, duplicated translate) this scan
exists to prevent. Route reads through KVIndexTranslator views and
writes through the ForwardBatch rebind instead."""
"""No backend calls the unified translate surfaces, and none probes an
allocator for translate capability through getattr."""
banned = re.compile(r"\.translate_kv_loc(_kernel_id)?\(")
hits = [
f"{rel}: {m.group(0)}"
for rel, src in _iter_sources()
for m in banned.finditer(src)
]
self.assertEqual(hits, [])
def test_no_translate_capability_probing(self):
"""No backend probes an allocator for translate capability -- the
getattr-hook pattern is how per-backend translation grew the first
time."""
probing = re.compile(r"""getattr\([^)]*['"]translate_kv_loc""")
hits = [rel for rel, src in _iter_sources() if probing.search(src)]
self.assertEqual(hits, [])
calls, probes = [], []
for rel, src in _iter_sources():
calls += [f"{rel}: {m.group(0)}" for m in banned.finditer(src)]
if probing.search(src):
probes.append(rel)
self.assertEqual(calls, [])
self.assertEqual(probes, [])
def test_hooks_module_deleted_and_unimported(self):
"""The per-backend hooks module (the previous owner of backend-side
v2p knowledge) stays deleted, and nothing imports it."""
"""The per-backend v2p hooks module stays deleted, and nothing imports
it."""
self.assertFalse(
os.path.exists(os.path.join(_ATTN_DIR, "unified_mem_hooks.py"))
)
@@ -96,10 +83,9 @@ class TestUnifiedTranslateBanned(CustomTestCase):
def _derive_wrapper_names():
"""Wrapper classes, read off the source so a NEW one shows up the day it is
written: an AttentionBackend subclass whose own __init__ takes another
backend. Per class, not per file -- a file-wide scan passes as soon as any
one class in it forwards."""
"""Wrapper classes discovered from source, so a new one shows up the day it
is written: an AttentionBackend subclass whose own __init__ takes another
backend. Per class, not per file."""
backend = re.compile(r"Att(?:ention|n)Backend")
names = set()
for _rel, src in _iter_sources():
@@ -155,8 +141,8 @@ class _Runner:
def _build_wrappers(translator):
"""One live instance per wrapper. Only the inner that MUST supply the
translator carries it; every other inner carries None, so a wrapper that
copies from the linear / sparse / DSA side ends up with None and fails."""
translator carries it, so a wrapper copying off the linear / sparse / DSA
side ends up with None and fails."""
carrier = _Inner(translator)
# HybridAttnBackend reads the spec bag in __init__; the bag is unpublished
# outside a launched server.
@@ -176,10 +162,9 @@ def _build_wrappers(translator):
class TestWrapperBackendsForwardTranslator(CustomTestCase):
"""BUG REGRESSION. `AttentionBackend.kv_index_translator` defaults to None,
so a wrapper that does not re-expose its inner's copy reads as "needs no
translation" and producers that fetch it off `get_attn_backend()` skip the
virtual->kernel-facing translation instead of failing."""
"""Bug regression: `AttentionBackend.kv_index_translator` defaults to None,
so a wrapper that does not re-expose its inner's copy makes producers skip
the virtual->kernel-facing translation instead of failing."""
def test_every_wrapper_is_constructed_here(self):
self.assertEqual(
test/registered/unit/mem_cache/test_full_loc_fast_path.py
+53 -140
View File
@@ -15,21 +15,9 @@
`HybridLinearKVPool`).
All write-location info travels in the attention metadata (`KVWriteLoc`); the
pools hold none and never translate — the write loc reaching `set_kv_buffer` is
always PHYSICAL. Two routing contracts are pinned here:
1. Full-attention. The full-physical loc is carried in `KVWriteLoc.full_loc`
(from `ForwardMetadata.out_cache_loc_full_physical`) and written directly.
`UnifiedSWAKVPool` asserts it's present (the unified memory pool always precomputes
it); `HybridLinearKVPool` falls back to `loc` for a static (non-shared) pool,
where `loc` is itself already physical.
2. SWA. The swa-physical loc rides the backend `swa_out_cache_loc` slot
(`KVWriteLoc.swa_loc`) and is written directly.
Pure dispatch tests: the inner sub-pools are recording stubs, so no GPU / real
buffers are needed (CPU CI).
python -m pytest test/registered/unit/mem_cache/test_full_loc_fast_path.py -v
pools hold none and never translate, so the loc reaching `set_kv_buffer` is
always PHYSICAL. Pure dispatch: the inner sub-pools are recording stubs, so no
GPU and no real buffers are needed.
"""
import types
@@ -61,9 +49,8 @@ class _RecordingPool:
class TestUnifiedSWARouting(unittest.TestCase):
"""`UnifiedSWAKVPool.set_kv_buffer` routing: full layers write `full_loc`
when present (triton's capture-stable buffer), else the rebound generic
`loc` -- the same id space once the loc is rebound; SWA layers write the swa-physical
`swa_loc`, which has no fallback (a different id space). The pool never
translates."""
`loc` -- the same id space; SWA layers write the swa-physical `swa_loc`,
which has no fallback (a different id space)."""
def _make_bare_pool(self):
from sglang.srt.mem_cache.unified_memory_pool import UnifiedSWAKVPool
@@ -93,18 +80,15 @@ class TestUnifiedSWARouting(unittest.TestCase):
self.assertEqual(len(pool.full_kv_pool.calls), 1)
forwarded, kwargs = pool.full_kv_pool.calls[0]
# Forward the full-physical tensor from the write metadata, NOT the
# virtual loc. No `already_physical` — the pool only ever gets physical.
# virtual loc; no `already_physical`, the pool only ever gets physical.
self.assertIs(forwarded, full_phys)
self.assertIsNot(forwarded, virtual_loc)
self.assertNotIn("already_physical", kwargs)
def test_full_layer_falls_back_to_generic_loc(self):
"""Bug regression: fa3 x unified-SWA crashed at gpt-oss
cuda-graph capture because every backend except triton bundles the
2-arg KVWriteLoc(loc, swa) and the full-layer door demanded an explicit
full_loc. Once the loc is rebound the generic `loc` IS the full-side kernel-facing id
(rebind_write_loc runs at ForwardBatch construction),
so the door must fall back to it -- the pool still never translates."""
"""Bug regression: a 2-arg `KVWriteLoc(loc, swa)` with no explicit
`full_loc` must fall back to the rebound `loc` -- which IS the full-side
kernel-facing id -- instead of failing the full-layer door."""
pool = self._make_bare_pool()
rebound_loc = torch.tensor([10, 11, 12], dtype=torch.int64)
swa_phys = torch.tensor([1, 2, 0], dtype=torch.int64)
@@ -163,10 +147,9 @@ class TestUnifiedSWATombstoneClamp(unittest.TestCase):
"""`UnifiedSWAKVPool.translate_loc_from_full_to_swa` must clamp tombstoned
ids to the reserved padding sink (0).
A token whose swa page was freed carries -1 in `virtual_to_physical`. Before
the clamp, that produced a negative id, which a captured graph stores at a
negative offset from the buffer base. The composite allocator's method of
the same name already clamped; this path did not.
A token whose swa page was freed carries -1 in `virtual_to_physical`, and a
negative id makes a captured graph store at a negative offset from the
buffer base.
"""
def _make_bare_pool(self, page_size, v2p, multiplier=1):
@@ -203,83 +186,66 @@ class TestUnifiedSWATombstoneClamp(unittest.TestCase):
class TestHybridLinearFullLocRouting(unittest.TestCase):
"""`HybridLinearKVPool.set_kv_buffer` (non-MLA) writes the full-physical
`full_loc` from the write metadata when present (unified memory pool), else the
already-physical `loc` (static pool). No translate, no `already_physical`."""
"""`HybridLinearKVPool.set_kv_buffer` writes the full-physical `full_loc`
when present (unified memory pool), else the already-physical `loc` (static
pool) -- the same rule on the MHA and MLA branches."""
def _make_bare_pool(self):
def _make_bare_pool(self, use_mla):
from sglang.srt.mem_cache.memory_pool import HybridLinearKVPool
pool = object.__new__(HybridLinearKVPool)
pool.full_kv_pool = _RecordingPool()
pool.use_mla = False
pool.use_mla = use_mla
pool.full_attention_layer_id_mapping = {0: 0}
return pool
def test_writes_full_loc_from_write_loc(self):
pool = self._make_bare_pool()
virtual_loc = torch.tensor([7, 8, 9], dtype=torch.int64)
full_phys = torch.tensor([2, 3, 4], dtype=torch.int64)
for use_mla in (False, True):
for has_full_loc in (True, False):
with self.subTest(use_mla=use_mla, has_full_loc=has_full_loc):
pool = self._make_bare_pool(use_mla)
loc = torch.tensor([7, 8, 9], dtype=torch.int64)
full_phys = (
torch.tensor([2, 3, 4], dtype=torch.int64)
if has_full_loc
else None
)
layer = types.SimpleNamespace(layer_id=0)
pool.set_kv_buffer(
layer,
_loc_info(virtual_loc, full_phys=full_phys),
torch.zeros(3, 4, 8),
torch.zeros(3, 4, 8),
)
layer = types.SimpleNamespace(layer_id=0)
pool.set_kv_buffer(
layer,
_loc_info(loc, full_phys=full_phys),
torch.zeros(3, 4, 8),
None if use_mla else torch.zeros(3, 4, 8),
)
self.assertEqual(len(pool.full_kv_pool.calls), 1)
forwarded, kwargs = pool.full_kv_pool.calls[0]
self.assertIs(forwarded, full_phys)
self.assertIsNot(forwarded, virtual_loc)
self.assertNotIn("already_physical", kwargs)
def test_falls_back_to_loc_when_absent(self):
# Static (non-shared) pool: no full_loc bundled; `loc` is already
# physical, so write it directly.
pool = self._make_bare_pool()
phys_loc = torch.tensor([7, 8, 9], dtype=torch.int64)
layer = types.SimpleNamespace(layer_id=0)
pool.set_kv_buffer(
layer,
_loc_info(phys_loc),
torch.zeros(3, 4, 8),
torch.zeros(3, 4, 8),
)
self.assertEqual(len(pool.full_kv_pool.calls), 1)
forwarded, kwargs = pool.full_kv_pool.calls[0]
self.assertIs(forwarded, phys_loc)
self.assertNotIn("already_physical", kwargs)
self.assertEqual(len(pool.full_kv_pool.calls), 1)
forwarded, kwargs = pool.full_kv_pool.calls[0]
if has_full_loc:
self.assertIs(forwarded, full_phys)
self.assertIsNot(forwarded, loc)
else:
# Static (non-shared) pool: no full_loc bundled; `loc`
# is already physical, so write it directly.
self.assertIs(forwarded, loc)
self.assertNotIn("already_physical", kwargs)
class _RecordingMLAPool(_RecordingPool):
"""Also records the model-level MLA entry points."""
"""Also records the model-level MLA write entry point."""
def __init__(self):
super().__init__()
self.mla_set_calls = []
self.mla_get_calls = []
def set_mla_kv_buffer(self, layer, loc, cache_k_nope, cache_k_rope):
self.mla_set_calls.append(loc)
def get_mla_kv_buffer(self, layer, loc, dst_dtype=None):
self.mla_get_calls.append(loc)
return None, None
class TestHybridLinearMLARouting(unittest.TestCase):
"""MLA-side routing contracts of `HybridLinearKVPool`:
- `set_kv_buffer` (MLA branch) mirrors the MHA branch — write the
pre-translated `KVWriteLoc.full_loc` when present (unified pool, where it
carries the kernel-facing loc), else the raw `loc` (static pool, already physical).
- `set_mla_kv_buffer` / `get_mla_kv_buffer` forward `loc` untouched:
writes are kernel-facing since the ForwardBatch rebind, and read
indices are translated at their production sites."""
"""MLA-side door contract of `HybridLinearKVPool`: `set_mla_kv_buffer`
forwards `loc` untouched -- writes are kernel-facing since the ForwardBatch
rebind."""
def _make_bare_pool(self):
from sglang.srt.mem_cache.memory_pool import HybridLinearKVPool
@@ -290,46 +256,10 @@ class TestHybridLinearMLARouting(unittest.TestCase):
pool.full_attention_layer_id_mapping = {0: 0}
return pool
def test_mla_writes_full_loc_from_write_loc(self):
pool = self._make_bare_pool()
virtual_loc = torch.tensor([7, 8, 9], dtype=torch.int64)
kernel_phys = torch.tensor([21, 24, 27], dtype=torch.int64)
layer = types.SimpleNamespace(layer_id=0)
pool.set_kv_buffer(
layer,
_loc_info(virtual_loc, full_phys=kernel_phys),
torch.zeros(3, 1, 8),
None,
)
self.assertEqual(len(pool.full_kv_pool.calls), 1)
forwarded, _ = pool.full_kv_pool.calls[0]
self.assertIs(forwarded, kernel_phys)
self.assertIsNot(forwarded, virtual_loc)
def test_mla_falls_back_to_loc_when_absent(self):
pool = self._make_bare_pool()
phys_loc = torch.tensor([7, 8, 9], dtype=torch.int64)
layer = types.SimpleNamespace(layer_id=0)
pool.set_kv_buffer(
layer,
_loc_info(phys_loc),
torch.zeros(3, 1, 8),
None,
)
self.assertEqual(len(pool.full_kv_pool.calls), 1)
forwarded, _ = pool.full_kv_pool.calls[0]
self.assertIs(forwarded, phys_loc)
def test_set_mla_kv_buffer_door_never_translates(self):
"""Physical-loc contract: the write door forwards `loc` UNTOUCHED.
The translate happens exactly once at ForwardBatch construction
(rebind_write_loc, kernel-facing-first); a door that translated
again would double-translate every unified MLA write. Deleting the
forward (or re-adding a door translate) turns this red."""
"""The translate happens exactly once, at ForwardBatch construction
(`rebind_write_loc`); a door that translated again would
double-translate every unified MLA write."""
pool = self._make_bare_pool()
loc = torch.tensor([107, 108, 109], dtype=torch.int64)
layer = types.SimpleNamespace(layer_id=0)
@@ -339,21 +269,6 @@ class TestHybridLinearMLARouting(unittest.TestCase):
self.assertEqual(len(pool.full_kv_pool.mla_set_calls), 1)
self.assertIs(pool.full_kv_pool.mla_set_calls[0], loc)
def test_get_mla_kv_buffer_door_never_translates(self):
"""Kernel-facing contract, read side: `loc` is a read-index tensor
already translated at its production site
(fetch_mha_one_shot_kv_indices / prepare_chunked_kv_indices); the
door forwards it UNTOUCHED -- a re-added door translate would
double-translate every unified MLA prefix read."""
pool = self._make_bare_pool()
loc = torch.tensor([104, 105], dtype=torch.int64)
layer = types.SimpleNamespace(layer_id=0)
pool.get_mla_kv_buffer(layer, loc)
self.assertEqual(len(pool.full_kv_pool.mla_get_calls), 1)
self.assertIs(pool.full_kv_pool.mla_get_calls[0], loc)
class TestMlaWriteDoorsUnderDcp(unittest.TestCase):
"""Which MLA write door is DCP-aware, and which refuses.
@@ -362,10 +277,8 @@ class TestMlaWriteDoorsUnderDcp(unittest.TestCase):
the DCP write. `set_kv_buffer` (the combined latent+rope row) cannot: the
two backends that could reach it disagree on the loc space -- flashinfer's
`k_rope is None` branch passes a WIDENED loc, the Triton backend one it
already collapsed -- so there is no single correct translation. It used to
select `loc % dcp_size == dcp_rank` and then write WITHOUT dividing, i.e.
widened ids straight into a rank-local buffer. Refusing is the contract;
a re-added masked-but-undivided write is what this guards."""
already collapsed -- so there is no single correct translation and refusing
is the contract."""
def _bare_mla_pool(self):
from sglang.srt.mem_cache.memory_pool import MLATokenToKVPool
test/registered/unit/mem_cache/test_kv_index_translator.py
+46 -98
View File
@@ -13,29 +13,8 @@
# ==============================================================================
"""KVIndexTranslator -- the read-path id translator.
Covers, CPU-only (the builder's pure-torch reference path; GPU parity of the
Triton kernel is a later CUDA CI pin):
- strict passthrough: a non-unified source returns the SAME req_to_token /
req_pool_indices objects -- zero tensor ops, no copies (the property that
makes backend re-pointing byte-identical for every non-unified server);
- static SWA pools keep their legacy full->swa mapping on the view;
- the read table matches the hand formula
entry[b, c] = clamp(v2p[req_to_token[req[b], c*ps] // ps] * mult, 0)
over the REAL SWA composite's tables (full AND swa, ps in {1, 4},
multiplier in {1, 2L}), with the swa table built from VIRTUAL ids;
- sink routing: dead lanes (seq_len 0), -1 req_to_token entries, and
tombstoned v2p pages all read entry 0;
- the capture contract: buffers are zero-filled and idempotent; a refresh
updates ONLY the live prefix (stale tails and rows beyond bs keep prior
contents); the returned table is the WHOLE buffer (pointer-stable);
- the eager-view memo: a single source-resident slot keyed by batch
identity (same batch shares one build; the next batch replaces it; a
dead batch never matches);
- the two-phase write contract: the rebind touches only the full side, and
the sliding-window write loc derives POINTWISE from the kernel-facing values
(pads, slices, and fresh copies included), for both pool families.
python -m pytest test/registered/unit/mem_cache/test_kv_index_translator.py -v
CPU-only: these exercise the builder's pure-torch reference path, not the
Triton kernel.
"""
from sglang.test.ci.ci_register import register_cpu_ci
@@ -137,9 +116,8 @@ def _reference_table(req_to_token, req_pool_indices, seq_lens, v2p, mult, ps, wi
class TestPassthrough(unittest.TestCase):
def test_non_unified_returns_same_objects(self):
"""The strict-passthrough property: no copy, no branch, the exact
tensors backends read today. A regression here (any tensor op on the
non-unified path) breaks byte-identity for every static-pool server."""
"""Strict passthrough: no copy, no branch. Any tensor op on the
non-unified path breaks byte-identity for every static-pool server."""
req_to_token = torch.arange(64, dtype=torch.int64).reshape(4, 16)
src = KVIndexTranslator(
req_to_token=req_to_token,
@@ -166,7 +144,7 @@ class TestPassthrough(unittest.TestCase):
def _alloc_and_fill(allocator, ps, lens):
"""Allocate per-request virtual runs and write them into a fake
req_to_token; returns (req_to_token, req_pool_indices, seq_lens)."""
req_to_token."""
width = 16 * ps
req_to_token = torch.full((len(lens), width), -1, dtype=torch.int64)
for r, n in enumerate(lens):
@@ -183,11 +161,10 @@ def _alloc_and_fill(allocator, ps, lens):
class TestReadTableBuild(unittest.TestCase):
def test_read_table_matches_reference_across_multipliers(self):
"""The load-bearing formula pin: full AND swa read tables equal
the independent per-element derivation, across page sizes and both
multiplier regimes (MLA=1, MHA=2L). The swa table agreeing with
a formula over VIRTUAL ids is also the never-chained-through-
full-physical proof."""
"""Both read tables must equal the independent per-element derivation,
across page sizes and both multiplier regimes (MLA=1, MHA=2L); the swa
table agreeing over VIRTUAL ids proves it is never chained through
full-physical."""
for ps in (1, 4):
for collapse in (True, False):
allocator = _build_composite(ps, collapse=collapse)
@@ -235,10 +212,8 @@ class TestReadTableBuild(unittest.TestCase):
)
def test_packed_stream_equals_the_rectangle_it_replaces(self):
"""The packed builder and the rectangle must agree element for element:
packed[indptr[b] + p] == ids[b, p // ps] * ps + p % ps. Consumers that
plan over the stream and consumers that read the page table have to see
the same KV, so a change to either builder alone turns this red."""
"""The two builders must agree element for element:
packed[indptr[b] + p] == ids[b, p // ps] * ps + p % ps."""
for ps in (1, 4):
allocator = _build_composite(ps)
req_to_token, rows, seq_lens = _alloc_and_fill(
@@ -269,9 +244,9 @@ class TestReadTableBuild(unittest.TestCase):
)
def test_sink_routing(self):
"""Dead lanes (seq_len 0), -1 slots inside the live prefix, and
tombstoned v2p pages must ALL read entry 0 -- one wild entry is a
captured-graph OOB read at replay."""
"""Dead lanes, -1 slots inside the live prefix, and tombstoned v2p
pages must ALL read entry 0; one wild entry is a captured-graph OOB
read at replay."""
ps = 4
allocator = _build_composite(ps)
req_to_token, rows, seq_lens = _alloc_and_fill(
@@ -300,12 +275,9 @@ class TestBuildInto(unittest.TestCase):
(their rows ARE the read table's rows)."""
def test_prefix_filled_tail_sentinel_preserved_width_capped(self):
"""Three contracts in one batch: entries equal the read-table formula,
lanes past each row's live pages keep the backend's -1 sentinel
(prefix-only -- a tail write scatters the trtllm sentinel contract),
and a table padded WIDER than the req_to_token page span (trtllm's
LCM alignment) is capped instead of tripping the builder's width
assert."""
"""The -1 tail sentinel belongs to the backend, and a table padded
WIDER than the req_to_token page span (trtllm's LCM alignment) must be
capped rather than trip the builder's width assert."""
ps = 4
allocator = _build_composite(ps)
full_mult = allocator.kernel_page_multiplier
@@ -339,8 +311,8 @@ class TestBuildInto(unittest.TestCase):
)
def test_passthrough_source_refuses(self):
"""Callers dispatch on `enabled`; a passthrough source has no v2p to
build from and must fail loud, not fill garbage."""
"""Callers dispatch on `reads_are_translated`; a passthrough source has
no v2p to build from and must fail loud, not fill garbage."""
src = KVIndexTranslator(
req_to_token=torch.zeros((2, 4), dtype=torch.int64),
token_to_kv_pool_allocator=SimpleNamespace(),
@@ -360,21 +332,18 @@ class TestPoolOwnership(unittest.TestCase):
"""A runner only gets the kernel-facing id space when the pool IT reads and
writes is the one the allocator's ids address.
Guarded shape: a runner handed a SHARED allocator (one slot index space,
one req_to_token) while owning a SEPARATE KV buffer sized to the
allocator's SLOT count. Probing the allocator alone reports "unified" for
that runner, so its indices would be mapped into the composite's
kernel-facing space (kernel-facing ids up to num_pages * multiplier) and then used
to address a buffer with only num_slots rows -- out of bounds on both the
read gather and the KV store.
Guarded shape: a runner handed a SHARED allocator while owning a SEPARATE
KV buffer sized to the allocator's SLOT count. Probing the allocator alone
reports "unified" for that runner, so its indices would be mapped into the
composite's kernel-facing space (ids up to num_pages * multiplier) and used
to address a buffer with only num_slots rows.
"""
def test_real_factory_bundle_satisfies_the_ownership_identity(self):
"""The guard rests on `allocator.get_kvcache() is token_to_kv_pool`
holding for a REAL target bundle. If a factory ever returned a pool
the allocator does not hold, the guard would silently disable the
unified path for EVERY model -- so pin it against the real factory
rather than against this file's own construction."""
"""The guard rests on `allocator.get_kvcache() is token_to_kv_pool`, so
a factory returning a pool the allocator does not hold would silently
disable the unified path for EVERY model. Pinned against the real
factory, not this file's own construction."""
from sglang.srt.mem_cache.unified_memory_pool import init_unified_swa_pools
bundle = init_unified_swa_pools(
@@ -424,9 +393,9 @@ class TestPoolOwnership(unittest.TestCase):
self.assertFalse(src.is_translating)
def test_disabled_source_is_strict_passthrough(self):
"""Consequence of the guard: such a runner must see RAW virtual ids on
the read table -- they index its own pool directly. A translate here is
the out-of-bounds bug the ownership identity exists to prevent."""
"""Such a runner must see RAW virtual ids: they index its own pool
directly, and a translate here is the out-of-bounds bug the ownership
identity exists to prevent."""
alloc = _build_composite(ps=1)
req_to_token = torch.arange(16, dtype=torch.int32, device=_DEV).view(2, 8)
src = KVIndexTranslator(
@@ -483,10 +452,9 @@ class TestCaptureContract(unittest.TestCase):
self.assertTrue(bool((cap[1:] == 7).all()), "rows beyond bs were touched")
def test_row_ids_not_reallocated_across_builds(self):
"""`row_ids` is a constant arange sized once from the request pool, so
builds at different batch sizes hand back slices of ONE buffer. A
per-build `torch.arange` would be correct but would spend an allocation
and a launch on every replay prep."""
"""`row_ids` is one arange sized from the request pool and sliced per
build; a per-build `torch.arange` would be correct but costs an
allocation and a launch on every replay prep."""
allocator = _build_composite(1)
req_to_token, rows, seq_lens = _alloc_and_fill(allocator, 1, lens=[4, 2, 3])
src = _make_source(allocator, req_to_token, 1)
@@ -531,10 +499,9 @@ class _FakeForwardBatch:
class TestViewMemo(unittest.TestCase):
"""The eager view is memoized ON THE SOURCE in a single slot keyed by
batch identity -- per-batch state stays out of the ForwardBatch (it does
not scale with the number of id spaces), and one metadata build's many
consumers still share one table build."""
"""The eager view is memoized ON THE SOURCE in a single slot keyed by batch
identity, so per-batch state stays out of the ForwardBatch while one
metadata build's many consumers still share one table build."""
def _fb(self, allocator, ps, lens):
req_to_token, rows, seq_lens = _alloc_and_fill(allocator, ps, lens=lens)
@@ -590,12 +557,10 @@ class TestViewMemo(unittest.TestCase):
class TestWriteLoc(unittest.TestCase):
"""The two-phase write contract: phase 1 (`rebind_write_loc`) rebinds the
full side once at ForwardBatch construction; phase 2 derives the
sliding-window write loc on demand, POINTWISE from the full-side
values. Value-based derivation is the property under test: pads, slices,
and fresh copies of the loc must all derive correctly with no handover
and no stored per-forward state."""
"""The two-phase write contract: `rebind_write_loc` rebinds the full side
once at ForwardBatch construction, and the sliding-window write loc derives
POINTWISE from the full-side values -- pads, slices, and fresh copies
included -- with no handover and no stored per-forward state."""
def _built(self, ps=1, n=4):
allocator = _build_composite(ps)
@@ -622,10 +587,8 @@ class TestWriteLoc(unittest.TestCase):
self.assertTrue(torch.equal(virt, keep))
def test_swa_write_loc_round_trips_from_full_side(self):
"""The derived property behind phase 2: for any virtual run t,
deriving from the kernel-facing full-side values must equal the direct
virtual->swa translate — `field(full(t)) == swa(t)` across page sizes
and multipliers."""
"""Derived property: `field(full(t)) == swa(t)` for any virtual run t,
across page sizes and multipliers."""
for ps in (1, 4, 64):
src, _, rows, seq_lens, _, want_full, want_swa = self._built(
ps=ps, n=3 * ps
@@ -634,30 +597,15 @@ class TestWriteLoc(unittest.TestCase):
self.assertTrue(torch.equal(got, want_swa))
def test_pad_lanes_derive_to_sink(self):
"""The DP pad appends zeros; kernel-facing 0 is the reserved padding slot in
every id space, so pad lanes must derive to swa slot 0 with no
`num_live` bookkeeping."""
"""The DP pad appends zeros, and kernel-facing 0 is the reserved
padding slot in every id space, so pad lanes derive to swa slot 0 with
no `num_live` bookkeeping."""
src, _, rows, seq_lens, _, want_full, want_swa = self._built(n=3)
padded = torch.cat([want_full, want_full.new_zeros(2)])
got = self._field(src, rows, seq_lens, padded)
self.assertTrue(torch.equal(got[:3], want_swa))
self.assertTrue(bool((got[3:] == 0).all()), "pad lanes must land on slot 0")
def test_slice_and_copy_derive_pointwise_without_handover(self):
"""REGRESSION (design): the retired identity-resolver refused any
tensor it had not been handed -- a TBO child's re-padded slice or a
registry's fresh copy raised. Value-based derivation must accept
both, pointwise, with no adopt/handover call."""
src, _, rows, seq_lens, _, want_full, want_swa = self._built(n=4)
padded = torch.cat([want_full, want_full.new_zeros(2)])
# TBO-child shape: a slice crossing the pad boundary.
got = self._field(src, rows, seq_lens, padded[2:6])
self.assertTrue(torch.equal(got[:2], want_swa[2:4]))
self.assertTrue(bool((got[2:] == 0).all()))
# Registry shape: a fresh equal-value copy.
got2 = self._field(src, rows, seq_lens, want_full.clone())
self.assertTrue(torch.equal(got2, want_swa))
def test_tombstoned_swa_page_clamps_to_sink(self):
src, allocator, rows, seq_lens, virt, want_full, _ = self._built(ps=1, n=2)
allocator.swa_v2p_page_table[int(virt[0])] = -1
test/registered/unit/mem_cache/test_layout_compat.py
+6 -73
View File
@@ -11,22 +11,12 @@
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Unit tests for the page-major envelope byte layout.
"""``move_kv_cache_native`` -- the stock per-layer 3-D move static-pool
compaction rides on -- must stay byte-exact. CPU-only.
The subject here is the ENVELOPE — the byte layout the unified pool stores its
KV in — pinned through ``MHASubPoolSpec``'s offset math. The 3-D per-layer views
the pool exposes over the same bytes are covered by
``test_unified_mha_views.py``, which also pins the view addressing
against the envelope formula byte for byte.
Verifies that:
1. ``MHASubPoolSpec.layer_k_offset_in_page`` / ``layer_v_offset_in_page`` math
matches the layout intent at ``page_size == 1`` and ``> 1``.
2. ``move_kv_cache_native`` (the stock per-layer 3-D move) stays byte-exact.
CPU-only — no GPU / Triton needed.
python -m pytest test/registered/unit/mem_cache/test_layout_compat.py -v
The page-major envelope layout and the per-layer views over it are covered by
``test_unified_mha_views.py``, which pins the view addressing against the
envelope formula byte for byte.
"""
from sglang.test.ci.ci_register import register_cpu_ci
@@ -38,69 +28,12 @@ import unittest
import torch
from sglang.srt.mem_cache.memory_pool import move_kv_cache_native
from sglang.srt.mem_cache.unified_memory_pool import MHASubPoolSpec
_DEV = "cpu"
def _make_mha_spec(name, grow, layer_num=2, head_num=2, head_dim=4):
return MHASubPoolSpec(
name=name,
layer_num=layer_num,
head_num=head_num,
head_dim=head_dim,
store_dtype=torch.float16,
grow_direction=grow,
)
class TestMHASpecLayerOffsets(unittest.TestCase):
"""Verify ``layer_k_offset_in_page`` / ``layer_v_offset_in_page`` math."""
def test_offsets_at_page_size_1_match_envelope(self):
spec = _make_mha_spec("full", "up", layer_num=3, head_num=2, head_dim=4)
# At ps=1, layer-major within a 1-token page IS envelope-per-token.
# Layer L's K offset = L * (k_row + v_row); V offset = +k_row.
k_row = spec.k_row_bytes()
v_row = spec.v_row_bytes()
for L in range(spec.layer_num):
self.assertEqual(
spec.layer_k_offset_in_page(L, page_size=1),
L * (k_row + v_row),
)
self.assertEqual(
spec.layer_v_offset_in_page(L, page_size=1),
L * (k_row + v_row) + k_row,
)
def test_offsets_at_page_size_gt_1(self):
spec = _make_mha_spec("full", "up", layer_num=3, head_num=2, head_dim=4)
ps = 8
k_row = spec.k_row_bytes()
v_row = spec.v_row_bytes()
# Layer L's K block within the page starts at L * ps * (k_row+v_row).
# V block starts at +ps * k_row.
for L in range(spec.layer_num):
self.assertEqual(
spec.layer_k_offset_in_page(L, page_size=ps),
L * ps * (k_row + v_row),
)
self.assertEqual(
spec.layer_v_offset_in_page(L, page_size=ps),
L * ps * (k_row + v_row) + ps * k_row,
)
def test_page_bytes(self):
spec = _make_mha_spec("full", "up", layer_num=3, head_num=2, head_dim=4)
# page_bytes = page_size * entry_bytes (preserved invariant)
for ps in [1, 8, 64, 256]:
self.assertEqual(spec.page_bytes(ps), ps * spec.entry_bytes())
class TestMoveKVCacheNative(unittest.TestCase):
def test_move_kv_cache_3d_path_unchanged(self):
"""The stock per-layer 3-D move must relocate exactly the named token
rows, byte-identically — compaction on static pools rides on it."""
rows, byte-identically; compaction on static pools rides on it."""
k = [torch.zeros((32, 2, 4), dtype=torch.float16) for _ in range(2)]
v = [torch.zeros((32, 2, 4), dtype=torch.float16) for _ in range(2)]
for L in range(2):
test/registered/unit/mem_cache/test_multi_ended_allocator.py
+205 -675
View File
File diff suppressed because it is too large Load Diff
test/registered/unit/mem_cache/test_swa_locked_full_recover_unified.py
+29 -57
View File
@@ -11,31 +11,20 @@
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Locked-full SWA tombstone-recovery under the unified pool (action handler).
"""Locked-full SWA tombstone recovery under the unified pool (action handler).
`RecoverSWAWithLockedFull` recovers a tombstoned SWA node whose full value is
LOCKED: the node cannot adopt the incoming request's ids wholesale, so the
static-pool recipe hands the node the INCOMING ids' swa pages, frees only their
FULL pages, and re-points the locked ids through `full_to_swa_index_mapping`.
`RecoverSWAWithLockedFull` hands a tombstoned SWA node the INCOMING request's
swa pages and frees only their FULL pages. The static recipe re-points the
locked ids through `full_to_swa_index_mapping`; the unified composite has no
such tensor -- the swa sub-pool's v2p IS the mapping -- so it expresses the
same move as a page-ownership rebind and then frees the incoming ids through
the composite, whose `swa_v2p_pages > 0` filter skips the tombstoned swa side.
The unified composite has no mapping tensor — the swa sub-pool's v2p IS the
mapping — and its `set_full_to_swa_mapping` is an explicit no-op stub. The
pre-fix handler therefore raised AttributeError on `full_to_swa_index_mapping`
(and, had that line been removed, would have silently skipped the rebind while
line 1 freed swa pages the kept ids still referenced). The fix expresses the
same move as a page-ownership REBIND: bind the node's virtual pages to the
incoming pages' physical pages, tombstone the incoming ones, then free the
incoming ids through the composite — whose `swa_v2p_pages > 0` filter skips the
tombstoned swa side, releasing ONLY the full side.
Why the recovery must succeed rather than decline (the v1 lesson, still true on
this branch): the TreeCore insert walk counts the node in `prefix_len`
regardless of component consumption, while the SWA match validator rejects a
`value is None` node — a declined recovery makes `insert` report a prefix the
follow-up `match_prefix` cannot honor, tripping
The recovery must succeed rather than decline: the TreeCore insert walk counts
the node in `prefix_len` regardless of component consumption, while the SWA
match validator rejects a `value is None` node, so a declined recovery reports
a prefix `match_prefix` cannot honor and trips
`new_prefix_len <= len(new_indices)` in `cache_unfinished_req`.
python -m pytest test/registered/unit/mem_cache/test_swa_locked_full_recover_unified.py -v
"""
import unittest
@@ -121,8 +110,8 @@ class _Probe(SWAComponent):
class _StaticAllocRecorder:
"""Stands in for the STATIC SWATokenToKVPoolAllocator: has the mapping
tensor and a real set_full_to_swa_mapping. The handler must keep routing
static pools through the original recipe."""
tensor and a real `set_full_to_swa_mapping`, so the handler must keep
routing static pools through the original recipe."""
def __init__(self, n=16):
self.full_to_swa_index_mapping = torch.arange(n, dtype=torch.int64)
@@ -133,9 +122,8 @@ class _StaticAllocRecorder:
self.full_attn_allocator = self
def set_full_to_swa_mapping(self, full, swa):
# Honour the write like the real static allocator: the handler routes
# every mapping write THROUGH the API (never by indexing the tensor),
# so the fake must apply it for the mapping asserts to observe it.
# The handler routes every mapping write THROUGH the API, so the fake
# must apply it for the mapping asserts to observe anything.
self.mapping_calls.append((full, swa))
self.full_to_swa_index_mapping[full.to(torch.int64)] = swa.to(torch.int64)
@@ -171,9 +159,9 @@ class _RecoverTestBase(unittest.TestCase):
class TestPagePairing(_RecoverTestBase):
def test_pairs_positionally_not_by_sorted_id(self):
"""Allocation hands out virtual ids in no particular order; deduping
with `torch.unique` (which sorts) would bind the node's page k to an
unrelated incoming page — silent wrong-KV."""
"""Allocation hands out virtual ids in no particular order, so a
`torch.unique` dedup (which sorts) would bind the node's page k to an
unrelated incoming page."""
probe, _ = self._probe()
kept = torch.tensor([9, 7, 5], dtype=torch.int64) # descending
incoming = torch.tensor([2, 4, 6], dtype=torch.int64) # ascending
@@ -225,15 +213,17 @@ class TestOwnershipTransfer(_RecoverTestBase):
class TestRecoverActionHandler(_RecoverTestBase):
def test_recovery_sets_a_live_device_value_and_frees_only_the_full_side(self):
"""End-to-end through apply_component_action — the pre-fix handler
raises AttributeError (`full_to_swa_index_mapping`) on this exact
call. Post-fix: the node gets a LIVE swa value, the HANDLER neither
allocates nor frees any swa page (ownership only moves), and the
incoming ids' FULL side returns to the pool."""
"""Bug regression: recovery must give the node a LIVE swa value and
return only the incoming ids' FULL side, allocating and freeing no swa
page of its own."""
probe, allocator = self._probe()
swa = allocator.swa_attn_allocator
kept, incoming = self._two_ranges(allocator)
allocator.free_swa(kept) # what eviction does when it tombstones
self.assertTrue(
bool((allocator.translate_loc_from_full_to_swa(kept) == 0).all()),
"precondition: a tombstoned range translates to the sink",
)
# Snapshot AFTER the setup traffic: the invariant under test is that
# the recovery handler itself moves ownership without moving capacity.
swa_live = swa.allocated_count()
@@ -260,21 +250,6 @@ class TestRecoverActionHandler(_RecoverTestBase):
full_avail + len(incoming),
"the incoming ids' FULL side must come back",
)
def test_recovered_ids_translate_to_live_pages_not_the_sink(self):
"""The tombstoned range translates to the clamped sink before the
recovery and to real pages after — recovering from the node's OWN
already-freed ids (instead of the donated ones) reintroduces the sink."""
probe, allocator = self._probe()
kept, incoming = self._two_ranges(allocator)
allocator.free_swa(kept)
self.assertTrue(
bool((allocator.translate_loc_from_full_to_swa(kept) == 0).all()),
"precondition: a tombstoned range translates to the sink",
)
probe.apply_component_action(
RecoverSWAWithLockedFull(node_id=1, kept_full=kept, incoming_full=incoming)
)
self.assertTrue(
bool((allocator.translate_loc_from_full_to_swa(kept) > 0).all()),
"after recovery the node's ids must address live swa pages",
@@ -283,8 +258,8 @@ class TestRecoverActionHandler(_RecoverTestBase):
class TestStaticPoolPathUnchanged(unittest.TestCase):
def test_static_allocator_keeps_the_mapping_recipe(self):
"""A static SWA allocator (has the mapping tensor) must keep the
original recipe — the unified branch must not hijack it."""
"""A static SWA allocator must keep the mapping recipe; the unified
branch must not hijack it."""
static = _StaticAllocRecorder()
probe = _Probe.__new__(_Probe)
probe.cache = _Cache(static)
@@ -296,11 +271,8 @@ class TestStaticPoolPathUnchanged(unittest.TestCase):
RecoverSWAWithLockedFull(node_id=3, kept_full=kept, incoming_full=incoming)
)
# Both mapping writes go through the allocator API -- the kept remap
# via set_full_to_swa_mapping, the incoming tombstone via
# clear_full_to_swa_mapping -- never by indexing
# `full_to_swa_index_mapping` (the tensor is absent on the unified
# composite by design).
# Both mapping writes go through the allocator API, never by indexing
# `full_to_swa_index_mapping` (absent on the unified composite).
self.assertEqual(len(static.mapping_calls), 1, "static recipe must run")
self.assertEqual(len(static.clear_calls), 1, "incoming must be tombstoned")
self.assertTrue(
test/registered/unit/mem_cache/test_swa_pool_v_head_dim.py
+9 -72
View File
@@ -11,55 +11,32 @@
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""`SWAKVPool.get_v_head_dim()` — the pool method a mambaish + hybrid-SWA
model reaches on boot.
"""`SWAKVPool.get_v_head_dim()` -- the pool method a mambaish + hybrid-SWA model
reaches on boot.
`TritonAttnBackend.__init__` picks its `v_head_dim` from one of three
branches, and the middle one asks the POOL:
if sliding_window_size is not None and swa_v_head_dim != v_head_dim:
... from model_config ... # asymmetric hybrid SWA
elif mambaish_config(model_config) is not None:
v_head_dim = token_to_kv_pool.get_v_head_dim() # <-- this one
else:
... from get_value_buffer(start_layer) ...
A model that is BOTH mambaish AND hybrid-SWA with MATCHING full/SWA value
head dims (Inkling-class) skips the first branch and lands in the second —
where its pool is an SWA-shaped pool, which had no `get_v_head_dim`. The
server died at backend construction with
AttributeError: 'SWAKVPool' object has no attribute 'get_v_head_dim'
on the STATIC pool and, identically, on `UnifiedSWAKVPool`. Neither the
mamba-hybrid pools (`HybridLinearKVPool` has the method) nor pure hybrid-SWA
models (not mambaish, so the branch is never taken) can reach it, which is
why it went unnoticed.
python -m pytest test/registered/unit/mem_cache/test_swa_pool_v_head_dim.py -v
`TritonAttnBackend.__init__` asks the POOL for `v_head_dim` when the model is
mambaish and its full/SWA value head dims MATCH (Inkling-class), so an
SWA-shaped pool without the method kills backend construction with
AttributeError.
"""
import inspect
import unittest
import torch
from sglang.srt.mem_cache.memory_pool import HybridLinearKVPool
from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool
from sglang.srt.mem_cache.unified_memory_pool import UnifiedSWAKVPool
from sglang.test.ci.ci_register import register_cpu_ci
register_cpu_ci(est_time=15, suite="base-a-test-cpu")
_DEV = "cpu"
_FULL_V_HEAD_DIM = 8
_SWA_V_HEAD_DIM = 8 # MATCHING — this is what routes Inkling into the branch
_SWA_V_HEAD_DIM = 8 # MATCHING: this is what routes Inkling into the branch
def _swa_pool():
"""A static SWAKVPool with the Inkling-class layer split: full and SWA
layers interleaved, layer 0 NOT a full-attention layer (which is exactly
why the backend asks the pool instead of indexing layer 0)."""
"""Inkling-class layer split: full and SWA layers interleaved, layer 0 NOT a
full-attention layer, which is why the backend asks the pool at all."""
return SWAKVPool(
size=32,
size_swa=16,
@@ -76,49 +53,9 @@ def _swa_pool():
class TestSWAPoolVHeadDim(unittest.TestCase):
def test_static_pool_reports_the_full_side_value_head_dim(self):
"""Red before the fix with AttributeError; the value must be the FULL
side's, since that is the geometry the caller means."""
pool = _swa_pool()
self.assertEqual(pool.get_v_head_dim(), _FULL_V_HEAD_DIM)
def test_answer_matches_the_full_pool_buffer_not_layer_zero(self):
"""Layer 0 is an SWA layer here, so a naive `get_value_buffer(0)`
would read the SWA side. Pin that the method routes through the FULL
sub-pool at its own start_layer — the property that makes it correct
under pipeline parallelism too."""
pool = _swa_pool()
want = pool.full_kv_pool.get_value_buffer(pool.full_kv_pool.start_layer).shape[
-1
]
self.assertEqual(pool.get_v_head_dim(), want)
# And layer 0 really is the SWA side in this fixture.
_, is_swa = pool.layers_mapping[0]
self.assertTrue(is_swa, "fixture must keep layer 0 on the SWA side")
def test_unified_swa_pool_inherits_it(self):
"""`UnifiedSWAKVPool` subclasses `SWAKVPool`, so the unified tri-pool
path (mambaish + hybrid SWA in one buffer) is covered by the same
method — no second implementation to drift."""
self.assertTrue(issubclass(UnifiedSWAKVPool, SWAKVPool))
self.assertIs(
UnifiedSWAKVPool.get_v_head_dim,
SWAKVPool.get_v_head_dim,
"the unified pool must inherit the method, not shadow it",
)
def test_signature_matches_the_hybrid_linear_precedent(self):
"""The backend calls this method on whichever pool it holds, so every
pool reachable from the mambaish branch must expose the SAME
zero-argument shape. `HybridLinearKVPool` is the precedent this one
mirrors; a future pool added to that branch has to match too."""
for cls in (SWAKVPool, HybridLinearKVPool):
sig = inspect.signature(cls.get_v_head_dim)
self.assertEqual(
[p for p in sig.parameters if p != "self"],
[],
f"{cls.__name__}.get_v_head_dim must take no arguments",
)
if __name__ == "__main__":
unittest.main()
test/registered/unit/mem_cache/test_unified_byte_accounting.py
+36 -61
View File
@@ -11,28 +11,17 @@
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Byte-conservation verifier for the unified 2-pool composites.
"""Byte-conservation verifier (`verify_byte_accounting`) for the unified
2-pool composites.
`verify_byte_accounting` is the idle-time tripwire the token-identity leak
check cannot provide: the unified pool's correctness rests on BYTE bookkeeping
(watermark spans, holes, pending compaction, frontier ordering inside one
shared buffer), and a drifted counter admits requests into memory that is not
actually free — silent corruption territory, not a crash.
The unified pool's correctness rests on BYTE bookkeeping -- watermark spans,
holes, pending compaction, frontier ordering inside one shared buffer -- and
a drifted counter admits requests into memory that is not actually free.
Nothing about that is visible to the token-identity leak check, so this
verifier is the only idle-time tripwire for it.
Derived properties pinned here:
* Conservation: on a lazy end pool the watermark span must equal
live + holes + pending pages at EVERY point of a healthy lifecycle
(alloc, partial free, group free, flush) — not just at rest.
* The check is not vacuous: drifting any single term (live count, watermark,
a leaked hole) reports loudly, naming the sub-pool.
* Chain order: one member's low frontier clearing the other's high frontier
is what "two pools share one buffer without overlap" MEANS; the pair check
must hold regardless of which member grows up.
* The strict escalation env defaults OFF: promoting the diagnostic to a
RuntimeError is a validation posture, not the production one.
python -m pytest test/registered/unit/mem_cache/test_unified_byte_accounting.py -v
The conservation identity: on a lazy end pool the watermark span equals
live + holes + pending pages at EVERY point of a lifecycle, not just at rest.
"""
import unittest
@@ -62,18 +51,6 @@ def _paged_pair(lazy: bool):
class TestHealthyLifecycleReportsClean(unittest.TestCase):
def test_swa_composite_clean_at_every_step(self):
inst, allocator, kvcache = _swa_composite()
self.assertEqual(allocator.verify_byte_accounting(), [])
v = inst._alloc(allocator, kvcache, 8)
self.assertEqual(allocator.verify_byte_accounting(), [])
allocator.free_swa(v[:4]) # tombstone half the swa side
self.assertEqual(allocator.verify_byte_accounting(), [])
inst._free(allocator, kvcache, v)
self.assertEqual(allocator.verify_byte_accounting(), [])
allocator.clear()
self.assertEqual(allocator.verify_byte_accounting(), [])
def test_lazy_end_pool_clean_through_free_and_flush(self):
full, _swa = _paged_pair(lazy=True)
self.assertEqual(full._byte_accounting_violations(), [])
@@ -86,10 +63,10 @@ class TestHealthyLifecycleReportsClean(unittest.TestCase):
class TestDriftReportsLoudly(unittest.TestCase):
"""Each mutation below models a distinct bookkeeping bug; the verifier
must name the drifted sub-pool. Without these, a regression in any single
counter passes every other test (the pool still 'works' — it just lies
about capacity)."""
"""Each mutation models a distinct bookkeeping bug; the verifier must name
the drifted sub-pool. Without these, a regression in any single counter
passes every other test -- the pool still 'works', it just lies about
capacity."""
def _lazy_full(self):
full, _swa = _paged_pair(lazy=True)
@@ -98,20 +75,27 @@ class TestDriftReportsLoudly(unittest.TestCase):
self.assertEqual(full._byte_accounting_violations(), [])
return full
def test_drifted_live_count(self):
full = self._lazy_full()
full.live_page_count += 1
self.assertTrue(any("span" in s for s in full._byte_accounting_violations()))
def test_any_drifted_term_reports(self):
def drift_live_count(full):
full.live_page_count += 1
def test_leaked_hole(self):
full = self._lazy_full()
full._free_phys_pages = full._free_phys_pages[:-1] # hole vanished
self.assertTrue(any("span" in s for s in full._byte_accounting_violations()))
def leak_hole(full):
full._free_phys_pages = full._free_phys_pages[:-1] # hole vanished
def test_drifted_watermark(self):
full = self._lazy_full()
full.watermark_physical += 1
self.assertTrue(any("span" in s for s in full._byte_accounting_violations()))
def drift_watermark(full):
full.watermark_physical += 1
for term, mutate in (
("live_count", drift_live_count),
("leaked_hole", leak_hole),
("watermark", drift_watermark),
):
with self.subTest(term=term):
full = self._lazy_full()
mutate(full)
self.assertTrue(
any("span" in s for s in full._byte_accounting_violations())
)
def test_composite_report_names_the_sub_pool(self):
"""Frontier-bounds drift (checked in BOTH lazy and eager modes): push
@@ -129,11 +113,10 @@ class TestDriftReportsLoudly(unittest.TestCase):
class TestChainFrontierOrder(unittest.TestCase):
def test_overlapping_frontiers_report(self):
"""Both bands hold pages, then the up member's watermark is pushed past
the down member's LIVE low frontier: the two bands now claim the same
bytes of one buffer. (An empty down band cannot overlap — its low
frontier IS the buffer top — so both sides must be populated for the
scenario to be a real corruption.)"""
"""Both bands hold pages, then the up member's watermark is pushed
past the down member's LIVE low frontier. Both sides must be populated
for this to be real corruption: an empty down band cannot overlap,
its low frontier IS the buffer top."""
full, swa = _paged_pair(lazy=False)
chain = mea._end_pair_chain(full, swa)
up, down = chain
@@ -144,14 +127,6 @@ class TestChainFrontierOrder(unittest.TestCase):
out = mea._chain_byte_accounting_violations(chain)
self.assertTrue(any("overlap" in s for s in out), out)
def test_pair_order_is_direction_agnostic(self):
"""The factories and the unit fixtures orient the pair differently;
the check must order by grow direction, not by argument position."""
full, swa = _paged_pair(lazy=False)
a = mea._end_pair_chain(full, swa)
b = mea._end_pair_chain(swa, full)
self.assertEqual([x.sub_pool_name for x in a], [x.sub_pool_name for x in b])
if __name__ == "__main__":
unittest.main()
test/registered/unit/mem_cache/test_unified_byte_budget_sizing.py
+40 -83
View File
@@ -13,26 +13,17 @@
# ==============================================================================
"""Byte-budget buffer sizing for the unified 2-pool factories.
Derived properties pinned here:
With ``unified_total_bytes`` set the buffer is that many bytes exactly (the
mamba pair adds the state pool's bytes on top -- the budget is captured AFTER
the state carve-out); without it, sizing falls back to the token-count re-sum.
Sizing from the ratio-derived token counts instead would re-introduce the
configurator's rounding, which floors by the cell size and then page-aligns
EACH side, losing up to about a page of tokens per side.
* Budget honored EXACTLY: with ``unified_total_bytes`` set, the swa pair's
buffer is that many bytes (the mamba pair adds the state pool's bytes on
top — the budget is captured AFTER the state carve-out). Sizing from the
ratio-derived token counts instead re-introduces the configurator's
rounding: the swa split floors the budget by the cell size and then
page-aligns EACH side's token count, so the re-sum reconstructs less
than the profiled budget by up to about one page of tokens per side.
* Fallback: without the budget, sizing is the historical token-count re-sum,
bit-for-bit.
* bs=1 feasibility floor: a budget that cannot fit ONE worst-case request
(full KV at max context, plus one SWA window / the state slots a single
running request locks) raises at BOOT, before any pool construction —
under-sizing is a retract LIVELOCK at runtime, not a perf bug.
* The 4096-byte alignment exists because the factories ``.view()`` the whole
uint8 buffer as the KV dtype; an unaligned budget must be floored, never
rounded up (rounding up overcommits profiled memory).
python -m pytest test/registered/unit/mem_cache/test_unified_byte_budget_sizing.py -v
The bs=1 feasibility floor raises at BOOT, before any pool construction: a
budget that cannot fit ONE worst-case request (full KV at max context, plus
one SWA window / the state slots a running request locks) is a retract
LIVELOCK at runtime, not a perf bug.
"""
import unittest
@@ -92,37 +83,32 @@ def _entry_bytes():
class TestBudgetSizing(unittest.TestCase):
def test_swa_factory_honors_the_budget_exactly(self):
"""The buffer IS the budget, including values the token-count re-sum
cannot represent (not a whole-token multiple per side)."""
e = _entry_bytes()
budget = 96 * e + 512 # deliberately NOT a token-count multiple
bundle = _swa_factory(unified_total_bytes=budget)
self.assertEqual(bundle.unified_memory_pool.total_bytes, budget)
for budget in (
96 * e + 512, # deliberately NOT a token-count multiple
(64 + 32) * e + (e - 2), # almost one more entry
):
with self.subTest(budget=budget):
bundle = _swa_factory(unified_total_bytes=budget)
self.assertEqual(bundle.unified_memory_pool.total_bytes, budget)
def test_fallback_is_the_token_count_resum(self):
e = _entry_bytes()
bundle = _swa_factory()
self.assertEqual(bundle.unified_memory_pool.total_bytes, (64 + 32) * e)
def test_budget_beats_resum_on_rounding(self):
"""The property that motivates the whole phase: the re-sum cannot
represent a budget that is not a whole-token multiple per side, so it
strands bytes the buffer could have held."""
e = _entry_bytes()
budget = (64 + 32) * e + (e - 2) # almost one more entry
bundle = _swa_factory(unified_total_bytes=budget)
self.assertEqual(bundle.unified_memory_pool.total_bytes, budget)
self.assertGreater(budget, (64 + 32) * e)
class TestReservedFloorIsOneSourceOfTruth(unittest.TestCase):
"""The bs=1 floor charges the slot-0 sink, and MUST charge exactly what
`UnifiedKVPool` actually reserves.
Regression (GPU eval_434/436, Falcon-H1 boot): the floor hand-copied the
formula as `page_size * max(entry_bytes)`, applying the page multiplier to
the MAMBA spec. The pool deliberately excludes mamba (it is page_size=1),
so with page_size=256 and a ~139 MB state entry the floor over-charged the
sink by 256x — ~33 GiB of phantom requirement — and a healthy config
failed to boot with 25 GiB of real headroom.
Regression (Falcon-H1 boot): the floor hand-copied the formula as
`page_size * max(entry_bytes)`, applying the page multiplier to the MAMBA
spec, which the pool deliberately excludes because it is page_size=1. At
page_size=256 that over-charged the sink 256x and a healthy config failed
to boot with real headroom to spare.
"""
def _specs(self, page_size):
@@ -134,7 +120,7 @@ class TestReservedFloorIsOneSourceOfTruth(unittest.TestCase):
store_dtype=torch.float16,
grow_direction="down",
)
# A state entry vastly larger than a KV token entry — the real ratio
# A state entry vastly larger than a KV token entry: the real ratio
# (~139 MB vs ~45 KB) is what made the over-charge fatal.
mamba = MambaSubPoolSpec(
name="mamba",
@@ -187,51 +173,22 @@ class TestBs1FeasibilityFloor(unittest.TestCase):
self.assertIn("bs=1 floor", str(ctx.exception))
self.assertIn("swa_window_kv", str(ctx.exception))
def test_context_longer_than_the_pool_is_not_rejected(self):
"""REGRESSION: the floor must NOT charge the full-attention token side.
`TpModelWorker.get_worker_info` clamps max_req_len to the pool, so a
context far larger than the buffer is refused at admission, not a
livelock -- and it is an ordinary way to serve a long-context model on
one GPU. Charging it here made such configs fail at boot."""
def test_feasible_floor_inputs_are_not_rejected(self):
"""Floor inputs that must NOT raise: charging the full-attention token
side, or a window the context never clamps, failed these at boot."""
e = _entry_bytes()
bundle = _swa_factory(
unified_total_bytes=200 * e,
model_context_len=1_000_000, # far beyond what the buffer holds
sliding_window_size=16,
)
self.assertEqual(bundle.unified_memory_pool.total_bytes, 200 * e)
def test_feasible_config_boots_with_floor_inputs_present(self):
e = _entry_bytes()
bundle = _swa_factory(
unified_total_bytes=200 * e,
model_context_len=64,
sliding_window_size=16,
)
self.assertEqual(bundle.unified_memory_pool.total_bytes, 200 * e)
def test_window_term_is_clamped_to_context(self):
"""A window larger than the context must charge at most the context —
otherwise short-context models over-raise."""
e = _entry_bytes()
bundle = _swa_factory(
unified_total_bytes=200 * e,
model_context_len=64,
sliding_window_size=10_000, # window >> context
)
self.assertIsNotNone(bundle)
def test_floor_message_itemizes_terms(self):
with self.assertRaises(RuntimeError) as ctx:
_check_bs1_feasibility_floor(
total_bytes=10,
floor_terms=[("a", 8), ("b", 8)],
factory="test",
)
msg = str(ctx.exception)
self.assertIn("a=8", msg)
self.assertIn("b=8", msg)
self.assertIn("16", msg)
for case, model_context_len, sliding_window_size in (
("context_longer_than_the_pool", 1_000_000, 16),
("feasible_config", 64, 16),
("window_larger_than_context", 64, 10_000),
):
with self.subTest(case=case):
bundle = _swa_factory(
unified_total_bytes=200 * e,
model_context_len=model_context_len,
sliding_window_size=sliding_window_size,
)
self.assertEqual(bundle.unified_memory_pool.total_bytes, 200 * e)
def test_exact_floor_passes(self):
"""Boundary: total == floor must NOT raise (>= is the contract)."""
test/registered/unit/mem_cache/test_unified_capacity_memo.py
+19 -71
View File
@@ -11,29 +11,17 @@
# 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).
"""Epoch-memoized capacity views on the allocator chain.
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 per-band and composite capacity views are pure functions of a handful of
CPU-resident fields that schedulers read 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
The guarded failure mode is a memo serving a STALE value after a mutation the
epoch machinery missed -- a new mutation site writing an uncovered field, or
an in-place write that bypasses `__set__`. Readers cannot detect either one;
stale capacity is silent over-/under-admission, not a crash.
"""
import random
@@ -114,10 +102,8 @@ class TestCapacityMemoCoherence(unittest.TestCase):
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."""
"""Property: after ANY mutation sequence the memoized views equal
fresh recomputes. Seeded, so the interleavings are deterministic."""
rng = random.Random(0xC0FFEE)
for lazy in (False, True):
with self.subTest(lazy_compaction=lazy):
@@ -162,29 +148,13 @@ class TestCapacityMemoCoherence(unittest.TestCase):
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."""
"""A hole-free float alloc rebinds NO free-list and has no watermark:
the span fields are its ONLY capacity state, so unless they are
`_CapacityField` descriptors the float's own memo AND both neighbours'
(the span flips transparency, walling off their gaps) keep serving
pre-move values. Exercised on a hand-wired end+float+end chain so no
end-pool descriptor write can mask a missing span bump."""
from test_multi_ended_allocator import TestFloatMultiEndedAllocator
inst = TestFloatMultiEndedAllocator(
@@ -212,21 +182,12 @@ class TestCapacityMemoCoherence(unittest.TestCase):
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."""
invalidate the composite memo -- including the two 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
@@ -281,19 +242,6 @@ class TestTriCapacityMemoCoherence(unittest.TestCase):
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()
test/registered/unit/mem_cache/test_unified_free_no_host_sync.py
+59 -137
View File
@@ -11,28 +11,17 @@
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""The unified free path must not move anything across the PCIe bus.
"""The unified allocator free path must not sync the host.
Two independent host syncs lived in `MultiEndedAllocator`'s free path:
Two syncs are guarded. A tombstone written as ``t[idx] = -1`` makes torch
materialise the scalar as a CPU tensor and copy it H2D, which BLOCKS the host
until the stream drains. And `torch.unique`, recovering distinct PAGE ids from
freed TOKEN ids, has a data-dependent output shape and so must D2H the count;
`free_segment` instead takes stride slices off the caller's `start_pos`, a
page's tokens sitting consecutively in the kv row.
1. Tombstone scatters written as ``t[idx] = -1``. The scalar RHS makes torch
materialise ``-1`` as a CPU tensor and copy it H2D, and a pageable H2D
copy BLOCKS the host until the stream drains. Invisible on decode-shaped
work; ~16 ms per free behind an 8192-token prefill.
2. `torch.unique` recovering distinct PAGE ids from freed TOKEN ids. Its
output shape is data-dependent, so it must D2H the count
(``_unique2 -> item -> _local_scalar_dense -> cudaStreamSynchronize``).
`PagedTokenToKVPoolAllocator` already solved this with `free_segment`:
a page's tokens sit consecutively in the kv row, so given `start_pos` the
page representatives are stride slices. The unified allocators simply
never implemented it and so were permanently on the syncing path.
These tests mirror `test_paged_free_segment.py` -- the same sweep against the
`torch.unique` reference, the same free-group deferral -- because the unified
allocators now mirror that allocator's design rather than a parallel one.
python -m pytest test/registered/unit/mem_cache/test_unified_free_no_host_sync.py -v
Mirrors `test_paged_free_segment.py`, which pins the same properties for
`PagedTokenToKVPoolAllocator`.
"""
import ast
@@ -68,10 +57,9 @@ def _paged_allocator(lazy: bool):
_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.
# Methods that MUST tombstone through index_fill_; hand-listed because "writes
# a tombstone" is a per-method design fact a scan cannot infer. Completeness is
# guarded by `test_every_allocator_free_path_is_listed` below.
_TOMBSTONE_METHODS = [
(mea.MultiEndedAllocator, "_free_lazy"),
(mea.MultiEndedAllocator, "free"),
@@ -82,9 +70,8 @@ _TOMBSTONE_METHODS = [
]
# Ways to write `-1` into a table without the value crossing the bus:
# `index_fill_` takes the scalar as an argument torch keeps off the host, and
# the fused launchers store it from inside the kernel.
# Sanctioned no-sync tombstone forms: `index_fill_` takes the scalar as an
# argument torch keeps off the host; the fused launcher stores it in-kernel.
_NO_SYNC_TOMBSTONE_FORMS = ("index_fill_", "free_unbind_inplace")
@@ -109,9 +96,8 @@ def _allocators_in_module():
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.
Discovery rather than a hand list, so a new allocator class with its own
free path is covered the day it lands.
"""
out = []
for cls in _allocators_in_module():
@@ -130,15 +116,13 @@ def _table_touching_methods():
def _scalar_index_assignments(fn):
"""`self.<table>[<tensor idx>] = <scalar>` occurrences in fn's source.
Slice assignments (``t[a:b] = -1``) are excluded: a slice is a view, so the
fill needs no index tensor. Tensor-valued scatters are excluded too -- only
the scalar RHS materialises a CPU value tensor.
Slices (``t[a:b] = -1``) and tensor-valued scatters are excluded: only a
scalar RHS behind a tensor index materialises a CPU value tensor.
"""
def _is_scalar_literal(node):
# NOTE: `-1` parses as UnaryOp(USub, Constant(1)), NOT Constant. Testing
# only for Constant silently skips every negative literal -- i.e. every
# tombstone this scan exists to find.
# `-1` parses as UnaryOp(USub, Constant(1)), not Constant; matching
# only Constant silently skips every negative literal.
if isinstance(node, ast.Constant):
return True
return isinstance(node, ast.UnaryOp) and isinstance(node.operand, ast.Constant)
@@ -156,9 +140,8 @@ 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()`.
# A constant index (`t[0] = 0`) is a single-element sentinel write,
# not a tensor-index tombstone; only `clear()` does it.
if _is_scalar_literal(tgt.slice):
continue
bad.append(ast.unparse(node))
@@ -167,6 +150,13 @@ def _scalar_index_assignments(fn):
class TestTombstonesDoNotCrossTheBus(unittest.TestCase):
def test_no_scalar_index_assignment(self):
# Self-check first: a scan whose AST matching has drifted reports
# clean on every method below.
def _offender(self):
self.virtual_to_physical[free_v_pages] = -1 # noqa: F821
self.assertEqual(len(_scalar_index_assignments(_offender)), 1)
discovered = _table_touching_methods()
self.assertGreaterEqual(
len(discovered), len(_TOMBSTONE_METHODS), "discovery scan went blind"
@@ -185,24 +175,10 @@ class TestTombstonesDoNotCrossTheBus(unittest.TestCase):
),
)
def test_the_scan_detects_the_scalar_form_it_guards(self):
"""Self-check. The scan is only as good as its AST matching, and it
silently missed every tombstone until `-1` was recognised as
UnaryOp(USub, Constant) rather than Constant. Pin that."""
def _offender(self):
self.virtual_to_physical[free_v_pages] = -1 # noqa: F821
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.
"""
"""Bug regression: an allocator that owns a free path but is missing
from `_TOMBSTONE_METHODS` must fail loudly here rather than drop out of
tombstone coverage."""
listed = {(cls.__name__, name) for cls, name in _TOMBSTONE_METHODS}
for cls in _allocators_in_module():
for name, fn in vars(cls).items():
@@ -212,9 +188,8 @@ class TestTombstonesDoNotCrossTheBus(unittest.TestCase):
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.
# Only a method that WRITES a page table needs a tombstone;
# one that merely READS has nothing to guard.
if not (
any(f"{t}.index_fill_" in src for t in _TABLES)
or _scalar_index_assignments(fn)
@@ -231,10 +206,8 @@ class TestTombstonesDoNotCrossTheBus(unittest.TestCase):
)
def test_free_paths_actually_write_a_tombstone(self):
"""Positive form, so deleting the scatter entirely cannot pass. The
mechanism is not the point -- keeping the tombstone value off the host
is -- so this lists the sanctioned ways to do that and a new one is
added here deliberately."""
"""Positive form, so deleting the scatter entirely cannot pass; a new
no-sync mechanism is added to `_NO_SYNC_TOMBSTONE_FORMS` deliberately."""
for cls, name in _TOMBSTONE_METHODS:
with self.subTest(method=f"{cls.__name__}.{name}"):
src = inspect.getsource(getattr(cls, name))
@@ -244,22 +217,6 @@ class TestTombstonesDoNotCrossTheBus(unittest.TestCase):
f"{_NO_SYNC_TOMBSTONE_FORMS}",
)
def test_index_fill_matches_scalar_assign_semantics(self):
"""Behaviour-preserving, including the edge cases the free path hands
it: empty index, duplicate pages, full table."""
for idx in (
torch.tensor([], dtype=torch.int64),
torch.tensor([1, 3, 5], dtype=torch.int64),
torch.tensor([2, 2, 3], dtype=torch.int64), # duplicates
torch.arange(6, dtype=torch.int64),
):
with self.subTest(n=int(idx.numel())):
a = torch.arange(6, dtype=torch.int64)
b = a.clone()
a[idx] = -1
b.index_fill_(0, idx, -1)
self.assertTrue(torch.equal(a, b))
# --------------------------------------------------------------------------
# 2. free_segment: stride page extraction instead of torch.unique
@@ -270,7 +227,6 @@ class TestFreeSegment(unittest.TestCase):
"""Mirrors `test_paged_free_segment.TestFreeSegment`."""
def test_matches_unique_over_tail_alignments(self):
"""Page-aligned starts against every tail alignment."""
for num_tokens in (1, PAGE_SIZE, PAGE_SIZE + 1, 3 * PAGE_SIZE - 1):
for start in range(0, num_tokens, PAGE_SIZE):
for end in (start + 1, num_tokens):
@@ -288,8 +244,6 @@ class TestFreeSegment(unittest.TestCase):
self.assertEqual(freed.numel(), expected.numel())
def test_never_calls_unique(self):
"""The decisive check -- make `torch.unique` explode. A textual guard
can be fooled; this cannot."""
for start in (0, PAGE_SIZE, 2 * PAGE_SIZE):
alloc = _paged_allocator(lazy=True)
row = alloc.alloc(3 * PAGE_SIZE)
@@ -299,12 +253,6 @@ class TestFreeSegment(unittest.TestCase):
):
alloc.free_segment(row[start : start + PAGE_SIZE], start_pos=start)
def test_unaligned_start_is_rejected(self):
alloc = _paged_allocator(lazy=True)
row = alloc.alloc(3 * PAGE_SIZE)
with self.assertRaises(AssertionError):
alloc.free_segment(row[1 : PAGE_SIZE + 1], start_pos=1)
def test_empty_segment_is_noop(self):
alloc = _paged_allocator(lazy=True)
before = alloc._free_phys_pages.numel()
@@ -312,8 +260,8 @@ class TestFreeSegment(unittest.TestCase):
self.assertEqual(alloc._free_phys_pages.numel(), before)
def test_page_size_one_takes_the_plain_path(self):
"""token == page: nothing to dedup, so `free_segment` must not invent
a stride slice that would drop tokens."""
"""token == page: a stride slice would drop tokens, so take the plain
path."""
alloc = _paged_allocator(lazy=True)
alloc.page_size = 1
v = alloc.alloc(PAGE_SIZE)
@@ -325,12 +273,9 @@ class TestFreeSegment(unittest.TestCase):
class TestFreeGroupKeepsPositions(unittest.TestCase):
"""Mirrors `test_paged_free_segment.test_group_defers_until_group_end`.
Bug regression: buffering RAW tokens and `torch.cat`-ing them at
`free_group_end` destroys each segment's shape, so the merged tensor has no
recoverable page structure and falls back to `torch.unique`. Measured as 71
of 77 `_free_lazy` calls still syncing on gpt-oss and Qwen3.5 ps=256
(eval_429), all attributed to `free_group_end` via the decode path. The fix
buffers page REPRESENTATIVES, so the merge concatenates page ids.
Bug regression: a free group must buffer page REPRESENTATIVES, not raw
tokens -- concatenating raw tokens loses the page structure and sends
`free_group_end` back to `torch.unique`.
"""
def test_group_defers_until_group_end(self):
@@ -346,8 +291,6 @@ class TestFreeGroupKeepsPositions(unittest.TestCase):
self.assertEqual(alloc._free_phys_pages.numel(), before + 2)
def test_group_end_does_not_sync(self):
"""The property the whole fix exists for: a grouped segment free must
complete with `torch.unique` disabled."""
alloc = _paged_allocator(lazy=True)
row = alloc.alloc(3 * PAGE_SIZE)
alloc.free_group_begin()
@@ -360,8 +303,8 @@ class TestFreeGroupKeepsPositions(unittest.TestCase):
self.assertGreater(alloc._free_phys_pages.numel(), 0)
def test_positionless_group_still_uses_the_unique_path(self):
"""Plain `free()` inside a group has no position to keep, so it must
still go through the (syncing) dedup -- correctness over speed."""
"""Plain `free()` carries no position to keep, so it must still take
the syncing dedup -- correctness over speed."""
alloc = _paged_allocator(lazy=True)
row = alloc.alloc(2 * PAGE_SIZE)
alloc.free_group_begin()
@@ -374,12 +317,9 @@ class TestFreeGroupKeepsPositions(unittest.TestCase):
class TestEveryUnifiedAllocatorOverridesFreeSegment(unittest.TestCase):
"""Completeness guard. The base `free_segment` DISCARDS `start_pos` and
"""Completeness guard: the base `free_segment` DISCARDS `start_pos` and
calls plain `free`, so an allocator that inherits it sends every segment
free into the syncing dedup -- silently, with no error and no wrong answer,
just a stalled scheduler thread. That is exactly what happened: the SWA
composite was overridden and the Mamba composite was not, and 77 of 77
`_free_lazy` calls on Qwen3.5 ps=256 still synced (eval_428).
free into the syncing dedup -- silently, with no error and no wrong answer.
"""
def test_all_overridden(self):
@@ -400,8 +340,8 @@ class TestEveryUnifiedAllocatorOverridesFreeSegment(unittest.TestCase):
)
def test_composites_buffer_reps_not_tokens_in_a_group(self):
"""The group buffer must exist on every allocator that can receive a
segment free, or `free_segment` raises inside a group."""
"""Every allocator that can receive a segment free needs the group
buffer, or `free_segment` raises inside a group."""
for cls in (
mea.MultiEndedAllocator,
unified_mamba.UnifiedMambaTokenToKVPoolAllocator,
@@ -432,13 +372,10 @@ class TestUnifiedSwaFullSideGroup(unittest.TestCase):
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).
"""The per-decode-step SWA window ratchet frees a CONTIGUOUS row slice with
host-int, page-aligned bounds, so `free_swa(..., start_pos=)` must reach the
swa side with caller-derived page ids: no `torch.unique` and no stale-slot
`.item()` on the per-step path.
"""
PS = 4
@@ -480,8 +417,7 @@ class TestFreeSwaWindowRatchetNoHostSync(unittest.TestCase):
)
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."""
"""Aligned bounds, which the ratchet guarantees at ps > 1."""
alloc = self._swa_composite(lazy=True)
v = alloc.alloc(8 * self.PS)
self.assertIsNotNone(v)
@@ -497,8 +433,8 @@ class TestFreeSwaWindowRatchetNoHostSync(unittest.TestCase):
alloc.free_swa(v[4 * self.PS :], start_pos=4 * self.PS)
def test_full_only_segment_free_never_syncs(self):
"""The request-finish dead half (swa already tombstoned) frees the full
side by page reps: no unique from free_full's token dedup."""
"""Request-finish shape: the swa side is already tombstoned, so the
full side must free by page reps rather than `free_full`'s dedup."""
alloc = self._swa_composite(lazy=True)
v = alloc.alloc(8 * self.PS)
alloc.free_swa(v, start_pos=0)
@@ -523,8 +459,8 @@ class TestFreeSwaWindowRatchetNoHostSync(unittest.TestCase):
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)."""
"""Derived property: the stride-rep path and the dedup fallback leave
identical v2p tombstones and capacity."""
for lazy in (True, False):
with self.subTest(lazy=lazy):
a1 = self._swa_composite(lazy=lazy)
@@ -559,10 +495,9 @@ class TestFreeSwaWindowRatchetNoHostSync(unittest.TestCase):
torch.cuda.is_available(), "the fused tombstone is a Triton kernel"
)
class TestFusedTombstoneWritesBothTables(unittest.TestCase):
"""The source scan above accepts `free_unbind_inplace` as a no-sync
mechanism; this is what makes that acceptance mean something. On CPU the
launcher takes its pure-torch reference path, so nothing else in the suite
ever runs the kernel that does the tombstoning.
"""The source scan accepts `free_unbind_inplace` as a no-sync mechanism; on
CPU the launcher takes its pure-torch reference path, so nothing else in
the suite ever runs the kernel that does the tombstoning.
"""
def test_matches_the_reference_over_randomized_bindings(self):
@@ -602,19 +537,6 @@ class TestFusedTombstoneWritesBothTables(unittest.TestCase):
f"a live binding was disturbed, trial {trial}",
)
def test_cuda_agrees_with_the_cpu_reference(self):
from sglang.kernels.ops.memory.virtual_slot import free_unbind_inplace
v = torch.tensor([3, 0, 5], dtype=torch.int64)
cpu_v2p = torch.tensor([1, 2, 3, 4, 5, 0], dtype=torch.int64)
cpu_p2v = torch.tensor([5, 0, 1, 2, 3, 4], dtype=torch.int64)
cu_v2p, cu_p2v = cpu_v2p.cuda(), cpu_p2v.cuda()
cpu_out = free_unbind_inplace(v, cpu_v2p, cpu_p2v)
cu_out = free_unbind_inplace(v.cuda(), cu_v2p, cu_p2v)
self.assertTrue(torch.equal(cpu_out, cu_out.cpu()))
self.assertTrue(torch.equal(cpu_v2p, cu_v2p.cpu()))
self.assertTrue(torch.equal(cpu_p2v, cu_p2v.cpu()))
def test_empty_free_is_a_noop(self):
from sglang.kernels.ops.memory.virtual_slot import free_unbind_inplace
test/registered/unit/mem_cache/test_unified_mamba_views.py
+103 -231
View File
@@ -11,49 +11,27 @@
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Round-trip correctness of ``UnifiedKVPool._build_mamba_views`` — the
"""Round-trip correctness of ``UnifiedKVPool._build_mamba_views`` -- the
envelope-strided conv/temporal (SSM) state views that back ``UnifiedMambaPool``.
This isolates the unified-memory-pool Mamba STATE layout from the full model. It guards
against a class of correctness defect where Falcon-H1 greedy decode is garbled
under the unified memory pool, isolated to the Mamba conv/temporal state path: a
stride/offset/alignment bug in the view construction (analogous to the fixed
`_extract_kv_strides` MHA bug).
Guards against garbled greedy decode from a stride/offset/alignment bug in the
state views. CPU-only.
Within one slot's envelope the bytes are
``[conv[0]·L0 | conv[0]·L1 | ... | conv[1]·L0 | ... | temporal·L0 | ...]`` and
across slots the layout is envelope (slot stride == entry_bytes). Each returned
view is ``(num_layers, max_slots, *inner_shape)``. The conv dtype (bf16, 2 B)
and temporal dtype (fp32, 4 B) DIFFER, so the temporal view's byte offset must
be a multiple of the temporal itemsize — an alignment hazard that
``_build_mamba_views`` now asserts.
These tests prove the views:
- round-trip every (tensor, layer, slot) element with the Falcon-like
bf16-conv / fp32-temporal dtype mix (catches stride/offset/alignment bugs);
- do NOT alias each other (conv[i] vs conv[j] vs temporal) or across
layers/slots (catches envelope-overlap);
- match a contiguous ``(num_layers, max_slots, *inner)`` reference exactly
(the shape `MambaPool.State.conv[i]` / `.temporal` expose);
- reject a deliberately mis-aligned spec via the alignment assert.
The round-trip class is skipped on CPU — those views back GPU kernels and we
mirror the GPU path. ``TestKDAFlashInferEnvelopeStateContract`` is pure stride
arithmetic and runs everywhere.
python -m pytest test/registered/unit/mem_cache/test_shared_mamba_views.py -v
``[conv[0] L0 | conv[0] L1 | ... | conv[1] L0 | ... | temporal L0 | ...]``,
the slot stride is ``entry_bytes``, and each returned view is
``(num_layers, max_slots, *inner_shape)``. The conv dtype (bf16, 2 B) and
temporal dtype (fp32, 4 B) may DIFFER, so the temporal view's byte offset must
be a multiple of the temporal itemsize -- an alignment hazard that
``_build_mamba_views`` asserts.
"""
import unittest
import torch
from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
from sglang.test.ci.ci_register import register_amd_ci, register_cpu_ci
_HAS_CUDA = torch.cuda.is_available()
_DEV = "cuda" if _HAS_CUDA else "cpu"
register_cuda_ci(est_time=30, stage="base-b", runner_config="1-gpu-small")
register_cpu_ci(est_time=30, suite="base-a-test-cpu")
register_amd_ci(est_time=30, stage="stage-b", runner_config="1-gpu-small-amd")
@@ -65,11 +43,10 @@ def _make_pool(
temporal_state_shape,
temporal_dtype,
want_slots=8,
device=_DEV,
device="cpu",
):
"""Build a minimal 2-sub-pool ``UnifiedKVPool`` (a small MHA grow-up peer
+ the Mamba grow-down pool under test) sized to hold >= ``want_slots`` Mamba
slots, and return ``(pool, mamba_spec)``."""
"""Minimal 2-sub-pool ``UnifiedKVPool`` sized to hold >= ``want_slots``
Mamba slots; returns ``(pool, mamba_spec)``."""
from sglang.srt.mem_cache.unified_memory_pool import (
MambaSubPoolSpec,
MHASubPoolSpec,
@@ -97,9 +74,8 @@ def _make_pool(
entry_mamba = mamba_spec.entry_bytes()
entry_full = full_spec.entry_bytes()
entry_max = max(entry_mamba, entry_full)
# Need max_slots("mamba") = total // entry_mamba >= want_slots, and total
# large enough that BOTH pools clear their min_slot_index. Add generous
# headroom, then round up to a multiple of 8 (covers bf16/fp32 .view()).
# Headroom so BOTH pools clear their min_slot_index, then round up to a
# multiple of 8 so bf16/fp32 ``.view()`` stays legal.
total_bytes = want_slots * entry_mamba + 8 * entry_max
total_bytes = ((total_bytes + 7) // 8) * 8
pool = UnifiedKVPool(
@@ -111,11 +87,8 @@ def _make_pool(
return pool, mamba_spec
@unittest.skipUnless(_HAS_CUDA, "shared Mamba views back GPU kernels")
class TestUnifiedMambaViews(unittest.TestCase):
# Falcon-H1-like dims: even conv_dim, bf16 conv, fp32 temporal, several
# layers. (Mamba2 conv state is (conv_dim, kernel-1); temporal/SSM state is
# (nheads, head_dim, ssm_state_size).)
# Falcon-H1-like dims: even conv_dim, bf16 conv, fp32 temporal.
FALCON_KW = dict(
mamba_layer_num=5, # odd, to stress the temporal-offset alignment
conv_state_shapes=[(48, 3)], # conv_dim=48, kernel-1=3
@@ -125,11 +98,9 @@ class TestUnifiedMambaViews(unittest.TestCase):
)
def _fill_and_roundtrip(self, pool, mamba_spec):
"""Write a distinct random tensor to each conv view + the temporal view
(in their own dtypes), then read all back and assert exact equality.
Writing ALL views first and reading ALL after means any envelope overlap
(conv[i]/conv[j]/temporal aliasing) corrupts an earlier write → mismatch.
"""
"""Write a distinct random tensor to every view, then read all of them
back: writing all first makes any envelope overlap corrupt an earlier
write."""
conv_views, temporal_view = pool.mamba_views_for("mamba")
torch.manual_seed(0)
refs = []
@@ -155,34 +126,48 @@ class TestUnifiedMambaViews(unittest.TestCase):
f"stride={temporal_view.stride()}",
)
def test_roundtrip_falcon_like(self):
pool, spec = _make_pool(**self.FALCON_KW)
self._fill_and_roundtrip(pool, spec)
def test_roundtrip_single_layer_single_slot_edges(self):
OTHER_GEOMETRIES = {
# 1 layer, multiple conv tensors, same-dtype conv/temporal.
pool, spec = _make_pool(
"single_layer_single_slot_edges": dict(
mamba_layer_num=1,
conv_state_shapes=[(16, 3), (8, 3)],
conv_dtype=torch.float32,
temporal_state_shape=(4, 8, 16),
temporal_dtype=torch.float32,
want_slots=4,
)
self._fill_and_roundtrip(pool, spec)
def test_roundtrip_multi_conv_tensors(self):
# Two conv tensors + bf16/fp32 mix — exercises the per-conv-tensor offset
),
# Two conv tensors + bf16/fp32 mix: exercises the per-conv-tensor offset
# accumulation in _build_mamba_views.
pool, spec = _make_pool(
"multi_conv_tensors": dict(
mamba_layer_num=3,
conv_state_shapes=[(32, 3), (16, 3)],
conv_dtype=torch.bfloat16,
temporal_state_shape=(8, 8, 16),
temporal_dtype=torch.float32,
want_slots=6,
)
self._fill_and_roundtrip(pool, spec)
),
# ALIGNED spec: conv region = 2 layers * 2*3*2 B = 24 B, a multiple of
# the fp32 temporal itemsize, so it must build and round-trip.
"alignment_ok": dict(
mamba_layer_num=2,
conv_state_shapes=[(2, 3)],
conv_dtype=torch.bfloat16,
temporal_state_shape=(2, 4, 4),
temporal_dtype=torch.float32,
want_slots=4,
),
}
def test_roundtrip_geometries(self):
"""Round-trip every (tensor, layer, slot) element across the geometry
shapes the view builder has to handle."""
for name, kwargs in (
("falcon_like", self.FALCON_KW),
*self.OTHER_GEOMETRIES.items(),
):
with self.subTest(geometry=name):
pool, spec = _make_pool(**kwargs)
self._fill_and_roundtrip(pool, spec)
def test_no_cross_region_overlap(self):
"""Zero buffer; write a sentinel to ONE view; every OTHER view must read
@@ -207,62 +192,10 @@ class TestUnifiedMambaViews(unittest.TestCase):
f"(envelope regions overlap)",
)
def test_per_layer_per_slot_addressing(self):
"""Distinct value per (layer, slot) on the temporal view; verify exact
addressing (no layer/slot aliasing). Uses small integers exactly
representable in the view dtype.
NB: ``temporal_view`` is a non-contiguous strided view, so we must NOT
``.reshape()`` it (that would COPY, breaking the alias) — we
broadcast-assign into the view in place and read back via basic
indexing (which keeps the view)."""
pool, spec = _make_pool(**self.FALCON_KW)
_, temporal_view = pool.mamba_views_for("mamba")
N, S = temporal_view.shape[0], temporal_view.shape[1]
inner_ndim = temporal_view.dim() - 2
# value = layer*S + slot (< N*S, small → exact in fp32)
base = (
torch.arange(N, device=temporal_view.device)[:, None] * S
+ torch.arange(S, device=temporal_view.device)[None, :]
).to(temporal_view.dtype)
# Broadcast (N, S) over the inner dims, in place into the strided view.
temporal_view[:] = base.view(N, S, *([1] * inner_ndim))
# Read back the first inner element of every (layer, slot) via basic
# indexing (stays a view).
readback = temporal_view[(slice(None), slice(None)) + (0,) * inner_ndim]
self.assertTrue(
torch.equal(readback, base),
"temporal (layer, slot) addressing wrong — layer/slot stride bug",
)
def test_matches_contiguous_reference(self):
"""The shared view must be a faithful relabeling of a contiguous
``(num_layers, max_slots, *inner)`` tensor: identical data written by the
same logical index reads back identically."""
pool, spec = _make_pool(**self.FALCON_KW)
conv_views, temporal_view = pool.mamba_views_for("mamba")
for v in conv_views + [temporal_view]:
ref = torch.randn(v.shape, device=v.device).to(v.dtype)
contig = ref.clone().contiguous()
v.copy_(ref)
self.assertEqual(tuple(v.shape), tuple(contig.shape))
self.assertTrue(
torch.equal(v.contiguous(), contig),
"shared view not equivalent to its contiguous counterpart",
)
def test_alignment_guard_fires_on_misaligned_spec(self):
"""A spec whose conv region (bf16) is an odd multiple of 2 B makes the
per-slot entry (= conv_region + N*temporal_row = 2 B + 4 B = 6 B) NOT a
multiple of the temporal itemsize (fp32, 4 B). The temporal/SSM-state
view's storage_offset is computed by integer-dividing a byte offset by
the temporal itemsize, so this would silently mis-offset the view.
``_build_mamba_views`` must reject it with a loud alignment assert.
NOTE: the ``entry_bytes % itemsize`` guard is what fires here, and it
subsumes the conv-region offset check (see the comment in
``_build_mamba_views``). We assert on the shared "misaligned" wording
rather than on which specific guard trips."""
"""A per-slot entry that is not a multiple of the temporal itemsize would
mis-offset the view and must trip the alignment assert; either guard can
fire first, so match on the shared "misaligned" wording."""
with self.assertRaises(AssertionError) as cm:
_make_pool(
mamba_layer_num=1, # entry = 2 B conv + 4 B temporal = 6 B, not %4
@@ -274,30 +207,12 @@ class TestUnifiedMambaViews(unittest.TestCase):
)
self.assertIn("misalign", str(cm.exception).lower())
def test_alignment_ok_for_aligned_spec(self):
"""An aligned spec (conv region a multiple of the temporal itemsize)
must build and round-trip cleanly."""
# conv region = N * conv_dim*(k-1) * 2 ; with conv_dim=2 -> per-layer 2*3*2=12,
# times N=2 = 24, divisible by 4. Aligned.
pool, spec = _make_pool(
mamba_layer_num=2,
conv_state_shapes=[(2, 3)],
conv_dtype=torch.bfloat16,
temporal_state_shape=(2, 4, 4),
temporal_dtype=torch.float32,
want_slots=4,
)
self._fill_and_roundtrip(pool, spec)
def _k3_kda_mamba_geometry(heads_per_rank: int) -> dict:
"""Kimi K3 KDA per-rank state geometry: 69 KDA layers, K = V = 128,
conv width 4 (=> 3 cached tokens), conv row ``(kernel-1, q+k+v dim)``
in the KimiLinear layout (``KimiLinearStateShape.create`` with
num_k_heads == num_heads, head_k_dim == head_dim — see
``models/kimi_linear.py``), temporal/SSM state ``(HV, V, K)``.
``heads_per_rank`` = 96 total KDA heads / attn_tp (12 at the TP8
deployment shape, cf. ``kernels/ops/attention/kda_fused_decode.py``)."""
"""Kimi K3 KDA per-rank state geometry: 69 KDA layers, K = V = 128, conv
width 4 (=> 3 cached tokens), conv row ``(kernel-1, q+k+v dim)`` per
``KimiLinearStateShape.create``, temporal/SSM state ``(HV, V, K)``, and
``heads_per_rank`` = 96 total KDA heads / attn_tp."""
h = heads_per_rank
return dict(
layer_num=69,
@@ -311,46 +226,35 @@ def _k3_kda_mamba_geometry(heads_per_rank: int) -> dict:
class TestKDAFlashInferEnvelopeStateContract(unittest.TestCase):
"""Derived property: the envelope-strided KDA temporal view (unified memory
/ page-major layout) must satisfy the state contract of FlashInfer
``recurrent_kda``, because the KDA flashinfer decode wrapper
(``linear/kernels/kda_flashinfer.py``) passes the committed per-layer pool
view straight into the kernel (in-place state update on the cu_seqlens path
— no gather/scatter copy around the call).
"""Derived property: the envelope-strided KDA temporal view must satisfy the
state contract of FlashInfer ``recurrent_kda``, because the KDA flashinfer
decode wrapper (``linear/kernels/kda_flashinfer.py``) passes the per-layer
pool view straight into the kernel, with no gather/scatter copy.
The kernel compiles its state argument as a CuTe fake tensor of shape
``[N, HV, V, K]`` with stride ``(sym_int64(divisibility=16), V*K, K, 1)``
and ``assumed_align=32`` (flashinfer ``kda_kernels/recurrent_kda.py``), so
a per-layer pool view is only readable by the kernel when:
and ``assumed_align=32``, so a per-layer pool view is readable only when:
* its inner strides are exactly compact ``(V*K, K, 1)``;
* its slot stride — the per-slot envelope pitch, NOT ``HV*V*K`` — is a
* its slot stride -- the per-slot envelope pitch, NOT ``HV*V*K`` -- is a
multiple of 16 elements (32 bytes at bf16);
* its base byte offset is 32-byte aligned (for every layer).
Any envelope-layout change that breaks one of these (per-slot padding that
is not a 32 B multiple, a conv-shape change misaligning the temporal
region, a transposed/padded temporal inner layout) would silently
mis-address every KDA state read/write on SM100 flashinfer decode; this
test turns such a diff red without a GPU.
"""
# 32 B: recurrent_kda's assumed_align AND its slot-stride divisibility
# (16 elements * 2 B bf16). External-source literal from flashinfer
# kda_kernels/recurrent_kda.py (S_batch = cute.sym_int64(divisibility=16),
# make_fake_tensor(..., assumed_align=32)).
# recurrent_kda's assumed_align AND its slot-stride divisibility (16
# elements * 2 B bf16), from flashinfer kda_kernels/recurrent_kda.py.
_KERNEL_ALIGN_BYTES = 32
@staticmethod
def _build_tp8_views():
"""Real TP8 K3 KDA envelope views on CPU (2 slots suffice — the
per-slot geometry is slot-count independent)."""
def _build_kda_views(heads_per_rank: int = 12):
"""Real K3 KDA envelope views on CPU for one attn-TP shard; 2 slots
suffice, the per-slot geometry is slot-count independent."""
from sglang.srt.mem_cache.layout.page_major import (
build_page_major_mamba_views,
mamba_entry_bytes,
)
geom = _k3_kda_mamba_geometry(12) # 96 heads / TP8
geom = _k3_kda_mamba_geometry(heads_per_rank)
entry_bytes = mamba_entry_bytes(**geom)
max_slots = 2
raw = torch.empty(max_slots * entry_bytes, dtype=torch.uint8, device="cpu")
@@ -360,76 +264,44 @@ class TestKDAFlashInferEnvelopeStateContract(unittest.TestCase):
return geom, entry_bytes, temporal_view
def test_k3_tp8_envelope_view_matches_recurrent_kda_contract(self):
"""Check every per-layer temporal view against the kernel contract."""
geom, entry_bytes, temporal_view = self._build_tp8_views()
itemsize = temporal_view.element_size()
_, v, k = geom["temporal_state_shape"]
for layer in (0, geom["layer_num"] - 1):
view = temporal_view[layer] # [slots, HV, V, K], what decode() gets
self.assertEqual(
view.stride()[1:],
(v * k, k, 1),
"temporal inner strides must stay compact (V*K, K, 1): "
"recurrent_kda compiles them as constants",
)
self.assertEqual(
view.stride(0),
entry_bytes // itemsize,
"slot stride must be the envelope pitch (entry_bytes)",
)
self.assertEqual(
view.stride(0) % (self._KERNEL_ALIGN_BYTES // itemsize),
0,
"slot stride must satisfy recurrent_kda's "
"sym_int64(divisibility=16) — 16 elements = 32 B at bf16",
)
self.assertEqual(
(view.storage_offset() * itemsize) % self._KERNEL_ALIGN_BYTES,
0,
f"layer {layer} temporal view base is not 32 B aligned "
"(recurrent_kda assumed_align=32)",
)
def test_k3_entry_and_temporal_offset_32B_multiples_across_tp(self):
"""The two byte quantities that feed the contract above — the per-slot
envelope pitch and the temporal region's offset inside the envelope
(= all-layers conv region, temporal comes last) — must be 32 B
multiples for every plausible attn-TP shard of K3's 96 KDA heads."""
import math
from sglang.srt.mem_cache.layout.page_major import mamba_entry_bytes
"""Check every per-layer temporal view against the kernel contract, for
every plausible attn-TP shard of K3's 96 KDA heads."""
for heads_per_rank in (96, 48, 24, 12): # attn_tp 1 / 2 / 4 / 8
geom = _k3_kda_mamba_geometry(heads_per_rank)
entry_bytes = mamba_entry_bytes(**geom)
conv_region_bytes = (
geom["layer_num"]
* math.prod(geom["conv_state_shapes"][0])
* geom["conv_dtype"].itemsize
)
self.assertEqual(
entry_bytes % self._KERNEL_ALIGN_BYTES,
0,
f"tp shard h={heads_per_rank}: envelope pitch {entry_bytes} B "
"breaks recurrent_kda's slot-stride divisibility",
)
self.assertEqual(
conv_region_bytes % self._KERNEL_ALIGN_BYTES,
0,
f"tp shard h={heads_per_rank}: temporal region offset "
f"{conv_region_bytes} B breaks assumed_align=32",
)
with self.subTest(heads_per_rank=heads_per_rank):
geom, entry_bytes, temporal_view = self._build_kda_views(heads_per_rank)
itemsize = temporal_view.element_size()
_, v, k = geom["temporal_state_shape"]
for layer in (0, geom["layer_num"] - 1):
# [slots, HV, V, K], what decode() gets
view = temporal_view[layer]
self.assertEqual(
view.stride()[1:],
(v * k, k, 1),
"temporal inner strides must stay compact (V*K, K, 1): "
"recurrent_kda compiles them as constants",
)
self.assertEqual(
view.stride(0),
entry_bytes // itemsize,
"slot stride must be the envelope pitch (entry_bytes)",
)
self.assertEqual(
view.stride(0) % (self._KERNEL_ALIGN_BYTES // itemsize),
0,
"slot stride must satisfy recurrent_kda's "
"sym_int64(divisibility=16) — 16 elements = 32 B at bf16",
)
self.assertEqual(
(view.storage_offset() * itemsize) % self._KERNEL_ALIGN_BYTES,
0,
f"layer {layer} temporal view base is not 32 B aligned "
"(recurrent_kda assumed_align=32)",
)
def test_wrapper_state_contract_check_matches_layout(self):
"""The KDA flashinfer decode wrapper enforces this same contract at
runtime (``FlashInferKDAKernel._check_state_stride_contract``, called
once per pool view before handing the pool to ``recurrent_kda``). A
regression in that check would only surface on SM100 hardware, so pin
its accept/reject behavior here: it must ACCEPT exactly what the
layouts produce — the envelope-strided per-layer view and a plain
contiguous pool — and REJECT views the kernel would silently
mis-address (wrong inner strides; a slot stride off the divisibility)."""
"""``FlashInferKDAKernel._check_state_stride_contract`` enforces the same
contract at runtime, and a regression there would only surface on SM100
hardware, so pin its accept/reject behavior on CPU."""
import types
from sglang.srt.layers.attention.linear.kernels.kda_flashinfer import (
@@ -442,7 +314,7 @@ class TestKDAFlashInferEnvelopeStateContract(unittest.TestCase):
# Fresh stub per call: the real kernel caches approvals by id().
check(types.SimpleNamespace(_state_contract_ok=set()), view)
_, _, temporal_view = self._build_tp8_views()
_, _, temporal_view = self._build_kda_views()
envelope = temporal_view[0] # what forward_decode hands to the kernel
run(envelope) # must not raise
test/registered/unit/mem_cache/test_unified_mha_views.py
+25 -146
View File
@@ -11,26 +11,14 @@
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""MHA K/V views for the unified memory pool (uniform-row hybrid models).
"""MHA K/V views for the unified memory pool (uniform-row hybrid models), CPU-only.
Covers, CPU-only (pure torch — no GPU / Triton kernels):
- `build_mha_views` refuses an asymmetric-KV spec: its addressing
assumes one uniform row width, so it is the boundary that checks;
- `build_mha_views` addressing: view_l[kernel_id(t)] must land exactly at
the page-major envelope byte offset the STRIDED builder assigns to the same
(page, slot, layer, K|V) cell — the two builders are views over one truth;
- K and V of one token share ONE kernel-facing id (per-layer origin shift does the
disambiguation), with no aliasing across the 2*L overlapping views;
- the missing-tail-pad and asymmetric-dims cases fail loud at construction.
Addressing law under test (the derived property everything else builds on):
Addressing law under test:
kernel_id(t) = (t // ps) * (ps * 2L) + t % ps
K of layer l at block 2l, V at block 2l+1, blocks are ps rows of
head_num*head_dim elements — offsets identical to
K of layer l at block 2l, V at block 2l+1; blocks are ps rows of
head_num*head_dim elements, at offsets identical to
MHASubPoolSpec.layer_k/v_offset_in_page when rows are uniform.
python -m pytest test/registered/unit/mem_cache/test_unified_mha_views.py -v
"""
from sglang.test.ci.ci_register import register_cpu_ci
@@ -43,10 +31,7 @@ from types import SimpleNamespace
import torch
from sglang.srt.environ import envs
from sglang.srt.mem_cache.layout.page_major import (
build_mha_views,
mha_entry_bytes,
)
from sglang.srt.mem_cache.layout.page_major import build_mha_views
from sglang.srt.mem_cache.unified_memory_pool import (
MHASubPoolSpec,
UnifiedKVPool,
@@ -55,12 +40,10 @@ from sglang.srt.mem_cache.unified_memory_pool import (
_DEV = "cpu"
# `set_kv_buffer` dispatches on the PLATFORM (memory_pool._is_cuda, resolved at
# import), not on the tensors it is handed, so cases driving it must build on
# the platform's device. The rest of this file is byte arithmetic, so CPU.
# import), not on the tensors it is handed, so cases driving it build there.
_STORE_DEV = "cuda" if torch.cuda.is_available() else "cpu"
# Small-but-nontrivial MHA geometry: L=2 layers, H=2 heads, D=4, so every byte
# offset is hand-checkable. blocks = 2L = 4 per page.
# Geometry kept tiny so every byte offset is hand-checkable.
_L = 2
_H = 2
_D = 4
@@ -108,13 +91,8 @@ def _build_views(raw, ps, num_pages, head_dim=_D, v_head_dim=_D, layer_num=_L):
def _reference_strided_views(raw, *, page_size, num_pages, anchor_bytes=0):
"""Independent 4-D strided description of the page-major envelope.
This is the retired production strided builder, kept here as the oracle:
per-layer ``(num_pages, page_size, head_num, head_dim)`` views addressed by
``(page, slot)``, so the builder's addressing can be cross-checked
against a second, independently-derived description of the same bytes.
"""
"""Independent 4-D strided description of the same page-major envelope,
addressed by ``(page, slot)`` -- the oracle for the view builder."""
k_row_bytes = _ROW * _ITEM
v_row_bytes = _ROW * _ITEM
page_bytes = page_size * _L * (k_row_bytes + v_row_bytes)
@@ -147,11 +125,8 @@ def _reference_strided_views(raw, *, page_size, num_pages, anchor_bytes=0):
class TestMHASpecSurface(unittest.TestCase):
def test_asymmetric_rows_refused_by_the_view_builder(self):
"""The row-block array exists only for uniform rows, so the builder
whose addressing depends on it is the one that refuses (the MiMoV2
shape, scaled down). ServerArgs screens such models out of
--enable-unified-memory long before we get here; this is the check for
a caller that reaches the builder directly."""
"""ServerArgs screens asymmetric-KV models out of
--enable-unified-memory; this guards a caller reaching the builder."""
spec = _mha_spec()
raw = torch.zeros(1 << 16, dtype=torch.uint8)
with self.assertRaises(AssertionError):
@@ -166,28 +141,6 @@ class TestMHASpecSurface(unittest.TestCase):
num_pages=4,
)
def test_spec_offsets_equal_block_origins(self):
"""The spec's byte math and the view builder's origins are two
independent derivations of the envelope; under uniform rows they must
agree: layer_k_offset(l) == (2l)*ps*row, layer_v_offset(l) == (2l+1)*ps*row."""
spec = _mha_spec()
for ps in (1, 4):
row = spec.k_row_bytes()
for l in range(_L):
self.assertEqual(spec.layer_k_offset_in_page(l, ps), (2 * l) * ps * row)
self.assertEqual(
spec.layer_v_offset_in_page(l, ps), (2 * l + 1) * ps * row
)
def test_entry_bytes_matches_layout_helper(self):
spec = _mha_spec()
self.assertEqual(
spec.entry_bytes(),
mha_entry_bytes(
layer_num=_L, head_num=_H, head_dim=_D, v_head_dim=_D, itemsize=_ITEM
),
)
class TestMHAViews(unittest.TestCase):
def test_view_shapes_are_stock_mha(self):
@@ -202,11 +155,8 @@ class TestMHAViews(unittest.TestCase):
self.assertEqual(v.stride(), (_ROW, _D, 1))
def test_addressing_matches_strided_reference(self):
"""Cross-readback: bytes written through the reference STRIDED views at
(page, slot) must be read back through the views at kernel_id(t),
for both K and V of every layer — and vice versa. This pins that the
view builder and the independent strided description agree on the
same physical envelope."""
"""Cross-readback both ways: a write through the strided oracle at
(page, slot) must read back through the view at kernel_id(t)."""
for ps in (1, 4):
num_pages = 5
raw = _make_raw(ps, num_pages)
@@ -240,37 +190,6 @@ class TestMHAViews(unittest.TestCase):
torch.all(sv[l][p, s] == float(p * 100 + l * 10 + s + 4))
)
def test_byte_addresses_match_envelope_formula(self):
"""The per-layer view's byte address for token ``t``, layer ``L`` must equal
the hand-computed envelope formula: page origin + layer-block origin +
slot offset. Independent of any view builder — this is the raw layout
contract every envelope consumer (moves, sizing, transfer math) relies
on."""
k_row = _ROW * _ITEM
v_row = _ROW * _ITEM
for ps in (1, 4):
num_pages = 5
page_bytes = ps * _L * (k_row + v_row)
dk, dv = _build_views(_make_raw(ps, num_pages), ps, num_pages)
for t in (0, 1, ps, 3 * ps + (ps - 1), 4 * ps):
d = _kernel_id(t, ps)
for L in range(_L):
expected_k = (
(t // ps) * page_bytes
+ L * ps * (k_row + v_row)
+ (t % ps) * k_row
)
expected_v = (
(t // ps) * page_bytes
+ L * ps * (k_row + v_row)
+ ps * k_row
+ (t % ps) * v_row
)
got_k = (dk[L].storage_offset() + d * dk[L].stride(0)) * _ITEM
got_v = (dv[L].storage_offset() + d * dv[L].stride(0)) * _ITEM
self.assertEqual(got_k, expected_k, f"K t={t} L={L} ps={ps}")
self.assertEqual(got_v, expected_v, f"V t={t} L={L} ps={ps}")
def test_k_and_v_share_one_kernel_id_without_aliasing(self):
"""One kernel-facing id, 2L distinct cells (K and V of every layer): writes
through all 2L views at the SAME id must not clobber each other."""
@@ -291,12 +210,6 @@ class TestMHAViews(unittest.TestCase):
with self.assertRaises(AssertionError):
_build_views(raw, ps, num_pages)
def test_asymmetric_dims_rejected(self):
ps, num_pages = 2, 4
raw = _make_raw(ps, num_pages)
with self.assertRaises(AssertionError):
_build_views(raw, ps, num_pages, head_dim=6, v_head_dim=4)
# ---- pool level ----
@@ -341,26 +254,6 @@ class TestUnifiedKVPoolViews(unittest.TestCase):
self.assertEqual(v[0].dim(), 3, f"{name} V at ps={ps}")
self.assertTrue(k[0].is_contiguous())
def test_tail_pad_is_derived_from_the_specs(self):
"""The per-layer views hang past the last page envelope, so the pool
over-allocates one envelope of the widest sub-pool. Derived here, not
passed in, so no construction site can under-allocate it."""
for ps in (1, 4):
kv = _make_pool(ps)
full, swa = _mha_spec(), _swa_spec()
self.assertEqual(
kv.view_tail_pad_bytes,
ps * max(full.entry_bytes(), swa.entry_bytes()),
f"tail pad at ps={ps}",
)
self.assertEqual(
kv._raw.numel(),
full.entry_bytes() * _N_FULL
+ swa.entry_bytes() * _N_SWA
+ kv.view_tail_pad_bytes,
"the pad extends the allocation only",
)
def _layer(l):
return SimpleNamespace(layer_id=l)
@@ -386,12 +279,8 @@ class TestUnifiedMHATokenToKVPool(unittest.TestCase):
self.assertEqual(pool_under_test.size, n_rows - ps)
def test_stock_write_lands_on_envelope_truth(self):
"""Byte-identity: the pool's stock inherited `set_kv_buffer` at kernel-facing
locs must produce exactly the bytes that direct writes through STRIDED
views over the same envelope produce at the same (page, slot, layer)
cells. The strided views are built here purely as the independent
description of the envelope — pins the whole write path (loc -> view ->
raw bytes) end to end."""
"""Byte-identity: the inherited `set_kv_buffer` at kernel-facing locs must
produce the same bytes as writes through the strided oracle."""
for ps in (1, 4):
kv, pool = _make_pool_and_kv(ps, device=_STORE_DEV)
# An independent strided view of the SAME sub-pool region.
@@ -431,10 +320,8 @@ class TestUnifiedMHATokenToKVPool(unittest.TestCase):
)
def test_move_kv_cache_relocates_whole_envelopes(self):
"""Compaction hands PHYSICAL token runs, not kernel-facing ids. The override
must relocate exactly the page envelopes those runs name — red if it is
lost, since the inherited per-layer move would apply physical ids to
the row space."""
"""Compaction hands PHYSICAL token runs, not kernel-facing ids; the
override must relocate exactly the page envelopes those runs name."""
ps = 4
kv, pool = _make_pool_and_kv(ps)
live = kv._raw.numel() - kv.view_tail_pad_bytes
@@ -459,8 +346,7 @@ class TestUnifiedMHATokenToKVPool(unittest.TestCase):
def test_transfer_entry_points_fail_loud(self):
"""PD / CPU-copy entry points assume per-layer buffers indexed by TOKEN
id; against the row space they would silently mis-index (or hit a
missing-attr AttributeError). Every one of them must raise."""
id and would silently mis-index the row space, so each must raise."""
_, pool = _make_pool_and_kv(1)
with self.assertRaises(NotImplementedError):
pool.get_contiguous_buf_infos()
@@ -472,10 +358,8 @@ class TestUnifiedMHATokenToKVPool(unittest.TestCase):
pool.set_kv_buffer_prefix_valid()
def test_hnd_env_cannot_hijack_layout(self):
"""SGLANG_USE_HND_KVCACHE=1 used to flip the inherited env-driven
layout selector, putting the pool in a mode whose code paths do not
match its buffers (HND indexes 4-D; the per-layer views are 3-D). The
pinned label must win."""
"""SGLANG_USE_HND_KVCACHE must not flip this pool's layout: HND indexes
4-D while the per-layer views are 3-D, so the pinned label has to win."""
with envs.SGLANG_USE_HND_KVCACHE.override(True):
_, pool = _make_pool_and_kv(1)
self.assertFalse(pool.use_hnd)
@@ -483,17 +367,15 @@ class TestUnifiedMHATokenToKVPool(unittest.TestCase):
class TestFactoryViews(unittest.TestCase):
"""The real SWA factory builds the sub-pools and wires the matching
kernel-facing multipliers into the composite allocator. End-to-end over
that factory, the rebind must emit BOTH kernel-facing write locs."""
"""Over the real SWA factory: matching kernel-facing multipliers in the
composite allocator, and a rebind that emits both write locs."""
# _swa_factory geometry: L_full = L_swa = 2, uniform 8/8 dims, ps = 1.
FULL_MULT = 4 # 2 * L_full
SWA_MULT = 4 # 2 * L_swa
def _bundle(self):
# Self-contained tiny SWA-factory bundle (L_full = L_swa = 2, uniform
# 8/8 dims, ps = 1) — small enough that per-layer views build on CPU.
# Kept tiny so the per-layer views build on CPU.
from sglang.srt.mem_cache.unified_memory_pool import init_unified_swa_pools
return init_unified_swa_pools(
@@ -528,11 +410,8 @@ class TestFactoryViews(unittest.TestCase):
self.assertGreater(pool.view_tail_pad_bytes, 0)
def test_rebind_emits_kernel_facing_full_and_build_derives_swa(self):
"""End-to-end over the real factory: rebind_write_loc rebinds
out_cache_loc to FULL-kernel-facing ids (phase 1), and the per-batch build
derives the SWA write loc pointwise from those kernel-facing values
(phase 2) — both checked against the formulas over the VIRTUAL
ids."""
"""rebind_write_loc rebinds out_cache_loc to FULL-kernel-facing ids, and
the SWA write loc is derived pointwise from those kernel-facing values."""
from sglang.srt.mem_cache.kv_index_translator import KVIndexTranslator
b = self._bundle()
test/registered/unit/mem_cache/test_unified_mla_gpu_parity.py
+11 -47
View File
@@ -14,19 +14,11 @@
"""GPU parity of the per-layer-view `UnifiedMLATokenToKVPool` against the stock
`MLATokenToKVPool` on real K3 MLA geometry (L=24, D=512+64).
The unified pool receives kernel-facing locs (kernel_id(t) = (t//ps)*(ps*L) + t%ps); the
reference pool receives the raw token ids. Every (layer, token) cell must hold
identical bytes afterwards. Covers:
- `set_mla_kv_buffer` under BOTH kernel paths — the Triton fallback
(n_loc < 768) and the TMA JIT fast path (n_loc >= 768, which flattens the
buffer via `.view(shape[0], -1)`, only legal because per-layer views are
contiguous);
- `set_kv_buffer` (combined pre-concatenated write, the Triton-backend path);
- `get_mla_kv_buffer` roundtrip;
- page_size 1 and 64.
python -m pytest test/registered/unit/mem_cache/test_unified_mla_gpu_parity.py -v
The unified pool receives kernel-facing locs (kernel_id(t) = (t//ps)*(ps*L) +
t%ps) where the reference pool receives raw token ids; every (layer, token)
cell must hold identical bytes afterwards. The TMA JIT fast path (n_loc >= 768)
flattens the buffer via `.view(shape[0], -1)`, which is legal only because the
per-layer views are contiguous.
"""
import types
@@ -134,17 +126,12 @@ class TestUnifiedMLAPoolGPUParity(unittest.TestCase):
torch.cuda.synchronize()
self._assert_parity(unified, ref, locs, ps)
def test_set_mla_kv_buffer_triton_fallback_ps1(self):
self._run_set_mla(ps=1, n_loc=256) # < 768 -> Triton fallback kernel
def test_set_mla_kv_buffer_tma_jit_ps1(self):
self._run_set_mla(ps=1, n_loc=1024) # >= 768 -> TMA JIT fast path
def test_set_mla_kv_buffer_triton_fallback_ps64(self):
self._run_set_mla(ps=64, n_loc=256)
def test_set_mla_kv_buffer_tma_jit_ps64(self):
self._run_set_mla(ps=64, n_loc=1024)
def test_set_mla_kv_buffer_matches_stock_pool(self):
"""Both kernel paths at both page sizes: n_loc < 768 takes the Triton
fallback, n_loc >= 768 the TMA JIT fast path."""
for ps, n_loc in ((1, 256), (1, 1024), (64, 256), (64, 1024)):
with self.subTest(page_size=ps, n_loc=n_loc):
self._run_set_mla(ps=ps, n_loc=n_loc)
def test_set_kv_buffer_combined_write(self):
for ps in (1, 64):
@@ -175,29 +162,6 @@ class TestUnifiedMLAPoolGPUParity(unittest.TestCase):
torch.testing.assert_close(got_nope, nope, rtol=0, atol=0)
torch.testing.assert_close(got_rope, rope, rtol=0, atol=0)
def test_move_kv_cache_page_envelope_gpu(self):
for ps in (1, 64):
unified, ref, max_tokens = _make_pools(ps)
num_pages = max_tokens // ps
n_loc = ps # one full page of tokens
src_page, dst_page = num_pages - 2, 2
src_t = torch.arange(ps, device=_DEV, dtype=torch.int64) + src_page * ps
dst_t = torch.arange(ps, device=_DEV, dtype=torch.int64) + dst_page * ps
torch.manual_seed(17)
for l in range(_L):
layer = types.SimpleNamespace(layer_id=l)
k = torch.randn(n_loc, 1, _D, dtype=_DTYPE, device=_DEV)
unified.set_kv_buffer(layer, _kernel_id(src_t, ps), k, None)
before = [
unified.get_key_buffer(l)[_kernel_id(src_t, ps)].clone()
for l in range(_L)
]
unified.move_kv_cache(dst_t, src_t)
torch.cuda.synchronize()
for l in range(_L):
got = unified.get_key_buffer(l)[_kernel_id(dst_t, ps)]
torch.testing.assert_close(got, before[l], rtol=0, atol=0)
if __name__ == "__main__":
unittest.main()
test/registered/unit/mem_cache/test_unified_mla_views.py
+52 -104
View File
@@ -11,26 +11,14 @@
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""MLA views for the unified memory pool (MLA-hybrid-Mamba, Kimi K3).
"""MLA views for the unified memory pool (MLA-hybrid-Mamba, Kimi K3), CPU-only.
Covers, CPU-only (pure torch — no GPU / Triton kernels):
- `MLASubPoolSpec` byte math;
- `build_mla_views` addressing: view_l[kernel_id(t)] must land exactly at
the page-major envelope byte offset `p*(L*ps*D) + l*(ps*D) + s*D`, the
overlapping per-layer views must not alias at equal kernel-facing ids, and the
missing-tail-pad case must fail loud;
- `UnifiedKVPool` MLA plumbing: `view_tail_pad_bytes` extends the allocation
only, and the reserved sink floor covers the whole page-0 envelope;
- `UnifiedMLATokenToKVPool`: buffer wiring, V-as-prefix-slice, and the
page-envelope `move_kv_cache` (REAL physical token ids, page-major runs);
- `MultiEndedAllocator.translate_kv_loc_for_kernel`: kernel id = v2p-page * (ps*L) +
offset, tombstone clamp to the sink, `out=` contract, multiplier-1
fallback, and correctness across eager compaction.
Addressing law under test: the (page, layer, slot) cell sits at envelope byte
offset `p*(L*ps*D) + l*(ps*D) + s*D`, reached through the kernel-facing id
`(t // ps) * (ps * L) + t % ps`.
GPU parity of the actual read/write kernels (set_mla_kv_buffer TMA path etc.)
lives in the server-level tests, not here.
python -m pytest test/registered/unit/mem_cache/test_unified_mla_views.py -v
GPU parity of the read/write kernels (set_mla_kv_buffer TMA path etc.) lives in
`test_unified_mla_gpu_parity.py`.
"""
from sglang.test.ci.ci_register import register_cpu_ci
@@ -42,10 +30,7 @@ import unittest
import torch
from sglang.srt.mem_cache.allocator.unified_sub_pool import MultiEndedAllocator
from sglang.srt.mem_cache.layout.page_major import (
build_mla_views,
mla_entry_bytes,
)
from sglang.srt.mem_cache.layout.page_major import build_mla_views
from sglang.srt.mem_cache.unified_memory_pool import (
MambaSubPoolSpec,
MLASubPoolSpec,
@@ -55,8 +40,8 @@ from sglang.srt.mem_cache.unified_memory_pool import (
_DEV = "cpu"
# Small-but-nontrivial MLA geometry: L=3 layers, D=8 (=6+2), so every byte
# offset is hand-checkable. Real K3 is L=24, D=576 (=512+64).
# Geometry kept tiny so every byte offset is hand-checkable; real K3 is
# L=24, D=576 (=512+64).
_L = 3
_LORA = 6
_ROPE = 2
@@ -107,16 +92,6 @@ def _kernel_id(t, ps, layer_num):
class TestMLASubPoolSpec(unittest.TestCase):
def test_entry_bytes_and_dim(self):
spec = _mla_spec()
self.assertEqual(spec.kv_cache_dim, _D)
self.assertEqual(spec.entry_bytes(), _L * _D * _ITEM)
self.assertEqual(
spec.entry_bytes(),
mla_entry_bytes(layer_num=_L, kv_cache_dim=_D, itemsize=_ITEM),
)
self.assertEqual(spec.get_dtype(), _DTYPE)
def test_rejects_nonpositive_dims(self):
with self.assertRaises(AssertionError):
MLASubPoolSpec(
@@ -153,7 +128,7 @@ class TestMLAViews(unittest.TestCase):
n_rows = num_pages * _L * ps
for v in views:
self.assertEqual(tuple(v.shape), (n_rows, 1, _D))
# contiguous in the (row, dim) sense — .view(-1, ps, D) legality
# contiguous in the (row, dim) sense: .view(-1, ps, D) legality
self.assertEqual(v.stride(0), _D)
self.assertEqual(v.stride(2), 1)
flat = raw.view(_DTYPE)
@@ -203,14 +178,6 @@ class TestMLAViews(unittest.TestCase):
class TestUnifiedKVPoolMLA(unittest.TestCase):
def test_max_slots_ignore_tail_pad(self):
pool, full, mamba = _make_unified(page_size=4)
total = full.entry_bytes() * 64 + mamba.entry_bytes() * 8
self.assertEqual(pool.max_slots("full"), total // full.entry_bytes())
self.assertEqual(pool.max_slots("mamba"), total // mamba.entry_bytes())
# allocation actually carries the pad
self.assertEqual(pool._raw.numel(), total + 4 * full.entry_bytes())
def test_reserved_floor_covers_page0_envelope(self):
ps = 4
pool, full, mamba = _make_unified(page_size=ps)
@@ -253,15 +220,24 @@ class TestUnifiedMLATokenToKVPool(unittest.TestCase):
self.assertTrue(torch.all(v[7] == 2.5))
def test_move_kv_cache_moves_page_envelopes(self):
"""Whole page envelopes relocate, in raw bytes and (at ps=4) as read
back through the per-layer views at the destination kernel ids."""
for ps in (1, 4):
pool, kv_pool = self._make(ps=ps)
num_pages = pool.max_slots("full") // ps
page_bytes = ps * _L * _D * _ITEM
env = pool._raw[: num_pages * page_bytes].view(num_pages, page_bytes)
src_pages = torch.tensor([num_pages - 2, num_pages - 4])
dst_pages = torch.tensor([2, 3])
src_pages = torch.tensor([num_pages - 2, num_pages - 4, num_pages - 3])
dst_pages = torch.tensor([2, 3, 5])
env[src_pages[0]] = 7
env[src_pages[1]] = 9
if ps == 4:
# write through the views at src, expect it at dst after the move
for l in range(_L):
for s in range(ps):
kv_pool.kv_buffer[l][
_kernel_id(int(src_pages[2]) * ps + s, ps, _L)
] = float(l * ps + s + 1)
# page-major token runs, exactly how compaction expands pages
offsets = torch.arange(ps, dtype=torch.int64)
src_t = (src_pages[:, None] * ps + offsets).reshape(-1)
@@ -269,29 +245,15 @@ class TestUnifiedMLATokenToKVPool(unittest.TestCase):
kv_pool.move_kv_cache(dst_t, src_t)
self.assertTrue(torch.all(env[dst_pages[0]] == 7), f"ps={ps}")
self.assertTrue(torch.all(env[dst_pages[1]] == 9), f"ps={ps}")
def test_move_then_readback(self):
ps = 4
pool, kv_pool = self._make(ps=ps)
num_pages = pool.max_slots("full") // ps
src_page, dst_page = num_pages - 3, 5
# write through the views at src, expect it at dst after the move
for l in range(_L):
for s in range(ps):
kv_pool.kv_buffer[l][_kernel_id(src_page * ps + s, ps, _L)] = float(
l * ps + s + 1
)
offsets = torch.arange(ps, dtype=torch.int64)
kv_pool.move_kv_cache(
(torch.tensor([dst_page])[:, None] * ps + offsets).reshape(-1),
(torch.tensor([src_page])[:, None] * ps + offsets).reshape(-1),
)
for l in range(_L):
for s in range(ps):
got = kv_pool.kv_buffer[l][_kernel_id(dst_page * ps + s, ps, _L)]
self.assertTrue(
torch.all(got == float(l * ps + s + 1)), f"(l={l}, s={s})"
)
if ps == 4:
for l in range(_L):
for s in range(ps):
got = kv_pool.kv_buffer[l][
_kernel_id(int(dst_pages[2]) * ps + s, ps, _L)
]
self.assertTrue(
torch.all(got == float(l * ps + s + 1)), f"(l={l}, s={s})"
)
class _FakeKVCache:
@@ -325,23 +287,29 @@ class TestTranslateKvLocForKernel(unittest.TestCase):
mamba_alloc.bind_peer(full_alloc)
return full_alloc
def test_kernel_id_matches_formula_ps1(self):
alloc = self._build(ps=1)
v = alloc.alloc(8)
self.assertIsNotNone(v)
phys = alloc.translate_kv_loc(v)
kernel = alloc.translate_kv_loc_for_kernel(v)
self.assertTrue(torch.all(kernel == phys * _L))
def test_kernel_id_matches_formula(self):
"""kernel id = (phys // ps) * (ps * multiplier) + phys % ps, across
page sizes, the multiplier-1 physical fallback, and eager compaction."""
for ps, multiplier in ((1, _L), (4, _L), (1, 1)):
with self.subTest(page_size=ps, multiplier=multiplier):
alloc = self._build(ps=ps, multiplier=multiplier)
a = alloc.alloc(4 * ps)
b = alloc.alloc(4 * ps)
c = alloc.alloc(4 * ps)
self.assertIsNotNone(c)
def test_kernel_id_matches_formula_paged(self):
ps = 4
alloc = self._build(ps=ps)
v = alloc.alloc(3 * ps)
self.assertIsNotNone(v)
phys = alloc.translate_kv_loc(v)
kernel = alloc.translate_kv_loc_for_kernel(v)
expected = (phys // ps) * (ps * _L) + phys % ps
self.assertTrue(torch.all(kernel == expected))
def check(virt):
phys = alloc.translate_kv_loc(virt)
expected = (phys // ps) * (ps * multiplier) + phys % ps
self.assertTrue(
torch.all(alloc.translate_kv_loc_for_kernel(virt) == expected)
)
for virt in (a, b, c):
check(virt)
alloc.free(b) # eager compaction relocates survivors
for virt in (a, c):
check(virt)
def test_tombstone_clamps_to_sink(self):
alloc = self._build(ps=1)
@@ -365,26 +333,6 @@ class TestTranslateKvLocForKernel(unittest.TestCase):
alloc.translate_kv_loc_for_kernel(x, out=x)
self.assertTrue(torch.all(x == no_out))
def test_multiplier_one_falls_back_to_physical(self):
alloc = self._build(ps=1, multiplier=1)
v = alloc.alloc(4)
self.assertIsNotNone(v)
self.assertTrue(
torch.all(alloc.translate_kv_loc_for_kernel(v) == alloc.translate_kv_loc(v))
)
def test_kernel_id_follows_compaction(self):
alloc = self._build(ps=1)
a = alloc.alloc(4)
b = alloc.alloc(4)
c = alloc.alloc(4)
self.assertIsNotNone(c)
alloc.free(b) # eager compaction relocates survivors
phys_a = alloc.translate_kv_loc(a)
phys_c = alloc.translate_kv_loc(c)
self.assertTrue(torch.all(alloc.translate_kv_loc_for_kernel(a) == phys_a * _L))
self.assertTrue(torch.all(alloc.translate_kv_loc_for_kernel(c) == phys_c * _L))
if __name__ == "__main__":
unittest.main()
test/registered/unit/mem_cache/test_unified_npool_sweep.py
+61 -88
View File
@@ -13,25 +13,13 @@
# ==============================================================================
"""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:
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 -- and
sorts them into the canonical chain order
``[up end, floats (input order), down end]``. Every sub-pool view spans the
whole buffer at anchor 0; keeping the bands disjoint is the allocators' job.
- 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
Pure CPU geometry -- no allocator, no GPU.
"""
import unittest
@@ -46,10 +34,8 @@ from sglang.srt.mem_cache.unified_memory_pool import (
)
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.
# Hermetic convention of this directory's pool tests: plain unittest.TestCase,
# only ci_register imported (no heavy sglang.test.test_utils chain).
register_cpu_ci(est_time=30, suite="base-a-test-cpu")
_DEV = "cpu"
@@ -111,34 +97,46 @@ def _chain_names(pool: UnifiedKVPool):
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")],
def test_chain_order_is_canonical(self):
"""Ends canonical (up first, down last) whatever the input order;
floats keep INPUT order between them, on every cache-spec kind."""
for specs, expect in (
([_mha("full", "down"), _mamba("mamba", "up")], ["mamba", "full"]),
([_mamba("mamba", "up"), _mha("full", "down")], ["mamba", "full"]),
(
[_mha("full", "down"), _mha("swa", "float"), _mamba("conv", "up")],
["conv", "swa", "full"],
),
(
[_mha("swa", "float"), _mamba("conv", "up"), _mha("full", "down")],
["conv", "swa", "full"],
),
(
[_mamba("conv", "up"), _mha("full", "down"), _mha("swa", "float")],
["conv", "swa", "full"],
),
(
[
_mha("full", "down"),
_mha("f1", "float"),
_mamba("state", "up"),
_mha("f0", "float", layer_num=1),
],
["state", "f1", "f0", "full"],
),
(
[
_mamba("state", "up"),
_mha("f_mha", "float"),
_mla("f_mla", "float", layer_num=1),
_mamba("f_mamba", "float", layer_num=1),
_mha("full", "down"),
],
["state", "f_mha", "f_mla", "f_mamba", "full"],
),
):
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"])
with self.subTest(inputs=[s.name for s in specs]):
self.assertEqual(_chain_names(_make_pool(specs)), expect)
def test_by_name_geometry_independent_of_n(self):
two = _make_pool([_mha("full", "down"), _mamba("mamba", "up")])
@@ -151,9 +149,6 @@ class TestNPoolCanonicalOrder(unittest.TestCase):
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):
@@ -165,37 +160,24 @@ class TestNPoolValidation(unittest.TestCase):
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_end_direction_counts_rejected(self):
"""Exactly one grow-up END and one grow-down END, no more, no less."""
for case, specs in (
("two_ups", [_mha("a", "up"), _mamba("b", "up")]),
("missing_down", [_mha("a", "up"), _mha("b", "float")]),
(
"missing_up",
[_mha("a", "down"), _mha("b", "float"), _mha("c", "float")],
),
):
with self.subTest(case=case):
with self.assertRaisesRegex(AssertionError, "exactly one grow-up"):
_make_pool(specs)
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):
@@ -217,8 +199,7 @@ class TestReservedFloorWithFloats(unittest.TestCase):
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.
# (16*128 = 2048 B) consumes the whole buffer, leaving no slot.
with self.assertRaisesRegex(RuntimeError, "no room"):
_make_pool(
[_mamba("state", "up"), _mha("swa", "float"), _mha("full", "down")],
@@ -258,14 +239,6 @@ class TestFloatViews(unittest.TestCase):
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()
test/registered/unit/mem_cache/test_unified_tri_pool.py
+98 -243
View File
@@ -15,23 +15,7 @@
SWA KV + mamba/conv state in ONE unified byte buffer, chain
``[mamba (up END) | swa (FLOAT) | full (down END)]``.
Pinned contracts (each guards a distinct failure mode):
- chain wiring + the swa side being a `FloatMultiEndedAllocator`;
- the JOINT `available_size()` feasibility contract: `alloc(N)` for
N == available_size() must succeed (full extends into the high band
first, then the float extends a single side — the predicate models
exactly that order; over-promising here is the fail-loud
`alloc_with_virtual` assert, i.e. a crash in production);
- `free_swa` tombstones become float HOLES recycled IN PLACE by later
allocs (steady-state SWA churn == zero copies), while the full side
keeps the token;
- the per-request state surface (`UnifiedMambaSlotAllocator` over the
mamba END) and its independence from the token surface;
- urgent flushes drain the two ENDS but never touch the float's holes.
Pure CPU; fakes stand in for the KV pools (data markers verify moves).
python -m pytest test/registered/unit/mem_cache/test_unified_tri_pool.py -v
"""
import inspect
@@ -39,11 +23,13 @@ import unittest
import torch
import sglang.srt.mem_cache.allocator.unified_sub_pool as mea
from sglang.srt.mem_cache.allocator.unified_hybrid_swa import (
UnifiedMambaSWATokenToKVPoolAllocator,
)
from sglang.srt.mem_cache.allocator.unified_sub_pool import FloatMultiEndedAllocator
from sglang.srt.mem_cache.allocator.unified_sub_pool import (
FloatMultiEndedAllocator,
MultiEndedAllocator,
)
from sglang.srt.mem_cache.unified_memory_pool import (
MambaSubPoolSpec,
MHASubPoolSpec,
@@ -195,22 +181,6 @@ class TestUnifiedTriPool(unittest.TestCase):
# -- the joint availability contract --
def test_available_size_alloc_contract(self):
for lazy in (False, True):
_, allocator, kvcache, _ = self._build(lazy_compaction=lazy)
avail = allocator.available_size()
self.assertGreater(avail, 0)
v = allocator.alloc(avail)
self.assertIsNotNone(
v, f"alloc(available_size()={avail}) must succeed (lazy={lazy})"
)
self.assertEqual(int(v.numel()), avail)
# Both sides bound for every allocated virtual id.
fa = allocator.full_attn_allocator
sa = allocator.swa_attn_allocator
self.assertTrue(bool((fa.virtual_to_physical[v] >= 0).all()))
self.assertTrue(bool((sa.virtual_to_physical[v] >= 0).all()))
def test_available_shrinks_as_state_slots_grow(self):
_, allocator, _, _ = self._build()
before = allocator.available_size()
@@ -224,9 +194,8 @@ class TestUnifiedTriPool(unittest.TestCase):
# -- steady-state SWA churn: tombstones -> holes -> in-place reuse --
def _swa_interior_block(self, allocator, blocks):
"""The block whose SWA-physical pages touch neither float boundary --
`free_swa` on it must create interior holes (a boundary block would be
absorbed instead; both are zero-copy, different mechanisms)."""
"""The block whose SWA-physical pages touch neither float boundary:
`free_swa` on it holes; a boundary block would be absorbed instead."""
sa = allocator.swa_attn_allocator
for v in blocks:
pages = set(int(x) for x in sa.virtual_to_physical[v].tolist())
@@ -262,10 +231,9 @@ class TestUnifiedTriPool(unittest.TestCase):
self.assertEqual(len(sa._inverse_history), 0) # zero copies
def test_free_swa_boundary_block_absorbed_zero_copy(self):
# The OTHER zero-copy mechanism: a boundary block's tombstones shrink
# the span, handing bytes back to the neighbours. The shrink is
# DEFERRED out of the per-step free (it needs the hole set on the
# host); the per-step opportunistic flush is where it lands.
# A boundary block's tombstones shrink the span instead of holing it.
# The shrink needs the hole set on the host, so it is deferred out of
# the per-step free and lands in the opportunistic flush.
_, allocator, kvcache, _ = self._build()
blocks = [allocator.alloc(4) for _ in range(2)]
for v in blocks:
@@ -329,16 +297,7 @@ class TestUnifiedTriPool(unittest.TestCase):
(pool.max_slots("mamba") - 1) - 1,
)
# -- cost + flush semantics --
def test_mamba_slot_full_token_cost_formula(self):
_, allocator, _, _ = self._build()
e_tok = (
allocator.full_attn_allocator.entry_bytes
+ allocator.swa_attn_allocator.entry_bytes
)
m = allocator.mamba_allocator.entry_bytes_per_page
self.assertEqual(allocator.mamba_slot_full_token_cost(), -(-m // e_tok))
# -- flush semantics --
def test_urgent_flush_preserves_float_holes(self):
_, allocator, kvcache, _ = self._build(lazy_compaction=True)
@@ -359,18 +318,11 @@ class TestTriPagedFreeGroup(unittest.TestCase):
"""The tri composite at PAGE SIZE > 1, driven through the production free
path: free_group_begin -> free_segment -> free_group_end.
Regression (GPU eval_440, Inkling ps=128 boot crash): every other tri test
runs at page_size=1, where the page-REPRESENTATIVE machinery
(`free_page_reps_group` / `_release_page_reps`, added when the free path
was made host-sync-free) is entirely dead code — `free_segment` frees
directly. At ps>1 the composite releases reps by calling
`swa_attn_allocator.free(..., _pages=...)`, and the float allocator was
ported from a base that predates that keyword, so the first real decode
batch died with `TypeError: unexpected keyword argument '_pages'`.
`_pages` is not cosmetic: honouring it is what keeps the free path free of
the data-dependent `torch.unique` host sync, so this also pins that the
float takes the caller's page ids rather than re-deriving them.
Regression: every other tri test runs at page_size=1, where the
page-representative machinery is dead code. At ps>1 the composite releases
reps via `swa_attn_allocator.free(..., _pages=...)`, which the float must
accept or the first decode batch dies; honouring `_pages` is also what
keeps the free path off the data-dependent `torch.unique` host sync.
"""
def _build_paged(self, page_size=4, n_full=64, n_swa=32, n_state=8):
@@ -418,22 +370,7 @@ class TestTriPagedFreeGroup(unittest.TestCase):
# Capacity fully recovered: the float parked, both ends rewound.
self.assertTrue(allocator.swa_attn_allocator._is_frontier_transparent())
def test_float_free_honours_caller_supplied_pages(self):
"""`_pages` must be USED, not merely accepted -- re-deriving it is the
host sync the paged free path exists to avoid."""
pool, allocator = self._build_paged()
v = allocator.alloc(8)
self.assertIsNotNone(v)
sa = allocator.swa_attn_allocator
ps = allocator.page_size
pages = (v[::ps] // ps).clone()
live_before = sa._live_pages()
sa.free(v[::ps] * 0 + v[::ps], _pages=pages)
self.assertEqual(sa._live_pages(), live_before - pages.numel())
def test_ungrouped_segment_free_also_reaches_the_float(self):
"""`free_segment` outside a free group releases reps immediately --
the same float call, one frame shallower."""
pool, allocator = self._build_paged()
v = allocator.alloc(8)
self.assertIsNotNone(v)
@@ -442,13 +379,10 @@ class TestTriPagedFreeGroup(unittest.TestCase):
class TestTriFreeSwaNoHostSync(unittest.TestCase):
"""The tri's swa side is the FLOAT, and the float can never run the lazy
event pipeline — so unless the per-step frees carry caller-derived page
ids, the tri silently reintroduces the host syncs the sync-free free
path removed. Poison the ops to pin the property.
(Fixtures at page_size > 1 on purpose: ps==1 short-circuits the whole
page machinery and hides exactly this class of bug.)
"""The tri's swa side is the FLOAT, which can never run the lazy event
pipeline: unless the per-step frees carry caller-derived page ids, the tri
reintroduces the host syncs the sync-free free path removed. Fixtures run
at page_size > 1; ps==1 short-circuits the page machinery and hides this.
"""
PS = 4
@@ -476,24 +410,9 @@ class TestTriFreeSwaNoHostSync(unittest.TestCase):
alloc.free_swa(v[: 4 * self.PS], start_pos=0)
self.assertEqual(alloc.verify_byte_accounting(), [])
def test_float_free_has_no_stale_slot_item_sync(self):
"""The float's free must not `.item()`-assert per free (the lazy-path
contract: callers must not double-free; the idle span == p2v-bound +
holes conservation catches violations without a per-free sync)."""
alloc = self._tri()
v = alloc.alloc(4 * self.PS)
sa = alloc.swa_attn_allocator
from unittest import mock
with mock.patch.object(
torch.Tensor, "item", side_effect=AssertionError("item = host sync")
):
sa.free(v[:: self.PS], _pages=v[:: self.PS] // self.PS)
def test_fallback_free_swa_still_correct_for_radix_shapes(self):
"""Radix eviction hands arbitrary node values (no start_pos): the
dedup fallback must keep working and end in the same state as the
stride path."""
dedup fallback must end in the same state as the stride path."""
a1, a2 = self._tri(), self._tri()
v1, v2 = a1.alloc(6 * self.PS), a2.alloc(6 * self.PS)
self.assertTrue(torch.equal(v1, v2))
@@ -512,12 +431,9 @@ class TestTriFreeSwaNoHostSync(unittest.TestCase):
class TestGeneralizedRebalance(unittest.TestCase):
"""The float must yield to WHICHEVER end is short, with the direction
computed from the layout — not only to the token path's hard-coded side.
The mechanism (`make_room`) was always side-agnostic; these pin the
POLICY: any end pool's own-alloc shortfall reaches
`_ask_float_for_room`, which derives the side from the caller's
growth direction."""
derived from the caller's growth direction rather than hard-coded to the
token path's side. `make_room` was always side-agnostic; these pin the
POLICY in `_ask_float_for_room`."""
PS = 4
@@ -527,34 +443,10 @@ class TestGeneralizedRebalance(unittest.TestCase):
)
return inst._build_paged(page_size=self.PS)[1]
def test_state_end_shortfall_slides_the_float_low(self):
"""The previously-missing direction: mamba (grow-up END) starved while
free bytes idle ABOVE the float. The remedy must slide the float up
(open its LOW side) and let the state alloc succeed."""
alloc = self._tri()
v = alloc.alloc(4 * self.PS) # places the float mid-region
self.assertIsNotNone(v)
ma = alloc.mamba_allocator
sa = alloc.swa_attn_allocator
self.assertFalse(sa._is_frontier_transparent())
# Fill the LOW band exactly: as many state slots as fit below the
# float's low frontier.
e_m = ma.entry_bytes_per_page
fit = (sa._byte_low_frontier() - ma._byte_high_frontier()) // e_m
self.assertGreater(fit, 0)
got = ma.alloc(int(fit) * ma.page_size)
self.assertIsNotNone(got)
low_before = sa.low_wm_page
# One more slot does NOT fit below the float -- only a rebalance helps.
more = ma.alloc(ma.page_size)
self.assertIsNotNone(more, "state alloc must succeed via float rebalance")
self.assertGreater(sa.low_wm_page, low_before) # float slid UP
self.assertEqual(alloc.verify_byte_accounting(), [])
def test_direction_is_derived_from_growth_on_both_ends(self):
"""Raw end+float+end chain, BOTH orientations in one fixture: the
up-growing end opens the float's LOW side; the down-growing end opens
its HIGH side. No layout assumption survives."""
up-growing end opens the float's LOW side, the down-growing end its
HIGH side."""
from test_multi_ended_allocator import TestFloatMultiEndedAllocator
inst = TestFloatMultiEndedAllocator(
@@ -583,8 +475,6 @@ class TestGeneralizedRebalance(unittest.TestCase):
self.assertLess(fla.high_wm_page, high_before) # opened HIGH side
def test_two_pool_chain_rebalance_is_a_noop(self):
"""No float in the chain => the remedy must change nothing (the
2-pool composites keep their exact pre-existing behavior)."""
from test_multi_ended_allocator import (
TestPagedMultiEndedAllocator as _PagedFixture,
)
@@ -617,9 +507,8 @@ class TestGeneralizedRebalance(unittest.TestCase):
class TestComputedShortSide(unittest.TestCase):
"""`_ask_float_for_room` must open the side that MEASURES short -- never
"the side facing full". These pin the per-side computation, including
the coupled-ends-on-both-sides shape a DSV4-style composite
(C128 | swa-float | C4) will need.
"the side facing full" -- including the shape with coupled ends on BOTH
sides of the float.
"""
PS = 4
@@ -638,9 +527,9 @@ class TestComputedShortSide(unittest.TestCase):
def test_float_share_short_opens_the_state_side(self):
"""RED-LINE: full's demand fits its band, the float's own share fits
NEITHER band, and the state side has the larger surplus — the policy
must open the STATE side (the float slides toward full during a
TOKEN alloc), which the old "side facing full" policy could never do.
NEITHER band, and the state side has the larger surplus -- so the
policy must open the STATE side, i.e. slide the float toward full
during a TOKEN alloc.
"""
from unittest import mock
@@ -654,10 +543,9 @@ class TestComputedShortSide(unittest.TestCase):
# Position: slide the float LOW (setup uses the mechanism directly),
# so the low band is small and the geometry below is expressible.
b_low0, b_high0 = self._sides(alloc)
# Two positioning moves: pack the float low (leapfrog over-opens by
# design), then open the LOW side back to ~2 full-pages -- small
# enough that F outgrows it, wide enough that the integer need_n
# window below is non-empty.
# Pack the float low (leapfrog over-opens by design), then reopen LOW
# to ~2 full-pages: F must outgrow it, yet the need_n window below
# must stay non-empty.
sa.make_room(side="high", min_bytes=b_low0 + b_high0 - 2 * e_f)
sa.make_room(side="low", min_bytes=2 * e_f)
@@ -686,10 +574,8 @@ class TestComputedShortSide(unittest.TestCase):
self.assertEqual(calls[0]["side"], "low") # the STATE side
def test_full_side_short_target_matches_the_closed_form(self):
"""Equivalence: when the full side is the short one (today's only
reachable end-shortage), the ask must equal the documented formula
demand + max(0, F - far_surplus) + slack — i.e. the historical
behavior is the special case, preserved."""
"""When the full side is the short one, the ask must equal the
closed form demand + max(0, F - far_surplus) + slack."""
from unittest import mock
alloc = self._tri()
@@ -717,9 +603,9 @@ class TestComputedShortSide(unittest.TestCase):
self.assertEqual(calls[0]["min_bytes"], want)
def test_two_coupled_ends_lands_demand_on_both_sides(self):
"""DSV4 shape (C128 | float | C4): a coupled set with ends on BOTH
sides. One-side-short must open that side; BOTH-sides-short must not
move at all (relocation is zero-sum between the bands)."""
"""A coupled set with ends on BOTH sides of the float: one-side-short
must open that side; BOTH-sides-short must not move at all
(relocation is zero-sum between the bands)."""
from unittest import mock
alloc = self._tri()
@@ -730,8 +616,8 @@ class TestComputedShortSide(unittest.TestCase):
alloc.full_attn_allocator,
alloc.mamba_allocator,
)
# Synthetic coupling: the state end joins the demand vector, exactly
# the override a DSV4-style composite would ship.
# Synthetic coupling: the state end joins the demand vector, the
# override a composite with ends on both sides would ship.
need = lambda self, t: {
fa: -(-t // self.page_size),
sa: -(-t // self.page_size),
@@ -764,11 +650,14 @@ class TestComputedShortSide(unittest.TestCase):
self.assertEqual(calls[0]["side"], "high")
def test_nothing_short_means_no_relocation(self):
"""Everything fits -> the policy must not move a single page."""
"""Everything fits -> no page moves. The tri's token vector carries
{mamba: 0}: a zero entry must not move the float either."""
from unittest import mock
alloc = self._tri()
alloc.alloc(4 * self.PS)
demand = alloc._alloc_demand(2 * self.PS)
self.assertEqual(demand[alloc.mamba_allocator], 0)
sa = alloc.swa_attn_allocator
with mock.patch.object(
sa, "make_room", side_effect=AssertionError("needless move")
@@ -779,12 +668,9 @@ class TestComputedShortSide(unittest.TestCase):
class TestFloatPolicyTotalTarget(unittest.TestCase):
"""`make_room`'s min_bytes is a TARGET for the whole band, not a delta.
Regression: the band-level policy passed `deficit + one page` — with a
PARTIALLY free band that is below the current gap, so `make_room`
no-oped and the allocation failed even though the float had room to
slide. (Its own test missed this because it filled the band exactly,
making deficit ≈ the whole need.) The demand-vector policy computes the
total target, so a partial gap under-asks never.
Regression: an ask shaped as `deficit + one page` lands BELOW the current
gap when the band is only PARTIALLY free, so `make_room` no-ops and the
allocation fails though the float had room to slide.
"""
PS = 4
@@ -804,34 +690,17 @@ class TestFloatPolicyTotalTarget(unittest.TestCase):
gap_slots = int((sa._byte_low_frontier() - ma._byte_high_frontier()) // e_m)
self.assertGreater(gap_slots, 2)
low_before = sa.low_wm_page
# Need = partial-gap + 3: the old delta-ask was BELOW the current
# gap, so nothing moved and this returned None.
# partial-gap + 3: a delta-shaped ask lands below the current gap.
got = ma.alloc((gap_slots + 3) * ma.page_size)
self.assertIsNotNone(got, "partial-gap shortfall must relocate, not fail")
self.assertGreater(sa.low_wm_page, low_before)
self.assertEqual(alloc.verify_byte_accounting(), [])
def test_zero_demand_bands_are_inert(self):
"""The tri's token vector carries {mamba: 0}: a zero entry must
neither move the float for mamba's sake nor trip the index guard."""
from unittest import mock
alloc = self._tri()
alloc.alloc(4 * self.PS)
demand = alloc._alloc_demand(2 * self.PS)
self.assertEqual(demand[alloc.mamba_allocator], 0)
sa = alloc.swa_attn_allocator
with mock.patch.object(
sa, "make_room", side_effect=AssertionError("needless move")
):
alloc._ask_float_for_room(1) # nothing short -> no relocation
class TestTriDeferredAbsorption(unittest.TestCase):
"""Boundary absorption is deferred out of the per-step free and paid once
at a quiescent point — the base allocator's model (its lazy free does "no
boundary absorb" and `_flush` pays a single D2H). These pin WHERE it is
now paid, and that skipping it stays merely conservative."""
at a quiescent point. These pin WHERE it is paid, and that skipping it
stays merely conservative."""
PS = 4
@@ -855,9 +724,8 @@ class TestTriDeferredAbsorption(unittest.TestCase):
self.assertEqual(alloc.verify_byte_accounting(), [])
def test_shortfall_ladder_absorbs_before_the_deficit_math(self):
"""The zero-copy rung must run FIRST: a stale-wide span would inflate
the rebalance deficit and buy a `make_room` relocation the shrink
already covers."""
"""The zero-copy rung must run FIRST: a stale-wide span inflates the
rebalance deficit and buys a relocation the shrink already covers."""
alloc = self._tri()
v = alloc.alloc(8 * self.PS)
sa = alloc.swa_attn_allocator
@@ -884,9 +752,9 @@ class TestTriDeferredAbsorption(unittest.TestCase):
def test_clean_flush_skips_the_d2h_entirely(self):
"""Only `free` can put a hole ON a boundary (alloc DRAINS holes into
live pages; extension adds live pages), so with nothing freed since
the last absorb the walk provably finds nothing — and must not pay
the D2H. Steady churn with only interior holes then costs no sync."""
live pages, extension adds live pages), so with nothing freed since
the last absorb the walk provably finds nothing and must not pay the
D2H."""
from unittest import mock
alloc = self._tri()
@@ -902,9 +770,8 @@ class TestTriDeferredAbsorption(unittest.TestCase):
self.assertEqual(sa._flush(urgent=False), 0)
def test_alloc_between_frees_cannot_hide_a_boundary_hole(self):
"""Soundness of the skip: an alloc drains holes and can change the
hole COUNT back to a previously-seen value, so the flag must be armed
by `free`, not inferred from `numel()`."""
"""An alloc drains holes and can restore a previously-seen hole
COUNT, so the flag must be armed by `free`, not read off `numel()`."""
alloc = self._tri()
v = alloc.alloc(8 * self.PS)
sa = alloc.swa_attn_allocator
@@ -1019,9 +886,9 @@ class TestTriFactorySizing(unittest.TestCase):
class TestTriPoolHardening(unittest.TestCase):
"""C1.7 pressure lanes: the planned-rebalance remedy in the alloc path
(a mis-positioned float must not fail an alloc that fits in total bytes),
retract-loop convergence through check_decode_capacity, and bounded copy
"""Pressure lanes: the planned-rebalance remedy in the alloc path (a
mis-positioned float must not fail an alloc that fits in total bytes),
retract-loop convergence through `check_decode_capacity`, and bounded copy
traffic under alternating end pressure.
"""
@@ -1029,9 +896,9 @@ class TestTriPoolHardening(unittest.TestCase):
return TestUnifiedTriPool._build(self, **kw)
def test_alloc_rebalances_a_blocking_float(self):
# Fill much of the high band so the float (midpoint-placed) walls off
# the low band's free bytes from `full`; the next alloc must succeed
# by SLIDING the float, not fail while total bytes suffice.
# Fill much of the high band so the midpoint-placed float walls the
# low band's free bytes off from `full`; the next alloc must SLIDE it
# rather than fail while total bytes suffice.
_, allocator, kvcache, _ = self._build(n_full=32, n_swa=24, n_state=8)
sa = allocator.swa_attn_allocator
v0 = allocator.alloc(4) # places the float at the region midpoint
@@ -1043,9 +910,8 @@ class TestTriPoolHardening(unittest.TestCase):
b_high_pages = fa._current_gap_bytes() // fa.entry_bytes_per_page
grab = fa.alloc(max(0, (b_high_pages - 2)))
self.assertIsNotNone(grab)
# The honest gate under-reports (no slide credit) -- asking BEYOND it
# is what fires the rebalance remedy; the ask still fits total free
# bytes because the LOW band holds them behind the float.
# The gate under-reports (no slide credit), so asking BEYOND it fires
# the remedy; the ask still fits the free bytes the LOW band holds.
avail = allocator.available_size()
need = avail + 4
live_before = sa._live_pages()
@@ -1063,9 +929,8 @@ class TestTriPoolHardening(unittest.TestCase):
self.assertEqual(allocator.verify_byte_accounting(), [])
def test_check_decode_capacity_retract_convergence(self):
# Simulated retract loop: requests' token blocks freed one at a time
# until the next-step allocation fits; must converge before bs=1 and
# never report capacity while the gate is short.
# Retract loop: token blocks freed one at a time until the next step
# fits; must converge before bs=1 and never report capacity early.
_, allocator, _, _ = self._build(n_full=32, n_swa=24, n_state=8)
reqs = []
while True:
@@ -1085,9 +950,8 @@ class TestTriPoolHardening(unittest.TestCase):
self.assertEqual(allocator.verify_byte_accounting(), [])
def test_alternating_pressure_copy_traffic_bounded(self):
# Alternating full-grow / swa-churn cycles: total float moves stay
# bounded (hole recycling + absorption do the steady-state work; the
# rebalance fires only on real positional deficits).
# Alternating full-grow / swa-churn: hole recycling and absorption do
# the steady-state work, so total float moves stay bounded.
_, allocator, kvcache, _ = self._build(n_full=48, n_swa=32, n_state=8)
sa = allocator.swa_attn_allocator
fa = allocator.full_attn_allocator
@@ -1124,13 +988,11 @@ class TestJointCapacityIsHonoured(unittest.TestCase):
"""`alloc(available_size())` must never fail.
REGRESSION: the joint predicate priced the swa float's extension in RAW
BYTES while `take_physical_pages` can only use whole pages on the float's
OWN grid -- `_region_bounds_pages` rounds the band's low edge UP. The
bounding frontier is a multiple of the NEIGHBOUR's entry size, which is
unrelated to the float's, so the byte budget credited a page the grid could
not yield and the very first alloc tripped `alloc_with_virtual`'s backstop
assert. Swept over geometries rather than pinned to one, so a symmetric
mistake on the FULL side would surface here too.
BYTES, but `take_physical_pages` yields only whole pages on the float's
OWN grid, whose low edge is rounded up to a multiple of the NEIGHBOUR's
entry size -- so the budget credited a page the grid could not yield.
Swept over geometries so a symmetric mistake on the FULL side surfaces
here too.
"""
def _build(self, *, page_size, n_full, n_swa, n_state, lazy, specs):
@@ -1162,8 +1024,7 @@ class TestJointCapacityIsHonoured(unittest.TestCase):
def test_fresh_boot_alloc_of_available_size_succeeds(self):
# Geometries chosen so the mamba end's frontier (a multiple of the
# STATE entry size) lands off the swa float's page grid -- the
# misalignment the byte budget used to ignore.
# STATE entry size) lands off the swa float's page grid.
for page_size in (1, 2, 4):
for fl, sl, ml in ((4, 3, 1), (4, 2, 2), (6, 3, 1), (3, 5, 2)):
for n_full, n_swa, n_state in ((24, 16, 4), (32, 16, 8), (20, 12, 6)):
@@ -1200,6 +1061,15 @@ class TestJointCapacityIsHonoured(unittest.TestCase):
f"alloc(available_size()={n}) returned None",
)
self.assertEqual(out.numel(), n)
# Both sides bound for every allocated virtual id.
fa = alloc.full_attn_allocator
sa = alloc.swa_attn_allocator
self.assertTrue(
bool((fa.virtual_to_physical[out] >= 0).all())
)
self.assertTrue(
bool((sa.virtual_to_physical[out] >= 0).all())
)
def test_available_size_never_exceeds_the_float_page_grid(self):
"""Direct form: the joint answer, converted to float pages, must fit
@@ -1235,13 +1105,11 @@ class TestJointCapacityIsHonoured(unittest.TestCase):
class TestFloatRelocationIsOrderedAgainstTheForward(unittest.TestCase):
"""Float relocation must settle the in-flight forward BEFORE its first copy.
REGRESSION: `make_room` / `compact_holes` issued `move_kv_cache` and rebound
`virtual_to_physical` with no ordering against the running forward, so the
copy could carry pre-write bytes and the rebind then pointed every later
reader at a destination that never received those writes -- silently wrong
KV, no crash. The END pools guard exactly this hazard in
`_flush(urgent=True)` via `_settle_inflight_forward`; the float had no
`forward_stream` / `wait_event` / settle call anywhere in its body.
REGRESSION: `make_room` / `compact_holes` issued `move_kv_cache` and
rebound `virtual_to_physical` with no ordering against the running
forward, so a copy could carry pre-write bytes and every later reader saw
a destination that never received those writes -- silently wrong KV, no
crash. The END pools guard the same hazard via `_settle_inflight_forward`.
"""
def _tri(self, lazy=True):
@@ -1318,11 +1186,9 @@ class TestFloatRelocationIsOrderedAgainstTheForward(unittest.TestCase):
self.assertEqual(order[0], "settle", f"first action was not a settle: {order}")
def test_the_settle_is_a_stream_wait_not_a_host_sync(self):
"""Pin the mechanism: `_settle_inflight_forward` must stream-wait, so
the fix costs no host sync on the shortfall path."""
src = inspect.getsource(
mea.MultiEndedAllocator._settle_inflight_forward # noqa: SLF001
)
"""The settle before a float move must be a stream wait: a host sync
there would land on the alloc-shortfall path every time the float moves."""
src = inspect.getsource(MultiEndedAllocator._settle_inflight_forward)
self.assertIn("wait_event", src)
self.assertNotIn(".item()", src)
self.assertNotIn("synchronize()", src)
@@ -1331,14 +1197,12 @@ class TestFloatRelocationIsOrderedAgainstTheForward(unittest.TestCase):
class TestFloatHoleCreditIsPerSide(unittest.TestCase):
"""A float's schedulable credit must follow the side the holes are on.
REGRESSION: the base `_peer_drainable_hole_bytes` asks
`_growth_side_neighbor()`, which reads `grow_direction`. A float's is
"float", so the base fell through to `low_peer` -- it never saw the HIGH
neighbour, and the single scalar it returned was then added to
`max(gap_low, gap_high)`, landing a LOW neighbour's holes on the HIGH gap.
Over-reporting `schedulable_available_size` makes the scheduler admit work
the shortfall ladder cannot satisfy, which the caller treats as a
memory-estimation bug.
REGRESSION: the base `_peer_drainable_hole_bytes` picks its neighbour from
`grow_direction`, which for a float is "float" -- it fell through to
`low_peer`, never saw the HIGH neighbour, and its single scalar was added
to `max(gap_low, gap_high)`, landing a LOW neighbour's holes on the HIGH
gap. Over-reporting `schedulable_available_size` admits work the shortfall
ladder cannot satisfy.
"""
def _float(self):
@@ -1365,15 +1229,6 @@ class TestFloatHoleCreditIsPerSide(unittest.TestCase):
)
return alloc, alloc.swa_attn_allocator
def test_credit_sees_both_neighbours(self):
_alloc, flt = self._float()
self.assertIsInstance(flt, FloatMultiEndedAllocator)
low = flt._side_drainable_hole_bytes("low")
high = flt._side_drainable_hole_bytes("high")
self.assertEqual(flt._peer_drainable_hole_bytes(), max(low, high))
# The base would have answered with the LOW side alone.
self.assertGreaterEqual(flt._peer_drainable_hole_bytes(), high)
def test_schedulable_never_exceeds_the_sum_of_the_two_sides(self):
"""Upper bound that the undirected scalar could violate: no side may be
credited with the other side's holes on top of its own gap."""
test/registered/unit/server_args/test_page_major_backend_allowlist.py
+19 -48
View File
@@ -13,30 +13,19 @@
# ==============================================================================
"""`--enable-page-major-kv-layout` full-attention backend allowlist.
Two-way gate (see `_handle_page_major_kv_layout`), because the unified pool
exposes per-layer views and nothing else:
* unified-memory MLA models (`build_mla_views`) allow the whole wired
paged MLA family -- `fa3`, `flashinfer`'s MLA backend, `trtllm_mla` with
its `cutedsl_mla` / `tokenspeed_mla` subclasses, and `flashmla` (ps=64
snap);
* unified-memory MHA/SWA models (`build_mha_views`) allow `fa3` /
`fa4` / `flashinfer` / `trtllm_mha` alongside Triton;
* plain `--enable-page-major-kv-layout` without the unified pool keeps the
envelope-strided 4-D views only the stride-aware Triton kernels read.
`handle_page_major_kv_layout` gates two ways, because the per-layer views the
unified pool exposes are all the allowlisted backends can read: an MLA arm
(`build_mla_views`), an MHA/SWA arm (`build_mha_views`), and no page-major arm
at all without the unified pool. `fa3` is the resolved default on pre-Blackwell
hosts, so its absence from an arm makes `--enable-unified-memory` fail to boot
under its own default configuration.
The same handler also screens the pool itself: the MHA/SWA per-layer views need
The same handler screens the pool itself: the MHA/SWA per-layer views need
uniform K/V rows, so an asymmetric-K/V model (MiMoV2: head_dim 192 !=
v_head_dim 128) cannot run `--enable-unified-memory` at all and is rejected on
EVERY backend, Triton included. MLA models are exempt -- their sub-pool keeps
one latent row per layer, and several MLA configs (Kimi-Linear: head_dim 72,
v_head_dim 128) report asymmetric dims while running the unified pool today.
Pinned here so no arm silently widens to an unwired backend (`cutlass_mla`,
`aiter`) and no arm silently narrows: `fa3` is the resolved default on
pre-Blackwell hosts, so its absence from an arm makes `--enable-unified-memory`
fail to boot under its own default configuration.
python -m pytest test/registered/unit/server_args/test_page_major_backend_allowlist.py -v
v_head_dim 128) is rejected on EVERY backend, Triton included. MLA models are
exempt -- their sub-pool keeps one latent row per layer, and real MLA configs
(Kimi-Linear: head_dim 72, v_head_dim 128) report asymmetric dims while running
the unified pool today.
"""
import unittest
@@ -59,11 +48,8 @@ def _accepts(
linear_prefill: str | None = None,
has_asymmetric_kv: bool = False,
) -> bool:
"""Run just `_handle_page_major_kv_layout` against a minimal stand-in.
ServerArgs' real constructor pulls in a model config; this exercises the
single handler under test with the fields it reads.
"""
"""Run just `handle_page_major_kv_layout` against a minimal stand-in, since
ServerArgs' real constructor pulls in a model config."""
sa = ServerArgs.__new__(ServerArgs)
for name, value in {
"enable_unified_memory": unified,
@@ -145,21 +131,9 @@ class TestPageMajorBackendAllowlist(unittest.TestCase):
f"{backend} is an MLA kernel and must stay out of the MHA arm",
)
def test_plain_page_major_arm_is_gated_at_boot(self):
"""The strided views were removed: --enable-page-major-kv-layout
without --enable-unified-memory must be rejected up front for EVERY
backend, Triton included, until the per-layer-view reimplementation."""
for backend in ("triton",) + self.PER_LAYER_VIEW_MLA_BACKENDS:
for use_mla in (True, False):
self.assertFalse(
_accepts(backend, use_mla=use_mla, unified=False),
f"{backend} must be rejected on the static page-major arm",
)
def test_asymmetric_kv_mha_model_cannot_use_unified_memory(self):
"""head_dim != v_head_dim (MiMoV2): no uniform rows, so no per-layer views
and no unified pool. The rejection is the POOL's, not a backend's, so
it must fire on every backend -- Triton included."""
"""The rejection is the POOL's, not a backend's, so it must fire on
every backend -- Triton included."""
for backend in ("triton",) + self.PER_LAYER_VIEW_MHA_BACKENDS:
self.assertFalse(
_accepts(backend, use_mla=False, has_asymmetric_kv=True),
@@ -168,10 +142,8 @@ class TestPageMajorBackendAllowlist(unittest.TestCase):
)
def test_asymmetric_dims_do_not_screen_out_mla(self):
"""MLA stores one latent row per layer, so its K/V head dims never have
to agree -- and real MLA configs report them as unequal (Kimi-Linear:
head_dim 72, v_head_dim 128). Screening on `has_asymmetric_kv` alone
would lock every one of them out of the unified pool."""
"""Screening on `has_asymmetric_kv` alone would lock every real MLA
config out of the unified pool."""
for backend in ("triton",) + self.PER_LAYER_VIEW_MLA_BACKENDS:
self.assertTrue(
_accepts(backend, use_mla=True, has_asymmetric_kv=True),
@@ -180,9 +152,8 @@ class TestPageMajorBackendAllowlist(unittest.TestCase):
)
def test_page_major_rejected_without_unified_memory(self):
"""--enable-page-major-kv-layout without the unified pool is rejected
outright, Triton included: the static page-major arm went away with
the strided views and awaits its per-layer-view reimplementation."""
"""There is no static page-major arm today, so the flag alone is
rejected outright -- Triton included."""
for backend in ("triton",) + tuple(
set(self.PER_LAYER_VIEW_MLA_BACKENDS + self.PER_LAYER_VIEW_MHA_BACKENDS)
):