refactor(kv-cache): one declaration per fact about the KV pools - #2147
Open
valarLip wants to merge 14 commits into
Open
refactor(kv-cache): one declaration per fact about the KV pools#2147valarLip wants to merge 14 commits into
valarLip wants to merge 14 commits into
Conversation
The lightning indexer carried two implementations of the same selection behind a flag -- a paired bitonic sort and a packed-key tl.topk -- plus a second decode score kernel, and two kernels that were already dead before any of this. Timed alone over the deployed shapes the packed selector is 2.5-4.1x the bitonic one, and the tiled decode score reads each index-K block once per request instead of once per (query row, index head), which is 8x at tp2 under MTP-4. Keep the winners, delete the rest: 1155 -> 930 lines, and ruff drops from 4 pre-existing violations to none. The two selectors were checked to emit byte-identical ids, order and sparse_bt across 100+ shapes before the loser was removed, with the comparison armed first so a green result meant something. They part only where two blocks hold the same fp32 score -- about once per 25k rows on a 3-mantissa-bit fp8 index cache -- and that is an ordering, not a selection: both blocks are in either way. The survivor resolves it to the higher id, which the sort network it replaces never did in any statable direction. Also in here: - Decode's score chunk count is now derived from the size it implies. 800 blocks over 64 chunks rounds to 13 blocks each, which covers 832, so the last two chunks used to launch only to find themselves past the end. - The fused emit scatters a [slot, page] rectangle in one masked store instead of looping the page axis, on a path that runs per (query token, kv head). - The two wrappers shared 15 lines of buffer allocation verbatim and a 24-argument launch that differed in five places; both are now one helper, so what actually separates prefill from decode is legible. First tests for this module, 86 of them: the launch policies run in CI without a GPU, the kernels sit behind a skip. They assert the op's definition rather than a second implementation -- a row with no more candidates than topk must select all of them and report its own causal length, and sparse_bt must hold the full blocks by descending score with the partial tail last. One of them found the chunk-count overshoot above. GSM8K (1319, 5-shot, fp8 cache) scores 0.9469 twice here against 0.9393-0.9447 for the code this replaces. That is a no-regression check and not a win: at that prompt length a row has 9-14 blocks against a topk of 16, so every block is selected and the selector is not what is being measured.
Every consumer takes max(floor, rows), so the floor only adds split-K where the grid is too small to fill the machine and is an identity everywhere else. The splits are numerically inert: each CTA owns a disjoint KV-column range, with no cross-CTA partial sums. 512 left the machine idle on long contexts, which is exactly where rows are scarce -- the logits buffer widens as the context grows, so the Q chunk and with it the row count shrink. Decode rows=128 W~32768 54.1us -> 51.2us; prefill rows=1024 224.6us -> 206.2us. 4096 is not any single shape's optimum, because CTA-count quantization makes the ordering shape-specific, but it has the smallest worst-case regret of the values tried. Worth re-tuning against a real workload mix. The comment at the prefill call site said the floor was for decode. It is not: a long-context prefill is short on rows for the same reason.
… serves
`StateArena` was never state-specific. Its `buf=` / `slot_stride=` /
`live_entries=` hooks already let DeepSeek-V4 carve it out of a plane it
shares with paged blocks, so its only tie to the STATE pool was the name --
and a paged KV pool is about to read the same field declarations through a
sibling materializer. What separates the two is where the entry axis sits,
which its own docstring already said ("laid out entry-major"), so that is
what it is called now. Left as it was, the pair would have read as "the
state one" against "the by-layer one", implying the first cannot serve PAGE.
`entry` is not a new word for this: `sub_pool_spec` already sizes both pools
in entries -- `entry_bytes`, `entries_per_req`, "entry class" -- because it
is the unit-agnostic term. A PAGE entry is a block, a STATE entry is one
request's state.
StateField -> EntryField
StateArena -> EntryMajorArena
SplitStateArena -> SplitEntryMajorArena
state_arena.py -> entry_arena.py
Splitting the declarations out of the arena was tried and reverted, and the
module docstring now records why, because it is the obvious thing to try
next: the `pool_layout` guard bans relative imports too (`imported_roots`
scores a relative import as reaching `atom`), so every member of the package
is a leaf, and a field list cannot live apart from an arena over it.
No behaviour change. Bodies are byte-identical modulo the renames except for
three strings that name the old type.
Seven places spelled out the SHUFFLE shapes of a paged MHA block. They agreed
on every byte and disagreed on a view: the Eagle3 draft's builder declared V
non-transposed where the others declared it 5-D, and `_bind_kv_cache_to_modules`
did the same. Nothing could catch it -- the byte count is identical either way,
and the only consumer that cares reads `v_cache.ndim` at runtime, densifying
anything that is not already 5-D. On MiniMax-M3 that relaid the whole 1.4 GiB
draft pool once per chunked-prefill forward: 7.1% of GPU time, 61% of what the
draft model cost.
`MhaKvPool` is that declaration and the two arenas over it. It takes five
numbers -- layers, block_size, num_kv_heads, head_dim, kv_dtype -- and never
sees a ModelRunner, a config or a forward_vars, which is what makes a second
one askable for. That was the actual obstacle: a backend also owns per-step
metadata, and `CommonAttentionBuilder.__init__` ends by writing into
`model_runner.forward_vars`, so building a second builder to get a second pool
would have overwritten the target's buffers. Splitting the two jobs is what
lets the draft have a pool without anyone copying shapes.
So a draft no longer carries an attention. `Eagle3DraftBuilder` hardcoded MHA;
`DraftKvBuilder` is handed a pool and keeps only the part that is about being a
draft -- its blocks are the target's blocks, re-paged at its own block size.
Which pool it gets comes from the draft's own config through the same selector
the target uses, and the backend answers `make_kv_pool` for its own flavor
(None means "a draft of this flavor binds into the target's pool", which is
MLA and the default). Adding a flavor is one method on that backend; nothing in
spec_decode imports an attention any more.
Sizing derives from the same declaration, so a block cannot be charged for at
one size and built at another. M3's `entry_bytes` is 3502080 before and after.
Also folded in:
- `LayerMajorArena`, the sibling of `EntryMajorArena` for a pool whose layer
axis is outermost. It has no `entry(i)` on purpose: a block's bytes are
`blocks` apart, so nothing can copy or register one as a unit. That is the
deletion condition too -- turn the pool block-major and one class serves
both.
- `TritonMHAMetadataBuilder.build_kv_cache_tensor` shrinks to the parent's
bind plus the layout flag, which was the only thing it added.
- `GDNAttentionMetadataBuilder` allocated its own tensor but inherited the
bind, so it needed a pool as well; its shorter layer axis is now one
`_kv_pool_layers` hook that sizing and allocation share.
- `_bind_kv_cache_to_modules` runs the builders' own bind loop instead of a
second copy of it, after `adopt_imported_kv_pool` re-derives the pool over
the IPC-imported buffers. Handing those buffers to an arena is also the
check that they fit the declaration.
Verified on MiniMax-M3-MXFP4 + EAGLE3, tp8 unless noted:
- trace (tp4, ISL 100k): float8_copy 416 -> 0 on all four ranks, with
cp_mha_gather_cache still at 104, so the path runs rather than being
skipped.
- GSM8K 5-shot chat: 0.9431 flexible / 0.9439 strict against a 0.9469
baseline and a 0.93 threshold. Run with --level 0, which the CI entry does
not use, so it is a band check rather than a point comparison.
- MTP acceptance 73.55% over 123k draft tokens against a 0.715 threshold --
the sensitive one here, since a draft KV write the verifier reads back
differently collapses it rather than nudging it.
- Unit tests unchanged at 4 pre-existing failures, all aiter-side.
Not exercised: P/D, MiMo-V2, MLA targets. GDN was caught by reading the
inheritance chain, not by a test.
`MhaKvPool` left MLA as the asymmetric half, and MLA's copy of the problem was
worse: `sub_pool_specs`, `allocate_kv_cache_tensors` and `build_kv_cache_tensor`
each wrote the literal `576`, while `mla_kv_entry_dim` -- the expression that
exists precisely because a NoPE model declares a padded width that need not
equal `kv_lora_rank + qk_rope_head_dim` -- was read only by `mla_qk_head_dim`
and the Kimi backend. Every model in the tree happens to be 576, so it is a
latent split rather than a live one; it is closed here as a substitution that
is provably an identity today, not as a bugfix.
`MlaKvPool` is that declaration. MLA packs k_c and k_pe into one row and
absorbs V, so a block is one field rather than four, plus an indexer field the
sparse variants add -- owned by some layers only, and holding one row per token
or, where the indexer pools, fewer. Differing layer counts and row counts are
what `EntryField` already says, so both are fields and neither needs a second
pool.
An entry is the *scheduler* block, which is what `page_pool` charges. MLA
usually pages at size 1, so `block_ratio` rows of the allocation make one
entry; keeping the entry at the scheduler block is what lets a row of
`region_tensors` be the unit a transfer registers. Three things fall out:
- `get_kv_transfer_tensors` drops its `* block_ratio`, because `stride(0)`
is already the bytes a block owns;
- `_index_cache_block_bytes` goes, base and Kimi override both, since the
field's own arithmetic answers it;
- `runner.physical_block_size` goes. It was the builder's own `block_size`
mirrored onto the runner, with five readers, all of them inside builders.
`kimi_mla_gdn` gets the same pool over a shorter layer axis (`_kv_pool_layers`),
and `PageUnitGeometryMixin` reads the pool's view through a new
`_page_unit_kv_cache` hook, symmetric with the index hook it already had.
The 39 tests this broke were all real: MLA's fixtures pinned the allocator's
call shape and the tensor doubles it sliced. They now assert the views a reader
binds, and the doubles shrank to a fake pool -- `layer(name, i)` and
`region_tensors()` are a much narrower interface than a tensor that has to
behave like a tensor. One of them earned its keep immediately, catching that
the transfer path must use the layer count the pool was *built* with: a hybrid
caches for fewer layers than the model has.
Also here, found by the M3 gate rather than by tests: the block count a pool is
built at is `config.num_kvcache_blocks`, not `pool_plan.paged_entries`. Sizing
runs in every runner subprocess and they disagree by a few blocks, so EngineCore
takes one answer and broadcasts it; the plan is this rank's own estimate. Using
it made the KV pool larger than the sparse index cache sized alongside it, and
`fused_qknorm_idxrqknorm` refused the mismatch at capture time. The comment says
so at the one place that reads it.
`wait_server_ready.sh` learns `openai_server.py: error:` -- argparse refusing
the command line kills the process in under a second, and the poll used to wait
out its whole budget. A bare `Traceback` was tried and reverted: torch logs
whole tracebacks at WARNING during Inductor warmup, hundreds of them, and it
failed a healthy start on first use.
Verified:
- MLA byte parity, measured both ways: Kimi-K3 `entry_bytes` is 1769472 on
this commit and on HEAD, and R1-0528's 571392 decomposes as 62 layers x 16
tokens x 576 x 1B, which is the old formula's expression.
- DeepSeek-R1-0528-MXFP4 MTP3 tp4: GSM8K 0.9507 flexible against a 0.93
threshold, MTP acceptance 64.74% against 0.593 over 150k draft tokens.
- MiniMax-M3 + EAGLE3 re-run after the `physical_block_size` removal:
float8_copy still 0 on all four ranks, cp_mha_gather_cache still 104,
`entry_bytes` still 3502080, throughput flat.
- Unit tests unchanged at 4 pre-existing failures, all aiter-side.
Not exercised: Kimi-K3 accuracy. It scores 0 with `!!!!` output on this box,
and reproduces identically on HEAD down to the first token, so it is upstream
of this change -- the aiter working tree is mid-rebuild. P/D is untested as
before.
…s block count Two changes, in one commit because they interleave inside `AiterAttentionMetadataBuilder.allocate_kv_cache_tensors`; splitting by file would leave a first commit that does not import. 1. MiniMax-M3's indexer key cache becomes a field of `MhaKvPool` It was the last hand-written declaration on the MHA side, and it was written twice: a formula added to `block_bytes` in `sub_pool_specs`, and a `torch.zeros` in `allocate_kv_cache_tensors`. The two even read different block-count attributes -- exactly the shape that produced fused_qknorm_idxrqknorm.cu:851 index_cache must contain at least num_blocks * block_size * 128 elements when one of them moved. Now `index_layers` / `index_dim` / `index_dtype` are constructor arguments, symmetric with `MlaKvPool.index`, so one expression answers sizing and allocation and a divergence is unrepresentable. Three things fall out, none of them cosmetic: - The old sizing charged `sparse_layers * backend_page * dim` while the allocation was `blocks * block_ratio` pages deep, i.e. it undercharged by `block_ratio`. A field is declared per scheduler block, so the count is now right for any ratio. Every shipped M3 config runs the two block sizes equal, which is why no budget ever came up short. - `runner.sparse_attention_index_cache` is gone rather than re-created. Nothing outside this builder read it, and neither `ModelRunner.release` nor the rollout memory manager named it -- so it was a bare attribute that survived a rollout sleep. The pool is now its only owner and `MhaKvPool.release` frees it. - `_sparse_attention_topk_cache_state` moves out from under the indexer gate. It is a per-run memo keyed off `use_index_cache`; being nested under the cache claimed a dependency that does not exist. P/D now refuses this model instead of guessing: the IPC export ships `kv_cache` and `kv_scale` only, so a decode side has no imported buffer for the indexer cache, and allocating a fresh one would give it zeroed keys and a top-k that picks the wrong blocks without saying so. 2. The block count arrives as an argument `ModelRunner` used to reach into the builder for `block_ratio`, multiply, hang the product back on itself as `num_physical_kvcache_blocks`, and let five backends read it off again -- two of them reading `config.num_kvcache_blocks` instead, for the same quantity. Sizing runs per subprocess and the answers differ (48425 / 48339 / 48310 on M3 tp4) while EngineCore broadcasts one, so reading the wrong attribute builds one pool at one count and something sized off another at a second. `allocate_kv_cache_tensors` now takes a keyword-only `blocks`, the same argument `adopt_imported_kv_pool` already took, and each builder records `self.num_blocks = blocks * self.block_ratio` -- its own page, which is what its kernels index. `num_physical_kvcache_blocks` is gone. The units are the point, and they are kept apart even where they are the same number: `blocks`, a pool entry and `req.block_ids` are the scheduler's; `num_blocks` is the backend's. Conversions are written out; none is elided because a ratio happens to be 1. `KVTransferTensors.num_blocks` becomes `field(init=False)`. Three backends used to state it themselves, in two different units, while the field's consumers address `req.block_ids`. It is now set once by `ModelRunner.set_kv_cache_data`'s caller -- after the draft appends its regions, and the only place holding both the complete region list and the scheduler's count -- and `set_block_count` checks every region divides into exactly that many units. The failure it catches has identical total bytes: a region registered per token rather than per scheduler block has `block_ratio` times the units at `1/block_ratio` the stride, so both ends agree on the address arithmetic and disagree on the stride. No byte count distinguishes them. Tests `tests/test_kv_pool_contract.py` is new: an AST check that every implementation of the allocate hook takes `blocks` keyword-only with no default. A source grep would have asserted on the text of the code it tests; the signature is the invariant that makes the runner attribute unnecessary. Also 8 cases pinning the indexer field, 4 pinning the block id space, and one pinning that an indexer region is one scheduler block wide -- the single place the fold is deliberately not byte-for-byte with what it replaced. `tests/test_mla_index_cache.py`'s allocate fixture now runs at `block_ratio` 16, not 1: at 1 it could not tell the two counts apart. Verification - unit tests 5443 passed / 4 failed, against a 5415 / 4 baseline; the 4 are pre-existing aiter-side failures (`_top_k_per_row_decode` signature, `_is_gfx1250`), neither of which CI runs - `entry_bytes` byte parity, offline and in the server log: 2961408 == 2961408 for the target pool, `sub-pool kv: entry_bytes=3502080` unchanged - MiniMax-M3-MXFP4 + EAGLE3, tp4: GSM8K 0.9363 and 0.9431 over two runs, acceptance 75.37% / 73.00%. Five runs on this branch span 0.9363-0.9462 and 0.9431 appears in three of them, across both this change and the one before it, so the arms overlap - DeepSeek-V4-Flash-DSpark, tp1: GSM8K 0.9484, acceptance 65.29%, 4.26 toks/fwd -- the recorded baseline on all three - revert checks: swapping arena order reddens the region-order test; removing GDN's `blocks` parameter reddens that implementation's two contract cases and nothing else Not exercised: P/D (`set_block_count` has unit coverage only), MiMo-V2, and Kimi-K3, whose accuracy is blocked on an unrelated aiter-side fault. `tests/test_mla_index_cache.py` skips at module level without aiter, so CI collects it and runs none of it.
The target's KV, its dequantization scales, MiniMax-M3's indexer keys and an
EAGLE3 draft's sibling pool were up to five `torch.zeros` under five runner
attributes, while the budget had long treated them as one entry class -- two
specs naming the same class sum, which is how a draft rides the target's block
ids. Sizing said one pool; memory was five.
They are one buffer now. `ModelRunner._carve_paged_pool` asks each builder how
many bytes its pool takes, places them with `plan_regions`, and hands each its
region. `paged_pool_bytes` answering zero means "I allocate my own", which is
where DeepSeek-V4 stays: half of what its PAGE entry is priced for lives in
the plane pool `allocate_per_req_cache` makes later.
What that buys:
- `allocate_kv_cache_tensors(..., buf=)` and `adopt_imported_kv_pool(blocks,
buf)` collapse into one expression, `pool.allocate(blocks, dev, buf=...)`.
The P/D decode side carves with the same walk instead of being told the
offsets, so the two ends cannot disagree about where a region starts.
- The IPC export ships one handle instead of two, and the sparse indexer
cache travels inside it. The refusal in the MHA adopt path went with that:
its reason -- "the export ships kv_cache and kv_scale and nothing else" --
stopped being true.
- `runner.kv_scale` and `runner.index_cache` are gone, with their entries in
both release lists. `kv_cache` is the pool.
- MiMo-V2's per-module pools take successive slices of their builder's
region, so they are exported and freed like every other model's.
Each pool carries its groups as one list that the price and the allocation
both walk (`field_groups` -> `entry_bytes` / `carve_layer_major`), so a fourth
group cannot be added to one and missed by the other.
The one thing checked on every start: a builder either allocates its own pool
or takes exactly what its own PAGE spec was charged. A region sized off a
second reading of the layout surfaces as a wrong answer while serving, never
as a bad number at startup.
Unit tests 5460 passed / 4 failed, the four being aiter-side failures already
on HEAD. MiniMax-M3 tp4: `entry_bytes=3502080`, byte-identical to before;
GSM8K 0.9454 and 0.9447 at 73.28% / 73.08% MTP acceptance, against five
earlier runs on this branch spanning 0.9363-0.9462; trace
`float8_copy_kernel_cuda` 0 on all four ranks and `cp_mha_gather_cache_kernel`
104. DeepSeek-V4-Flash-DSpark tp2, which opts out: 0.9507, 65.12%, 4.26
toks/fwd against a 0.9484 / 65.09% / 4.25 baseline.
Not exercised: P/D (the single-handle format has unit coverage only), MiMo-V2,
Kimi-K3.
…ules declare the pools A layer's row in a pool was recovered from `layer_id` by knowing the model's shape: `layer_id // full_attention_interval` for a Qwen3-Next full-attention layer, its complement for a linear one, `num_full_attn + (layer_id - mtp_start)` past a draft's boundary, a dict keyed by layer for Kimi's two kinds. Four spellings of a counter, each carrying an assumption about a model it did not name -- and `layer_id` is already a counter, of every layer the walk registered. It is `take_slot(kind)` now: a module's row is its position among the modules of its kind, which is the order the walk visits them in. The runner clears the counters before each walk, because the walk is the runner's -- a pool is bound again on a P/D import and on a rollout wake, and that also fixes a draft builder whose own counter was never reset (a wake would have handed its one-layer pool row 1). The pools themselves come from the modules too. A module says what it caches, so `MhaKvPool`s are the modules counted by `(num_kv_heads, head_dim)`: one for almost every model, two for MiMo-V2-Flash, whose sliding-window layers carry a second head count -- a second geometry in one block-id space, which is what sizing had been summing for it all along. Which layers of a hybrid are full attention, where a draft's stack begins, what a PP stage holds: all of it is now a property of what modules exist, and one predicate answers it for the sizing walk and the bind walk both, so a layer cannot be charged for and left unbound. `is_mimo_v2` is gone from the tree with the last of the per-layer allocation it gated -- and with it MiMo-V2's cursor into its own region, its `_kv_layer_cache_store`, the P/D refusal, the transfer path's second branch, and the Triton backend's refusal, which named that allocation path. The GDN-hybrid builder loses both of its pool overrides: counting the modules it owns leaves the linear-attention layers out without it saying so. Also gone, having lost their last reader: `runner.mtp_start_layer_idx`, `runner.full_attention_interval`, `runner.kda_attention_layers`, Kimi's two layer-index dicts, and `_kv_pool_layers` on the MHA side. `entry_arena.carve` is the placement rule the runner's pool, a pool's field groups and a builder's several pools now share. Unit tests 5460 passed / 4 failed, the four being aiter-side failures already on HEAD. `tests/test_bind_walk_slots.py` pins the counter against the deleted formulas on the shape they were written for. MiniMax-M3 tp4: `entry_bytes=3502080` unchanged, GSM8K 0.9401 at 73.14% acceptance. DeepSeek-V4-Flash-DSpark tp2: 0.9500, 65.13%, 4.26 toks/fwd, against 0.9507 / 65.12% / 4.26 measured on the parent commit the same day. Qwen3.5-27B tp2, the hybrid that exercises both counters: sizing identical at `entry_bytes=266240`, and a bind probe on both arms agrees line for line on all 64 layers -- shape, stride, dtype and storage offset -- so the mapping is unchanged rather than merely close. GSM8K interleaved two runs per arm overlaps (0.8767 / 0.8633 against 0.8333 / 0.8667); this configuration's run-to-run spread is 3.3pp. Not exercised: P/D, MiMo-V2, Kimi-K3.
ModelRunner carried four predicates over `model_type` -- `is_deepseek_mla`,
`is_qwen_next`, `is_kimi_linear`, `is_deepseek_v4` -- and set four booleans
from them, which the selector then re-composed into one answer through a chain
of ifs. The priority between them (V4 beats a latent rank, a hybrid beats one
too) existed only in the order of that chain, and the runner had to know
strings like `qwen3_5_moe_text` to feed it.
It is `Family`, one value, decided by `attn_family(hf_text_config)` next to the
dispatch it feeds. An enum rather than a string, so a name that does not exist
raises instead of falling through to MHA.
Two of the four predicates were also wrong, and the shape says so:
- MLA is the presence of a latent rank, not membership of a list. The list
named six model types and missed `k3_dspark`, whose config carries
`kv_lora_rank=512`. That config reaches the engine only as a draft, and
`draft_kv` already asked the right question -- `bool(kv_lora_rank)` -- so
the mistake sat on a path nothing ran. It is one question now, asked the
same way for a target and a draft, which also retires that fifth spelling.
- An EAGLE draft is not an attention. `eagle` names a decoding algorithm and
the old predicate branched on it, then hard-coded `deepseek_v2`/`deepseek_v3`
for the stack underneath -- so an EAGLE trained on any other MLA base read
as MHA. Two drafts of one target need not share a stack: the deployed
MiniMax-M3 draft declares `model_type=llama`. Nothing about the algorithm
appears here now.
`is_deepseek_v32` goes with them, renamed `has_mla_indexer`: every one of its
twelve readers asks whether this model's MLA rows carry an indexer key cache,
and GLM-5.2 is sparse too, so the name was already wrong for one of its own
users -- `aiter_mla` reassigned it to `self.is_sparse` before using it. One of
those readers was redundant and is gone: `forward_vars["sparse_kv_indptr"]`
exists exactly when the model has an indexer, so the buffer's own presence is
the test, which is how `eagle_proposer` had always spelled it.
`block_size` leaves the selector entirely. It was threaded through all three
functions and read by none.
Verified by dispatch parity over every config on the box -- family and indexer,
old spelling against new, 16 configs. Fifteen agree; the sixteenth is
`k3_dspark`, mha -> mla, on the path described above. Unit tests 5482 passed /
4 failed, the four being aiter-side failures already on HEAD, with 22 new ones
pinning the two corrections; reverting either makes them red.
GSM8K per family: MiniMax-M3 tp4 0.9356 at 73.24% acceptance (nine runs on this
branch span 0.9356-0.9462); DeepSeek-V4-Flash-DSpark tp2 0.9500, 64.86%, 4.24
toks/fwd against 0.9484/65.09%/4.25; GLM-5.2-MXFP4 tp4 0.9340, CI threshold
0.92 -- the only config on the box with an indexer, so the rename's riskiest
path; Qwen3.5-27B tp2 0.8533, inside the 0.8333-0.8767 its own noise spans.
Kimi-K3 could not be run: aiter-side breakage unrelated to this.
…once
Four hand-written byte formulas in the DeepSeek-V4 backend each described a
layout that was also described somewhere else, and nothing tied the copies
together:
- how many planes this build has was answered independently five times --
`rope_row_bytes` returning 0, two truthiness filters, `_kv_planes`'
`is not None`, and an `if self._kv_fp8` in the carve. Only the last one
fails loudly when it disagrees; the carve just builds one plane fewer.
- what one indexer block costs was two formulas, one for sizing and one for
allocation. They agreed by inspection, and a disagreement would have made
sizing quietly wrong.
- where the two regions sit inside that block was spelled out four times,
the fourth being a scale view recomputing `rows * head_dim` from the pool
shape. A view that starts one region early reads the tail of the data --
valid-looking numbers, not a crash.
- `head_dim` appeared in both the state window's shape and its row width,
which have to agree and had nothing making them.
Each becomes one `EntryField` list in `pool_layout/v4_pool_fields.py`, the
sibling of `v4_pool_geometry`: that one says where rows go, this one says what
a row is. The backend keeps the policy -- the gfx fallback off fp8,
`--index_cache_dtype`, whether this GPU has an FP4 indexer -- and the shapes
follow from the answer with no branch of their own. A bf16 build declares one
plane instead of declaring two and setting one to zero bytes.
`indexer_block_regions` packs with nothing between regions rather than reusing
`entry_arena`'s aligned walk. Every consumer recomputes a region's start from
`rows` and the head dim rather than reading one, so borrowing the aligned walk
would have introduced exactly the gap that goes unnoticed. Both shipped
layouts land on 256 B anyway, which is why the tests scan head dims and pin
one case that does not.
Relaxes one existing invariant, deliberately: `pool_layout/` members may now
import a sibling member. The rule that matters is that a member is reachable
without aiter or the rest of atom so CI can import it, and a sibling is held
to that same rule by that same test -- nothing new is reached. Reading every
relative import as "leaves the tree" was the cheaper check and was already
charging for itself, in `entry_arena`'s note that a field list and an arena
over it could not be split apart. `test_a_sibling_import_names_a_member` is
what keeps the exemption to files the rule actually covers.
Verified: sub-pool bytes byte-for-byte unchanged on V4-Flash-DSpark tp2
dspark-5 (`kv 977152 / state 18652160`); GSM8K 1319 0.9522, DSpark acceptance
4.25-4.35 toks/fwd, which is the probe that fails first if the scale region
moved; unit tests 5527 passed with the four pre-existing failures; ruff
unchanged at 905.
…he walk
A pool row was a counter the bind walk advanced: `take_slot(kind)` on the
attention builders, `_next_layer_id` on the draft's own pool, and `layer_id`
itself where a dense model made the two agree. Three implementations of one
idea, and each of them assumed the runner's walk visits in the same order the
sizing walk did.
It is `pool_rows[kind][module]` now, assigned by one walk over
`_pooled_modules` that also answers how many rows to buy. Sizing and binding
read the same map rather than agreeing by convention, and the case that
convention did not cover is the reason: the runner's bind walk offers a draft's
modules to the target builder whenever a sibling builder declines them, and a
counter hands one the next row -- past the end of a pool sized without it --
without saying anything. A map raises.
With the row a property of the module rather than of the count, every
config-derived layer count that predicted the same thing goes:
`_get_total_num_layers`, `_num_draft_kv_layers`, `_num_cache_rows`,
`_owns_module`, `runner.num_full_attn`, `runner.num_gdn_attn_state`, and the
`num_kv_heads` / `num_draft_layers` parameters of `allocate_kv_cache_tensors`,
which no implementation read. Counting the modules is also what PP and a draft
stack needed those predictions to special-case.
`build_kv_cache_tensor`'s `layer_id` goes with them. It was the walk's ordinal
under a name that says otherwise, and two backends had already worked around
it -- DeepSeek-V4 re-deriving the layer id from the module prefix, MLA reading
`module.layer_num` and calling it `global_layer_id`. Two things it was standing
in for were wrong:
- Kimi keyed `index_cache_layer_map` with it. That map is keyed by GLOBAL
layer ids, which is why `aiter_mla` deliberately uses `module.layer_num`
there; on a non-first PP stage local row 0 may be global layer 39.
- the RapidServe import path keyed its whole `kv_cache_data` by
`enumerate(...)` while the allocating path keyed by `module.layer_num`.
One registry, two keyings, agreeing only while the first bound module was
layer 0.
`PoolRowsMixin` lives in `pool_layout/`, which imports nothing: a draft's own
pool needs it too, and `spec_decode` promises in its own docstring to import no
attention.
Verified: unit tests 5530 passed with the four pre-existing failures; ruff
unchanged at 905. DeepSeek-R1-0528-MXFP4 MTP3 tp4 GSM8K 0.9439 (CI threshold
0.93) at 2.93-2.99 toks/fwd, binding 62 mla rows = 61 target + 1 draft, the
count the deleted arithmetic gave. Qwen3.5-27B sizing byte-identical to HEAD
(kv 877755 x 266240, state 256 x 39223296) and 48 + 16 rows, again the counts
the deleted formulas gave -- but its GSM8K read 0.85 against a same-session
HEAD control of 0.89, one run each, and that gap is not yet explained. K3 does
not start on this machine, so its half is covered by unit tests only.
Contributor
🏷️ CI GuideRuns automatically on every eligible PR before approval:
Heavy model tests:
|
A pre-existing RUF012 that CI had never been asked about: reviewdog filters to the pull request's diff context, and until this branch moved the lines around it, nothing in that context reached it.
`test_minimax_m3_index_topk` imports `index_topk`, which defines `@triton.jit`
kernels -- and a decorator runs at import. CI has no triton, so the import was
a collection error, and one of those takes the whole run down rather than the
one file. The tests now `importorskip("triton")`, which is a visible skip.
Their module also took triton from `vllm.triton_utils`, as does `sparse_attn`.
That is how they arrived from vLLM; the repository decided against it when the
same files were vendored into `fla_ops` -- "couples to vLLM internals. ATOM
uses plain triton imports throughout" -- and 73 files import it directly. These
two were the last that did not.
The skip means CI runs none of these, including the launch-policy tests that
need no device. Splitting the plain-Python half into a module of its own is
what would get them run there.
The PostToolUse hook ran `ruff check --fix` and captured its output without ever reading it. Rules `--fix` cannot fix -- RUF012 is one it only advises on -- therefore vanished. CI does not let them vanish: reviewdog filters to the pull request's diff context, so a finding on a line an edit merely moved is a red check whatever the repository's total was before, which is why counting `ruff` findings before and after a change is not a gate for it. Whatever survives `--fix` now goes back on stderr with exit 2, which for a PostToolUse hook means "show this, the tool already ran". Pre-existing findings in the edited file are reported too, on purpose: those are exactly the ones CI picks up once an edit lands beside them. Then everything it reported in the files this branch touches: SIM102, RUF012, RUF013, RUF059, PLR1730, G201/TRY401, and a debug log for S112. Three BLE001 sites are `noqa` with a reason rather than narrowed -- each is a deliberate best-effort probe of an optional dependency, and guessing which exception it raises is how a probe becomes a startup crash.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What this is
The same fact about the KV pools was declared in several places, and two of
those places disagreed. This branch collapses each fact to one declaration and
deletes the copies. The measured cost of one such disagreement was 7.1% of a
MiniMax-M3 step; the rest were latent.
The declarations
What one entry holds. MHA blocks, MLA blocks and DeepSeek-V4's rows and
indexer blocks each used to be spelled out by whoever priced the pool, whoever
allocated it, and whoever bound a view into it. Each is now one
EntryFieldlist (
mha_kv_pool,mla_kv_pool,pool_layout/v4_pool_fields), read by allthree. For V4 that also collapsed five spellings of "how many planes does this
build have" -- one of which, the carve, built one plane fewer without failing.
Where those bytes go. One
torch.zerosfor every paged pool, one regionper builder, placed by
plan_regions. A draft that cannot share the target'spool is the next region rather than a second allocation.
Which pool row a module gets. A counter the bind walk advanced becomes
pool_rows[kind][module], assigned by the same walk that counts the rows tobuy. A module the walk never saw now raises instead of being handed the next
row -- which is what happened to a draft's modules whenever a sibling builder
declined them.
Which attention a model has. Four
model_typepredicates onModelRunnerplus four
use_*flags plus an if-chain that re-combined them become oneFamilyvalue decided where the backends are. Two live misclassifications fellout:
k3_dspark(haskv_lora_rank, was read as MHA) and an EAGLE draft overa non-DeepSeek MLA base.
Bugs found on the way
index_cache_layer_mapwith the bind ordinal. That map is keyedby global layer ids --
aiter_mlausesmodule.layer_numthere deliberately,because on a non-first PP stage local row 0 may be global layer 39.
kv_cache_databyenumerate(...)whilethe allocating path keyed it by
module.layer_num. One registry, twokeyings, agreeing only while the first bound module was layer 0.
_declare_kv_poolstold geometries apart from named row spaces byexclusion, so the first hybrid to add one crashed at startup. Caught by the
Qwen3.5 gate; the key is a
KvGeometrytype now.Also here, unrelated to the above
6487407dakeeps one MiniMax-M3 indexer selector rather than two, ande9b49342fraises the FP4 MQA grid floor to 4096.Test plan
pre-existing on
main(test_deepseek_v4_wo_a_dequantx3,test_dcp_topkx1) and fail identically with the branch stashed.ruffunchanged at 905;blackclean.(
kv 977152,state 18652160); GSM8K 0.9522; DSpark acceptance4.25-4.35 toks/fwd.
0.93, 2.93-2.99 toks/fwd. Binds 62 MLA rows = 61 target + 1 draft, the
count the deleted arithmetic gave.
float8_copy_kernel_cudagone from the trace,cp_mha_gather_cache_kernelpresent -- the 7.1% this started from.main(kv 877755 x 266240,state 256 x 39223296), 48 linear-state + 16 KV rows, again the countsthe deleted formulas gave.
maincontrol of0.89, one run each. Not explained yet. The suspect is this branch's
change of the registry key from
getattr(module, "layer_num", layer_id)to
module.layer_num; the check is a per-module key diff, not moresamples. Flagging rather than sitting on it.
including the
index_cache_layer_mapfix above -- is covered by unittests only.