From d50e9a9756f9178049113c65bd1d15c15b1fd442 Mon Sep 17 00:00:00 2001 From: Liangsheng Yin Date: Fri, 4 Sep 2026 20:47:01 -0700 Subject: [PATCH] [Test] Prune redundant unified-memory allocator and pool tests (#38093) --- .../test_kimi_linear_unified_memory.py | 59 +- .../page_major/test_page_major_gpt_oss.py | 16 +- .../page_major/test_page_major_qwen_hybrid.py | 20 +- .../attention/test_kv_translate_ownership.py | 67 +- .../unit/mem_cache/test_full_loc_fast_path.py | 193 ++-- .../mem_cache/test_kv_index_translator.py | 144 +-- .../unit/mem_cache/test_layout_compat.py | 79 +- .../mem_cache/test_multi_ended_allocator.py | 880 ++++-------------- .../test_swa_locked_full_recover_unified.py | 86 +- .../mem_cache/test_swa_pool_v_head_dim.py | 81 +- .../mem_cache/test_unified_byte_accounting.py | 97 +- .../test_unified_byte_budget_sizing.py | 123 +-- .../mem_cache/test_unified_capacity_memo.py | 90 +- .../test_unified_free_no_host_sync.py | 196 ++-- .../mem_cache/test_unified_mamba_views.py | 334 ++----- .../unit/mem_cache/test_unified_mha_views.py | 171 +--- .../mem_cache/test_unified_mla_gpu_parity.py | 58 +- .../unit/mem_cache/test_unified_mla_views.py | 156 ++-- .../mem_cache/test_unified_npool_sweep.py | 149 ++- .../unit/mem_cache/test_unified_tri_pool.py | 341 ++----- .../test_page_major_backend_allowlist.py | 67 +- 21 files changed, 920 insertions(+), 2487 deletions(-) diff --git a/test/registered/models_e2e/test_kimi_linear_unified_memory.py b/test/registered/models_e2e/test_kimi_linear_unified_memory.py index b8e733a28..72c907316 100644 --- a/test/registered/models_e2e/test_kimi_linear_unified_memory.py +++ b/test/registered/models_e2e/test_kimi_linear_unified_memory.py @@ -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 diff --git a/test/registered/page_major/test_page_major_gpt_oss.py b/test/registered/page_major/test_page_major_gpt_oss.py index 4bb46a411..54992e93e 100644 --- a/test/registered/page_major/test_page_major_gpt_oss.py +++ b/test/registered/page_major/test_page_major_gpt_oss.py @@ -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() diff --git a/test/registered/page_major/test_page_major_qwen_hybrid.py b/test/registered/page_major/test_page_major_qwen_hybrid.py index a2e6462a7..ea2a1d7f6 100644 --- a/test/registered/page_major/test_page_major_qwen_hybrid.py +++ b/test/registered/page_major/test_page_major_qwen_hybrid.py @@ -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() diff --git a/test/registered/unit/layers/attention/test_kv_translate_ownership.py b/test/registered/unit/layers/attention/test_kv_translate_ownership.py index 020823063..6a78b9de0 100644 --- a/test/registered/unit/layers/attention/test_kv_translate_ownership.py +++ b/test/registered/unit/layers/attention/test_kv_translate_ownership.py @@ -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( diff --git a/test/registered/unit/mem_cache/test_full_loc_fast_path.py b/test/registered/unit/mem_cache/test_full_loc_fast_path.py index eefb87115..438f65453 100644 --- a/test/registered/unit/mem_cache/test_full_loc_fast_path.py +++ b/test/registered/unit/mem_cache/test_full_loc_fast_path.py @@ -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 diff --git a/test/registered/unit/mem_cache/test_kv_index_translator.py b/test/registered/unit/mem_cache/test_kv_index_translator.py index cdb010f31..e6b40408e 100644 --- a/test/registered/unit/mem_cache/test_kv_index_translator.py +++ b/test/registered/unit/mem_cache/test_kv_index_translator.py @@ -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 diff --git a/test/registered/unit/mem_cache/test_layout_compat.py b/test/registered/unit/mem_cache/test_layout_compat.py index 3b5ca7e05..506c505ad 100644 --- a/test/registered/unit/mem_cache/test_layout_compat.py +++ b/test/registered/unit/mem_cache/test_layout_compat.py @@ -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): diff --git a/test/registered/unit/mem_cache/test_multi_ended_allocator.py b/test/registered/unit/mem_cache/test_multi_ended_allocator.py index b7284974d..363f1c25e 100644 --- a/test/registered/unit/mem_cache/test_multi_ended_allocator.py +++ b/test/registered/unit/mem_cache/test_multi_ended_allocator.py @@ -11,13 +11,11 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -"""Unit tests for the shared-KV-pool v2 core: UnifiedKVPool views and -MultiEndedAllocator (virtual<->physical slot ids + eager compaction). +"""Unit tests for UnifiedKVPool views and the unified sub-pool allocators: +virtual<->physical id mapping, compaction, and the SWA / Mamba composites. -CPU-only — no GPU / Triton needed (the allocator's data-copy delegates to a -fake kvcache here; the UnifiedKVPool view math is pure torch). - - python -m pytest test/registered/unit/mem_cache/test_multi_ended_allocator.py -v +Data copies go through a fake kvcache and the view math is pure torch, so the +suite runs on CPU apart from the cases that skip themselves without CUDA. """ from sglang.test.ci.ci_register import register_cpu_ci @@ -75,10 +73,8 @@ def _make_mamba_spec(name, grow, layer_num=2): class _FakeKVCache: - """Tracks, per *physical* slot, the virtual id whose data lives there. - `move_kv_cache(dst, src)` copies the marker — so after compaction we can - check that the data followed the relocation. - """ + """Tracks, per physical slot, the virtual id whose data lives there, so a + test can assert the data followed a compaction move.""" def __init__(self, max_slots: int): # buf[p] == virtual id currently stored at physical slot p (-1 if free). @@ -138,12 +134,8 @@ class TestUnifiedKVPoolViews(unittest.TestCase): k_swa, v_swa = pool.mha_views_for("swa") self.assertEqual(len(k_full), 3) self.assertEqual(len(k_swa), 2) - # Write distinct patterns into a couple of slots/layers of "full" and - # confirm they read back, and that "swa" was not disturbed. "full" and - # "swa" share the buffer from byte 0; "full" grows up (low slots) and - # "swa" grows down, so the allocator places "swa" at the high slots. - # Mirror that here — a low "swa" slot would byte-overlap "full" slot 5 - # (a configuration the byte-frontier coordination never produces). + # "full" and "swa" share the buffer from byte 0 and grow toward each + # other, so a low "swa" slot would byte-overlap "full" slot 5. swa_slot = pool.max_slots("swa") - 1 for lyr in range(3): k_full[lyr][5] = float(lyr + 1) @@ -232,7 +224,7 @@ class TestMultiEndedAllocator(unittest.TestCase): self.assertEqual(alloc_hi - alloc_lo, len(live_v)) for p in range(alloc_lo, alloc_hi): self.assertNotEqual(int(p2v[p].item()), -1, f"hole at physical {p}") - # free virtual ids ∪ live = [min_slot_index, max_slots) + # free virtual ids | live = [min_slot_index, max_slots) free_set = set(int(x) for x in alloc.free_virtual_ids.tolist()) self.assertEqual( free_set | set(live_v), @@ -278,19 +270,6 @@ class TestMultiEndedAllocator(unittest.TestCase): self._check_invariants(full_alloc, full_kv) self.assertEqual(full_alloc.allocated_count(), 0) - def test_grow_down_side(self): - _, full_alloc, mamba_alloc, full_kv, mamba_kv = self._build_pair() - a = self._alloc(mamba_alloc, mamba_kv, 2) - b = self._alloc(mamba_alloc, mamba_kv, 3) - c = self._alloc(mamba_alloc, mamba_kv, 1) - self._check_invariants(mamba_alloc, mamba_kv) - self._free(mamba_alloc, mamba_kv, b) # interior -> compaction - self._check_invariants(mamba_alloc, mamba_kv) - self._free(mamba_alloc, mamba_kv, a) - self._free(mamba_alloc, mamba_kv, c) - self._check_invariants(mamba_alloc, mamba_kv) - self.assertEqual(mamba_alloc.allocated_count(), 0) - def test_byte_frontier_coordination(self): # full has 8 slots' worth of bytes; mamba's entry is larger, so a few # mamba allocs should shrink full's available_size below its slot headroom. @@ -339,9 +318,8 @@ class TestMultiEndedAllocator(unittest.TestCase): self.assertEqual(alloc.allocated_count(), 0) def _build_lazy_full(self, n_full_slots=64, n_mamba_slots=16, move_cap=2): - """A LAZY-compaction 'full' (grow-up) allocator + its peer, with a - small per-call move cap so flushes are partial (the retract-pressure - regime where the ghost bug appears).""" + """A lazy-compaction 'full' (grow-up) allocator + its peer, with a + small per-call move cap so every flush is partial.""" full = _make_mha_spec("full", "up", layer_num=2) mamba = _make_mamba_spec("mamba", "down", layer_num=2) total = full.entry_bytes() * n_full_slots + mamba.entry_bytes() * n_mamba_slots @@ -375,22 +353,12 @@ class TestMultiEndedAllocator(unittest.TestCase): return pool, full_alloc, full_kv def test_lazy_retract_churn_no_ghost(self): - """Regression: heavy free/flush churn in LAZY mode must never - leave a 'ghost' page (p2v<0 not registered in _free_phys_pages or - _pending_reuse). This reproduces the retract×lazy-compaction pattern. - The fail-fast ghost invariant is enabled so any ghost trips at the - CREATING op (free_lazy / flush_exit), not at a later survivor walk. - - CPU scope: the GPU write-race *urgent* path (unfired CUDA events -> - _pending_reuse) needs a real device; this covers the no-event - move/absorb/free/released_fired bookkeeping. - """ + """Regression: heavy free/flush churn in lazy mode must never leave a + 'ghost' page -- p2v < 0 in neither _free_phys_pages nor _pending_reuse.""" rng = random.Random(0x5373) _, alloc, kv = self._build_lazy_full(n_full_slots=64, move_cap=2) live = [] - # Bias toward near-capacity occupancy (watermark high, holes pile - # up) then churn free/flush — the conditions that surface the bug. - for _ in range(3000): + for _ in range(600): avail = alloc.available_size() if (rng.random() < 0.55 and avail > 0) or not live: n = rng.randint(1, min(5, max(1, avail))) @@ -419,9 +387,8 @@ class TestMultiEndedAllocator(unittest.TestCase): # -- `out=` parameter regression tests -- def test_translate_kv_loc_with_out_writes_inplace(self): - """REGRESSION: `translate_kv_loc(virt, out=buf)` must - modify `buf` in place AND preserve `buf.data_ptr()` — the buffer- - stability invariant for cuda-graph capture.""" + """Regression: `translate_kv_loc(virt, out=buf)` must write into `buf` + and preserve `buf.data_ptr()` -- the cuda-graph capture invariant.""" _, full_alloc, _, full_kv, _ = self._build_pair() v = self._alloc(full_alloc, full_kv, 5) buf = torch.empty(v.shape, dtype=torch.int64, device=_DEV) @@ -435,45 +402,22 @@ class TestMultiEndedAllocator(unittest.TestCase): expected = full_alloc.virtual_to_physical[v] self.assertTrue(bool((buf == expected).all().item())) - def test_translate_kv_loc_out_matches_no_out(self): - """REGRESSION: result of translate_kv_loc(v, out=buf) byte-equals - translate_kv_loc(v).""" + def test_translate_kv_loc_out_guards(self): + """REGRESSION: a malformed `out=` buffer raises AssertionError -- it + must match both the v2p dtype the gather writes and the input shape.""" _, full_alloc, _, full_kv, _ = self._build_pair() v = self._alloc(full_alloc, full_kv, 5) - buf = torch.empty(v.shape, dtype=torch.int64, device=_DEV) - with_out = full_alloc.translate_kv_loc(v, out=buf) - no_out = full_alloc.translate_kv_loc(v) - self.assertTrue(bool((with_out == no_out).all().item())) + with self.subTest(guard="dtype"): + wrong_dtype = torch.empty(v.shape, dtype=torch.int32, device=_DEV) + with self.assertRaises(AssertionError): + full_alloc.translate_kv_loc(v, out=wrong_dtype) + with self.subTest(guard="shape"): + wrong_shape = torch.empty((v.numel() + 1,), dtype=torch.int64, device=_DEV) + with self.assertRaises(AssertionError): + full_alloc.translate_kv_loc(v, out=wrong_shape) - def test_translate_kv_loc_dtype_assertion(self): - """REGRESSION: wrong-dtype `out=` (int32 instead of int64) raises - AssertionError -- `out=` must match the v2p dtype the gather writes.""" - _, full_alloc, _, full_kv, _ = self._build_pair() - v = self._alloc(full_alloc, full_kv, 5) - wrong_dtype = torch.empty(v.shape, dtype=torch.int32, device=_DEV) - with self.assertRaises(AssertionError): - full_alloc.translate_kv_loc(v, out=wrong_dtype) - - def test_translate_kv_loc_shape_assertion(self): - """REGRESSION: mismatched `out=` shape raises AssertionError.""" - _, full_alloc, _, full_kv, _ = self._build_pair() - v = self._alloc(full_alloc, full_kv, 5) - wrong_shape = torch.empty((v.numel() + 1,), dtype=torch.int64, device=_DEV) - with self.assertRaises(AssertionError): - full_alloc.translate_kv_loc(v, out=wrong_shape) - - # REGRESSION: `translate_kv_loc(buf, out=buf)` — same - # tensor for input and output — is the canonical in-place form used by - # the cuda-graph capture/replay paths in `triton_backend.py`: - # - # self._translate_kv_loc(kv_indices, out=kv_indices) - # - # A naive implementation routes this through - # `torch.index_select(v2p, 0, virt_tokens, out=out)`, which crashes with - # "unsupported operation: some elements of the input tensor and the - # written-to tensor refer to a single memory location" - # because index_select does NOT support aliasing between `index` and - # `out`. Fix: gather into a transient buffer then `out.copy_(tmp)`. + # `index_select(v2p, 0, virt, out=out)` rejects aliasing between `index` + # and `out`, so the in-place form must gather into a temporary first. def test_translate_kv_loc_with_out_aliasing_input(self): """REGRESSION: in-place form `translate_kv_loc(buf, out=buf)` must succeed and produce identical results to the no-out form.""" @@ -496,29 +440,15 @@ class TestMultiEndedAllocator(unittest.TestCase): "in-place result must equal no-out result", ) - # Tombstone-safety clamp regression. - # - # The captured cuda-graph paths (full-layer set_kv_buffer elif, - # init_forward_metadata_*_cuda_graph kv_indices translate, init_new - # precompute) all eventually call `translate_kv_loc` against `v2p_full`. - # Padded / stale-tail entries in the cuda-graph input buffers can carry - # virtual ids whose v2p entries got tombstoned (-1) by free/compaction - # between replays. Without a clamp, the captured `k_buffer[result[i]]` - # would index at -1 (illegal memory access). The clamp routes those to - # physical slot 0 (the reserved padding sink under the `min_slot_index` - # invariant — bytes [0, entry_max) hold no real data). - # These tests lock in the clamp contract so a future refactor can't - # quietly remove it and re-introduce the crash. + # A captured cuda-graph input buffer can hold virtual ids tombstoned (-1) + # between replays; slot 0 is the sink, as bytes [0, entry_max) hold no data. def test_translate_kv_loc_clamps_tombstoned_v2p(self): """`translate_kv_loc` must clamp `v2p[v] == -1` entries to 0 (the padding sink). Required for cuda-graph capture safety.""" _, full_alloc, _, full_kv, _ = self._build_pair() v = self._alloc(full_alloc, full_kv, 5) - # Inject a tombstone at one of the live virtual id positions WITHOUT - # going through `free` (which would also touch p2v / compaction). - # This emulates the steady-state where a captured-graph input - # buffer's padded/stale entries reference virtual ids that have - # since been tombstoned by free/compaction. + # Inject a tombstone directly, not through `free` (which would also + # touch p2v and run compaction). v_tombstoned = int(v[2].item()) full_alloc.virtual_to_physical[v_tombstoned] = -1 # No-out form: result must clamp. @@ -532,13 +462,7 @@ class TestMultiEndedAllocator(unittest.TestCase): 0, "tombstoned virtual id must map to slot 0 (padding sink)", ) - - def test_translate_kv_loc_with_out_clamps_tombstoned_v2p(self): - """`translate_kv_loc(..., out=buf)` (the captured-graph path) must - clamp tombstoned entries in-place.""" - _, full_alloc, _, full_kv, _ = self._build_pair() - v = self._alloc(full_alloc, full_kv, 5).clone() - full_alloc.virtual_to_physical[int(v[1].item())] = -1 + # out= form (the captured-graph path) must clamp in place too. buf = torch.empty_like(v) ret = full_alloc.translate_kv_loc(v, out=buf) self.assertIs(ret, buf) @@ -546,16 +470,12 @@ class TestMultiEndedAllocator(unittest.TestCase): bool((buf >= 0).all().item()), "out= path must clamp tombstoned entries", ) - self.assertEqual(int(buf[1].item()), 0) + self.assertEqual(int(buf[2].item()), 0) def test_slot_zero_sink_invariant_survives_churn(self): - """PINNED INVARIANT: virtual 0 <-> physical 0 (the padding sink), so - `translate_kv_loc(zeros) == zeros` -- after init AND after alloc/free/ - compaction churn. The cuda-graph capture path RELIES on this: the - physical-loc contract replaced capture-time translate with a plain - copy of the zero-filled static buffer, which is only equivalent while - v2p[0] == 0. If an allocator change breaks this, captured stores would - write pad lanes to a live slot.""" + """Virtual 0 must stay mapped to physical 0 through alloc/free/compaction + churn: cuda-graph capture copies a zero-filled static buffer instead of + translating, so a nonzero v2p[0] would send pad lanes to a live slot.""" _, full_alloc, _, full_kv, _ = self._build_pair() zeros = torch.zeros(4, dtype=torch.int64) @@ -575,18 +495,13 @@ class TestMultiEndedAllocator(unittest.TestCase): # --------------------------------------------------------------------------- -# Shared SWA composite — unit tests +# Shared SWA composite -- unit tests # --------------------------------------------------------------------------- class _FakeUnifiedSWAKVPool: - """Minimal stand-in for `UnifiedSWAKVPool` that the composite allocator - needs. Exposes the two sub-pool views (each a `_FakeKVCache` with an - `attach_allocator` no-op) and an `attach_allocators` setter. - - CPU-only — avoids constructing a real `UnifiedMHATokenToKVPool` (which - instantiates `MHATokenToKVPool` and is heavier than these tests need). - """ + """Minimal stand-in for `UnifiedSWAKVPool`: a real one would construct + `UnifiedMHATokenToKVPool`, which is heavier than these tests need.""" class _SubKV(_FakeKVCache): def __init__(self, max_slots): @@ -608,13 +523,8 @@ class _FakeUnifiedSWAKVPool: class TestUnifiedSWATokenToKVPoolAllocator(unittest.TestCase): - """Tests for the SWA composite — joint byte-budget, slot-conservation - leak invariant, tombstone semantics for `free_swa`, divergent compaction - of the two sub-pools, and the alloc-rollback path. - - These tests cover the core invariants: joint byte-budget, - slot-conservation, the `schedulable_*` split, and watermark - rollback.""" + """The SWA composite: joint byte-budget, slot-conservation, `free_swa` + tombstone semantics, and divergent compaction of the two sub-pools.""" def _build( self, @@ -675,8 +585,8 @@ class TestUnifiedSWATokenToKVPoolAllocator(unittest.TestCase): return v def _free(self, allocator, kvcache, v): - """Erase markers on both sub-pools (mirror compaction's no-data-at - -freed-slot invariant), then call the composite's free.""" + """Erase markers on both sub-pools (mirroring compaction's + no-data-at-a-freed-slot invariant), then call the composite's free.""" full_phys = allocator.full_attn_allocator.virtual_to_physical[v] swa_phys = allocator.swa_attn_allocator.virtual_to_physical[v] # erase only the LIVE swa entries (`free_swa` may have already @@ -687,7 +597,7 @@ class TestUnifiedSWATokenToKVPoolAllocator(unittest.TestCase): allocator.free(v) def _check_sub_pool_invariants(self, sub, kv): - """Per-sub-pool: v2p ∘ p2v identity on the live set, hole-free + """Per-sub-pool: v2p/p2v mutual inverse on the live set, hole-free allocated band, data followed relocations.""" v2p = sub.virtual_to_physical p2v = sub.physical_to_virtual @@ -724,8 +634,7 @@ class TestUnifiedSWATokenToKVPoolAllocator(unittest.TestCase): # Swa sub-pool is non-owner -> free_virtual_ids is None. self.assertIsNone(allocator.swa_attn_allocator.free_virtual_ids) - # 2. Composite `free` releases both sub-pools' v2p; the virtual goes back - # to the full id-owner's free list. + # 2. Composite `free` releases both sub-pools' v2p and recycles the virtual. def test_swa_free_releases_both(self): _, allocator, kvcache = self._build() v = self._alloc(allocator, kvcache, 3) @@ -822,9 +731,8 @@ class TestUnifiedSWATokenToKVPoolAllocator(unittest.TestCase): int(allocator.full_attn_allocator.virtual_to_physical[c].item()), c_full_before, ) - # c's swa-physical MUST have moved (b was interior to swa's - # allocated band on grow-down: a then b then c means b is between - # them; freeing b triggers compaction relocating c into b's slot). + # c's swa-physical MUST have moved: b is interior to swa's allocated + # band, so freeing it relocates c into b's slot. c_swa_after = int(allocator.swa_attn_allocator.virtual_to_physical[c].item()) self.assertNotEqual(c_swa_after, c_swa_before) # Per-sub-pool invariants still hold. @@ -835,8 +743,7 @@ class TestUnifiedSWATokenToKVPoolAllocator(unittest.TestCase): allocator.swa_attn_allocator, kvcache.swa_kv_pool ) - # 5. Byte-frontier coordination — peer-aware available_size shrinks as - # the peer grows. + # 5. Byte-frontier coordination: available_size shrinks as the peer grows. def test_swa_byte_frontier_coordination(self): _, allocator, kvcache = self._build(n_full_slots=8, n_swa_slots=8) avail0 = allocator.available_size() @@ -846,7 +753,7 @@ class TestUnifiedSWATokenToKVPoolAllocator(unittest.TestCase): # Joint budget enforcement: over-alloc returns None. self.assertIsNone(allocator.alloc(allocator.available_size() + 1)) - # 6. Randomized stress — invariants under mixed alloc / free / free_swa. + # 6. Randomized stress -- invariants under mixed alloc / free / free_swa. def test_swa_randomized_alloc_free_freeswa(self): rng = random.Random(0xBADBEE) _, allocator, kvcache = self._build( @@ -881,10 +788,8 @@ class TestUnifiedSWATokenToKVPoolAllocator(unittest.TestCase): self._check_sub_pool_invariants( allocator.swa_attn_allocator, kvcache.swa_kv_pool ) - # Slot-conservation invariant balances at all times. NOTE: the - # leak view is now `_conserve_*` — the public `full/swa_available_size()` - # returns `min(conserve, schedulable)` (physical), which can be - # strictly smaller (e.g. the reserved sink page). + # The leak view is `_conserve_*`; public `*_available_size()` is + # `min(conserve, schedulable)` and can be strictly smaller. self.assertEqual( allocator._conserve_full_available_size(), allocator._full_max_total_num_tokens @@ -901,50 +806,11 @@ class TestUnifiedSWATokenToKVPoolAllocator(unittest.TestCase): self.assertEqual(allocator.full_attn_allocator.allocated_count(), 0) self.assertEqual(allocator.swa_attn_allocator.allocated_count(), 0) - # 8. Watermark rollback on partial alloc failure. - def test_swa_alloc_swa_failure_is_fail_loud(self): - """The SWA composite runs a tight JOINT pre-check before allocating, so - a swa-side ``alloc_with_virtual`` failure after the full-side alloc can - only mean an internal-state inconsistency. By design (``UnifiedSWA.alloc``: - "assert rather than silently rollback") that surfaces as a loud error, - NOT a silent ``None`` / rollback — masking it would hide the bug. The - real ``alloc_with_virtual`` self-asserts on shortfall, so the production - path is fail-loud too; here we force the failure to prove it propagates. - """ - _, allocator, kvcache = self._build() - sa = allocator.swa_attn_allocator - original = sa.alloc_with_virtual - - def _bomb(virtual_ids): - raise AssertionError("synthetic alloc_with_virtual failure") - - sa.alloc_with_virtual = _bomb - try: - with self.assertRaises(AssertionError): - allocator.alloc(3) - finally: - sa.alloc_with_virtual = original - # -- `out=` parameter regression tests for the SWA composite -- - def test_swa_translate_kv_loc_with_out_writes_inplace(self): - """REGRESSION: composite delegates to base-class - translate_kv_loc with `out=` passthrough. Result lands in `buf`.""" - _, allocator, _ = self._build() - v = allocator.alloc(4) - self.assertIsNotNone(v) - buf = torch.empty(v.shape, dtype=torch.int64, device=_DEV) - ptr_before = buf.data_ptr() - ret = allocator.translate_kv_loc(v, out=buf) - self.assertIs(ret, buf) - self.assertEqual(buf.data_ptr(), ptr_before) - expected = allocator.translate_kv_loc(v) - self.assertTrue(bool((buf == expected).all().item())) - def test_swa_translate_loc_from_full_to_swa_with_out_writes_inplace(self): - """REGRESSION: `translate_loc_from_full_to_swa(v, out=buf)` - must modify `buf` in place AND preserve `buf.data_ptr()`. `out=` - buffer is int64 — every id the allocator emits is.""" + """Regression: `translate_loc_from_full_to_swa(v, out=buf)` must write + into `buf` and preserve its data_ptr; `out=` is int64, as every id is.""" _, allocator, _ = self._build() v = allocator.alloc(4) self.assertIsNotNone(v) @@ -959,9 +825,8 @@ class TestUnifiedSWATokenToKVPoolAllocator(unittest.TestCase): self.assertTrue(bool((buf == no_out).all().item())) def test_swa_translate_loc_from_full_to_swa_dtype_assertion(self): - """REGRESSION: wrong-dtype `out=` (int32 instead of int64) raises - AssertionError. Guards against reintroducing a narrowed SWA write loc: the - allocator emits int64 and consumers narrow at their own buffer.""" + """Regression: a wrong-dtype `out=` (int32) must raise; the allocator + emits int64 and consumers narrow at their own buffer.""" _, allocator, _ = self._build() v = allocator.alloc(4) self.assertIsNotNone(v) @@ -969,10 +834,8 @@ class TestUnifiedSWATokenToKVPoolAllocator(unittest.TestCase): with self.assertRaises(AssertionError): allocator.translate_loc_from_full_to_swa(v, out=wrong_dtype) - # Tombstone-safety clamp for SWA — mirrors the full-side test - # in `TestMultiEndedAllocator`. The captured SWA attention kernel reads - # `swa_k_buffer[result[i]]` at replay; without the clamp, a tombstoned - # `v2p_swa[v] == -1` would index at `swa_k_buffer[-1]` (illegal access). + # SWA side of the tombstone clamp: a tombstoned `v2p_swa[v] == -1` would + # make the captured kernel read `swa_k_buffer[-1]`. def test_swa_translate_loc_from_full_to_swa_clamps_tombstoned(self): _, allocator, _ = self._build() v = allocator.alloc(4) @@ -996,13 +859,9 @@ class TestUnifiedSWATokenToKVPoolAllocator(unittest.TestCase): self.assertEqual(int(buf[1].item()), 0) def test_swa_slot_zero_sink_invariant_survives_churn(self): - """PINNED INVARIANT (swa side of the physical-loc contract): BOTH maps - send virtual 0 to physical 0 — `translate_kv_loc(zeros) == zeros` AND - `translate_loc_from_full_to_swa(zeros) == zeros` — after init and - after alloc/free/free_swa churn. Cuda-graph capture replaced the - capture-time translate with zero-fill/copy of the zero-filled static - buffers; that is only equivalent while slot 0 stays the sink in both - sub-pools.""" + """Both maps must send virtual 0 to physical 0 through churn: cuda-graph + capture zero-fills its static buffers instead of translating, which is + only equivalent while slot 0 stays the sink in both sub-pools.""" _, allocator, kvcache = self._build() zeros64 = torch.zeros(4, dtype=torch.int64) @@ -1024,18 +883,13 @@ class TestUnifiedSWATokenToKVPoolAllocator(unittest.TestCase): # --------------------------------------------------------------------------- -# page_size > 1 — paged unit tests +# page_size > 1 -- paged unit tests # --------------------------------------------------------------------------- class TestPagedMultiEndedAllocator(unittest.TestCase): - """Per-sub-pool paged tests for `MultiEndedAllocator(page_size=...)`. - - All tests use ``page_size = 8`` against a buffer sized for ~16 pages per - sub-pool. Invariants are page-granular: free-list, v2p/p2v tables, and - compaction operate on pages. The external API (alloc → token ids, free - takes token ids) is byte-identical to the page_size == 1 case. - """ + """`MultiEndedAllocator(page_size=8)`: free-list, v2p/p2v and compaction are + page-granular, while the external API stays in token ids as at page_size 1.""" PAGE_SIZE = 8 @@ -1094,8 +948,7 @@ class TestPagedMultiEndedAllocator(unittest.TestCase): def _stamp_tokens( self, alloc: MultiEndedAllocator, kv: _FakeKVCache, v_tokens: torch.Tensor ): - """Mark `kv.buf[phys_token] = some_unique_id` for every returned - token. Uses the alloc's v2p_page table to compute physical tokens.""" + """Stamp each returned token's own id into `kv.buf` at its physical token.""" if v_tokens.numel() == 0: return ps = alloc.page_size @@ -1136,7 +989,7 @@ class TestPagedMultiEndedAllocator(unittest.TestCase): self.assertNotEqual( int(p2v[p_page].item()), -1, f"hole at physical page {p_page}" ) - # Free virtual page ids ∪ live = [min_page_index, num_pages). + # Free virtual page ids | live = [min_page_index, num_pages). free_set = set(int(x) for x in alloc.free_virtual_ids.tolist()) self.assertEqual( free_set | set(live_v_pages), @@ -1161,7 +1014,7 @@ class TestPagedMultiEndedAllocator(unittest.TestCase): # 1. alloc(N) returns N TOKEN ids that are page-aligned. def test_paged_alloc_token_aligned(self): _, full_alloc, swa_alloc, full_kv, swa_kv = self._build() - v = full_alloc.alloc(16) # 2 pages × 8 tokens + v = full_alloc.alloc(16) # 2 pages x 8 tokens self.assertIsNotNone(v) self.assertEqual(int(v.numel()), 16) # The output must consist of exactly 2 contiguous page-ranges. @@ -1181,21 +1034,6 @@ class TestPagedMultiEndedAllocator(unittest.TestCase): with self.assertRaises(AssertionError): full_alloc.alloc(5) # not a multiple of 8 - # 3. v2p / p2v tables are sized by PAGES. - def test_paged_v2p_sized_by_pages(self): - pool, full_alloc, _, _, _ = self._build(n_full_pages=10) - # +1 for the trailing -1 sentinel row. - self.assertEqual( - int(full_alloc.virtual_to_physical.numel()), - full_alloc.num_pages + 1, - ) - self.assertEqual( - int(full_alloc.physical_to_virtual.numel()), - full_alloc.num_pages + 1, - ) - # `num_pages` should be > 1 to be a meaningful test. - self.assertGreater(full_alloc.num_pages, 1) - # 4. Compaction relocates a whole page at once (data follows). def test_paged_compaction_relocates_whole_pages(self): _, full_alloc, _, full_kv, _ = self._build() @@ -1205,9 +1043,7 @@ class TestPagedMultiEndedAllocator(unittest.TestCase): b = full_alloc.alloc(self.PAGE_SIZE) # tokens of page Y (middle) c = full_alloc.alloc(self.PAGE_SIZE) # tokens of page Z - # Stamp each token with a UNIQUE marker. (alloc returns unique virtuals, - # but we want each token to be distinguishable from its in-page - # siblings, so we use the virtual-token value itself.) + # Stamp each token with its own virtual id so in-page siblings differ. for v in (a, b, c): self._stamp_tokens(full_alloc, full_kv, v) for t in v.tolist(): @@ -1218,11 +1054,6 @@ class TestPagedMultiEndedAllocator(unittest.TestCase): full_alloc.free(b) for t in b.tolist(): stamped.pop(t, None) - # Erase markers for the freed page in the fake kv buf so the - # invariant check doesn't see stale data. - # (The compaction kernel moved the survivor's data; we don't manually - # touch full_kv.buf for the freed page — the test below verifies that - # `c`'s data followed the relocation.) self._check_invariants(full_alloc, full_kv, stamped) # `a` and `c` pages must still be live. for t in a.tolist(): @@ -1232,12 +1063,8 @@ class TestPagedMultiEndedAllocator(unittest.TestCase): v_page = t // self.PAGE_SIZE self.assertNotEqual(int(full_alloc.virtual_to_physical[v_page].item()), -1) - # 5. free() recovers pages via unique(// page_size) — matches upstream. - # REGRESSION: `allocated_count()` MUST return - # TOKENS, not pages -- matching upstream's convention that all external - # capacity methods report tokens. At page_size > 1, returning pages - # here breaks the leak invariant - # (`available + evictable + ... == total`, with all terms in tokens). + # 5. Regression: `allocated_count()` reports TOKENS, not pages; the leak + # invariant `available + evictable + ... == total` is all in tokens. def test_paged_free_unique_by_page(self): _, full_alloc, _, full_kv, _ = self._build() a = full_alloc.alloc(self.PAGE_SIZE * 2) # 2 pages = 2*PS tokens @@ -1246,8 +1073,7 @@ class TestPagedMultiEndedAllocator(unittest.TestCase): self.assertEqual(allocated_count_before, 2 * self.PAGE_SIZE) # Internal page count. self.assertEqual(full_alloc._allocated_pages(), 2) - # Free a SUBSET of tokens — but covering all tokens of both pages. - # (Matches the upstream contract: caller passes coherent ranges.) + # The caller must pass page-coherent ranges; `a` covers both pages whole. full_alloc.free(a) self.assertEqual(full_alloc.allocated_count(), 0) self.assertEqual(full_alloc._allocated_pages(), 0) @@ -1255,8 +1081,6 @@ class TestPagedMultiEndedAllocator(unittest.TestCase): # 6. take_physical overflow check (grow-up direction). def test_paged_take_physical_overflow_check(self): _, full_alloc, _, _, _ = self._build(n_full_pages=4) - # Try to take more pages than the buffer can hold; should return None. - # First, fill normally up to the available_size, then over-alloc by 1. avail = full_alloc.available_size() n_pages = avail // self.PAGE_SIZE result = full_alloc.take_physical(n_pages * self.PAGE_SIZE) @@ -1309,9 +1133,7 @@ class TestPagedMultiEndedAllocator(unittest.TestCase): need_sort=False, forward_stream=None, ) - # available_size() returns TOKENS. The joint byte-budget at page - # granularity uses `entry_sum_per_page = entry_full_per_page + - # entry_swa_per_page`. Pre-check: + # available_size() is in TOKENS; the joint budget charges both sub-pools. fa = allocator.full_attn_allocator sa = allocator.swa_attn_allocator entry_sum_pp = fa.entry_bytes_per_page + sa.entry_bytes_per_page @@ -1326,17 +1148,15 @@ class TestPagedMultiEndedAllocator(unittest.TestCase): * self.PAGE_SIZE ) self.assertEqual(allocator.available_size(), expected) - # And it's strictly less than min(fa.available_size, sa.available_size) - # (since the joint cost is heavier than either single-side cost). + # The joint cost is heavier than either single-side cost, so the joint + # budget can never exceed either sub-pool's own available_size. self.assertLessEqual( allocator.available_size(), min(fa.available_size(), sa.available_size()), ) - # 9. REGRESSION: alloc_extend must bind v2p / p2v on - # this allocator. Without binding, `virtual_to_physical[virt_page]` - # stays -1 and `translate_kv_loc(virt_token)` returns negative token - # ids → CUDA OOB in the Triton attention kernel. + # 9. Regression: alloc_extend must bind v2p / p2v; an unbound page leaves + # v2p at -1, and translate_kv_loc then emits negative (OOB) token ids. def test_paged_alloc_extend_binds_v2p_p2v(self): from sglang.srt.mem_cache.allocator import unified_sub_pool as mea_mod @@ -1346,10 +1166,8 @@ class TestPagedMultiEndedAllocator(unittest.TestCase): watermark_before = full_alloc.watermark_physical allocated_count_before = full_alloc.allocated_count() - # Stub the kernel — we only need to verify the BINDING contract. - # (Driving the real Triton kernel needs a GPU; the contract we're - # checking is that the v2p/p2v tables get updated regardless of - # what the kernel writes into out_indices.) + # Stub the kernel: driving the real Triton one needs a GPU, and the + # binding contract holds regardless of what it writes to out_indices. original_kernel = mea_mod.alloc_extend_kernel class _NoOpKernelGrid: @@ -1381,8 +1199,6 @@ class TestPagedMultiEndedAllocator(unittest.TestCase): mea_mod.alloc_extend_kernel = original_kernel self.assertIsNotNone(out) - # The two virtual pages consumed from the front of free_virtual_ids - # must now be BOUND in v2p_page (not -1). consumed_pages = free_before[:2] v2p_values = full_alloc.virtual_to_physical[consumed_pages] for v_page, p_page in zip(consumed_pages.tolist(), v2p_values.tolist()): @@ -1395,9 +1211,8 @@ class TestPagedMultiEndedAllocator(unittest.TestCase): # And p2v_page must round-trip. for v_page, p_page in zip(consumed_pages.tolist(), v2p_values.tolist()): self.assertEqual(int(full_alloc.physical_to_virtual[p_page].item()), v_page) - # Watermark must have advanced by 2 pages. - # `allocated_count()` returns TOKENS, so it - # advances by 2 * PAGE_SIZE; `_allocated_pages()` is the page count. + # `allocated_count()` is in TOKENS (advances by 2 * PAGE_SIZE); + # `_allocated_pages()` is the page count. self.assertEqual( full_alloc.allocated_count(), allocated_count_before + 2 * PS, @@ -1416,10 +1231,8 @@ class TestPagedMultiEndedAllocator(unittest.TestCase): int(free_before.numel()) - 2, ) - # 10. REGRESSION: alloc_decode must bind v2p / p2v on - # this allocator when num_new_pages > 0. Most decode steps reuse the - # prefix's tail page (num_new_pages == 0), but the page-wrapping case - # must update tables. + # 10. Regression: alloc_decode must bind v2p / p2v when it wraps to a new + # page; most decode steps reuse the prefix's tail page and bind nothing. def test_paged_alloc_decode_binds_v2p_p2v_on_page_wrap(self): from sglang.srt.mem_cache.allocator import unified_sub_pool as mea_mod @@ -1433,9 +1246,8 @@ class TestPagedMultiEndedAllocator(unittest.TestCase): watermark_before = full_alloc.watermark_physical allocated_count_before = full_alloc.allocated_count() - # Build a decode that wraps to a new page: seq_len % page_size == 1 - # (one req that just stepped past a page boundary). The kernel will - # consume 1 new page from `free_virtual_ids[0]`. + # A decode wraps to a new page exactly when seq_len % page_size == 1; + # the kernel then consumes `free_virtual_ids[0]`. seq_lens = torch.tensor([PS + 1], dtype=torch.int64, device=_DEV) seq_lens_cpu = torch.tensor([PS + 1], dtype=torch.int64) last_loc = torch.tensor( @@ -1475,8 +1287,7 @@ class TestPagedMultiEndedAllocator(unittest.TestCase): self.assertEqual( int(full_alloc.physical_to_virtual[p_page].item()), consumed_page ) - # Watermark must have advanced by 1 page. - # `allocated_count()` returns TOKENS (advance by PAGE_SIZE); + # `allocated_count()` is in TOKENS (advances by PAGE_SIZE); # `_allocated_pages()` is the page count. self.assertEqual( full_alloc.allocated_count(), @@ -1496,16 +1307,14 @@ class TestPagedMultiEndedAllocator(unittest.TestCase): int(free_before.numel()) - 1, ) - # 11. REGRESSION: alloc_decode with num_new_pages == 0 - # (the common case — the decode token reuses the prefix's tail page) - # must NOT advance the watermark and NOT touch v2p / p2v. + # 11. Regression: alloc_decode with num_new_pages == 0 must not advance + # the watermark or touch v2p / p2v. def test_paged_alloc_decode_no_op_when_no_new_page(self): from sglang.srt.mem_cache.allocator import unified_sub_pool as mea_mod _, full_alloc, _, _, _ = self._build() PS = self.PAGE_SIZE - # Pre-allocate 2 pages worth. We'll simulate a decode where seq_len - # advances WITHIN the existing tail page (no new page consumed). + # Alloc one page, then decode within it: no new page is consumed. v = full_alloc.alloc(PS) free_before = full_alloc.free_virtual_ids.clone() watermark_before = full_alloc.watermark_physical @@ -1537,7 +1346,7 @@ class TestPagedMultiEndedAllocator(unittest.TestCase): mea_mod.alloc_decode_kernel = original_kernel self.assertIsNotNone(out) - # Nothing should have moved — no new page consumed. + # Nothing should have moved -- no new page consumed. self.assertEqual(full_alloc.watermark_physical, watermark_before) self.assertEqual(full_alloc.allocated_count(), allocated_count_before) self.assertEqual( @@ -1545,11 +1354,8 @@ class TestPagedMultiEndedAllocator(unittest.TestCase): int(free_before.numel()), ) - # REGRESSION: `translate_kv_loc(virt, out=buf)` must work - # under page_size > 1 — the page-math branch writes via - # `index_select(out=out)` + in-place `mul_` / `add_` and must match the - # no-`out=` form byte-for-byte. Tests the actual page-math path of the - # base-class implementation. + # Regression: at page_size > 1 the `out=` page-math branch must match the + # no-`out=` form byte-for-byte. def test_paged_translate_kv_loc_with_out(self): _, full_alloc, _, _, _ = self._build() ps = self.PAGE_SIZE @@ -1573,12 +1379,8 @@ class TestPagedMultiEndedAllocator(unittest.TestCase): self.assertTrue(bool((buf == expected).all().item())) self.assertTrue(bool((with_out == no_out).all().item())) - # REGRESSION: the in-place aliasing form - # `translate_kv_loc(buf, out=buf)` must work at page_size > 1 too. The - # page-math branch computes `virt_pages` and `offsets` BEFORE writing - # into `out`, so those fresh tensors capture the pre-mutation values of - # virt_tokens. The final result must equal `phys_page * ps + offset` - # against the original input. + # Regression: at page_size > 1 the aliasing form `translate_kv_loc(buf, + # out=buf)` must derive `virt_pages` / `offsets` before writing into `out`. def test_paged_translate_kv_loc_with_out_aliasing_input(self): _, full_alloc, _, _, _ = self._build() ps = self.PAGE_SIZE @@ -1597,66 +1399,8 @@ class TestPagedMultiEndedAllocator(unittest.TestCase): "page>1 in-place result must equal no-out result", ) - # REGRESSION: the stale-tail scenario. - # - # A naive cuda-graph replay path in `triton_backend.py` would call - # `_translate_kv_loc(kv_indices, out=kv_indices)` on the WHOLE - # pre-allocated buffer (`self.cuda_graph_kv_indices`), even though only - # the `kv_indptr[-1]`-length prefix was freshly written by - # `create_flashinfer_kv_indices_triton`. Stale tail data, when fed - # through `v2p` repeatedly across replays, eventually produced negative - # values (via `v2p[unbound] = -1 → -1 * page_size + offset` ∈ [-ps,-1]), - # and the NEXT translation's `// page_size` produced `-1`, which CUDA's - # `index_select` rejects with a scatter-gather OOB device-side assert. - # - # The fix in `triton_backend.py` slices the translate to the valid - # prefix `kv_indices[:kv_indptr[-1]]`. This test confirms slicing is - # transparent to the translate — the in-place result on a contiguous - # slice of a larger buffer matches the standalone-tensor result. - def test_paged_translate_kv_loc_on_buffer_slice(self): - _, full_alloc, _, _, _ = self._build() - ps = self.PAGE_SIZE - v = full_alloc.alloc(2 * ps) - self.assertIsNotNone(v) - # Simulate the cuda-graph buffer pattern: a large pre-allocated - # buffer where only a prefix is freshly written. - N = v.numel() - big_buf = torch.zeros((N * 4,), dtype=torch.int64, device=_DEV) - big_buf[:N] = v - # Translate the valid prefix slice in-place (this is what the - # post-fix triton_backend call does). - slice_view = big_buf[:N] - ptr_before = slice_view.data_ptr() - ret = full_alloc.translate_kv_loc(slice_view, out=slice_view) - self.assertIs(ret, slice_view) - self.assertEqual( - slice_view.data_ptr(), - ptr_before, - "slice in-place write must preserve data_ptr", - ) - # The slice's translation must match a standalone translate of v. - expected = full_alloc.translate_kv_loc(v) - self.assertTrue( - bool((slice_view == expected).all().item()), - "slice in-place translate must equal standalone translate", - ) - # And the tail [N:] must remain UNTOUCHED — zeros, not corrupted - # by the translate. - tail = big_buf[N:] - self.assertTrue( - bool((tail == 0).all().item()), - "translating a slice must NOT touch the buffer tail; if this " - "test fails, the translate is reading/writing past the slice " - "bound — the same regression that caused a scatter-gather OOB " - "after several replays.", - ) - - # REGRESSION: tombstone-safety clamp at page_size > 1. - # - # For ps > 1, `v2p_page[vpage] == -1` produces `-1 * ps + offset` for - # output tokens — a negative value in `[-ps, -1]`. Without clamping, - # the captured `k_buffer[result[i]]` is an illegal access. The clamp - # must produce `>= 0` for every output token. + # At ps > 1 a tombstoned page yields `-1 * ps + offset`, i.e. a value in + # [-ps, -1]; the clamp must lift every output token to >= 0. def test_paged_translate_kv_loc_clamps_tombstoned_v2p(self): _, full_alloc, _, _, _ = self._build() ps = self.PAGE_SIZE @@ -1672,7 +1416,7 @@ class TestPagedMultiEndedAllocator(unittest.TestCase): bool((out >= 0).all().item()), f"paged translate_kv_loc must clamp tombstoned to >=0; got {out.tolist()}", ) - # The first `ps` tokens belong to the tombstoned page → all 0. + # The first `ps` tokens belong to the tombstoned page -> all 0. self.assertTrue( bool((out[:ps] == 0).all().item()), "all tokens in a tombstoned page must map to slot 0 (padding sink)", @@ -1682,13 +1426,7 @@ class TestPagedMultiEndedAllocator(unittest.TestCase): bool((out[ps:] > 0).all().item()), "non-tombstoned pages must still translate to live physical slots", ) - - def test_paged_translate_kv_loc_with_out_clamps_tombstoned_v2p(self): - _, full_alloc, _, _, _ = self._build() - ps = self.PAGE_SIZE - v = full_alloc.alloc(2 * ps).clone() - tomb_page = int((v[0] // ps).item()) - full_alloc.virtual_to_physical[tomb_page] = -1 + # out= form must also clamp. buf = torch.empty_like(v) ret = full_alloc.translate_kv_loc(v, out=buf) self.assertIs(ret, buf) @@ -1698,10 +1436,8 @@ class TestPagedMultiEndedAllocator(unittest.TestCase): ) self.assertTrue(bool((buf[:ps] == 0).all().item())) - # 14. REGRESSION: the leak-invariant terms used by the - # scheduler runtime checker must all be in TOKENS. Specifically - # `full_available_size() + allocated_tokens == static_cap` must hold for - # the SWA composite. + # 14. Regression: the leak-invariant terms the scheduler checks are all in + # TOKENS -- `full_available_size() + allocated_tokens == static_cap`. def test_paged_swa_full_available_size_in_tokens(self): from sglang.srt.mem_cache.allocator.unified_hybrid_swa import ( UnifiedSWATokenToKVPoolAllocator, @@ -1748,9 +1484,8 @@ class TestPagedMultiEndedAllocator(unittest.TestCase): need_sort=False, forward_stream=None, ) - # Idle: conserve view == cap (in tokens). (The leak invariant reads - # `_conserve_*`; the public `full/swa_available_size()` is now - # `min(conserve, schedulable)` and may be smaller — reserved sink page.) + # The leak invariant reads `_conserve_*`; public `*_available_size()` + # is `min(conserve, schedulable)` and may be smaller. self.assertEqual(allocator._conserve_full_available_size(), full_max) self.assertEqual(allocator._conserve_swa_available_size(), swa_max) @@ -1771,10 +1506,7 @@ class TestPagedMultiEndedAllocator(unittest.TestCase): swa_max - 2 * PS, ) - # First-principles leak invariant: at this point, allocated tokens - # are all "live" (no eviction yet). So: - # total = conserve + allocated_tokens - # where allocated_tokens = full_max - conserve. + # No eviction yet, so total == conserve + allocated_tokens. allocated_tokens = full_max - allocator._conserve_full_available_size() self.assertEqual(allocated_tokens, 2 * PS) self.assertEqual( @@ -1782,20 +1514,15 @@ class TestPagedMultiEndedAllocator(unittest.TestCase): full_max, ) - # 15. REGRESSION: UnifiedMambaTokenToKVPoolAllocator.size - # must be TOTAL TOKENS (available + allocated, both in tokens). At - # page_size > 1, the earlier `available + allocated_pages` formula gave - # `tokens + pages` which silently broke the chunk-cache Mamba log lines - # (`#full token`, `full token usage`) and would have crashed Mamba+radix - # if radix weren't auto-downgraded to page=1. + # 15. Regression: `UnifiedMambaTokenToKVPoolAllocator.size` must be TOTAL + # TOKENS -- available + allocated, both in tokens, never a page count. def test_paged_mamba_size_in_tokens(self): from sglang.srt.mem_cache.allocator.unified_mamba import ( UnifiedMambaTokenToKVPoolAllocator, ) - # Build a minimal Mamba composite: one MHA spec for full + one - # Mamba spec. The mamba sub-allocator always uses page_size=1, but - # the full sub-allocator uses self.PAGE_SIZE. + # The mamba sub-allocator always uses page_size=1; only the full + # sub-allocator takes self.PAGE_SIZE. PS = self.PAGE_SIZE full_spec = MHASubPoolSpec( name="full", @@ -1825,9 +1552,6 @@ class TestPagedMultiEndedAllocator(unittest.TestCase): device=_DEV, enable_memory_saver=False, ) - # Build a fake HybridLinearKVPool-like object with two sub-pool kv - # caches. We only need `.full_kv_pool` and `.mamba_pool` with - # `attach_allocator` / `move_kv_cache` stubs. full_kv = _FakeKVCache(pool.max_slots("full")) full_kv.attach_allocator = lambda allocator: None mamba_kv = _FakeKVCache(pool.max_slots("mamba")) @@ -1858,10 +1582,8 @@ class TestPagedMultiEndedAllocator(unittest.TestCase): v = allocator.alloc(2 * PS) self.assertIsNotNone(v) - # size should be CONSERVED in tokens: (available + allocated_tokens) - # stays at the initial total. (For the Mamba composite, `.size` is - # dynamic — it shrinks as the peer consumes bytes — but at this - # point the peer is idle so we should see `size == full_avail_before`.) + # `.size` is dynamic -- it shrinks as the peer consumes bytes -- but the + # peer is idle here, so the conserved sum still equals the initial total. self.assertEqual( allocator.full_attn_allocator.available_size() + allocator.full_attn_allocator.allocated_count(), @@ -1873,23 +1595,10 @@ class TestPagedMultiEndedAllocator(unittest.TestCase): # And .size matches this conserved sum. self.assertEqual(allocator.size, full_avail_before) - # 16. REGRESSION: the page-math helper used by - # `UnifiedSWAKVPool.translate_loc_from_full_to_swa`, - # `UnifiedSWAKVPool.get_cpu_copy`, and `load_cpu_copy` must do - # `virt_pages = loc // page_size; offsets = loc % page_size; - # phys_tokens = v2p_page[virt_pages] * page_size + offsets`. - # - # A naive implementation does `v2p[loc]` directly — - # indexing a page-granular table with token-granular ids, producing - # wrong physical token ids and (when used as Triton kernel inputs) - # OOB reads (the same bug class as the alloc_extend/alloc_decode - # binding regressions above). - # - # We can't easily construct a real UnifiedSWAKVPool in the CPU test shim - # (it inherits SWAKVPool which builds MHATokenToKVPool sub-pools), so - # we exercise the static helper `_virt_tokens_to_phys_tokens` directly. - # The instance methods in production wrap this helper, so the same - # math is covered. + # 16. Regression: `_virt_tokens_to_phys_tokens` must do the page math -- + # indexing the page-granular v2p with token ids gives OOB kernel inputs. + # A real `UnifiedSWAKVPool` cannot be built in this CPU shim, so the + # production instance methods are covered through the static helper. def test_paged_pool_translate_helper_returns_physical_tokens(self): from sglang.srt.mem_cache.allocator.unified_hybrid_swa import ( UnifiedSWATokenToKVPoolAllocator, @@ -1936,8 +1645,6 @@ class TestPagedMultiEndedAllocator(unittest.TestCase): forward_stream=None, ) - # Alloc 2 pages worth of tokens — the swa allocator's v2p_page table - # now has bindings for the consumed virtual pages. v_tokens = allocator.alloc(2 * PS) self.assertIsNotNone(v_tokens) @@ -1946,10 +1653,6 @@ class TestPagedMultiEndedAllocator(unittest.TestCase): v_tokens, allocator.swa_attn_allocator ) - # Output must: - # 1. Be non-negative for every input (none unbound at this point). - # 2. Be distinct (one-to-one mapping). - # 3. Match `swa_phys_page * page_size + offset` reconstructed directly. self.assertTrue( bool((swa_phys >= 0).all().item()), "REGRESSION: _virt_tokens_to_phys_tokens returned negative " @@ -1972,9 +1675,8 @@ class TestPagedMultiEndedAllocator(unittest.TestCase): "v2p_page[virt_pages] * page_size + offsets.", ) - # The composite emits KERNEL-FACING ids, not the physical token ids - # this helper returns; they coincide only at multiplier 1, which no - # sub-pool uses. + # The composite emits KERNEL-FACING ids, not the physical token ids this + # helper returns; they coincide only at multiplier 1, which nothing uses. swa_mult = allocator.swa_kernel_page_multiplier self.assertEqual(swa_mult, 2 * swa_spec.layer_num) composite_out = allocator.translate_loc_from_full_to_swa(v_tokens) @@ -1987,11 +1689,9 @@ class TestPagedMultiEndedAllocator(unittest.TestCase): class TestLazyCompaction(unittest.TestCase): - """Lazy compaction invariants and lazy-vs-eager - equivalence harness. CPU-only (no GPU events; the conservative Phase A - `_flush` uses `wait_stream(forward_stream)` only when `forward_stream is - not None`, so passing `forward_stream=None` keeps it a no-op). - """ + """Lazy compaction invariants and the lazy-vs-eager equivalence harness. + `_flush` only waits on `forward_stream` when one is passed, so these + allocators stay CPU-only.""" def _make_full(self, *, lazy: bool, n_full_slots=64, n_mamba_slots=16): full = _make_mha_spec("full", "up", layer_num=2) @@ -2033,52 +1733,20 @@ class TestLazyCompaction(unittest.TestCase): p = int(alloc.virtual_to_physical[v].item()) kv.buf[p] = int(v) - def test_lazy_free_boundary_shortcut(self): - """Boundary absorption is DEFERRED to `_flush` (the hot - path `_free_lazy` does only a `torch.cat`, no watermark mutation). - After `_flush`, the freed boundary page is absorbed into the - watermark. - """ - _pool, fa, _kv = self._make_full(lazy=True) - a = fa.alloc(3) # virtual tokens - before_wm = fa.watermark_physical - # Free the last-alloced virtual id (its physical IS the boundary). - last = a[-1:].clone() - fa.free(last) - # Watermark is NOT shrunk inline; freed page is in the - # free list. - self.assertEqual(fa.watermark_physical, before_wm) - self.assertEqual(len(fa._free_phys_pages), 1) - # `live_page_count` is decremented at free time (CPU-side metadata). - self.assertEqual(fa.live_page_count, 2) - # `_flush` runs the complete boundary absorb → watermark shrinks - # by 1, free list emptied. - fa._flush(urgent=True) - self.assertEqual(fa.watermark_physical, before_wm - 1) - self.assertEqual(len(fa._free_phys_pages), 0) - self.assertEqual(fa.live_page_count, 2) - def test_lazy_free_inward_walk(self): - """The inward walk (multiple contiguous holes absorbed - into the watermark in one pass) is DEFERRED to `_flush`. After - flush, the watermark shrinks past all contiguous-from-boundary - holes, regardless of the order they were freed. - """ + """The inward walk is deferred to `_flush`, which then absorbs every + contiguous-from-boundary hole in one pass, whatever the free order.""" _pool, fa, _kv = self._make_full(lazy=True) a = fa.alloc(5) wm_before = fa.watermark_physical - # Free a middle slot first → hole. + # Free a middle slot first -> hole. fa.free(a[2:3].clone()) - # Free the topmost ids — in eager mode these would absorb inline, + # Free the topmost ids -- in eager mode these would absorb inline, # but in lazy mode they cat onto the free list. fa.free(a[4:5].clone()) fa.free(a[3:4].clone()) - # 3 entries in the free list now; watermark still at the - # pre-free position. self.assertEqual(fa.watermark_physical, wm_before) self.assertEqual(len(fa._free_phys_pages), 3) - # `_flush` runs the complete CPU-side boundary absorb. All 3 - # contiguous holes near the boundary get absorbed in one pass. fa._flush(urgent=True) self.assertEqual(fa.watermark_physical, wm_before - 3) self.assertEqual(len(fa._free_phys_pages), 0) @@ -2094,11 +1762,11 @@ class TestLazyCompaction(unittest.TestCase): n_holes_before = len(fa._free_phys_pages) self.assertEqual(n_holes_before, 2) wm_before = fa.watermark_physical - # Allocate 2 more — both should come from holes, watermark unchanged. + # Allocate 2 more -- both should come from holes, watermark unchanged. a2 = fa.alloc(2) self.assertEqual(fa.watermark_physical, wm_before) self.assertEqual(len(fa._free_phys_pages), 0) - # Allocate 1 more — must extend (no holes left). + # Allocate 1 more -- must extend (no holes left). a3 = fa.alloc(1) self.assertEqual(fa.watermark_physical, wm_before + 1) # All three new alloc batches are non-None. @@ -2110,15 +1778,14 @@ class TestLazyCompaction(unittest.TestCase): _pool, fa, _kv = self._make_full(lazy=True) avail_initial = fa.available_size() a = fa.alloc(5) - # 3 non-boundary frees → 3 holes; watermark unchanged. + # 3 non-boundary frees -> 3 holes; watermark unchanged. fa.free(a[0:1].clone()) fa.free(a[1:2].clone()) fa.free(a[2:3].clone()) self.assertEqual(len(fa._free_phys_pages), 3) avail_after = fa.available_size() - # holes (3) + remaining extension capacity == original capacity - # adjusted by the 2 still-live tokens at the top. - # Concretely: avail_after = avail_initial - 2 (live). + # Holes stay available capacity, so only the 2 tokens that stayed live + # are subtracted from the initial availability. self.assertEqual(avail_after, avail_initial - 2) def test_lazy_flush_compacts_holes_into_gap(self): @@ -2134,7 +1801,7 @@ class TestLazyCompaction(unittest.TestCase): self.assertEqual(len(fa._free_phys_pages), 1) wm_before = fa.watermark_physical n_moves = fa._flush(urgent=True) - # At least one move should have happened (topmost survivor → hole). + # At least one move should have happened (topmost survivor -> hole). self.assertGreaterEqual(n_moves, 1) # Watermark shrunk; hole list is now empty. self.assertLess(fa.watermark_physical, wm_before) @@ -2142,22 +1809,6 @@ class TestLazyCompaction(unittest.TestCase): # live_page_count invariant under compaction. self.assertEqual(fa.live_page_count, 4) - def test_lazy_v2p_p2v_identity_after_flush(self): - """After a flush, v2p ∘ p2v == identity on the live set.""" - _pool, fa, _kv = self._make_full(lazy=True) - a = fa.alloc(8) - # Free a scattered set of virtuals (not at boundary). - fa.free(a[1:2].clone()) - fa.free(a[3:4].clone()) - fa.free(a[5:6].clone()) - fa._flush(urgent=True) - # For every still-live virtual token, v2p ∘ p2v == identity. - for v in a.tolist(): - p = int(fa.virtual_to_physical[v].item()) - if p == -1: - continue # freed - self.assertEqual(int(fa.physical_to_virtual[p].item()), v) - def test_lazy_flush_gathers_survivor_mappings_as_one_batch(self): """Compaction must not synchronize once per relocated survivor.""" _pool, fa, kv = self._make_full(lazy=True) @@ -2238,10 +1889,9 @@ class TestLazyCompaction(unittest.TestCase): self.assertEqual(lazy_data[v], lazy_stamps[v], f"lazy: KV[v={v}] != stamp") def test_lazy_non_urgent_stops_at_write_set_blocker(self): - """Write-race case: when the topmost survivor IS in an - in-flight batch's write-set, non-urgent `_flush` STOPS the - boundary walk (skipping past would shuffle holes without - shrinking the watermark — wasted work).""" + """When the topmost survivor is in an in-flight batch's write-set, a + non-urgent `_flush` stops the boundary walk: skipping past it would + shuffle holes without shrinking the watermark.""" _pool, fa, _kv = self._make_full(lazy=True) a = fa.alloc(5) @@ -2257,14 +1907,12 @@ class TestLazyCompaction(unittest.TestCase): # Free a non-boundary slot to create a compactable hole. fa.free(a[1:2].clone()) self.assertEqual(len(fa._free_phys_pages), 1) - # Register an in-flight forward whose write-set INCLUDES the - # topmost survivor. We pass the virtual `out_cache_loc` TENSOR - # (not a materialized physical set) — `_flush` translates it - # lazily on the scheduler thread when classifying survivors. + # `set_inflight_forward` takes the virtual `out_cache_loc` tensor, not a + # materialized physical set; `_flush` translates it when it classifies. topmost_phys = int(fa.virtual_to_physical[int(a[-1].item())].item()) oclv = a[-1:].clone() # virtual id that translates to topmost_phys fa.set_inflight_forward(ev, oclv) - # Non-urgent flush → case A blocker at the top → STOP. + # Non-urgent flush -> case A blocker at the top -> STOP. n_moves = fa._flush(urgent=False) self.assertEqual( n_moves, @@ -2275,21 +1923,17 @@ class TestLazyCompaction(unittest.TestCase): # State untouched: hole still present. self.assertEqual(len(fa._free_phys_pages), 1) self.assertEqual(len(fa._pending_reuse), 0) - # Fire the event (forward done) so the write-set entry prunes; a - # subsequent flush proceeds and releases src directly (event fired - # → no pending reuse entry needed). + # Firing the event prunes the write-set entry; the next flush then + # releases src directly, with no pending-reuse entry needed. ev.fired = True n_moves2 = fa._flush(urgent=False) self.assertGreaterEqual(n_moves2, 1) self.assertEqual(len(fa._pending_reuse), 0) def test_lazy_non_urgent_read_race_uses_pending_reuse(self): - """Read-race case (read race, no write race): when the - topmost survivor is NOT in any in-flight write-set, non-urgent - `_flush` compacts immediately (read+read on KV[src] is safe) and - pushes `(src, latest_event)` to `_pending_reuse` so a future - alloc can't write KV[src] while iter N+1's read is still pending. - """ + """With no in-flight write-set, a non-urgent `_flush` compacts at once + (read+read on KV[src] is safe) and parks `(src, latest_event)` in + `_pending_reuse`, so no later alloc writes KV[src] under a live read.""" _pool, fa, _kv = self._make_full(lazy=True) a = fa.alloc(5) @@ -2302,23 +1946,19 @@ class TestLazyCompaction(unittest.TestCase): ev = _FakeEvent() fa.set_latest_forward_done_event(ev) - # Free a non-boundary slot → 1 hole. + # Free a non-boundary slot -> 1 hole. fa.free(a[1:2].clone()) - # Empty in-flight write-set for the in-flight forward → topmost - # survivor is NOT in the write-set → compact proceeds, src goes - # to pending. Pass `None` (or an empty tensor) for - # `out_cache_loc_virtual` to signal "no write race on this pool" - # — this is the same path Mamba uses (forward writes mamba state - # via its own kernels, not via `out_cache_loc`). + # `out_cache_loc_virtual=None` means "no write race on this pool" -- the + # path Mamba uses, since its forward writes state through its own + # kernels rather than `out_cache_loc`. fa.set_inflight_forward(ev, None) n_moves = fa._flush(urgent=False) self.assertGreaterEqual(n_moves, 1) - # `_pending_reuse` has ONE entry per BATCH (keyed by - # event), not per src. So len(_pending_reuse) == 1 here. The - # total pages held is tracked in `_pending_reuse_pages_cpu`. + # `_pending_reuse` holds one entry per BATCH (keyed by event), not per + # src; the page count lives in `_pending_reuse_pages_cpu`. self.assertEqual(len(fa._pending_reuse), 1) self.assertEqual(len(fa._pending_reuse_pages_cpu), n_moves) - # Fire the event and drain — srcs return to availability. + # Fire the event and drain -- srcs return to availability. ev.fired = True fa._drain_pending_reuse(urgent=False) self.assertEqual(len(fa._pending_reuse), 0) @@ -2326,13 +1966,8 @@ class TestLazyCompaction(unittest.TestCase): class TestO3FusedAllocBind(unittest.TestCase): - """Fused take_physical_pages + bind_pages. - - GPU-only tests (Triton kernel requires CUDA). Exercise the helper - `_alloc_bind_fast_or_slow` directly: fast-path correctness, slow- - path fallback when holes exist (Invariant B), overflow handling, - eager vs lazy modes, grow-up vs grow-down, and page_size > 1. - """ + """Fused take_physical_pages + bind_pages via `_alloc_bind_fast_or_slow`. + GPU-only: the fused kernel is Triton.""" @classmethod def setUpClass(cls): @@ -2381,9 +2016,7 @@ class TestO3FusedAllocBind(unittest.TestCase): return pool, fa, full_kv def test_fast_path_when_no_holes(self): - """When `_free_phys_pages` is empty, the fast path fires. - Verifies: watermark advanced, v2p and p2v scattered correctly, - return tensor matches the kernel's arange.""" + """When `_free_phys_pages` is empty, the fused fast path fires.""" _pool, fa, _kv = self._make_full(lazy=True) # Sanity: empty holeset. self.assertEqual(len(fa._free_phys_pages), 0) @@ -2398,26 +2031,30 @@ class TestO3FusedAllocBind(unittest.TestCase): wm_before, wm_before + 4, dtype=torch.int64, device="cuda" ) self.assertTrue(torch.equal(phys, expected_phys)) - # v2p table: each virtual → its physical. + # v2p table: each virtual -> its physical. for v, p in zip(v_pages.tolist(), expected_phys.tolist()): self.assertEqual(int(fa.virtual_to_physical[v].item()), p) - # p2v table: each physical → its virtual. + # p2v table: each physical -> its virtual. for v, p in zip(v_pages.tolist(), expected_phys.tolist()): self.assertEqual(int(fa.physical_to_virtual[p].item()), v) # live_page_count updated. self.assertEqual(fa.live_page_count, 4) + # Another fast-path call accumulates. + v_pages2 = torch.tensor([24, 25], dtype=torch.int64, device="cuda") + fa._alloc_bind_fast_or_slow(v_pages2, 2) + self.assertEqual(fa.live_page_count, 6) def test_slow_path_when_holes_exist(self): - """Invariant B (greedy hole reuse): when a hole exists, alloc - drains it BEFORE extending the watermark. The fast path MUST - NOT fire.""" + """Greedy hole reuse: an existing hole must be drained before the + watermark extends, so the fast path must not fire.""" _pool, fa, _kv = self._make_full(lazy=True) # Build a hole by alloc-then-free-non-boundary. a = fa.alloc(3) - fa.free(a[0:1].clone()) # frees a non-boundary slot → enters holeset + self.assertEqual(fa.live_page_count, 3) + fa.free(a[0:1].clone()) # frees a non-boundary slot -> enters holeset + self.assertEqual(fa.live_page_count, 2) self.assertEqual(len(fa._free_phys_pages), 1) - # `_free_phys_pages` is a torch.Tensor; read the single - # hole position via `.tolist()` (`._alive` no longer exists). + # `_free_phys_pages` is a torch.Tensor; read the hole via `.tolist()`. hole_pos = int(fa._free_phys_pages.tolist()[0]) wm_before = fa.watermark_physical # Alloc 1 page via the helper. Slow path should drain the hole. @@ -2430,12 +2067,15 @@ class TestO3FusedAllocBind(unittest.TestCase): # v2p/p2v updated. self.assertEqual(int(fa.virtual_to_physical[42].item()), hole_pos) self.assertEqual(int(fa.physical_to_virtual[hole_pos].item()), 42) + # The slow path advances live_page_count via take_physical_pages. + self.assertEqual(fa.live_page_count, 3) def test_fast_path_in_eager_mode(self): - """Eager mode (no lazy compaction) ALWAYS uses the fast path — - no holes ever accumulate.""" + """Eager mode never accumulates holes, so it always takes the fast path; + `live_page_count` is not maintained there and stays 0.""" _pool, fa, _kv = self._make_full(lazy=False) self.assertFalse(fa.lazy_compaction) + self.assertEqual(fa.live_page_count, 0) wm_before = fa.watermark_physical v_pages = torch.tensor([30, 31, 32], dtype=torch.int64, device="cuda") phys = fa._alloc_bind_fast_or_slow(v_pages, 3) @@ -2444,6 +2084,8 @@ class TestO3FusedAllocBind(unittest.TestCase): wm_before, wm_before + 3, dtype=torch.int64, device="cuda" ) self.assertTrue(torch.equal(phys, expected_phys)) + # live_page_count UNCHANGED (eager mode invariant). + self.assertEqual(fa.live_page_count, 0) def test_index_space_overflow_returns_none(self): """When the requested allocation would overflow `num_pages`, @@ -2452,8 +2094,7 @@ class TestO3FusedAllocBind(unittest.TestCase): # Try to alloc more pages than exist. N = fa.num_pages + 100 wm_before = fa.watermark_physical - # Note: we need v_pages of size N; but only its NUMEL matters for - # the helper. We pass a dummy tensor of the right shape. + # Only v_pages' numel matters to the helper. v_pages = torch.zeros(N, dtype=torch.int64, device="cuda") phys = fa._alloc_bind_fast_or_slow(v_pages, N) self.assertIsNone(phys) @@ -2472,11 +2113,8 @@ class TestO3FusedAllocBind(unittest.TestCase): self.assertEqual(fa.watermark_physical, wm_before) def test_fast_path_equivalent_to_slow_path(self): - """For the same input on an empty-holeset allocator, the fast - path produces byte-identical v2p / p2v / return-tensor to the - slow path (which is the unfused take_physical_pages + bind - sequence). Verifies the kernel correctness against the - reference implementation.""" + """On an empty holeset, the fast path must produce identical v2p / p2v / + return tensors to the unfused take_physical_pages + bind sequence.""" # Two identical allocators; one takes fast path, one takes slow. _pool_a, fa_a, _kv_a = self._make_full(lazy=True) _pool_b, fa_b, _kv_b = self._make_full(lazy=True) @@ -2498,59 +2136,9 @@ class TestO3FusedAllocBind(unittest.TestCase): self.assertEqual(fa_a.watermark_physical, fa_b.watermark_physical) self.assertEqual(fa_a.live_page_count, fa_b.live_page_count) - def test_eager_mode_live_page_count_not_updated(self): - """Match `_take_physical_eager` semantics: in eager mode, - `live_page_count` is NOT maintained (the leak-checker uses - `allocated_count()` based on the watermark span). The helper's - fast path must respect this — updating it would break the - invariant that eager-mode `live_page_count == 0` always.""" - _pool, fa, _kv = self._make_full(lazy=False) - self.assertFalse(fa.lazy_compaction) - self.assertEqual(fa.live_page_count, 0) - v_pages = torch.tensor([10, 11, 12], dtype=torch.int64, device="cuda") - fa._alloc_bind_fast_or_slow(v_pages, 3) - # live_page_count UNCHANGED (eager mode invariant). - self.assertEqual(fa.live_page_count, 0) - - def test_lazy_mode_live_page_count_updated_on_fast_path(self): - """Lazy mode: the fast path advances `live_page_count` by N - (matches `take_physical`'s lazy-path bookkeeping).""" - _pool, fa, _kv = self._make_full(lazy=True) - self.assertEqual(fa.live_page_count, 0) - v_pages = torch.tensor([20, 21, 22], dtype=torch.int64, device="cuda") - fa._alloc_bind_fast_or_slow(v_pages, 3) - self.assertEqual(fa.live_page_count, 3) - # Another fast-path call accumulates. - v_pages2 = torch.tensor([23, 24], dtype=torch.int64, device="cuda") - fa._alloc_bind_fast_or_slow(v_pages2, 2) - self.assertEqual(fa.live_page_count, 5) - - def test_lazy_mode_live_page_count_updated_on_slow_path(self): - """Lazy mode + holes exist: the slow path advances - `live_page_count` via the existing `take_physical_pages` call - (which updates it internally). End state must match the fast - path's accumulation.""" - _pool, fa, _kv = self._make_full(lazy=True) - a = fa.alloc(3) - # alloc(3) used the fast path (no holes at the time). - self.assertEqual(fa.live_page_count, 3) - # Free one non-boundary → creates a hole; subsequent alloc takes - # the slow path. - fa.free(a[0:1].clone()) - self.assertEqual(fa.live_page_count, 2) - self.assertEqual(len(fa._free_phys_pages), 1) - # Now alloc(1) takes the slow path (holes exist). Should still - # update live_page_count back to 3. - b = fa.alloc(1) - self.assertIsNotNone(b) - self.assertEqual(fa.live_page_count, 3) - # Verify slow path actually fired: hole drained, watermark unchanged. - self.assertEqual(len(fa._free_phys_pages), 0) - def test_page_size_gt_1(self): - """Helper works at page_size > 1: virtual ids and table indices - are page-granular. Verifies kernel scatters one v2p entry per - PAGE (not per token).""" + """At page_size > 1 the kernel must scatter one v2p entry per PAGE, + not per token.""" _pool, fa, _kv = self._make_full( lazy=True, n_full_slots=64, n_mamba_slots=16, page_size=4 ) @@ -2610,9 +2198,8 @@ class TestO3FusedAllocBind(unittest.TestCase): wm_before = sa.watermark_physical v_pages = torch.tensor([5, 6, 7], dtype=torch.int64, device="cuda") phys = sa._alloc_bind_fast_or_slow(v_pages, 3) - # Grow-down: kernel emits ASCENDING (matches `_take_physical_eager`'s - # `torch.arange(wm-N+1, wm+1)` output). For wm=wm_before, N=3: - # range is [wm_before-2, wm_before-1, wm_before]. + # Grow-down: the kernel emits ASCENDING, matching + # `_take_physical_eager`'s `torch.arange(wm - N + 1, wm + 1)`. expected = torch.tensor( [wm_before - 2, wm_before - 1, wm_before], dtype=torch.int64, @@ -2621,24 +2208,15 @@ class TestO3FusedAllocBind(unittest.TestCase): self.assertTrue(torch.equal(phys, expected)) # Watermark decreased by N. self.assertEqual(sa.watermark_physical, wm_before - 3) - # v2p / p2v consistent with the ascending mapping: - # v_pages[0] → wm_before - 2 (lowest of the new range) - # v_pages[2] → wm_before (highest of the new range) for v, p in zip(v_pages.tolist(), expected.tolist()): self.assertEqual(int(sa.virtual_to_physical[v].item()), p) self.assertEqual(int(sa.physical_to_virtual[p].item()), v) class TestSWACompositeKernelIdSurface(unittest.TestCase): - """The SWA composite's kernel-facing id surface. - - Presence of `translate_kv_loc_for_kernel` / `full_v2p_page_table` is what flips - the attention backends' kernel-facing-first probes, and the `page_stride` scale in - `translate_loc_from_full_to_swa` is what carries the swa kernel-facing space. - Everything must collapse - byte-identically at multiplier 1 — the strided arm every existing SWA model - runs — and follow `kernel_id(t) = v2p[t//ps]*(ps*mult) + t%ps` otherwise. - """ + """The SWA composite's kernel-facing id surface. Attention backends probe for + `translate_kv_loc_for_kernel` / `full_v2p_page_table`, and every id must + follow `kernel_id(t) = v2p[t // ps] * (ps * mult) + t % ps`.""" PS = 4 FULL_L = 4 @@ -2682,36 +2260,26 @@ class TestSWACompositeKernelIdSurface(unittest.TestCase): forward_stream=None, ) - def test_multipliers_come_from_the_specs(self): - """Both sides scale by their OWN sub-pool's block count, and the - composite exposes the raw v2p tables unwrapped. Nothing injects the - scale: a spec whose views carry every layer cannot be paired with a - physical-id multiplier, which is the state that writes physical ids - into view rows.""" + def test_full_kernel_translate_matches_formula(self): + """Both sides scale by their OWN sub-pool's block count, and the full + kernel id follows v2p[t // ps] * (ps * mult) + t % ps.""" + mult = 2 * self.FULL_L a = self._build() self.assertEqual(a.kernel_page_multiplier, 2 * self.FULL_L) self.assertEqual(a.swa_kernel_page_multiplier, 2 * self.SWA_L) - self.assertIs(a.full_v2p_page_table, a.full_attn_allocator.virtual_to_physical) - self.assertIs(a.swa_v2p_page_table, a.swa_attn_allocator.virtual_to_physical) - - def test_full_kernel_translate_matches_formula(self): - mult = 2 * self.FULL_L - a = self._build() v = a.alloc(3 * self.PS) self.assertIsNotNone(v) v2p = a.full_attn_allocator.virtual_to_physical expected = v2p[v // self.PS] * (self.PS * mult) + v % self.PS self.assertTrue(torch.equal(a.translate_kv_loc_for_kernel(v), expected)) - # The PHYSICAL translate must stay unscaled — compaction and the byte + # The PHYSICAL translate must stay unscaled -- compaction and the byte # machinery depend on it staying in physical space. phys = v2p[v // self.PS] * self.PS + v % self.PS self.assertTrue(torch.equal(a.translate_kv_loc(v), phys)) def test_kernel_translate_accepts_an_int32_page_table(self): - """REGRESSION: fa3 translates its own page table, which is int32 and - 2-D. A gather that requires an int64 index (`torch.take`) crashes the - scheduler there while every int64 caller stays green. Both page sizes: - at ps == 1 the index IS the caller's tensor, at ps > 1 it is derived.""" + """Regression: fa3 passes its own page table, which is int32 and 2-D, so + the gather must not require an int64 index.""" for ps in (1, 4): with self.subTest(page_size=ps): self.PS = ps @@ -2756,12 +2324,9 @@ class TestSWACompositeKernelIdSurface(unittest.TestCase): class TestPs64MLACompositeFeasibility(unittest.TestCase): - """The Kimi/flashmla shape: MLA + mamba composite at page_size=64 (the - flashmla arg snap). Large pages stress every sizing derivation at once — - the 64-token sink-page floor, the ps*entry_bytes per-layer-view tail pad, and - the page-granular alloc — so this pins that the factory-shaped - construction stays FEASIBLE and the kernel-facing surface stays on-formula when - the page size jumps from the usual 1..4 to 64.""" + """MLA + mamba composite at page_size=64, the size flashmla snaps to. Large + pages stress the sink-page floor, the per-layer-view tail pad and the + page-granular alloc at once.""" PS = 64 LAYERS = 3 @@ -2830,10 +2395,8 @@ class TestPs64MLACompositeFeasibility(unittest.TestCase): class _ChainStub: - """Duck-typed chain member for frontier-walk tests: only the attributes the - walk itself touches. Lets the tests position an (opaque|transparent) middle - at exact byte coordinates without the full allocator machinery (the real - float allocator lands in a later commit).""" + """Duck-typed chain member carrying only what the frontier walk touches, so a + test can place an opaque or transparent middle at exact byte coordinates.""" def __init__(self, *, low_byte: int, high_byte: int, transparent: bool): self._low_byte = low_byte @@ -2858,14 +2421,8 @@ class _ChainStub: class TestChainFrontierWalk(unittest.TestCase): - """N-pool chain walk: 2-pool byte-identity to the old single-peer formulas, - transparent-middle skipping, and growth-side-neighbor credit routing. - - Guarded failure modes: (a) a rewrite of the walk silently changes the - 2-pool gap math (golden identity); (b) an empty/parked middle walls off - free space it does not occupy; (c) drainable-hole credit reads the wrong - chain member. - """ + """N-pool chain walk: 2-pool byte-identity with the single-peer formulas, + transparent-middle skipping, and growth-side-neighbor credit routing.""" def _build_pair(self): full = _make_mha_spec("full", "up", layer_num=2) @@ -2895,13 +2452,6 @@ class TestChainFrontierWalk(unittest.TestCase): ma.bind_peer(fa) return pool, fa, ma - def test_bind_peer_mirrors_growth_side(self): - _, fa, ma = self._build_pair() - self.assertIs(fa.high_peer, ma) # grow-up's neighbor sits above - self.assertIsNone(fa.low_peer) - self.assertIs(ma.low_peer, fa) # grow-down's neighbor sits below - self.assertIsNone(ma.high_peer) - def test_bind_peer_rejects_float_members(self): _, fa, ma = self._build_pair() stub = _ChainStub(low_byte=0, high_byte=0, transparent=True) @@ -2915,9 +2465,8 @@ class TestChainFrontierWalk(unittest.TestCase): fa.grow_direction = "up" def test_two_pool_gap_equals_old_single_peer_formula(self): - # Golden identity: with no middles, the chain walk must reproduce the - # pre-chain closed form gap_up = peer_low - my_high ; - # gap_down = my_low - peer_high at every allocation state. + # With no middles the walk must equal the closed form + # gap_up = peer_low - my_high; gap_down = my_low - peer_high. _, fa, ma = self._build_pair() for n_full, n_mamba in ((0, 0), (8, 0), (8, 4), (32, 16)): fa.clear() @@ -3000,16 +2549,8 @@ class TestChainFrontierWalk(unittest.TestCase): class TestFloatMultiEndedAllocator(unittest.TestCase): """Holes-first float middle: midpoint placement, in-place hole recycling, - larger-gap boundary extension, boundary absorption + park-on-empty - transparency, and the on-demand data movers (`make_room` boundary - relocation / leapfrog, `compact_holes` ordered pack) with data-move - verification through the fake KV marker. - - Guarded failure modes: a float that copies on steady-state churn, walls - off free space when empty, extends toward the tighter gap, moves more - than min(L_live, G) pages, loses data across relocation, or corrupts - v2p/p2v span/hole bookkeeping. - """ + larger-gap extension, boundary absorption with park-on-empty transparency, + and the on-demand movers `make_room` and `compact_holes`.""" def _build_tri(self, n_state=8, n_float=32, n_full=32): state = _make_mamba_spec("state", "up", layer_num=2) @@ -3113,11 +2654,9 @@ class TestFloatMultiEndedAllocator(unittest.TestCase): self._check_float_state(fla, kv) def test_boundary_free_absorbed_at_the_deferred_point(self): - """Boundary holes shrink the span ZERO-COPY -- but the shrink is - DEFERRED out of `free`, which must stay host-sync-free (deciding how - far to walk needs the hole set on the host). `free` records the holes; - `_absorb_span_boundary_holes` (per-step flush / shortfall ladder) reclaims - the span. Skipping it is only ever conservative.""" + """Boundary holes shrink the span zero-copy, but the shrink is deferred + out of `free`, which must stay host-sync-free; skipping the deferred + absorb is only ever conservative.""" _, _, fla, _, kv = self._build_tri() va = fla.alloc(2) vb = fla.alloc(2) # extends one side; frees at that edge absorb @@ -3137,9 +2676,8 @@ class TestFloatMultiEndedAllocator(unittest.TestCase): self._check_float_state(fla, kv) def test_free_is_host_sync_free(self): - """The per-decode-step property: no D2H anywhere in the float's free - (the base's lazy free earns this by deferring absorption to `_flush`; - the float now follows the same model).""" + """No D2H anywhere in the float's free: absorption is deferred to + `_flush`, so the per-decode-step path never syncs with the host.""" from unittest import mock _, _, fla, _, kv = self._build_tri() @@ -3159,9 +2697,7 @@ class TestFloatMultiEndedAllocator(unittest.TestCase): fla.free(v[:2], _pages=v[:2]) def test_deferred_absorption_reaches_the_same_state_as_eager(self): - """Derived property: deferring must not change WHERE the span lands, - only when. Two floats, identical ops; one absorbs after every free, - one only at the end.""" + """Deferring absorption must not change WHERE the span lands, only when.""" _, _, f1, _, kv1 = self._build_tri() _, _, f2, _, kv2 = self._build_tri() for f, kv, absorb_each in ((f1, kv1, True), (f2, kv2, False)): @@ -3303,11 +2839,6 @@ class TestFloatMultiEndedAllocator(unittest.TestCase): self.assertGreater(moved, 0) self._check_float_state(fla, kv) - def test_bind_peer_raises_on_float(self): - _, sa, fla, _, _ = self._build_tri() - with self.assertRaises(AssertionError): - fla.bind_peer(sa) - class TestDcpWidening(unittest.TestCase): """`dcp_size > 1`: the alloc surface speaks a widened virtual id space while @@ -3359,10 +2890,9 @@ class TestDcpWidening(unittest.TestCase): return self._build_pair(page_size=page_size, n_full_slots=n_full_slots)[0] def test_replicated_peer_stays_slot_granular_under_dcp(self): - """BUG REGRESSION. Widening every sub-allocator off the process DCP - width, rather than only the sharding one, makes the Mamba page size - dcp_size; state is allocated one slot per request, so the first request - fails the page-multiple check.""" + """Regression: only the sharding sub-allocator widens under DCP. A + widened Mamba page size fails the page-multiple check, since state is + allocated one slot per request.""" with self._dcp(4): _, mamba = self._build_pair(page_size=2) self.assertEqual(mamba.page_size, 1) diff --git a/test/registered/unit/mem_cache/test_swa_locked_full_recover_unified.py b/test/registered/unit/mem_cache/test_swa_locked_full_recover_unified.py index 2e4eaadb0..3ccf34783 100644 --- a/test/registered/unit/mem_cache/test_swa_locked_full_recover_unified.py +++ b/test/registered/unit/mem_cache/test_swa_locked_full_recover_unified.py @@ -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( diff --git a/test/registered/unit/mem_cache/test_swa_pool_v_head_dim.py b/test/registered/unit/mem_cache/test_swa_pool_v_head_dim.py index a7667fe21..e9dca5423 100644 --- a/test/registered/unit/mem_cache/test_swa_pool_v_head_dim.py +++ b/test/registered/unit/mem_cache/test_swa_pool_v_head_dim.py @@ -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() diff --git a/test/registered/unit/mem_cache/test_unified_byte_accounting.py b/test/registered/unit/mem_cache/test_unified_byte_accounting.py index b5e5b163e..b309e7d61 100644 --- a/test/registered/unit/mem_cache/test_unified_byte_accounting.py +++ b/test/registered/unit/mem_cache/test_unified_byte_accounting.py @@ -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() diff --git a/test/registered/unit/mem_cache/test_unified_byte_budget_sizing.py b/test/registered/unit/mem_cache/test_unified_byte_budget_sizing.py index d17c99280..b8d0060e7 100644 --- a/test/registered/unit/mem_cache/test_unified_byte_budget_sizing.py +++ b/test/registered/unit/mem_cache/test_unified_byte_budget_sizing.py @@ -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).""" diff --git a/test/registered/unit/mem_cache/test_unified_capacity_memo.py b/test/registered/unit/mem_cache/test_unified_capacity_memo.py index c3611f5c9..837f2c8fe 100644 --- a/test/registered/unit/mem_cache/test_unified_capacity_memo.py +++ b/test/registered/unit/mem_cache/test_unified_capacity_memo.py @@ -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() diff --git a/test/registered/unit/mem_cache/test_unified_free_no_host_sync.py b/test/registered/unit/mem_cache/test_unified_free_no_host_sync.py index cb0395085..f92def73d 100644 --- a/test/registered/unit/mem_cache/test_unified_free_no_host_sync.py +++ b/test/registered/unit/mem_cache/test_unified_free_no_host_sync.py @@ -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.[] = ` 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 diff --git a/test/registered/unit/mem_cache/test_unified_mamba_views.py b/test/registered/unit/mem_cache/test_unified_mamba_views.py index 2c7bdb49e..840c40d0b 100644 --- a/test/registered/unit/mem_cache/test_unified_mamba_views.py +++ b/test/registered/unit/mem_cache/test_unified_mamba_views.py @@ -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 diff --git a/test/registered/unit/mem_cache/test_unified_mha_views.py b/test/registered/unit/mem_cache/test_unified_mha_views.py index b0be816ce..e415c9ff9 100644 --- a/test/registered/unit/mem_cache/test_unified_mha_views.py +++ b/test/registered/unit/mem_cache/test_unified_mha_views.py @@ -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() diff --git a/test/registered/unit/mem_cache/test_unified_mla_gpu_parity.py b/test/registered/unit/mem_cache/test_unified_mla_gpu_parity.py index a6d3cd669..68b695ffd 100644 --- a/test/registered/unit/mem_cache/test_unified_mla_gpu_parity.py +++ b/test/registered/unit/mem_cache/test_unified_mla_gpu_parity.py @@ -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() diff --git a/test/registered/unit/mem_cache/test_unified_mla_views.py b/test/registered/unit/mem_cache/test_unified_mla_views.py index 5e7ebf508..ab5313509 100644 --- a/test/registered/unit/mem_cache/test_unified_mla_views.py +++ b/test/registered/unit/mem_cache/test_unified_mla_views.py @@ -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() diff --git a/test/registered/unit/mem_cache/test_unified_npool_sweep.py b/test/registered/unit/mem_cache/test_unified_npool_sweep.py index 4e81da222..fea776dc3 100644 --- a/test/registered/unit/mem_cache/test_unified_npool_sweep.py +++ b/test/registered/unit/mem_cache/test_unified_npool_sweep.py @@ -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() diff --git a/test/registered/unit/mem_cache/test_unified_tri_pool.py b/test/registered/unit/mem_cache/test_unified_tri_pool.py index c56950362..834e00722 100644 --- a/test/registered/unit/mem_cache/test_unified_tri_pool.py +++ b/test/registered/unit/mem_cache/test_unified_tri_pool.py @@ -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.""" diff --git a/test/registered/unit/server_args/test_page_major_backend_allowlist.py b/test/registered/unit/server_args/test_page_major_backend_allowlist.py index 86fc08fae..94e80282a 100644 --- a/test/registered/unit/server_args/test_page_major_backend_allowlist.py +++ b/test/registered/unit/server_args/test_page_major_backend_allowlist.py @@ -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) ):