From 4eb57b815e904c6b76ee413de86dcbd0f550990c Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Fri, 7 Aug 2026 09:03:42 +0000 Subject: [PATCH] =?UTF-8?q?perf(kimi-linear):=20per-channel-decay=20KDA=20?= =?UTF-8?q?device=20kernel=20vt::KdaGatedDeltaRule=20=E2=80=94=20106?= =?UTF-8?q?=E2=86=92122/128=20AND=203.1x=20speed=20(=C2=A715)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The §14-named residual — our vt::GdnDecode/GdnPrefill carry only a per-HEAD scalar decay g[T,Hv], so a NEW per-channel-decay GDN kernel g[T,H,D] is required — is now the additive device op vt::KdaGatedDeltaRule, and it is BOTH the STRICT lever and the speed lever, one object. MEASURED on GB10: 106->122/128 (7/8 prompts token-exact) AND 1.35->4.24 tok/s (3.1x). Grounding 1:1 (file:line both sides @ pin 555967922): KDA decode REUSES the GDN recurrence kernel — fused_recurrent_kda (kda.py:109-146) calls fused_recurrent_gated_delta_rule_fwd_kernel with IS_KDA=True (fused_recurrent.py:88-175). The SOLE net-new numeric: GDN does b_h *= exp(b_g) (per-HEAD scalar, :132-134), KDA does b_h *= exp(b_gk[None,:]) (per-K-CHANNEL, :136-137) — g is [T,Hv,Dk], broadcast across the Dv state rows. Everything else (decay -> predict -> beta -> rank-1 update -> read-out, f32 on bf16 loads) is byte-for-byte GDN's recurrence, so the shared GDN kernels are UNTOUCHED. Landed: OpId kKdaGatedDeltaRule + Fn typedef + wrapper/per-channel-g validation (ops.h, ops.cpp); CPU KdaHeadTokenStep/KdaGatedDeltaRuleKernel (cpu_ops.cpp, GdnHeadTokenStep with a per-ki decay vector); CUDA KdaScanKernel/ KdaGatedDeltaRuleKernelCuda (cuda_gdn.cu, GdnScanKernel + a 3rd shared-mem dk-array for the per-K decay). Dual-registered CPU+CUDA. Unit gate RED-first (tests/vt/test_ops_kda_recurrence.cpp) 4/4x8 GPU-green on the GB10 CUDA binary: (1) broadcast-g == vt::GdnPrefill BIT-IDENTICAL (exact float ==) — ties the net-new op to a landed+gated reference with zero new numerics; (2) distinct per-channel decay vs a from-first-principles f64 island reference; (3) validation; (4) CPU<->CUDA parity (KdaScanKernel == CPU kernel on Blackwell). GDN untouched (test_ops_gdn 66/66x4242), test_kimi_kda 14/14, test_kimi_linear_forward 13/13x656 unchanged. 23 KDA symbols linked. Wiring (opt-in, default OFF): KdaRecurrenceIsland (kimi_linear_device.cpp) gains a VT_KIMI_DEVICE_KDA branch — device-resident q_n/k_n/v feed vt::KdaGatedDeltaRule (fresh zero state, qsl=[0,T]); only the elementwise decay gate + beta stay host. Requires VT_KIMI_DEVICE_COMPUTE=1. env-doc: VT_KIMI_DEVICE_KDA + the two pre-existing H3 VT_H3_ACT_* documented. FULL 48.9B 128-token gate vs the §12 STRICT golden (GB10 sm_121a, single-load per config, memory-safe: host RSS peak 1.7 GiB, min-avail 21 GiB, freed cleanly, NO reboot): control (f64 host recurrence) 106/128 1.35 tok/s device-KDA 122/128 4.24 tok/s <- NEW BEST both axes device-KDA + bf16 knobs 90/128 4.19 tok/s (regression) RESULT (the §14 thesis CONFIRMED): the device recurrence — vLLM's ACTUAL f32-on-bf16 arithmetic — moves 106->122/128 (prompts 0-6 all 16/16; only p7 diverges at pos-6, a comma near-tie) AND is 3.1x faster. It beats BOTH the control (106) AND §14's host-precision best (120, which needed both bf16 knobs), and FIXES the p2 divergence the f64 host path had — right arithmetic, not a coin-flip. The §14 bf16 knobs are now SUPERSEDED + counterproductive (device-KDA + bf16 regresses 122->90). The speed win = the device recurrence kills the per-step host Download/f64-recompute/upload round-trip and runs O(T^2) in parallel on the GPU. Default: VT_KIMI_DEVICE_KDA STAYS OFF (122/128 is a DIVERGENCE, not STRICT; parity-enablers). Residual = the SINGLE p7 near-tie -> the recorded next brick: the KDA chunked-prefill kernel family (vLLM's prompt uses chunk_kda; regen a Triton-AOT cubin for sm_121a via scripts/regen-triton-aot.sh, or a native port) + paged mla::ForwardMlaAttentionBlock for the 7 NoPE-MLA layers + paged-incremental decode. Row STAYS ACTIVE. Records: spec §15, benchmark-record, state, NOW, STATUS/BENCHMARKS/FEATURES, ENVIRONMENT. FOLLOWING_AGENTS_PROTOCOL Assisted-by: Claude Code:claude-opus-4-8 [ClaudeCode] --- .agents/NOW.md | 4 +- .agents/benchmark-record.md | 17 +- .agents/specs/kimi-linear.md | 83 ++++ .agents/state.md | 59 +++ docs/BENCHMARKS.md | 2 +- docs/ENVIRONMENT.md | 3 + docs/FEATURES.md | 4 +- docs/STATUS.md | 2 +- include/vt/ops.h | 29 ++ .../models/kimi_linear_device.cpp | 54 +++ src/vt/cpu/cpu_ops.cpp | 74 ++++ src/vt/cuda/cuda_gdn.cu | 112 ++++++ src/vt/ops.cpp | 43 ++ tests/CMakeLists.txt | 2 + tests/vt/test_ops_kda_recurrence.cpp | 371 ++++++++++++++++++ 15 files changed, 852 insertions(+), 7 deletions(-) create mode 100644 tests/vt/test_ops_kda_recurrence.cpp diff --git a/.agents/NOW.md b/.agents/NOW.md index 342823425..d384b170e 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 NVFP4 grid DIAGNOSED (#95): NO loader bug** | weights/islands/RoPE all quant-noise-close to coherent GGUF; residual = community-NVFP4 quant fidelity §8.12 | -| Kimi-Linear-48B (KDA+NoPE-MLA+MoE) | bf16 knobs **106→120/128**, NOT STRICT (§14, `row/KIMI-LINEAR-STRICT-SPEED`); default OFF | residual = device islands; 1.30 tok/s | +| Kimi-Linear-48B (KDA/NoPE-MLA/MoE) | **`vt::KdaGatedDeltaRule` LANDED + GB10-MEASURED** (§15, `row/KIMI-KDA-DEVICE-KERNEL`): **106→122/128 + 1.35→4.24 tok/s (3.1×)**, beats §14's 120. OFF (122≠STRICT) | close p7: chunked-prefill/paged-MLA/incremental | | 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 | @@ -67,7 +67,7 @@ throughput ⇒ audit the context; per-shape MEASUREMENT arbitrates). into a lock or worktree+PR; operator merges PRs first and does features only via sub-agents; helpers use worktrees on `row/` and open a DRAFT PR at the START, which IS the claim. **W0-W5 LANDED**; role discipline ENFORCING, -`--require-role` is the DEFAULT. Queue: 10 rows (6 audit-vacated, LANDED gate anchors; READ before picking). Backfill: 79 rows, 30 anchored; blocker is claim FAMILIES. +`--require-role` is the DEFAULT. Queue: 10 rows; backfill 79 rows, 30 anchored. **Upstream inventory** ([spec](specs/upstream-derived-inventory-2026-08-05.md), drift-gated, arch parity BOTH ways): SM060/061/070 below vLLM's floor = OUT-OF-SCOPE; COMP-*/DISTRIBUTED-* are REAL unported work; **all 362 archs now have rows**; llama.cpp's 11 extra devices are IN SCOPE, spike-gated diff --git a/.agents/benchmark-record.md b/.agents/benchmark-record.md index 170c2b811..e98b25126 100644 --- a/.agents/benchmark-record.md +++ b/.agents/benchmark-record.md @@ -14593,4 +14593,19 @@ Full 48.9B model on GB10 (sm_121a clean CUDA build, cutlass-4.5.0, 14 GDN AOT sy FLIP LEDGER (deterministic golden arbitrates): the two levers interact — island bf16-input rounding fixes p2 but repeats p3; the bf16 residual stream (vLLM `fused_add_rms_norm` order) re-stabilizes p3. At 120/128 the SOLE divergence is p7 position 8 (golden deterministically `18705`, ours `58084`), a single near-tie that cascades to 8/16 on p7. Further precision-"matching" (output bf16, f32 accum) is a coin-flip that regresses — it is not vLLM's actual GDN-Triton/FA2 kernel arithmetic. Host-precision-matching PLATEAUS at 120/128. VERDICT: NO arm STRICT (K=3-deterministic golden → STRICT required, not distributional); default STAYS OFF (parity-enablers). NAMED residual (= also the speed lever): the device islands — a NEW per-channel-decay GDN kernel (`g[T,H,D]`; `vt::GdnDecode`/`GdnPrefill` carry only per-head `g[T,Hv]`, ops.h:1797/1846 — NOT a drop-in) + paged `mla::ForwardMlaAttentionBlock` (FA2). Speed HW-forced-indirect: 1.30 tok/s (O(n²) recompute + host islands, invariant to the numeric knobs); vLLM cannot serve Kimi-Linear-48B at bf16 on one GB10 (oracle capture needed util 0.82 for a single-seq eager run) so a direct `vllm bench throughput` arm is infeasible. Row STAYS ACTIVE. ->>>>>>> origin/main + +## 2026-08-07 — Kimi-Linear-48B: per-channel-decay KDA device kernel `vt::KdaGatedDeltaRule` — 106→122/128 AND 3.1× speed (the §14 residual, MEASURED on GB10) + +The §14 named residual ("a NEW per-channel-decay GDN kernel `g[T,H,D]`; `vt::GdnDecode`/`GdnPrefill` carry only per-head `g[T,Hv]` — NOT a drop-in") was BUILT as the additive device op `vt::KdaGatedDeltaRule` (`row/KIMI-KDA-DEVICE-KERNEL`, #104) and MEASURED. Grounded 1:1 in FLA `fused_recurrent_gated_delta_rule_fwd_kernel` IS_KDA=True (`b_h *= exp(b_gk[None,:])` per-K-channel, `third_party/flash_linear_attention/ops/fused_recurrent.py:136-137` @ 555967922); the shared GDN kernels are untouched (`test_ops_gdn` 66/66·4242 on GPU). Unit-gate `test_ops_kda_recurrence` (broadcast-g == `vt::GdnPrefill` BIT-IDENTICAL; distinct per-channel vs f64 ref; CPU↔CUDA parity) 4/4·8 GREEN on the GB10 CUDA binary. + +**Full 48.9B 128-token gate vs the §12 STRICT golden** (GB10 sm_121a, clean Release CUDA build, Triton-AOT vendored, cutlass-4.5.0; single-load per config; memory-safe: host RSS peak 1.7 GiB, min-avail 21 GiB, freed cleanly between configs, NO reboot): + +| Config | env | /128 | tok/s | +|---|---|---|---| +| control (f64 host recurrence) | `VT_KIMI_DEVICE_COMPUTE=1` | 106 | 1.35 | +| **device-KDA** | `…DEVICE_KDA=1` | **122** | **4.24** | +| device-KDA + bf16 knobs | `…DEVICE_KDA=1 BF16_RESIDUAL BF16_ISLANDS` | 90 | 4.19 | + +RESULT (the §14 thesis CONFIRMED): the device recurrence — vLLM's ACTUAL f32-on-bf16 arithmetic — moves **106→122/128** (prompts 0-6 all 16/16; only p7 diverges at pos-6, got `387` vs golden `11`, a comma near-tie) AND is **3.1× FASTER (1.35→4.24 tok/s)**. It beats BOTH the control (106) AND §14's host-precision best (120, which needed both bf16 knobs), and FIXES the p2 divergence the f64 host path had — because it runs the right arithmetic, not a coin-flip. The §14 bf16 knobs are now SUPERSEDED + COUNTERPRODUCTIVE: device-KDA + bf16 REGRESSES 122→90 (reintroducing p3's `163586×` repeat loop) — they were tuned to compensate for the f64 host island's over-precision. The speed win = the device recurrence kills the per-step host Download/f64-recompute/upload round-trip and runs the O(T²) recurrence in parallel on the GPU. + +VERDICT: device-KDA (122/128, 4.24 tok/s) is the NEW BEST on BOTH axes but STILL a DIVERGENCE (STRICT required, K=3-deterministic golden) → `VT_KIMI_DEVICE_KDA` STAYS OFF (parity-enablers). The residual is now a SINGLE near-tie (p7 pos-6). NAMED next brick to STRICT (+ more speed): the KDA chunked-prefill kernel family (vLLM processes the PROMPT with `chunk_kda`, we still run the recurrent form — regen a Triton-AOT cubin for sm_121a via `scripts/regen-triton-aot.sh`, or a native `chunk_kda` port) + paged `mla::ForwardMlaAttentionBlock` for the 7 NoPE-MLA layers + paged-incremental decode (persistent KDA state + MLA-KV, kills the remaining O(n²)). Row STAYS ACTIVE. diff --git a/.agents/specs/kimi-linear.md b/.agents/specs/kimi-linear.md index 32e0d82b5..acdbb5c04 100644 --- a/.agents/specs/kimi-linear.md +++ b/.agents/specs/kimi-linear.md @@ -830,6 +830,89 @@ rewrite, which is the SAME W7-speed residual. Scoped as the named follow-up, not --- +## 15. PER-CHANNEL-DECAY KDA DEVICE KERNEL LANDED (2026-08-07, `row/KIMI-KDA-DEVICE-KERNEL`) + +The §14-named residual — "our `vt::GdnDecode`/`GdnPrefill` carry only a per-HEAD scalar decay +`g[T,Hv]`, so a NEW per-channel-decay GDN kernel (`g[T,H,D]`) is required" — is now IMPLEMENTED +as the additive device op **`vt::KdaGatedDeltaRule`**, the genuinely-net-new-vs-GDN primitive. + +**Grounding (file:line, BOTH sides @ pin 555967922).** KDA's decode path REUSES the exact GDN +recurrence kernel — `fused_recurrent_kda` (`third_party/flash_linear_attention/ops/kda.py:109-146`) +calls `fused_recurrent_gated_delta_rule_fwd_kernel` with `IS_KDA=True` +(`ops/fused_recurrent.py:88-175`). The SOLE net-new numeric is the decay application: plain GDN +does `b_h *= exp(b_g)` (per-HEAD scalar, `fused_recurrent.py:132-134`), KDA does +`b_h *= exp(b_gk[None, :])` (per-K-CHANNEL, `:136-137`) — `g` is `[T,Hv,Dk]`, one log-decay per K +channel of the value head's `[Dv,Dk]` state, broadcast across the Dv rows. Everything else +(decay → predict → beta → rank-1 update → read-out, all `tl.float32` on bf16 loads) is byte-for-byte +GDN's recurrence. The op is thus GdnPrefill's per-channel twin; the shared GDN kernels are UNTOUCHED +(Qwen3.6 27B/35B gate byte-identical — `test_ops_gdn` 58/58·1825 unchanged). + +**Implementation (ours, additive).** OpId `kKdaGatedDeltaRule` + `KdaGatedDeltaRuleFn` +(`include/vt/ops.h`); wrapper + per-channel-g validation (`src/vt/ops.cpp`); CPU +`KdaHeadTokenStep`/`KdaGatedDeltaRuleKernel` (`src/vt/cpu/cpu_ops.cpp`, GdnHeadTokenStep with a +per-`ki` decay vector); CUDA `KdaScanKernel` + `KdaGatedDeltaRuleKernelCuda` (`src/vt/cuda/cuda_gdn.cu`, +GdnScanKernel staging the per-K decay in shared memory, +1 dk-array). Both dual-registered CPU+CUDA. + +**Unit gate (RED-first, `tests/vt/test_ops_kda_recurrence.cpp`) — 3/3·6 CPU-green:** +(1) EQUIVALENCE — with `g` broadcast from a per-head scalar, the per-channel op reduces +**BIT-IDENTICALLY** to the landed+gated `vt::GdnPrefill` (out & state exact-float-equal), tying the +net-new op to a proven reference with zero new numerics; (2) PER-CHANNEL — distinct per-channel +decay vs the from-first-principles f64 island reference (`KdaRecurrenceIsland` math) at documented +f32 tolerance (atol 1e-4, rtol 3e-3); (3) VALIDATION — rejects per-head `g` and unset scale; plus a +CPU↔CUDA parity case (GPU-pending). GDN untouched, `test_kimi_kda` 14/14, `test_kimi_linear_forward` +13/13·656 unchanged. + +**Wiring (opt-in, default OFF).** `KdaRecurrenceIsland` (`kimi_linear_device.cpp`) gains a +`VT_KIMI_DEVICE_KDA` branch: q_n/k_n/v (already device-resident) feed `vt::KdaGatedDeltaRule` with a +fresh zero state + qsl=[0,T]; only the elementwise decay gate (`KdaDecayGate`) + beta = sigmoid(b) +stay host (numerically stable; the numerically-sensitive object is the RECURRENCE). Requires +`VT_KIMI_DEVICE_COMPUTE=1`. CPU whole-forward gate passes with the flag ON (`test_kimi_linear_forward` +13/13·656, f32 device recurrence within the forward's rtol 5e-3 vs the f64 ref) — the WIRING is +correct. Default OFF (parity-enabler) keeps the f64 host path as production. + +**Why this is the STRICT path (spec §14 razor).** §14 proved host-precision-matching PLATEAUS at +120/128 because the f64 island is MORE precise than vLLM and coin-flips near-ties (f32-accumulation +knob regressed 120→91-106). This op runs vLLM's ACTUAL f32-on-bf16 recurrence arithmetic on device, +not a host approximation — the principled STRICT lever AND the speed lever (it is the per-step device +recurrence the paged-incremental-decode rewrite needs). + +**GPU-VERIFIED + FULL-MODEL GATE MEASURED on GB10 (2026-08-07, sm_121a, clean Release CUDA build, +Triton-AOT vendored, cutlass-4.5.0).** Kernel GPU-verify: `test_ops_kda_recurrence` **4/4·8 GREEN** +on the CUDA binary (the CPU↔CUDA parity case confirms `KdaScanKernel` == the CPU kernel on Blackwell); +GDN untouched `test_ops_gdn` 66/66·4242; `test_kimi_kda` 14/14; 23 KDA symbols linked. Full 48.9B +128-token gate vs the §12 STRICT golden, single-load per config, memory-safe throughout (host RSS peak +1.7 GiB, min-avail 21 GiB, freed cleanly between configs, NO reboot): + +| Config | env | /128 | tok/s | verdict | +|---|---|---|---|---| +| control (f64 host recurrence) | `DEVICE_COMPUTE=1` | 106 | 1.35 | reproduces §13/§14 baseline | +| **device-KDA** | `DEVICE_COMPUTE=1 DEVICE_KDA=1` | **122** | **4.24** | **NEW BEST on BOTH axes** | +| device-KDA + bf16 knobs | `…DEVICE_KDA=1 BF16_RESIDUAL BF16_ISLANDS` | 90 | 4.19 | REGRESSION (reverted) | + +**RESULT (the §14 thesis CONFIRMED).** The device recurrence — vLLM's ACTUAL f32-on-bf16 arithmetic — +moves **106→122/128** (prompts 0-6 all 16/16; only p7 diverges at pos-6, `387` vs golden `11`, a comma +near-tie) AND is **3.1× FASTER (1.35→4.24 tok/s)**. It beats BOTH the control (106) AND §14's +host-precision best (120, which needed both bf16 knobs). It FIXES the p2 divergence the f64 host path +had — because it runs the right arithmetic, not a coin-flip. The §14 bf16 knobs are now SUPERSEDED and +COUNTERPRODUCTIVE (device-KDA + bf16 REGRESSES 122→90, reintroducing p3's `163586×` repeat loop) — they +were tuned to compensate for the f64 host island's over-precision; on the already-correct device +arithmetic they perturb the wrong way. The speed win is because the device recurrence kills the host +Download/f64-recompute/upload round-trip and runs the O(T²) recurrence in parallel on the GPU. + +**Default + parity-enabler.** `VT_KIMI_DEVICE_KDA` STAYS OFF (122/128 is still a DIVERGENCE, not STRICT; +parity-enablers flip only with the token gate green). But the result reframes the residual: it is now a +SINGLE near-tie (p7 pos-6) and the recorded next brick is the clear path to STRICT + more speed. + +**NAMED residuals to STRICT (the p7 near-tie).** vLLM processes the PROMPT with the CHUNKED prefill +kernel (`chunk_kda`), we still run the RECURRENT form over the whole sequence; and the 7 NoPE-MLA layers +still use a host f64 softmax island. Closing p7 needs (c) the KDA chunked-prefill kernel family + +(d) paged `mla::ForwardMlaAttentionBlock` for the NoPE-MLA layers + (e) paged-incremental decode +(persistent KDA state + MLA-KV) to kill the remaining O(n²) recompute (more speed still). Options for +(c) mirror-first: regen a Triton-AOT cubin from FLA's KDA kernels for sm_121a (`scripts/regen-triton- +aot.sh`), or a native `chunk_kda` port. Row STAYS `ACTIVE`. + +--- + ## Structured contract (machine-readable — mirrors deepseek-v4-flash.md) ## Scope diff --git a/.agents/state.md b/.agents/state.md index 81a59e1f9..e1c128055 100644 --- a/.agents/state.md +++ b/.agents/state.md @@ -40896,3 +40896,62 @@ Records: spec §14, STATUS/BENCHMARKS/FEATURES Kimi rows, benchmark-record, NOW. fails to create. `sudo` on the host is the working path. No container was holding the GPU (all exited), so nothing was stopped and nothing needed restoring. + +## 2026-08-07T12:30 — Kimi-Linear: per-channel-decay KDA device kernel `vt::KdaGatedDeltaRule` LANDED (the §14 STRICT+speed residual, one object) + +KIMI-KDA-DEVICE-KERNEL (`row/KIMI-KDA-DEVICE-KERNEL`, helper) — the §14 named residual ("`vt::GdnDecode`/ +`GdnPrefill` carry only a per-HEAD scalar decay `g[T,Hv]`; a NEW per-channel-decay GDN kernel `g[T,H,D]` +is required") is now IMPLEMENTED as the additive device op `vt::KdaGatedDeltaRule`. + + **Grounding (1:1, file:line both sides @ 555967922).** KDA decode REUSES the GDN recurrence kernel — + `fused_recurrent_kda` (kda.py:109-146) calls `fused_recurrent_gated_delta_rule_fwd_kernel` with + `IS_KDA=True` (fused_recurrent.py:88-175). The SOLE net-new numeric: GDN does `b_h *= exp(b_g)` + (per-HEAD scalar, :132-134), KDA does `b_h *= exp(b_gk[None,:])` (per-K-CHANNEL, :136-137) — `g` is + `[T,Hv,Dk]`, broadcast across Dv rows. Everything else (decay→predict→beta→rank-1→read-out, f32 on + bf16 loads) is byte-for-byte GDN's recurrence. Shared GDN kernels UNTOUCHED (`test_ops_gdn` 58/58·1825). + + **Landed.** OpId `kKdaGatedDeltaRule` + Fn typedef + wrapper/validation (ops.h, ops.cpp); CPU + `KdaHeadTokenStep`/`KdaGatedDeltaRuleKernel` (cpu_ops.cpp, GdnHeadTokenStep with a per-`ki` decay + vector); CUDA `KdaScanKernel`/`KdaGatedDeltaRuleKernelCuda` (cuda_gdn.cu, GdnScanKernel + a 3rd + shared-mem dk-array for the per-K decay). Dual-registered CPU+CUDA. + + **Unit gate RED-first (`tests/vt/test_ops_kda_recurrence.cpp`) 3/3·6 CPU-green:** (1) broadcast-`g` + == `vt::GdnPrefill` BIT-IDENTICAL (out+state exact float ==) — ties the op to a landed reference with + zero new numerics; (2) distinct per-channel vs the f64 island reference (atol 1e-4 rtol 3e-3); + (3) validation rejects per-head g + unset scale; (+CPU↔CUDA parity, GPU-pending). `test_kimi_kda` + 14/14, `test_kimi_linear_forward` 13/13·656 unchanged (default AND with `VT_KIMI_DEVICE_KDA=1` on CPU). + + **Wiring (opt-in OFF).** `KdaRecurrenceIsland` gains a `VT_KIMI_DEVICE_KDA` branch: device-resident + q_n/k_n/v feed `vt::KdaGatedDeltaRule` (fresh zero state, qsl=[0,T]); decay gate + beta stay host + (stable). Requires `VT_KIMI_DEVICE_COMPUTE=1`. Default OFF (parity-enabler) = f64 host path production. + env-doc: `VT_KIMI_DEVICE_KDA` + the two pre-existing H3 `VT_H3_ACT_*` documented (env-doc gate green). + + **Why STRICT (the §14 razor).** §14 proved host-precision-matching PLATEAUS at 120/128 (f64 island is + MORE precise than vLLM, coin-flips near-ties; f32-acc knob regressed 120→91-106). This op runs vLLM's + ACTUAL f32-on-bf16 recurrence on device, not a host approximation — the principled STRICT lever AND + the per-step device recurrence the paged-incremental-decode speed rewrite needs. + + **GPU-VERIFIED + FULL GATE MEASURED on GB10 (2026-08-07, sm_121a, clean Release CUDA build, Triton-AOT + vendored, cutlass-4.5.0).** Kernel GPU-verify: `test_ops_kda_recurrence` **4/4·8** on the CUDA binary + (CPU↔CUDA parity confirms KdaScanKernel == CPU kernel on Blackwell); `test_ops_gdn` 66/66·4242 (GDN + untouched); `test_kimi_kda` 14/14; 23 KDA symbols linked. Full 48.9B 128-token gate vs the §12 STRICT + golden, single-load per config, memory-safe (host RSS peak 1.7 GiB, min-avail 21 GiB, freed cleanly + between configs, NO reboot): + + | config | env | /128 | tok/s | + |---|---|---|---| + | control (f64 host) | `DEVICE_COMPUTE=1` | 106 | 1.35 | + | **device-KDA** | `DEVICE_COMPUTE=1 DEVICE_KDA=1` | **122** | **4.24** | + | device-KDA + bf16 | `…DEVICE_KDA=1 BF16_RESIDUAL BF16_ISLANDS` | 90 | 4.19 | + + **RESULT (§14 thesis CONFIRMED):** the device recurrence (vLLM's ACTUAL f32-on-bf16 arithmetic) moves + **106→122/128** (prompts 0-6 all 16/16; only p7 diverges pos-6, `387` vs golden `11`, a comma near-tie) + AND is **3.1× faster (1.35→4.24 tok/s)** — beats BOTH the control (106) AND §14's host-precision best + (120). It FIXES the p2 divergence the f64 path had (right arithmetic, not a coin-flip). The §14 bf16 + knobs are now SUPERSEDED + COUNTERPRODUCTIVE (device-KDA + bf16 REGRESSES 122→90, reintroduces p3's + `163586×` repeat loop). Speed win = the device recurrence kills the host Download/f64-recompute/upload + round-trip and runs O(T²) in parallel on GPU. `VT_KIMI_DEVICE_KDA` STAYS OFF (122 is a DIVERGENCE, not + STRICT). Residual = the SINGLE p7 near-tie → the recorded next brick: KDA chunked-prefill family + (`chunk_kda`; regen a Triton-AOT cubin for sm_121a via `scripts/regen-triton-aot.sh`, or native port) + + paged `mla::ForwardMlaAttentionBlock` (7 NoPE-MLA layers) + paged-incremental decode. Row `ACTIVE`. + Box left clean (artifacts removed, memory restored, worker parked, no reboot). diff --git a/docs/BENCHMARKS.md b/docs/BENCHMARKS.md index 487c8594e..431fa82e8 100644 --- a/docs/BENCHMARKS.md +++ b/docs/BENCHMARKS.md @@ -309,7 +309,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); bf16-regime knobs 106→120/128 (7/8 exact), NOT STRICT; default OFF | bf16 residual+island-inputs → 120/128 best (control/each-alone 106; output-bf16 & f32-accum NEGATIVE); 1 near-tie left. 1.30 tok/s (O(n²)); vLLM HW-can't-serve bf16 on 1 GB10. Residual = device islands. §14 | +| Kimi-Linear-48B-A3B (KDA+MLA+MoE) | e2e RUNS (bf16-resident §13); per-channel-decay KDA device kernel `vt::KdaGatedDeltaRule` LANDED+GB10: **106→122/128 + 4.24 tok/s (3.1×)**, beats §14's 120 on both axes, NOT STRICT; default OFF | device-KDA runs vLLM's actual GPU recurrence; bf16 knobs now hurt (122→90). p7 near-tie; residual = KDA chunked-prefill + paged FA2 MLA + incremental. §15 | | 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. 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 6a938ebc6..6be88d14e 100644 --- a/docs/ENVIRONMENT.md +++ b/docs/ENVIRONMENT.md @@ -67,6 +67,7 @@ portable/reference path. In normal operation leave them unset. | `VT_KIMI_BF16_RESIDUAL` | off (opt-in) | `=1` carries the Kimi-Linear device-compute residual stream in bf16 like vLLM's `fused_add_rms_norm` (residual/block-outputs bf16, RMSNorm variance over the f32 pre-store sum), via in-place f32→bf16→f32 rounds. W7-speed STRICT-lever A/B (spec §14). Default OFF → byte-identical. MEASURED: alone net-zero; WITH `VT_KIMI_BF16_ISLANDS` → 120/128 (best, still a near-tie, NOT STRICT) | | `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_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 | @@ -109,6 +110,8 @@ Read-only observability; none change output. | `VT_H3_TRACE_MOTION` | unset | `=1` prints one `[h3-motion] step ...` line per MiniMax-H3 denoise step to stderr: the step's velocity stats (`v_rms`/`v_amax`/`v_mean` of the DiT output), the per-step latent motion over the denoise-target rows (`drows_rms`), and the running latent norm (`rows_rms`). Because the rectified-flow Euler integration telescopes to `(sigma0 - sigmaN) * v`, a velocity that does not EVOLVE across steps produces a step-count-invariant result; this trace measures exactly that (added for the render-coherence bisection). Byte-identical when unset — every read is guarded and it only reads buffers the loop already holds | | `VT_H3_VAE_PROBE` | unset | `=1` runs a video-VAE receptive-field probe after the normal decode: it perturbs ONE interior spatial latent cell (across all channels and temporal frames), re-decodes, and prints a per-16px-block RMS-change map (`[h3-vae-probe]`) over output frame 0. If only the perturbed cell's block moves, the ViT3D decoder is not mixing tokens spatially. Byte-identical to production when unset (no second decode) | | `VT_H3_DUMP_DIR` | unset | Directory into which the MiniMax-H3 denoise loop writes the initial and final video latent rows (`init_video_rows.f32`, `final_video_rows.f32`) and the pipeline writes the exact VAE-input latent (`vae_input_video_latent.f32`), all raw little-endian f32. Lets two runs (e.g. 12 vs 50 steps, conditioned vs not) be byte/stat-compared, and the video VAE decode be replayed on a KNOWN latent, without re-running the denoise. Byte-identical to production when unset (no file is opened) | +| `VT_H3_ACT_DUMP` | unset | `=` writes a per-STAGE activation fingerprint of the MiniMax-H3 DiT forward (`MiniMaxH3DitForwardDevice`) — stats (mean/rms/absmax/finite) plus fixed positional sample values — so two weight arms running the SAME graph on the SAME inputs (e.g. the NVFP4-bf16 stream vs the FL2VA-GGUF-bf16 control) can be diffed stage-by-stage to localise where they diverge. Byte-identical to production when unset (no file opened) | +| `VT_H3_ACT_CALL` | `0` | Selects WHICH forward `VT_H3_ACT_DUMP` captures: only the forward whose 0-based call index equals this value dumps (default `0`), so a single small render (`--denoise-only --steps 1`) captures exactly one clean forward. No effect unless `VT_H3_ACT_DUMP` is set | | `VT_H3_DUMP_INPUTS` | unset | Directory into which the MiniMax-H3 denoise loop writes EVERY DiT input at step 0 as raw little-endian binary plus a `manifest.txt` — the packed layout (`input_ids`/`image_mask`/`audio_mask`/`img_pos`/`audio_pos`/`text_pos`/`update_mask`/`cu_seqlens`/`document_id`), the fp64 position grid (`img_position_ids.f64`), the per-token modality tags (`token_tags.i64`), the per-token pre-unique timesteps and their `unique_timesteps`/`inverse_indices`/`combined_indices` AdaLN selection, both sigma schedules, and the raw `prompt_embeds`; the `minimax-h3-gen` driver additionally writes `prompt_token_ids.i32`. Lets the REAL-scale DiT inputs be diffed EXACTLY against upstream `pipeline_minimax_h3.py` (the render-coherence S1 surface the reduced-dim ladder never fed real values into). Byte-identical to production when unset (no file is opened) | ## Kernel-internal knobs (deferred) diff --git a/docs/FEATURES.md b/docs/FEATURES.md index c188e198e..5c5ec0bf3 100644 --- a/docs/FEATURES.md +++ b/docs/FEATURES.md @@ -119,7 +119,7 @@ speed-pending, which [BENCHMARKS.md](BENCHMARKS.md) tracks. | `Glm4ForCausalLM` | GLM-4-9B-0414 | near-tie 16/16 vs vLLM 0.25.0 | pending | | `Glm4MoeLiteForCausalLM` | zai-org/GLM-4.7-Flash (31.2B, MLA MoE) | near-tie 8/8 vs vLLM 0.25.0 | pending | | `LagunaForCausalLM` | poolside/Laguna-S-2.1-NVFP4, GGUF-Q4_K, Laguna-XS | byte-exact near-tie (distributional vs vLLM) | vLLM parity+ 1.03x, default on | -| `KimiLinearForCausalLM` | Kimi-Linear-48B-A3B (KDA + NoPE-MLA + MoE) | e2e runs bf16-resident; bf16-regime knobs 106→120/128 (7/8 exact), not STRICT | 1.30 tok/s, default off; residual = device islands (§14) | +| `KimiLinearForCausalLM` | Kimi-Linear-48B-A3B (KDA + NoPE-MLA + MoE) | e2e runs bf16-resident; per-channel-decay KDA device kernel `vt::KdaGatedDeltaRule` landed+GB10-measured: 106→122/128 (7/8 exact) + 1.35→4.24 tok/s (3.1×), not STRICT | default off; beats §14's bf16-knob 120; residual = KDA chunked-prefill + paged FA2 MLA + incremental (§15) | | `KimiK3ForConditionalGeneration` | Kimi-K3 (2.8T MoE) | scaffold: registry+config+enumeration gated, forward refuses | HW-infeasible (~1.56 TB); no run | | `CohereForCausalLM` | Command-R / Cohere (and Cohere2) | scaffold: W0 tiny-random oracle run-verified; real-checkpoint gate blocked | no run | @@ -250,7 +250,7 @@ abstraction, and `world_size == 1` stays byte-identical. | Gap | State | Detail | |---|---|---| -| Kimi-Linear-48B-A3B (KDA + NoPE-MLA + MoE hybrid) | e2e RUNS (bf16-resident §13); bf16-regime knobs 106→120/128 (7/8 exact), NOT STRICT; default OFF | bf16 residual+island-inputs → 120/128 best (§14); 1 near-tie left. STRICT+speed residual = device islands (per-channel GDN kernel + paged FA2). 1.30 tok/s; vLLM HW-can't-serve bf16 on 1 GB10 | +| Kimi-Linear-48B-A3B (KDA + NoPE-MLA + MoE hybrid) | e2e RUNS (bf16-resident §13); per-channel-decay KDA device kernel `vt::KdaGatedDeltaRule` LANDED+GB10: 106→122/128 + 4.24 tok/s (3.1×), beats §14's 120, NOT STRICT; default OFF | p7 near-tie left; residual = KDA chunked-prefill + paged FA2 MLA + incremental decode (§15) | | Multi-GPU execution | Hardware-blocked | TP proven equal to tp=1 on CPU; no 2-GPU box to run it | | LoRA end to end | CPU brick landed | Unwired standalone; not usable through the server | | Multimodal over HTTP | Architecturally blocked | Vision tower lives outside the registered engine forward | diff --git a/docs/STATUS.md b/docs/STATUS.md index 79c209d6c..26644d76f 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); bf16-regime knobs 106→120/128 (7/8 exact), NOT STRICT; default OFF | CPU+CUDA 13/13·656. bf16 residual+island-inputs → 120/128 (§14); 1 near-tie left (p7). STRICT+speed residual = device islands (per-channel GDN kernel + paged FA2). 1.30 tok/s; vLLM HW-can't-serve bf16 on 1 GB10 | +| Kimi-Linear-48B-A3B (KDA + NoPE-MLA + MoE hybrid) | e2e RUNS (bf16-resident §13); per-channel-decay KDA device kernel `vt::KdaGatedDeltaRule` LANDED+GB10: **106→122/128 + 4.24 tok/s (3.1×)**, beats §14's 120 on both axes, NOT STRICT; default OFF | CPU+CUDA 13/13·656 + `test_ops_kda_recurrence` 4/4 (§15). p7 near-tie left; residual = KDA chunked-prefill + paged FA2 MLA + incremental (§15) | | 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** | diff --git a/include/vt/ops.h b/include/vt/ops.h index 0bc04f300..63a900180 100644 --- a/include/vt/ops.h +++ b/include/vt/ops.h @@ -101,6 +101,7 @@ enum class OpId : uint8_t { kGdnDecode, kGdnSpecDecode, kGdnPackedDecode, + kKdaGatedDeltaRule, kMoeRouterTopK, kMoeCombine, kAttention, @@ -851,6 +852,11 @@ using GdnPackedDecodeFn = void (*)(Queue&, Tensor&, const Tensor&, const Tensor&, const Tensor&, const Tensor&, const Tensor&, Tensor&, const Tensor&, const GdnArgs&); +// Per-k-channel-decay gated-delta recurrence (KDA). Same shape as GdnPrefillFn; +// the ONLY difference is g is [T,Hv,Dk] (per-channel) not [T,Hv] (per-head). +using KdaGatedDeltaRuleFn = void (*)(Queue&, Tensor&, const Tensor&, const Tensor&, const Tensor&, + const Tensor&, const Tensor&, Tensor&, const Tensor&, + const GdnArgs&); using GdnStateGatherFn = void (*)(Queue&, Tensor&, const Tensor&, const Tensor&, const Tensor*); using GdnStateScatterFn = @@ -1805,6 +1811,29 @@ void GdnPrefill(Queue& q, Tensor& out, const Tensor& q_in, const Tensor& k, cons const Tensor& g, const Tensor& beta, Tensor& state, const Tensor& query_start_loc, const GdnArgs& args); +// Kimi Delta Attention (KDA) gated-delta recurrence — the PER-K-CHANNEL-DECAY +// variant of GdnPrefill. Ported 1:1 from FLA's +// fused_recurrent_gated_delta_rule_fwd_kernel with IS_KDA=True (third_party/ +// flash_linear_attention/ops/fused_recurrent.py:88-175 @ pin 555967922; the KDA +// wrapper is ops/kda.py:109-146 fused_recurrent_kda). The plain-GDN kernel +// (IS_KDA=False) applies a PER-HEAD scalar decay `b_h *= exp(b_g)`; KDA applies a +// PER-K-CHANNEL decay `b_h *= exp(b_gk[None, :])` — so `g` here is [T,Hv,Dk] +// (one log-decay per K channel of the value head's [Dv,Dk] state), broadcast +// across the Dv rows. Everything else is byte-for-byte GdnPrefill's recurrence +// (decay -> predict -> beta -> rank-1 update -> read-out, all in f32): +// S[:,k] *= exp(g[hv,k]); v' = (v - S @ k) * beta[hv]; S += outer(v',k); +// out = S @ (q*scale) +// The shared GDN kernels are UNTOUCHED (Qwen3.6 27B/35B GDN gate byte-identical); +// KDA lands as this additive per-channel op. q_in/k [T,Hk,Dk] MUST be +// l2-normalized by the caller (as GdnPrefill; upstream fuses it via +// USE_QK_L2NORM_IN_KERNEL, exact per gdn-semantics.md §4). v/out [T,Hv,Dv], +// g [T,Hv,Dk] f32 per-channel log-decay, beta [T,Hv] f32 (per-head sigmoid(b)), +// state [N,Hv,Dv,Dk] f32 in/out (zeros for fresh sequences), +// query_start_loc [N+1] i32. For KDA Hk==Hv; GQA broadcast supported for reuse. +void KdaGatedDeltaRule(Queue& q, Tensor& out, const Tensor& q_in, const Tensor& k, + const Tensor& v, const Tensor& g, const Tensor& beta, Tensor& state, + const Tensor& query_start_loc, const GdnArgs& args); + // Single-token gated-delta-rule step, one token per sequence // (gdn-semantics.md §7 decode path). Same math as GdnPrefill with T == B and // state[B,Hv,Dv,Dk] row b for token b. q_in/k must be l2-normalized by the diff --git a/src/vllm/model_executor/models/kimi_linear_device.cpp b/src/vllm/model_executor/models/kimi_linear_device.cpp index 662400ced..4c2cf93bc 100644 --- a/src/vllm/model_executor/models/kimi_linear_device.cpp +++ b/src/vllm/model_executor/models/kimi_linear_device.cpp @@ -156,6 +156,24 @@ bool IslandF32Acc() { }(); return on; } + +// (4) VT_KIMI_DEVICE_KDA — run the KDA per-k-channel gated-delta RECURRENCE through the +// device op vt::KdaGatedDeltaRule (the net-new per-channel-decay GDN kernel, cuda_gdn.cu +// KdaScanKernel) instead of the f64 host recompute. This is the principled path to STRICT +// AND the speed lever (spec §14): the recurrence runs vLLM's actual f32-on-bf16 arithmetic +// (FLA fused_recurrent_gated_delta_rule_fwd_kernel IS_KDA=True) on device rather than a +// host f64 recompute that is MORE precise than vLLM and coin-flips near-ties. The decay +// gate `g = -exp(A_log)*softplus(f_b(f_a(x))+dt_bias)` and beta = sigmoid(b) stay host +// (elementwise, numerically stable — the numerically-sensitive object is the recurrence). +// Default OFF (parity-enabler: flip ON only with the token gate green). Independent of the +// bf16-precision knobs (BF16_ISLANDS still rounds the gate inputs when both are set). +bool DeviceKda() { + static const bool on = [] { + const char* e = std::getenv("VT_KIMI_DEVICE_KDA"); + return e != nullptr && e[0] == '1'; + }(); + return on; +} // Round a running f64 accumulator to f32 precision when the knob is on (identity else). inline double AccR(double x) { static const bool f32 = IslandF32Acc(); @@ -281,6 +299,42 @@ DBuf KdaRecurrenceIsland(const Dev& d, DBuf& qn, DBuf& kn, DBuf& vc, DBuf& g1, const int64_t nh = p.kda_num_heads; const int64_t hd = p.kda_head_dim; const int64_t proj = nh * hd; + + // ── DEVICE RECURRENCE (VT_KIMI_DEVICE_KDA): run the per-k-channel gated-delta + // recurrence on device via vt::KdaGatedDeltaRule (KdaScanKernel), vLLM's actual + // f32-on-bf16 arithmetic, instead of the host f64 recompute below. q_n/k_n/v are + // ALREADY device-resident; only the elementwise gate (KdaDecayGate) + beta = sigmoid(b) + // are computed on host and uploaded (small, numerically stable). Fresh zero state, + // single sequence, qsl=[0,T] — the stateless full-sequence recurrence the island needs. + if (DeviceKda()) { + std::vector dhg1(static_cast(T) * proj), + dhbraw(static_cast(T) * nh); + g1.Download(d, dhg1.data()); + braw.Download(d, dhbraw.data()); + RoundHostBf16(dhg1); // honor BF16_ISLANDS on the gate inputs + RoundHostBf16(dhbraw); + const std::vector gch = + kimi_kda::KdaDecayGate(dhg1, a_log, dt_bias, T, nh, hd); // [T,nh,hd] per-channel + std::vector hbeta(static_cast(T) * nh); + for (size_t i = 0; i < hbeta.size(); ++i) + hbeta[i] = static_cast(Sigmoid(dhbraw[i])); + DBuf dg(d, DType::kF32, {T, nh, hd}, gch.data()); + DBuf dbeta(d, DType::kF32, {T, nh}, hbeta.data()); + DBuf dstate(d, DType::kF32, {1, nh, hd, hd}); + dstate.Zero(d); + DBuf dcore(d, DType::kF32, {T, proj}); + const int32_t qsl[2] = {0, static_cast(T)}; + DBuf dqsl(d, DType::kI32, {2}, qsl); + Tensor qn3 = MakeTensor(qn.ptr(), DType::kF32, d.q.device, {T, nh, hd}); + Tensor kn3 = MakeTensor(kn.ptr(), DType::kF32, d.q.device, {T, nh, hd}); + Tensor vc3 = MakeTensor(vc.ptr(), DType::kF32, d.q.device, {T, nh, hd}); + Tensor out3 = MakeTensor(dcore.ptr(), DType::kF32, d.q.device, {T, nh, hd}); + const float scale = static_cast(std::pow(static_cast(hd), -0.5)); + vt::KdaGatedDeltaRule(d.q, out3, qn3, kn3, vc3, dg.t(), dbeta.t(), dstate.t(), dqsl.t(), + vt::GdnArgs{scale}); + return dcore; + } + std::vector hqn(static_cast(T) * proj), hkn(hqn.size()), hv(hqn.size()), hg1(hqn.size()), hbraw(static_cast(T) * nh); qn.Download(d, hqn.data()); diff --git a/src/vt/cpu/cpu_ops.cpp b/src/vt/cpu/cpu_ops.cpp index c0a518dc2..b64daf2dc 100644 --- a/src/vt/cpu/cpu_ops.cpp +++ b/src/vt/cpu/cpu_ops.cpp @@ -1343,6 +1343,77 @@ void GdnDecodeKernel(Queue&, Tensor& out, const Tensor& q_in, const Tensor& k, c }); } +// ── KDA per-K-channel-decay gated-delta recurrence (kKdaGatedDeltaRule) ──────── +// Byte-for-byte GdnHeadTokenStep EXCEPT the decay is per-K-channel: plain GDN +// does `s_row[ki] *= exp(g_head)` (one scalar per value head), KDA does +// `s_row[ki] *= exp(g[.,hv,ki])` (one log-decay per K channel), broadcast across +// the Dv rows. Ported 1:1 from FLA fused_recurrent_gated_delta_rule_fwd_kernel +// IS_KDA=True (`b_h *= exp(b_gk[None, :])`, fused_recurrent.py:136-137 @ 555967922). +// All arithmetic f32 (FLA loads bf16 -> tl.float32); same reduction order as GDN. +void KdaHeadTokenStep(Tensor& out, const Tensor& q_in, const Tensor& k_in, const Tensor& v_in, + const Tensor& g, const Tensor& beta, float* s_head, int64_t tok, + int64_t hv, int64_t hk, int64_t hk_n, int64_t hv_n, int64_t dk, int64_t dv, + float scale, std::vector& qbuf, std::vector& kbuf, + std::vector& vbuf, std::vector& decaybuf) { + const float beta_t = beta.Ptr()[tok * hv_n + hv]; + const float* g_row = g.Ptr() + (tok * hv_n + hv) * dk; + for (int64_t ki = 0; ki < dk; ++ki) { + qbuf[static_cast(ki)] = LoadF32(q_in, (tok * hk_n + hk) * dk + ki) * scale; + kbuf[static_cast(ki)] = LoadF32(k_in, (tok * hk_n + hk) * dk + ki); + decaybuf[static_cast(ki)] = std::exp(g_row[ki]); + } + for (int64_t vi = 0; vi < dv; ++vi) { + float* s_row = s_head + vi * dk; + float dot = 0.0f; // (S * exp(g_channel)) @ k, fused with the per-channel decay + for (int64_t ki = 0; ki < dk; ++ki) { + s_row[ki] *= decaybuf[static_cast(ki)]; + dot += s_row[ki] * kbuf[static_cast(ki)]; + } + vbuf[static_cast(vi)] = + (LoadF32(v_in, (tok * hv_n + hv) * dv + vi) - dot) * beta_t; + } + for (int64_t vi = 0; vi < dv; ++vi) { + float* s_row = s_head + vi * dk; + float o = 0.0f; // (S + outer(v',k)) @ q', fused with the rank-1 update + for (int64_t ki = 0; ki < dk; ++ki) { + s_row[ki] += vbuf[static_cast(vi)] * kbuf[static_cast(ki)]; + o += s_row[ki] * qbuf[static_cast(ki)]; + } + StoreF32(out, (tok * hv_n + hv) * dv + vi, o); + } +} + +void KdaGatedDeltaRuleKernel(Queue&, Tensor& out, const Tensor& q_in, const Tensor& k, + const Tensor& v, const Tensor& g, const Tensor& beta, Tensor& state, + const Tensor& qsl, const GdnArgs& args) { + const int64_t n = state.shape[0], hv_n = state.shape[1], dv = state.shape[2], + dk = state.shape[3]; + const int32_t* qslp = qsl.Ptr(); + VT_CHECK(qslp[0] == 0 && qslp[n] == q_in.shape[0], + "kda_gated_delta_rule: bad query_start_loc bounds"); + for (int64_t s = 0; s < n; ++s) { + VT_CHECK(qslp[s + 1] >= qslp[s], "kda_gated_delta_rule: query_start_loc not monotonic"); + } + const int64_t hk_n = q_in.shape[1]; + const int64_t ratio = hv_n / hk_n; + const int64_t nitems = n * hv_n; + // Row-chunked over (SEQUENCE, VALUE-HEAD) exactly as GdnPrefillKernel — each + // (s, hv) owns a disjoint state block + output rows; sequential in tok. + ForRows(nitems, [&](int64_t r0, int64_t r1) { + std::vector qbuf(static_cast(dk)), kbuf(static_cast(dk)), + vbuf(static_cast(dv)), decaybuf(static_cast(dk)); + for (int64_t item = r0; item < r1; ++item) { + const int64_t s = item / hv_n; + const int64_t hv = item % hv_n; + const int64_t hk = hv / ratio; + float* s_head = state.Ptr() + (s * hv_n + hv) * dv * dk; + for (int64_t t = qslp[s]; t < qslp[s + 1]; ++t) + KdaHeadTokenStep(out, q_in, k, v, g, beta, s_head, t, hv, hk, hk_n, hv_n, dk, dv, + args.scale, qbuf, kbuf, vbuf, decaybuf); + } + }); +} + // SPECULATIVE (multi-token, slot-snapshotting) gated-delta-rule step. // Ported from vllm/model_executor/layers/fla/ops/fused_sigmoid_gating.py @ // e24d1b24 — fused_recurrent_gated_delta_rule_fwd_kernel with IS_VARLEN @@ -2412,6 +2483,9 @@ struct Registrar { OpId::kGdnPackedDecode, DeviceType::kCPU, reinterpret_cast(static_cast( &GdnPackedDecodeKernel))); + RegisterOp( + OpId::kKdaGatedDeltaRule, DeviceType::kCPU, + reinterpret_cast(static_cast(&KdaGatedDeltaRuleKernel))); RegisterOp( OpId::kGdnStateGather, DeviceType::kCPU, reinterpret_cast( diff --git a/src/vt/cuda/cuda_gdn.cu b/src/vt/cuda/cuda_gdn.cu index b0264ca76..08f02abb3 100644 --- a/src/vt/cuda/cuda_gdn.cu +++ b/src/vt/cuda/cuda_gdn.cu @@ -2642,6 +2642,115 @@ void GdnScanCuda(Queue& q, Tensor& out, const Tensor& q_in, const Tensor& k, con } } +// =========================================================================== +// KDA per-K-channel-decay sequential scan (kKdaGatedDeltaRule). +// +// Byte-for-byte GdnScanKernel EXCEPT the state decay is per-K-channel: GDN uses +// one scalar `decay = expf(g[t*hv_n+hv])` for the whole [Dv,Dk] state; KDA stages +// a per-K vector `decay[ki] = expf(g[(t*hv_n+hv)*dk + ki])` in shared memory and +// applies `s_row[ki] *= decay[ki]`. Ported 1:1 from FLA +// fused_recurrent_gated_delta_rule_fwd_kernel IS_KDA=True +// (`b_h *= exp(b_gk[None, :])`, third_party/flash_linear_attention/ops/ +// fused_recurrent.py:136-137 @ pin 555967922). All state math f32; q/k must be +// l2-normalized by the caller (as GdnScan). The shared GDN kernels are untouched. +template +__global__ void KdaScanKernel(Tout* out, const Tin* q, const Tin* k, const Tin* v, + const float* g, const float* beta, TState* state, + const int32_t* qsl, const int32_t* state_idx, + int64_t state_slots, int64_t hk_n, int64_t dk, + int64_t hv_n, int64_t dv, float scale) { + const int64_t s = blockIdx.y; // sequence + const int64_t hv = blockIdx.x; // v-head + const int64_t hk = hv / (hv_n / hk_n); + const int64_t state_slot = state_idx != nullptr ? state_idx[s] : s; + if (state_slot < 0 || state_slot >= state_slots) { + const int64_t begin = qsl != nullptr ? qsl[s] : s; + const int64_t end = qsl != nullptr ? qsl[s + 1] : s + 1; + for (int64_t t = begin; t < end; ++t) + for (int64_t vi = threadIdx.x; vi < dv; vi += blockDim.x) + Store(out, (t * hv_n + hv) * dv + vi, 0.0f); + return; + } + extern __shared__ float smem[]; // [dk] q' then [dk] k then [dk] per-K decay + float* q_sh = smem; + float* k_sh = smem + dk; + float* d_sh = smem + 2 * dk; + TState* s_head = state + (state_slot * hv_n + hv) * dv * dk; // [Dv, Dk] + const int64_t begin = qsl != nullptr ? qsl[s] : s; + const int64_t end = qsl != nullptr ? qsl[s + 1] : s + 1; + for (int64_t t = begin; t < end; ++t) { + for (int64_t i = threadIdx.x; i < dk; i += blockDim.x) { + q_sh[i] = Load(q, (t * hk_n + hk) * dk + i) * scale; + k_sh[i] = Load(k, (t * hk_n + hk) * dk + i); + d_sh[i] = expf(g[(t * hv_n + hv) * dk + i]); // per-K-channel decay + } + __syncthreads(); + const float beta_t = beta[t * hv_n + hv]; + for (int64_t vi = threadIdx.x; vi < dv; vi += blockDim.x) { + TState* s_row = s_head + vi * dk; + float dot = 0.0f; // (S * exp(g_channel)) @ k, fused with the per-channel decay + for (int64_t ki = 0; ki < dk; ++ki) { + const float decayed = Load(s_row, ki) * d_sh[ki]; + dot += decayed * k_sh[ki]; + } + const float vp = (Load(v, (t * hv_n + hv) * dv + vi) - dot) * beta_t; + float o = 0.0f; // (S + outer(v',k)) @ q', fused with the rank-1 update + for (int64_t ki = 0; ki < dk; ++ki) { + const float updated = Load(s_row, ki) * d_sh[ki] + vp * k_sh[ki]; + Store(s_row, ki, updated); + o += updated * q_sh[ki]; + } + Store(out, (t * hv_n + hv) * dv + vi, o); + } + __syncthreads(); // all reads of the staged rows done before next token's load + } +} + +template +void LaunchKdaScan(cudaStream_t s, Tensor& out, const Tensor& q_in, const Tensor& k, + const Tensor& v, const Tensor& g, const Tensor& beta, Tensor& state, + const int32_t* qsl, int64_t n, const GdnArgs& args) { + const int64_t hk_n = q_in.shape[1], dk = q_in.shape[2]; + const int64_t hv_n = v.shape[1], dv = v.shape[2]; + const dim3 grid(static_cast(hv_n), static_cast(n)); + const size_t shmem = 3 * static_cast(dk) * sizeof(float); // q' + k + decay + KdaScanKernel<<>>( + out.Ptr(), q_in.Ptr(), k.Ptr(), v.Ptr(), g.Ptr(), + beta.Ptr(), state.Ptr(), qsl, nullptr, state.shape[0], hk_n, dk, hv_n, dv, + args.scale); + Check(cudaGetLastError(), "kda scan launch"); +} + +void KdaGatedDeltaRuleKernelCuda(Queue& q, Tensor& out, const Tensor& q_in, const Tensor& k, + const Tensor& v, const Tensor& g, const Tensor& beta, + Tensor& state, const Tensor& qsl, const GdnArgs& args) { + constexpr const char* name = "kda_gated_delta_rule"; + VT_CHECK(q_in.dtype == DType::kF32 || q_in.dtype == DType::kBF16, + std::string("cuda ") + name + ": unsupported q dtype (f32/bf16 only)"); + VT_CHECK(k.dtype == q_in.dtype && v.dtype == q_in.dtype, + std::string("cuda ") + name + ": q/k/v dtypes must match"); + const int64_t n = state.shape[0]; + const int64_t hv_n = state.shape[1], dv = state.shape[2], dk = state.shape[3]; + if (n == 0 || hv_n == 0 || dv == 0) return; + VT_CHECK(n <= kMaxGridY, std::string("cuda ") + name + ": too many sequences (grid.y limit)"); + VT_CHECK(3 * static_cast(dk) * sizeof(float) <= 48 * 1024, + std::string("cuda ") + name + ": Dk too large for the shared q'/k/decay staging"); + cudaStream_t s = AsStream(q); + const int32_t* qsl_ptr = qsl.Ptr(); + if (q_in.dtype == DType::kF32) { + if (out.dtype == DType::kF32) + LaunchKdaScan(s, out, q_in, k, v, g, beta, state, qsl_ptr, n, args); + else + LaunchKdaScan(s, out, q_in, k, v, g, beta, state, qsl_ptr, n, args); + } else { + if (out.dtype == DType::kF32) + LaunchKdaScan<__nv_bfloat16, float>(s, out, q_in, k, v, g, beta, state, qsl_ptr, n, args); + else + LaunchKdaScan<__nv_bfloat16, __nv_bfloat16>(s, out, q_in, k, v, g, beta, state, qsl_ptr, n, + args); + } +} + // =========================================================================== // Chunk-parallel GDN prefill scan (gdn-semantics.md §7 "chunked oracle"). // @@ -5800,6 +5909,9 @@ struct Registrar { OpId::kGdnPackedDecode, DeviceType::kCUDA, reinterpret_cast(static_cast( &GdnPackedDecodeKernelCuda))); + RegisterOp( + OpId::kKdaGatedDeltaRule, DeviceType::kCUDA, + reinterpret_cast(static_cast(&KdaGatedDeltaRuleKernelCuda))); RegisterOp( OpId::kGdnStateGather, DeviceType::kCUDA, reinterpret_cast( diff --git a/src/vt/ops.cpp b/src/vt/ops.cpp index 89a569976..ab63e9f6a 100644 --- a/src/vt/ops.cpp +++ b/src/vt/ops.cpp @@ -1826,6 +1826,49 @@ void GdnPrefill(Queue& q, Tensor& out, const Tensor& q_in, const Tensor& k, cons q, out, q_in, k, v, g, beta, state, query_start_loc, args); } +void KdaGatedDeltaRule(Queue& q, Tensor& out, const Tensor& q_in, const Tensor& k, + const Tensor& v, const Tensor& g, const Tensor& beta, Tensor& state, + const Tensor& query_start_loc, const GdnArgs& args) { + constexpr const char* name = "kda_gated_delta_rule"; + // Same contracts as GdnPrefill EXCEPT g is per-K-channel [T,Hv,Dk]. + VT_CHECK(q_in.rank == 3 && k.rank == 3 && v.rank == 3 && out.rank == 3 && g.rank == 3 && + beta.rank == 2 && state.rank == 4, + std::string(name) + + ": q/k [T,Hk,Dk], v/out [T,Hv,Dv], g [T,Hv,Dk], beta [T,Hv], state [N,Hv,Dv,Dk]"); + const int64_t t = q_in.shape[0], hk = q_in.shape[1], dk = q_in.shape[2]; + const int64_t hv = v.shape[1], dv = v.shape[2]; + VT_CHECK(k.shape[0] == t && k.shape[1] == hk && k.shape[2] == dk, + std::string(name) + ": k shape must match q"); + VT_CHECK(v.shape[0] == t, std::string(name) + ": v token count must match q"); + VT_CHECK(out.shape[0] == t && out.shape[1] == hv && out.shape[2] == dv, + std::string(name) + ": out must be [T,Hv,Dv]"); + VT_CHECK(g.shape[0] == t && g.shape[1] == hv && g.shape[2] == dk, + std::string(name) + ": g must be [T,Hv,Dk] (per-K-channel decay)"); + VT_CHECK(beta.shape[0] == t && beta.shape[1] == hv, + std::string(name) + ": beta must be [T,Hv]"); + VT_CHECK(hk >= 1 && hv % hk == 0, + std::string(name) + ": Hv must be a multiple of Hk (GQA broadcast)"); + VT_CHECK(state.shape[1] == hv && state.shape[2] == dv && state.shape[3] == dk, + std::string(name) + ": state must be [N,Hv,Dv,Dk]"); + VT_CHECK(IsFloat(q_in.dtype) && IsFloat(k.dtype) && IsFloat(v.dtype) && IsOutFloat(out.dtype), + std::string(name) + ": float q/k/v, f32/bf16 out"); + VT_CHECK(g.dtype == DType::kF32 && beta.dtype == DType::kF32, + std::string(name) + ": g/beta must be f32 (upstream keeps them f32)"); + VT_CHECK(state.dtype == DType::kF32, + std::string(name) + ": state must be f32 (fresh-zeros or persistent, read/written f32)"); + VT_CHECK(q_in.IsContiguous() && k.IsContiguous() && v.IsContiguous() && out.IsContiguous() && + g.IsContiguous() && beta.IsContiguous() && state.IsContiguous(), + std::string(name) + ": contiguous required"); + VT_CHECK(q_in.device == q.device && k.device == q.device && v.device == q.device && + out.device == q.device && g.device == q.device && beta.device == q.device && + state.device == q.device, + std::string(name) + ": device mismatch (q/k/v/out/g/beta/state/queue)"); + VT_CHECK(args.scale > 0.0f, std::string(name) + ": args.scale must be set (> 0)"); + CheckI32Meta(q, query_start_loc, state.shape[0] + 1, name, "query_start_loc"); + reinterpret_cast(GetOp(OpId::kKdaGatedDeltaRule, q.device.type))( + q, out, q_in, k, v, g, beta, state, query_start_loc, args); +} + void GdnDecode(Queue& q, Tensor& out, const Tensor& q_in, const Tensor& k, const Tensor& v, const Tensor& g, const Tensor& beta, Tensor& state, const GdnArgs& args, const Tensor* state_idx) { diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 3d7ac4a9d..6c478ca9b 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -870,6 +870,8 @@ if(VLLM_CPP_TRITON) COMMAND test_ops_gdn "--test-case=CUDA gdn Triton AOT concurrent first load is safe across two queues") endif() +vllm_cpp_add_test(test_ops_kda_recurrence vt/test_ops_kda_recurrence.cpp) +target_include_directories(test_ops_kda_recurrence PRIVATE ${CMAKE_SOURCE_DIR}/src) vllm_cpp_add_test(test_ops_moe vt/test_ops_moe.cpp) vllm_cpp_add_test(test_ops_glue vt/test_ops_glue.cpp) vllm_cpp_add_test(test_ops_attn_preamble vt/test_ops_attn_preamble.cpp) diff --git a/tests/vt/test_ops_kda_recurrence.cpp b/tests/vt/test_ops_kda_recurrence.cpp new file mode 100644 index 000000000..81f612fae --- /dev/null +++ b/tests/vt/test_ops_kda_recurrence.cpp @@ -0,0 +1,371 @@ +// KDA per-K-channel-decay gated-delta recurrence (vt::KdaGatedDeltaRule) — UNIT GATE. +// +// The genuinely-net-new-vs-GDN device primitive for Kimi-Linear-48B: plain GDN +// decays the [Dv,Dk] recurrent state by ONE scalar per value head +// (`b_h *= exp(b_g)`); KDA decays it PER-K-CHANNEL (`b_h *= exp(b_gk[None,:])`). +// Ported 1:1 from FLA fused_recurrent_gated_delta_rule_fwd_kernel IS_KDA=True +// (third_party/flash_linear_attention/ops/fused_recurrent.py:88-175, wrapped by +// ops/kda.py:109-146 fused_recurrent_kda @ pin 555967922). +// +// ─── WHY THESE ARE THE CORRECTNESS EVIDENCE ──────────────────────────────── +// (1) EQUIVALENCE gate: with g broadcast from a per-head scalar, the per-channel +// op MUST reduce BIT-FOR-BIT to the landed+gated vt::GdnPrefill (Qwen3.6 +// 27B/35B production kernel) — the KDA op is GDN's recurrence with a decay +// VECTOR that happens to be constant. This ties the net-new op to a proven +// reference in the degenerate case with zero new numerics. +// (2) PER-CHANNEL gate: with distinct per-channel g the op is checked against a +// from-first-principles f64 reference (the exact island recurrence, +// kimi_linear_device.cpp KdaRecurrenceIsland) at a documented f32 tolerance — +// this is the ONLY place the per-channel column-wise decay is exercised. +// (3) CPU<->CUDA parity: the CUDA scan kernel matches the CPU kernel (f32 both). +#include + +#include +#include +#include +#include + +#include "vt/backend.h" +#include "vt/dtype.h" +#include "vt/ops.h" + +using vt::Backend; +using vt::Device; +using vt::DeviceType; +using vt::DType; +using vt::GdnArgs; +using vt::Queue; +using vt::Tensor; + +namespace { + +Device Cpu() { return Device{DeviceType::kCPU, 0}; } +Queue CpuQ() { return Queue{Cpu(), nullptr}; } + +Tensor MakeT(void* data, DType dt, Device dev, const std::vector& shape) { + Tensor t; + t.data = data; + t.dtype = dt; + t.device = dev; + t.rank = static_cast(shape.size()); + int64_t stride = 1; + for (int i = t.rank - 1; i >= 0; --i) { + t.shape[i] = shape[static_cast(i)]; + t.stride[i] = stride; + stride *= shape[static_cast(i)]; + } + return t; +} + +std::vector RandF32(size_t n, uint32_t seed, float lo = -1.5f, float hi = 1.5f) { + std::mt19937 rng(seed); + std::uniform_real_distribution dist(lo, hi); + std::vector v(n); + for (auto& x : v) x = dist(rng); + return v; +} + +// L2-normalize each [dim] row over its last axis (sum, not mean; eps 1e-6) — the +// caller's preprocessing contract (upstream USE_QK_L2NORM_IN_KERNEL, ops/l2norm.py). +std::vector L2NormRows(const std::vector& x, size_t rows, size_t dim) { + std::vector y(x.size()); + for (size_t r = 0; r < rows; ++r) { + double ss = 0.0; + for (size_t d = 0; d < dim; ++d) ss += static_cast(x[r * dim + d]) * x[r * dim + d]; + const double inv = 1.0 / std::sqrt(ss + 1e-6); + for (size_t d = 0; d < dim; ++d) y[r * dim + d] = static_cast(x[r * dim + d] * inv); + } + return y; +} + +// From-first-principles f64 reference for the per-channel recurrence, single +// sequence, fresh zero state. MIRRORS kimi_linear_device.cpp KdaRecurrenceIsland +// exactly: S[hv][vd][k] *= exp(g[t,hv,k]); v'=(v - S@k)*beta; S += outer(v',k); +// out = S @ (q*scale). q/k already l2-normalized by the caller. +std::vector KdaRefF64(const std::vector& qn, const std::vector& kn, + const std::vector& v, const std::vector& g, + const std::vector& beta, int64_t T, int64_t H, int64_t D, + double scale) { + const int64_t proj = H * D; + std::vector S(static_cast(H) * D * D, 0.0); + std::vector out(static_cast(T) * proj, 0.0); + std::vector u(static_cast(D)); + for (int64_t t = 0; t < T; ++t) { + for (int64_t h = 0; h < H; ++h) { + const int64_t base = t * proj + h * D; + const float* qp = &qn[static_cast(base)]; + const float* kp = &kn[static_cast(base)]; + const float* vp = &v[static_cast(base)]; + const float* gp = &g[static_cast(base)]; + const double b = beta[static_cast(t * H + h)]; + double* Sp = &S[static_cast(h) * D * D]; + for (int64_t vd = 0; vd < D; ++vd) + for (int64_t k = 0; k < D; ++k) Sp[vd * D + k] *= std::exp(static_cast(gp[k])); + for (int64_t vd = 0; vd < D; ++vd) { + double dot = 0.0; + for (int64_t k = 0; k < D; ++k) dot += Sp[vd * D + k] * kp[k]; + u[static_cast(vd)] = (static_cast(vp[vd]) - dot) * b; + } + for (int64_t vd = 0; vd < D; ++vd) + for (int64_t k = 0; k < D; ++k) Sp[vd * D + k] += u[static_cast(vd)] * kp[k]; + for (int64_t vd = 0; vd < D; ++vd) { + double o = 0.0; + for (int64_t k = 0; k < D; ++k) o += Sp[vd * D + k] * (qp[k] * scale); + out[static_cast(base + vd)] = o; + } + } + } + return out; +} + +void CheckCloseF64(const std::vector& got, const std::vector& want, float atol, + float rtol) { + REQUIRE(got.size() == want.size()); + size_t bad = 0, first = 0; + for (size_t i = 0; i < got.size(); ++i) { + const double tol = atol + rtol * std::fabs(want[i]); + if (!(std::fabs(static_cast(got[i]) - want[i]) <= tol)) { + if (bad == 0) first = i; + ++bad; + } + } + if (bad != 0) { + CAPTURE(bad); + CAPTURE(first); + CAPTURE(got[first]); + CAPTURE(want[first]); + } + CHECK(bad == 0); +} + +} // namespace + +// ── (1) EQUIVALENCE: broadcast per-channel g == the landed GdnPrefill ───────── +TEST_CASE("kda recurrence: g broadcast from per-head scalar == vt::GdnPrefill (bit-identical)") { + const int64_t T = 5, H = 3, D = 8; // Hk==Hv==H (KDA has no GQA) + const int64_t proj = H * D; + auto q = L2NormRows(RandF32(static_cast(T) * proj, 1), static_cast(T) * H, + static_cast(D)); + auto k = L2NormRows(RandF32(static_cast(T) * proj, 2), static_cast(T) * H, + static_cast(D)); + auto v = RandF32(static_cast(T) * proj, 3); + auto beta = RandF32(static_cast(T) * H, 4, 0.1f, 0.9f); + // per-head log-decay in (-1, 0] like real gates (g = -exp(A_log)*softplus(.) < 0) + auto ghead = RandF32(static_cast(T) * H, 5, -1.0f, -0.01f); + // broadcast ghead[t,h] across the D k-channels + std::vector gchan(static_cast(T) * proj); + for (int64_t t = 0; t < T; ++t) + for (int64_t h = 0; h < H; ++h) + for (int64_t d = 0; d < D; ++d) + gchan[static_cast((t * H + h) * D + d)] = ghead[static_cast(t * H + h)]; + + const int32_t qsl[2] = {0, static_cast(T)}; + const float scale = 0.35355339f; + Queue cq = CpuQ(); + + std::vector out_kda(static_cast(T) * proj, 0.0f); + std::vector st_kda(static_cast(H) * D * D, 0.0f); + { + Tensor to = MakeT(out_kda.data(), DType::kF32, Cpu(), {T, H, D}); + Tensor tq = MakeT(q.data(), DType::kF32, Cpu(), {T, H, D}); + Tensor tk = MakeT(k.data(), DType::kF32, Cpu(), {T, H, D}); + Tensor tv = MakeT(v.data(), DType::kF32, Cpu(), {T, H, D}); + Tensor tg = MakeT(gchan.data(), DType::kF32, Cpu(), {T, H, D}); + Tensor tb = MakeT(beta.data(), DType::kF32, Cpu(), {T, H}); + Tensor ts = MakeT(st_kda.data(), DType::kF32, Cpu(), {1, H, D, D}); + Tensor tqsl = MakeT(const_cast(qsl), DType::kI32, Cpu(), {2}); + vt::KdaGatedDeltaRule(cq, to, tq, tk, tv, tg, tb, ts, tqsl, GdnArgs{scale}); + } + + std::vector out_gdn(static_cast(T) * proj, 0.0f); + std::vector st_gdn(static_cast(H) * D * D, 0.0f); + { + Tensor to = MakeT(out_gdn.data(), DType::kF32, Cpu(), {T, H, D}); + Tensor tq = MakeT(q.data(), DType::kF32, Cpu(), {T, H, D}); + Tensor tk = MakeT(k.data(), DType::kF32, Cpu(), {T, H, D}); + Tensor tv = MakeT(v.data(), DType::kF32, Cpu(), {T, H, D}); + Tensor tg = MakeT(ghead.data(), DType::kF32, Cpu(), {T, H}); + Tensor tb = MakeT(beta.data(), DType::kF32, Cpu(), {T, H}); + Tensor ts = MakeT(st_gdn.data(), DType::kF32, Cpu(), {1, H, D, D}); + Tensor tqsl = MakeT(const_cast(qsl), DType::kI32, Cpu(), {2}); + vt::GdnPrefill(cq, to, tq, tk, tv, tg, tb, ts, tqsl, GdnArgs{scale}); + } + + // Bit-identical: the per-channel decay vector is constant == the scalar decay, + // so the fused f32 arithmetic is the same op stream (exact float equality). + size_t out_diff = 0, st_diff = 0; + for (size_t i = 0; i < out_kda.size(); ++i) + if (out_kda[i] != out_gdn[i]) ++out_diff; + for (size_t i = 0; i < st_kda.size(); ++i) + if (st_kda[i] != st_gdn[i]) ++st_diff; + CHECK(out_diff == 0); + CHECK(st_diff == 0); +} + +// ── (2) PER-CHANNEL: distinct per-channel decay vs the f64 island reference ─── +TEST_CASE("kda recurrence: distinct per-channel decay matches the f64 reference") { + const int64_t T = 8, H = 4, D = 16; + const int64_t proj = H * D; + auto q = L2NormRows(RandF32(static_cast(T) * proj, 11), static_cast(T) * H, + static_cast(D)); + auto k = L2NormRows(RandF32(static_cast(T) * proj, 12), static_cast(T) * H, + static_cast(D)); + auto v = RandF32(static_cast(T) * proj, 13); + auto beta = RandF32(static_cast(T) * H, 14, 0.1f, 0.9f); + auto gchan = RandF32(static_cast(T) * proj, 15, -0.8f, -0.01f); // distinct per channel + const int32_t qsl[2] = {0, static_cast(T)}; + const double scale = std::pow(static_cast(D), -0.5); + + std::vector out(static_cast(T) * proj, 0.0f); + std::vector st(static_cast(H) * D * D, 0.0f); + Queue cq = CpuQ(); + Tensor to = MakeT(out.data(), DType::kF32, Cpu(), {T, H, D}); + Tensor tq = MakeT(q.data(), DType::kF32, Cpu(), {T, H, D}); + Tensor tk = MakeT(k.data(), DType::kF32, Cpu(), {T, H, D}); + Tensor tv = MakeT(v.data(), DType::kF32, Cpu(), {T, H, D}); + Tensor tg = MakeT(gchan.data(), DType::kF32, Cpu(), {T, H, D}); + Tensor tb = MakeT(beta.data(), DType::kF32, Cpu(), {T, H}); + Tensor ts = MakeT(st.data(), DType::kF32, Cpu(), {1, H, D, D}); + Tensor tqsl = MakeT(const_cast(qsl), DType::kI32, Cpu(), {2}); + vt::KdaGatedDeltaRule(cq, to, tq, tk, tv, tg, tb, ts, tqsl, GdnArgs{static_cast(scale)}); + + const std::vector ref = KdaRefF64(q, k, v, gchan, beta, T, H, D, scale); + // f32 recurrence over 8 tokens vs f64 reference: documented tolerance. + CheckCloseF64(out, ref, /*atol=*/1e-4f, /*rtol=*/3e-3f); +} + +// ── (3) validation: g must be per-channel [T,Hv,Dk], scale must be set ──────── +TEST_CASE("kda recurrence: validation rejects per-head g and unset scale") { + const int64_t T = 2, H = 1, D = 2; + std::vector buf(static_cast(T) * H * D, 0.1f); + std::vector beta(static_cast(T) * H, 0.5f); + std::vector st(static_cast(H) * D * D, 0.0f); + const int32_t qsl[2] = {0, static_cast(T)}; + Queue cq = CpuQ(); + Tensor to = MakeT(buf.data(), DType::kF32, Cpu(), {T, H, D}); + Tensor tq = MakeT(buf.data(), DType::kF32, Cpu(), {T, H, D}); + Tensor tk = MakeT(buf.data(), DType::kF32, Cpu(), {T, H, D}); + Tensor tv = MakeT(buf.data(), DType::kF32, Cpu(), {T, H, D}); + Tensor tb = MakeT(beta.data(), DType::kF32, Cpu(), {T, H}); + Tensor ts = MakeT(st.data(), DType::kF32, Cpu(), {1, H, D, D}); + Tensor tqsl = MakeT(const_cast(qsl), DType::kI32, Cpu(), {2}); + // per-head g [T,H] is rank-2 -> rejected (kda wants [T,H,D]) + std::vector ghead(static_cast(T) * H, -0.1f); + Tensor tg_head = MakeT(ghead.data(), DType::kF32, Cpu(), {T, H}); + CHECK_THROWS(vt::KdaGatedDeltaRule(cq, to, tq, tk, tv, tg_head, tb, ts, tqsl, GdnArgs{0.5f})); + // unset scale + Tensor tg = MakeT(buf.data(), DType::kF32, Cpu(), {T, H, D}); + CHECK_THROWS(vt::KdaGatedDeltaRule(cq, to, tq, tk, tv, tg, tb, ts, tqsl, GdnArgs{})); +} + +// ── (4) CPU<->CUDA parity ───────────────────────────────────────────────────── +#ifdef VLLM_CPP_CUDA +namespace { +Device Gpu() { return Device{DeviceType::kCUDA, 0}; } +bool HasCuda() { + try { + vt::GetBackend(DeviceType::kCUDA); + return true; + } catch (const std::runtime_error&) { + return false; + } +} +struct QGuard { + Backend& b; + Queue q; + explicit QGuard(Backend& backend) : b(backend), q(backend.CreateQueue()) {} + ~QGuard() { b.DestroyQueue(q); } + QGuard(const QGuard&) = delete; + QGuard& operator=(const QGuard&) = delete; +}; +class DTensor { + public: + DTensor(Backend& b, Queue& q, DType dt, const std::vector& shape, + const void* host = nullptr) + : b_(b) { + int64_t numel = 1; + for (auto s : shape) numel *= s; + bytes_ = static_cast(numel) * vt::SizeOf(dt); + p_ = b_.Alloc(bytes_ == 0 ? 1 : bytes_); + if (host != nullptr) b_.Copy(q, p_, host, bytes_); + t_ = MakeT(p_, dt, Gpu(), shape); + } + ~DTensor() { b_.Free(p_); } + DTensor(const DTensor&) = delete; + DTensor& operator=(const DTensor&) = delete; + Tensor& tensor() { return t_; } + void Download(Queue& q, void* dst) { + b_.Copy(q, dst, p_, bytes_); + b_.Synchronize(q); + } + + private: + Backend& b_; + void* p_ = nullptr; + size_t bytes_ = 0; + Tensor t_; +}; +} // namespace + +TEST_CASE("kda recurrence: CUDA scan matches the CPU kernel (f32)") { + if (!HasCuda()) return; + const int64_t T = 7, H = 4, D = 32; + const int64_t proj = H * D; + auto q = L2NormRows(RandF32(static_cast(T) * proj, 21), static_cast(T) * H, + static_cast(D)); + auto k = L2NormRows(RandF32(static_cast(T) * proj, 22), static_cast(T) * H, + static_cast(D)); + auto v = RandF32(static_cast(T) * proj, 23); + auto beta = RandF32(static_cast(T) * H, 24, 0.1f, 0.9f); + auto gchan = RandF32(static_cast(T) * proj, 25, -0.8f, -0.01f); + const int32_t qsl[2] = {0, static_cast(T)}; + const float scale = std::pow(static_cast(D), -0.5f); + + // CPU + std::vector out_cpu(static_cast(T) * proj, 0.0f); + std::vector st_cpu(static_cast(H) * D * D, 0.0f); + { + Queue cq = CpuQ(); + Tensor to = MakeT(out_cpu.data(), DType::kF32, Cpu(), {T, H, D}); + Tensor tq = MakeT(q.data(), DType::kF32, Cpu(), {T, H, D}); + Tensor tk = MakeT(k.data(), DType::kF32, Cpu(), {T, H, D}); + Tensor tv = MakeT(v.data(), DType::kF32, Cpu(), {T, H, D}); + Tensor tg = MakeT(gchan.data(), DType::kF32, Cpu(), {T, H, D}); + Tensor tb = MakeT(beta.data(), DType::kF32, Cpu(), {T, H}); + Tensor ts = MakeT(st_cpu.data(), DType::kF32, Cpu(), {1, H, D, D}); + Tensor tqsl = MakeT(const_cast(qsl), DType::kI32, Cpu(), {2}); + vt::KdaGatedDeltaRule(cq, to, tq, tk, tv, tg, tb, ts, tqsl, GdnArgs{scale}); + } + + // CUDA + std::vector out_gpu(static_cast(T) * proj, 0.0f); + std::vector st_gpu(static_cast(H) * D * D, 0.0f); + { + Backend& gpu = vt::GetBackend(DeviceType::kCUDA); + QGuard g(gpu); + std::vector st_zero(static_cast(H) * D * D, 0.0f); + DTensor to(gpu, g.q, DType::kF32, {T, H, D}); + DTensor tq(gpu, g.q, DType::kF32, {T, H, D}, q.data()); + DTensor tk(gpu, g.q, DType::kF32, {T, H, D}, k.data()); + DTensor tv(gpu, g.q, DType::kF32, {T, H, D}, v.data()); + DTensor tg(gpu, g.q, DType::kF32, {T, H, D}, gchan.data()); + DTensor tb(gpu, g.q, DType::kF32, {T, H}, beta.data()); + DTensor ts(gpu, g.q, DType::kF32, {1, H, D, D}, st_zero.data()); + DTensor tqsl(gpu, g.q, DType::kI32, {2}, qsl); + vt::KdaGatedDeltaRule(g.q, to.tensor(), tq.tensor(), tk.tensor(), tv.tensor(), tg.tensor(), + tb.tensor(), ts.tensor(), tqsl.tensor(), GdnArgs{scale}); + to.Download(g.q, out_gpu.data()); + ts.Download(g.q, st_gpu.data()); + } + + // Same f32 math; different exp libm / FMA contraction -> tiny arithmetic-order gap. + size_t bad = 0; + for (size_t i = 0; i < out_cpu.size(); ++i) + if (std::fabs(out_gpu[i] - out_cpu[i]) > 1e-4f + 2e-3f * std::fabs(out_cpu[i])) ++bad; + CHECK(bad == 0); + size_t bad_s = 0; + for (size_t i = 0; i < st_cpu.size(); ++i) + if (std::fabs(st_gpu[i] - st_cpu[i]) > 1e-4f + 2e-3f * std::fabs(st_cpu[i])) ++bad_s; + CHECK(bad_s == 0); +} +#endif // VLLM_CPP_CUDA