diff --git a/docs/models/k3/bring-up.md b/docs/models/k3/bring-up.md index b80528f7..d903d211 100644 --- a/docs/models/k3/bring-up.md +++ b/docs/models/k3/bring-up.md @@ -12,10 +12,15 @@ and bitwise-invariant to peer traffic. Single-rank executor: buckets up to `B = 128`, per-bucket CUDA graphs default-on, token-matching a certified 4-layer greedy golden (39/40 exact under the fused kernel, 38/40 under the masked-chain anchor, misses inside a structural ≤2-ULP noise floor — see -below). Kernel surface: twelve batched TileLang decode families + DeepGEMM -FP8xFP4 AOT shims (fused MegaMoE and the masked grouped GEMM) behind -`pegainfer-kernels`'s `k3` feature; dense projections on cuBLASLt. Next: CUDA -graphs over the EP4 fused path. +below). MLA KV is a **paged latent cache** — 576-wide `[kv_lora|rope]` rows in +64-token pages behind per-slot block tables, 27.6 KB/token — decoded by an +absorbed-MLA CUDA kernel with no compile-time context cap (the old 1.47 +MB/token expanded slot cache and its `max_ctx = 128` are gone). Kernel surface: +eleven batched TileLang decode families + the hand-written paged-attention +kernel + DeepGEMM FP8xFP4 AOT shims (fused MegaMoE and the masked grouped GEMM) +behind `pegainfer-kernels`'s `k3` feature; dense projections on cuBLASLt. Next: +CUDA graphs over the EP4 fused path, real (chunked) prefill, kv-store +integration. Last touched: 2026-08 @@ -51,20 +56,23 @@ model at EP16 — that is the development vehicle. 1e-6 against a dequant reference on real weights), so the loader uploads payload and e8m0 scale bytes as-is; the SF relayout to DeepGEMM's packed i32 form is a device-side build step (`k3_fp4_sf_prepare`). -- **KV story is dual-pool**: paged KV (kv-store `BlockPool`) for the 24 MLA - layers plus a qwen35-style fixed-size slot pool for KDA recurrent state. +- **KV story is dual-pool**: paged KV for the 24 MLA layers (executor-owned + latent page pool today; kv-store `BlockPool` integration is a later + milestone) plus a qwen35-style fixed-size slot pool for KDA recurrent state. Prefix caching ships disabled: KDA state is not recomputable from tokens (`docs/subsystems/kv-cache/design.md`, bounded class). - **TileLang kernels are generated at build time** from vendored, certified kernel definitions in `pegainfer-k3/kernels/` (three tiers: live generation → pre-generated dir → NOT_SUPPORTED stubs). The definitions' spelling is certified against the HF reference in a separate harness and - must not drift. The set covers every non-GEMM step of a decode iteration — - norms, the bf16 landings of the framework GEMMs, conv+silu, the KDA delta - rule, MLA attention, router, situ, combine and the attention-residual mix — - so the executor composes launches rather than reimplementing spellings. - Dense projections run on cuBLASLt and the routed experts on the DeepGEMM - masked grouped-GEMM chain, so no GEMV is generated. Every shape is a static + must not drift. The set covers every non-GEMM, non-attention step of a + decode iteration — norms, the bf16 landings of the framework GEMMs, + conv+silu, the KDA delta rule, router, situ, combine and the + attention-residual mix — so the executor composes launches rather than + reimplementing spellings. Dense projections run on cuBLASLt, the routed + experts on the DeepGEMM masked grouped-GEMM chain, and MLA decode on the + hand-written paged kernel below, so neither a GEMV nor an attention family + is generated. Every shape is a static compile dimension, batch size included: `B = 1` is a bucket whose per-row spelling is the certified single-row kernel, gated bitwise upstream, so single-stream and high-concurrency decode share one kernel set. See @@ -75,22 +83,24 @@ model at EP16 — that is the development vehicle. `pegainfer-k3/src/executor/` composes the certified kernels into the decode step: `step.rs` is a line-by-line port of the certified reference engine's launch sequence (dense projections on cuBLASLt with banded/offset landings, -the 12 batched TileLang families, the fused MoE launch), `buffers.rs` owns the -state pools, `mod.rs` owns graphs and the `StepExecutor` impl. +the 11 batched TileLang families, the paged MLA attention, the fused MoE +launch), `buffers.rs` owns the state pools and the paged KV pool, `mod.rs` +owns graphs and the `StepExecutor` impl. - **Batching**: seat `i` is row `i` of every state pool. Buckets `{1,2,4,8,16,32,48,64,96,128}`; one CUDA graph per `(bucket, parity)` (KDA recurrent state ping-pongs across two slabs, so parity is part of the graph identity). Graphs are default-on; `PEGAINFER_K3_CUDA_GRAPH=0` escapes - to eager. H2D feed (`token_ids`, `context_len`, `cache_row`) and the single - argmax D2H stay outside capture. + to eager. H2D feed (`token_ids`, `context_len`, `kv_row`, the KV block + table) and the single argmax D2H stay outside capture. - **State contract**: a padding row is stepped like any live row, so its recurrent state advances — a seat's state is only meaningful while the seat is in *every* batch. The scheduler preserves this (running requests decode every step; `prefill` resets the seat at admission). Prefill runs on a separate one-row pool and hands its state over by row copy. - **Bring-up flags**: `PEGAINFER_K3_LAYERS` (layer truncation), - `PEGAINFER_K3_MAX_BATCH`, `PEGAINFER_K3_CUDA_GRAPH`. + `PEGAINFER_K3_MAX_BATCH`, `PEGAINFER_K3_CUDA_GRAPH`, + `PEGAINFER_K3_MAX_CTX` (per-slot context ceiling, default 4096). ### Gates and the noise floor @@ -114,6 +124,79 @@ dominate — known to lose to a GEMV below ~8 rows/expert). Graphs buy ~4.5%; the step is not launch-bound. Absolute numbers move a few percent with box load, so compare within a session, not across. +### MLA KV: paged latent cache + absorbed decode + +The MLA cache holds the **latent**, not the expanded heads: one 576-wide bf16 +row per token per MLA layer — the post-norm kv latent (512, `kv_lora_rank`) +next to the raw shared rope half (64; K3 is NoPE, nothing is ever rotated). +That is 27.6 KB/token against the 1.47 MB/token the expanded `[96×192 K | +96×128 V]` slot cache used to pin, and it is what lifted the fixed +`max_ctx = 128`. + +- **Layout**: one slab per rank, `[page][layer][token][576]`, 64 tokens per + page, every MLA layer's slice inside the same page (layer offset = + `mla_index × 64 × 576` elements). Pages come from a plain free list — + claimed when a slot's position crosses a 64 boundary, freed together when + the slot retires; no content addressing, no reuse (kv-store integration is + a later milestone). Per-slot block tables live host-side and ride to the + device with the step feed, outside graph capture; the captured kernels read + the device table by pointer. `K3ExecutorConfig::kv_pages` sizes the pool + (`0` = full coverage, `max_batch × ceil(max_ctx/64)`). +- **Write path**: the step computes `kv_norm` and `rope` exactly as before + and lands them into the mapped page row (`kv_row` is the device-fed + destination index; `-1` rows are skipped). A page is zeroed when claimed + and every position written once, so the indexed add is an exact indexed + copy. +- **Decode is absorbed MLA** (`csrc/k3/k3_mla_paged_attn.cu`, hand CUDA — + TileLang 0.1.12 cannot express the runtime-length page walk without a + compile-time capacity, which is the thing being removed): per head, + `score(t) = (W_UKᵀ q_nope)ᵀ c_t + q_ropeᵀ k_rope_t` and + `o = W_UV (Σ p_t c_t)`, with `W_UK`/`W_UV` read straight out of the + checkpoint's `w_kv_b` — no expansion GEMM, no expanded cache. One + (row, head) block; softmax is a 3-sweep recompute (max / sum / + probs+attend), so there is no O(ctx) storage and **no compile-time context + cap**. The certified rounding chain is preserved landing for landing and + documented step-by-step in the kernel header: q absorption lands bf16 (f32 + dot per latent column); f32 score dot over the 576 columns ascending; the + dot lands bf16 and multiplies the bf16 scale in bf16; max and Σexp in f32 + fixed order; probabilities land bf16 after normalizing; the latent + accumulation is f32 per column over ascending t, landing bf16; the W_UV + expansion is an f32 dot landing bf16. +- **Physical page ids never enter the arithmetic** — the kernel walks the + block table by logical position — so any page permutation is bit-identical, + and an unmapped (`-1`) entry reads as zero latent, which is exactly the + zeroed padding row of the old cache. + +Gates (`pegainfer-k3/tests/paged_kv.rs`, plus the golden suite): + +```bash +# Paged gates: page-permutation bitwise, long-context (2048-ctx, 1100 steps, +# rerun + graphs-vs-eager bitwise), on a GPU box with the checkpoint: +PEGAINFER_K3_TEST_224= cargo test --release -p pegainfer-k3 \ + --test paged_kv -- --ignored --test-threads 1 + +# Absorbed-vs-expanded certification: dump per-step logits on the expanded +# revision (M1) and on this one, then compare per step in bf16 ULP: +PEGAINFER_K3_TEST_224= PEGAINFER_K3_LOGIT_DUMP=/tmp/logits.bin \ + cargo test --release -p pegainfer-k3 --test paged_kv \ + dump_forced_replay_logits -- --ignored + +# And the existing suites must stay green: +PEGAINFER_K3_TEST_224= cargo test --release -p pegainfer-k3 \ + --test golden_decode -- --ignored --test-threads 1 +``` + +Certified (4-layer golden inputs, absorbed vs expanded per step): 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 of the FP8 expert path; whole-row max deviation at the top-logit +magnitude median 1.75 / max 4.0 ULP. (The absorbed kernel associates the +score dot as `(W_UKᵀ q)ᵀ c` rather than `qᵀ (W_UK c)`, so bit equality is not +expected — the floor is.) The golden fixture holds at 39/40 (fused) and 38/40 +(masked chain) with misses only on the documented coin-flip steps; page +permutation (200 steps) and the long-context reruns and graphs-vs-eager +(1100 steps, ctx 2048) are bit-identical; the EP4 oracle is bitwise green. + ## EP4: free-running ranks The full model does not fit one GPU (routed experts alone are ~354 GB of @@ -285,5 +368,7 @@ recorded here as the measurement, not as a configuration you can still select. ## Next Graphs over the EP4 fused path (the ranks=1 path already captures), then -launch-ahead, paged KV for MLA (current MLA cache is fixed `max_ctx = 128` per -seat) and real prefill. +launch-ahead, real (chunked) prefill over the paged KV, kv-store `BlockPool` +integration (content addressing / reuse for the MLA pages), and a perf pass on +the paged attention kernel (the 3-sweep recompute reads the cache three times; +fine at bring-up depth, worth a fused pass at 24 MLA layers x long contexts). diff --git a/pegainfer-k3/kernels/README.md b/pegainfer-k3/kernels/README.md index f22d78b5..84e28539 100644 --- a/pegainfer-k3/kernels/README.md +++ b/pegainfer-k3/kernels/README.md @@ -2,9 +2,12 @@ **TL;DR**: `generate.py` AOT-compiles the vendored TileLang kernel definitions in `tilelang_defs.py` into CUDA that `pegainfer-kernels/build.rs` hands to -nvcc under the `k3` feature. Thirteen batched families cover every non-GEMM -step of a K3 decode iteration — three tiers (generate, pre-generated, stub), -and the generated CUDA is a Cargo `OUT_DIR` artifact that is never checked in. +nvcc under the `k3` feature. Eleven batched families cover every non-GEMM, +non-attention step of a K3 decode iteration — three tiers (generate, +pre-generated, stub), and the generated CUDA is a Cargo `OUT_DIR` artifact +that is never checked in. MLA decode is a hand-written absorbed paged-KV CUDA +kernel (`pegainfer-kernels/csrc/k3/k3_mla_paged_attn.cu`), not a TileLang +family. ## What lives here @@ -35,36 +38,36 @@ every row below is its shape count × 10. | `situ_batched` | N ∈ {6144, 33792} | 20 | `k3_situ_batched` | | `conv_silu_batched` | KP = 12288, W = 4, SK = 1 | 10 | `k3_conv_silu_batched` | | `kda_core_batched` | 96 heads × 128 head_dim | 10 | `k3_kda_core_batched` | -| `mla_attn_batched` | 96 heads, qk 192, v 128, CAP ∈ {128} | 10 | `k3_mla_attn_batched` | | `router_topk_batched` | E ∈ {896, 224}, TOPK = 16 | 20 | `k3_router_topk_batched` | | `attnres_scores_batched` | NB ∈ 1..8, H = 7168 | 80 | `k3_attnres_scores_batched` | | `attnres_mix_batched` | NB ∈ 1..8, H = 7168 | 80 | `k3_attnres_mix_batched` | -**430 instantiations**, about 20 seconds of generation and 7 seconds of nvcc. +**420 instantiations**, about 20 seconds of generation and 7 seconds of nvcc. The pool fans out at *instantiation* granularity, not family granularity — the families differ by more than an order of magnitude in size, so a family-granular pool would be bound by `land_batched` alone. It defaults to one worker per CPU capped at 32; each worker holds a TileLang lowering, so lower it with `PEGAINFER_K3_TILELANG_JOBS` on memory-tight hosts. -Two lists in `generate.py` are deliberately narrow and are one-line changes: +One list in `generate.py` is deliberately narrow and is a one-line change: -* `MAX_CTX` — the MLA context-capacity list. Both the instantiation count and - the per-block shared memory scale with the capacity, so serving longer - contexts is a deliberate widening, not a free one. * `SPLIT_K` — the segment counts the partial consumers (`land`, `land_rms_norm_rbs`, `conv_silu`) accept. Only `1` — the single partial a framework GEMM produces — has a launch site; the reference engine's split-K-8 GEMV shapes are not generated. -Neither needs a launcher edit. +It needs no launcher edit. ### What is deliberately *not* here -No GEMV family. Dense projections run on cuBLASLt and the routed experts on -the DeepGEMM masked grouped-GEMM chain (`csrc/k3/k3_moe_chain.cu` and the -`k3.deepgemm.*` ops), so the upstream `gemv`, `expert_gemv` and -`packed_expert_gemv` kernels would be dead weight. Their consumers are still +No GEMV family, and no attention family. Dense projections run on cuBLASLt, +the routed experts on the DeepGEMM masked grouped-GEMM chain +(`csrc/k3/k3_moe_chain.cu` and the `k3.deepgemm.*` ops), and MLA decode on the +absorbed paged-KV kernel (`csrc/k3/k3_mla_paged_attn.cu` — a runtime page walk +needs no per-capacity instantiation, which is what retired the upstream +`mla_attn` family and its `MAX_CTX` list), so the upstream `gemv`, +`expert_gemv`, `packed_expert_gemv` and `mla_attn` kernels would be dead +weight. Their consumers are still here — `land_batched` at `SK = 1` is exactly the bf16 landing of a cuBLASLt f32 output. @@ -115,18 +118,18 @@ trim them. The same parse also asserts that no body lowers to TMA. A warp-specialized TileLang kernel takes `CUtensorMap` descriptors instead of pointers and adds a producer warpgroup to the block, so the launchers — which bind plain pointers -and the requested thread count — would be silently wrong. None of the batched -bodies do this today (their only bulk copy is `mla_attn`'s one-dimensional -query load), but that is a property of TileLang, not of the kernels, so it is -asserted rather than assumed. +and the requested thread count — would be silently wrong. None of the batched bodies use a bulk +copy at all today, but that is a property of TileLang, not of the kernels, so +it is asserted rather than assumed. ## TileLang codegen is not byte-reproducible Re-running the same instantiation can swap the names of two *aliases of the -same* dynamic-shared-memory offset — `mla_attn` emits `void* workspace` and -`void* workspace_1`, both `buf_dyn_shmem + 0`, and which one each `AllReduce` -scratch argument gets flips run to run (measured 8/10 vs 2/10 on an otherwise -identical invocation). Nothing else moves: same offsets, same instructions. +same* dynamic-shared-memory offset — the retired `mla_attn` family emitted +`void* workspace` and `void* workspace_1`, both `buf_dyn_shmem + 0`, and which +one each `AllReduce` scratch argument got flipped run to run (measured 8/10 vs +2/10 on an otherwise identical invocation). Nothing else moves: same offsets, +same instructions. The practical consequence is only for tooling: a byte-diff of a generated body against a fresh upstream `get_kernel_source()` dump has to canonicalize @@ -175,6 +178,7 @@ Allocating at `K3_MAX_BATCH` once and reusing keeps the pointers stable for CUDA Graph capture. Per-slot state is `[b, ...]` contiguous with each row holding exactly the -single-row layout: the `conv_silu` windows, the `kda_core` recurrent state and -the MLA `Kc`/`Vc` caches. That contract is what the upstream bitwise gate -proves, and it is why the caches are slot-indexed rather than paged. +single-row layout: the `conv_silu` windows and the `kda_core` recurrent state. +That contract is what the upstream bitwise gate proves. (The MLA KV cache is +*not* per-slot state anymore — it is the paged latent pool owned by the +executor and consumed by `csrc/k3/k3_mla_paged_attn.cu`.) diff --git a/pegainfer-k3/kernels/generate.py b/pegainfer-k3/kernels/generate.py index 309bbbe5..d3f9e7da 100644 --- a/pegainfer-k3/kernels/generate.py +++ b/pegainfer-k3/kernels/generate.py @@ -6,17 +6,19 @@ k3_rms_norm_rbs_batched.cu k3_conv_silu_batched.cu k3_land_batched.cu k3_kda_core_batched.cu - k3_land_rms_norm_rbs_batched.cu k3_mla_attn_batched.cu - k3_add2_batched.cu k3_router_topk_batched.cu - k3_mul_sigmoid_batched.cu k3_attnres_scores_batched.cu - k3_situ_batched.cu k3_attnres_mix_batched.cu + k3_land_rms_norm_rbs_batched.cu k3_router_topk_batched.cu + k3_add2_batched.cu k3_attnres_scores_batched.cu + k3_mul_sigmoid_batched.cu k3_attnres_mix_batched.cu + k3_situ_batched.cu The batch size is a static compile-time dimension, so a single-stream step is served by the `B = 1` instantiation of the same family — its per-row spelling is word for word the certified single-row kernel, which is what the upstream bitwise gate proves. There is therefore no separate bs=1 kernel set. The dense -projections are served by cuBLASLt and the routed experts by the DeepGEMM -masked grouped-GEMM chain, so no GEMV family is generated here either. +projections are served by cuBLASLt, the routed experts by the DeepGEMM +masked grouped-GEMM chain, and MLA decode by the hand-written absorbed +paged-attention kernel (`csrc/k3/k3_mla_paged_attn.cu`), so neither a GEMV +family nor an attention family is generated here. Each file is [shared TileLang preamble] + [one renamed `main_kernel` per instantiation] + [one hand-written `extern "C"` dispatch launcher]. TileLang @@ -111,13 +113,6 @@ # shapes have no launch site here and are not generated. SPLIT_K = [1] -# `mla_attn`'s compile-time cache capacity; `engine.cap`, which the engine -# asserts is a multiple of 128. Serving will want larger capacities, but both -# the instantiation count and the per-block shared memory scale with it, so -# only the certification capacity is generated by default. Extending serving -# to longer contexts is this one-line change. -MAX_CTX = [128] - # `engine._wsm_n`: the beta / low-rank-gate projection is padded up to a # multiple of 64 so its tile divides evenly. WSM_N = (KDA_HEADS + KDA_HEAD_DIM + 63) // 64 * 64 @@ -130,8 +125,8 @@ ATTNRES_BLOCKS = (NUM_LAYERS + ATTNRES_BLOCK_SIZE - 1) // ATTNRES_BLOCK_SIZE ATTNRES_NB = list(range(1, ATTNRES_BLOCKS + 1)) -# Default thread count of the vendored kernels; `kda_core` and `mla_attn` -# override it (head_dim and 128 respectively). +# Default thread count of the vendored kernels; `kda_core` overrides it +# (head_dim). THREADS = 256 # Batch buckets. A step with `rows` live rows runs the next bucket up and @@ -172,9 +167,9 @@ # a round-before-scale norm are fused. LAND_RMS_NORM_CONFIGS = [(MLA_FUSED, Q_LORA, 0)] -# add2 / mul_sigmoid / situ / conv_silu / kda_core / mla_attn take a single -# width each; situ has two (shared, dense) — the routed-expert situ is fused -# into the masked-GEMM chain and the mega kernel. +# add2 / mul_sigmoid / situ / conv_silu / kda_core take a single width each; +# situ has two (shared, dense) — the routed-expert situ is fused into the +# masked-GEMM chain and the mega kernel. ADD2_N = [HIDDEN] MUL_SIGMOID_N = [KDA_DIM] SITU_N = [SHARED_INTER, DENSE_INTER] @@ -219,11 +214,6 @@ f"const {BF16}* __restrict__ Q, const float* __restrict__ State, " f"float* __restrict__ StateN, const {BF16}* __restrict__ V)" ) -MLA_ATTN_PARAMS = ( - f"(const {BF16}* __restrict__ Kc, const int* __restrict__ N, " - f"{BF16}* __restrict__ O, const {BF16}* __restrict__ Q, " - f"const {BF16}* __restrict__ Sc, const {BF16}* __restrict__ Vc)" -) ROUTER_PARAMS = ( "(const float* __restrict__ Bias, int* __restrict__ Idx, " f"const {BF16}* __restrict__ Rs, const float* __restrict__ S, " @@ -250,10 +240,10 @@ # TileLang lowers a TMA copy into a warp-specialized kernel that takes # `CUtensorMap` descriptors instead of pointers and adds a producer warpgroup -# to the block. None of the batched bodies do — their only bulk copy is -# `mla_attn`'s one-dimensional query load — so the launchers bind plain -# pointers and the requested thread count. That is an assumption about -# TileLang, not about the kernels, so it is asserted rather than trusted. +# to the block. None of the batched bodies use a bulk copy at all, so the +# launchers bind plain pointers and the requested thread count. That is an +# assumption about TileLang, not about the kernels, so it is asserted rather +# than trusted. TENSORMAP_BUILDER = "__tvm_tensormap_create_tiled" @@ -872,61 +862,6 @@ def plan_kda_core() -> Plan: ) -def plan_mla_attn() -> Plan: - insts = [] - for cap in MAX_CTX: - for batch in B_BUCKETS: - insts.append(Inst( - family="mla_attn", - order=len(insts), - label=f"mla_attn_batched NH={MLA_HEADS} CAP={cap} B={batch}", - factory="mla_attn_batched", - args=(MLA_HEADS, QK_DIM, V_DIM, cap, batch, 128), - num_params=6, - params=MLA_ATTN_PARAMS, - symbol=f"k3_mla_attn_b{batch}_nh{MLA_HEADS}_cap{cap}_kernel", - grid=(batch, MLA_HEADS), - threads=128, - guard=( - f"b == {batch} && num_heads == {MLA_HEADS} && " - f"qk_dim == {QK_DIM} && v_dim == {V_DIM} && max_ctx == {cap}" - ), - call_args=( - _bf16("Kc"), "N", _bf16("O", False), - _bf16("Q"), _bf16("Sc"), _bf16("Vc"), - ), - )) - return Plan( - stem=_STEM.format("mla_attn"), - signature=( - "k3_mla_attn_batched(\n" - " const void* Q,\n" - " const void* Kc,\n" - " const void* Vc,\n" - " const int* N,\n" - " const void* Sc,\n" - " void* O,\n" - " int b,\n" - " int num_heads,\n" - " int qk_dim,\n" - " int v_dim,\n" - " int max_ctx,\n" - " cudaStream_t stream)" - ), - doc=( - "// NoPE full-context attention, one (row, head) per block; every head\n" - "// holds its own K, so there is no GQA sharing. The cache is slot\n" - "// indexed: each row owns a fixed max_ctx window and its cache is the\n" - "// contiguous single-row block -- this is not a paged cache and there is\n" - "// no block table. N is a per-row device i32 length; slots at or past it\n" - "// score NEG, so a row at length n is bit-for-bit the single-row kernel\n" - "// at that n and the step needs no host sync. Sc is the shared bf16\n" - "// softmax scale." - ), - insts=tuple(insts), - ) - - def plan_router_topk() -> Plan: insts = [] for experts in EXPERTS: @@ -1060,7 +995,6 @@ def plan_attnres_mix() -> Plan: plan_situ, plan_conv_silu, plan_kda_core, - plan_mla_attn, plan_router_topk, plan_attnres_scores, plan_attnres_mix, diff --git a/pegainfer-k3/kernels/tilelang_defs.py b/pegainfer-k3/kernels/tilelang_defs.py index 1a11aa37..aeacde5e 100644 --- a/pegainfer-k3/kernels/tilelang_defs.py +++ b/pegainfer-k3/kernels/tilelang_defs.py @@ -1,7 +1,7 @@ """Vendored TileLang kernel definitions for the K3 batched decode step. This file is a **verbatim** subset of the certified upstream kernel module: the -shared prologue and the thirteen batched kernel factories, copied character for +shared prologue and the eleven batched kernel factories, copied character for character. Nothing here is re-spelled, re-indented or "cleaned up", and no kernel body is edited to fit this repository. The upstream module is the authority on what these kernels compute; it carries the bitwise parity gates @@ -467,69 +467,3 @@ def main( ) * T.Cast(DT, T.sigmoid(G2[bb, bh * KD + d].astype(ACC))) return _compile(main) - - -# --- MLA ------------------------------------------------------------------- # - - -@lru_cache(maxsize=None) -def mla_attn_batched(NH: int, QK: int, VD: int, CAP: int, B: int, - threads: int = 128): - """Batched ``mla_attn``: block (b, h) attends row b's head h over row b's - own cache. The cache is **slot indexed** -- every sequence owns a fixed - CAP-slot window, so Kc/Vc simply gain a leading batch axis and a row's - cache is the contiguous bs=1 [CAP, NH*QK] / [CAP, NH*VD] block (this is not - a paged cache; there is no block table). - - The context length is **per slot**: N is an i32 tensor of shape (B,) and - the mask predicate becomes ``s < N[bb]``, which is the only difference from - the bs=1 kernel's ``s < N[0]``. Out-of-range slots still score NEG, so a - row at length n is bit-for-bit the bs=1 kernel run at that n, whatever the - other rows' lengths are. Sc (the bf16 softmax scale) is a config constant - shared by every row. Landings (f32 dot -> bf16 -> times the bf16 scale -> - bf16 -> f32 softmax -> bf16 probabilities -> f32 V accumulation -> one - bf16 landing) are the bs=1 body verbatim.""" - @T.prim_func - def main( - Q: T.Tensor((B, NH * QK), DT), - Kc: T.Tensor((B, CAP, NH * QK), DT), - Vc: T.Tensor((B, CAP, NH * VD), DT), - N: T.Tensor((B,), "int32"), - Sc: T.Tensor((1,), DT), - O: T.Tensor((B, NH * VD), DT), - ): - with T.Kernel(B, NH, threads=threads) as (bb, bh): - Qs = T.alloc_shared((QK,), DT) - probs = T.alloc_shared((CAP,), DT) - dot = T.alloc_fragment((CAP,), ACC) - scl = T.alloc_fragment((CAP,), ACC) - pr = T.alloc_fragment((CAP,), ACC) - oac = T.alloc_fragment((VD,), ACC) - mx = T.alloc_fragment((1,), ACC) - tot = T.alloc_fragment((1,), ACC) - T.copy(Q[bb, bh * QK:(bh + 1) * QK], Qs) - T.sync_threads() - T.clear(dot) - for s in T.Parallel(CAP): - for d in T.serial(QK): - dot[s] += Qs[d].astype(ACC) * Kc[bb, s, bh * QK + d].astype(ACC) - for s in T.Parallel(CAP): - # Land bf16, times bf16 scale (product lands bf16), widen f32. - scl[s] = T.if_then_else( - s < N[bb], (T.Cast(DT, dot[s]) * Sc[0]).astype(ACC), NEG - ) - T.reduce_max(scl, mx, dim=0) - for s in T.Parallel(CAP): - pr[s] = T.exp(scl[s] - mx[0]) - T.reduce_sum(pr, tot, dim=0) - for s in T.Parallel(CAP): - probs[s] = T.Cast(DT, pr[s] / tot[0]) # bf16 after normalizing - T.sync_threads() - T.clear(oac) - for dv in T.Parallel(VD): - for s in T.serial(CAP): - oac[dv] += probs[s].astype(ACC) * Vc[bb, s, bh * VD + dv].astype(ACC) - for dv in T.Parallel(VD): - O[bb, bh * VD + dv] = T.Cast(DT, oac[dv]) - - return _compile(main) diff --git a/pegainfer-k3/src/executor/buffers.rs b/pegainfer-k3/src/executor/buffers.rs index f2b4423f..34d538d4 100644 --- a/pegainfer-k3/src/executor/buffers.rs +++ b/pegainfer-k3/src/executor/buffers.rs @@ -24,7 +24,6 @@ use pegainfer_kernels::ops::K3_KDA_HEAD_DIM; use pegainfer_kernels::ops::K3_KDA_HEADS; use pegainfer_kernels::ops::K3_MLA_HEADS; use pegainfer_kernels::ops::K3_MOE_QUANT_GROUP; -use pegainfer_kernels::ops::K3_QK_DIM; use pegainfer_kernels::ops::K3_ROUTER_TOPK; use pegainfer_kernels::ops::K3_V_DIM; use pegainfer_kernels::ops::K3MegaSymmLayout; @@ -33,15 +32,14 @@ use pegainfer_kernels::ops::k3_mega_open_peer_access; use pegainfer_kernels::ops::k3_mega_symm_buffer_layout; use pegainfer_kernels::ops::k3_mega_token_alignment; use pegainfer_kernels::tensor::DeviceContext; -use pegainfer_kernels::tensor::HiddenStates; +use super::paged_kv::K3PagedKv; use crate::config::K3_ATTN_INNER; use crate::config::K3_DENSE_INTERMEDIATE; use crate::config::K3_EXPERT_INTERMEDIATE; use crate::config::K3_HEAD_DIM; use crate::config::K3_HIDDEN; use crate::config::K3_KV_A_OUT; -use crate::config::K3_KV_B_OUT; use crate::config::K3_KV_LORA_RANK; use crate::config::K3_Q_B_OUT; use crate::config::K3_Q_LORA_RANK; @@ -65,8 +63,7 @@ pub(crate) const K3_MLA_FUSED: usize = pub(crate) const K3_CONV_STATE: usize = K3_CONV_WIDTH - 1; /// Elements of one row's KDA recurrent state. pub(crate) const K3_KDA_STATE: usize = K3_KDA_HEADS * K3_KDA_HEAD_DIM * K3_KDA_HEAD_DIM; -/// Per-slot MLA cache row widths. -pub(crate) const K3_MLA_K_ROW: usize = K3_MLA_HEADS * K3_QK_DIM; +/// Width of the MLA attention output, `heads * v_head_dim`. pub(crate) const K3_MLA_V_ROW: usize = K3_MLA_HEADS * K3_V_DIM; /// One KDA layer's per-slot state: recurrent matrix plus the three convolution @@ -78,17 +75,11 @@ pub(crate) struct K3KdaState { pub(crate) conv: [[CudaSlice; 3]; 2], } -/// One MLA layer's per-slot slot-indexed cache. Each slot owns a fixed -/// `max_ctx` window, so the buffers are the batched kernel's `[rows, cap, w]` -/// and, seen as `[rows * cap, w]`, the indexed row write's destination. -pub(crate) struct K3MlaState { - pub(crate) k_cache: HiddenStates, - pub(crate) v_cache: HiddenStates, -} - pub(crate) enum K3LayerState { Kda(Box), - Mla(Box), + /// MLA state lives in the pool-wide paged latent cache ([`K3PagedKv`]), + /// not per layer. + Mla, } /// Everything about a slot that outlives a step. @@ -96,6 +87,8 @@ pub(crate) struct K3StatePool { pub(crate) rows: usize, pub(crate) max_ctx: usize, pub(crate) layers: Vec, + /// The paged MLA latent cache all MLA layers share. + pub(crate) kv: K3PagedKv, /// Attention-residual snapshot history, `[rows, blocks, hidden]` bf16. pub(crate) blocks: CudaSlice, pub(crate) block_count: usize, @@ -111,8 +104,14 @@ impl K3StatePool { max_ctx: usize, num_layers: usize, block_count: usize, + kv_pages: usize, ) -> Result { let stream = &ctx.stream; + let mla_layers = (0..num_layers) + .filter(|layer| k3_layer_kind(*layer) == K3LayerKind::Mla) + .count() + .max(1); + let kv = K3PagedKv::new(ctx, rows, max_ctx, mla_layers, kv_pages)?; let mut layers = Vec::with_capacity(num_layers); for layer in 0..num_layers { layers.push(match k3_layer_kind(layer) { @@ -144,28 +143,14 @@ impl K3StatePool { conv: [conv_even, conv_odd], })) } - K3LayerKind::Mla => K3LayerState::Mla(Box::new(K3MlaState { - k_cache: HiddenStates { - data: stream - .alloc_zeros::(rows * max_ctx * K3_MLA_K_ROW) - .context("alloc K3 MLA key cache")?, - hidden_dim: K3_MLA_K_ROW, - seq_len: rows * max_ctx, - }, - v_cache: HiddenStates { - data: stream - .alloc_zeros::(rows * max_ctx * K3_MLA_V_ROW) - .context("alloc K3 MLA value cache")?, - hidden_dim: K3_MLA_V_ROW, - seq_len: rows * max_ctx, - }, - })), + K3LayerKind::Mla => K3LayerState::Mla, }); } Ok(Self { rows, max_ctx, layers, + kv, blocks: stream .alloc_zeros::(rows * block_count * K3_HIDDEN) .context("alloc K3 attention-residual snapshots")?, @@ -189,25 +174,13 @@ impl K3StatePool { } } } - K3LayerState::Mla(mla) => { - zero_rows( - ctx, - &mut mla.k_cache.data, - row * self.max_ctx, - self.max_ctx, - K3_MLA_K_ROW, - )?; - zero_rows( - ctx, - &mut mla.v_cache.data, - row * self.max_ctx, - self.max_ctx, - K3_MLA_V_ROW, - )?; - } + // The paged latent cache is released below; freed pages are + // zeroed when next claimed, not here. + K3LayerState::Mla => {} } } zero_rows(ctx, &mut self.blocks, row, 1, self.block_count * K3_HIDDEN)?; + self.kv.release_row(row); self.positions[row] = 0; Ok(()) } @@ -261,26 +234,9 @@ impl K3StatePool { )?; } } - (K3LayerState::Mla(target), K3LayerState::Mla(origin)) => { - copy_rows( - ctx, - &origin.k_cache.data, - source_row * self.max_ctx, - &mut target.k_cache.data, - row * self.max_ctx, - self.max_ctx, - K3_MLA_K_ROW, - )?; - copy_rows( - ctx, - &origin.v_cache.data, - source_row * self.max_ctx, - &mut target.v_cache.data, - row * self.max_ctx, - self.max_ctx, - K3_MLA_V_ROW, - )?; - } + // The paged latent cache is adopted once for the whole pool, + // below the layer walk. + (K3LayerState::Mla, K3LayerState::Mla) => {} _ => anyhow::bail!("K3 state pools disagree on layer kinds"), } } @@ -293,6 +249,13 @@ impl K3StatePool { 1, self.block_count * K3_HIDDEN, )?; + self.kv.adopt_row( + ctx, + &source.kv, + source_row, + row, + source.positions[source_row], + )?; self.positions[row] = source.positions[source_row]; Ok(()) } @@ -315,7 +278,7 @@ fn zero_rows( .context("zero a K3 state row") } -fn copy_rows( +pub(super) fn copy_rows( ctx: &DeviceContext, source: &CudaSlice, source_row: usize, @@ -579,12 +542,9 @@ pub(crate) struct K3Scratch { pub(crate) token_ids: CudaSlice, /// Per-row MLA context length, i.e. valid cache slots including this step. pub(crate) context_len: CudaSlice, - /// Per-row destination of this step's cache write, `row * cap + position`, - /// or `-1` for a row this step does not own. - pub(crate) cache_row: CudaSlice, - /// `head_row[r * heads + h] = r`, the broadcast of one row's shared rope - /// half to every MLA head. Static. - pub(crate) head_row: CudaSlice, + /// Per-row destination of this step's paged latent write + /// ([`K3PagedKv::write_index`]), or `-1` for a row this step does not own. + pub(crate) kv_row: CudaSlice, // Residual stream. pub(crate) hidden: CudaSlice, pub(crate) prefix: CudaSlice, @@ -617,13 +577,8 @@ pub(crate) struct K3Scratch { pub(crate) mla_gate: CudaSlice, pub(crate) q_partial: CudaSlice, pub(crate) query: CudaSlice, - pub(crate) kv_partial: CudaSlice, - pub(crate) kv: CudaSlice, - pub(crate) k_nope: CudaSlice, - pub(crate) k_new: HiddenStates, - pub(crate) v_new: HiddenStates, - pub(crate) rope: HiddenStates, - pub(crate) rope_heads: HiddenStates, + /// The shared per-token rope half, `[rows, 64]` — cached verbatim (NoPE). + pub(crate) rope: CudaSlice, pub(crate) attn: CudaSlice, // MLP / MoE. pub(crate) hidden_partial: CudaSlice, @@ -673,18 +628,13 @@ impl K3Scratch { mega: Option, ) -> Result { let stream = &ctx.stream; - let heads = K3_MLA_HEADS; - let head_row: Vec = (0..rows * heads) - .map(|entry| (entry / heads) as i32) - .collect(); let wide = |width: usize| stream.alloc_zeros::(rows * width); let partial = |width: usize| stream.alloc_zeros::(rows * width); let argmax_partials = argmax_batch_bf16_split_partials_len(rows, K3_VOCAB); Ok(Self { token_ids: stream.alloc_zeros(rows)?, context_len: stream.alloc_zeros(rows)?, - cache_row: stream.alloc_zeros(rows)?, - head_row: stream.clone_htod(&head_row)?, + kv_row: stream.clone_htod(&vec![-1i32; rows])?, hidden: wide(K3_HIDDEN)?, prefix: wide(K3_HIDDEN)?, mixed: wide(K3_HIDDEN)?, @@ -714,29 +664,7 @@ impl K3Scratch { mla_gate: wide(K3_ATTN_INNER)?, q_partial: partial(K3_Q_B_OUT)?, query: wide(K3_Q_B_OUT)?, - kv_partial: partial(K3_KV_B_OUT)?, - kv: wide(K3_KV_B_OUT)?, - k_nope: wide(heads * K3_HEAD_DIM)?, - k_new: HiddenStates { - data: wide(K3_MLA_K_ROW)?, - hidden_dim: K3_MLA_K_ROW, - seq_len: rows, - }, - v_new: HiddenStates { - data: wide(K3_MLA_V_ROW)?, - hidden_dim: K3_MLA_V_ROW, - seq_len: rows, - }, - rope: HiddenStates { - data: wide(K3_QK_ROPE_HEAD_DIM)?, - hidden_dim: K3_QK_ROPE_HEAD_DIM, - seq_len: rows, - }, - rope_heads: HiddenStates { - data: wide(heads * K3_QK_ROPE_HEAD_DIM)?, - hidden_dim: K3_QK_ROPE_HEAD_DIM, - seq_len: rows * heads, - }, + rope: wide(K3_QK_ROPE_HEAD_DIM)?, attn: wide(K3_MLA_V_ROW)?, hidden_partial: partial(K3_HIDDEN)?, router_partial: partial(routed_experts)?, diff --git a/pegainfer-k3/src/executor/mod.rs b/pegainfer-k3/src/executor/mod.rs index e2751a2c..342015aa 100644 --- a/pegainfer-k3/src/executor/mod.rs +++ b/pegainfer-k3/src/executor/mod.rs @@ -56,6 +56,7 @@ mod buffers; pub mod ep; mod gemm; +mod paged_kv; mod step; use std::path::Path; @@ -71,7 +72,6 @@ use pegainfer_frontend::sampler::SamplingParams; use pegainfer_kernels::ops::K3_BATCH_BUCKETS; use pegainfer_kernels::ops::K3_DEEPGEMM_SM100_GROUPS; use pegainfer_kernels::ops::K3_MAX_BATCH; -use pegainfer_kernels::ops::K3_MAX_CTX; use pegainfer_kernels::ops::k3_batch_bucket; use pegainfer_kernels::tensor::DeviceContext; @@ -111,6 +111,12 @@ const K3_CUDA_GRAPH_ENV: &str = "PEGAINFER_K3_CUDA_GRAPH"; /// Concurrent slots per rank; rounded up to a compiled bucket, capped at the /// widest one. const K3_MAX_BATCH_ENV: &str = "PEGAINFER_K3_MAX_BATCH"; +/// Context ceiling per slot (tokens); the paged pool is sized from it when +/// `kv_pages` is not set explicitly. +const K3_MAX_CTX_ENV: &str = "PEGAINFER_K3_MAX_CTX"; +/// Default per-slot context ceiling. Free to raise: the cost is pool pages +/// (27.6 KB per token across the 24 MLA layers), not compiled kernels. +const K3_DEFAULT_MAX_CTX: usize = 4096; /// The only SM count the fused MegaMoE kernel is AOT-instantiated for. Its /// grid sync spans the whole grid, so the launch geometry is baked in. const K3_MEGA_SMS: usize = 152; @@ -120,9 +126,14 @@ const K3_MEGA_SMS: usize = 152; const K3_MEGA_EP_SIZES: [usize; 2] = [1, 4]; /// Slots per rank an expert-parallel launch takes when nothing says otherwise. -/// Well inside the fused kernel's 384-row protocol maximum, and the value every -/// EP gate and serve run has been measured at. -const K3_EP_DEFAULT_MAX_BATCH: usize = 16; +/// +/// The fused kernel's protocol maximum is 384 rows per rank, so the compiled +/// bucket ceiling (128) is the target once the backbone goes FP8. Today the +/// binding constraint is the KDA state slab: ~929 MB per slot (f32 recurrent +/// x2 parity + conv windows across 69 layers), so 64 slots cost ~58 GiB — +/// what fits next to the 224-expert rank's weights with room left for the +/// paged MLA pool. An explicit `PEGAINFER_K3_MAX_BATCH` still wins. +const K3_EP_DEFAULT_MAX_BATCH: usize = 64; /// What a launch decides about an executor before its weights are read. #[derive(Clone, Copy, Debug)] @@ -130,8 +141,15 @@ pub struct K3ExecutorConfig { /// Concurrent slots, i.e. the row capacity of every state slab. Rounded up /// to a compiled bucket. pub max_batch: usize, - /// Cache slots per sequence. Must be a compiled MLA capacity. + /// Context ceiling per slot, in tokens. A runtime number: the paged + /// attention kernel walks block tables, so nothing is compiled per + /// capacity. pub max_ctx: usize, + /// Pages in the MLA latent KV pool (64 tokens per page, all MLA layers' + /// slices inside one page). `0` derives full coverage — every slot can + /// reach `max_ctx` — so allocation can only fail when this is set lower + /// (oversubscription is the caller's explicit choice). + pub kv_pages: usize, /// Layers to build; `K3_LAYERS` for the whole model. pub num_layers: usize, /// Capture and replay the step, rather than launching it eagerly. @@ -188,7 +206,8 @@ impl Default for K3ExecutorConfig { fn default() -> Self { Self { max_batch: K3_MAX_BATCH, - max_ctx: K3_MAX_CTX[0], + max_ctx: K3_DEFAULT_MAX_CTX, + kv_pages: 0, num_layers: K3_LAYERS, cuda_graph: true, moe_transport: K3MoeTransport::MEGA, @@ -215,6 +234,12 @@ impl K3ExecutorConfig { { self.max_batch = slots; } + if let Ok(raw) = std::env::var(K3_MAX_CTX_ENV) + && let Ok(tokens) = raw.parse::() + && (1..=crate::config::K3_MAX_CONTEXT).contains(&tokens) + { + self.max_ctx = tokens; + } self } @@ -261,7 +286,7 @@ pub struct K3Executor { /// Step inputs, staged on the host and copied in before every step. token_host: Vec, context_len_host: Vec, - cache_row_host: Vec, + kv_row_host: Vec, sampled_host: Vec, thread_bound: bool, /// Present exactly when `ep_size > 1`: this rank's slab handshake with its @@ -368,9 +393,10 @@ impl K3Executor { ) -> Result { let max_batch = k3_batch_bucket(config.max_batch)?; ensure!( - K3_MAX_CTX.contains(&config.max_ctx), - "K3 max_ctx {} is not a compiled MLA capacity {K3_MAX_CTX:?}", - config.max_ctx + (1..=crate::config::K3_MAX_CONTEXT).contains(&config.max_ctx), + "K3 max_ctx {} is outside 1..={}", + config.max_ctx, + crate::config::K3_MAX_CONTEXT ); let num_sms = ctx .ctx @@ -401,8 +427,23 @@ impl K3Executor { cuda_graph = false; } - let decode_state = K3StatePool::new(&ctx, max_batch, config.max_ctx, num_layers, blocks)?; - let prefill_state = K3StatePool::new(&ctx, 1, config.max_ctx, num_layers, blocks)?; + let slot_pages = config.max_ctx.div_ceil(paged_kv::K3_KV_PAGE_TOKENS); + let kv_pages = if config.kv_pages == 0 { + max_batch * slot_pages + } else { + config.kv_pages + }; + let decode_state = K3StatePool::new( + &ctx, + max_batch, + config.max_ctx, + num_layers, + blocks, + kv_pages, + )?; + // The prefill pool is one row, so full coverage is one slot's pages. + let prefill_state = + K3StatePool::new(&ctx, 1, config.max_ctx, num_layers, blocks, slot_pages)?; if mega { ensure!( K3_MEGA_EP_SIZES.contains(&ep_size), @@ -496,7 +537,7 @@ impl K3Executor { prefill_graphs: (0..2).map(|_| CudaGraphState::new()).collect(), token_host: vec![0; max_batch], context_len_host: vec![1; max_batch], - cache_row_host: vec![-1; max_batch], + kv_row_host: vec![-1; max_batch], sampled_host: vec![0; max_batch], thread_bound: false, ep, @@ -550,7 +591,6 @@ impl K3Executor { bucket, live_rows, parity, - max_ctx: self.max_ctx, groups: self.groups, masked_cap: K3_MASKED_CAP, num_sms: self.num_sms, @@ -568,8 +608,8 @@ impl K3Executor { .memcpy_htod(&self.context_len_host, &mut self.scratch.context_len) .map_err(|error| anyhow::anyhow!("K3 context-length feed failed: {error}"))?; stream - .memcpy_htod(&self.cache_row_host, &mut self.scratch.cache_row) - .map_err(|error| anyhow::anyhow!("K3 cache-row feed failed: {error}")) + .memcpy_htod(&self.kv_row_host, &mut self.scratch.kv_row) + .map_err(|error| anyhow::anyhow!("K3 KV-row feed failed: {error}")) } /// Run one step against `pool`, through its graph when graphs are on. @@ -593,6 +633,9 @@ impl K3Executor { let ctx = &self.ctx; let model = &self.model; let scratch = &mut self.scratch; + // The block table rides outside capture with the rest of the step + // inputs; the captured kernels read the device table by pointer. + pool.kv.sync_table(ctx)?; if !self.cuda_graph { let launches = self.mega_launches_per_step; if let Some(mega) = scratch.mega.as_mut() { @@ -681,6 +724,15 @@ impl K3Executor { Ok(sampled) } + /// Test hook: reverse the decode pool's free page list, so the next + /// sequence's pages land at different physical ids in a different order. + /// The paged cache's core gate (`tests/paged_kv.rs`) is that no page + /// permutation can move a single logit bit. + #[doc(hidden)] + pub fn scramble_kv_pages(&mut self) { + self.decode_state.kv.reverse_free_list(); + } + /// Bring-up diagnostics: the logit row the most recent step left for /// `row`, widened to f32. Costs a device round trip; not a serving path. pub fn last_logits(&mut self, row: usize) -> Result> { @@ -701,11 +753,14 @@ impl K3Executor { fn prefill_token(&mut self, token: u32, position: usize, parity: usize) -> Result { self.token_host[0] = token; self.context_len_host[0] = i32::try_from(position + 1)?; - self.cache_row_host[0] = i32::try_from(position)?; + self.prefill_state + .kv + .ensure_mapped(&self.ctx, 0, position)?; + self.kv_row_host[0] = self.prefill_state.kv.write_index(0, position)?; for row in 1..self.max_batch { self.token_host[row] = 0; self.context_len_host[row] = 1; - self.cache_row_host[row] = -1; + self.kv_row_host[row] = -1; } self.feed()?; self.run_step(true, 1, parity, 1)?; @@ -825,7 +880,7 @@ impl K3Executor { for row in 0..self.max_batch { self.token_host[row] = 0; self.context_len_host[row] = 1; - self.cache_row_host[row] = -1; + self.kv_row_host[row] = -1; } for entry in batch { let position = self.decode_state.positions[entry.slot]; @@ -837,7 +892,11 @@ impl K3Executor { ); self.token_host[entry.slot] = entry.last_token; self.context_len_host[entry.slot] = i32::try_from(position + 1)?; - self.cache_row_host[entry.slot] = i32::try_from(entry.slot * self.max_ctx + position)?; + self.decode_state + .kv + .ensure_mapped(&self.ctx, entry.slot, position)?; + self.kv_row_host[entry.slot] = + self.decode_state.kv.write_index(entry.slot, position)?; } self.feed()?; diff --git a/pegainfer-k3/src/executor/paged_kv.rs b/pegainfer-k3/src/executor/paged_kv.rs new file mode 100644 index 00000000..77d167ca --- /dev/null +++ b/pegainfer-k3/src/executor/paged_kv.rs @@ -0,0 +1,264 @@ +//! The paged MLA latent KV cache — the pool-wide slab, its block tables, and +//! the free-list page allocator. Everything else about a slot's state stays in +//! [`super::buffers`]. + +use anyhow::Context; +use anyhow::Result; +use anyhow::ensure; +use cudarc::driver::CudaSlice; +use cudarc::driver::DevicePtr; +use cudarc::driver::DevicePtrMut; +use half::bf16; +use pegainfer_kernels::tensor::DeviceContext; + +use super::buffers::copy_rows; +use crate::config::K3_KV_LORA_RANK; +use crate::config::K3_QK_ROPE_HEAD_DIM; + +/// Tokens per MLA KV page. +pub(crate) const K3_KV_PAGE_TOKENS: usize = 64; +/// One token's cached MLA latent: the post-norm kv latent and the shared +/// per-token rope half, exactly what `kv_b` expands K and V from. NoPE — no +/// rotary is ever applied, so the row is position-independent. +pub(crate) const K3_MLA_LATENT_ROW: usize = K3_KV_LORA_RANK + K3_QK_ROPE_HEAD_DIM; + +/// The paged MLA latent cache: one slab per pool, shared by every MLA layer. +/// +/// Layout is page-first, `[page][mla_layer][token][K3_MLA_LATENT_ROW]` bf16 — +/// all of a pool's MLA layers keep their 64-token slices at fixed offsets +/// inside one page, so a page is the transfer/interchange unit and one +/// allocation serves the whole model. Seen as rows of `K3_MLA_LATENT_ROW`, +/// layer `l`'s row for `(page, token)` is `page * mla_layers * 64 + l * 64 + +/// token`; the per-layer term is a constant row shift, which is how the write +/// path and the attention kernel address one layer's slice out of the shared +/// slab. +/// +/// Pages come from a plain free-list pool: allocated when a slot's position +/// crosses a 64 boundary, freed together when the slot resets. No content +/// addressing and no reuse — this is the local pool only. +pub(crate) struct K3PagedKv { + /// `[num_pages, mla_layers, K3_KV_PAGE_TOKENS, K3_MLA_LATENT_ROW]` bf16. + pub(crate) slab: CudaSlice, + pub(crate) mla_layers: usize, + pub(crate) num_pages: usize, + pub(crate) max_pages_per_slot: usize, + /// Device block table, `[rows, max_pages_per_slot]` i32, `-1` unmapped. + pub(crate) table_dev: CudaSlice, + /// Host mirror; the executor mutates this and re-uploads before a step. + table_host: Vec, + /// Free page ids; `pop` hands them out in ascending order from fresh. + free: Vec, +} + +impl K3PagedKv { + pub(crate) fn new( + ctx: &DeviceContext, + rows: usize, + max_ctx: usize, + mla_layers: usize, + num_pages: usize, + ) -> Result { + let max_pages_per_slot = max_ctx.div_ceil(K3_KV_PAGE_TOKENS); + anyhow::ensure!( + num_pages >= max_pages_per_slot, + "K3 KV pool of {num_pages} pages cannot hold even one {max_ctx}-token slot" + ); + let slab_len = num_pages * mla_layers * K3_KV_PAGE_TOKENS * K3_MLA_LATENT_ROW; + Ok(Self { + slab: ctx + .stream + .alloc_zeros::(slab_len) + .context("alloc K3 paged MLA latent slab")?, + mla_layers, + num_pages, + max_pages_per_slot, + table_dev: ctx + .stream + .clone_htod(&vec![-1i32; rows * max_pages_per_slot]) + .context("alloc K3 KV block table")?, + table_host: vec![-1; rows * max_pages_per_slot], + free: (0..num_pages as i32).rev().collect(), + }) + } + + /// Rows (of `K3_MLA_LATENT_ROW`) one page spans across its layer slices. + fn page_rows(&self) -> usize { + self.mla_layers * K3_KV_PAGE_TOKENS + } + + /// Elements from one page to the next — the attention kernel's page walk + /// stride. + pub(crate) fn page_stride(&self) -> usize { + self.page_rows() * K3_MLA_LATENT_ROW + } + + /// Make sure the page holding `position` is mapped for `row`, zeroing a + /// freshly claimed page (the step's cache write is an indexed *add*, so + /// the destination must be zero until the write). + pub(crate) fn ensure_mapped( + &mut self, + ctx: &DeviceContext, + row: usize, + position: usize, + ) -> Result<()> { + let page_slot = position / K3_KV_PAGE_TOKENS; + anyhow::ensure!( + page_slot < self.max_pages_per_slot, + "K3 KV row {row} position {position} exceeds its {} pages", + self.max_pages_per_slot + ); + let entry = row * self.max_pages_per_slot + page_slot; + if self.table_host[entry] >= 0 { + return Ok(()); + } + let page = self + .free + .pop() + .context("K3 KV pool is out of pages; raise kv_pages")?; + self.table_host[entry] = page; + let start = page as usize * self.page_rows() * K3_MLA_LATENT_ROW; + let mut window = self + .slab + .slice_mut(start..start + self.page_rows() * K3_MLA_LATENT_ROW); + ctx.stream + .memset_zeros(&mut window) + .context("zero a fresh K3 KV page") + } + + /// Row index (in the `K3_MLA_LATENT_ROW`-wide view, before the per-layer + /// shift) of `row`'s cache write at `position`. The page must be mapped. + pub(crate) fn write_index(&self, row: usize, position: usize) -> Result { + let page = self.table_host[row * self.max_pages_per_slot + position / K3_KV_PAGE_TOKENS]; + anyhow::ensure!( + page >= 0, + "K3 KV row {row} position {position} has no mapped page" + ); + Ok(page * self.page_rows() as i32 + (position % K3_KV_PAGE_TOKENS) as i32) + } + + /// Return every page `row` holds to the pool. Contents are not cleared — + /// a page is zeroed when it is next claimed. + pub(crate) fn release_row(&mut self, row: usize) { + let base = row * self.max_pages_per_slot; + for entry in &mut self.table_host[base..base + self.max_pages_per_slot] { + if *entry >= 0 { + self.free.push(*entry); + *entry = -1; + } + } + } + + /// Test hook: reverse the free list, so the next claims come from the + /// opposite end of the pool. Physical page ids never enter the attention + /// arithmetic — the kernel walks the block table by logical position — so + /// any permutation must leave every logit bit-identical, and + /// `tests/paged_kv.rs` holds the executor to exactly that. + #[doc(hidden)] + pub(crate) fn reverse_free_list(&mut self) { + self.free.reverse(); + } + + /// Map fresh pages for `row` covering `tokens` positions and copy the + /// source row's pages into them wholesale. The caller releases first. + pub(crate) fn adopt_row( + &mut self, + ctx: &DeviceContext, + source: &K3PagedKv, + source_row: usize, + row: usize, + tokens: usize, + ) -> Result<()> { + anyhow::ensure!( + source.mla_layers == self.mla_layers + && source.max_pages_per_slot == self.max_pages_per_slot, + "K3 KV pools disagree on geometry" + ); + let width = self.page_rows() * K3_MLA_LATENT_ROW; + for page_slot in 0..tokens.div_ceil(K3_KV_PAGE_TOKENS) { + self.ensure_mapped(ctx, row, page_slot * K3_KV_PAGE_TOKENS)?; + let origin = source.table_host[source_row * source.max_pages_per_slot + page_slot]; + let target = self.table_host[row * self.max_pages_per_slot + page_slot]; + anyhow::ensure!( + origin >= 0, + "K3 KV adoption found no source page for slot {page_slot}" + ); + copy_rows( + ctx, + &source.slab, + origin as usize, + &mut self.slab, + target as usize, + 1, + width, + )?; + } + Ok(()) + } + + /// Upload the host block table. Cheap enough to do before every step, and + /// outside graph capture like the rest of the step inputs. + pub(crate) fn sync_table(&mut self, ctx: &DeviceContext) -> Result<()> { + ctx.stream + .memcpy_htod(&self.table_host, &mut self.table_dev) + .map_err(|error| anyhow::anyhow!("K3 KV block-table feed failed: {error}")) + } + + /// Write one step's latent rows into layer `mla_index`'s page slices. + /// + /// `kv_row` is the per-batch-row destination from [`Self::write_index`] + /// (`-1` for rows the step does not own — the indexed write skips them). + /// The destination is zero until this step writes it — a page is zeroed + /// when claimed and every position is written once — so the indexed add is + /// an exact indexed copy, taking the destination row from the device. + pub(crate) fn append_latent( + &mut self, + ctx: &DeviceContext, + mla_index: usize, + rows: usize, + kv_row: &CudaSlice, + kv_norm: &CudaSlice, + rope: &CudaSlice, + ) -> Result<()> { + ensure!( + mla_index < self.mla_layers, + "K3 KV layer index out of range" + ); + ensure!( + kv_norm.len() >= rows * K3_KV_LORA_RANK + && rope.len() >= rows * K3_QK_ROPE_HEAD_DIM + && kv_row.len() >= rows, + "K3 KV append buffers too small for {rows} rows" + ); + // Layer `mla_index`'s slice of every page starts `mla_index * 64` rows + // into the page, so a base pointer shifted by that many rows makes the + // layer-independent `write_index` address this layer's slice. + let shift_rows = mla_index * K3_KV_PAGE_TOKENS; + let out_rows = self.num_pages * self.page_rows() - shift_rows; + let (slab_ptr, _slab_guard) = self.slab.device_ptr_mut(&ctx.stream); + let base = slab_ptr + (shift_rows * K3_MLA_LATENT_ROW * size_of::()) as u64; + let (kv_row_ptr, _kv_row_guard) = kv_row.device_ptr(&ctx.stream); + for (delta, width, column) in [ + (kv_norm, K3_KV_LORA_RANK, 0usize), + (rope, K3_QK_ROPE_HEAD_DIM, K3_KV_LORA_RANK), + ] { + let (delta_ptr, _delta_guard) = delta.device_ptr(&ctx.stream); + unsafe { + pegainfer_kernels::ffi::scaled_add_rows_indexed_cuda( + delta_ptr as *const pegainfer_kernels::ffi::Half, + 1.0, + kv_row_ptr as *const i32, + base as *mut pegainfer_kernels::ffi::Half, + K3_MLA_LATENT_ROW as i32, + column as i32, + width as i32, + rows as i32, + out_rows as i32, + pegainfer_kernels::tensor::active_cu_stream(ctx), + ) + } + .result() + .map_err(|error| anyhow::anyhow!("K3 KV latent append failed: {error}"))?; + } + Ok(()) + } +} diff --git a/pegainfer-k3/src/executor/step.rs b/pegainfer-k3/src/executor/step.rs index 17237fda..0a68548f 100644 --- a/pegainfer-k3/src/executor/step.rs +++ b/pegainfer-k3/src/executor/step.rs @@ -31,9 +31,7 @@ use half::bf16; use pegainfer_kernels::ops::K3_CONV_WIDTH; use pegainfer_kernels::ops::K3_MLA_HEADS; use pegainfer_kernels::ops::K3_MOE_QUANT_GROUP; -use pegainfer_kernels::ops::K3_QK_DIM; use pegainfer_kernels::ops::K3_ROUTER_TOPK; -use pegainfer_kernels::ops::K3_V_DIM; use pegainfer_kernels::ops::K3DeepGemmFp8Fp4Kind; use pegainfer_kernels::ops::K3MegaActivation; use pegainfer_kernels::ops::K3MegaShape; @@ -42,7 +40,6 @@ use pegainfer_kernels::ops::argmax_bf16_split_into; use pegainfer_kernels::ops::copy_hidden_rows_raw_into; use pegainfer_kernels::ops::embedding_rows_into; use pegainfer_kernels::ops::extract_hidden_rows_raw_into; -use pegainfer_kernels::ops::gather_hidden_tokens_into; use pegainfer_kernels::ops::k3_add2_batched_launch; use pegainfer_kernels::ops::k3_attnres_mix_batched_launch; use pegainfer_kernels::ops::k3_attnres_scores_batched_launch; @@ -54,7 +51,7 @@ use pegainfer_kernels::ops::k3_land_batched_launch; use pegainfer_kernels::ops::k3_land_rms_norm_rbs_batched_launch; use pegainfer_kernels::ops::k3_mega_moe_launch; use pegainfer_kernels::ops::k3_mega_write_inputs_launch; -use pegainfer_kernels::ops::k3_mla_attn_batched_launch; +use pegainfer_kernels::ops::k3_mla_paged_attn_launch; use pegainfer_kernels::ops::k3_moe_gather_fp8_quant_masked_launch; use pegainfer_kernels::ops::k3_moe_local_route_metadata_launch; use pegainfer_kernels::ops::k3_moe_weighted_combine_launch; @@ -63,7 +60,6 @@ use pegainfer_kernels::ops::k3_rms_norm_rbs_batched_launch; use pegainfer_kernels::ops::k3_router_topk_batched_launch; use pegainfer_kernels::ops::k3_situ_and_mul_fp8_quant_masked_launch; use pegainfer_kernels::ops::k3_situ_batched_launch; -use pegainfer_kernels::ops::scaled_add_rows_indexed_into; use pegainfer_kernels::tensor::DeviceContext; use pegainfer_kernels::tensor::DeviceMatrix; @@ -77,6 +73,9 @@ use super::buffers::parity_pair; use super::gemm::K3PartialSpan; use super::gemm::k3_gemm_full; use super::gemm::k3_gemm_partial; +use super::paged_kv::K3_KV_PAGE_TOKENS; +use super::paged_kv::K3_MLA_LATENT_ROW; +use super::paged_kv::K3PagedKv; use crate::config::K3_ATTN_INNER; use crate::config::K3_DENSE_INTERMEDIATE; use crate::config::K3_EXPERT_INTERMEDIATE; @@ -84,11 +83,9 @@ use crate::config::K3_HEAD_DIM; use crate::config::K3_HEADS; use crate::config::K3_HIDDEN; use crate::config::K3_KV_A_OUT; -use crate::config::K3_KV_B_OUT; use crate::config::K3_KV_LORA_RANK; use crate::config::K3_Q_B_OUT; use crate::config::K3_Q_LORA_RANK; -use crate::config::K3_QK_NOPE_HEAD_DIM; use crate::config::K3_QK_ROPE_HEAD_DIM; use crate::config::K3_ROUTED_EXPERT_HIDDEN; use crate::config::K3_SHARED_INTERMEDIATE; @@ -113,8 +110,6 @@ pub(crate) struct K3StepShape { pub(crate) live_rows: usize, /// Which half of each ping-pong state slab this step reads. pub(crate) parity: usize, - /// Slots per sequence in the MLA cache. - pub(crate) max_ctx: usize, /// Rank-local expert groups (the masked GEMM's instantiation). pub(crate) groups: usize, /// Rows reserved per expert in the masked layout. @@ -140,7 +135,7 @@ impl K3StepShape { /// Advance every row of the bucket by one token and leave the sampled ids in /// `scratch.argmax_indices`. /// -/// Reads `scratch.token_ids`, `scratch.context_len` and `scratch.cache_row`; +/// Reads `scratch.token_ids`, `scratch.context_len` and `scratch.kv_row`; /// the caller fills those before the step (or before the graph replay). /// /// Every MoE layer issues the same launches in the same order on every rank — @@ -156,11 +151,14 @@ pub(crate) fn k3_decode_step( let b = shape.bucket; let K3StatePool { layers: layer_state, + kv, blocks: snapshots, block_count, .. } = state; let block_count = *block_count; + // Index of the current MLA layer within the paged pool's layer slices. + let mut mla_index = 0usize; embedding_rows_into( ctx, @@ -217,8 +215,9 @@ pub(crate) fn k3_decode_step( scratch, )?; } - (K3LayerAttention::Mla(mla), K3LayerState::Mla(mla_state)) => { - mla_attention(ctx, b, shape, layer, mla, mla_state, scratch)?; + (K3LayerAttention::Mla(mla), K3LayerState::Mla) => { + mla_attention(ctx, b, layer, mla, kv, mla_index, scratch)?; + mla_index += 1; } _ => anyhow::bail!("K3 layer state does not match the layer's attention kind"), } @@ -591,13 +590,14 @@ fn kda_conv_stream( // ── MLA ───────────────────────────────────────────────────────────────── +#[allow(clippy::too_many_arguments)] fn mla_attention( ctx: &DeviceContext, b: usize, - shape: K3StepShape, layer: &K3LayerWeights, w: &K3MlaWeights, - state: &mut super::buffers::K3MlaState, + kv: &mut K3PagedKv, + mla_index: usize, s: &mut K3Scratch, ) -> Result<()> { k3_rms_norm_rbs_batched_launch( @@ -655,7 +655,7 @@ fn mla_attention( ctx, &s.kv_a, K3_KV_A_OUT, - &mut s.rope.data, + &mut s.rope, K3_QK_ROPE_HEAD_DIM, K3_KV_LORA_RANK, b, @@ -668,6 +668,12 @@ fn mla_attention( &w.gamma_kv_a.data, &mut s.kv_norm, )?; + // Paged latent append: the post-norm kv latent and the shared rope half + // are the whole cached quantity (NoPE — nothing here is + // position-dependent), written into this layer's slice of the row's + // current page. The expanded K/V the reference builds from it is folded + // into the absorbed attention below. + kv.append_latent(ctx, mla_index, b, &s.kv_row, &s.kv_norm, &s.rope)?; k3_gemm_full(ctx, &w.w_q_b, &s.q_norm, b, &mut s.q_partial)?; k3_land_batched_launch( ctx, @@ -679,81 +685,21 @@ fn mla_attention( &s.q_partial, &mut s.query, )?; - k3_gemm_full(ctx, &w.w_kv_b, &s.kv_norm, b, &mut s.kv_partial)?; - k3_land_batched_launch( - ctx, - b, - K3_KV_B_OUT, - K3_KV_B_OUT, - 0, - 1, - &s.kv_partial, - &mut s.kv, - )?; - - // `kv` is `[rows, heads, nope | value]`. The cache row wants - // `[heads, nope | rope]` for K and `[heads, value]` for V, with the one - // rope half shared by every head. - let per_head = K3_QK_NOPE_HEAD_DIM + K3_HEAD_DIM; - let head_rows = b * K3_MLA_HEADS; - extract_hidden_rows_raw_into( - ctx, - &s.kv, - per_head, - &mut s.k_nope, - K3_QK_NOPE_HEAD_DIM, - 0, - head_rows, - )?; - copy_hidden_rows_raw_into( - ctx, - &s.k_nope, - K3_QK_NOPE_HEAD_DIM, - &mut s.k_new.data, - K3_QK_DIM, - 0, - head_rows, - )?; - s.rope.seq_len = b; - s.rope_heads.seq_len = head_rows; - gather_hidden_tokens_into(ctx, &s.rope, &s.head_row, head_rows, &mut s.rope_heads)?; - copy_hidden_rows_raw_into( - ctx, - &s.rope_heads.data, - K3_QK_ROPE_HEAD_DIM, - &mut s.k_new.data, - K3_QK_DIM, - K3_QK_NOPE_HEAD_DIM, - head_rows, - )?; - extract_hidden_rows_raw_into( - ctx, - &s.kv, - per_head, - &mut s.v_new.data, - K3_V_DIM, - K3_QK_NOPE_HEAD_DIM, - head_rows, - )?; - - // The destination slot is zero until this step writes it — every position - // is written once and `release` clears the row — so an indexed add is an - // exact indexed copy, and it takes the destination row from the device. - s.k_new.seq_len = b; - s.v_new.seq_len = b; - scaled_add_rows_indexed_into(ctx, &s.k_new, 1.0, &s.cache_row, b, &mut state.k_cache, 0)?; - scaled_add_rows_indexed_into(ctx, &s.v_new, 1.0, &s.cache_row, b, &mut state.v_cache, 0)?; - k3_mla_attn_batched_launch( + // Absorbed MLA over the paged latent: the kernel folds `w_kv_b`'s per-head + // W_UK into the query and expands the attended latent with W_UV, so the + // per-step kv_b expansion and the expanded K/V cache no longer exist. + k3_mla_paged_attn_launch( ctx, b, K3_MLA_HEADS, - K3_QK_DIM, - K3_V_DIM, - shape.max_ctx, &s.query, - &state.k_cache.data, - &state.v_cache.data, + &w.w_kv_b.data, + &kv.slab, + mla_index * K3_KV_PAGE_TOKENS * K3_MLA_LATENT_ROW, + kv.page_stride(), + &kv.table_dev, + kv.max_pages_per_slot, &s.context_len, &w.scale.data, &mut s.attn, diff --git a/pegainfer-k3/tests/ep_mega_oracle.rs b/pegainfer-k3/tests/ep_mega_oracle.rs index 60328e07..35dd0d68 100644 --- a/pegainfer-k3/tests/ep_mega_oracle.rs +++ b/pegainfer-k3/tests/ep_mega_oracle.rs @@ -154,6 +154,7 @@ fn config(fixture: &Fixture) -> K3ExecutorConfig { K3ExecutorConfig { max_batch: MAX_BATCH, max_ctx: fixture.max_ctx, + kv_pages: 0, num_layers: fixture.num_layers, // Eager either way: EP forces capture off and the reference has to // match. (The single-rank mega path does capture; this pins it off so diff --git a/pegainfer-k3/tests/golden_decode.rs b/pegainfer-k3/tests/golden_decode.rs index 4bcaa9a7..433e921c 100644 --- a/pegainfer-k3/tests/golden_decode.rs +++ b/pegainfer-k3/tests/golden_decode.rs @@ -133,6 +133,7 @@ fn executor( let config = K3ExecutorConfig { max_batch, max_ctx: golden.max_ctx, + kv_pages: 0, num_layers: golden.num_layers, cuda_graph, moe_transport, diff --git a/pegainfer-k3/tests/paged_kv.rs b/pegainfer-k3/tests/paged_kv.rs new file mode 100644 index 00000000..10c6cfb6 --- /dev/null +++ b/pegainfer-k3/tests/paged_kv.rs @@ -0,0 +1,295 @@ +//! Paged MLA KV cache gates. +//! +//! The MLA cache is a pool of 64-token latent pages behind a per-slot block +//! table, and the absorbed decode kernel walks that table by logical position +//! — physical page ids never enter the arithmetic. That claim is checkable, +//! so these gates check it: +//! +//! 1. **Page permutation is bitwise invisible.** The same replay runs twice, +//! the second time after scrambling the free list so every page lands at a +//! different physical id in a different order. Every logit of every step +//! must be bit-identical. +//! 2. **Long context is self-consistent.** A context far past the old 128-slot +//! cache (and past one page, and past one *table row* of pages) decodes to +//! finite logits, reproduces itself bitwise on a fresh executor, and +//! produces the same trajectory through CUDA graphs. +//! 3. **A/B logit dump.** Not a gate by itself but the instrument for one: +//! `dump_forced_replay_logits` writes every step's logit row to +//! `PEGAINFER_K3_LOGIT_DUMP`, so the absorbed kernel can be held to the +//! expanded kernel it replaced by running the dump on both builds and +//! comparing the files offline (the M1 revision carries the expanded +//! attention with the paged write already in place, which makes the diff a +//! verdict on the attention kernel alone). +//! +//! Manual gates like the golden suite: CI compiles them, a Blackwell box with +//! the checkpoint runs them with `--ignored`. `PEGAINFER_K3_TEST_224` points +//! at the 224-expert checkpoint, `PEGAINFER_K3_TEST_DEVICE` picks the GPU. + +use std::io::Write; +use std::path::PathBuf; + +use pegainfer_k3::DecodeSlot; +use pegainfer_k3::K3Executor; +use pegainfer_k3::K3ExecutorConfig; +use pegainfer_k3::K3MoeTransport; +use pegainfer_k3::StepExecutor; + +const FIXTURE: &str = include_str!("fixtures/k3_4l_greedy.json"); +const CHECKPOINT_ENV: &str = "PEGAINFER_K3_TEST_224"; +const DEVICE_ENV: &str = "PEGAINFER_K3_TEST_DEVICE"; +const DUMP_ENV: &str = "PEGAINFER_K3_LOGIT_DUMP"; + +/// The fixture's prompt-then-argmax feed and layer truncation — the same +/// forced-replay inputs the golden gate uses, so the A/B dump below compares +/// the two attention kernels on certified ground. +struct Fixture { + feed: Vec, + num_layers: usize, +} + +fn fixture() -> Fixture { + let json: serde_json::Value = serde_json::from_str(FIXTURE).expect("fixture parses"); + let prompt: Vec = json["prompt"] + .as_array() + .expect("prompt array") + .iter() + .map(|entry| entry.as_u64().expect("token id") as u32) + .collect(); + let argmax: Vec = json["steps"] + .as_array() + .expect("steps array") + .iter() + .map(|step| step["argmax"].as_u64().expect("argmax") as u32) + .collect(); + let mut feed = prompt.clone(); + feed.extend(argmax[prompt.len() - 1..argmax.len() - 1].iter().copied()); + Fixture { + feed, + num_layers: json["num_layers"].as_u64().expect("num_layers") as usize, + } +} + +/// The checkpoint directory, or `None` when this box does not have it mounted. +fn checkpoint() -> Option { + let path = PathBuf::from(std::env::var(CHECKPOINT_ENV).ok()?); + path.join("config.json").exists().then_some(path) +} + +fn device() -> usize { + std::env::var(DEVICE_ENV) + .ok() + .and_then(|raw| raw.parse().ok()) + .unwrap_or(0) +} + +fn executor(config: K3ExecutorConfig) -> Option { + let path = checkpoint()?; + Some( + K3Executor::load(&path, device(), 0, 1, config) + .expect("the truncated rank model should load"), + ) +} + +/// Diverse but deterministic filler tokens, clear of the fixture's ids. +fn filler_tokens(count: usize, seed: u64) -> Vec { + let mut state = seed; + (0..count) + .map(|_| { + state = state + .wrapping_mul(6_364_136_223_846_793_005) + .wrapping_add(1_442_695_040_888_963_407); + ((state >> 33) % 100_000) as u32 + 1_000 + }) + .collect() +} + +/// Forced replay of `feed` on slot 0, one decode step per token, returning +/// each step's sampled token and its logit row as raw f32 bit patterns. +/// The logits land in bf16, so bit equality here *is* bf16 bit equality. +fn replay_with_logits(executor: &mut K3Executor, feed: &[u32]) -> (Vec, Vec>) { + executor.release(0); + let mut sampled = Vec::with_capacity(feed.len()); + let mut logits = Vec::with_capacity(feed.len()); + for &last_token in feed { + let step = executor + .decode(&[DecodeSlot { + slot: 0, + last_token, + }]) + .expect("the decode step should run"); + sampled.push(step[0]); + logits.push( + executor + .last_logits(0) + .expect("logit readback") + .into_iter() + .map(f32::to_bits) + .collect(), + ); + } + (sampled, logits) +} + +fn assert_bitwise(baseline: &[Vec], other: &[Vec], what: &str) { + assert_eq!(baseline.len(), other.len(), "{what}: step counts differ"); + for (step, (a, b)) in baseline.iter().zip(other).enumerate() { + assert_eq!(a.len(), b.len(), "{what}: step {step} logit widths differ"); + if let Some(id) = (0..a.len()).find(|&id| a[id] != b[id]) { + panic!( + "{what}: step {step} logit {id} differs — {:e} vs {:e} \ + ({:#010x} vs {:#010x})", + f32::from_bits(a[id]), + f32::from_bits(b[id]), + a[id], + b[id] + ); + } + } + eprintln!("{what}: {} steps bit-identical", baseline.len()); +} + +/// Gate: scrambling the page pool must not move a single logit bit. The feed +/// spans four pages, so the walk crosses page boundaries in both runs; the +/// scrambled run claims disjoint physical pages in the opposite order. +#[test] +#[ignore = "requires a Blackwell GPU and the K3 checkpoint"] +fn page_permutation_leaves_every_logit_bit_identical() { + let fixture = fixture(); + let config = K3ExecutorConfig { + max_batch: 1, + max_ctx: 256, + kv_pages: 16, + num_layers: fixture.num_layers, + cuda_graph: false, + moe_transport: K3MoeTransport::MEGA, + }; + let Some(mut executor) = executor(config) else { + eprintln!("skipping: {CHECKPOINT_ENV} is not set to a mounted checkpoint"); + return; + }; + let mut feed = fixture.feed; + feed.extend(filler_tokens(200 - feed.len(), 7)); + + // Run 1 claims pages 0,1,2,3 (a fresh pool hands them out ascending). + let (tokens_a, logits_a) = replay_with_logits(&mut executor, &feed); + // Run 2: the release inside the replay returns run 1's pages, and the + // scramble reverses the whole list, so the claims are 15,14,13,12 — + // disjoint pages, opposite order. + executor.release(0); + executor.scramble_kv_pages(); + let (tokens_b, logits_b) = replay_with_logits(&mut executor, &feed); + + assert_eq!(tokens_a, tokens_b, "sampled trajectories differ"); + assert_bitwise(&logits_a, &logits_b, "page permutation"); +} + +/// Gate: a context past the old 128-token cap (and past 1024) decodes to +/// finite logits, reproduces itself bitwise on a fresh executor, and takes +/// the same trajectory through CUDA graphs. Logits are compared at every +/// 128th step and the last one; tokens at every step. +#[test] +#[ignore = "requires a Blackwell GPU and the K3 checkpoint"] +fn long_context_decode_is_self_consistent() { + let fixture = fixture(); + let config = K3ExecutorConfig { + max_batch: 1, + max_ctx: 2048, + kv_pages: 0, + num_layers: fixture.num_layers, + cuda_graph: false, + moe_transport: K3MoeTransport::MEGA, + }; + let mut feed = fixture.feed; + feed.extend(filler_tokens(1100 - feed.len(), 23)); + + let sample_step = |step: usize| -> bool { step == feed.len() - 1 || step.is_multiple_of(128) }; + let run = |executor: &mut K3Executor| -> (Vec, Vec>) { + executor.release(0); + let mut sampled = Vec::with_capacity(feed.len()); + let mut logits = Vec::new(); + for (step, &last_token) in feed.iter().enumerate() { + let tokens = executor + .decode(&[DecodeSlot { + slot: 0, + last_token, + }]) + .expect("the decode step should run"); + sampled.push(tokens[0]); + if sample_step(step) { + let row = executor.last_logits(0).expect("logit readback"); + assert!( + row.iter().all(|logit| logit.is_finite()), + "step {step} produced a non-finite logit" + ); + logits.push(row.into_iter().map(f32::to_bits).collect()); + } + } + (sampled, logits) + }; + + let Some(mut eager) = executor(config) else { + eprintln!("skipping: {CHECKPOINT_ENV} is not set to a mounted checkpoint"); + return; + }; + let (tokens_a, logits_a) = run(&mut eager); + let (tokens_b, logits_b) = run(&mut eager); + assert_eq!( + tokens_a, tokens_b, + "eager rerun took a different trajectory" + ); + assert_bitwise(&logits_a, &logits_b, "long context, eager rerun"); + drop(eager); + + let mut graphs = executor(K3ExecutorConfig { + cuda_graph: true, + ..config + }) + .expect("the checkpoint was there a moment ago"); + let (tokens_c, logits_c) = run(&mut graphs); + assert_eq!(tokens_a, tokens_c, "graphs took a different trajectory"); + assert_bitwise(&logits_a, &logits_c, "long context, graphs vs eager"); + eprintln!( + "long context: {} steps, final context {} tokens", + feed.len(), + feed.len() + ); +} + +/// Instrument for the absorbed-vs-expanded certification: replay the golden +/// fixture's forced feed and write every step's logit row (widened from bf16) +/// to `PEGAINFER_K3_LOGIT_DUMP` as flat little-endian f32, `[steps, vocab]` +/// row-major. Run once on the expanded-attention revision and once on this +/// one, then compare the files per step in bf16 ULP. +#[test] +#[ignore = "requires a Blackwell GPU and the K3 checkpoint"] +fn dump_forced_replay_logits() { + let Ok(dump_path) = std::env::var(DUMP_ENV) else { + eprintln!("skipping: {DUMP_ENV} is not set to an output path"); + return; + }; + let fixture = fixture(); + let config = K3ExecutorConfig { + max_batch: 1, + max_ctx: 128, + kv_pages: 0, + num_layers: fixture.num_layers, + cuda_graph: false, + moe_transport: K3MoeTransport::MEGA, + }; + let Some(mut executor) = executor(config) else { + eprintln!("skipping: {CHECKPOINT_ENV} is not set to a mounted checkpoint"); + return; + }; + let (sampled, logits) = replay_with_logits(&mut executor, &fixture.feed); + let mut file = std::fs::File::create(&dump_path).expect("create the dump file"); + for row in &logits { + let bytes: Vec = row.iter().flat_map(|bits| bits.to_le_bytes()).collect(); + file.write_all(&bytes).expect("write the dump file"); + } + eprintln!( + "dumped {} steps x {} logits to {dump_path}; sampled {:?}", + logits.len(), + logits.first().map_or(0, Vec::len), + sampled + ); +} diff --git a/pegainfer-kernels/KERNELS.md b/pegainfer-kernels/KERNELS.md index 5e7cfad6..307dc125 100644 --- a/pegainfer-kernels/KERNELS.md +++ b/pegainfer-kernels/KERNELS.md @@ -131,11 +131,16 @@ f32, and narrowing any of them to bf16 measurably flips routing decisions. | `k3.act.situ` | `pegainfer-k3` | `ops::k3_situ_batched_launch` | `k3_situ_batched` | `k3_situ_batched.cu` | TileLang AOT CUDA | `4*tanh(g/4)*sigmoid(g) * 25*tanh(u/25)` in f32, landing bf16 once; the betas are compiled in. n ∈ {6144 (shared), 33792 (dense)} × 10 buckets; the routed-expert situ is fused into the masked chain and the mega kernel, so no wide routed instantiation exists here. | | `k3.kda.conv_silu` | `pegainfer-k3` | `ops::k3_conv_silu_batched_launch` | `k3_conv_silu_batched` | `k3_conv_silu_batched.cu` | TileLang AOT CUDA | Causal depthwise convolution over the 4-slot window plus silu, one token per row. Consumes the projection's `[b, split_k, 12288]` f32 partial: its bf16 landing is `X`, the newest window slot; `Sn` is the shifted state the caller carries. Conv weights `[4, 12288]` are **f32** and have no batch axis; the window state is `[b, 3, 12288]`, one independent window per row. split_k = 1 × 10 buckets. | | `k3.kda.core` | `pegainfer-k3` | `ops::k3_kda_core_batched_launch` | `k3_kda_core_batched` | `k3_kda_core_batched.cu` | TileLang AOT CUDA | One delta-rule step per row, one (row, head) per block, `threads = head_dim`. State `[b, 96, 128, 128]` f32 laid out `[head, v_dim, k_dim]` per row with decay along k, read from `State` and written to `StateN` (must not alias). `Dt`/`Alog`/`Go` f32 weights with no batch axis, `Bt`/`G2` bf16; gate lower bound and eps compiled in. Gate partial uses split-K 1. 10 buckets. | -| `k3.attn.mla_decode` | `pegainfer-k3` | `ops::k3_mla_attn_batched_launch` | `k3_mla_attn_batched` | `k3_mla_attn_batched.cu` | TileLang AOT CUDA | NoPE full-context attention, one (row, head) per block (no GQA — all 96 heads hold their own K). The cache is **slot indexed**: each row owns a fixed `max_ctx` window and `Kc`/`Vc` are the contiguous single-row blocks, so this is not a paged cache and there is no block table. `N` is a per-row device i32 length and slots at or past it score NEG, so the step needs no host sync and a row at length `l` is bit-identical to the single-row kernel at `l`. Score landings follow the eager chain (f32 dot → bf16 → × bf16 scale → bf16 → f32 softmax → bf16 probabilities). max_ctx ∈ {128} × 10 buckets; both the instantiation count and the per-block shared memory scale with the capacity, so extending the list is a one-line but deliberate change in `generate.py`. | | `k3.moe.router_topk` | `pegainfer-k3` | `ops::k3_router_topk_batched_launch` | `k3_router_topk_batched` | `k3_router_topk_batched.cu` | TileLang AOT CUDA | Sigmoid router + biased top-k over already-merged `[b, E]` f32 score rows, one row per block. Serial O(topk*E) scan by thread 0 with lowest-index tie-break; weights gathered from the **un-biased** scores, denominator `+1e-20`, scaled by the bf16 `Rs[0]`. E ∈ {896 (full table), 224 (4-way EP shard)}, TOPK = 16, × 10 buckets. | | `k3.attnres.scores` | `pegainfer-k3` | `ops::k3_attnres_scores_batched_launch` | `k3_attnres_scores_batched` | `k3_attnres_scores_batched.cu` | TileLang AOT CUDA | Attention-residual candidate scoring, one block per (row, candidate): weightless RMS normalization then a dot with the pre-fused f32 scoring vector `[7168]`. Candidate `NB` is that row's prefix sum, below it its own snapshot history `[b, NB, 7168]`. NB ∈ 1..8 (the history grows one entry per 12 layers over 93 layers) × 10 buckets. | | `k3.attnres.mix` | `pegainfer-k3` | `ops::k3_attnres_mix_batched_launch` | `k3_attnres_mix_batched` | `k3_attnres_mix_batched.cu` | TileLang AOT CUDA | Softmax over each row's `NB+1` scores, then a probability-weighted mix of the **un-normalized** candidates landing bf16 once. Grid `(b, 7168/256)`; each block redoes its row's softmax. NB ∈ 1..8 × 10 buckets. | +## Kimi-K3 MLA Paged Attention + +| op_id | Runtime owner | Rust wrapper | FFI symbol | Source | Backend | Shape / layout notes | +| --- | --- | --- | --- | --- | --- | --- | +| `k3.attn.mla_paged_absorbed` | `pegainfer-k3` | `ops::k3_mla_paged_attn_launch` | `k3_mla_paged_attn_cuda` | `csrc/k3/k3_mla_paged_attn.cu` | CUDA | Absorbed-MLA NoPE decode over the paged **latent** cache, one (row, head) per block. The cache row is `[kv_lora 512 | rope 64]` bf16 in 64-token pages (`[page][layer][token][576]`, per-layer byte offset passed in); the block walks the per-row device block table by logical position, so page permutation is bit-identical and a `-1` page reads as zero latent. Query is absorbed against `w_kv_b`'s W_UK rows and the attended latent expanded with W_UV; the softmax is a 3-sweep recompute (max / sum / probs+attend), so there is no O(ctx) storage and **no compile-time context cap**. `N` is a per-row device i32 length — no host sync. Score landings replay the certified chain (f32 dot over 576 → bf16 → × bf16 scale in bf16 → f32 softmax → bf16 probabilities → f32 latent accumulation → bf16 → f32 W_UV expansion → bf16), documented step-by-step in the source header. Batch is a plain launch dimension (no per-bucket instantiation). | + ## Kimi-K3 MoE Bring-Up Surface Kimi-K3 uses the `pegainfer-kernels/k3` feature. Unlike `glm52`, `k3` does not diff --git a/pegainfer-kernels/build.rs b/pegainfer-kernels/build.rs index 0c4f209d..c6a1a634 100644 --- a/pegainfer-kernels/build.rs +++ b/pegainfer-kernels/build.rs @@ -1507,11 +1507,6 @@ const K3_TILELANG_LAUNCHERS: &[(&str, &str)] = &[ const void*, const void*, const float*, const float*, float*, void*, \ int, int, int, int", ), - ( - "k3_mla_attn_batched", - "const void*, const void*, const void*, const int*, const void*, void*, \ - int, int, int, int, int", - ), ( "k3_router_topk_batched", "const float*, const float*, const void*, int*, float*, int, int, int", diff --git a/pegainfer-kernels/csrc/k3/k3_mla_paged_attn.cu b/pegainfer-kernels/csrc/k3/k3_mla_paged_attn.cu new file mode 100644 index 00000000..154beb5d --- /dev/null +++ b/pegainfer-kernels/csrc/k3/k3_mla_paged_attn.cu @@ -0,0 +1,267 @@ +// Kimi-K3 absorbed-MLA decode over the paged latent KV cache. +// +// One kernel replaces the expanded-cache chain (kv_b expansion + slot-indexed +// attention): per (row, head) block it absorbs the query into latent space, +// walks the row's block table page by page, and expands the attended latent +// back to a value head — so the cache holds 576 bf16 per token per layer (the +// post-norm kv latent | the shared rope half) instead of the 96-head expanded +// K/V. NoPE: nothing here is position-dependent, which is what makes the +// latent row cacheable at all. +// +// --------------------------------------------------------------------------- +// Math (standard MLA absorption) +// --------------------------------------------------------------------------- +// The expanded path scores head h at position t as +// score = q_nope[h] . (W_UK[h] c_t) + q_rope[h] . rope_t +// which regroups to +// score = (W_UK[h]^T q_nope[h]) . c_t + q_rope[h] . rope_t +// so the absorbed query q_abs[h] = [W_UK[h]^T q_nope[h] | q_rope[h]] is one +// 576-wide row dotted against the cached row — MQA over a shared cache. The +// output regroups the same way: +// o[h] = sum_t p_t (W_UV[h] c_t) = W_UV[h] (sum_t p_t c_t) +// with the probs applied in 512-wide latent space and one W_UV expansion per +// head at the end. W_UK/W_UV are read straight out of the checkpoint's +// `w_kv_b` ([96 heads x (128 nope | 128 value)] x 512): rows [h*256, h*256+128) +// are W_UK[h], rows [h*256+128, h*256+256) are W_UV[h]. +// +// --------------------------------------------------------------------------- +// Rounding chain (each landing deliberate; 1-2 mirror every projection's +// f32-matmul-then-one-bf16-landing, 3-6 are the certified slot-indexed +// kernel's spelling) +// --------------------------------------------------------------------------- +// 1. q_abs[0..512) = bf16(f32 sum_d q_nope[d] * W_UK[d, j]), d ascending; +// q_abs[512..576) = q_rope, copied bf16. +// 2. dot(t) = f32 sum_d q_abs[d] * c_t[d] over 576, d ascending, one thread. +// 3. scl(t) = f32( bf16(dot) * scale ), the product taken in bf16. +// 4. m = max_t scl; tot = sum_t exp(scl - m) in f32 (per-thread strided +// partials in ascending t, then a fixed-order tree reduction). +// 5. p_t = bf16( exp(scl - m) / tot ). +// 6. o_lat[j] = bf16( f32 sum_t p_t * c_t[j] ), t ascending (chunk-major). +// 7. o[dv] = bf16( f32 sum_j W_UV[dv, j] * o_lat[j] ), j ascending. +// +// The context walk is three sweeps over the pages (max, sum, probs+attend): +// scores are recomputed rather than stored, so nothing is sized by the context +// length and there is no compile-time cap. A recomputed score is the same +// expression over the same operands in the same order, hence bit-identical. +// The walk is by *logical* position — the block table only selects which +// physical page backs a 64-token window — so any permutation of physical +// pages produces bit-identical output. A page id below zero (padding rows) +// reads as a zero latent row, which is what the retired slot-indexed kernel's +// zeroed cache produced. +// +// CUDA-graph safety: no allocation, no host readback, no device-side launch; +// launch geometry is (b, heads) with everything else read from device tensors. + +#include "../common.cuh" +#include "../shared/ffi_guard.cuh" + +#include +#include + +namespace { + +constexpr int kNope = 128; // qk_nope_head_dim +constexpr int kRope = 64; // qk_rope_head_dim +constexpr int kLatent = 512; // kv_lora_rank +constexpr int kRow = kLatent + kRope; // cached latent row width +constexpr int kVd = 128; // v_head_dim +constexpr int kPageTokens = 64; +constexpr int kThreads = 128; +constexpr int kWarps = kThreads / WARP_SIZE; +constexpr int kDimsPerThread = kLatent / kThreads; +constexpr float kNeg = -1.0e30f; + +// Fixed-order block reductions: warp shuffle trees, then thread 0 folds the +// warp partials in ascending warp order. The result is broadcast via stage[0]. +__device__ __forceinline__ float block_max(float value, float* stage) { + value = warp_reduce_max(value); + if ((threadIdx.x & (WARP_SIZE - 1)) == 0) stage[threadIdx.x / WARP_SIZE] = value; + __syncthreads(); + if (threadIdx.x == 0) { + float folded = stage[0]; + for (int w = 1; w < kWarps; ++w) folded = fmaxf(folded, stage[w]); + stage[0] = folded; + } + __syncthreads(); + float out = stage[0]; + __syncthreads(); + return out; +} + +__device__ __forceinline__ float block_sum(float value, float* stage) { + value = warp_reduce_sum(value); + if ((threadIdx.x & (WARP_SIZE - 1)) == 0) stage[threadIdx.x / WARP_SIZE] = value; + __syncthreads(); + if (threadIdx.x == 0) { + float folded = stage[0]; + for (int w = 1; w < kWarps; ++w) folded += stage[w]; + stage[0] = folded; + } + __syncthreads(); + float out = stage[0]; + __syncthreads(); + return out; +} + +// Landed score of logical position `s`: chain steps 2-3. Deterministic in the +// operands alone, so the three sweeps recompute it bit-identically. +__device__ __forceinline__ float position_score( + const __nv_bfloat16* __restrict__ q_abs, + const __nv_bfloat16* __restrict__ cache, const int* __restrict__ bt, + long long page_stride, int s, __nv_bfloat16 sc) { + const int page = bt[s / kPageTokens]; + float acc = 0.0f; + if (page >= 0) { + const __nv_bfloat16* c = cache + (long long)page * page_stride + + (long long)(s % kPageTokens) * kRow; + for (int d = 0; d < kRow; ++d) { + acc += __bfloat162float(q_abs[d]) * __bfloat162float(c[d]); + } + } + return __bfloat162float(__hmul(__float2bfloat16_rn(acc), sc)); +} + +__global__ void mla_paged_absorbed_attn_kernel( + const __nv_bfloat16* __restrict__ q, // [b, heads * 192] + const __nv_bfloat16* __restrict__ w_kv_b, // [heads * 256, 512] + const __nv_bfloat16* __restrict__ cache, // layer-shifted slab base + const int* __restrict__ table, // [b, max_pages] + int max_pages, + long long page_stride, // elements from page to page + const int* __restrict__ n, // [b] context lengths + const __nv_bfloat16* __restrict__ scale, // [1] softmax scale + __nv_bfloat16* __restrict__ o) { // [b, heads * 128] + const int bb = blockIdx.x; + const int bh = blockIdx.y; + const int heads = gridDim.y; + const int tid = threadIdx.x; + const int ctx = n[bb]; + const __nv_bfloat16* qh = + q + ((size_t)bb * heads + bh) * (size_t)(kNope + kRope); + const int* bt = table + (size_t)bb * max_pages; + const __nv_bfloat16 sc = scale[0]; + + __shared__ __nv_bfloat16 q_abs[kRow]; + __shared__ __nv_bfloat16 probs[kPageTokens]; + __shared__ __nv_bfloat16 o_lat[kLatent]; + __shared__ float stage[kWarps]; + + // Chain step 1: the absorbed query. + const __nv_bfloat16* w_uk = w_kv_b + (size_t)bh * 2 * kVd * kLatent; + for (int j = tid; j < kLatent; j += kThreads) { + float acc = 0.0f; + for (int d = 0; d < kNope; ++d) { + acc += __bfloat162float(qh[d]) * __bfloat162float(w_uk[(size_t)d * kLatent + j]); + } + q_abs[j] = __float2bfloat16_rn(acc); + } + for (int j = tid; j < kRope; j += kThreads) { + q_abs[kLatent + j] = qh[kNope + j]; + } + __syncthreads(); + + // Sweep 1 (chain step 4a): the score maximum. + float local = kNeg; + for (int s = tid; s < ctx; s += kThreads) { + local = fmaxf(local, position_score(q_abs, cache, bt, page_stride, s, sc)); + } + const float mx = block_max(local, stage); + + // Sweep 2 (chain step 4b): the softmax denominator. + local = 0.0f; + for (int s = tid; s < ctx; s += kThreads) { + local += expf(position_score(q_abs, cache, bt, page_stride, s, sc) - mx); + } + const float tot = block_sum(local, stage); + + // Sweep 3 (chain steps 5-6): bf16 probs per page chunk, then the latent + // accumulation, each thread owning kDimsPerThread strided latent dims. + float oacc[kDimsPerThread]; + for (int i = 0; i < kDimsPerThread; ++i) oacc[i] = 0.0f; + const int chunks = (ctx + kPageTokens - 1) / kPageTokens; + for (int chunk = 0; chunk < chunks; ++chunk) { + const int base = chunk * kPageTokens; + const int len = min(kPageTokens, ctx - base); + const int page = bt[chunk]; + __syncthreads(); // probs from the previous chunk are consumed + if (tid < len) { + const float scl = + position_score(q_abs, cache, bt, page_stride, base + tid, sc); + probs[tid] = __float2bfloat16_rn(expf(scl - mx) / tot); + } + __syncthreads(); + if (page >= 0) { + const __nv_bfloat16* cpage = cache + (long long)page * page_stride; + for (int i = 0; i < kDimsPerThread; ++i) { + const int j = i * kThreads + tid; + float acc = oacc[i]; + for (int t = 0; t < len; ++t) { + acc += __bfloat162float(probs[t]) * + __bfloat162float(cpage[(size_t)t * kRow + j]); + } + oacc[i] = acc; + } + } + } + for (int i = 0; i < kDimsPerThread; ++i) { + o_lat[i * kThreads + tid] = __float2bfloat16_rn(oacc[i]); + } + __syncthreads(); + + // Chain step 7: the W_UV expansion, one value dim per thread. + const __nv_bfloat16* w_uv = w_kv_b + ((size_t)bh * 2 * kVd + kNope) * kLatent; + const int dv = tid; // kThreads == kVd + float acc = 0.0f; + for (int j = 0; j < kLatent; ++j) { + acc += __bfloat162float(w_uv[(size_t)dv * kLatent + j]) * + __bfloat162float(o_lat[j]); + } + o[((size_t)bb * heads + bh) * (size_t)kVd + dv] = __float2bfloat16_rn(acc); +} + +CUresult map_cuda_error(cudaError_t err) { + if (err == cudaSuccess) return CUDA_SUCCESS; + if (err == cudaErrorInvalidValue || err == cudaErrorInvalidDevicePointer) { + return CUDA_ERROR_INVALID_VALUE; + } + if (err == cudaErrorMemoryAllocation) return CUDA_ERROR_OUT_OF_MEMORY; + if (err == cudaErrorNotSupported) return CUDA_ERROR_NOT_SUPPORTED; + return CUDA_ERROR_LAUNCH_FAILED; +} + +} // namespace + +extern "C" { + +// Absorbed-MLA decode over the paged latent cache. `cache` is the pool slab; +// `layer_offset` (elements) shifts it to this layer's slice inside every page +// and `page_stride` (elements) is the page-to-page distance, so the kernel is +// layout-agnostic beyond "64-token slices of 576-wide rows". `table` is the +// device block table (`-1` = unmapped, read as zero latent), `n` the per-row +// device context lengths — no host sync anywhere. Geometry is pinned: qk_dim +// 192 (128 nope + 64 rope), v_dim 128, 96-head `w_kv_b` rows per head +// `[128 nope | 128 value] x 512`. +CUresult k3_mla_paged_attn_cuda(const __nv_bfloat16* q, + const __nv_bfloat16* w_kv_b, + const __nv_bfloat16* cache, + long long layer_offset, long long page_stride, + const int* table, int max_pages, const int* n, + const __nv_bfloat16* scale, __nv_bfloat16* o, + int b, int num_heads, int qk_dim, int v_dim, + cudaStream_t stream) { + PEGAINFER_FFI_GUARD_BEGIN + if (q == nullptr || w_kv_b == nullptr || cache == nullptr || + table == nullptr || n == nullptr || scale == nullptr || o == nullptr || + b <= 0 || num_heads <= 0 || qk_dim != kNope + kRope || v_dim != kVd || + max_pages <= 0 || layer_offset < 0 || page_stride < kPageTokens * kRow) { + return CUDA_ERROR_INVALID_VALUE; + } + dim3 grid(b, num_heads); + mla_paged_absorbed_attn_kernel<<>>( + q, w_kv_b, cache + layer_offset, table, max_pages, page_stride, n, scale, + o); + return map_cuda_error(cudaGetLastError()); + PEGAINFER_FFI_GUARD_END(CUDA_ERROR_UNKNOWN) +} + +} // extern "C" diff --git a/pegainfer-kernels/csrc/shared/elementwise.cu b/pegainfer-kernels/csrc/shared/elementwise.cu index 8fe0db6e..48a3f303 100644 --- a/pegainfer-kernels/csrc/shared/elementwise.cu +++ b/pegainfer-kernels/csrc/shared/elementwise.cu @@ -177,7 +177,10 @@ __global__ void scaled_add_rows_indexed_kernel( row < rows; row += gridDim.x * blockDim.x) { int delta_idx = token * rows + row; - int out_idx = out_token * out_hidden_dim + row_offset + row; + // out_token * out_hidden_dim overflows i32 for large row pools (the K3 + // paged-KV slab exceeds 2^31 elements), so the element index is size_t. + size_t out_idx = + (size_t)out_token * out_hidden_dim + row_offset + row; float base = __bfloat162float(out[out_idx]); float add = __bfloat162float(delta[delta_idx]) * scale; out[out_idx] = __float2bfloat16(base + add); diff --git a/pegainfer-kernels/src/ffi/k3.rs b/pegainfer-kernels/src/ffi/k3.rs index c0178d5a..e123b4c3 100644 --- a/pegainfer-kernels/src/ffi/k3.rs +++ b/pegainfer-kernels/src/ffi/k3.rs @@ -129,6 +129,30 @@ unsafe extern "C" { stream: CUstream, ) -> CUresult; + /// Absorbed-MLA decode over the paged latent cache + /// (`csrc/k3/k3_mla_paged_attn.cu`). One (row, head) block absorbs the + /// query against `w_kv_b`'s W_UK, walks the row's block table, and expands + /// the attended latent with W_UV. `layer_offset`/`page_stride` are in + /// elements; `table` is `[b, max_pages]` i32 (`-1` = unmapped, read as + /// zero latent) and `n` the per-row device context length. + pub fn k3_mla_paged_attn_cuda( + q: *const Half, + w_kv_b: *const Half, + cache: *const Half, + layer_offset: i64, + page_stride: i64, + table: *const i32, + max_pages: i32, + n: *const i32, + scale: *const Half, + o: *mut Half, + b: i32, + num_heads: i32, + qk_dim: i32, + v_dim: i32, + stream: CUstream, + ) -> CUresult; + // --- fused MegaMoE (see `csrc/k3/k3_mega_moe_sm100.cu`) --- /// Token-count alignment the MegaMoE API enforces on diff --git a/pegainfer-kernels/src/ffi/k3_tilelang.rs b/pegainfer-kernels/src/ffi/k3_tilelang.rs index c68e8afe..6cce97f1 100644 --- a/pegainfer-kernels/src/ffi/k3_tilelang.rs +++ b/pegainfer-kernels/src/ffi/k3_tilelang.rs @@ -137,27 +137,6 @@ unsafe extern "C" { stream: CUstream, ) -> i32; - /// NoPE full-context MLA decode over a slot-indexed cache: each row owns a - /// fixed `max_ctx` window, so `Kc [b, max_ctx, num_heads * qk_dim]` and - /// `Vc [b, max_ctx, num_heads * v_dim]` are dense, not paged. `N [b]` is - /// the per-row device-side context length; slots at or past it score - /// negative infinity, so no host sync is needed. `Sc [1]` is the shared - /// bf16 softmax scale. - pub fn k3_mla_attn_batched( - q: *const c_void, - kc: *const c_void, - vc: *const c_void, - n: *const i32, - sc: *const c_void, - o: *mut c_void, - b: i32, - num_heads: i32, - qk_dim: i32, - v_dim: i32, - max_ctx: i32, - stream: CUstream, - ) -> i32; - /// Sigmoid router plus biased top-k over merged f32 score rows /// `S [b, num_experts]`, with `Bias [num_experts]` f32 and the bf16 routed /// scale `Rs [1]`. Writes `Idx [b, topk]` i32 and `Wts [b, topk]` f32. diff --git a/pegainfer-kernels/src/ops/k3/mla_paged.rs b/pegainfer-kernels/src/ops/k3/mla_paged.rs new file mode 100644 index 00000000..5351a5fe --- /dev/null +++ b/pegainfer-kernels/src/ops/k3/mla_paged.rs @@ -0,0 +1,123 @@ +//! Kimi-K3 absorbed-MLA decode over the paged latent KV cache. +//! +//! The cache holds one 576-wide bf16 latent row per token per MLA layer (the +//! post-norm kv latent | the shared rope half — NoPE, nothing rotated), in a +//! page-first slab of 64-token pages addressed through a per-row block table. +//! The kernel absorbs the query into latent space against `w_kv_b`'s W_UK +//! rows, attends MQA-style over the shared rows, and expands the attended +//! latent with the W_UV rows — see `csrc/k3/k3_mla_paged_attn.cu` for the +//! documented rounding chain. +//! +//! There is no compile-time context cap: the page walk is a runtime loop and +//! the per-row length `N` lives on the device, so the step needs no host sync. +//! Physical page ids never enter the arithmetic — the walk is by logical +//! position — so any page permutation is bit-identical. Batch is a plain +//! launch dimension here (no per-bucket instantiation), but callers still run +//! the compiled buckets: every other kernel in the step is bucket-shaped. + +use anyhow::Result; +use anyhow::anyhow; +use anyhow::ensure; +use cudarc::driver::CudaSlice; +use cudarc::driver::DevicePtr; +use cudarc::driver::DevicePtrMut; +use half::bf16; + +use super::super::k3_tilelang::K3_QK_DIM; +use super::super::k3_tilelang::K3_V_DIM; +use crate::ffi; +use crate::tensor::DeviceContext; + +/// Tokens per MLA KV page — the kernel's only compile-time context term. +pub const K3_KV_PAGE_TOKENS: usize = 64; +/// kv_lora_rank, the latent width the absorption works in. +const K3_KV_LORA: usize = 512; +/// Cached latent row width: kv_lora_rank 512 | rope 64. +pub const K3_KV_LATENT_ROW: usize = K3_KV_LORA + 64; + +/// Absorbed-MLA decode: `q [b, heads * 192]`, `w_kv_b [heads * 256, 512]` +/// (checkpoint orientation — per head `[128 nope | 128 value] x 512`), the +/// pool `slab` with this layer's slice at `layer_offset` elements and pages +/// `page_stride` elements apart, the device block `table [rows, max_pages]` +/// (`-1` = unmapped, read as zero latent), per-row device lengths `n [b]`, +/// the shared bf16 softmax `scale [1]`, out `o [b, heads * 128]`. +#[allow(clippy::too_many_arguments)] +pub fn k3_mla_paged_attn_launch( + ctx: &DeviceContext, + b: usize, + num_heads: usize, + q: &CudaSlice, + w_kv_b: &CudaSlice, + slab: &CudaSlice, + layer_offset: usize, + page_stride: usize, + table: &CudaSlice, + max_pages: usize, + n: &CudaSlice, + scale: &CudaSlice, + o: &mut CudaSlice, +) -> Result<()> { + ensure!(b > 0 && num_heads > 0, "K3 paged MLA needs rows and heads"); + ensure!( + page_stride >= K3_KV_PAGE_TOKENS * K3_KV_LATENT_ROW + && layer_offset + K3_KV_PAGE_TOKENS * K3_KV_LATENT_ROW <= page_stride, + "K3 paged MLA layer slice [{layer_offset}..) does not fit the {page_stride}-element page" + ); + ensure!( + q.len() >= b * num_heads * K3_QK_DIM + && w_kv_b.len() >= num_heads * 2 * K3_V_DIM * K3_KV_LORA + && table.len() >= b * max_pages + && n.len() >= b + && !scale.is_empty() + && o.len() >= b * num_heads * K3_V_DIM, + "K3 paged MLA buffers too small for b={b}, heads={num_heads}, max_pages={max_pages}: \ + q {}, w_kv_b {}, table {}, n {}, o {}", + q.len(), + w_kv_b.len(), + table.len(), + n.len(), + o.len() + ); + let (q_ptr, _q_guard) = q.device_ptr(&ctx.stream); + let (w_ptr, _w_guard) = w_kv_b.device_ptr(&ctx.stream); + let (slab_ptr, _slab_guard) = slab.device_ptr(&ctx.stream); + let (table_ptr, _table_guard) = table.device_ptr(&ctx.stream); + let (n_ptr, _n_guard) = n.device_ptr(&ctx.stream); + let (scale_ptr, _scale_guard) = scale.device_ptr(&ctx.stream); + let (o_ptr, _o_guard) = o.device_ptr_mut(&ctx.stream); + unsafe { + ffi::k3_mla_paged_attn_cuda( + q_ptr as *const ffi::Half, + w_ptr as *const ffi::Half, + slab_ptr as *const ffi::Half, + i64::try_from(layer_offset)?, + i64::try_from(page_stride)?, + table_ptr as *const i32, + i32::try_from(max_pages)?, + n_ptr as *const i32, + scale_ptr as *const ffi::Half, + o_ptr as *mut ffi::Half, + i32::try_from(b)?, + i32::try_from(num_heads)?, + i32::try_from(K3_QK_DIM)?, + i32::try_from(K3_V_DIM)?, + crate::tensor::active_cu_stream(ctx), + ) + } + .result() + .map_err(|err| anyhow!("K3 paged MLA attention launch failed: {err}")) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn geometry_agrees_with_the_cache_row() { + // The cached row is [kv_lora 512 | rope 64]; the query per head is + // [nope 128 | rope 64]; the value head is 128 wide. + assert_eq!(K3_KV_LATENT_ROW, K3_KV_LORA + (K3_QK_DIM - 128)); + assert_eq!(K3_V_DIM, 128); + assert_eq!(K3_KV_PAGE_TOKENS, 64); + } +} diff --git a/pegainfer-kernels/src/ops/k3/mod.rs b/pegainfer-kernels/src/ops/k3/mod.rs index 0c25d69f..51d0ec50 100644 --- a/pegainfer-kernels/src/ops/k3/mod.rs +++ b/pegainfer-kernels/src/ops/k3/mod.rs @@ -2,8 +2,10 @@ mod deepgemm; mod mega_moe; +mod mla_paged; mod moe_chain; pub use deepgemm::*; pub use mega_moe::*; +pub use mla_paged::*; pub use moe_chain::*; diff --git a/pegainfer-kernels/src/ops/k3_tilelang.rs b/pegainfer-kernels/src/ops/k3_tilelang.rs index bc6128ed..af89fc27 100644 --- a/pegainfer-kernels/src/ops/k3_tilelang.rs +++ b/pegainfer-kernels/src/ops/k3_tilelang.rs @@ -1,12 +1,13 @@ //! K3 TileLang batched decode kernels: safe wrappers over the AOT dispatch //! launchers in `ffi::k3_tilelang`. //! -//! The set covers one whole K3 decode step that is not a GEMM — norms and the -//! bf16 landings of the framework GEMMs, the KDA convolution and delta rule, -//! MLA attention, the MoE router and expert combine, the situ activation and -//! the attention-residual mix. Dense projections are served by cuBLASLt and -//! the routed experts by the DeepGEMM masked grouped-GEMM chain, so no GEMV -//! lives here. The wrappers keep the certified kernels' operand names, so an +//! The set covers one whole K3 decode step that is not a GEMM or attention — +//! norms and the bf16 landings of the framework GEMMs, the KDA convolution and +//! delta rule, the MoE router and expert combine, the situ activation and the +//! attention-residual mix. Dense projections are served by cuBLASLt, the +//! routed experts by the DeepGEMM masked grouped-GEMM chain, and MLA decode by +//! the hand-written absorbed paged kernel (`ops::k3::mla_paged`), so neither a +//! GEMV nor an attention family lives here. The wrappers keep the certified kernels' operand names, so an //! executor written against them reads like the Python engine's launch //! sequence. //! @@ -64,14 +65,13 @@ pub const K3_KDA_DIM: usize = K3_KDA_HEADS * K3_KDA_HEAD_DIM; /// Short-convolution window; the carried state is `K3_CONV_WIDTH - 1` slots. pub const K3_CONV_WIDTH: usize = 4; -/// MLA head count and the query/key and value widths per head. +/// MLA head count and the query/key and value widths per head. MLA decode +/// itself is not a TileLang family — it is the hand-written absorbed paged +/// kernel in `ops::k3::mla_paged` — but its head geometry is shared with the +/// landings here. pub const K3_MLA_HEADS: usize = 96; pub const K3_QK_DIM: usize = 192; pub const K3_V_DIM: usize = 128; -/// Cache capacities `k3_mla_attn_batched` is instantiated for. Both the -/// instantiation count and the per-block shared memory scale with the -/// capacity, so extending this list is a deliberate generator change. -pub const K3_MAX_CTX: [usize; 1] = [128]; /// Round a live row count up to the bucket that will run it. /// @@ -512,73 +512,6 @@ pub fn k3_kda_core_batched_launch( ) } -/// NoPE full-context MLA decode over a slot-indexed cache. -/// -/// Each row owns a fixed `max_ctx` window, so `kc`/`vc` are dense per row and -/// there is no block table. `n` is the per-row context length *on the device*; -/// slots at or past it score negative infinity, so the step needs no host -/// sync and a row at length `l` is bit-identical to the single-row kernel run -/// at `l`. `sc` is the shared bf16 softmax scale. -#[allow(clippy::too_many_arguments)] -pub fn k3_mla_attn_batched_launch( - ctx: &DeviceContext, - b: usize, - num_heads: usize, - qk_dim: usize, - v_dim: usize, - max_ctx: usize, - q: &CudaSlice, - kc: &CudaSlice, - vc: &CudaSlice, - n: &CudaSlice, - sc: &CudaSlice, - o: &mut CudaSlice, -) -> Result<()> { - check_bucket(b)?; - ensure!( - q.len() >= b * num_heads * qk_dim - && kc.len() >= b * max_ctx * num_heads * qk_dim - && vc.len() >= b * max_ctx * num_heads * v_dim - && n.len() >= b - && !sc.is_empty() - && o.len() >= b * num_heads * v_dim, - "K3 mla_attn buffers too small for b={b}, heads={num_heads}, max_ctx={max_ctx}: \ - q {}, kc {}, vc {}, n {}, sc {}, o {}", - q.len(), - kc.len(), - vc.len(), - n.len(), - sc.len(), - o.len() - ); - let (q_ptr, _q_guard) = q.device_ptr(&ctx.stream); - let (kc_ptr, _kc_guard) = kc.device_ptr(&ctx.stream); - let (vc_ptr, _vc_guard) = vc.device_ptr(&ctx.stream); - let (n_ptr, _n_guard) = n.device_ptr(&ctx.stream); - let (sc_ptr, _sc_guard) = sc.device_ptr(&ctx.stream); - let (o_ptr, _o_guard) = o.device_ptr_mut(&ctx.stream); - let rc = unsafe { - ffi::k3_mla_attn_batched( - q_ptr as *const c_void, - kc_ptr as *const c_void, - vc_ptr as *const c_void, - n_ptr as *const i32, - sc_ptr as *const c_void, - o_ptr as *mut c_void, - b as i32, - num_heads as i32, - qk_dim as i32, - v_dim as i32, - max_ctx as i32, - ctx.stream.cu_stream(), - ) - }; - check( - rc, - &format!("K3 mla_attn_batched (B={b}, NH={num_heads}, CAP={max_ctx})"), - ) -} - /// Sigmoid router plus biased top-k over already-merged f32 score rows. /// /// The weights come from the *un-biased* scores, are normalized with a @@ -784,6 +717,5 @@ mod tests { assert_eq!(K3_ATTNRES_MAX_BLOCKS, 93_usize.div_ceil(12)); // The 4-way expert-parallel shard of the full table. assert_eq!(K3_ROUTER_EXPERTS[0] / 4, K3_ROUTER_EXPERTS[1]); - assert!(K3_MAX_CTX.iter().all(|cap| cap % 128 == 0)); } }