feat(gguf,dsv4): GGUF block-dequant + DSv4 architecture kernel surface

[TheTom]

Summary

Companion to thewafflehaus/FFAI#17 — adds the GPU kernels FFAI's GGUF loader + DSv4 family forward path dispatch into. Started as the three GGUF block-dequant kernels; grew to cover the full DSv4-Flash architecture surface as the FFAI side caught up.

GGUF block-dequant kernels

• `ffai_gguf_dequant_q8_0` — 1-byte-per-weight + fp16 block-scale, block=32
• `ffai_gguf_dequant_q2_K` — 2-bit + two-level (4-bit min + 4-bit scale) scales, block=256
• `ffai_gguf_dequant_iq2_xxs` — 256-codebook (`iq2xxs_grid[256][8]`) + ksigns sign-mask LUT, block=256

All three generic over `T = f32 / f16 / bf16`. Host produces the GPU-resident split (packed quants + per-block scales + LUT tables), kernel does `out[i] = scale_block × codebook_decode_at(packed_bits)`. Output dtype follows `T` via implicit Store coercion.

DeepSeek-V4 architecture kernels

• `ffai_moe_router_sqrtsoftplus` — V4's `scoring_func: sqrtsoftplus` MoE router (replaces V3's sigmoid+bias gate). Numerically-stable softplus (`max(x,0) + log(1 + exp(-|x|))`) + sqrt + noaux_tc bias correction.
• `ffai_dsv4_mxfp4_dequant` — OCP FP4 e2m1 (BPW=4.25) with E8M0 scale, 16-LUT.
• `ffai_dsv4_fp8_block_dequant` — FP8 e4m3 (1 B/weight) + per-(128×128)-block fp32 scales, 256-LUT.
• `ffai_dsv4_compressor_pool` — CSA / HCA softmax-gated weighted pool with APE add. Single-pass online-softmax — the three-sequential-for-loop form tripped a DSL codegen bug (see "DSL gotchas" below).
• `ffai_sdpa_decode_d512_sink` — clone of `sdpa_decode_d512` with GPT-OSS-style per-head learnable softmax sink folded into the cross-simdgroup denominator. Used by DSv4 full-attention layers (the `attn_sinks` parameter is per-head learnable, all layers carry it).
• `ffai_dsv4_mhc_sinkhorn_split` — converts `hc_*_fn @ flatten(H)` 24-mix into `(pre, post, comb)` per-token control tensors. 4×4 comb is softmax-then-Sinkhorn-Knopp normalized; iters constexpr.
• `ffai_dsv4_mhc_collapse` / `ffai_dsv4_mhc_expand` — runtime 4-channel residual mix kernels. The DSv4 mHC pre/post/comb are DYNAMIC per-token (computed by sinkhorn_split above), not stored model weights — initial static pre/post placeholder kernels were wrong and replaced.
• `ffai_dsv4_partial_rope` — head-dim-tail-only RoPE rotation (DSv4 splits per-head dim as nope=448 / rope=64). Forward + inverse modes (inverse needed on attn output before grouped O-LoRA since K == V in MQA mode and V is effectively double-rotated).
• `ffai_dsv4_indexer_score` — Lightning Indexer per-position aggregate: `score[t] = Σ_h w[h] · ReLU(q_idx[h] · k_idx[t, h])`.
• `ffai_dsv4_indexer_topk_block` — single-block bitonic top-K over score[n_kv≤1024] with parallel `(score, original_index)` tag survival.
• `ffai_dsv4_csa_sdpa_decode` — sparse-gather SDPA over a caller-supplied gather list (Lightning Indexer top-K ∪ sliding-window). Clone of d512 with the dense inner loop replaced by an index-gather load.

Plus utility: `ffai_vector_add` (generic elementwise out[i] = a[i] + b[i]) — `Ops.add` already existed but routed through a different binary kernel; this one is used by Ops.add's f32/f16/bf16 paths for the per-expert MoE accumulator.

Validation

All 14 kernels have CPU oracle + GPU correctness tests in `kernel_tests_harness`. Tolerances within the standard f32 / f16 / bf16 bands (elementwise 1e-4 / 5e-3 / 5e-2; SDPA 1e-3 / 3e-3 / 1.5e-2). Validated end-to-end via the FFAI integration tests against the real 86 GB DSv4-Flash GGUF on M5 Max.

DSL gotchas surfaced + documented

Three codegen pitfalls hit while writing the DSv4 kernels (notes in `feedback_metaltile_dsl_codegen_pitfalls.md` in the project memory):

1. Sequential `for _w in range(...)` loops sharing the loop index name — DSL collapses them and references variables out of scope. Fuse into a single online-softmax-style loop. Burned ~30min on the compressor pool before pattern-matching.
2. Chained method-form `.exp() / .ln() / .sqrt()` inside larger expressions drop intermediate bindings. Use the free-function forms `exp(...) / log(...) / sqrt(...)` and single-op `let` lines. Caught on `moe_router_sqrtsoftplus`.
3. Let-binding names that prefix-collide with Tensor-parameter names (`score` vs `score_unbiased`) silently elide the local. Rename locals to non-prefix-overlapping identifiers.

IR codegen tests (`every_registered_benchspec_codegens`) pass for all three pitfalls, but real-GPU dispatch via `kernel_tests_harness` fails with `MSL library compilation failed: use of undeclared identifier vXX`. Fix patterns documented for future kernel work.

Target

Cross-fork PR: `TheTom/metaltile:tom/wip/gguf-dequant` → `0xClandestine/metaltile:dev`. Branch is on a personal fork because cross-repo CI pins in the FFAI side need to point at a stable target (org branch).

[TheTom]

Status — review pass complete

Audited this branch against the conventions in #200+ / CONTRIBUTING / docs and fixed the findings. 3 logical commits (fix(macros) · feat(codegen) · feat(ffai)), rebased on current dev.

Changes from the review:

• +9 GPU correctness tests for the previously-untested kernels (grouped_gemm_q8, gemm_q8_mpp + grouped, grouped_gemv_q8_rows_tiled, moe_gemv_rows_view_iq2xxs u8+u16, moe_gemv_ws_iq2xxs, moe_gemv_ws_q2k, moe_bgemm_iq2xxs_view_sg32, mt_remap_u32) — each vs a CPU reference, all passing (cosine 1.0 / relDiff ≤ 3e-5). Every kernel now ships its test in the kernel commit, per CONTRIBUTING.
• + emit unit test pinning the new dispatchThreadgroups Swift-wrapper output (binding-by-slot + dispatch shape).
• Removed debug eprintln! from the runtime bind hot path; normalized copyright headers to the repo's multi-author form.

Green locally: make fmt-check, make typos, full-workspace make clippy -D warnings, and the new + existing GPU tests.

One open design question for @ekryski above (gating the dispatchThreadgroups variant). DSv4 kernel surface backs FFAI DeepSeek-V4-Flash prefill (thewafflehaus/FFAI#17).

[0xClandestine]

Should have a constant somewhere to avoid brittle match statement.

[TheTom]

Good call. Switched is_float/is_integer from a non-exhaustive matches! to an exhaustive match — now adding a DType variant fails to compile until it's classified in both, so they can't silently disagree or miss a new type (which is exactly how U16 could've slipped through). Went with the compiler-enforced match rather than a shared const since that catches the omission at build time instead of relying on the list staying in sync; happy to swap to a const set instead if you'd rather have one source of truth to reuse elsewhere. Pushed.

[ekryski]

I might already have ffai_sdpa_decode_d512_sink in PR #250. Will check.

[TheTom]

@ekryski — one design decision I'd like your call on before this merges:

dispatchThreadgroups Swift-wrapper variant — gate it, or keep it unconditional?

This branch adds emit_swift_wrapper_threadgroups to codegen, currently emitted for every kernel. Only the coop_tile / simdgroup-matrix / bm64 MoE kernels actually need a threadgroups-grid launch; elementwise and reduction kernels never use it. The sibling emit_swift_wrapper_indirect is gated behind a per-kernel wants_indirect_variant flag, so the consistent design would be a matching wants_threadgroups_variant set only on the MMA/coop_tile kernels.

I left it unconditional because the generated binding surface is already consumed broadly on the FFAI side, so gating is a cross-repo change. Which do you prefer — gate it (I'll set the flag on the MMA/bm64 kernels and coordinate the FFAI binding regen), or keep it emitted for all?

(FYI, unrelated: I pulled an opportunistic BTreeMap→FxHashMap codegen-pass refactor out of this PR to keep it scoped — it'll come as its own perf(codegen): PR with bench rows.)

[TheTom]

Merge coordination — assuming #247 lands first

Updating my earlier note: since ek/bench-metrics (#247) goes in ahead of this PR, #243 will rebase onto a dev that already has #247 (and likely #251, already merged). Plan so the rebase is clean:

Conflict surface is small — two files:

• crates/metaltile/src/runner/gpu.rs — feat(bench): wall-clock latency, GFLOP/s, roofline %-of-peak + bottleneck verdict #247 adds the profiling/roofline path here; my only edit is U16 joining the existing DType::I4 | DType::U64 | DType::I64 => unimplemented!() arm in buffer_typed/read_typed. Localized, one-line resolve on rebase.
• crates/metaltile-std/src/ffai/mod.rs — pub mod additions, alphabetical → trivial auto-merge (this is the overlap with feat(quant): complete op × quant-format matrix — 30 Track-1 formats across every weight-bearing kernel #250/feat(ffai): vision/audio/TTS/video kernel batch + kernel_tests_harness dead-strip fix #252, not feat(bench): wall-clock latency, GFLOP/s, roofline %-of-peak + bottleneck verdict #247).

Structural note (unchanged): #243 adds DType::U16 to the core enum — any new exhaustive DType match on a branch rebasing after #243 needs a U16 arm (fully wired everywhere else already).

Post-#247-rebase to-dos I'll do (depend on #247's .flops() API, so not in this PR yet):

• Add .flops(...) to the DSv4 compute-kernel benches so they populate the new GFLOP/s + roofline columns: MoE GEMMs (2·m_total·n_out·k_in), Q8 GEMMs + grouped O-LoRA (2·rows·out·in), SDPA d512-sink prefill/decode (4·q_heads·tokens·dim). Memory-bound kernels (partial-rope, swiglu, router, dequant) stay blank, per the spec.
• Normalize the DSv4 bench names to the ffai/<op>/<subop> scheme while I'm in there.

Cleanup since the last push: removed two dead negative-result kernels (moe_bgemm_iq2xxs_view_sg, moe_bgemm_iq2xxs_view_sg32 — 0 dispatch refs, lose to bm64) and their tests, to de-dup the MoE-GEMM surface. The *_mpp gather-BGEMMs stay — they're the live prefill kernels + the oracle for the bm64/view correctness tests.

[TheTom]

@ekryski re ffai_sdpa_decode_d512_sink — checked #250 for you: it doesn't add it (only sdpa_decode_d512.rs, the non-sink). The d512 sink variant is unique to this PR, so no collision. (For reference: #252's sdpa_decode_sink_buf is also distinct — that's the GPT-OSS buffer-sink, mine is the DSv4 scalar-sink at head_dim=512.)

[TheTom]

@ekryski CI is fully green now ✅ (Test 39s — the 57 NAX/MPP cooperative-tensor tests correctly [SKIP] via your #258, no longer failing).

Heads-up on a real bug I found and fixed while getting here — worth knowing about:

Q2_K dequant test was silently validating the wrong layout

test_gguf_q2_k_single_block / _many_blocks were failing with err≈0.2 (dtype-invariant → structural, not precision). Root cause: the test was internally inconsistent.

• The GPU kernel reads quants with the correct canonical llama.cpp Q2_K layout — the dequantize_row_q2_K walk: half/jg/sub_half → q_byte = half*32 + sub_half*16 + l, shift = jg*2. I re-verified it against llama.cpp; it's right.
• But the test's quantizer wrote quants in the naive byte = i/4, shift = (i&3)*2 layout, and the CPU reference read them naive too. So quantizer + reference agreed with each other but disagreed with the kernel → kernel pulled quant bits from the wrong byte/shift.

So the production kernel was always correct — the test scaffolding was wrong, which is the dangerous kind: a green-looking harness that doesn't exercise the real on-disk layout. Fix: a shared q2_k_qpos() helper so quantizer, CPU reference, and kernel all use the identical canonical mapping. Now err=0.00e0 exact. (Scales were already fine — canonical scale_idx == i/16.)

It hid from my local runs because my filter (tile test ffai) didn't match test_gguf_q2_k_* — only surfaced once the full-suite Test job ran it. Lesson logged on my end.

Also in this push: bumped test_axpy_decode_shape f32 tol 1e-3→5e-3 (k=2048 simdgroup tree-reduce vs sequential CPU sum — reduction-order gap, not logic).

And a correction to my earlier NAX comment on #247

I was wrong that those were M5-hardware-gated. Your #256/#258 archaeology shows it's the macOS 26.4 runtime Metal compiler choking on Apple's own MPP cooperative_tensor (deferred-static-alloca-size); they pass on a plain M1 Max + 26.5. #258's build-failure→SKIP handles it and self-re-enables when the runner image hits 26.5. Nothing hardware about it.

[ekryski]

Merging! We'll do two follow up PRs once this and #252 merge.

1. Ensure all new kernels support precisions we added in feat(quant): complete op × quant-format matrix — 30 Track-1 formats across every weight-bearing kernel #250 (where required)
2. A clean up pass to clean up file structure and reduce loc. Will be dependent on refactor: dissolve metaltile-std::bench_types and trim cross-crate deps #260 and feat(macros): #[kernel(variants(...))] compile-time kernel specialisation #261 (and other PRs).