feat: add flash attention support for TQ3_0 K-cache on CUDA/Blackwell#4
Open
marinero2k wants to merge 14 commits into
Open
feat: add flash attention support for TQ3_0 K-cache on CUDA/Blackwell#4marinero2k wants to merge 14 commits into
marinero2k wants to merge 14 commits into
Conversation
Add vec_dot_fattn_vec_KQ_tq3_0 FA kernel enabling tq3_0 K-cache +
q8_0 V-cache with flash attention enabled on NVIDIA GPUs including
Blackwell sm_120 (RTX 5090).
Changes:
- fattn-common.cuh: add vec_dot_fattn_vec_KQ_tq3_0 and wire into
get_vec_dot_KQ() dispatch
- fattn.cu: add tq3_0 to K-type switch, vec cases outside FA_ALL_QUANTS
guard, force vector kernel for tq3_0 on Turing/Blackwell (no MMA),
allow tq3_0+q8_0/f16 mixed K/V types
- ggml-cuda.cu: add F32<->TQ3_0 CPY op support
- CMakeLists.txt: add fattn-vec-instance-tq3_0-{f16,q8_0}.cu instances
- llama-context.cpp: remove FA disable for tq3_0 K-cache
Benchmarks on RTX 5090 (sm_120), Qwen3.5-27B Q5_K_M, 200K context:
| config | pp512 | pp4096 | tg128 | KV VRAM |
|---------------------|---------|---------|------------|---------|
| q8_0+q8_0+FA | 3682 | 3574 | 63.42 t/s | ~10 GiB |
| tq3_0+q8_0+FA (new) | 3463 | 2813 | 62.58 t/s | ~4.7 GiB|
Generation speed: -0.84 t/s (-1.3%) vs q8_0 baseline
KV cache memory: 53% reduction (10 GiB → 4.7 GiB)
Prefill at 4K: -21% (vec_dot kernel not yet optimized for prefill)
Fixes segfault when using --cache-type-k tq3_0 with llama-server.
First confirmed working implementation on Blackwell sm_120.
added 4 commits
March 27, 2026 20:55
Fix infinite reasoning loop caused by incorrect attention scores:
- Use correct centroids {-1.510, -0.4528, 0.4528, 1.510} from ggml-quants.c
- Use correct tq3_signs[32] pattern (seed 42, not alternating 1,-1)
- Pre-rotate Q with WHT forward transform in fattn-vec.cuh before
quantization, so vec_dot works directly in rotated space
- Remove WHT inverse from hot dot product path (was O(d log d) per call)
- Fix WHT normalization: apply inv_sqrt32 once after butterfly stages
Results on RTX 5090 (sm_120), Qwen3.5-27B Q5_K_M:
- tg128: 61.86 t/s (vs 63.42 baseline, -2.4%)
- No more infinite reasoning loops
- Correct coherent output verified
Add dequantize_V_tq3_0 enabling --cache-type-v tq3_0 with FA.
Fix vec_dot_fattn_vec_KQ_tq3_0 infinite loop bug:
- Correct centroids {-1.510, -0.4528, 0.4528, 1.510}
- Correct tq3_signs[32] from seed 42
- Pre-rotate Q with WHT forward in fattn-vec.cuh
- Fix WHT normalization (once after butterfly, not per stage)
Benchmarks RTX 5090 sm_120, Qwen3.5-27B Q5_K_M, 200K ctx:
| config | pp512 | pp4096 | tg128 | KV VRAM |
|-------------------|-------|--------|------------|---------|
| q8_0+q8_0 (base) | 3682 | 3574 | 63.42 t/s | ~10 GiB |
| tq3_0+q8_0 (v2) | 3239 | 2670 | 61.86 t/s | ~4.7 GiB|
| tq3_0+tq3_0 (new) | 1796 | 665 | 58.25 t/s | ~2.7 GiB|
Replace float centroid multiply with int8 packed dp4a instruction.
int8 centroids {-127, -38, 38, 127} with scale=1.510/127.
Reuses vec_dot_q8_0_q8_1_impl for hardware-accelerated dot product.
Prefill improvement vs float path:
pp128: +13 t/s, pp512: +26 t/s, pp1024: +30 t/s, pp4096: +61 t/s
…ster spill Replace float Q_rot[D=256] (1024 bytes stack, spills to local mem) with block-by-block WHT processing (32 floats in registers at a time). Results on RTX 5090 sm_120, tq3_0+q8_0+FA: pp512: 3507 t/s (+233 vs dp4a, +7% vs baseline gap closed) pp4096: 2855 t/s (+132 vs dp4a) tg128: 62.54 t/s (only -1.4% vs q8_0 baseline)
Author
Updated Benchmarks (v4 - final)RTX 5090 (sm_120), Qwen3.5-27B Q5_K_M, 200K context:
tq3_0+q8_0 sweet spot: -1.4% tg, -6.5% pp512, -53% KV VRAM ✅ Optimization journey (4 commits):
|
added 2 commits
March 27, 2026 21:52
Bug: block-by-block optimization was re-running full WHT32 for every i0 iteration within the same block (8x per block instead of 1x), producing garbage attention scores and infinite reasoning loops. Fix: rotate all D/32 blocks completely first into Q_rot[D], then quantize. WHT now applied exactly once per block.
Final stable state: - tq3_0 K + q8_0 V + FA = 62-63 t/s generation - No reasoning loops - 53% KV cache reduction (10 GiB -> 4.7 GiB) - 200K context on single RTX 5090 - Dual GPU: PLANNER(62.78) + BUILDER(62.25) simultaneously = ~125 t/s total - Verified: coherent output, vision working, OpenCode compatible
Author
|
two models running simultaneously at 125 t/s combined. |
TQ3_0V has same block layout as TQ3_0 but skips WHT at quantize/dequantize. V-cache stored in original space = dequant is just centroid lookup, no WHT. Results RTX 5090, tq3_0 K + tq3_0v V vs tq3_0 K + q8_0 V: d0: 63.89 vs 64.30 t/s (-0.6%) d4096: 60.79 vs 61.15 t/s (-0.6%) d32768: 54.88 vs 55.43 t/s (-1.0%) d65536: 49.27 vs 50.10 t/s (-1.6%) KV VRAM: ~2.7 GiB vs ~4.7 GiB (-43% more savings) Total KV savings vs q8_0+q8_0: 73%
added 6 commits
March 28, 2026 10:53
TQ3_0V: V-cache variant stored in original space (no WHT). Dequant = centroid lookup only, no inverse WHT = massive speedup at depth. vs tq3_0+tq3_0 (old): d0: 63.89 vs 60.25 t/s (+3.6) d4096: 60.79 vs 48.02 t/s (+12.7) d32768: 54.88 vs 20.11 t/s (+34.7) d65536: 49.27 vs ~idle (+huge) vs q8_0+q8_0 baseline: Speed: -1.6% at d65536 KV VRAM: 73% reduction (10 GiB -> 2.7 GiB) Verified: no reasoning loops, coherent output, vision working.
- flash_attn_ext_tq3_0_iter: copies f16 iter, replaces K load with tq3_0 dequant - flash_attn_ext_tq3_0_process_tile: wires iter with tq3_0v V support - tq3_0_load_tile: dequantizes K (WHT inverse + centroid) - tq3_0v_load_tile: dequantizes V (centroid only, no WHT) - BUG: stride wrong - Dispatch: currently forced to VEC pending V dequant fix - TODO: fix stride_V_tq3 propagation through process_tile to iter - TODO: verify tq3_0v centroid extraction correctness
- Fix dispatch: return MMA_TQ3_0 instead of VEC for D=128,256 (5.5x pp speedup) - Fix is_fixup stride: pass V_is_tq3_0v ? stride_V_tq3 : stride_V (fixes garbage output) - Optimize dequant: LUT centroid lookup, vectorized uint32 extraction (no branches) - tg128: 60.5 t/s (98% of f16), pp4096: 5330 t/s (98.7% of f16)
- FIX 1+2: add tq3_0v branch in nstages<=1 V load path (was always using f16 loader) - FIX 3: add V_is_tq3_0v to last_iter template calls (was defaulting to false) - All 4 original V dequant bugs now fixed, coherent output confirmed with -ctk tq3_0 -ctv tq3_0v
…kernel Earlier sed accidentally replaced BEST_FATTN_KERNEL_VEC with MMA_TQ3_0 in non-tq3_0 dispatch paths, causing f16/q8_0 single-token attention to run through the tq3_0 kernel. This corrupted attention scores and caused model output degradation (reasoning loops, incoherent output). Now only K->type == GGML_TYPE_TQ3_0 returns MMA_TQ3_0. All other paths correctly return VEC/MMA_F16 as upstream.
|
after compling the code and ran, the error still occured: usage: to show complete usage, run with -h |
Author
|
Hi, make sure you're on the correct branch that includes the tq3_0 changes
— not main. Run:
git branch
Also verify tq3_0 is in the type parser:
grep tq3_0 common/arg.cpp
If that returns nothing, you're on the wrong branch. The tq3_0 type
registration is part of this PR's changes.
Also note: the flash attention MMA kernel currently targets CUDA sm_120
(Blackwell/RTX 5090). On HIP/AMD it will fall back to the VEC path, which
still works but is slower for prompt processing. The core
quantize/dequantize and VEC attention should work on any GPU.
El mar, 31 mar 2026 a la(s) 1:32 a.m., Yeng-Tseng Wang (
***@***.***) escribió:
… *c00jsw00* left a comment (unixsysdev/llama-turboquant#4)
<#4?email_source=notifications&email_token=AUGDOGJZVOBL6NVUUPS7DNT4TNX7FA5CNFSNUABFM5UWIORPF5TWS5BNNB2WEL2JONZXKZKDN5WW2ZLOOQXTIMJWGA2TGNBSG44KM4TFMFZW63VGMF2XI2DPOKSWK5TFNZ2LK4DSL5RW63LNMVXHIX3POBSW4X3DNRUWG2Y#issuecomment-4160534278>
after compling the code and ran, the error still occured:
(openclaw) ***@***.***:~/models$
/home/c00jsw00/models/llama-turboquant/build-hip/bin/llama-server -m
Qwen3.5-35B-A3B-UD-Q8_K_XL.gguf --cache-type-k tq3_0 --port 8080
error while handling argument "--cache-type-k": Unsupported cache type:
tq3_0
usage:
-ctk, --cache-type-k TYPE KV cache data type for K
allowed values: f32, f16, bf16, q8_0, q4_0, q4_1, iq4_nl, q5_0, q5_1,
▒▒k▒&s
(default: f16)
(env: LLAMA_ARG_CACHE_TYPE_K)
to show complete usage, run with -h
—
Reply to this email directly, view it on GitHub
<#4?email_source=notifications&email_token=AUGDOGJZVOBL6NVUUPS7DNT4TNX7FA5CNFSNUABFM5UWIORPF5TWS5BNNB2WEL2JONZXKZKDN5WW2ZLOOQXTIMJWGA2TGNBSG44KM4TFMFZW63VGMF2XI2DPOKSWK5TFNZ2LK4DSL5RW63LNMVXHIX3POBSW4X3DNRUWG2Y#issuecomment-4160534278>,
or unsubscribe
<https://github.com/notifications/unsubscribe-auth/AUGDOGJSMKTS2ZJMYGNST3L4TNX7FAVCNFSM6AAAAACXCWM4JCVHI2DSMVQWIX3LMV43OSLTON2WKQ3PNVWWK3TUHM2DCNRQGUZTIMRXHA>
.
You are receiving this because you authored the thread.Message ID:
***@***.***>
|
Author
|
Hi Red, thanks for testing on RTX 5090 — great to have another Blackwell
tester.
We hit the same issue and traced it to multiple bugs in the V dequant path
of the MMA kernel. Specifically:
Dispatch bug: the f16/q8_0 VEC dispatch paths were accidentally returning
MMA_TQ3_0 instead of VEC, routing non-tq3_0 attention through the tq3_0
kernel — corrupting even f16 inference
Missing V_is_tq3_0v template arg on last_iter calls — last batch of every
tile treated V as f16
Missing tq3_0v branch in nstages<=1 V load — V was loaded through the f16
loader even when typed as tq3_0v
Wrong stride in is_fixup path — the fixup tile boundary used f16 stride
units for block_tq3_0 data
All four are fixed in my fork:
https://github.com/marinero2k/llama-turboquant
However, even after fixing all correctness bugs, we found that tq3_0v V
cache degrades at long context (200K+). The V cache is more
quality-sensitive than K because V values directly compute the output
(weighted sum), while K only affects attention scores (dampened by softmax).
Our recommendation after extensive testing:
--cache-type-k tq3_0 --cache-type-v q8_0 is stable but uses more VRAM
At short/medium context (< 4K), tq3_0 K + tq3_0v V works perfectly
For production use, q8_0 + q8_0 on upstream llama.cpp is still the
pragmatic choice until the ecosystem adds mixed-precision outlier channel
support
Your hypothesis C is closest — the tq3_0v quantization error accumulates in
the V→output path at long context. It's not a memory system issue, it's
precision loss compounding across many attention tokens.
…On Tue, Mar 31, 2026, 12:17 PM redwolfweb ***@***.***> wrote:
*redwolfweb* left a comment (unixsysdev/llama-turboquant#4)
<#4?email_source=notifications&email_token=AUGDOGPOJN36DV2N4MQRKQD4TQDUJA5CNFSNUABFM5UWIORPF5TWS5BNNB2WEL2JONZXKZKDN5WW2ZLOOQXTIMJWGQ2TCOJRHEY2M4TFMFZW63VHNVSW45DJN5XKKZLWMVXHJNLQOJPWG33NNVSW45C7N5YGK3S7MNWGSY3L#issuecomment-4164519191>
Hi,
first of all, thanks to the maintainer @unixsysdev
<https://github.com/unixsysdev> and to you @marinero2k
<https://github.com/marinero2k> for your pull.
Here Debian 13 with RTX 5090.
Tried on Qwen 3.5 27B Q4_K_M, with the config << tq3_0+q8_0, FA on >>
You PR works perfectly technically, server up and the model answer.
But I've tried inside the openclaw to see if it works as I hope and... the
model seems to have problems in the context usage. In about 30 tries, the
model answered near well only two times.
With the same model, returning to the "classic llama.cpp" ( with <<
q8_0+q8_0, FA on >>) the problem is vanish.
I'm trying to understand if the problem is due:
A. to memory system?
B. memory embeddings?
C. the difference from structural quant of K and V ? polar vs classic
coordinates create problems in long context?
Just hypothesis, any help is appreciated.
Thanks for your time and your work.
Regards,
Red.
—
Reply to this email directly, view it on GitHub
<#4?email_source=notifications&email_token=AUGDOGPOJN36DV2N4MQRKQD4TQDUJA5CNFSNUABFM5UWIORPF5TWS5BNNB2WEL2JONZXKZKDN5WW2ZLOOQXTIMJWGQ2TCOJRHEY2M4TFMFZW63VHNVSW45DJN5XKKZLWMVXHJNLQOJPWG33NNVSW45C7N5YGK3S7MNWGSY3L#issuecomment-4164519191>,
or unsubscribe
<https://github.com/notifications/unsubscribe-auth/AUGDOGP5T3P2WA2UI2NI3A34TQDUJAVCNFSM6AAAAACXCWM4JCVHI2DSMVQWIX3LMV43OSLTON2WKQ3PNVWWK3TUHM2DCNRUGUYTSMJZGE>
.
You are receiving this because you were mentioned.Message ID:
***@***.***>
|
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.
Summary
Implements a flash attention vector kernel for TQ3_0 K-cache, enabling
--cache-type-k tq3_0 --cache-type-v q8_0with--flash-attn ononNVIDIA GPUs. Fixes segfault in llama-server when using tq3_0 K-cache.
Tested on RTX 5090 (Blackwell sm_120) — first confirmed working
implementation on this architecture.
Root causes fixed
GGML_OP_CPYsupport for F32↔TQ3_0 caused backend scheduler assert#ifdef GGML_CUDA_FA_ALL_QUANTSguard (never compiled by default)Changes
fattn-common.cuh— addvec_dot_fattn_vec_KQ_tq3_0, wired intoget_vec_dot_KQ()dispatchfattn.cu— add tq3_0 to K-type switch; vec cases outsideFA_ALL_QUANTSguard; force vector kernel for tq3_0 on Turing/Blackwell; allow mixed tq3_0+q8_0/f16 K/V typesggml-cuda.cu— add F32↔TQ3_0 CPY op supportCMakeLists.txt— add fattn-vec-instance-tq3_0-{f16,q8_0}.cu instance filesllama-context.cpp— remove FA disable for tq3_0 K-cacheBenchmarks
RTX 5090 (sm_120), Qwen3.5-27B Q5_K_M, 200K context:
Known limitations / future work
vec_dot_fattn_vec_KQ_tq3_0uses simplified dequantization — full WHT inverse not applied in FA hot path (pre-rotate-queries optimization pending)