feat(gguf,dsv4): GGUF v3 reader + DeepSeek-V4-Flash GGUF loader & forward path#17
Conversation
90d700f to
4a8e7ac
Compare
61b0b37 to
70b1e21
Compare
Commit message hygiene checkAll commit messages and PR text are clean. ✅ |
e84732b to
070a87c
Compare
Status — rebased & squashed onto
|
| metric | measured |
|---|---|
| decode (warm steady-state) | ~30–32 tok/s |
| prefill @8k — reliable (prewarmed) | ~320 tok/s |
| prefill @8k — peak warm | ~382 tok/s |
| prefill @8k — cold-cache plateau | ~232 tok/s |
Prefill path: NAX (matmul2d) neural-accelerator MoE GEMMs, zero-repack view-bm64 reading raw resident IQ2/Q2_K blocks in place (no per-layer repack), and an expert-tensor page-cache prewarm (the cold→reliable gap is page-fault/eviction bound, not compute). GPU kernels live in 0xClandestine/metaltile#243.
8cbcb53 to
afa47d6
Compare
TheTom
changed the title
|
@ekryski — your call on one thing here: the I trimmed its external bug-tracker references (ollama #15594 / #14432, llama.cpp PR #16634) down to "a known class of MPP header-skew bug" as part of removing external-repo mentions from this PR. Those links were the provenance for the bit-determinism bug you originally documented — so this is a judgment call I'd rather you make: keep them removed, or restore the citations? Happy to put them back verbatim if you want the provenance preserved. |
4f489eb to
c70b0bf
Compare
#243) ## Summary Companion to [thewafflehaus/butter#17](thewafflehaus/butter#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).
There was a problem hiding this comment.
I'm thinking we should pair down this file. The stuff I've been testing in the common integration test pattern for a model is:
- Loader (does it load)
- Does it have the shapes and configs we expect
- Does it have the default parameters for it's model interface
- Does it produce coherent output (in this case coherent text).
See https://github.com/thewafflehaus/FFAI/blob/ek/fix-all-the-things/Tests/ModelIntegrationTests/Text/Gemma4TextIntegrationTests.swift for example. We should separate the GGUF testing stuff and keep it in a separate file ideally with a focus on being able to load a GGUF model (and a smaller one).
This test will also need to be skipped by default because the model is just so big.
There was a problem hiding this comment.
Agreed. Reworking the suite to the 4-point pattern (loads / shapes+configs / default params / coherent output) per your Gemma4 example, and splitting the GGUF-loading tests into their own file. It already skip-guards when the model isn't staged at $FFAI_DSV4_GGUF_PATH. (in progress, separate commit)
|
|
||
| /// System-wide free memory as a percentage (free+inactive vs hw.memsize), or | ||
| /// nil if the mach query fails. Used by the prefill freeze guard. | ||
| func ffaiSystemFreePercent() -> Double? { |
There was a problem hiding this comment.
We have a system memory stuff in https://github.com/thewafflehaus/FFAI/blob/ek/fix-all-the-things/Sources/FFAI/Stats/MemoryStats.swift. Alongside our PR #19 we should modify it if we don't have what you need. Ideally we have a single memory management module that all models can take advantage of to ensure they are using memory wisely (and fit in) before being loaded.
There was a problem hiding this comment.
Generally I've been putting all stuff into one model family file. So in this case DeepSeekV4Forward.swift and DeepSeekV4Prefill.swift would be inside DeepSeekV4Text.swift. Unless they are shared components that would be used for vision or something.
Yes, this does make the single model file bigger, and we may split them apart later, but it at least keeps things contained so we don't have too much new file sprawl. So maybe let's consolidate all 3 of these files into DeepSeekV4Text.swift so it follows the same convention as other models? I think it will make it easier to split things out across all models consistently in the future when we decide how we should split them. Happy to discuss.
There was a problem hiding this comment.
Done (e720481) — DeepSeekV4Forward.swift + DeepSeekV4Prefill.swift are now consolidated into DeepSeekV4Text.swift, matching the one-file-per-family convention.
There was a problem hiding this comment.
Need matching unit test files for this file and the other GGUF files that aren't covered.
There was a problem hiding this comment.
Agreed — adding unit-test coverage for GGUFDequant.swift and the other uncovered GGUF loader files. (in progress, separate commit)
| /// throw keeps control flow well-formed and documents the gap. Once | ||
| /// `DeepSeekV4Model: LanguageModel` lands, replace the trailing | ||
| /// throw with the standard `return Loaded(engine:...)` wrap. | ||
| public static func loadDeepSeekV4( |
There was a problem hiding this comment.
This has always been gross and needs a refactor I think. This model is sprawling. Not a blocker but flagging in case you have some ideas. Should probably coincide with a "ModelRegistry" refactor. We have 3+ different places that look and feel like a model registry right now. Should just have one.
There was a problem hiding this comment.
Ack — agree it's sprawling (3+ registry-ish spots). Out of scope for this PR; flagging for a dedicated ModelRegistry consolidation. Left untouched here.
There was a problem hiding this comment.
What is this needed for? What is it actually testing?
There was a problem hiding this comment.
It covers Device fundamentals (shared device/queue, makeBuffer length, makeCommandBuffer usability) plus the scratch bump-allocator — allocScratch 16-byte alignment + offset bumping + the diagnostic counters. See the scratch-mode answer below for what that allocator is; if we drop scratch mode these assertions go with it.
There was a problem hiding this comment.
All these kernels here probably shouldn't be checked in. They should come from metaltile. Likely an artifact of local AI iteration. The dev/ folder shouldn't even exist. Please remove.
There was a problem hiding this comment.
Done (e720481) — removed the whole dev/ folder. It was local-iteration scratch; the real kernels come from metaltile.
| // amortise the MLA absorb-W_UK setup over many positions, | ||
| // small enough to fit on a 96 GB Apple Silicon machine. | ||
| GenerationParameters( | ||
| maxTokens: 256, prefillStepSize: 4096, |
There was a problem hiding this comment.
Don't think we should have default maxTokens. Other default parameters are fine. Might want to double check what Deepseek recommends.
There was a problem hiding this comment.
Done (e720481) — dropped the DSv4-specific default maxTokens (falls to the GenerationParameters default). Also set temperature 0.6 / top-p 0.95 per DeepSeek's recommended sampling.
| let dev = Device.shared | ||
| switch out.dtype { | ||
| case .f32: | ||
| let startBuf = dev.makeBuffer(length: 4) |
There was a problem hiding this comment.
Are we sure we don't need this for other models?
There was a problem hiding this comment.
Today only the f32 arange has a production caller, so that's the only variant wired; f16/bf16 arange is a few lines against the matching kernel when a model needs it (noted in the comment). Not DSv4-specific — happy to add the half variants eagerly if you'd prefer.
| #expect(t.buffer.length >= 128) | ||
| } | ||
|
|
||
| @Test("empty routes through the scratch slab when scratchModeActive") |
There was a problem hiding this comment.
What is "scratch mode"? What is it used for?
There was a problem hiding this comment.
Scratch mode = a per-scope bump allocator on Device. Inside device.withScratch { ... }, Tensor.empty routes transient allocations through one reusable 16-byte-aligned slab instead of makeBuffer-per-tensor. It exists for the per-token decode forward: allocating many short-lived Tensors per step hammered Metal's internal driver pool and RSS grew ~3 GB/min until OOM on the M5. The slab resets at scope end; carry-over state still uses plain Tensor.empty. It's load-bearing for long-context forward, but if you'd rather not carry it I can pull it and lean on the driver pool — your call.
TheTom
changed the title
GGUF v3 mmap reader, DSv4 tensor-name map, IQ2_XXS/Q2_K block dequant tables, zero-copy model views, and tokenizer. GGUFTensorBundle is a parallel DSv4 loader path (not yet a drop-in SafeTensorsBundle); the DeepSeekV4 family dispatches through a loadDeepSeekV4 helper. Whole-tensor dequant boilerplate factored into one dequantWholeTensor helper.
Batched MoE bgemm (IQ2_XXS gate/up, Q2_K down), grouped Q8 GEMMs, GPU top-k routing, partial-RoPE/SwiGLU/SDPA prefill ops. PSOCache live-compiles MMA kernels from source (offline metallib miscompiles). Adds a Device scratch-slab allocator (Tensor.empty routes through it).
Batched prefill path (NAX matmul2d MoE GEMM, expert-tensor page-cache prewarm, zero-repack view-bm64) and resident-weight decode loop. Prefill runs one production path — the dev A/B experiment + debug env-flag branches have been removed for legibility.
Authoritative .metal/.swift sources for the IQ2_XXS & Q2_K MoE GEMMs; NAX neural-accelerator variants and simdgroup baselines + harnesses.
… default path
Reword the 'WIP'-tagged status/doc comments to factual phrasing
('not yet implemented' / 'deferred to follow-ups' / 'scaffold') — the
described state (stubbed safetensors forward, unimplemented CSA/HCA,
known-incorrect numerical shortcuts) is unchanged, only the labelling.
Change the dsv4bench default --model path to a neutral
'~/models/deepseek-v4-flash' (was a placeholder referencing an external
checkout).
… Swift 6.1 The IQ2_XXS / Q2_K resident-gather paths capture pool pointers (d / dmin / scales / qs) into a DispatchQueue.concurrentPerform @sendable closure. Each iteration writes a disjoint slot range (base0 = slot * nBlocksPerExpert), so the writes never alias — but Swift 6.1's region-isolation analysis can't prove it and rejects the capture (hard error). Swift 6.3 proves it safe, which is why local builds were clean while CI (6.1.2) failed to compile. Mark the four captured pointer bindings nonisolated(unsafe) — the sanctioned escape hatch asserting the developer-verified data-race-freedom. No runtime change; builds clean on both 6.1 and 6.3.
… dsv4 bench command + maxTokens default - Remove dev/moe_mma/ (local-iteration artifact; kernels live in metaltile). - Consolidate DeepSeekV4Forward.swift + DeepSeekV4Prefill.swift into DeepSeekV4Text.swift — one file per model family, matching convention. - Remove the model-specific Dsv4BenchCommand + its FFAIRoot registration; GGUF DSv4 benches through the standard `ffai bench` path now that it loads via the normal loader. - Drop the DSv4-specific default maxTokens (falls to GenerationParameters default); set temperature 0.6 / top-p 0.95 per DeepSeek's recommendation.
Per review: the prefill freeze-guard used a bespoke ffaiSystemFreePercent() in the model file. Move it to MemorySnapshot.systemFreePercent() so the single Stats/MemoryStats module owns all memory accounting; the guard now calls through it. No behavior change.
…rage Per review: - DeepSeekV4IntegrationTests pared to the common model pattern — loads / shapes+configs / default params / coherent-output (finite NaN-free logits). Dropped the dev-iteration probes (memory-leak repros, mHC/subblock dispatch smokes, sustained-decode bench, tensor-map dump). Skip-by-default (guards on $FFAI_DSV4_GGUF_PATH — the model is ~86 GB). - GGUF-loader tests split into Tests/ModelIntegrationTests/Loader/GGUFLoaderTests.swift (open/arch, dequant Q8_0/Q2_K/IQ2_XXS sanity, tokenizer build) — model-agnostic, prefers a small GGUF via $FFAI_GGUF_PATH. - New unit tests: Tests/FFAITests/Loader/GGUFDequantTests.swift — block-format constants + a deterministic Q8_0 round-trip (runs in CI). - Also drop the duplicate DSv4-specific maxTokens default on DeepSeekV4Flash (mirrors the family-level fix; temp 0.6 / top-p 0.95).
The model column recorded the absolute /Users/<name>/models/... load path. Strip to the bare model name, matching the HF-id style of the existing m1-max log.
b563676 to
977a1d0
Compare
2 participants
Summary
Two pieces, reviewed together — the GPU kernels they target are merged in metaltile (#242 Step-3 surface, #243 GGUF/DSv4 surface).
Sources/FFAI/Loader/GGUF/) — pure-Swift parser for the GGUF checkpoint format. Header + metadata KV + tensor-info table decoded eagerly; tensor data stays mmap'd. Q8_0 / Q2_K / IQ2_XXS dequant dispatch from aGGUFTensorBundle, so family-loader dispatch can consume GGUF or safetensors.Sources/FFAI/Models/Text/DeepSeekV4Text.swift) — mHC + MQA + attn-sinks + grouped O-LoRA + MoE + LM head, loaded lazily per-layer from the real ~86 GB DSv4-Flash GGUF.What works end-to-end on the real GGUF
DeepSeekV4Flash.loadModelFromGGUF(...)opens the checkpoint and exposes a lazy per-layer handle (eagerly loads only the 6 non-block tensors; per-layer weights load on firstmodel.layer(N)).forwardAllLayers(inputTokenId:state:)runs all 43 layers — mHC sinkhorn-split → Q low-rank chain + per-head Q-norm → partial RoPE → MQA SDPA with learnableattn_sinks→ inverse RoPE → 8-way grouped O-LoRA → mHC expand → FFN (sqrtsoftplus router → top-6 experts + shared expert) → output mHC head + LM head. Finite, NaN-free logits against real weights.intermediateScratchbuffers and zero-copy mmap reads, validated by a cross-layer load/release probe (held flat at ~1 GB vs the naive +12 GB/layer growth).Review changes addressed
DeepSeekV4Forward.swift+DeepSeekV4Prefill.swiftintoDeepSeekV4Text.swift(one file per family).Dsv4BenchCommand— GGUF DSv4 benches via the standardffai bench.MemorySnapshot.systemFreePercent()(single memory module).maxTokens(uses the framework default; temp 0.6 / top-p 0.95 per DeepSeek).dev/local-iteration folder.GGUFDequantunit tests. Integration tests skip by default (the checkpoint is ~86 GB).Known follow-ups (separate PRs)
forwardFullAttnSubblockcurrently runs on all layers; the compressed/sparse attention path for CSA/HCA layers isn't implemented, so generated-token quality is off until it lands.loadModelthrowsnotYetImplemented; only the GGUF path is wired today.Test surface
FFAITests+MetalTileSwiftTestsgreen — incl. newGGUFDequantTests(block-format constants + a Q8_0 round-trip), plusGGUFReaderTests,DeviceTests,TensorTests.DeepSeekV4IntegrationTests— loads + layer geometry, default generation parameters, forward produces finite NaN-free logits.GGUFLoaderTests— open/arch, Q8_0/Q2_K/IQ2_XXS dequant sanity, tokenizer build (prefers a small GGUF via$FFAI_GGUF_PATH).