K3: paged latent MLA KV + absorbed decode; EP slot default 64 - #869
Merged
Conversation
The 16-slot default enforced the retired stepwise chain's ep_size x max_batch <= masked_cap constraint. The fused transport's protocol maximum is 384 rows per rank; the binding constraint today is the KDA state slab (~929 MB/slot: f32 recurrent x2 parity + conv windows across 69 layers), which puts 64 slots (~58 GiB) next to the 224-expert rank's weights with room left for the paged MLA pool. The compiled 128-row bucket ceiling becomes reachable once the backbone goes FP8. Signed-off-by: xiaguan <751080330@qq.com>
Add the paged latent cache the K3 MLA layers will decode from: one page-first slab per state pool ([page][mla_layer][64 tokens][576] bf16 — the post-norm kv latent and the shared rope half, the exact quantity kv_b expands K and V from; NoPE, so the row is position-independent), a per-slot block table with a host mirror, and a plain free-list page allocator (claim on 64-boundary crossing, release-all on slot reset, whole-page copy on prefill adoption). Pool size comes from the new K3ExecutorConfig::kv_pages (0 derives full coverage: every slot can reach max_ctx). The decode step now appends every MLA layer's new-token latent into the current page slice (two indexed writes through the per-layer row shift) alongside the existing expanded K/V append. This is a transitional dual write: attention still reads the expanded cache, so behavior is unchanged and every gate stays green; M2 lands the absorbed-MLA kernel over the pages and deletes the expanded cache and its kernel family. The shared indexed row write widens its element index to size_t — the K3 paged slab is the first destination past 2^31 elements. Gates: cargo test -p pegainfer-k3 --lib (34 passed) and cargo check --tests green; GPU gates unaffected by construction (attention path untouched). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: xiaguan <751080330@qq.com>
Decode now attends over the paged latent cache M1 introduced, and the expanded slot-indexed MLA cache is gone. The step lands the 576-wide [kv_lora 512 | rope 64] row into the mapped page and runs one absorbed kernel per MLA layer: score(t) = (W_UK[h]^T q_nope)^T c_t + q_rope^T k_rope_t, o[h] = W_UV[h] (sum p_t c_t), with W_UK/W_UV read straight out of the checkpoint's w_kv_b — no expansion GEMM, no per-head K/V rows, 27.6 KB/token instead of 1.47 MB/token, and no compile-time context cap (max_ctx is a runtime number now: default 4096, PEGAINFER_K3_MAX_CTX override, ceiling K3_MAX_CONTEXT). The kernel is hand CUDA (csrc/k3/k3_mla_paged_attn.cu), not TileLang: the design goals are a runtime-length block-table walk and no O(ctx) storage (3-sweep recomputed softmax), and TileLang 0.1.12 expresses neither without a static capacity — which is the thing being removed. The certified rounding chain is preserved landing for landing and documented step-by-step in the kernel header: q absorption lands bf16; f32 score dot over 576 ascending; dot lands bf16 and multiplies the bf16 scale in bf16; f32 fixed-order max and sum(exp); probabilities land bf16 after normalizing; f32 latent accumulation over ascending t landing bf16; f32 W_UV expansion landing bf16. Physical page ids never enter the arithmetic and a -1 table entry reads as zero latent, so page permutation is bitwise invisible and padding rows behave exactly like the old zeroed cache. The mla_attn_batched TileLang family is deleted in lockstep: vendored factory, generator plan + MAX_CTX list, FFI decl, ops wrapper + K3_MAX_CTX, build.rs stub launcher, KERNELS.md and kernels/README rows. Eleven batched families remain, none of them attention. Gates (224-expert checkpoint, PEGAINFER_K3_LAYERS via fixture, GB300): - Absorbed-vs-expanded certification (tests/paged_kv.rs logit dump on the M1 revision vs this one, same forced feed): argmax identical on all 40 steps; deviation over the fixture-published top-5 logits median 1.0 / max 3.0 bf16 ULP — inside the measured <=2-ULP structural noise floor; whole-row max deviation at the top-logit magnitude median 1.75 / max 4.0 ULP. Bit equality is not expected: the absorbed kernel associates the score dot as (W_UK^T q)^T c rather than q^T (W_UK c). - golden_decode: all 8 gates green; 39/40 exact under the fused kernel (miss only at step 11, a 0.00-ULP coin flip), 38/40 under the masked chain (steps 25 and 33, the documented coin flips). - Paged-vs-permuted (tests/paged_kv.rs): 200 steps across 4 pages, scrambled free list claims disjoint pages in opposite order — every logit of every step bit-identical. - Long context (tests/paged_kv.rs): 1100 steps at max_ctx 2048, finite logits throughout; eager rerun and graphs-vs-eager both bit-identical (tokens every step, logits every 128th and the last). - ep_mega_oracle (4 GPUs): EP4 rank 0 vs single-rank mega bitwise (all 40 tokens, all 163840 final logits), peer-traffic invariance bitwise. - cargo test -p pegainfer-k3 --lib: 34 passed. fmt + clippy clean. Signed-off-by: xiaguan <751080330@qq.com>
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: fc551f894c
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
Comment on lines
+337
to
+338
| /// The paged MLA latent cache all MLA layers share. | ||
| pub(crate) kv: K3PagedKv, |
There was a problem hiding this comment.
When K3 serves requests, this newly owned pool allocates and releases pages, but K3Scheduler::load in pegainfer-k3/src/scheduler/mod.rs still hard-codes kv_used_blocks: 0, while model_line.rs supplies kv_capacity: None. Consequently the load snapshot and Prometheus/load-aware consumers always see zero K3 KV capacity and usage even as this pool fills; expose the pool's total and allocated page counts through the scheduler.
Useful? React with 👍 / 👎.
Pure code motion: the paged latent cache (struct, page allocator, block tables, append/adopt paths) and its two geometry constants move from executor/buffers.rs to executor/paged_kv.rs; copy_rows stays in buffers as pub(super). No behavior change. Signed-off-by: xiaguan <751080330@qq.com>
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
Replaces K3's fully-expanded, statically-sized MLA decode cache with a paged latent KV and an absorbed-MLA decode kernel, and raises the EP slot default now that the chain-era cap is gone.
Three commits:
EP slot default 16 → 64 — the old ceiling came from the retired NCCL-chain constraint (
ep × batch ≤ masked_cap 128). The mega protocol admits up to 384 rows/rank; the practical bound is now the KDA state slab (~929 MB/slot), which fits 64 slots next to the weights on GB300. 128 becomes reachable once the FP8 backbone lands.Paged MLA latent KV pool + write path (M1) — decode stores the 576-wide latent rows (KV-LoRA 512 + RoPE 64, bf16) in 64-token pages,
[page][layer][token][576]page-first slab, per-slot block tables (host mirror + device i32) backed by a free-list allocator.K3ExecutorConfiggainskv_pages(0 = full coverage derived frommax_ctx), andmax_ctxbecomes a runtime knob (default 4096,PEGAINFER_K3_MAX_CTX, ceiling 1M) instead of a compile-time constant. 27.6 KB/token across the 24 MLA layers, vs ~1.47 MB/token for the expanded cache this replaces.Absorbed-MLA decode over the paged KV (M2) — hand CUDA kernel (
k3_mla_paged_attn.cu): per step,q_abs = W_UKᵀ · q_nopeis precomputed and attention runs as MQA directly over the latent rows via a 3-sweep online softmax with no O(ctx) scratch and a runtime page-walk (no compile-time ctx instantiation — TileLang can't express the runtime-length page loop, hence hand CUDA). The expanded cache,K3MlaState, the kv_b expansion, and themla_attn_batchedTileLang family are deleted in lockstep.Verification
pegainfer-k3/tests/paged_kv.rs.Follow-ups (out of scope here)
Chunked prefill (removes the token loop and the second state pool), kv-store BlockPool integration for content addressing/reuse, EP4 CUDA graphs, and a fusion pass over the 3-sweep kernel.
🤖 Generated with Claude Code