From f61bb7b40a4e647a19f35ae5589e68ec11a4c52e Mon Sep 17 00:00:00 2001 From: Cheng Wan <54331508+ch-wan@users.noreply.github.com> Date: Mon, 31 Aug 2026 00:54:13 -0700 Subject: [PATCH] [unified-memory] Drop the vacated 'dense' qualifier and the restating comments (#37170) --- .../kernels/ops/kvcache/kv_read_table.py | 4 +- .../layers/attention/trtllm_mla_backend.py | 4 +- python/sglang/srt/mem_cache/memory_pool.py | 6 +-- .../srt/mem_cache/unified_memory_pool.py | 2 +- .../test_kimi_linear_unified_memory.py | 7 ++-- .../page_major/test_page_major_gpt_oss.py | 4 +- .../page_major/test_page_major_qwen_hybrid.py | 2 +- .../test_flashattention_graph_metadata.py | 1 - .../unit/mem_cache/test_full_loc_fast_path.py | 8 ++-- .../mem_cache/test_kv_index_translator.py | 10 ++--- .../unit/mem_cache/test_layout_compat.py | 2 +- .../mem_cache/test_multi_ended_allocator.py | 22 +++++----- .../test_pd_envelope_transfer_layout.py | 6 +-- .../unit/mem_cache/test_unified_mha_views.py | 18 ++++----- .../mem_cache/test_unified_mla_block_table.py | 40 +++++++++---------- .../unit/mem_cache/test_unified_mla_views.py | 28 ++++++------- .../test_unified_out_cache_loc_rebind.py | 8 ++-- .../test_page_major_backend_allowlist.py | 30 +++++++------- 18 files changed, 101 insertions(+), 101 deletions(-) diff --git a/python/sglang/kernels/ops/kvcache/kv_read_table.py b/python/sglang/kernels/ops/kvcache/kv_read_table.py index 6909c1287..04e4b90ae 100644 --- a/python/sglang/kernels/ops/kvcache/kv_read_table.py +++ b/python/sglang/kernels/ops/kvcache/kv_read_table.py @@ -21,8 +21,8 @@ the result into `out`: for c < ceil(seq_lens[b] / ps) -- the row's LIVE prefix `v2p` is the pool's virtual->physical page table and `multiplier` scales a -physical page into the id space the per-layer views use (1 when they are not -dense). Since only the page number is rewritten, a token-level consumer can +physical page into the id space the per-layer views use (1 when one page maps +to one row-block). Since only the page number is rewritten, a token-level consumer can rebuild flat ids as `entry * ps + offset`. PREFIX-ONLY per row: columns past the live prefix are never written, so a diff --git a/python/sglang/srt/layers/attention/trtllm_mla_backend.py b/python/sglang/srt/layers/attention/trtllm_mla_backend.py index 4cac93462..d3de3bb0d 100755 --- a/python/sglang/srt/layers/attention/trtllm_mla_backend.py +++ b/python/sglang/srt/layers/attention/trtllm_mla_backend.py @@ -396,7 +396,7 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend): self.decode_cuda_graph_kv_indices = torch.full( (max_bs, max_blocks_per_seq), -1, dtype=torch.int32, device=self.device ) - # Unified pool: capture-stable buffer for the DENSE KV write loc, filled + # Unified pool: capture-stable buffer for the kernel-facing KV write loc, filled # out-of-graph in init_forward_metadata_out_graph so the in-graph # set_mla_kv_buffer captures no translate. if self.kv_index_translator.is_translating: @@ -635,7 +635,7 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend): # Replay-prep receives the RAW (unpadded) out_cache_loc # (build_replay_fb_view), but the captured write kernel consumes the # full captured tier of this buffer. Zero the tail so pad rows write - # to the dense sink (row 0) instead of stale dense locs left by + # to the sink (row 0) instead of stale kernel-facing locs left by # earlier larger replays — a stale tail scatters pad-row garbage into # live KV pages. Mirrors the runner's PaddingPolicy.ZERO on its own # out_cache_loc slot. diff --git a/python/sglang/srt/mem_cache/memory_pool.py b/python/sglang/srt/mem_cache/memory_pool.py index 1cbce3c17..9e9ede5cc 100644 --- a/python/sglang/srt/mem_cache/memory_pool.py +++ b/python/sglang/srt/mem_cache/memory_pool.py @@ -1697,8 +1697,8 @@ class KVCache(abc.ABC): self.size = size self.page_size = page_size # Row-blocks one page holds in this pool's kernel-facing id space; >1 - # only where the per-layer views are dense (the unified pool), and then - # a write loc must have been translated into that space first. + # only for the unified pool's per-layer views, and then a write loc must + # have been translated into that space first. self.kernel_page_blocks = 1 self.dtype = dtype self.device = device @@ -2790,7 +2790,7 @@ class MHATokenToKVPool(KVCache): num_rows = int(loc_2d.numel()) if cache_k.shape[0] != num_rows or cache_v.shape[0] != num_rows: raise ValueError( - "dense KV rows must match loc_2d size: " + "KV rows must match loc_2d size: " f"{tuple(cache_k.shape)=} {tuple(cache_v.shape)=} {tuple(loc_2d.shape)=}." ) diff --git a/python/sglang/srt/mem_cache/unified_memory_pool.py b/python/sglang/srt/mem_cache/unified_memory_pool.py index fb84d9263..afd4d3c7d 100644 --- a/python/sglang/srt/mem_cache/unified_memory_pool.py +++ b/python/sglang/srt/mem_cache/unified_memory_pool.py @@ -1192,7 +1192,7 @@ def init_unified_mamba_pools( pre_alloc_size=decode_pre_alloc_size, ) if use_mla_backend: - # start_layer stays 0: HybridLinearKVPool patches layer ids to the dense + # start_layer stays 0: HybridLinearKVPool patches layer ids to the contiguous # 0..N-1 index via _transfer_id_context before every MLA pool call. unified_full_kv_pool = UnifiedMLATokenToKVPool( unified_buffer=shared_pool, 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 d77854d28..e253798ce 100644 --- a/test/registered/models_e2e/test_kimi_linear_unified_memory.py +++ b/test/registered/models_e2e/test_kimi_linear_unified_memory.py @@ -1,7 +1,7 @@ """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 DENSE per-layer +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 @@ -57,8 +57,9 @@ class TestKimiLinearUnifiedMemory( class TestKimiLinearUnifiedMemoryFlashMLA(TestKimiLinearUnifiedMemory): """flashmla at its ps=64 snap: the canonical block-table route - (KVIndexTranslator.build_into into flashmla's padded tables) plus the ps=64 - sub-pool sizing (64-token sink floor, dense-view tail pad) end to end. + (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 + [ 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 1c60f48f1..4bb46a411 100644 --- a/test/registered/page_major/test_page_major_gpt_oss.py +++ b/test/registered/page_major/test_page_major_gpt_oss.py @@ -28,8 +28,8 @@ _UNIFIED_COMMON_ARGS = [ class TestUnifiedGptOssTriton(DefaultServerBase): - """Unified pool on gpt-oss-20b (hybrid-SWA MoE), Triton pinned: dense - MHA/SWA views through the reference backend.""" + """Unified pool on gpt-oss-20b (hybrid-SWA MoE), Triton pinned: the MHA/SWA + per-layer views through the reference backend.""" model = DEFAULT_MODEL_NAME_FOR_TEST_MXFP4_WITH_MOE 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 d0096b0e4..a2e6462a7 100644 --- a/test/registered/page_major/test_page_major_qwen_hybrid.py +++ b/test/registered/page_major/test_page_major_qwen_hybrid.py @@ -34,7 +34,7 @@ _UNIFIED_COMMON_ARGS = [ class TestUnifiedQwenHybridTriton(DefaultServerBase): - """Unified pool on Qwen3.5-4B (GDN-hybrid), Triton pinned: dense + """Unified pool on Qwen3.5-4B (GDN-hybrid), Triton pinned: contiguous full-attention views + strided conv/SSM state through the reference backends.""" diff --git a/test/registered/unit/layers/attention/test_flashattention_graph_metadata.py b/test/registered/unit/layers/attention/test_flashattention_graph_metadata.py index 93e2a679d..4b3996333 100644 --- a/test/registered/unit/layers/attention/test_flashattention_graph_metadata.py +++ b/test/registered/unit/layers/attention/test_flashattention_graph_metadata.py @@ -59,7 +59,6 @@ class TestFlashAttentionGraphMetadata(CustomTestCase): backend.is_prefill_aware_swa = False backend.has_swa = False backend.use_sliding_window_kv_pool = False - backend._unified_dense = False backend.page_size = 1 backend._compute_scheduler_metadata = lambda *_: None backend._maybe_init_local_attn_metadata = lambda *_: None 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 d04ab4875..498eb6bc3 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 @@ -273,7 +273,7 @@ class TestHybridLinearMLARouting(unittest.TestCase): - `set_kv_buffer` (MLA branch) mirrors the MHA branch — write the pre-translated `KVWriteLoc.full_loc` when present (unified pool, where it - carries the DENSE loc), else the raw `loc` (static pool, already physical). + 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.""" @@ -290,19 +290,19 @@ class TestHybridLinearMLARouting(unittest.TestCase): 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) - dense_phys = torch.tensor([21, 24, 27], 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=dense_phys), + _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, dense_phys) + self.assertIs(forwarded, kernel_phys) self.assertIsNot(forwarded, virtual_loc) def test_mla_falls_back_to_loc_when_absent(self): 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 a69a217e6..c9a37054e 100644 --- a/test/registered/unit/mem_cache/test_kv_index_translator.py +++ b/test/registered/unit/mem_cache/test_kv_index_translator.py @@ -183,10 +183,10 @@ def _alloc_and_fill(allocator, ps, lens): class TestReadTableBuild(unittest.TestCase): - def test_read_table_matches_reference_dense_and_strided(self): + 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 (strided=1, dense=2L). The swa table agreeing with + 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.""" for ps in (1, 4): @@ -587,9 +587,9 @@ class TestWriteLoc(unittest.TestCase): self.assertTrue(torch.equal(fb.out_cache_loc, want_full)) self.assertTrue(torch.equal(virt, keep)) - def test_swa_write_loc_round_trips_from_dense(self): + 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 dense full-side values must equal the direct + 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.""" for ps in (1, 4, 64): @@ -600,7 +600,7 @@ class TestWriteLoc(unittest.TestCase): self.assertTrue(torch.equal(got, want_swa)) def test_pad_lanes_derive_to_sink(self): - """The DP pad appends zeros; dense 0 is the reserved padding slot in + """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.""" src, _, rows, seq_lens, _, want_full, want_swa = self._built(n=3) diff --git a/test/registered/unit/mem_cache/test_layout_compat.py b/test/registered/unit/mem_cache/test_layout_compat.py index 93b002397..3b5ca7e05 100644 --- a/test/registered/unit/mem_cache/test_layout_compat.py +++ b/test/registered/unit/mem_cache/test_layout_compat.py @@ -14,7 +14,7 @@ """Unit tests for the page-major envelope byte layout. The subject here is the ENVELOPE — the byte layout the unified pool stores its -KV in — pinned through ``MHASubPoolSpec``'s offset math. The dense 3-D views +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. 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 f927892b8..263222b24 100644 --- a/test/registered/unit/mem_cache/test_multi_ended_allocator.py +++ b/test/registered/unit/mem_cache/test_multi_ended_allocator.py @@ -1971,9 +1971,9 @@ class TestPagedMultiEndedAllocator(unittest.TestCase): 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) - expected_dense = swa_phys_pages_direct * (PS * swa_mult) + offsets_in + expected_kernel = swa_phys_pages_direct * (PS * swa_mult) + offsets_in self.assertTrue( - bool((composite_out.long() == expected_dense.long()).all().item()), + bool((composite_out.long() == expected_kernel.long()).all().item()), "REGRESSION: translate_loc_from_full_to_swa must emit the swa " "sub-pool's kernel-facing ids (phys_page * ps * blocks_per_page + offset).", ) @@ -2622,8 +2622,8 @@ class TestO3FusedAllocBind(unittest.TestCase): self.assertEqual(int(sa.physical_to_virtual[p].item()), v) -class TestSWACompositeDenseSurface(unittest.TestCase): - """The SWA composite's dense (kernel-facing) id surface. +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 @@ -2678,7 +2678,7 @@ class TestSWACompositeDenseSurface(unittest.TestCase): 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 are dense cannot be paired with a + 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.""" a = self._build() @@ -2687,7 +2687,7 @@ class TestSWACompositeDenseSurface(unittest.TestCase): 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_dense_translate_matches_formula(self): + def test_full_kernel_translate_matches_formula(self): mult = 2 * self.FULL_L a = self._build() v = a.alloc(3 * self.PS) @@ -2700,7 +2700,7 @@ class TestSWACompositeDenseSurface(unittest.TestCase): phys = v2p[v // self.PS] * self.PS + v % self.PS self.assertTrue(torch.equal(a.translate_kv_loc(v), phys)) - def test_dense_translate_accepts_an_int32_page_table(self): + 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: @@ -2732,7 +2732,7 @@ class TestSWACompositeDenseSurface(unittest.TestCase): expected = v2p_swa[v // self.PS] * (self.PS * mult) + v % self.PS self.assertTrue(torch.equal(a.translate_loc_from_full_to_swa(v), expected)) - def test_swa_dense_tombstone_still_lands_on_sink(self): + def test_swa_kernel_tombstone_still_lands_on_sink(self): """The scaled stride must not break the tombstone clamp: a tombstoned page's ids (v2p == -1 -> -stride + offset, negative for every in-page offset) still land on the sink, never negative.""" @@ -2753,7 +2753,7 @@ class TestPs64MLACompositeFeasibility(unittest.TestCase): 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 dense surface stays on-formula when + construction stays FEASIBLE and the kernel-facing surface stays on-formula when the page size jumps from the usual 1..4 to 64.""" PS = 64 @@ -2805,7 +2805,7 @@ class TestPs64MLACompositeFeasibility(unittest.TestCase): forward_stream=None, ) - def test_construction_alloc_and_dense_formula(self): + def test_construction_alloc_and_kernel_formula(self): a = self._build() # MLA: one latent row per layer, so the spec reports LAYERS blocks. self.assertEqual(a.kernel_page_multiplier, self.LAYERS) @@ -2813,7 +2813,7 @@ class TestPs64MLACompositeFeasibility(unittest.TestCase): self.assertIsNotNone(v, "2-page alloc infeasible at ps=64") # Page-aligned virtual run (page-granular allocator invariant). self.assertEqual(int(v[0].item()) % self.PS, 0) - # Dense translate follows the affine formula at ps=64, and every id + # The kernel translate follows the affine formula at ps=64, and every id # fits int32 (the canonical narrows on store). v2p = a.full_v2p_page_table want = v2p[v // self.PS] * (self.PS * self.LAYERS) + v % self.PS diff --git a/test/registered/unit/mem_cache/test_pd_envelope_transfer_layout.py b/test/registered/unit/mem_cache/test_pd_envelope_transfer_layout.py index 07147ccc5..eadb4b2d0 100644 --- a/test/registered/unit/mem_cache/test_pd_envelope_transfer_layout.py +++ b/test/registered/unit/mem_cache/test_pd_envelope_transfer_layout.py @@ -30,7 +30,7 @@ register_cpu_ci(est_time=60, suite="base-a-test-cpu") class TestMLAEnvelopeTransferAddressing(CustomTestCase): - def test_page_envelope_matches_dense_views(self): + def test_page_envelope_matches_per_layer_views(self): """Every (page, layer, slot) row written through the MLA views must land at raw_ptr + page * page_envelope_bytes + layer-block offset, i.e. inside the page's transfer envelope.""" @@ -62,9 +62,9 @@ class TestMLAEnvelopeTransferAddressing(CustomTestCase): for page in range(num_pages): for layer in range(layer_num): for off in range(page_size): - dense_id = page * layer_num * page_size + off + kernel_id = page * layer_num * page_size + off val = torch.randn(kv_dim, dtype=store_dtype) - views[layer][dense_id, 0] = val + views[layer][kernel_id, 0] = val start = ( page * page_bytes + layer * page_size * row_bytes 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 3d8a173ee..b0be816ce 100644 --- a/test/registered/unit/mem_cache/test_unified_mha_views.py +++ b/test/registered/unit/mem_cache/test_unified_mha_views.py @@ -11,7 +11,7 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -"""Dense 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). Covers, CPU-only (pure torch — no GPU / Triton kernels): - `build_mha_views` refuses an asymmetric-KV spec: its addressing @@ -145,7 +145,7 @@ def _reference_strided_views(raw, *, page_size, num_pages, anchor_bytes=0): return k_views, v_views -class TestMHADenseSpecSurface(unittest.TestCase): +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 @@ -189,7 +189,7 @@ class TestMHADenseSpecSurface(unittest.TestCase): ) -class TestDenseMHAViews(unittest.TestCase): +class TestMHAViews(unittest.TestCase): def test_view_shapes_are_stock_mha(self): ps, num_pages = 4, 6 k_views, v_views = _build_views(_make_raw(ps, num_pages), ps, num_pages) @@ -329,7 +329,7 @@ def _make_pool(ps=1, full_spec=None, device=_DEV): ) -class TestUnifiedKVPoolDenseViews(unittest.TestCase): +class TestUnifiedKVPoolViews(unittest.TestCase): def test_every_mha_sub_pool_is_per_layer_contiguous(self): """The unified pool has ONE MHA layout: both sub-pools come back as stock 3-D per-layer views, whatever their page size.""" @@ -482,8 +482,8 @@ class TestUnifiedMHATokenToKVPool(unittest.TestCase): self.assertEqual(pool.kv_cache_layout, "page_major") -class TestFactoryDenseViews(unittest.TestCase): - """The real SWA factory builds dense sub-pools and wires the matching +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.""" @@ -522,15 +522,15 @@ class TestFactoryDenseViews(unittest.TestCase): alloc = b.token_to_kv_pool_allocator self.assertEqual(alloc.kernel_page_multiplier, self.FULL_MULT) self.assertEqual(alloc.swa_kernel_page_multiplier, self.SWA_MULT) - # Sub-pools are the dense class exposing stock 3-D per-layer views. + # Sub-pools expose stock 3-D per-layer views. self.assertEqual(b.token_to_kv_pool.full_kv_pool.k_buffer[0].dim(), 3) self.assertEqual(b.token_to_kv_pool.swa_kv_pool.k_buffer[0].dim(), 3) self.assertGreater(pool.view_tail_pad_bytes, 0) - def test_rebind_emits_dense_full_and_build_derives_swa(self): + 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-DENSE write loc pointwise from those kernel-facing values + derives the SWA write loc pointwise from those kernel-facing values (phase 2) — both checked against the formulas over the VIRTUAL ids.""" from sglang.srt.mem_cache.kv_index_translator import KVIndexTranslator diff --git a/test/registered/unit/mem_cache/test_unified_mla_block_table.py b/test/registered/unit/mem_cache/test_unified_mla_block_table.py index 04a2a4bda..e4ad6f6dd 100644 --- a/test/registered/unit/mem_cache/test_unified_mla_block_table.py +++ b/test/registered/unit/mem_cache/test_unified_mla_block_table.py @@ -11,20 +11,20 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -"""DENSE block table / kv_indices for the paged MLA backends under the unified +"""Block table / kv_indices for the paged MLA backends under the unified memory pool (Kimi-Linear). -`req_to_token` holds VIRTUAL token ids, while the per-layer MLA views are dense -(`build_mla_views`). The paged MLA backends therefore need their page-level -block table filled with kernel-facing page ids: +`req_to_token` holds VIRTUAL token ids, while the per-layer MLA views are +contiguous (`build_mla_views`). The paged MLA backends therefore need their +page-level block table filled with kernel-facing page ids: - dense_page(virtual_page) = v2p[virtual_page] * layer_num + kernel_page(virtual_page) = v2p[virtual_page] * layer_num Since the read-path translator, ONE builder computes that formula for every -family — `build_kv_read_table` (the canonical) — and the backends only +family — `build_index_table` (the canonical) — and the backends only differ in how they consume it: - trtllm_mla / cutedsl_mla / tokenspeed_mla / flashmla: rows filled straight - into their padded block tables (`KVIndexTranslator.build_into`, prefix-only so + into their padded block tables (`KVIndexTranslator.fill_read_table`, prefix-only so the backends' own -1 / stale tail sentinels survive); - the flashinfer updaters: token ids reconstructed from the canonical by `create_flashinfer_kv_indices_triton[ENTRY_PAGE_SIZE=ps]`; @@ -34,7 +34,7 @@ differ in how they consume it: Covered here: - the static `create_flashmla_kv_indices_triton` (no id-space knowledge left) still matches the plain token//ps reference; - - the canonical route against the python dense reference, for several page + - the canonical route against the python reference, for several page sizes, ragged sequence lengths and a non-identity v2p permutation; - lanes past a row's live prefix keep the backend's -1 sentinel (prefix-only discipline — the trtllm/flashmla tail contract); @@ -108,7 +108,7 @@ def _fill_block_table_static(req_to_token, req_pool_indices, seq_lens, page_size def _reference(req_to_token, req_pool_indices, seq_lens, page_size, *, v2p, mult): - """Python reference: virtual token -> virtual page -> physical page -> dense.""" + """Python reference: virtual token -> virtual page -> physical page -> kernel id.""" bs = req_pool_indices.shape[0] max_blocks = (int(seq_lens.max().item()) + page_size - 1) // page_size ref = torch.full((bs, max_blocks), -1, dtype=torch.int64, device=_DEV) @@ -122,7 +122,7 @@ def _reference(req_to_token, req_pool_indices, seq_lens, page_size, *, v2p, mult @unittest.skipUnless(_HAS_CUDA, "requires CUDA") -class TestDenseBlockTable(unittest.TestCase): +class TestBlockTable(unittest.TestCase): def _make_batch(self, page_size, bs=5, max_ctx=2048, n_pages=512): """Ragged batch with a non-identity virtual->physical page permutation.""" g = torch.Generator(device="cpu").manual_seed(97 + page_size) @@ -159,7 +159,7 @@ class TestDenseBlockTable(unittest.TestCase): torch.equal(got.long(), want), f"page_size={page_size}: {got} != {want}" ) - def test_dense_block_table_matches_reference(self): + def test_block_table_matches_reference(self): for page_size in (1, 32, 64): rt, rpi, sl, v2p = self._make_batch(page_size) got = _fill_block_table(rt, rpi, sl, page_size, v2p=v2p, mult=_LAYERS) @@ -206,10 +206,10 @@ class TestDenseBlockTable(unittest.TestCase): f"row {r} padded lanes were written: {got[r]}", ) - def test_agrees_with_token_level_dense_translate(self): + def test_agrees_with_token_level_translate(self): """The flashinfer updaters translate TOKEN ids with `translate_kv_loc_for_kernel`; the trtllm path builds PAGE ids in-kernel. Both - must address the same dense page block.""" + must address the same kernel-facing page block.""" page_size = 64 rt, rpi, sl, v2p = self._make_batch(page_size) block_table = _fill_block_table( @@ -219,13 +219,13 @@ class TestDenseBlockTable(unittest.TestCase): n = int(sl[r].item()) virt_tokens = rt[r, :n].long() # translate_kv_loc_for_kernel's formula, applied to token ids. - dense_tokens = ( + kernel_tokens = ( v2p[virt_tokens // page_size] * (page_size * _LAYERS) + virt_tokens % page_size ) # The block-table entry scaled by page_size must be the kernel-facing id of # each page's first token. - first_of_page = dense_tokens[::page_size] + first_of_page = kernel_tokens[::page_size] n_pages = (n + page_size - 1) // page_size self.assertTrue( torch.equal(block_table[r, :n_pages] * page_size, first_of_page), @@ -234,7 +234,7 @@ class TestDenseBlockTable(unittest.TestCase): @unittest.skipUnless(_HAS_CUDA, "requires CUDA") -class TestFa3MetadataDenseBlockTable(unittest.TestCase): +class TestFa3MetadataBlockTable(unittest.TestCase): """fa3's captured-decode page table is written by `normal_decode_set_metadata` fed with the translator's read table kernel page table (src_is_read_table=True): the fused kernel copies the canonical @@ -251,7 +251,7 @@ class TestFa3MetadataDenseBlockTable(unittest.TestCase): build_kv_read_table, ) - maker = TestDenseBlockTable._make_batch + maker = TestBlockTable._make_batch rt, rpi, sl, v2p_full = maker(self, page_size, bs=bs, max_ctx=max_ctx) max_pages = (max_ctx + page_size - 1) // page_size @@ -325,11 +325,11 @@ class TestFa3MetadataDenseBlockTable(unittest.TestCase): got, want, sl = self._run(page_size, v2p=False, mult=1) self._assert_live_prefix(got, want, sl, page_size) - def test_dense_mapping_ps1_fast_path(self): + def test_translated_mapping_ps1_fast_path(self): got, want, sl = self._run(1, v2p=True, mult=_LAYERS) self._assert_live_prefix(got, want, sl, 1) - def test_dense_mapping_general_path(self): + def test_translated_mapping_general_path(self): got, want, sl = self._run(64, v2p=True, mult=_LAYERS) self._assert_live_prefix(got, want, sl, 64) @@ -339,7 +339,7 @@ class TestFa3MetadataDenseBlockTable(unittest.TestCase): got, want, sl = self._run(page_size, v2p=True, mult=1) self._assert_live_prefix(got, want, sl, page_size) virtual = _reference( - *TestDenseBlockTable._make_batch(self, page_size)[:3], + *TestBlockTable._make_batch(self, page_size)[:3], page_size, v2p=None, mult=1, 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 32b632ba9..72c3b71ad 100644 --- a/test/registered/unit/mem_cache/test_unified_mla_views.py +++ b/test/registered/unit/mem_cache/test_unified_mla_views.py @@ -11,7 +11,7 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -"""Dense MLA views for the unified memory pool (MLA-hybrid-Mamba, Kimi K3). +"""MLA views for the unified memory pool (MLA-hybrid-Mamba, Kimi K3). Covers, CPU-only (pure torch — no GPU / Triton kernels): - `MLASubPoolSpec` byte math; @@ -23,7 +23,7 @@ Covers, CPU-only (pure torch — no GPU / Triton kernels): 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`: dense = v2p-page * (ps*L) + + - `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. @@ -129,7 +129,7 @@ class TestMLASubPoolSpec(unittest.TestCase): ) -class TestDenseMLAViews(unittest.TestCase): +class TestMLAViews(unittest.TestCase): def _make_raw(self, ps, num_pages, pad_pages=1): page_bytes = ps * _L * _D * _ITEM raw = torch.zeros( @@ -270,7 +270,7 @@ class TestUnifiedMLATokenToKVPool(unittest.TestCase): 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_dense_readback(self): + def test_move_then_readback(self): ps = 4 pool, kv_pool = self._make(ps=ps) num_pages = pool.max_slots("full") // ps @@ -302,7 +302,7 @@ class _FakeKVCache: self.buf[dst_loc] = self.buf[src_loc].clone() -class TestTranslateKvLocDense(unittest.TestCase): +class TestTranslateKvLocForKernel(unittest.TestCase): def _build(self, ps=1, n_full_tokens=64, multiplier=_L): pool, full, mamba = _make_unified(page_size=ps, n_full_tokens=n_full_tokens) full_alloc = MultiEndedAllocator( @@ -325,30 +325,30 @@ class TestTranslateKvLocDense(unittest.TestCase): mamba_alloc.bind_peer(full_alloc) return full_alloc - def test_dense_matches_formula_ps1(self): + 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) - dense = alloc.translate_kv_loc_for_kernel(v) - self.assertTrue(torch.all(dense == phys * _L)) + kernel = alloc.translate_kv_loc_for_kernel(v) + self.assertTrue(torch.all(kernel == phys * _L)) - def test_dense_matches_formula_paged(self): + 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) - dense = alloc.translate_kv_loc_for_kernel(v) + kernel = alloc.translate_kv_loc_for_kernel(v) expected = (phys // ps) * (ps * _L) + phys % ps - self.assertTrue(torch.all(dense == expected)) + self.assertTrue(torch.all(kernel == expected)) def test_tombstone_clamps_to_sink(self): alloc = self._build(ps=1) # never-allocated virtual ids -> v2p == -1 -> kernel-facing id 0 virt = torch.tensor([alloc.min_slot_index + 1], dtype=torch.int64) - dense = alloc.translate_kv_loc_for_kernel(virt) - self.assertTrue(torch.all(dense == 0)) + kernel = alloc.translate_kv_loc_for_kernel(virt) + self.assertTrue(torch.all(kernel == 0)) def test_out_matches_and_aliases(self): for ps in (1, 4): @@ -373,7 +373,7 @@ class TestTranslateKvLocDense(unittest.TestCase): torch.all(alloc.translate_kv_loc_for_kernel(v) == alloc.translate_kv_loc(v)) ) - def test_dense_follows_compaction(self): + def test_kernel_id_follows_compaction(self): alloc = self._build(ps=1) a = alloc.alloc(4) b = alloc.alloc(4) diff --git a/test/registered/unit/model_executor/test_unified_out_cache_loc_rebind.py b/test/registered/unit/model_executor/test_unified_out_cache_loc_rebind.py index 3d0c19f37..31c657ae2 100644 --- a/test/registered/unit/model_executor/test_unified_out_cache_loc_rebind.py +++ b/test/registered/unit/model_executor/test_unified_out_cache_loc_rebind.py @@ -70,8 +70,8 @@ def _armed_source(v2p, swa_map): ) src.is_translating = True src._translate_full = lambda t, out=None: v2p[t.to(torch.int64)] - # Phase 2 derives from DENSE values through p2v + the swa v2p; arm the - # inverse of the fake v2p (ps=1, both multipliers 1: dense == physical, + # Phase 2 derives from kernel-facing values through p2v + the swa v2p; arm + # the inverse of the fake v2p (ps=1, both multipliers 1: kernel == physical, # and the expected swa loc for virtual t is swa_map[t]). p2v = torch.zeros(int(v2p.max()) + 1, dtype=torch.int64) p2v[v2p] = torch.arange(v2p.numel(), dtype=torch.int64) @@ -160,8 +160,8 @@ class TestPadComposesWithDerivation(CustomTestCase): with self.subTest(page_size=page_size, blocks=blocks): stride = page_size * blocks virt = torch.arange(1, 2 * stride, dtype=torch.int64) - dense = (virt // page_size) * stride + virt % page_size - in_space = dense % stride < page_size + kernel = (virt // page_size) * stride + virt % page_size + in_space = kernel % stride < page_size self.assertTrue(bool(in_space.all()), "kernel-facing ids must pass") # Virtual ids pass only in the first block; that is why the # probe needs a batch, not one id, to be conclusive. 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 d1f3e291a..86fc08fae 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 @@ -24,7 +24,7 @@ exposes per-layer views and nothing else: * plain `--enable-page-major-kv-layout` without the unified pool keeps the envelope-strided 4-D views only the stride-aware Triton kernels read. -The same handler also screens the pool itself: the dense MHA/SWA views need +The same handler also 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 @@ -99,8 +99,8 @@ def _accepts( class TestPageMajorBackendAllowlist(unittest.TestCase): - # Wired for the dense per-layer MLA views (see the module docstring). - DENSE_MLA_BACKENDS = ( + # Wired for the per-layer MLA views (see the module docstring). + PER_LAYER_VIEW_MLA_BACKENDS = ( "fa3", "trtllm_mla", "flashinfer", @@ -108,11 +108,11 @@ class TestPageMajorBackendAllowlist(unittest.TestCase): "tokenspeed_mla", "flashmla", ) - # Wired for the dense per-layer MHA/SWA views (uniform-row models). - DENSE_MHA_BACKENDS = ("fa3", "fa4", "flashinfer", "trtllm_mha") + # Wired for the per-layer MHA/SWA views (uniform-row models). + PER_LAYER_VIEW_MHA_BACKENDS = ("fa3", "fa4", "flashinfer", "trtllm_mha") # MLA-family kernels that must never leak into the MHA arm. MLA_ONLY_BACKENDS = ("trtllm_mla", "cutedsl_mla", "tokenspeed_mla", "flashmla") - # No dense-id wiring anywhere: must stay rejected until they get one. + # No kernel-facing-id wiring anywhere: must stay rejected until they get one. UNWIRED_BACKENDS = ("cutlass_mla", "aiter") def test_triton_allowed_on_every_arm(self): @@ -124,15 +124,15 @@ class TestPageMajorBackendAllowlist(unittest.TestCase): # backend, so it rejects even Triton. self.assertFalse(_accepts("triton", use_mla=False, has_asymmetric_kv=True)) - def test_dense_mla_backends_allowed_under_unified_mla(self): - for backend in self.DENSE_MLA_BACKENDS: + def test_per_layer_view_mla_backends_allowed_under_unified_mla(self): + for backend in self.PER_LAYER_VIEW_MLA_BACKENDS: self.assertTrue( _accepts(backend, use_mla=True), f"{backend} should be allowed with the unified-memory MLA pool", ) - def test_dense_mha_backends_allowed_for_uniform_row_models(self): - for backend in self.DENSE_MHA_BACKENDS: + def test_per_layer_view_mha_backends_allowed_for_uniform_row_models(self): + for backend in self.PER_LAYER_VIEW_MHA_BACKENDS: self.assertTrue( _accepts(backend, use_mla=False), f"{backend} should be allowed for a uniform-row MHA model", @@ -149,7 +149,7 @@ class TestPageMajorBackendAllowlist(unittest.TestCase): """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.DENSE_MLA_BACKENDS: + 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), @@ -160,7 +160,7 @@ class TestPageMajorBackendAllowlist(unittest.TestCase): """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.""" - for backend in ("triton",) + self.DENSE_MHA_BACKENDS: + for backend in ("triton",) + self.PER_LAYER_VIEW_MHA_BACKENDS: self.assertFalse( _accepts(backend, use_mla=False, has_asymmetric_kv=True), f"--enable-unified-memory + {backend} must be rejected for an " @@ -172,7 +172,7 @@ class TestPageMajorBackendAllowlist(unittest.TestCase): 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.""" - for backend in ("triton",) + self.DENSE_MLA_BACKENDS: + for backend in ("triton",) + self.PER_LAYER_VIEW_MLA_BACKENDS: self.assertTrue( _accepts(backend, use_mla=True, has_asymmetric_kv=True), f"{backend} must stay allowed for an MLA model with asymmetric " @@ -184,7 +184,7 @@ class TestPageMajorBackendAllowlist(unittest.TestCase): outright, Triton included: the static page-major arm went away with the strided views and awaits its per-layer-view reimplementation.""" for backend in ("triton",) + tuple( - set(self.DENSE_MLA_BACKENDS + self.DENSE_MHA_BACKENDS) + set(self.PER_LAYER_VIEW_MLA_BACKENDS + self.PER_LAYER_VIEW_MHA_BACKENDS) ): for use_mla in (True, False): self.assertFalse( @@ -197,7 +197,7 @@ class TestPageMajorBackendAllowlist(unittest.TestCase): for use_mla in (True, False): self.assertFalse( _accepts(backend, use_mla=use_mla), - f"{backend} has no dense-id wiring and must be rejected", + f"{backend} has no kernel-facing-id wiring and must be rejected", ) def test_helion_linear_attention_is_kda_only(self):