feat: fused turbo MMA flash-attention decode (turbo4/3/2, GQA tensor-core path) - #4
Merged
Conversation
…ritbuun fork MMA approach, adapted to graph-rotated Q
Port buun's fused turbo4 tensor-core (MMA) flash-attention DECODE kernel into
the turbo fork as an OPT-IN path (GGML_TURBO_MMA_FUSED=1, default OFF), routing
turbo4-K==turbo4-V, D in {128,256}, decode (Q->ne[1]<=4) onto the GQA-packed MMA
path so KV is read once per head-group instead of once per query head.
How it works:
- fattn-mma-f16.cuh: add flash_attn_ext_turbo4_load_tile (dequant turbo4 blocks
-> SRAM in the f16 row-major half2 layout) seeded with OUR Lloyd-Max centroids
(turbo-quant.cuh:297, NOT buun's -0.2415xx table). Thread type_K/type_V through
iter/process_tile/global flash_attn_ext_f16 (defaults F16 -> existing f16/q8 MMA
byte-identical). For turbo: force nstages=0 (the cp.async multi-stage path would
read raw turbo bytes as half2), pass RAW byte pitch nb11/nb21 (not /sizeof(half2)),
and dequantize K/V in the load tiles. sizeof(block_turbo4_0)-driven pointer math
(66B block, no rnorm) — never assumes a fixed offset.
- fattn-mma-turbo.cuh + 14 instance files (7 reachable ncols pairs x D in {128,256}):
host launcher reusing the f16 device kernel with type_K/type_V=TURBO4_0 and
need_f16_K/V=false so launch_fattn does NOT pre-convert KV to f16.
- fattn.cu: switch_ncols2 -> per-ncols2 dispatchers (only the 7 compiled instances),
env latch, and the gate (Q untouched — our fork already rotates Q at the graph
level, so NO inline FWHT / src swap, which would double-rotate Q).
Results (Qwen3.6-35B-A3B turbo4, GQA8, D=256, 5090):
depth MMA VEC+PDL buun
8192 213.8 188.5 200 (+6.9% vs buun)
16384 207.6 171.7 193 (+7.5% vs buun)
32768 191.3 141.1 187 (+2.3% vs buun)
MMA holds flat at depth where VEC collapses; the gap vs VEC grows with KV length.
Quality: Mean KLD vs f16 base = 0.008396 (== VEC baseline). The MMA path is NOT
strictly token-identical to VEC — MMA and VEC accumulate the P*V (VKQ) reduction in
f16 with different reduction trees, so a near-tie greedy token can flip (~1/25 on a
hard tie; same irreducible f16-order effect that exists between base f16-MMA and
f16-VEC, and within the VEC-vs-base "same-top-p" 96.3% noise floor). Because strict
token-identity does not hold, the gate DEFAULTS OFF and ships as opt-in; VEC remains
the untouched default/kill-switch.
(cherry picked from commit 5450669510a5cd0a356a863923b057f2552c6fbb)
(cherry picked from commit b3e51cf)
…eutral; GGML_TURBO_MMA_FUSED=0 kill-switch) (cherry picked from commit 726234b538194c644be4066e1413b91f9edf12e7) (cherry picked from commit 38dc665)
…extends the MMA path to turbo3/2) The GQA-packed tensor-core decode path was turbo4-only; turbo3/turbo2 fell back to VEC and collapsed at depth (turbo3 tg32 224->124 from d2048->d32768), losing to spiritbuun there. Added flash_attn_ext_turbo3_load_tile (3-bit split index: 2 low bits from qs + 1 high bit from signs) and flash_attn_ext_turbo2_load_tile (plain 2-bit), type-dispatched both the K and V MMA load sites, and routed the gate by K->type. is_turbo_kv already matches any non-f16 KV, so the kernel needed only the per-type tile loaders + instances (14 instance files now emit all three types). Correctness: MMA decode is bit-exact to VEC (Mean KLD @2048 chunks=8: turbo3 0.020278==0.020278, turbo2 0.041490==0.041490, turbo4 0.009296 unchanged). Decode tg32 (qwen3.6-35B-A3B, RTX 5090), MMA now flat at depth: turbo3 2k/8k/16k/32k = 224/225/217/203 (buun 193/193/183/178; was VEC 224/187/160/124) turbo2 2k/8k/16k/32k = 228/224/217/203 (buun 213/210/204/200; was VEC 229/218/203/183) turbo3/turbo2 now hold flat at depth (neither fork had them on the MMA path before). GGML_TURBO_MMA_FUSED=0 restores VEC. (cherry picked from commit 57ad5edd53a7fced0b99fecb99326132cbdb65d0) (cherry picked from commit 4e223ee)
Tester @everson isolated a turbo2-only decode regression on RTX 5060 Ti (Gemma-12B, head_dim 256): turbo2/turbo2 -2.46% at 8K to -1.17% at 200K, fully restored by GGML_TURBO_MMA_FUSED=0 at every depth, while turbo3 was neutral and turbo4 exact parity. At 2-bit KV the fused path's GQA-pack saving is tiny while its dequant/no-pipeline overhead is unchanged, so on head_dim 256 it is neutral on high-bandwidth GPUs (Ornith-35B and a dense Gemma-4-12B on RTX 5090, both within noise) and only costs on bandwidth-limited cards. turbo2 fused remains a large depth win on dense head_dim 128 (Llama-3.1-8B on RTX 5090: +7% at 8K, +24% at 32K, +69% at 131K), so this gates turbo2 fused to head_dim 128 only; head_dim 256 turbo2 falls through to VEC, which is the exact baseline path GGML_TURBO_MMA_FUSED=0 restored. turbo3 and turbo4 fused are unchanged at both head dims. Also corrects the stale "default OFF" gate comments to match the default-ON code. (cherry picked from commit 539ce5d)
thecodacus
pushed a commit
that referenced
this pull request
Sep 7, 2026
* Get started with Onyx * Add architecture * Skip keys handled in super() * Loading tensors * Shorten * Graph * Apply suggestion from @pcuenca * Remove norm now embedding in transformers weights * Add eot * Explicit output_multiplier * Handle post_norm_eps * No super call; unhardcode eot. The pattern `self._set_vocab_gpt2()` seems preferred throughout the codebase, and it allows `set_vocab()` to be called from a different part of the Python class hierarchy: the drafter model converter that we may need eventually. * Register for drafting * DFlash: inherit rope type from the linked target. Another option would be to store it in the gguf file itself. * mmproj conversion Note: some fields to be renamed after the implementation works. We are keeping compatibility with the reference Meta gguf for testing purposes. * "clip" header declarations * Load mmproj * Pre-processing * Graph * Go back to using delimiters. Otherwise our generations are worse. Transformers does not use them. We need to trace inputs to verify whether they are equivalent. * downsample_factor -> merge_size * Add vision graph lol, forgot from a previous commit * Additional renames, align with llama.cpp / transformers * Prefer _size instead of independent _h and _w * Fix token layout Co-authored-by: Young Han <younghan@fb.com> * onyx: bring the chat parser onto the onyx branch common/chat.cpp on this branch has no Onyx handling, so a converted model serves malformed chat: the assistant preamble leaks into content ("to=self<|message|>...") and tool calls fail with HTTP 500 "The model produced output that does not match the expected peg-native format" common_chat_params_init_onyx exists on onyx-fair-patch, added there by 8bb73dd3d. It was never on this branch, so this is not a regression -- the two lines developed independently. The code here is taken verbatim from that commit. It is the clean side of `git merge origin/onyx-fair-patch`: chat.cpp is one of the files that merges without conflict. The full merge is not viable -- it produces 13 conflicts, including add/add on conversion/onyx.py and src/models/onyx.cpp where the q_norm-folding and metadata-scale approaches contradict each other, and #4/#7 are stacked on this branch's side of that. Verified on this branch: builds with 0 errors, converts an Onyx checkpoint, and serving it gives "4" for "What is 2+2?" plus a correct get_weather {"city":"Paris"} tool call, where the unported branch gives the two failures above. No converter or runtime changes are included, so this should not interact with the q_norm work. Co-authored-by: Beto de Paola <betodepaola@meta.com> * Less params, bilinear pos-emb interpolation as a graph op instead of CPU * Map to symbolic V_MMPROJ instead of strings * Make a couple params explicit * Patchify via build_inp() * No param for rope_theta * Small cleanup * Restore blank line * Unpermute, to adapt to the latest transformers checkpoint * Apply norm after token embeddings This follows the latest transformers approach. * Remove duplicated function * build_vit * onyx: use the model rope theta on sliding-window layers * DFlash: conversion from transformers drafter * Revert rope_type derivation from target NOTE: this breaks compatibility with Meta's distributed DFlash GGUFs, as the Q/K are stored in "NEOX" (rotated half) format, like in transformers. * Apply suggestion from @pcuenca * Set model type * Remove comment that will become obsolete * Hardcode post_norm_rms_eps instead of new param * Derive SWA+RoPE pattern from gguf array or scalar * Fix model type <-> number of layers * Reorder * Rename * Fix typo * DFlash: seed the draft KV cache from multimodal embedding batches `common_speculative_impl_draft_dflash::process()` returned early on any batch carrying embeddings, so an image prefill never had its target-layer features fused through the DFlash encoder and injected into the draft's KV cache. That left a hole spanning the image's positions, and the next injection at a post-image position failed to initialize its batch: ``` decoding image batch 1/1, n_tokens_batch = 256 decode: failed to initialize batch llama_decode: failed to decode, ret = -1 process: llama_decode(ctx_dft) failed rc=-1 (n_tokens=17, offset=0) srv decode: failed to process speculative batch ``` Every image request with `--spec-type draft-dflash` failed with HTTP 500. Text-only was unaffected, since those batches carry token ids and were let through. Restore the earlier condition, which admits a batch that is either tokens or embeddings and skips only the degenerate neither/both cases. The rest of `process()` is already layout-agnostic -- it gathers features via `llama_get_embeddings_layer_inp()` and indexes `batch_in.pos[]` / `batch_in.seq_id[]`, none of which assume token ids -- so this is the whole fix. Validated against `muse-glimmer-30B-bf16.gguf` + `mmproj-muse-glimmer-30B-bf16.gguf` + a DFlash draft head, on an image describe-the-shapes request: - before: HTTP 500, `failed to process speculative batch` - after: HTTP 200, draft acceptance 0.34012 (167 accepted / 491 generated), mean len 3.04 Output equivalence holds, which is the property that matters: at temperature 0 the drafted response is byte-identical to the same request served with no draft attached (1213/1213 chars), so the draft is drafting correctly through the image context rather than merely not crashing. * Conversion: prefer rewrite to mapping * Revert "Conversion: prefer rewrite to mapping" This reverts commit a92d0ac. * fix lint * sliding_window metadata is not optional * disable state save/load * Apply suggestion from @pcuenca --------- Co-authored-by: Young Han <younghan@fb.com> Co-authored-by: Beto de Paola <betodepaola@meta.com> Co-authored-by: Daniel Han <michaelhan2050@gmail.com> Co-authored-by: ruanrms <ruanslv@gmail.com> Co-authored-by: Xuan Son Nguyen <son@huggingface.co> Co-authored-by: Sigbjørn Skjæret <sigbjorn.skjaeret@huggingface.co>
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
2 participants
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.
Ports the fused turbo MMA flash-attention decode path from TheTom/llama-cpp-turboquant (4 commits, cherry-picked with
-x). Until now turbo KV decode always used the vector kernel; the turbo MMA files in-tree were declaration-only scaffolding.What it does
For decode batches (
ne1 ≤ 4) with matching turbo K/V types (turbo4/3/2, head dims 128/256; turbo2 gated to 128), attention now runs on the GQA-packed tensor-core MMA kernel, dequantizing turbo blocks directly into shared memory — no F16 conversion pass, no VEC serialization. Q rotation stays at the graph level (no inline FWHT; this path must not double-rotate).GGML_TURBO_MMA_FUSED=0is the kill-switch back to VEC (default ON upstream-of-us: correctness-validated, KLD-parity with VEC per the reference repo; not bit-identical — f16 reduction-order differences can flip a hard-tie greedy token, same as base MMA-vs-VEC f16).Verified — RTX 3060
ctk/ctv turbo4, fused path → correct output ("The capital of France is Paris.")-ngl 99, K=V=turbo4, tg64 @ d16384): VEC 26.90 ± 0.46 → MMA 30.11 ± 0.57 t/s (+12%)ctk q8_0 + ctv turbo2, mixed → stays on existing path): pp2048 521.5 / tg64 49.1 — unchanged