[unified-memory] Drop the vacated 'dense' qualifier and the restating comments (#37170)
This commit is contained in:
Cheng Wan
@@ -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
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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)=}."
|
||||
)
|
||||
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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 + [
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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."""
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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):
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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):
|
||||
|
||||
Reference in New Issue
Block a user