Repository navigation
[unified-memory] Token-major dense views for the unified memory pool (3/7) - #38592
Conversation
caihuali95
changed the title
caihuali95
changed the title
caihuali95
changed the title
| out_cache_loc_dsv4: Optional[DSV4OutCacheLoc] = None | ||
| # Id space of `out_cache_loc`: "kernel" once a pool's write door may take it, | ||
| # set by KVIndexTranslator.rebind_write_loc; capture-time batches declare it. | ||
| out_cache_loc_id_space: str = "virtual" |
There was a problem hiding this comment.
[bug] out_cache_loc_id_space is a non-optional str (default "virtual", set to "kernel" by rebind_write_loc on every real batch). TBO's child builder copies a whitelist of fields and then fails if any parent field is non-None and missing from output_dict (two_batch_overlap.py around the dataclasses.fields(ForwardBatch) loop). This field is not on that whitelist, so --enable-two-batch-overlap will raise Field out_cache_loc_id_space has value, but is not yet supported on the first split, unified memory or not. Even if the completeness check were skipped, the child would default back to "virtual" and the unified write door would treat an already-rebound loc as unmarked.
Suggestion: Copy out_cache_loc_id_space (the parent's value, not the default) into the TBO child dict next to out_cache_loc / out_cache_loc_virtual. A unit case that splits a rebound parent and asserts the child keeps "kernel" would have caught this; TBO's field-completeness loop is the ratchet.
There was a problem hiding this comment.
Confirmed and fixed. `filter_batch` now copies the parent's value:
```python
output_dict["out_cache_loc_id_space"] = batch.out_cache_loc_id_space
```
You were right that this breaks TBO regardless of the pool: the field is a non-Optional `str`, so the completeness guard fires on every split. It was also worse than latent — five of the existing `batch_overlap` tests were already red on the branch, and our GPU lane could not see it because that directory was not in the lane's test list. Both are now fixed: the lane covers `test/registered/unit/batch_overlap/`, and two cases pin the behaviour (a rebound parent hands `"kernel"` to the child; an untranslated one stays `"virtual"`).
One note for navigation: the field is introduced by #40328 (4/7), so the fix landed there rather than here. Because a cross-fork PR renders the whole stack's union, several of these threads point at code owned by a different PR in the series.
| v_cache = v_cache_raw.view( | ||
| -1, self.page_size, layer.tp_v_head_num, layer.head_dim | ||
| ) | ||
| k_cache = paged_view(k_cache_raw, self.page_size) |
There was a problem hiding this comment.
[bug] — revised; supersedes my original note on this thread. The .view() here does not raise, and the real failure mode is worse than the one I described.
Checked against a unified-shaped per-layer view — (1024, 4, 128) stride (32768, 128, 1) (32 layers, 4 KV heads, hd 128, so entry_elems = 32*2*4*128):
ps=256 .view(-1,256,4,128) -> (4,256,4,128) (8388608, 32768, 128, 1)
paged_view(...) -> (4,256,4,128) (8388608, 32768, 128, 1)
identical strides; addressing matches the physical token id
Splitting dim 0 into (pages, ps) never crosses the entry stride, so _reshape_paged_kv_cache's .view() succeeds on a strided unified view and, for page_size > 1, returns exactly what paged_view returns. It is not a crash — and the paged_view swap in the fmha_v2 else-branch above is therefore a no-op at page_size > 1.
The problem is that both spellings are stride-silent. Same shape, different strides:
.view(-1,256,4,128) |
|
|---|---|
| main (contiguous per-layer view) | (131072, 512, 128, 1) |
| this PR (token-major) | (8388608, 32768, 128, 1) |
stride(1) — the in-page token stride — goes from H*D to entry_elems, a factor of 2*layer_num (64x here). On main, a kernel that derives that stride from the shape is correct by construction; under this PR it reads bytes 64x too close together, with no exception and no shape mismatch to catch it. A raise would have been the safe outcome.
So the open question is not the spelling, it is whether the trtllm-gen kernel consumes stride(0) / stride(1) at all, or computes the page address from tokens_per_page * num_heads * head_dim. canonicalize_stride sitting on this path is weak evidence that the wrapper reads strides; it says nothing about the kernel. trtllm_mha and trtllm_mla are both on the unified allowlist and neither appears in the reported accuracy or speed grid.
Suggestion: add a stride-sensitivity parity test for trtllm_mha / trtllm_mla — push identical page contents through a contiguous cache and through an over-strided view of the same data, assert equal output — and gate the allowlist on it. Rewriting .view to paged_view does not address this.
Two narrowings, for the record:
page_size == 1is the only place the two spellings differ:.viewgives the size-1 page dim strideH*Dwherepaged_viewgivesentry_elems. Addressing is identical, but after.permute(0, 2, 1, 3)it lands atstride(2), andcanonicalize_stridewill not repair it — its condition issizes[i] == 1 and strides[i] == strides[i+1], which is512 != 1here, and it only runs whentp_k_head_num == 1.- The NVFP4 arm is unreachable under the unified pool:
MHASubPoolSpec.layout()declares onlykandvparts and no scale parts, so an NVFP4 KV cache is not constructible there. Please disregard my earlier mention of it.
There was a problem hiding this comment.
Thank you for re-checking and correcting the first read — the revised version is the one that matters, and we agree with it. The `paged_view` swap is indeed a no-op at `page_size > 1`; what it buys is only the `page_size == 1` case, where `.view` gives the size-1 page dim the row stride instead of the slot stride.
Your actual point stands: nothing here proves the trtllm-gen kernel reads `stride()` rather than deriving the page address from `page_size * heads * head_dim`. We have added the parity test you describe — `test/registered/attention/test_trtllm_mha_stride_parity.py` (2/7) — which pushes identical page contents through a contiguous cache and through an `as_strided` view with the slot stride widened to a whole entry, runs both through `_reshape_paged_kv_cache` and `trtllm_batch_decode_with_kv_cache`, and asserts `torch.equal` (no tolerance, since a tolerance would hide exactly this).
To be straight about its current value: it is `skipUnless(is_sm100_supported())`, and our validation pool is SM90, so it reports a named skip rather than a verdict there. It will assert in CI on SM100. We have not gated the allowlist on it; if you would prefer `trtllm_mha`/`trtllm_mla` dropped from the unified allowlist until that test has run green somewhere, we are happy to do that.
| ) | ||
| if head_group_num != 1: | ||
| # Reinterpret each page's heads as `head_group_num` pseudo-pages. | ||
| key_cache = key_cache.view( |
There was a problem hiding this comment.
[suggestion] After get_paged_kv_buffer, the head_group_num != 1 arm still .view()s the 4-D result into extra pseudo-pages. Unlike the trtllm_mha path (see my revised note there — that .view succeeds), this one genuinely throws on a strided unified view, because it splits H and merges the factor into dim 0, which does have to cross the entry stride:
paged (4,256,8,128) stride (16777216, 65536, 128, 1)
head_group_num=1: OK
head_group_num=2: RuntimeError: view size is not compatible with input tensor's
size and stride (at least one dimension spans across two
contiguous subspaces)
head_group_num=4: RuntimeError (same)
contiguous, hg=2: OK -> (8,256,4,128) (131072, 512, 128, 1)
MiniCPM is not on the unified allowlist, so this is latent rather than a live boot path — but the helper is now the FA paged-cache door, and the head_group_num == 1 change is exactly what makes the != 1 arm layout-sensitive. Worth noting it is the better-behaved of the two: this arm is unreachable and fails loudly, while the shipping trtllm_mha path fails silently.
Suggestion: rebuild the grouped pages with as_strided from the slot stride, or make the exclusion explicit with assert head_group_num == 1 or key_cache.is_contiguous(). Do not .reshape() — that would pack a copy and drop the entry stride.
There was a problem hiding this comment.
Reproduced exactly, including your stride figures — `head_group_num` 2 and 4 both raise on the strided view, and the contiguous case gives `(8,256,4,128) (131072, 512, 128, 1)`.
On the suggested fix: we looked at rebuilding with `as_strided` and it cannot express this mapping. The `.view` here is not a head split — it is a flat re-chunk of the (token, head) plane, so walking the new `page_size` axis crosses a slot boundary every `head_group_num` steps. On a contiguous buffer both steps cost the same (`slot_stride == H*D`), which is why one constant stride works; under token-major they differ, and `as_strided` needs one stride per dimension. A head split would be expressible, but only as a 5-D tensor, and flattening it back to 4-D needs `page_size == 1`.
So we took your second option and made the exclusion explicit, with the strides in the message:
```python
assert key_cache.is_contiguous() and value_cache.is_contiguous(), (
f"head_group_num={head_group_num} needs a contiguous paged KV view; ..."
)
```
`is_contiguous()` is exact rather than approximate here: for the paged view it holds precisely when `stride(1) == head_num * head_dim`, which is the condition that makes the re-chunk expressible. Covered by a case that asserts the raise on a strided pool and that `head_group_num == 1` still passes. We did not use `.reshape()`, per your warning. Landed in #40327 (2/7).
| id_space: str = "virtual" | ||
|
|
||
| @classmethod | ||
| def for_batch( |
There was a problem hiding this comment.
[suggestion] KVWriteLoc.for_batch is the only way a unified pool accepts a write loc under SGLANG_ENABLE_ASYNC_ASSERT. Allowlisted attention backends were migrated, but model-layer writers still pass a bare forward_batch.out_cache_loc (or KVWriteLoc(...) without id_space): deepseek_common/.../forward_mha.py _set_mla_kv_buffer, layers/cp/zigzag.py materialize_full_*, models/dots3_common/modeling.py. After rebind the values are already kernel-facing, so production writes (assert off) are correct; with the assert on, DeepSeek/Kimi MLA + unified memory false-positives. The new test_write_loc_id_space.py only covers the pool door, not these producers.
Suggestion: Switch those call sites to KVWriteLoc.for_batch(forward_batch, ...). The tag is free on static pools (requires_translated_write_loc is false).
There was a problem hiding this comment.
Fixed, and there were more than the three you named. Writing a check to verify your list turned up seven further sites — the ROCm mirrors of the same DeepSeek MLA path (`forward_mha_rocm.py`, `forward_mla_fused_rope_rocm.py`, `forward_mla_rocm.py`), `sarvam_moe.py`, and a second call in `zigzag.py` just below the one you pointed at. Twelve sites across seven files now use `KVWriteLoc.for_batch(forward_batch, ...)`.
You were also right that the existing coverage stopped at the pool door. There is now a completeness check over `srt/models/`, `srt/layers/cp/` and `srt/speculative/` that fails on any write-door call passing `*.out_cache_loc` directly, or any `KVWriteLoc(...)` built without `id_space`. It is scoped to producers reachable under a unified pool — backends whose pools never set `requires_translated_write_loc` are left alone, since the tag is meaningless there.
Landed in #40328 (4/7), which owns the field.
|
|
||
| # -- translate (virtual TOKEN ids -> physical TOKEN ids) -- | ||
|
|
||
| # TODO(unified-memory): fold this onto the fused path. Under the token-major |
There was a problem hiding this comment.
[suggestion] The ten-line TODO(unified-memory) restates the current identity of three translate helpers, names callers, and narrates why a rename is unsafe. That is PR/design history sitting on the hot translate path; the load-bearing facts (fused path requires a contiguous loc; torch path wraps a negative index; SWA relies on the -1 handling) are buried in the paragraph.
Suggestion: Keep a two-line TODO with the retirement condition, e.g. do not fold onto _translate_loc_fused until a red test covers negative-loc sink routing and strided out= slices.
There was a problem hiding this comment.
Agreed, and the TODO is now gone entirely rather than condensed.
It was describing work that #40331 (7/7) performs: that PR folds `translate_kv_loc` onto the fused path and deletes both `_translate_kv_loc_impl` and `translate_kv_loc_for_kernel`, with the red-first tests the TODO named as its precondition (negative-loc sink routing, and strided `out=` slices via the stride commit that precedes it). With the work done and the helpers deleted, the comment had nothing left to point at, so we removed it at its source rather than shipping it and withdrawing it two commits later.
For navigation: the TODO lived in #40330 (6/7), not here.
|
⛔ Known groups: ⛔ Known groups: |
|
/rerun-group unified-memory |
|
Results for 🚀 🚀 🚀 🚀 🚀 🚀 |
7de746a to
537daa9
Compare
|
/rerun-test test_hicache_unified_memory.py |
|
Results for 🚀 |
537daa9 to
19e77e7
Compare
19e77e7 to
00a48b8
Compare
Introduce the description of a token-major slot entry, alongside the existing
page-major builders and used by nobody yet:
DensePart one row family inside the entry (K, V, or the MLA latent):
layer_num rows of row_shape, layer l at
offset_bytes + l * layer_stride_bytes.
DenseEntryLayout the entry's byte map -- its parts plus entry_bytes, which
is the slot stride of every view built over it. validate()
rejects overlapping parts, an unaligned offset/stride/row,
and any part reaching outside the entry.
build_dense_views per-layer views of one part: a strided
(num_pages * page_size, *row_shape) tensor indexed by the
PHYSICAL token id, with the whole entry as the slot stride.
Two alignments come with it: ENTRY_ALIGN_BYTES pads a slot so entries stay
vector-store aligned, ROW_ALIGN_BYTES is what validate() holds each part to.
Parts may differ in row width -- only their offsets differ, never the stride --
which is what lets one entry hold K and V rows of different widths. The
page-major builders cannot: they lay out one uniform row-block array per page
and address it by a re-numbered kernel-facing id.
build_page_major_mamba_views already had a nested contiguous_strides identical
to the one build_dense_views needs, so it is hoisted to module scope as
_contiguous_strides rather than copied.
Switch the unified pool's byte law from the page-major envelope to the
token-major entry, so that a kernel-facing id IS the physical token id.
Before, one page held a uniform array of row-blocks, one per layer per side:
page bytes = [L0_K * ps | L0_V * ps | L1_K * ps | L1_V * ps | ...]
Every per-layer view was contiguous, and the block origin was folded into its
storage offset, so layer l had to be addressed by a re-numbered id
`(t // ps) * (ps * blocks_per_page) + t % ps`. Views ran past the last page
envelope, which is what view_tail_pad_bytes sized.
Now one slot's entry holds every part of that token at a fixed offset:
page bytes = [ entry(slot 0) | entry(slot 1) | ... ]
entry bytes = [ K_0 | V_0 | K_1 | V_1 | ... ]
so the per-layer view is strided with the entry as its slot stride and is
indexed by `page * ps + slot` directly. Three consequences:
* No re-numbering. kernel_page_multiplier is pinned to 1 (asserted, not
read from the spec), and translate_kv_loc_for_kernel now computes what
translate_kv_loc does. The plumbing is removed in a later commit.
* No tail pad. The views end exactly at the last slot, so the raw buffer is
allocated at total_bytes.
* Parts may differ in row width. Only the startup screen still refuses
asymmetric K/V, because the backends' write and read paths are not
audited for it -- not because the views cannot hold it.
entry_bytes() is rounded up to ENTRY_ALIGN_BYTES here rather than in the
layout types: it feeds page_bytes(), buffer sizing, the sub-pool anchors,
move_kv_cache envelopes and PD transfer, so the alignment has to flip
together with them.
The view tests are rewritten against the new byte law
`byte(t, l, K) = t * entry + l * (k + v)`, each cross-checked against an
independent (page, slot) description at page sizes 1 and 4.
test_hicache.py's unified L2 case built its pool with 8 B rows (fp16, one head
of width 4), below the 16 B row alignment every token-major part needs; it
now uses 16 B rows, the smallest the layout admits. The entry, and so each
page the test copies, stays 32 B.
fa4, trtllm_mla, cutedsl_mla and tokenspeed_mla are admitted under
--enable-unified-memory and supported by this change, but not yet validated.
From this commit on their kernels read per-layer views whose slot stride is
a whole entry, and every validation run was on H20 (SM90), where those
kernels do not run. The same holds for trtllm_mha's SM100 kernels
(trtllm-gen context and decode) and for decode context parallelism
(--dcp-size > 1). Reviewers with SM100-class hardware: please run them
before merging.
Nothing reads any of it once the unified pool stores token-major; each symbol
below has zero remaining references:
build_mha_views / build_mla_views superseded by build_dense_views.
SubPoolSpec.view_tail_pad_bytes the views no longer overrun the last
page, so there is no pad to size.
SubPoolSpec.blocks_per_page kernel-facing ids are physical ids, so
there is no per-page block scale.
KVCache.kernel_page_blocks same, and it was the probe's only input.
maybe_detect_kernel_facing_loc the probe told a kernel-facing id from a
virtual one by its in-page offset. With
one id space left there is nothing to
separate, and it was vacuous at
blocks_per_page 1 anyway. The OOB probe
at the same write doors is unchanged.
The unified page-envelope host pool had one caller of `kernel_page_blocks`:
it divided the ids it is handed by the kernel-view page stride to recover
physical slots. Those ids are physical here, so the conversion goes, and
its unit test reads pages straight off them.
…nified pool Under the unified pool, trtllm_mha refills its cuda-graph page table out of graph, before each replay, bounded by `cache_seqlens_int32`. That buffer is written only by the fused metadata kernel recorded inside the graph, so at refill time it still holds the previous replay's lengths: once a sequence grows past them, the kernel reads page-table columns that were never translated. Decoding then produces garbage without any error (GSM8K drops to ~0 with graphs on and is at baseline with them off). Until the refill derives its lengths from the batch, refuse the combination at startup and point at `--disable-cuda-graph` or another decode attention backend. trtllm_mha as the prefill backend only, or with decode capture disabled, still boots.
The token-major test rewrites had been made on an older copy of these files and undid later upstream edits along the way. Restore them, keeping only the changes the layout itself requires: - `est_time` goes back to main's values in the four retimed files; - the tests pruned as redundant stay pruned: `TestMHASpecLayerOffsets` and `test_page_bytes` (both restate `test_layout_parts_match_the_entry_formula` and a one-line formula), the write-loc `..._without_handover` case, and the MLA `test_entry_bytes_and_dim` and `test_kernel_id_follows_compaction` (the view-addressing and `test_kernel_id_is_the_physical_id` cases pin the same things). The move readback is folded back into `test_move_kv_cache_moves_page_envelopes`, as on main; - docstrings keep main's wording where the layout does not change their meaning (`reads_are_translated`, not `enabled`; no plain page-major arm and no `cutlass_mla` in the allowlist docstring), and drop the retired 2L multiplier regime and the one-element loop it left behind. Also delete `mha_entry_bytes` / `mla_entry_bytes`: nothing outside the tests called them, and a test comparing the spec against a second copy of the same formula checks nothing. The tests now derive the entry size in closed form from the row sizes and `ENTRY_ALIGN_BYTES`.
With the kernel-to-physical id rewrite gone, `get_cpu_copy` and `load_cpu_copy` on `UnifiedMHATokenToKVPool` and `UnifiedMLATokenToKVPool` only forwarded to their parents. The parents already accept `mamba_indices` / `req_pool_index` and ignore them, so delete all four and let the inherited methods serve the unified pools.
…till said layer-major The token-major layout made several descriptions false: - the `--enable-page-major-kv-layout` help and the server-arguments page said a page is "layer-major within a page"; a page is now page_size token entries, each holding every layer's K/V; - `_reserved_floor_bytes` and the matching `UnifiedKVPool` comment justified reserving all of page 0 by saying a slot-0 write touches layer blocks spread across the page. It touches only slot 0's entry; the whole page is reserved because page-aware sub-pools allocate whole pages; - `handle_page_major_kv_layout` described the unified pool as the page-major envelope layout and kept an allowlist arm for plain page-major, which the assert above it already rejects. Drop the dead arm and its comment.
The token-major view tests use `paged_view` / `paged_row_view`, which now live in `mem_cache/layout/paged_view.py`. Import lines only.
- MHASubPoolSpec: the entry diagram documents `layout()`, not a gap between two methods. - _assert_kernel_id_bound: ids no longer scale with the layer count, so stop suggesting it as a remedy. - build_dense_views: give the itemsize assert a message; drop the per-layer offset assert that validate()'s row alignment already implies. - page_major: say that the Mamba views use the same per-slot entry. - MultiEndedAllocator: drop the note about future plumbing removal.
…dense_views `build_dense_views` took the part itself next to the layout it must come from, so the two could disagree, and `page_size` / `num_pages` only to multiply them: a token-major view is `num_slots` entries long whatever the page size. Take `part_name` (looked up in `layout`) and `num_slots`.
After the switch to per-slot entries, `mem_cache/layout/page_major.py` holds the token-major entry types (`DensePart`, `DenseEntryLayout`, `build_dense_views`) and the Mamba per-slot views; nothing in it lays pages out layer-major any more. Rename it to `token_major.py`, `build_page_major_mamba_views` to `build_mamba_entry_views` (next to `mamba_entry_bytes`), and the Mamba view test after the module. Pure rename: no body changes. The `--enable-page-major-kv-layout` flag keeps its name.
…ormat The unified host pool persists whole page envelopes. Their size is the same under the old layer-major layout and the token-major one, and storage keys record only the model and ranks, so a page written by an older build (or by a static pool) would be read back as a token-major page without any error. Give host pools a `storage_format_tag` (None by default) and set it on `UnifiedPageEnvelopeHostPool`. The controller appends it to the model name every storage backend namespaces its keys with, so pages of another format miss instead of loading. The page hash is untouched: it has to stay equal to the router's.
…s enforce The `--enable-page-major-kv-layout` and `--enable-unified-memory` help and the server-arguments page still said the Triton attention, linear-attention and Mamba backends are required, and the page also said unified memory supports neither PD disaggregation nor speculative decoding. Say what holds: the backends must be stride-aware, and the allowed set and the supported PD / speculative configurations are checked at startup. The page-major flag is only accepted together with unified memory.
The token-major entry lays every part out in 16-byte units, so a model whose per-rank K, V or latent row is not a multiple of 16 bytes cannot use the unified pool. That surfaced as a bare assertion inside `DenseEntryLayout.validate()`. Check it where the sub-pool builds its layout and raise a `ValueError` naming the sub-pool, the row and the way out; the assertion stays as the layout's own invariant.
…he retelling The token-major module docstring now states the entry and the byte formula once; the unified pool classes point at it instead of restating the layout, and the trtllm_mha gate comment says what is refused and why in three lines.
00a48b8 to
9423810
Compare
2 participants
3/7 — token-major dense views: a kernel-facing id IS the physical id
Base: #40327 (2/7). 15 commits, 36 files, +1215/−1108.
This is the core of the series and the PR that changes behaviour. The other six are
hardening, adoption sweeps, deletion or consolidation around it. If you review one, review
this one — and specifically its
layout-switchcommit.Why the layout has to change
Before, one page held a uniform array of row-blocks, one per layer per side:
Every per-layer view was contiguous, and the block origin was folded into its storage
offset, so layer
lhad to be addressed by a re-numbered id(t // ps) * (ps * blocks_per_page) + t % ps. Views ran past the last page envelope, whichis what
view_tail_pad_bytessized.That scheme is only valid when every block has the same row width. The old builder said
so itself, and asserted it:
One multiplier, one row shape, one stride can only address a page that is a uniform array of
equal-width blocks. Heterogeneous row widths are not hard under that scheme — they are
unrepresentable.
Now one slot's ENTRY holds every part of that token at a fixed offset:
so the per-layer view is strided with the entry as its slot stride and is indexed by
page * ps + slotdirectly.What that buys
Each part now carries its own geometry —
DenseParthas its ownrow_shape,offset_bytesandlayer_stride_bytes;MHASubPoolSpec.entry_bytes()sumsk_row_bytesand
v_row_bytesseparately. No uniform-row assert survives.DenseEntryLayout.validate()checks that every part stays inside the entry and that no two overlap, and it requires
alignment instead of equal widths: every row, part offset and layer stride is a multiple of
16 bytes (the store kernels' vector width) and the entry a multiple of 32. A model whose
per-rank K, V or latent row is not a multiple of 16 bytes is refused with a
ValueErrornaming the row when the unified pool is built.
because the backends' read and write paths are not audited for it; what this PR removes is
the structural obstacle, not the screen.
further per-token state — for example a draft model's K and V as extra lanes of the same
entry — can be added as parts without a new builder. Under a uniform-block page that is
not expressible at all. Nothing in this series adds such a part.
Three more consequences
kernel_page_multiplieris pinned to 1 (asserted, not read from thespec), and
translate_kv_loc_for_kernelnow computes whattranslate_kv_locdoes. Theplumbing that carried the multiplier is dead as of this PR and is removed in 5/7 — it
is left in place here so this PR stays the layout change and nothing else.
total_bytes.entry_bytes()is rounded up toENTRY_ALIGN_BYTESin the switch commit rather than inthe layout types, because it feeds
page_bytes(), buffer sizing, the sub-pool anchors,move_kv_cacheenvelopes and PD transfer — the alignment has to flip together with them.The commits, and where to spend review
layout-typesDensePart/DenseEntryLayout/build_dense_viewsalongside the old builders. No caller changes.layout-switchunified_memory_pool.pyto them:layout(),entry_bytes()alignment,view_rows = num_pages * ps, raw sizing, the sub-pool multiplier pin, the hook message, and the byte-law tests.layout-cleanuprefuse trtllm_mha decode cuda graphskeep main's test editsmain'sest_timevalues and docstrings and do not bring back testsmainpruned; the free*_entry_byteshelpers go, and tests derive the entry in closed form.drop the pass-through CPU-copy overridesdescribe the token-major entryimport the paged-view helpers …,tidy the token-major layout code,take a part name and a slot count …build_dense_views(raw, layout=, part_name=, num_slots=); comments.name the entry-layout module token_majormem_cache/layout/page_major.py→token_major.py,build_page_major_mamba_views→build_mamba_entry_views, and the Mamba view test after the module. Pure rename; the--enable-page-major-kv-layoutflag keeps its name.tag HiCache storage keys …refuse misaligned per-rank rows …,state the backend requirement …,keep the addressing rule in docstrings …ValueErrorabove; help text and docs that still said the Triton backends are required; shorter docstrings.Interaction with #39477 (PD page-envelope transfers)
That PR added
_physical_to_kernel_indicesto both unified pool classes, re-numberingphysical ids by
kernel_page_blocks. Under the token-major entry that formula is theidentity, so the method is removed together with the
get_cpu_copy/load_cpu_copyoverrides that called it: the parents' methods take physical ids directly. Verified numerically on both pool
classes at page sizes 1 and 4, including the round trip and that neighbouring slots are
untouched. #39477's own SWA tail-allocation tests (17 CUDA cases) pass under this series.
The one upstream test that asserted the scaled form
(
test_swa_transfer_page_is_physical_not_kernel_scaled) is adapted here to assert bothpaths are physical, since the multiplier it pins is what this PR removes.
Tests
The view tests are rewritten against the new byte law
byte(t, l, K) = t * entry + l * (k + v),each cross-checked against an independent (page, slot) description of the same bytes at
page sizes 1 and 4, for uniform and asymmetric K/V rows. That independent description is the
safety net for the whole PR: the flat builder's addressing is verified against a second
derivation rather than against itself.
test_byte_addresses_match_envelope_formulapins thesame law in closed form, independent of the builder.
Validation
Model families exercised — a hybrid of every pool shape the unified pool supports, plus a
non-hybrid control. Perf covers the gated-deltanet hybrid (Qwen3.5-9B), the MLA hybrid
(Kimi-Linear-48B-A3B), the sliding-window hybrid (gpt-oss-20b), the tri-pool hybrid
(Inkling and Inkling-Small), and dense Llama-3.1-8B as a control. Accuracy covers the same
set and additionally the Mamba hybrid, Falcon-H1-7B. Each model is run unified against a
baseline on the same model, backend and configuration, so every delta isolates the pool.
Median inter-token latency, unified vs baseline: +0.91% (measured run-to-run noise floor
for this metric: 1.9%). Median GSM8K delta at n=200: −0.25 pt (binomial delta-sigma ≈ 3.0 pt).
Re-validated on the previous base (0318a8d) after review feedback: the stack tip (7/7) gives +0.95% median inter-token latency over 40 paired configurations (median per-configuration shift against the previous run of the same 40: +0.12 pp) and −0.25 pt median GSM8K over 26, with the full unit scope green.
The accuracy median excludes trtllm_mha with decode cuda graphs, which is refused at startup
for the reason given below; including it the median is −0.50 pt.
End-to-end on B200
From this commit on, every allowlisted kernel reads per-layer views whose slot stride is a
whole entry. Besides the 2/7 kernel parity tests, each admitted SM100 backend was run with
unified memory against a baseline on the same model, flags and box (GSM8K 5-shot, n=200):
trtllm_mha with decode cuda graphs is refused at startup under unified memory (below);
without graphs it decodes 9–13% slower than baseline, from the translate launches an eager
step adds. Not covered: fa4 on the gated-deltanet hybrids (fa4 cannot run head_dim 256 with a
paged cache on SM100, with or without this PR) and trtllm_mha at tp > 1.
trtllm_mha decode cuda graphs are refused under unified memory
A defect that predates this series makes trtllm_mha decode under cuda graphs read
untranslated page-table columns: under the unified pool, the page table is refilled out of
graph before each replay, bounded by
cache_seqlens_int32, which only the in-graph metadatakernel writes — so the refill uses the previous replay's lengths. Measured: GSM8K
0.000/0.020 with graphs on against 0.900/0.900 with them off, baseline 0.905/0.900 in both.
Before this PR that configuration could not boot (its prefill path OOM'd copying a whole
sub-pool). This PR makes it bootable, so it also refuses it:
--enable-unified-memorywithtrtllm_mha as the decode backend and decode cuda graphs on fails at startup, pointing at
--disable-cuda-graphor another decode backend. trtllm_mha as the prefill backend onlystill boots. The gate goes once the refill derives its lengths from the batch.
HiCache storage
The unified host pool persists whole page envelopes. Their size does not change with this
PR, but their byte order does (layer-major → token-major), and storage keys recorded only the
model and ranks. Host pools now declare a
storage_format_tag;UnifiedPageEnvelopeHostPoolsets one, and the controller appends it to the model name every storage backend keys pages
on. Pages written by an older build, or by a static pool, therefore miss instead of being
read back in the wrong layout. The page hash is unchanged, since it has to match the
router's.
CI States
Latest PR Test (Base): ⏳ Run #36944546351
Latest PR Test (Extra): ⏳ Run #36944545943
Latest PR Test (AMD ROCm 10): ⏳ Run #36944546302