diff --git a/.agents/NOW.md b/.agents/NOW.md index 1fc0dd599..c796bff17 100644 --- a/.agents/NOW.md +++ b/.agents/NOW.md @@ -18,7 +18,7 @@ Working head: `row/backend-rocm-w0` (#41). Prior: benchmark checkpoint | f32-out GEMV audit | Only laguna + ds4 bf16 tower affected; gate models unaffected | Re-verify ds4 tower same-tool | | Invocation-parity prevention | CI guard + checklist landing | Merge; build-verify `kGemvHeuristicAlgos` on dgx | | MiniMax-H3 lane | **fl2va COHERENT; ref2va grid DIAGNOSED (#95): NO loader bug; bf16 13-shard DiT INDEXES** | residual = community-NVFP4 quant fidelity §8.12; no bf16 render yet | -| Kimi-Linear-48B (KDA/NoPE-MLA/MoE) | device-KDA **122/128 + 4.24 tok/s** best (§15); device NoPE-MLA `VT_KIMI_DEVICE_MLA` **MEASURED-NEGATIVE 122→109** (§16, #107 — f32-softmax≠FA2, coin-flips). Both OFF ≠STRICT | STRICT owes ACTUAL kernels: chunk_kda + paged FA2 MLA + incremental | +| Kimi-Linear-48B (KDA/NoPE-MLA/MoE) | device-KDA **122/128, 4.24 tok/s** best (§15); MLA device NEG (§16). chunk_kda prefill AOT **SPIKED**: 5 kernels authored+pinned+recipe (§17). Bar = MEET vLLM speed | Phase-2: regen harness, wire `vt::KdaChunkPrefill`, gate STRICT + vLLM 0.82 ladder | | 35B fresh grid | **BOUND** @`1ea26427`: 0.93-1.03x, c16 0.93x. INTAKE + Option A both NEGATIVE | Lever left: prefill glue (#61) | | Qwen3.5-4B revalidation | 0.9971x @`59674cf1` (#35); TTFT/PSS pass, TPOT/ITL open | `docs/bench-evidence/` | | MXFP4 parity | c1 1.020, c2-c8 0.962-0.969. **#82 CLOSED: ptxas-lineage REFUTED (A/B ties our+vLLM PTX all ptxas/JIT; +10us=engine context, not codegen)** | TERMINAL: at parity | diff --git a/.agents/benchmark-record.md b/.agents/benchmark-record.md index 4881cf80f..dd2adb85d 100644 --- a/.agents/benchmark-record.md +++ b/.agents/benchmark-record.md @@ -14673,3 +14673,17 @@ Full 48.9B GB10 gate vs the §12 STRICT `greedy_ids.npy` (single-load per config WHY NEGATIVE (the §14 razor, re-proven). device-KDA WORKS (106→122) because the recurrence is the SAME algorithm as vLLM's decode kernel, only f32-on-bf16 — it MATCHES. But vLLM's MLA prefill uses FA2 (a specific flash tiling/reduction ORDER); `vt::Attention`'s plain f32 online-softmax is the right MATH but a DIFFERENT reduction order, so — exactly like §14's host-precision plateau — it COIN-FLIPS near-ties: it BREAKS p3 16/16→3/16 (got `220,41938,382,1810,…163586,163586` — the same `163586×` degenerate repeat §14's bf16 knobs caused) while p7 stays diverged at 10/16. And it is SLOWER (4.24→3.89): the per-`(t,h)` key/value build copies + the 192-dim pad-V waste add overhead to the O(n²) recompute path. An approximation of vLLM's kernel is not enough — only the ACTUAL kernel matches. VERDICT: `VT_KIMI_DEVICE_MLA` STAYS OFF, kept as a documented-MEASURED-NEGATIVE A/B knob (parity-lever precedent: §14's `ISLAND_F32ACC`). device-KDA (122/128, 4.24 tok/s) remains the best config, itself default OFF (122 ≠ STRICT). The one-brick STRICT-close did NOT land. STRICT residual, SHARPENED: needs vLLM's ACTUAL kernels, not a device approximation — (c) the **chunk_kda** prefill family (`chunk_kda_scaled_dot_kkt`+`recompute_w_u`+`chunk_gla_fwd_o_gk`+`fused_kda_gate_chunk_cumsum`, FLA `ops/kda.py`) via a Triton-AOT regen for sm_121a (`scripts/regen-triton-aot.sh` + new `triton_kernels/*.py`), the named prime suspect; (d) the paged FA2 `mla::ForwardMlaAttentionBlock` (NOT the `vt::Attention` approximation tried here); (e) paged-incremental decode (needs a decode/paged-attn op with `query_len≠key_len`, which `vt::Attention` cannot express; kills the O(n²)). Each is a substantial multi-kernel brick, recorded as the named follow-on. Row STAYS ACTIVE. HONEST bar stays HW-forced-indirect (vLLM cannot serve this bf16 on one GB10 with KV headroom — §14). + +## 2026-08-07 — Kimi-Linear-48B: chunk_kda PREFILL AOT port SPIKED (kernel set + pinned-config record + regen recipe; NO gate yet) (`row/KIMI-CHUNK-KDA-AOT`) + +The §16 STRICT residual (c) — "vLLM processes the PROMPT with the CHUNKED `chunk_kda` kernel; the recurrent form we run is a different reduction order that coin-flips the p7 near-tie" — was SCOPED and de-risked to mechanical execution. This is a PLANNING/authoring spike; it MEASURES nothing on GPU and claims NO STRICT verdict. + +EXACT forward-only kernel set (`chunk_kda_with_fused_gate` → `_fwd` → `_chunk_kda_fwd_with_cumulative_g`, FLA `kda.py` @ 555967922), launch order: (1) `kda_gate_cumsum_fwd_kernel` `:1182` NEW; (2) `chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_inter` `:521` NEW; (3) `…intra_sub_intra` `:627` NEW; (4) `solve_tril` REUSE `gdn_tril_h32` (byte-identical sig at H=32,BT=64,A f32→bf16); (5) `recompute_w_u_fwd_kernel` `:817` NEW (KDA per-K-channel `exp2(gk)` + STORE_KG; ≠ GDN `wy_fast.py`); (6) `chunk_gated_delta_rule_fwd_h` REUSE `chunk_delta_h.py` with a NEW pin `USE_GK=1,USE_EXP2=1,USE_G=0,Hg=32` (the GDN `gdn_deltah` is `USE_G=1,USE_GK=0,USE_EXP2=0` — NOT reusable as-is); (7) `chunk_gla_fwd_kernel_o` `:1019` NEW. → 5 new Triton kernels + 1 new pin of an existing .py + 1 pure reuse. Decode STAYS the #104 recurrent `vt::KdaGatedDeltaRule` (mirrors vLLM's prefill=chunk / decode=recurrent split). Backward NOT owed. + +Pinned-config record (Kimi KDA: H=32, Hg=32, K=V=128, BT=FLA_CHUNK_SIZE=64, BC=16, NC=4): the full `_vllm_triton_aot_declare` recipe (bases, kernels, BK/BV/BD, warps/stages, grids, signatures) is in spec §17.3. num_warps/num_stages are correctness-invariant (pinned mirroring GDN); shape pins from FLA driver-fixed values + heuristic lists; dtypes MIRROR FLA (bf16 activations/intermediates, fp32 for gk-cumulative/A/Aqk/state). Scalar constants baked as literals (Triton AOT mis-packs fp32 scalars, per `chunk_o.py` note 3): `scale=K**-0.5`, softplus `beta=1.0`/`threshold=20.0`/`cumsum_scale=RCP_LN2`, `DOT_PRECISION="ieee"`. + +DELIVERED: 5 harness bodies authored (verbatim FLA ports, AOT-adapted) STAGED in `.agents/specs/kda-chunk-aot/` (CI-safe: drift check globs `triton_kernels/*.py` non-recursively); the `vt::KdaChunkPrefill` op design (buffer layout + 6-launch order + dispatch guard + prefill/decode split, mirrors `cuda_gdn.cu` GdnPrefill); the RED-first gate plan (unit vs #104 recurrent + #173 host refs + an FLA-python golden at Kimi shapes; then the full 48.9B GB10 gate vs the §12 STRICT golden with `DEVICE_KDA=1` + chunk-prefill). + +NOT YET (Phase-2, coupled — the harness signatures depend on the op's confirmed buffer dtypes so regen is premature until the op exists): move harness to `triton_kernels/`, add the §17.3 declarations, regen sm_121a cubins (`scripts/regen-triton-aot.sh`), wire `vt::KdaChunkPrefill`, run the RED-first unit + FLA golden + 48.9B STRICT gate + tok/s/TTFT ladder. Row STAYS ACTIVE. No default flips (nothing measured). + +**USER 2026-08-07 (mid-flight): the Kimi-Linear success bar is MEET vLLM SPEED — the §14/§16/#107 "HW-forced-indirect" framing is SUPERSEDED.** vLLM demonstrably RUNS Kimi-Linear-48B on ONE GB10: the §12 STRICT golden capture used it at `gpu_memory_utilization=0.82`, single-seq, eager. So the Phase-2 speed ladder MUST include a vLLM arm at that EXACT recipe (single-seq, eager, util 0.82, the §12 launch config) measuring steady decode tok/s + prefill TTFT on the SAME prompts as our arm — SEQUENTIAL after our runs, `local-ai-worker` PARKED, `drop_caches` before wall-clock, and PRE-WARM FlashInfer's autotune in a throwaway start at TINY util FIRST (cold autotune at util 0.82 with 91.5 GiB weights = the tightest vLLM config ever run on this box = a recorded OOM-reboot trigger; memory monitor mandatory, ONE attempt, if it OOMs record honestly and do NOT retry higher). The tok/s ladder then reads ours-vs-vLLM matched-config: the lane's distance-to-bar becomes a MEASURED number. Recorded in spec §17.5; below vLLM on any axis is an open gap, not done. diff --git a/.agents/specs/kda-chunk-aot/chunk_gla_o.py b/.agents/specs/kda-chunk-aot/chunk_gla_o.py new file mode 100644 index 000000000..010280b40 --- /dev/null +++ b/.agents/specs/kda-chunk-aot/chunk_gla_o.py @@ -0,0 +1,143 @@ +# triton_kernels/chunk_gla_o.py (STAGED — see .agents/specs/kimi-linear.md §17) +# +# KDA chunk-prefill kernel 5/5: the GLA-style output with per-K-channel gk decay. +# STAGED authored harness body (Phase-1 spike); Phase-2 moves it to +# triton_kernels/ + regenerates the sm_121a cubin. +# +# Ported VERBATIM FROM (vLLM oracle @ pin 555967922): +# vllm/third_party/flash_linear_attention/ops/kda.py:1019-1123 +# @triton.jit chunk_gla_fwd_kernel_o (driver chunk_gla_fwd_o_gk :1126) +# (upstream flash-linear-attention; MIT). Kernel BODY byte-for-byte FLA. AOT +# adaptations only: +# (1) @triton.heuristics / @triton.autotune removed — IS_VARLEN and dims +# (H,K,V,BT,BK,BV) PINNED via the compile SIGNATURE (see §17). +# (2) runtime `scale` REMOVED and pinned to K**-0.5 (the chunk_kda scale, +# kda.py:1472 `scale = k.shape[-1]**-0.5`; K=128). Triton AOT cannot take a +# reliable fp32 scalar (see chunk_o.py note 3); the KDA call site always +# passes scale == K**-0.5. +# (3) trailing runtime scalar `NT` (= total chunks) appended as the grid-y +# carrier: FLA grid is (cdiv(V,BV), NT, B*H); B*H == H (varlen B=1) is the +# constexpr grid-z. `NT` is unused by the body (dead arg). +# +# Pinned for Kimi-Linear KDA: H=32, K=128, V=128, BT=64, BK=64 (=> 2 K-iters), +# BV=64 (=> gx=cdiv(V,BV)=2), IS_VARLEN=1. +# Buffers (mirror FLA dtypes): q=[T,H,K] bf16; v(=v_new)=[T,H,V] bf16; +# g(=gk cumulative, exp2-space)=[T,H,K] fp32; h(=hstate snapshot)=[NT,H,V,K] bf16; +# A(=Aqk)=[T,H,BT] fp32; o(=out, reuses v buffer)=[T,H,V] bf16; +# cu_seqlens=[N+1] i32; chunk_indices=[NT,2] i32. +import triton +import triton.language as tl + +exp = tl.exp +exp2 = tl.exp2 + + +@triton.jit(do_not_specialize=["T", "NT"]) +def chunk_gla_fwd_kernel_o( + q, + v, + g, + h, + o, + A, + cu_seqlens, + chunk_indices, + T, + NT, # AOT grid-carrier (= total chunks): grid-y extent; unused by the body. + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + # scale = K**-0.5 (chunk_kda q-scale), pinned since Triton AOT can't take an + # fp32 scalar arg reliably; K == 128 for the Kimi KDA shape (see header note 2). + scale = K ** -0.5 + i_v, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_b, i_h = i_bh // H, i_bh % H + if IS_VARLEN: + i_tg = i_t + i_n, i_t = ( + tl.load(chunk_indices + i_t * 2).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + ) + bos, eos = ( + tl.load(cu_seqlens + i_n).to(tl.int32), + tl.load(cu_seqlens + i_n + 1).to(tl.int32), + ) + T = eos - bos + NT = tl.cdiv(T, BT) + else: + NT = tl.cdiv(T, BT) + i_tg = i_b * NT + i_t + bos, eos = i_b * T, i_b * T + T + + m_s = tl.arange(0, BT)[:, None] >= tl.arange(0, BT)[None, :] + + b_o = tl.zeros([BT, BV], dtype=tl.float32) + for i_k in range(tl.cdiv(K, BK)): + p_q = tl.make_block_ptr( + q + (bos * H + i_h) * K, + (T, K), + (H * K, 1), + (i_t * BT, i_k * BK), + (BT, BK), + (1, 0), + ) + p_g = tl.make_block_ptr( + g + (bos * H + i_h) * K, + (T, K), + (H * K, 1), + (i_t * BT, i_k * BK), + (BT, BK), + (1, 0), + ) + p_h = tl.make_block_ptr( + h + (i_tg * H + i_h) * K * V, + (V, K), + (K, 1), + (i_v * BV, i_k * BK), + (BV, BK), + (1, 0), + ) + + # [BT, BK] + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_q = (b_q * scale).to(b_q.dtype) + # [BT, BK] + b_g = tl.load(p_g, boundary_check=(0, 1)) + # [BT, BK] + b_qg = (b_q * exp2(b_g)).to(b_q.dtype) + # [BV, BK] + b_h = tl.load(p_h, boundary_check=(0, 1)) + # [BT, BV] + if i_k >= 0: + b_o += tl.dot(b_qg, tl.trans(b_h).to(b_qg.dtype)) + p_v = tl.make_block_ptr( + v + (bos * H + i_h) * V, + (T, V), + (H * V, 1), + (i_t * BT, i_v * BV), + (BT, BV), + (1, 0), + ) + p_o = tl.make_block_ptr( + o + (bos * H + i_h) * V, + (T, V), + (H * V, 1), + (i_t * BT, i_v * BV), + (BT, BV), + (1, 0), + ) + p_A = tl.make_block_ptr( + A + (bos * H + i_h) * BT, (T, BT), (H * BT, 1), (i_t * BT, 0), (BT, BT), (1, 0) + ) + # [BT, BV] + b_v = tl.load(p_v, boundary_check=(0, 1)) + # [BT, BT] + b_A = tl.load(p_A, boundary_check=(0, 1)) + b_A = tl.where(m_s, b_A, 0.0).to(b_v.dtype) + b_o += tl.dot(b_A, b_v, allow_tf32=False) + tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) diff --git a/.agents/specs/kda-chunk-aot/chunk_kda_kkt.py b/.agents/specs/kda-chunk-aot/chunk_kda_kkt.py new file mode 100644 index 000000000..d8f76ae99 --- /dev/null +++ b/.agents/specs/kda-chunk-aot/chunk_kda_kkt.py @@ -0,0 +1,229 @@ +# triton_kernels/chunk_kda_kkt.py (STAGED — see .agents/specs/kimi-linear.md §17) +# +# KDA chunk-prefill kernels 2/5 + 3/5: the per-K-channel-gated scaled dot K@Kᵀ +# (and q@kᵀ = Aqk), split intra-sub-inter + intra-sub-intra exactly as FLA. These +# produce the strictly-lower-tri A (which solve_tril inverts) and the causal Aqk +# (which chunk_gla_o consumes). STAGED authored harness bodies (Phase-1 spike); +# Phase-2 moves the file to triton_kernels/ + regenerates both sm_121a cubins. +# +# Ported VERBATIM FROM (vLLM oracle @ pin 555967922): +# vllm/third_party/flash_linear_attention/ops/kda.py +# :521-618 @triton.jit chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_inter +# :627-715 @triton.jit chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_intra +# (driver chunk_kda_scaled_dot_kkt_fwd :717; upstream flash-linear-attention; +# MIT). Kernel BODIES byte-for-byte FLA. AOT adaptations only: +# (1) @triton.heuristics / @triton.autotune removed — IS_VARLEN and dims +# (H,K,BT,BC,BK,NC) PINNED via each compile SIGNATURE (see §17). +# (2) runtime `scale` REMOVED and pinned to K**-0.5 (chunk_kda scale, +# kda.py:1472; K=128); Triton AOT cannot take a reliable fp32 scalar. +# (3) trailing runtime scalar `NT` appended as the grid-x carrier. FLA grids: +# inter (NT, NC*NC, B*H); intra (NT, NC, B*H). B*H == H (varlen B=1) is the +# constexpr grid-z; NC*NC=16 / NC=4 are the constexpr grid-y. `NT` unused. +# +# Pinned for Kimi-Linear KDA: H=32, K=128, BT=64, BC=16, NC=4. +# inter: BK=64 (=> 2 K-iters), grid (NT, 16, 32). +# intra: BK=128 (= max(next_pow2(K),16), driver :758), grid (NT, 4, 32). +# Buffers (mirror FLA dtypes): q=[T,H,K] bf16; k=[T,H,K] bf16; g(=gk cumulative, +# exp2-space)=[T,H,K] fp32; beta=[T,H] bf16; A=[T,H,BT] fp32 (out); +# Aqk=[T,H,BT] fp32 (out); cu_seqlens=[N+1] i32; chunk_indices=[NT,2] i32. +import triton +import triton.language as tl + +exp = tl.exp +exp2 = tl.exp2 + + +@triton.jit(do_not_specialize=["T", "NT"]) +def chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_inter( + q, + k, + g, + beta, + A, + Aqk, + cu_seqlens, + chunk_indices, + T, + NT, # AOT grid-carrier (grid-x extent = total chunks); unused by the body. + H: tl.constexpr, + K: tl.constexpr, + BT: tl.constexpr, + BC: tl.constexpr, + BK: tl.constexpr, + NC: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + # scale = K**-0.5 (chunk_kda scale), pinned (Triton AOT fp32-scalar limitation). + scale = K ** -0.5 + i_t, i_c, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_b, i_h = i_bh // H, i_bh % H + i_i, i_j = i_c // NC, i_c % NC + if IS_VARLEN: + i_n, i_t = ( + tl.load(chunk_indices + i_t * 2).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + ) + bos, eos = ( + tl.load(cu_seqlens + i_n).to(tl.int32), + tl.load(cu_seqlens + i_n + 1).to(tl.int32), + ) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + + if i_t * BT + i_i * BC >= T: + return + if i_i <= i_j: + return + + q += (bos * H + i_h) * K + k += (bos * H + i_h) * K + g += (bos * H + i_h) * K + A += (bos * H + i_h) * BT + Aqk += (bos * H + i_h) * BT + + p_b = tl.make_block_ptr( + beta + bos * H + i_h, (T,), (H,), (i_t * BT + i_i * BC,), (BC,), (0,) + ) + b_b = tl.load(p_b, boundary_check=(0,)) + + b_A = tl.zeros([BC, BC], dtype=tl.float32) + b_Aqk = tl.zeros([BC, BC], dtype=tl.float32) + for i_k in range(tl.cdiv(K, BK)): + p_q = tl.make_block_ptr( + q, (T, K), (H * K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0) + ) + p_k = tl.make_block_ptr( + k, (T, K), (H * K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0) + ) + p_g = tl.make_block_ptr( + g, (T, K), (H * K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0) + ) + b_kt = tl.make_block_ptr( + k, (K, T), (1, H * K), (i_k * BK, i_t * BT + i_j * BC), (BK, BC), (0, 1) + ) + p_gk = tl.make_block_ptr( + g, (K, T), (1, H * K), (i_k * BK, i_t * BT + i_j * BC), (BK, BC), (0, 1) + ) + + o_k = i_k * BK + tl.arange(0, BK) + m_k = o_k < K + # [BK,] + b_gn = tl.load(g + (i_t * BT + i_i * BC) * H * K + o_k, mask=m_k, other=0) + # [BC, BK] + b_g = tl.load(p_g, boundary_check=(0, 1)) + b_k = tl.load(p_k, boundary_check=(0, 1)) * exp2(b_g - b_gn[None, :]) + # [BK, BC] + b_gk = tl.load(p_gk, boundary_check=(0, 1)) + b_kt = tl.load(b_kt, boundary_check=(0, 1)) + # [BC, BC] + b_ktg = b_kt * exp2(b_gn[:, None] - b_gk) + b_A += tl.dot(b_k, b_ktg) + + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_qg = b_q * exp2(b_g - b_gn[None, :]) * scale + b_Aqk += tl.dot(b_qg, b_ktg) + + b_A *= b_b[:, None] + + p_A = tl.make_block_ptr( + A, (T, BT), (H * BT, 1), (i_t * BT + i_i * BC, i_j * BC), (BC, BC), (1, 0) + ) + tl.store(p_A, b_A.to(A.dtype.element_ty), boundary_check=(0, 1)) + p_Aqk = tl.make_block_ptr( + Aqk, (T, BT), (H * BT, 1), (i_t * BT + i_i * BC, i_j * BC), (BC, BC), (1, 0) + ) + tl.store(p_Aqk, b_Aqk.to(Aqk.dtype.element_ty), boundary_check=(0, 1)) + + +@triton.jit(do_not_specialize=["T", "NT"]) +def chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_intra( + q, + k, + g, + beta, + A, + Aqk, + cu_seqlens, + chunk_indices, + T, + NT, # AOT grid-carrier (grid-x extent = total chunks); unused by the body. + H: tl.constexpr, + K: tl.constexpr, + BT: tl.constexpr, + BC: tl.constexpr, + BK: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + # scale = K**-0.5 (chunk_kda scale), pinned (Triton AOT fp32-scalar limitation). + scale = K ** -0.5 + i_t, i_i, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_b, i_h = i_bh // H, i_bh % H + if IS_VARLEN: + i_n, i_t = ( + tl.load(chunk_indices + i_t * 2).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + ) + bos, eos = ( + tl.load(cu_seqlens + i_n).to(tl.int32), + tl.load(cu_seqlens + i_n + 1).to(tl.int32), + ) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + + if i_t * BT + i_i * BC >= T: + return + + o_i = tl.arange(0, BC) + o_k = tl.arange(0, BK) + m_k = o_k < K + m_A = (i_t * BT + i_i * BC + o_i) < T + o_A = (bos + i_t * BT + i_i * BC + o_i) * H * BT + i_h * BT + i_i * BC + + p_q = tl.make_block_ptr( + q + (bos * H + i_h) * K, + (T, K), + (H * K, 1), + (i_t * BT + i_i * BC, 0), + (BC, BK), + (1, 0), + ) + p_k = tl.make_block_ptr( + k + (bos * H + i_h) * K, + (T, K), + (H * K, 1), + (i_t * BT + i_i * BC, 0), + (BC, BK), + (1, 0), + ) + p_g = tl.make_block_ptr( + g + (bos * H + i_h) * K, + (T, K), + (H * K, 1), + (i_t * BT + i_i * BC, 0), + (BC, BK), + (1, 0), + ) + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_g = tl.load(p_g, boundary_check=(0, 1)) + + p_b = beta + (bos + i_t * BT + i_i * BC + o_i) * H + i_h + b_k = b_k * tl.load(p_b, mask=m_A, other=0)[:, None] + + p_kt = k + (bos + i_t * BT + i_i * BC) * H * K + i_h * K + o_k + p_gk = g + (bos + i_t * BT + i_i * BC) * H * K + i_h * K + o_k + + for j in range(0, min(BC, T - i_t * BT - i_i * BC)): + b_kt = tl.load(p_kt, mask=m_k, other=0).to(tl.float32) + b_gk = tl.load(p_gk, mask=m_k, other=0).to(tl.float32) + b_ktg = b_kt[None, :] * exp2(b_g - b_gk[None, :]) + b_A = tl.sum(b_k * b_ktg, 1) + b_A = tl.where(o_i > j, b_A, 0.0) + b_Aqk = tl.sum(b_q * b_ktg, 1) + b_Aqk = tl.where(o_i >= j, b_Aqk * scale, 0.0) + tl.store(A + o_A + j, b_A, mask=m_A) + tl.store(Aqk + o_A + j, b_Aqk, mask=m_A) + p_kt += H * K + p_gk += H * K diff --git a/.agents/specs/kda-chunk-aot/kda_gate_cumsum.py b/.agents/specs/kda-chunk-aot/kda_gate_cumsum.py new file mode 100644 index 000000000..af4b869a2 --- /dev/null +++ b/.agents/specs/kda-chunk-aot/kda_gate_cumsum.py @@ -0,0 +1,113 @@ +# triton_kernels/kda_gate_cumsum.py (STAGED — see .agents/specs/kimi-linear.md §17) +# +# KDA chunk-prefill kernel 1/5: the fused decay-gate + chunk-local cumsum. +# This is a STAGED authored harness body (Phase-1 spike). Phase-2 moves it to +# triton_kernels/ and regenerates the sm_121a cubin (scripts/regen-triton-aot.sh) +# together with the cmake/TritonAOTKernels.cmake declaration in §17. +# +# Ported VERBATIM FROM (vLLM oracle @ pin 555967922): +# vllm/third_party/flash_linear_attention/ops/kda.py:1182-1254 +# @triton.jit kda_gate_cumsum_fwd_kernel (driver fused_kda_gate_chunk_cumsum :1257) +# (upstream flash-linear-attention; MIT). The kernel BODY is byte-for-byte the +# FLA source. AOT adaptations only: +# (1) @triton.heuristics / @triton.autotune removed — HAS_BIAS/IS_VARLEN and +# dims (H,D,BT,BD) PINNED per-shape via the compile SIGNATURE (see §17). +# (2) the runtime scalars cumsum_scale/beta/threshold are REMOVED and pinned +# as module literals: the KDA driver ALWAYS passes cumsum_scale=RCP_LN2 +# (natural-log -> log2 fold, kda.py:1295), softplus beta=1.0 and +# threshold=20.0 (fused_kda_gate_chunk_cumsum defaults, kda.py:1261-1262). +# Triton's AOT launcher mis-packs an fp32 scalar as an 8-byte double, so a +# runtime float arg is unusable (same reason chunk_o.py pins `scale`). +# (3) trailing runtime scalar `NT` (= total chunks) appended as the grid-y +# carrier; B*H == H (varlen packing has B=1) is the constexpr grid-z. +# `NT` is unused by the body (dead arg). +# +# Pinned for Kimi-Linear KDA: H=32, D=128 (head_k_dim), BT=64, BD=64 (=> gx=2), +# HAS_BIAS=1 (g_bias = dt_bias[H*D]), IS_VARLEN=1. +# Buffers (mirror FLA dtypes): g(=raw_g projection)=[T,H,D] bf16; +# A(=A_log)=[H] fp32; y(=g cumulative, exp2-space)=[T,H,D] fp32; +# g_bias(=dt_bias)=[H*D] fp32; cu_seqlens=[N+1] i32; chunk_indices=[NT,2] i32. +import triton +import triton.language as tl + +# fla/ops/op.py: exp/log (FLA_USE_FAST_OPS=0 default -> tl.exp / tl.log). +exp = tl.exp +log = tl.log + +# fold + softplus constants pinned per the KDA driver (see header note 2). +RCP_LN2 = 1.4426950408889634 # 1 / ln(2) +SOFTPLUS_BETA = 1.0 +SOFTPLUS_THRESHOLD = 20.0 + + +@triton.jit(do_not_specialize=["T", "NT"]) +def kda_gate_cumsum_fwd_kernel( + g, + A, + y, + g_bias, + cu_seqlens, + chunk_indices, + T, + NT, # AOT grid-carrier (= total chunks): grid-y extent; unused by the body. + H: tl.constexpr, + D: tl.constexpr, + BT: tl.constexpr, + BD: tl.constexpr, + HAS_BIAS: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + beta = SOFTPLUS_BETA + threshold = SOFTPLUS_THRESHOLD + cumsum_scale = RCP_LN2 + i_d, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_b, i_h = i_bh // H, i_bh % H + if IS_VARLEN: + i_n, i_t = ( + tl.load(chunk_indices + i_t * 2).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + ) + bos, eos = ( + tl.load(cu_seqlens + i_n).to(tl.int32), + tl.load(cu_seqlens + i_n + 1).to(tl.int32), + ) + T = eos - bos + else: + bos = i_b * T + + p_g = tl.make_block_ptr( + g + (bos * H + i_h) * D, + (T, D), + (H * D, 1), + (i_t * BT, i_d * BD), + (BT, BD), + (1, 0), + ) + p_y = tl.make_block_ptr( + y + (bos * H + i_h) * D, + (T, D), + (H * D, 1), + (i_t * BT, i_d * BD), + (BT, BD), + (1, 0), + ) + + b_g = tl.load(p_g, boundary_check=(0, 1)).to(tl.float32) + if HAS_BIAS: + o_d = i_d * BD + tl.arange(0, BD) + b_bias = tl.load(g_bias + i_h * D + o_d, mask=o_d < D, other=0.0).to(tl.float32) + b_g = b_g + b_bias[None, :] + + b_a = -tl.exp(tl.load(A + i_h).to(tl.float32)) + b_g_scaled = b_g * beta + b_softplus = tl.where( + b_g_scaled > threshold, + b_g, + (1.0 / beta) * log(1.0 + tl.exp(b_g_scaled)), + ) + b_gate = b_a * b_softplus + + o_t = tl.arange(0, BT) + m_cumsum = tl.where(o_t[:, None] >= o_t[None, :], 1.0, 0.0) + b_y = tl.dot(m_cumsum, b_gate, allow_tf32=False) * cumsum_scale + tl.store(p_y, b_y.to(p_y.dtype.element_ty), boundary_check=(0, 1)) diff --git a/.agents/specs/kda-chunk-aot/recompute_w_u_kda.py b/.agents/specs/kda-chunk-aot/recompute_w_u_kda.py new file mode 100644 index 000000000..4166208e9 --- /dev/null +++ b/.agents/specs/kda-chunk-aot/recompute_w_u_kda.py @@ -0,0 +1,179 @@ +# triton_kernels/recompute_w_u_kda.py (STAGED — see .agents/specs/kimi-linear.md §17) +# +# KDA chunk-prefill kernel 4/5: recompute W and U (the WY-representation "apply"), +# the KDA PER-K-CHANNEL variant. Distinct from the GDN wy_fast.py kernel: KDA +# applies exp2(b_gk) per-K-channel and stores kg = k*exp2(gn - gk) (the STORE_KG +# branch). STAGED authored harness body (Phase-1 spike); Phase-2 moves it to +# triton_kernels/ + regenerates the sm_121a cubin. +# +# Ported VERBATIM FROM (vLLM oracle @ pin 555967922): +# vllm/third_party/flash_linear_attention/ops/kda.py:817-957 +# @triton.jit recompute_w_u_fwd_kernel (driver recompute_w_u_fwd :960) +# (upstream flash-linear-attention; MIT). Kernel BODY byte-for-byte FLA. AOT +# adaptations only: +# (1) @triton.heuristics / @triton.autotune removed — STORE_QG/STORE_KG/ +# IS_VARLEN and dims (H,K,V,BT,BK,BV) PINNED via the compile SIGNATURE. +# (2) DOT_PRECISION REMOVED and pinned to "ieee" (driver recompute_w_u_fwd +# always passes DOT_PRECISION="ieee", kda.py:1002); it is a compile-time +# constexpr string, baked as a module literal. +# (3) trailing runtime scalar `NT` appended as the grid-x carrier (grid is +# (NT, B*H); B*H == H is the constexpr grid-y). `NT` unused by the body. +# +# Pinned for Kimi-Linear KDA: H=32, K=128, V=128, BT=64, BK=64, BV=64, +# STORE_QG=0 (q/qg dead), STORE_KG=1 (kg written), IS_VARLEN=1. Grid (NT, 32, 1). +# Buffers (mirror FLA dtypes): k=[T,H,K] bf16; v=[T,H,V] bf16; beta=[T,H] bf16; +# w=[T,H,K] bf16 (out); u=[T,H,V] bf16 (out); A(=solve_tril inverse)=[T,H,BT] bf16; +# gk(=g cumulative, exp2-space)=[T,H,K] fp32; kg=[T,H,K] bf16 (out); q/qg dead +# (STORE_QG=0); cu_seqlens=[N+1] i32; chunk_indices=[NT,2] i32. +import triton +import triton.language as tl + +exp = tl.exp +exp2 = tl.exp2 + +# driver always passes DOT_PRECISION="ieee" (kda.py:1002); see header note 2. +DOT_PRECISION: tl.constexpr = "ieee" + + +@triton.jit(do_not_specialize=["T", "NT"]) +def recompute_w_u_fwd_kernel( + q, + k, + qg, + kg, + v, + beta, + w, + u, + A, + gk, + cu_seqlens, + chunk_indices, + T, + NT, # AOT grid-carrier (grid-x extent = total chunks); unused by the body. + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + STORE_QG: tl.constexpr, + STORE_KG: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_b, i_h = i_bh // H, i_bh % H + if IS_VARLEN: + i_n, i_t = ( + tl.load(chunk_indices + i_t * 2).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + ) + bos, eos = ( + tl.load(cu_seqlens + i_n).to(tl.int32), + tl.load(cu_seqlens + i_n + 1).to(tl.int32), + ) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + p_b = tl.make_block_ptr(beta + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) + b_b = tl.load(p_b, boundary_check=(0,)) + + p_A = tl.make_block_ptr( + A + (bos * H + i_h) * BT, (T, BT), (H * BT, 1), (i_t * BT, 0), (BT, BT), (1, 0) + ) + b_A = tl.load(p_A, boundary_check=(0, 1)) + + for i_v in range(tl.cdiv(V, BV)): + p_v = tl.make_block_ptr( + v + (bos * H + i_h) * V, + (T, V), + (H * V, 1), + (i_t * BT, i_v * BV), + (BT, BV), + (1, 0), + ) + p_u = tl.make_block_ptr( + u + (bos * H + i_h) * V, + (T, V), + (H * V, 1), + (i_t * BT, i_v * BV), + (BT, BV), + (1, 0), + ) + b_v = tl.load(p_v, boundary_check=(0, 1)) + b_vb = (b_v * b_b[:, None]).to(b_v.dtype) + b_u = tl.dot(b_A, b_vb, input_precision=DOT_PRECISION) + tl.store(p_u, b_u.to(p_u.dtype.element_ty), boundary_check=(0, 1)) + + for i_k in range(tl.cdiv(K, BK)): + p_w = tl.make_block_ptr( + w + (bos * H + i_h) * K, + (T, K), + (H * K, 1), + (i_t * BT, i_k * BK), + (BT, BK), + (1, 0), + ) + p_k = tl.make_block_ptr( + k + (bos * H + i_h) * K, + (T, K), + (H * K, 1), + (i_t * BT, i_k * BK), + (BT, BK), + (1, 0), + ) + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_kb = b_k * b_b[:, None] + + p_gk = tl.make_block_ptr( + gk + (bos * H + i_h) * K, + (T, K), + (H * K, 1), + (i_t * BT, i_k * BK), + (BT, BK), + (1, 0), + ) + b_gk = tl.load(p_gk, boundary_check=(0, 1)) + b_kb *= exp2(b_gk) + if STORE_QG: + p_q = tl.make_block_ptr( + q + (bos * H + i_h) * K, + (T, K), + (H * K, 1), + (i_t * BT, i_k * BK), + (BT, BK), + (1, 0), + ) + p_qg = tl.make_block_ptr( + qg + (bos * H + i_h) * K, + (T, K), + (H * K, 1), + (i_t * BT, i_k * BK), + (BT, BK), + (1, 0), + ) + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_qg = b_q * exp2(b_gk) + tl.store(p_qg, b_qg.to(p_qg.dtype.element_ty), boundary_check=(0, 1)) + if STORE_KG: + last_idx = min(i_t * BT + BT, T) - 1 + + o_k = i_k * BK + tl.arange(0, BK) + m_k = o_k < K + b_gn = tl.load( + gk + ((bos + last_idx) * H + i_h) * K + o_k, mask=m_k, other=0.0 + ) + b_kg = b_k * exp2(b_gn - b_gk) + + p_kg = tl.make_block_ptr( + kg + (bos * H + i_h) * K, + (T, K), + (H * K, 1), + (i_t * BT, i_k * BK), + (BT, BK), + (1, 0), + ) + tl.store(p_kg, b_kg.to(p_kg.dtype.element_ty), boundary_check=(0, 1)) + + b_w = tl.dot(b_A, b_kb.to(b_k.dtype)) + tl.store(p_w, b_w.to(p_w.dtype.element_ty), boundary_check=(0, 1)) diff --git a/.agents/specs/kimi-linear.md b/.agents/specs/kimi-linear.md index 741d5a1fa..f494a1eed 100644 --- a/.agents/specs/kimi-linear.md +++ b/.agents/specs/kimi-linear.md @@ -972,6 +972,146 @@ brick, not a one-shot; recorded as the named follow-on. --- +## 17. chunk_kda PREFILL AOT PORT — kernel set + pinned-config record + regen recipe (Phase-1 spike, 2026-08-07, `row/KIMI-CHUNK-KDA-AOT`) + +The §15/§16 STRICT residual (c) — "vLLM processes the PROMPT with the CHUNKED `chunk_kda` +kernel, we still run the RECURRENT form; a different reduction order coin-flips the p7 +near-tie" — is here scoped, grounded, and DE-RISKED to the point of mechanical execution. +This section is the **AOT regen recipe + pinned-config record** the mission asks for. +The authored Triton harness bodies are STAGED in +[`.agents/specs/kda-chunk-aot/`](kda-chunk-aot/) (CI-safe: the drift check globs +`triton_kernels/*.py` non-recursively, so a staged sibling directory does not gate). +**Phase-2** moves them into `triton_kernels/`, adds the declarations below to +`cmake/TritonAOTKernels.cmake`, regenerates the sm_121a cubins +(`scripts/regen-triton-aot.sh`), wires the `vt::KdaChunkPrefill` op, and runs the gates. + +### 17.1 The EXACT forward-only kernel set (`chunk_kda_with_fused_gate` → `_fwd`, file:line @ 555967922) +The prefill driver is `kimi_gdn_linear_attn.py:141` → `chunk_kda_with_fused_gate` +(`kda.py:1492`) → `chunk_kda_with_fused_gate_fwd` (`:1416`) → +`_chunk_kda_fwd_with_cumulative_g` (`:1306`). The kernel chain, in launch order: + +| # | Step | FLA kernel(s) `kda.py:line` | New/Reuse | +|---|---|---|---| +| 1 | fused decay-gate + chunk-local cumsum·RCP_LN2 | `kda_gate_cumsum_fwd_kernel` `:1182-1254` | **NEW** | +| 2 | per-channel-gated K·Kᵀ + q·kᵀ (A, Aqk), inter | `chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_inter` `:521-618` | **NEW** | +| 3 | …same, intra | `chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_intra` `:627-715` | **NEW** | +| 4 | invert the strictly-lower-tri A (WY solve) | `solve_tril` / `merge_16x16_to_64x64_inverse_kernel` | **REUSE `gdn_tril_h32`** | +| 5 | recompute W, U (+ kg) per-K-channel | `recompute_w_u_fwd_kernel` `:817-957` | **NEW** (KDA per-channel, ≠ GDN `wy_fast.py`) | +| 6 | chunked hidden-state scan (h, v_new, final) | `chunk_gated_delta_rule_fwd_h` (`chunk_delta_h.py`, imported `:19`) | **REUSE `chunk_delta_h.py`, NEW pin** | +| 7 | GLA-style output with per-K gk decay | `chunk_gla_fwd_kernel_o` `:1019-1123` | **NEW** | + +So **5 genuinely-new Triton kernels** (steps 1,2,3,5,7) + **1 new PIN of an existing .py** +(step 6: `chunk_delta_h.py` recompiled with `USE_GK=1, USE_EXP2=1, USE_G=0`, ≠ the GDN +`gdn_deltah` pin `USE_G=1, USE_GK=0, USE_EXP2=0`) + **1 pure reuse** (step 4: +`gdn_tril_h32`, byte-identical signature). The decode path is UNCHANGED — it stays the +#104 recurrent `vt::KdaGatedDeltaRule` (mirroring vLLM's own prefill=chunk / decode=recurrent +split). Backward kernels are NOT owed (forward-only inference). + +### 17.2 Pinned-config record (Kimi KDA shapes: H=32, Hg=32, K=V=128, BT=64, BC=16, NC=4) +Autotune metaparams cannot be expressed in AOT, so each is PINNED. `num_warps`/`num_stages` +are **correctness-invariant** (they change tiling/pipelining, not the numeric result), pinned +mirroring the GDN precedent; the shape pins (BK/BV/BD/BC/NC) follow FLA's driver-fixed values +and heuristic lists. Dtypes MIRROR FLA's exact buffer choices (bf16 activations/intermediates; +fp32 for gk-cumulative / A / Aqk / recurrent-state) — **Phase-2 confirms each against the +`vt::KdaChunkPrefill` buffer contract before regen**. + +| base | staged .py | kernel | BK | BV | BD | warps | stages | grid | +|---|---|---|---|---|---|---|---|---| +| `kda_gate_cumsum` | kda_gate_cumsum.py | `kda_gate_cumsum_fwd_kernel` | — | — | 64 | 4 | 2 | `2,NT,32` | +| `kda_kkt_inter` | chunk_kda_kkt.py | `…intra_sub_inter` | 64 | — | — | 4 | 3 | `NT,16,32` | +| `kda_kkt_intra` | chunk_kda_kkt.py | `…intra_sub_intra` | 128 | — | — | 4 | 2 | `NT,4,32` | +| `kda_wu` | recompute_w_u_kda.py | `recompute_w_u_fwd_kernel` | 64 | 64 | — | 4 | 3 | `NT,32,1` | +| `kda_deltah_h32` | chunk_delta_h.py (reuse) | `chunk_gated_delta_rule_fwd_kernel_h_blockdim64` | — | 64 | — | 4 | 3 | `2,NH,1` | +| `kda_gla_o` | chunk_gla_o.py | `chunk_gla_fwd_kernel_o` | 64 | 64 | — | 4 | 3 | `NT,2,32`† | + +†`kda_gla_o` grid is `(cdiv(V,BV), NT, H)` = `(2, NT, 32)`; expressed as `2,NT,32` with `NT` +the trailing carrier. **Scalar-constant pins baked as literals** (Triton AOT mis-packs fp32 +scalars — see `chunk_o.py` note 3): `scale = K**-0.5` (kkt ×2, gla_o); softplus `beta=1.0`, +`threshold=20.0`, `cumsum_scale=RCP_LN2` (gate_cumsum); `DOT_PRECISION="ieee"` (wu). + +### 17.3 The regen recipe — `cmake/TritonAOTKernels.cmake` declarations to ADD (Phase-2) +Signatures use the vendored `*dtype:align` / scalar / constexpr form. `NT`/`NH` are the +trailing grid carriers. Insert inside `vllm_triton_aot_declare_all()` after the GDN WY block: +``` +# KDA chunk-prefill family (Kimi-Linear; H=32). Mirrors the GDN WY pins. +_vllm_triton_aot_declare(kda_gate_cumsum kda_gate_cumsum.py kda_gate_cumsum_fwd_kernel 4 2 + "2,NT,32" + "*bf16:16, *fp32:16, *fp32:16, *fp32:16, *i32:16, *i32:16, i32, i32, 32, 128, 64, 64, 1, 1") +_vllm_triton_aot_declare(kda_kkt_inter chunk_kda_kkt.py chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_inter 4 3 + "NT,16,32" + "*bf16:16, *bf16:16, *fp32:16, *bf16:16, *fp32:16, *fp32:16, *i32:16, *i32:16, i32, i32, 32, 128, 64, 16, 64, 4, 1") +_vllm_triton_aot_declare(kda_kkt_intra chunk_kda_kkt.py chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_intra 4 2 + "NT,4,32" + "*bf16:16, *bf16:16, *fp32:16, *bf16:16, *fp32:16, *fp32:16, *i32:16, *i32:16, i32, i32, 32, 128, 64, 16, 128, 1") +_vllm_triton_aot_declare(kda_wu recompute_w_u_kda.py recompute_w_u_fwd_kernel 4 3 + "NT,32,1" + "*bf16:16, *bf16:16, *bf16:16, *bf16:16, *bf16:16, *bf16:16, *bf16:16, *bf16:16, *bf16:16, *fp32:16, *i32:16, *i32:16, i32, i32, 32, 128, 128, 64, 64, 64, 0, 1, 1") +_vllm_triton_aot_declare(kda_deltah_h32 chunk_delta_h.py chunk_gated_delta_rule_fwd_kernel_h_blockdim64 4 3 + "2,NH,1" + "*bf16:16, *bf16:16, *bf16:16, *bf16:16, *fp32, *fp32:16, *bf16:16, *fp32:16, *fp32:16, *i32:16, *i32:16, i32, i32, 32, 32, 128, 128, 64, 64, 0, 1, 1, 1, 1, 1, 1") +_vllm_triton_aot_declare(kda_gla_o chunk_gla_o.py chunk_gla_fwd_kernel_o 4 3 + "2,NT,32" + "*bf16:16, *bf16:16, *fp32:16, *bf16:16, *bf16:16, *fp32:16, *i32:16, *i32:16, i32, i32, 32, 128, 128, 64, 64, 64, 1") +``` +`kda_deltah_h32` note: the arg order is `k,v,w,v_new,g,gk,h,h0,ht,cu_seqlens,chunk_offsets, +T,NH,H,Hg,K,V,BT,BV,USE_G,USE_GK,USE_INITIAL_STATE,STORE_FINAL_STATE,SAVE_NEW_VALUE, +IS_VARLEN,USE_EXP2` — vs the GDN `gdn_deltah` the alignment marker MOVES from `g`(now dead, +`*fp32`) to `gk`(now used, `*fp32:16`), `Hg` flips 16→32 (KDA has no delta-rule GQA), and the +flag triple flips to `USE_G=0,USE_GK=1,…,USE_EXP2=1`. Steps 4 (`gdn_tril_h32`) and 6 +(`chunk_delta_h.py`) need NO new .py file — reuse the vendored source. + +### 17.4 The `vt::KdaChunkPrefill` op design (Phase-2 wiring, mirrors `cuda_gdn.cu` GdnPrefill) +A new additive op routing PROMPT-length KDA (`query_len == key_len`, the prefill step) through +the six cubins; decode (`query_len==1`) stays `vt::KdaGatedDeltaRule` (#104). Orchestration +(exactly `_chunk_kda_fwd_with_cumulative_g`): allocate per-`(T,H,·)` scratch — +`g_cum[T,H,128] f32`, `A/Aqk[T,H,64] f32`, `A_inv[T,H,64] bf16`, `w[T,H,128] bf16`, +`u[T,H,128] bf16`, `kg[T,H,128] bf16`, `h[NT,H,128,128] bf16`, `v_new[T,H,128] bf16`, then +launch (1) gate_cumsum → g_cum; (2)+(3) kkt → A,Aqk; (4) `gdn_tril_h32_default` → A_inv; +(5) wu(A_inv,gk=g_cum) → w,u,kg; (6) `kda_deltah_h32_default`(k=kg,w,u,gk=g_cum,h0=zeros, +ht=state) → h,v_new,final_state; (7) gla_o(q,v_new,g=g_cum,A=Aqk,h) → out. `cu_seqlens=[0,T]`, +`chunk_indices`/`chunk_offsets` from `prepare_chunk_indices`. Loader/launcher per the vendored +`std::call_once(load_gdn_*)` + `*_default(stream, …)` pattern (`cuda_gdn.cu:4594-4740`). +Dispatch guard: fire only when q/k/v are the pinned Kimi KDA geometry (H=32,K=V=128) and +`VLLM_CPP_TRITON`; else fall back to the recurrent island. The prefill/decode SPLIT mirrors +vLLM's `kimi_gdn_linear_attn.py:233-268` (decode `fused_recurrent_kda`, prefill +`chunk_kda_with_fused_gate`). + +### 17.5 Gate plan (Phase-2, RED-first) +1. **Unit** (`tests/vt/test_ops_kda_chunk_prefill.cpp`): (a) the chunk op == the recurrent + `vt::KdaGatedDeltaRule` (#104) to the documented chunked-vs-recurrent reduction-order delta + (NOT bit-exact — different order; assert the p-th token argmax stable, or an rtol band); (b) + vs the #173 host refs (`kimi_kda.cpp` `KdaDecayGateChunkCumsum` etc.) per intermediate; (c) + if feasible, an **FLA-python golden** captured on the oracle venv at Kimi shapes as the direct + oracle (the chunked result should match FLA's own `chunk_kda` output). RED-first: perturb a + pin/scale, see it fail. +2. **Full 48.9B GB10 correctness gate** (re-park worker FIRST): 128 vs the §12 STRICT golden with + `VT_KIMI_DEVICE_KDA=1` **+ chunk-prefill ON**. Target **STRICT** (close the p7 prefill-order + near-tie). +3. **Speed ladder — ours-vs-vLLM at MATCHED config (USER 2026-08-07: the success bar is MEET + vLLM SPEED; the §14/§16/#107 "HW-forced-indirect" framing is SUPERSEDED).** vLLM demonstrably + RUNS Kimi-Linear-48B on ONE GB10 — the §12 golden capture used it at `gpu_memory_utilization + =0.82`, single-seq, eager. So measure both arms at that EXACT recipe on the SAME prompts: + (a) OUR arm — steady decode tok/s + prefill TTFT with `DEVICE_KDA=1` + chunk-prefill; (b) the + **vLLM arm at the §12 launch config** (single-seq, eager, util 0.82) — steady decode tok/s + + prefill TTFT. Report the ladder as a MEASURED ours/vLLM ratio, not an indirect statement. + **OOM-REBOOT PROTOCOL (util 0.82 with 91.5 GiB weights is the tightest vLLM config ever run on + this box — treat as reboot-risk):** run the vLLM arm SEQUENTIAL after our runs; `local-ai-worker` + PARKED; `sudo drop_caches` before wall-clock; **PRE-WARM FlashInfer's autotune in a throwaway + start at TINY util FIRST** (cold autotune at high util is a recorded OOM-reboot trigger); memory + monitor MANDATORY; ONE attempt — if it OOMs, record the attempt honestly and do NOT retry at + higher risk. `flock` both GPU locks; single-load steady-state. +4. On STRICT **and** ≥ vLLM speed at matched config: default flips per parity-enablers with proofs; + model-matrix row moves. Below vLLM on any axis = an open gap (a MEASURED distance-to-bar), not done. + +### 17.6 Status +**Phase-1 (this spike) DONE**: kernel set enumerated + classified; 5 harness bodies authored +(verbatim FLA ports, AOT-adapted) + staged; pinned-config record + regen recipe committed; the +`vt::KdaChunkPrefill` op + gate plan specified. **NOT YET**: regen (Phase-2 — coupled to the +op's confirmed buffer dtypes), the C++ op, and the numeric gates. Row STAYS `ACTIVE`. No STRICT +verdict is claimed here. + +--- + ## Structured contract (machine-readable — mirrors deepseek-v4-flash.md) ## Scope diff --git a/.agents/state.md b/.agents/state.md index 95a55c958..156059a2a 100644 --- a/.agents/state.md +++ b/.agents/state.md @@ -41093,3 +41093,32 @@ is required") is now IMPLEMENTED as the additive device op `vt::KdaGatedDeltaRul (needs a decode/paged-attn op, query_len≠key_len, which `vt::Attention` cannot express; kills O(n²)). Each a substantial multi-kernel brick. Row STAYS ACTIVE. Box left clean (build tree/markers removed, memory restored to 112 GiB avail, worker restarted --restart=always, golden preserved, no reboot). + +## 2026-08-07T13:00 — Kimi-Linear chunk_kda PREFILL AOT port — Phase-1 SPIKE (kernel set + pins + recipe; no gate) (`row/KIMI-CHUNK-KDA-AOT`) + +The §16 STRICT residual (c) — vLLM runs the PROMPT through the CHUNKED `chunk_kda` kernel while we +still run the recurrent form (a different reduction order that coin-flips the p7 near-tie) — is now +SCOPED, GROUNDED and de-risked to mechanical execution. Spec §17 is the AOT regen recipe + +pinned-config record. +- **Exact forward-only kernel set** (`chunk_kda_with_fused_gate`→`_fwd`, FLA `kda.py` @ 555967922): + 5 NEW Triton kernels — `kda_gate_cumsum_fwd_kernel` `:1182`, the two `chunk_kda_scaled_dot_kkt` + sub-kernels `:521`/`:627`, `recompute_w_u_fwd_kernel` `:817` (KDA per-K-channel, ≠ GDN `wy_fast.py`), + `chunk_gla_fwd_kernel_o` `:1019`; + 1 NEW PIN of an existing .py (`chunk_delta_h.py` with + `USE_GK=1,USE_EXP2=1,USE_G=0,Hg=32` — GDN's `gdn_deltah` pin is NOT reusable); + 1 pure REUSE + (`solve_tril`→`gdn_tril_h32`, byte-identical sig). Decode STAYS #104 recurrent (vLLM's own split). +- **Delivered**: 5 harness bodies authored (verbatim FLA ports, AOT-adapted) STAGED in + `.agents/specs/kda-chunk-aot/` (CI-safe — drift globs `triton_kernels/*.py` non-recursively); the + full `_vllm_triton_aot_declare` recipe (§17.3); the `vt::KdaChunkPrefill` op design (buffer layout, + 6-launch order, dispatch guard, prefill/decode split, mirrors `cuda_gdn.cu` GdnPrefill); the + RED-first gate plan (unit vs #104 recurrent + #173 host refs + FLA-python golden; then 48.9B STRICT). +- **NOT YET (Phase-2, coupled)**: the harness signatures depend on the op's confirmed buffer dtypes, + so regen is premature until the op exists. Phase-2 = move harness → `triton_kernels/`, add the §17.3 + declarations, regen sm_121a cubins, wire `vt::KdaChunkPrefill`, run RED-first unit + FLA golden + + full 48.9B GB10 STRICT gate (`DEVICE_KDA=1` + chunk-prefill) + tok/s/TTFT ladder. Row STAYS ACTIVE. +- **USER 2026-08-07: the Kimi bar is MEET vLLM SPEED; the §14/§16 "HW-forced-indirect" framing is + SUPERSEDED.** vLLM DOES run this on one GB10 (§12 golden used util 0.82, single-seq, eager). The + Phase-2 speed ladder MUST add a vLLM arm at that EXACT recipe on the same prompts (SEQUENTIAL after + ours, worker parked, drop_caches, PRE-WARM FlashInfer autotune at tiny util first — cold autotune at + high util is an OOM-reboot trigger; ONE attempt, record honestly if it OOMs). Recorded in §17.5. +- No GPU work this session (pure authoring); box untouched, `local-ai-worker` left as-is. No STRICT + claimed. Branch `row/KIMI-CHUNK-KDA-AOT` off `origin/main` @ `ea89926f`; DRAFT PR opened. diff --git a/docs/BENCHMARKS.md b/docs/BENCHMARKS.md index a89d56213..652d24c44 100644 --- a/docs/BENCHMARKS.md +++ b/docs/BENCHMARKS.md @@ -311,7 +311,7 @@ built on it rather than keeping the flattering one. | Multimodal image, audio, video | Correctness gated, speed unmeasured | Per-modality speed grids | | `/v1/videos` OpenAI (Sora) shape | **No number owed**: a CPU serving-surface change (request aliases, the MP4 content route, and reference conditioning wiring), unit-gated only, no kernel or generation path touched | Video generation speed stays the MiniMax-H3 FP4 row below | | Qwen3-dense decode CUDA-graph | Token-exact pass, ~4.3% e2e directional | Steady-state per-step tok/s | -| Kimi-Linear-48B-A3B (KDA+MLA+MoE) | e2e RUNS (bf16-resident §13); KDA device op `vt::KdaGatedDeltaRule` GB10 **106→122/128 + 4.24 tok/s (3.1×)**, NOT STRICT, default OFF; device NoPE-MLA `VT_KIMI_DEVICE_MLA` MEASURED-NEGATIVE 122→109 (§16) | device-KDA = vLLM's actual GPU recurrence; bf16 knobs hurt (122→90); MLA vt::Attention f32-softmax ≠ FA2 order, coin-flips (breaks p3). p7 near-tie; residual = chunk_kda + paged FA2 MLA + incremental §15/§16 | +| Kimi-Linear-48B-A3B (KDA+MLA+MoE) | e2e RUNS (bf16-resident §13); KDA device op `vt::KdaGatedDeltaRule` GB10 **106→122/128 + 4.24 tok/s (3.1×)**, NOT STRICT, default OFF; NoPE-MLA `VT_KIMI_DEVICE_MLA` NEG 122→109 (§16); chunk_kda prefill AOT SPIKED (§17) | device-KDA = vLLM's actual GPU recurrence; MLA vt::Attention f32-softmax ≠ FA2 order coin-flips (§16). p7 near-tie; STRICT residual = chunk_kda prefill (5-kernel AOT recipe §17) + paged FA2 MLA + incremental | | vLLM 0.26 re-benchmark | Pending | Re-run the binding grids on the advanced pin | | MiniMax-H3 FP4 speed (W-FP4a) | **Measured GB10 (`row/H3-FP4-GPU-E2E`).** Marlin W4A16 byte-exact vs bf16; fp4 a memory win, 0.8x bf16/forward. Real-ckpt fp4-resident e2e RUNS (mp4/wav) | fp4 speed CLOSED. bf16-vs-quant A/B UNBLOCKED (the 13-shard bf16 DiT is now indexable) but NOT MEASURED: no bf16 render exists. Detail: benchmark-record + spec §8 | | MiniMax-H3 render coherence (`row/H3-RENDER-CLOSE` #77) | **CLOSED: a COHERENT scene on GB10.** #70/#74 white was wrong-PARTITION usage (t2va on the ref2va ckpt); t2va on the FL2VA GGUF renders a prompt-matched orange cat (adj-cos 0.95 vs 0.06, no patch-grid) | Verified first: t2va inputs byte-exact vs upstream; CUDA device==host at seq 1920. Follow-up `H3-TASK-PARTITION-GUARD`: the task/partition mismatch now RAISES 1:1 with `_resolve_task` (spec §8.6-8.7) | diff --git a/docs/ENVIRONMENT.md b/docs/ENVIRONMENT.md index 03f4c2c99..7e682f8c3 100644 --- a/docs/ENVIRONMENT.md +++ b/docs/ENVIRONMENT.md @@ -68,6 +68,7 @@ portable/reference path. In normal operation leave them unset. | `VT_KIMI_BF16_ISLANDS` | off (opt-in) | `=1` rounds the Kimi-Linear host-fallback island INPUTS (KDA q/k/v/g1/beta, NoPE-MLA q/kv/kpe) to bf16 (RNE) before the recurrence/softmax, toward vLLM's GDN-Triton/FA2 kernel precision. W7-speed STRICT-lever A/B (spec §14). Default OFF → byte-identical. MEASURED best config paired with `VT_KIMI_BF16_RESIDUAL` (106→120/128) | | `VT_KIMI_ISLAND_F32ACC` | off (opt-in) | `=1` computes the Kimi-Linear island recurrence/softmax in f32 accumulation (not f64). W7-speed A/B knob, **MEASURED NEGATIVE** (91–106/128; kept as a documented-negative A/B, spec §14). Default OFF → byte-identical | | `VT_KIMI_DEVICE_KDA` | off (opt-in) | `=1` runs the Kimi-Linear KDA per-k-channel gated-delta RECURRENCE through the net-new device op `vt::KdaGatedDeltaRule` (`cuda_gdn.cu` `KdaScanKernel`, the per-channel-decay `g[T,H,D]` variant of GDN's per-head-scalar decay; FLA `fused_recurrent_gated_delta_rule_fwd_kernel` IS_KDA=True) instead of the f64 host recompute. The principled path to STRICT AND the speed lever (spec §14/§15): the recurrence runs vLLM's actual f32-on-bf16 arithmetic on device rather than a host f64 recompute that is more precise than vLLM and coin-flips near-ties. Requires `VT_KIMI_DEVICE_COMPUTE=1`; the elementwise decay gate + beta stay host. Unit-gated (`test_ops_kda_recurrence`: broadcast-g == `vt::GdnPrefill` bit-identical, distinct-per-channel vs f64 ref, CPU↔CUDA parity). Default OFF (parity-enabler: flip only with the token gate green) → byte-identical | +| `VT_KIMI_DEVICE_MLA` | off (opt-in) | `=1` routes the Kimi-Linear 7 NoPE-MLA layers' causal-softmax core through the shared device op `vt::Attention` (f32 online max-subtracted softmax, value zero-padded qk_nope+qk_rope=192 vs v=128) instead of the f64 host `MlaSoftmaxIsland`. Requires `VT_KIMI_DEVICE_COMPUTE=1`. W7-speed STRICT-lever A/B (spec §16, #107), **MEASURED NEGATIVE** — `vt::Attention`'s f32 online-softmax is the right math but a DIFFERENT reduction order than vLLM's FA2, so it coin-flips near-ties (full 48.9B gate 122→109/128 AND 4.24→3.89 tok/s). Kept as a documented-MEASURED-NEGATIVE A/B knob; the STRICT path is the ACTUAL paged FA2 `mla::ForwardMlaAttentionBlock`, not this approximation. Default OFF → byte-identical | | `VT_WHISPER_ENC_EAGER` | off (flash-tiled attention on) | Set to `1` to force the naive per-key block-reduction attention in the Voxtral/Whisper audio encoder instead of the default flash-tiled kernel. Rollback / A-B knob; token-identical to the default path | | `VT_WHISPER_ENC_WARP` | off (flash-tiled attention on) | Set to `1` to force the warp-scoped online-softmax attention (`vt::AttentionDenseFast`, the pre-flash default) in the Voxtral/Whisper audio encoder instead of the default flash-tiled kernel (`vt::AttentionDenseFlash`). Rollback / A-B knob; the flash-tiled path is bit-identical to the warp path (encoder self-attention ~1.82x faster) | | `VT_WHISPER_ENC_REMARSHAL` | off (encoder weights resident) | Set to `1` to disable device-resident encoder weights and re-marshal (host f32->bf16 convert + H2D upload) all Whisper/Voxtral encoder weights on EVERY forward, restoring the pre-residency behavior. Rollback / A-B knob; byte-identical output (moves data only). Default residency uploads each encoder weight once and reuses it, removing ~648 ms of per-call host marshalling from the encoder forward | diff --git a/docs/STATUS.md b/docs/STATUS.md index cd8becfce..968e3027c 100644 --- a/docs/STATUS.md +++ b/docs/STATUS.md @@ -74,7 +74,7 @@ token-for-token correctness against the pinned oracle. | DeepSeek-V2 MLA | Correctness-complete, speed-pending | Token-exact 8/8 (DeepSeek-V2-Lite); 0.86-0.95x output rate, TTFT faster at c4/c8. A2+A5 MLA norm-rope fold default-ON (`VT_MLA_FUSED_NORM_ROPE`, bit-exact rollback, SACRED 8/8 unchanged; forensics in benchmark-record) — kimi_k3/kimi-linear inherit it | | GLM-4 dense (sandwich norms, partial rope) | Correctness-complete, speed-pending | Token-exact 16/16 (GLM-4-9B-0414); first GLM-family model; partial interleaved RoPE + Gemma2 sandwich norms + biased qkv | | GLM-4.7-Flash (MLA + GLM MoE) | Correctness-complete, speed-pending | Token-exact 8/8 (GLM-4.7-Flash, 31.2B); reuses the DeepSeek-V2 MLA stack; first e2e coverage of the q_lora query branch + noaux_tc sigmoid router with routed-scaling | -| Kimi-Linear-48B-A3B (KDA + NoPE-MLA + MoE hybrid) | e2e RUNS (bf16-resident §13); KDA device op `vt::KdaGatedDeltaRule` GB10 **122/128, 4.24 tok/s** best, NOT STRICT, default OFF. `VT_KIMI_DEVICE_MLA` (device pad-V attn) MEASURED-NEGATIVE 122→109, 4.24→3.89 (§16) | CPU+CUDA 14/14·825; p7 near-tie; STRICT residual = KDA chunked-prefill + paged FA2 MLA + incremental (§16) | +| Kimi-Linear-48B-A3B (KDA + NoPE-MLA + MoE hybrid) | e2e RUNS (§13); KDA op `vt::KdaGatedDeltaRule` GB10 **122/128, 4.24 tok/s**, NOT STRICT, OFF. `VT_KIMI_DEVICE_MLA` NEG (§16); chunk_kda prefill AOT SPIKED (§17) | CPU+CUDA 14/14·825; p7 near-tie; STRICT residual = chunk_kda prefill (5-kernel recipe §17) + paged FA2 MLA + incremental | | Gemma-3 dense (GeGLU, dual rope, sandwich norms) | Correctness-complete, speed-pending | STRICT token-exact 48/48 greedy (gemma-3-1b-it); first Gemma-family model; GeGLU (gelu_pytorch_tanh) + dual per-layer RoPE theta + Gemma-RMSNorm sandwich norms + sqrt(hidden) embed-scale + query_pre_attn_scalar scaling | | Gemma-2 dense (attn + final logit soft-cap) | Correctness-complete, speed-pending | Near-tie-band 48/48 (gemma-2-2b-it): 44/48 strict on vLLM's greedy + 4/48 at 0.0-nat ties in vLLM's own logits; proves the attention + final logit soft-cap primitives (attn_logit_softcapping 50 + final 30); the inverse of Gemma-3 (both soft-caps, no QK-norm) | | Gemma-1 dense (the original Gemma) | Correctness-complete, speed-pending | STRICT token-exact 48/48 greedy (gemma-2b); two fused norms/layer, head_dim scale, GeGLU + sqrt(hidden) embed-scale, tied lm_head; no soft-cap/QK-norm/sliding. **D1 (2026-07-31): the whole Gemma family (1/2/3/4) folded to the default-ON bf16 merged-QKV descriptor (`MergedQkvEnabled`); re-gated Gemma-2 SACRED 48/48 (global+sliding) + Gemma-4 STRICT 32/32 — its existing gate held** |