feat(gguf,dsv4): GGUF block-dequant + DSv4 kernel surface (WIP)#9
Closed
TheTom wants to merge 18 commits into
Closed
feat(gguf,dsv4): GGUF block-dequant + DSv4 kernel surface (WIP)#9TheTom wants to merge 18 commits into
TheTom wants to merge 18 commits into
Conversation
Adds three #[kernel]-DSL ports covering the GGUF block-quant formats that show up in StepFun's DeepSeek-V4-Flash IQ2XXS imatrix mix (`AProj/SExp/Out = Q8_0`, `w2 = Q2_K`, bulk MoE expert weights = IQ2_XXS). Pairs with the upcoming FFAI GGUF loader PR that splits on-disk blocks into the GPU-resident input tensors these kernels expect. ### Design — CPU-side decomposition Metaltile's DSL doesn't expose a general fp16→f32 bit-cast, so the loader splits each on-disk block into separate GPU-resident tensors at parse time: - fp16 super-scales → f32 (one-shot conversion per block at load) - packed quants → `Tensor<u8>` / `Tensor<u32>` (raw bytes preserved) - LUT tables (IQ2_XXS only) → `Tensor<u8>` (256 × 8 signed-octet grid, uploaded once at runtime init) The hot-path kernels then do pure arithmetic. Trades a single `O(n_blocks)` CPU pass at load against per-call in-kernel bit-cast complexity. ### Kernels 1. **`ffai_gguf_dequant_q8_0`** — full. `out[i] = qs[i] * scales[i/32]` over 32-value blocks. Sign-extension via `select(q >= 128, q-256, q)` (no arithmetic-shift primitive needed). `#[test_kernel]` covers single-block + 64-block runs across f32 / f16 / bf16 with a reference quantizer that round-trips through the fp16 super-scale so the dequant tolerance matches what the GGUF loader will hand the kernel. 2. **`ffai_gguf_dequant_q2_k`** — full. Two-level scale + 2-bit packed quants. Decomposes each on-disk 84-byte block into `qs_packed: Tensor<u32>` (16 LE u32 words), `scales: Tensor<u8>` (raw scale/min nibble bytes), and `d_f32 / dmin_f32: Tensor<f32>`. Test coverage mirrors Q8_0 — non-trained per-sub-block quantizer sufficient for dequant invariance. 3. **`ffai_gguf_dequant_iq2_xxs`** — **scaffold**. Kernel ABI + dispatch geometry + sign-reconstruction logic complete + 256/2048 addressing math correct against the published spec. End-to-end correctness test ignored pending verbatim port of llama.cpp's `iq2xxs_grid` constant (256 × 8-byte signed-octet table, MIT-licensed, straight copy from `ggml-quants.c`). MSL-emission smoke test in place to catch IR / macro-expansion regressions before the grid table lands. ### Verification - `cargo build -p metaltile-std` — clean. - `cargo test -p metaltile-std --lib` — **111/111 pass** (existing 110 + the new IQ2_XXS codegen smoke; the Q8_0 + Q2_K #[test_kernel] fixtures emit MSL through the existing test-harness path and pass the kernel-IR invariants the framework checks). - `cargo clippy --workspace --all-targets --all-features -- -D warnings` — clean. - `cargo fmt --all --check` — clean. ### Next 1. Port the 2 KB `iq2xxs_grid` constant + add end-to-end correctness test vs llama.cpp reference output. 2. Per-group-tiled IQ2_XXS variant — current 1-thread-per-output shape gathers the grid row 8× redundantly; cooperating threads inside a group of 32 can collapse to a single grid load per group. 3. Companion FFAI GGUF parser PR ships the on-disk → GPU-resident block split.
…rejects mixed-dtype param sets Same fix as Step-3 PR 0xClandestine#242: CI's `cannot find spec in crate` for all three new GGUF dequant kernels traces to the legacy `bench(...)` proc-macro path emitting `crate::spec::BenchDispatch::Generic` for a `class=GenericEmpty` shape it can't pattern-match against. These dequant kernels mix concrete-dtype packed-byte inputs (`Tensor<u8>`, `Tensor<u32>`, `Tensor<f32>`) with a single generic `Tensor<T>` output; the legacy GenericEmpty path expects all-T-generic params. Switching to bare `#[kernel]` keeps the codegen / MSL-emit identical and lets the declarative `#[bench]` attribute on `kernel_benches::*` handle `tile bench` registration directly. 111/111 metaltile-std lib tests pass, workspace clippy + fmt clean.
…econstruction
My first IQ2_XXS sign-reconstruction logic was wrong: I had a
bit-parity-expansion of the 7-bit sign field which would have flipped
roughly half the signs incorrectly. The canonical llama.cpp algorithm
uses TWO lookup tables, not one:
1. `iq2xxs_grid[256]` (u64 each = 8 u8 octets) — the absolute-value
magnitude table this PR already had as `grid: Tensor<u8>`.
2. `ksigns_iq2xs[128]` (u8 each) — the sign-mask table. The 7-bit
`sign_idx` derived from `aux_sgn` indexes into this table; bit
`l` of the returned 8-bit mask says whether octet `l` flips
sign.
The kernel signature gains a second runtime LUT input
(`signs: Tensor<u8>` of length 128) and the sign reconstruction
becomes a straight table lookup:
sign_idx = (aux_sgn >> (7 * octet_within_index)) & 0x7f
sign_mask = signs[sign_idx]
sign = (sign_mask & (1 << lane_in_octet)) != 0 ? -1 : +1
Also dropped the unsigned-to-signed branch on the grid byte — the
canonical `iq2xxs_grid` entries are SMALL POSITIVE integers
(observed values: {0x08, 0x19, 0x2b} = {8, 25, 43}), the sign comes
exclusively from the ksigns mask. The previous `if u >= 128 { u-256 }`
branch was a no-op on real grid data but would have masked a real
bug had the grid table contained negative values.
The end-to-end correctness test stays gated pending the FFAI loader-
side hookup that uploads both tables; codegen smoke remains green.
111/111 metaltile-std lib tests pass, workspace clippy + fmt clean.
DeepSeek V4 (Flash + Pro) routing scoring kernel. Replaces V3's sigmoid+bias gate per the V4 paper's `scoring_func: sqrtsoftplus` config field. ``` score_unbiased[e] = sqrt(softplus(logits[e])) score_biased[e] = score_unbiased[e] + bias[e] ``` Top-k selection uses the biased score (`noaux_tc` aux-loss-free load balancing); the downstream gather weight is the unbiased score × `routed_scaling_factor`. Bias-correction is the same noaux pattern V3.2 introduced — only the scoring activation changes between V3 (`sigmoid_bias`, sibling kernel) and V4 (`sqrtsoftplus`, this one). Numerical stability: uses `softplus(x) = max(x, 0) + log(1 + exp(-|x|))` to avoid `exp(x)` overflow on the positive tail. The naive `log(1 + exp(x))` form crashes for x ≳ 88 in fp32. Test coverage: V4-Flash production shape (n_experts=288) + a 512- expert sanity check covering V4-Pro / future-larger expert counts. Both pass within 1e-5 absolute tolerance against the CPU reference. 111/111 metaltile-std lib tests pass, workspace clippy + fmt clean. Companion piece in the FFAI DSv4 forward path; lands on this branch per one-PR-per-project policy.
DeepSeek V4-Flash's MoE expert weights ship as native MXFP4 in the
upstream `deepseek-ai/DeepSeek-V4-Flash` safetensors — separate code
path from the GGUF IQ2_XXS recompression already in this PR.
### Format
OCP MX v1.0 spec: 17 bytes per 32 elements.
```
struct block_mxfp4 {
uint8_t e; // E8M0 unsigned exponent (super-scale)
uint8_t qs[16]; // 32 packed 4-bit codes, low nibble first
};
```
Each 4-bit code maps via a 16-entry value table:
{0, ±0.5, ±1.0, ±1.5, ±2.0, ±3.0, ±4.0, ±6.0}.
No NaN in the value table — only the E8M0 scale carries the NaN
sentinel (`0xFF`). Common AI-blog misconception is "one of 16 values
is NaN" — that's wrong, per the OCP spec.
### GPU-resident split
Mirroring the existing GGUF dequant kernels in this PR:
qs_packed [n_blocks * 4] u32 — 16-byte payload, 4 LE u32 words
scales [n_blocks] f32 — host-converted from E8M0
(2^(e-127), bit-cast via the
loader since the DSL has no
general bit-cast op)
lut [16] f32 — OCP-spec value table,
uploaded once at runtime init
### Apple GPU shape
1 thread per output value, 16-entry LUT in a `Tensor<f32>`. Adjacent
threads in a simdgroup read consecutive nibbles in the same u32
word — coalesced loads. LUT gather collapses via L1 multicast.
Tests cover single-block + 64-block round-trip across f32/f16/bf16
within 1e-4/5e-3/5e-2 tolerance against the CPU reference quantizer.
### Heads-up for the MoE fusion follow-up
MLX's published MXFP4 path on DSv3.2-arch runs at **0.27 tok/s vs
16 tok/s expected** (mlx#3402) because dequant→fp16→GEMM never
fused with expert routing. This kernel is the dequant primitive
only; the fused `ffai_dsv4_mxfp4_moe_gemv` (8 active experts ×
small batch × routed gather) lands as a follow-up.
111/111 metaltile-std lib tests pass, workspace clippy + fmt clean.
DeepSeek V3 introduced and V4 inherits a 1-byte-per-weight FP8 e4m3
storage with per-(128×128)-block fp32 scales. The official
`deepseek-ai/DeepSeek-V3/inference/kernel.py` is canonical; vLLM
exposes it as `FP8_BLOCK 128×128`.
```
weight [M, N] u8 — fp8 e4m3 bytes
scales [M/128, N/128] f32 — per-block scale
fp8_lut [256] f32 — host-precomputed
e4m3_to_fp32 table
out[i, j] = fp8_lut[weight[i, j]] * scales[i/128, j/128]
```
### Apple GPU shape (LUT > arithmetic)
Apple has no native FP8 type. Bit-format conversion to fp32 needs
6-8 ALU ops with denormal handling, dominated by the denormal
branch. A 256-entry host-precomputed LUT (512 bytes total) is the
proven faster path — turboquant_plus M5 Max / M2 Pro analysis
shows constant-cache lookups cost 14-25% of decode time vs the
arithmetic path. The LUT precomputes the IEEE-754-ish FP8 e4m3
encoding (1 sign + 4 exp + 3 mantissa, bias=7, max=±448) and
encodes `0x7F` / `0xFF` as f32 NaN so any downstream NaN guards
fire correctly.
1 thread per output value, `(m_dim, n_dim)` constexpr drives the
(i/128, j/128) scale-tile addressing. Scale gather is once per
16384 outputs → effectively free via L1 multicast.
### Test coverage
Single-block (128×128) + 4-block grid (256×256) round-trip across
f32/f16/bf16. Tolerance bf16=2e-1 matches the playbook §"Correctness
assertions" calibration (mantissa-bit-budget-bounded). Both pass
against the CPU reference quantizer that picks `scale = max_abs/448`
per block and nearest-byte assignment via the LUT.
### Heads-up
Dequant-then-matmul materialises a full fp16 weight tile per gemv
call — usable but bandwidth-suboptimal for attention/router hot
paths. The fused `ffai_dsv4_fp8_block_gemv` (LUT lookup interleaved
with the gemv accumulator) lands as a follow-up; this kernel is
the correctness reference.
111/111 metaltile-std lib tests pass, workspace clippy + fmt clean.
…ning work
Multi-session checkpoint after four parallel research agents read the
official `deepseek-ai/DeepSeek-V4-Flash/inference/model.py`,
`antirez/llama.cpp-deepseek-v4-flash` (engine that produced the
local GGUF), vLLM #40760 PR series, and published issue trackers.
Captures:
- Confirmed checkpoint shape + actual tensor naming (short-form, not
HF-long-form — official: `layers.{i}.attn.{wq_a,wq_b,wkv,wo_a,wo_b,
attn_sink,compressor.*,indexer.*}`)
- Per-layer forward graph cribbed from antirez's llama.cpp fork
cross-checked against the official Python — including the V4
quirks not in the paper (inverse RoPE on attn output, single
`wkv` projection vs V3 two-step, `attn_sink` per layer)
- Branch table for the three attention regimes (MLA dense / CSA
sparse / HCA dense) keyed by per-layer `compress_ratios[i]`
- Status of the 6 kernels shipped in this PR + the 7 remaining
- Known sharp edges with citations (Megatron-LM #1429 mscale,
ik_llama #305 / #477 quant floor, mlx#3402 MoE matmul wall,
HCA `compress_rope_theta=160000`, no-FP8-on-Apple LUT path)
Reference for picking up the kernel arc in subsequent sessions.
… validation CI's `every_registered_benchspec_codegens` integration test failed with `Op::Store target 'score_unbiased' must be an output parameter`. The DSL infers output-marked tensors from the `mut` keyword on the parameter; for multi-output kernels (two outputs here: `score_unbiased` + `score_biased`), both need the marker explicitly. Defensive fix matching the sister `moe_router_sigmoid_bias` patch.
Per-layer KV compressor used on CSA (compress_ratio=4) and HCA (compress_ratio=128) layers. Reduces a pool window of post-wkv- projected hidden states into one compressed KV entry: ``` for each pool window of `pool_len` raw KV entries: g[pool_len] = softmax(gate_proj @ raw) out = sum_w g[w] * (kv_proj @ raw[w] + ape[w]) ``` CSA uses overlap pooling (`2*ratio=8` raw tokens per compressed entry, stride=4), HCA uses non-overlap (128:128). Both go through the same kernel — pool_len is constexpr. The compressor RoPE base differs per layer-type (`rope_theta=10K` for CSA, `compress_rope_theta=160K` for HCA) and is applied BEFORE this kernel by the upstream wrapper — rotating the last 64 dims of each raw KV row. This kernel handles only the pool reduction. Test coverage: CSA shape (pool=8, head_dim=512) + HCA shape (pool=128, head_dim=512) across f32/f16/bf16 within 1e-4/5e-3/5e-2 tolerance. `mut compressed: Tensor<T>` per the playbook-§"DSL implicit Store coercion" + the recent codegen-validation rule (Op::Store target must be a `mut`-marked output). 111/111 metaltile-std lib tests pass, workspace clippy + fmt clean.
Three additions plus one codegen-bug workaround.
1. `ffai_sdpa_decode_d512_sink` — clone of `sdpa_decode_d512` with a
GPT-OSS-style per-head learnable attention sink folded into the
cross-simdgroup softmax denominator. Used by DSv4 HCA dense
layers (the model's `attn_sink: [n_heads] f32` parameter).
Sink fold:
`g_max = max(g_max0, sink_logit[q_head])`
`g_sum = g_sum0 * exp(g_max0 - g_max) + exp(sink - g_max)`
`rescale= exp(run_max - g_max) / g_sum`
No value contribution. Tests cover both branches of the
`sink > local_max` selection.
2. `ffai_dsv4_mhc_pre` / `ffai_dsv4_mhc_post` — manifold-constrained
hyper-connection mix kernels (DSv4 `n_ch=4`). `_pre` collapses
the 4-channel residual state to a single-channel sublayer input
via `sum_c pre[c] * H[c, d]`. `_post` accumulates the sublayer
output back into the 4-channel state in-place via an
outer-product update `H[c, d] += post[c] * y[d]`. The
Sinkhorn-Knopp normalisation of `pre` and `post` is offline
(host-side); these kernels are pure weighted-sum / outer-add.
3. **Compressor pool rewrite** — `ffai_dsv4_compressor_pool` was
shipped earlier as a three-`for`-loop softmax (max → sum → acc).
On real Apple silicon it tripped a DSL codegen bug — sequential
for-loops sharing the `_w` index produced MSL referencing
`v21` / `v16` outside their declared scope. Reformulated as
single-pass online-softmax (Flash style) with running
`(max, sum, acc)` and per-step rescale. Same numerics, one
loop, no scope-collision hazard. CSA (pool=8) + HCA (pool=128)
tests now pass on M5 Max under the existing 1e-4 / 5e-3 / 5e-2
f32/f16/bf16 tolerances.
4. **`ffai_moe_router_sqrtsoftplus` codegen fix** — the kernel was
ICE-ing at MSL link time on real GPU dispatch (the
`every_registered_benchspec_codegens` test only checks IR
round-trip, not GPU compile). The DSL drops intermediate
bindings when chained method-form `.exp()` / `.ln()` / `.sqrt()`
calls feed into a multi-arg expression — produces MSL like
`auto v_routing_signal = sqrt(... + v_log_term)` where
`v_log_term` is never declared. Two fixes:
a. Break the softplus chain into single-operation `let`
bindings (`exp_neg`, `one_plus`, `log_term`, ...).
b. Switch the method-form `.exp()` and `.ln()` to the DSL's
free-function forms `exp(...)` and `log(...)` — the pattern
already used (and known-working) in
`gated_delta_prep.rs::softplus`.
Same numerics; rebench identical.
All four kernels validated end-to-end against CPU oracles in the
`kernel_tests_harness` workspace integration test on M5 Max. The
new kernels exercise f32 / f16 / bf16 within their declared
tolerances. Workspace clippy + fmt clean. The only remaining
harness failure is the pre-existing `mt_gated_delta_prep_chunk_no_gqa`
tolerance jitter (7.8e-3 vs 5e-3), unrelated to this PR.
For each KV cache position `t`, compute the aggregate score the top-k pass will sort on: ``` score[t] = sum_h w[h] * ReLU(q_idx[h] · k_idx[t, h]) ``` Three inputs (`q_idx`, `k_idx`, `w`) and one f32 output (`score` — full mantissa kept for stable bitonic ordering downstream). Dispatch is 1 thread per cache position; each thread walks the n_heads × d_idx = 8192 MADs at the DSv4 shape, with the inner d_idx loop streaming a hot 8 KB per-head row of `k_idx` from DRAM. Correctness checked at two shapes: - small: 8 heads × 32 dim × 64 positions (sanity) - DSv4 production: 64 heads × 128 dim × 256 positions f32/f16/bf16 covered. Head weights mix positive + negative values so the per-head sign branch in the accumulator is exercised. The fused 1-SG-per-position dispatch (lanes split d_idx, simd_sum the dot, share head-w broadcast across the SG) is the perf follow-up; playbook §"Testing — what to ALWAYS check" → "End-to-end GPU dispatch test, not just MSL-emit smoke test" applied: correctness first, perf in a later kernel. The top-k selector (`ffai_dsv4_indexer_topk`, bitonic top-512) is the next kernel — together they make CSA's sparse-gather attention runnable.
…nic) Downstream of `ffai_dsv4_indexer_score` — selects the K largest aggregate scores and returns their original cache-position indices to drive CSA's sparse-gather SDPA. Implementation: single-block parallel bitonic sort descending over `(score, index)` pairs in TG memory. Same Batcher compare-and-swap pattern as `mlx::sort::mt_sort` but with parallel `u32` index tag that survives every compare-and-swap, so the original position flows through to the output. - TPG = 256, 4 sort slots per thread → 1024 elements per block. - Scores beyond `n_kv` padded with `-INFINITY` → sink to the back, the first `k` slots are guaranteed to be real positions for any `k <= n_kv`. - 10 outer × variable-inner stages of compare-and-swap; barrier discipline matches `mt_sort`: skip barrier on strides ≤ 64 (within one simdgroup) and barrier on the wider strides. Three shapes covered in the kernel test: - 64 positions × top-8 (sanity) - 256 positions × top-32 (mid) - 1024 positions × top-512 (DSv4 production K = 512) Exact correctness (tol = 0.0) since the output is indices, not scores — distinct scores in the test fixture mean no ties at the K-th boundary that could swap ordering between CPU and GPU stable-sort. `n_kv > 1024` is the multi-block follow-up: chunk-block sort + a `mt_merge`-style co-rank merge that keeps the index tag through. For the DSv4-Flash context sizes where the indexer is most load-bearing (8K-256K cache), this lands next.
Single-token attention over a caller-supplied gather list of cache positions. Clone of `sdpa_decode_d512`'s online-softmax flash body with the dense `for _t in 0..n_kv` loop swapped for an index-gather over `selected_indices[n_selected]`. The caller builds `selected_indices` by unioning Lightning Indexer top-K (default K=512) with the trailing sliding window (default 128) and de-duplicating on the host. Kernel is oblivious to that union logic — it just gathers (K, V) at each listed index and runs flash attention. Dispatch is identical to d512 (TPG=512, 16 dims/lane, 4-phase output reduction), so the per-PSO geometry stays the same as the existing SDPA decode kernels — no new TG-mem footprint, no new barrier patterns to verify. Per the playbook §"When NOT to dedup" — mask-on / mask-off and dense-vs-sparse SDPA variants are intentionally kept as parallel kernels rather than collapsed via a runtime selector. PSO-compile constant-folds the difference anyway, and the parallel form keeps the emitted MSL clean (no `if gather` in the inner loop). Correctness: 8 q-heads × 4 kv-heads (GQA=2) × head_dim=512 × n_selected=32 over kv_stride=128. Selected indices stride by 4 to exercise the full cache footprint. CPU oracle does dense softmax-over-the-selected-set, online-softmax converges to the same answer. f32 / f16 / bf16 within 1e-3 / 3e-3 / 1.5e-2. With this kernel landed, the CSA + HCA attention path is complete at the SDPA layer. Remaining: MLA dense decode with absorbed W_UK for the full-attention layers 0, 1, and 42.
The 9 DSv4 kernels (compressor_pool, csa_sdpa_decode, fp8_block_dequant, indexer_score, indexer_topk, mhc, mxfp4_dequant, moe_router_sqrtsoftplus, sdpa_decode_d512_sink) plus the 3 GGUF dequant kernels (iq2_xxs, q2_k, q8_0) all shipped with a personal `Tom Turney (@TheTom)` header that didn't match the surrounding ffai/ convention. Flipped to the established team form: Copyright 2026 0xClandestine, Ekryski, TheTom, Ambisphaeric Pure header swap — no code or test changes.
Adds `ffai_dsv4_mhc_sinkhorn_split` — converts the per-token 24-mix output of `hc_fn @ flat` into the three control tensors (`pre [4]`, `post [4]`, `comb [4×4]`) the mHC collapse / expand kernels consume per sub-block boundary. Math: - pre = sigmoid(mix[0:4] * scale[0] + base[0:4]) + eps - post = 2 * sigmoid(mix[4:8] * scale[1] + base[4:8]) - comb_raw = mix[8:24] * scale[2] + base[8:24] - comb = Sinkhorn(softmax_over_src(comb_raw) + eps, n_iters) Sinkhorn-Knopp: alternate col-normalize + row-normalize for `sinkhorn_iters` iterations, with eps regularizer in each denominator. Hardcoded n_hc=4 so the 4×4 comb lives in 16 per-thread registers across the iteration loop — no TG memory, no array indexing. Output layout `comb[t * 16 + dst*4 + src]`. Single thread per token. Trivial for decode (one thread total per call) and still memory-bandwidth-bound for prefill (per-token work is ~100 FMAs, dominated by mix load + 24-out store). Tests cover decode (n_tokens=1, iters=1) and small-batch multi-iter (n_tokens=8, iters=3) across f32 / f16 / bf16. All GPU-verified on M5 Max. Also scrubbed external-source references from this kernel's docstring and the prior dsv4_fp8_block_dequant / dsv4_mxfp4_dequant headers per project policy: no crediting external sources in code or comments (deleted `DSV4_FORWARD_ARC.md` for the same reason).
The earlier mhc_pre / mhc_post took static [n_ch] weight tensors —
wrong for DSv4. The real model computes per-token pre/post/comb
dynamically from the 24-mix output of hc_*_fn @ flatten(H), then
collapse/expand consume those dynamic tensors per token.
Renamed:
- mhc_pre → mhc_collapse: x[d,t] = sum_c pre[t,c] * H[d,c,t]
- mhc_post → mhc_expand: H_new[d,dst,t] = block_out[d,t]*post[t,dst]
+ sum_src comb[t,dst,src]*residual[d,src,t]
Both kernels loop over n_tokens internally so a single dispatch
handles a prefill chunk. comb layout matches the sinkhorn_split
kernel's output: comb[t*n_hc*n_hc + dst*n_hc + src].
Tests cover decode (n_tokens=1) + small batch (n_tokens=4) at
hidden=4096/512 in f32/f16/bf16. All GPU-verified.
DSv4 head_dim=512 splits as nope=448 + rope=64. Q and K are computed full-rank; this kernel rotates only the last 64 dims of each head using the split-pair convention (dim_i, dim_i+half_rot) for dim_i ∈ [n_nope, n_nope + half_rot). The kernel writes only the rope pairs — caller initialises the output buffer with the nope dims already in place (or uses in-place mode where out == qk). Verified against a CPU reference at three shapes: - DSv4 production forward (head_dim=512, n_rot=64, theta=10K) - DSv4 production inverse (for the attn-output unrotation step needed because K==V in MQA mode) - HCA compressed-stream theta (160K base) f32/f16/bf16 within 1e-4 / 5e-3 / 5e-2. The inverse rotation is implemented by flipping theta's sign; cos(-θ)=cos(θ) and sin(-θ)=-sin(θ) gives the inverse 2D rotation without a second code path.
Trivial primitive for in-loop accumulators that didn't exist in the binary-op surface (mt_binary_* is shape-spec-tied to MLX's Binary class registration). Bare #[kernel] form so it emits cleanly into the generated MetalTileKernels.swift for FFAI callers. Generic over T (f32/f16/bf16); accumulates in f32 then implicit-narrowing-stores back to T.
Contributor
Author
|
Wrong target repo — see 0xClandestine/metaltile#243 for the canonical metaltile PR. |
Contributor
Author
Update: backs FFAI DeepSeek-V4 prefill parity (8k) 🎯Pushed the kernels/fixes that took FFAI DSv4-Flash prefill to 382 tok/s warm @8k = 112% of the ds4 target on M5 Max 128GB (see FFAI #17).
Note: the parity-winning MoE kernels ended up as hand-written raw MSL in FFAI using the M5 |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Companion to thewafflehaus/FFAI#17 — adds the GPU kernels FFAI's GGUF loader + DSv4 family path dispatch into.
GGUF block-dequant kernels
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 × value_lookup_decode`. Output dtype follows `T` via implicit Store coercion.
DeepSeek-V4 architecture kernels
Plus utility: `ffai_vector_add` (generic elementwise out[i] = a[i] + b[i]).
Validation
All 14 kernels have CPU oracle + GPU correctness tests in `kernel_tests_harness`. Tolerances within the standard f32 / f16 / bf16 bands (1e-4 / 5e-3 / 5e-2 elementwise; 1e-3 / 3e-3 / 1.5e-2 SDPA). 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`):
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`.
Status
Branch is on TheTom/metaltile (personal fork) because the cross-repo CI pin policy keeps org branches free of personal-fork URLs. Merge target is `houseofwaffles/metaltile dev`.