spike(kimi-linear): chunk_kda prefill AOT port — kernel set + pinned-config recipe (§17)

.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 |

.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.

.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))