diff --git a/docs/superpowers/specs/2026-08-10-turboquant-head-dim-256-design.md b/docs/superpowers/specs/2026-08-10-turboquant-head-dim-256-design.md new file mode 100644 index 000000000..eda6bad76 --- /dev/null +++ b/docs/superpowers/specs/2026-08-10-turboquant-head-dim-256-design.md @@ -0,0 +1,109 @@ +# TurboQuant für head_dim als Vielfaches von 128 + +**Datum:** 2026-08-10 +**Status:** Design abgenommen +**Ziel:** Qwen3.5-9B (head_dim 256) mit `--cache-type-k turbo3/turbo4` betreiben + +## Problem + +`llama-kv-cache.cpp` bricht die Kontext-Initialisierung ab, sobald ein Modell +head_dim ≠ 128 hat: + +``` +llama_kv_cache: TurboQuant requires head_dim=128, got 256 (layer 3) +llama_init_from_model: failed to initialize the context: turbo types require head_dim=128 +``` + +Qwen3.5 (`Qwen3_5ForConditionalGeneration`) verwendet head_dim 256 und fällt +damit heraus. Betroffen sind die 8 `full_attention`-Layer; die übrigen 24 +`linear_attention`-Layer halten einen rekurrenten State in F32 und berühren den +KV-Cache nicht. + +## Befund: der Datenpfad ist bereits chunk-basiert + +Die Beschränkung liegt ausschließlich im Guard, nicht in der Verarbeitung. +Alle drei Ebenen arbeiten in 128-Element-Chunks über einen flachen Speicher- +bereich und kennen den Begriff „Head" gar nicht: + +| Ebene | Datei | Verhalten | +|---|---|---| +| CPU-Referenz | `ggml/src/ggml-quants.c:612` | `for (offset = 0; offset < k; offset += TURBO_HEAD_DIM)` | +| GPU-Write | `ggml/src/ggml-cuda/set-rows.cu:444` | `GGML_ASSERT(ne00 % TURBO_HEAD_DIM_SR == 0)`, ein CUDA-Block je Chunk | +| GPU-Read | `ggml/src/ggml-cuda/convert.cu:731` | `num_chunks = k / TURBO_HEAD_DIM_GPU`, ein CUDA-Block je Chunk | + +Alle drei prüfen bereits auf Teilbarkeit, nicht auf Gleichheit. Ein Modell mit +head_dim 256 durchläuft sie unverändert korrekt — es entstehen zwei Chunks pro +Head statt einem. + +## Ansatz: 2×128-Chunking + +Ein 256-dimensionaler Head wird als zwei unabhängige 128er-Hälften behandelt. +Jede Hälfte wird separat L2-normalisiert, per FWHT rotiert und gegen das +bestehende Codebook quantisiert. + +``` +Head (256 dim) + ├─ Hälfte A (128) → norm_A → FWHT128 → Codebook → 4 Blöcke + └─ Hälfte B (128) → norm_B → FWHT128 → Codebook → 4 Blöcke +``` + +### Warum die Codebooks gültig bleiben + +Die Zentroide sind Lloyd-Max-optimal für die Beta((d−1)/2, (d−1)/2)-Verteilung, +die entsteht, wenn ein **Einheitsvektor in R^128** per FWHT rotiert wird. Nach +der separaten Normalisierung ist jede Hälfte genau das: ein Einheitsvektor in +R^128. Die Verteilungsannahme gilt also exakt weiter — im Gegensatz zu einer +FWHT über 256 Dimensionen, die Beta(127,5, 127,5) erzeugen würde und neue +Zentroide (um Faktor 1/√2 enger) erfordert hätte. + +### Warum kein zusätzlicher Speicher entsteht + +Die Norm wird ohnehin redundant in jedem 32er-Block abgelegt +(`ggml-quants.c:636`, „Store same norm in every block of this chunk"). Zwei +Halb-Normen pro 256er-Head belegen dieselben Bytes wie eine Norm pro 128er-Head. +Die Bitrate bleibt unverändert bei 3,5 bpw (turbo3) bzw. 4,5 bpw (turbo4). + +### Warum Chunks nie Head-Grenzen überspannen + +Heads liegen kontiguierlich in der Row (`ne00 = n_head_kv × head_dim`). Da +head_dim ein Vielfaches der Chunk-Größe 128 ist, fällt jede Chunk-Grenze mit +einer Head-Grenze oder einer Head-Mitte zusammen, nie über einen Head hinaus. + +## Änderungen + +| Datei | Änderung | +|---|---| +| `src/llama-kv-cache.cpp:136-149` | Guard `!= 128` → `% 128 != 0`, Fehlertext anpassen | +| `tests/test-turboquant.cpp` | head_dim parametrisieren, Fälle 256 und 384 ergänzen | +| `docs/turboquant.md` | Requirements/Limitations korrigieren, Benchmark-Zeilen für Qwen3.5 | + +`ggml-quants.c`, `set-rows.cu` und `convert.cu` bleiben unverändert. + +## Nicht im Umfang + +- head_dim 80/96/112 (kein Vielfaches von 128) — der Guard lehnt weiterhin ab +- Neu berechnete Codebooks für echte FWHT über 256 +- Fusionierter Flash-Attention-Kernel (die Pre-Dequantize-Strategie bleibt) + +## Verifikation + +1. **Unit-Tests** — `test-turboquant` mit head_dim 128, 256, 384: FWHT-Selbst- + inversion, Roundtrip-MSE·d gegen die Paper-Bereiche, Bitpack-Determinismus. + Erwartung: MSE·d bei 256 im selben Bereich wie bei 128. +2. **Modell lädt** — Qwen3.5-9B-Q8_0 mit `-ctk turbo3 -ctv turbo3` startet und + erzeugt kohärenten Text. +3. **Perplexität** — Qwen3.5-9B: f16 vs. turbo4 vs. turbo3. Erwartung: + Delta in der Größenordnung der dokumentierten Werte (turbo4 ≈ +0,01, + turbo3 ≈ +0,05). +4. **Durchsatz** — `llama-bench` pp512/tg128 gegen f16-Baseline. +5. **VRAM** — gemessener KV-Verbrauch gegen die Rechnung + (turbo3 bei head_dim 256: 7,3 KiB/Token über 8 Attention-Layer). +6. **Vision** — mmproj-Pfad funktioniert weiterhin mit aktivem turbo3. + +## Risiko und Rückfallweg + +Die zwei Halb-Normen pro Head sind ein realer Unterschied zum d=128-Fall. Sie +erfassen Varianzunterschiede zwischen Head-Hälften feiner, was die Qualität eher +verbessert; belegt ist das erst durch Schritt 3. Fällt die Perplexität deutlich +schlechter aus als erwartet, ist der Rückfallweg eine echte FWHT über 256 mit +neu berechneten Lloyd-Max-Codebooks für Beta(127,5, 127,5). diff --git a/docs/turboquant.md b/docs/turboquant.md index 95e1d1831..e03e20b0c 100644 --- a/docs/turboquant.md +++ b/docs/turboquant.md @@ -35,6 +35,133 @@ All benchmarks: Qwen3-14B Q4_K_M, AMD RX 9070 XT (16 GB VRAM, gfx1201), ROCm 6.1 | **turbo4/turbo4** | **1812** | **1816** | **1321** | **926** | **550** | **460** | **49.3** | **4.5** | | **turbo3/turbo3** | **1836** | **1816** | **1319** | **924** | **548** | — | **49.6** | **3.5** | +### Throughput — Qwen3.5-9B Q8_0 (head_dim 256) + +RX 9070 XT, `-fa 1 -ngl 99`, 2 runs each. Note the reversed picture compared to +head_dim 128: turbo is substantially *faster* at prompt processing here. + +| type_k | type_v | pp512 | tg128 | +|--------|--------|-------|-------| +| f16 | f16 | 2565 ± 175 | 56.4 | +| f16 | turbo4 | 4084 ± 118 | 55.8 | +| turbo4 | turbo4 | 3929 ± 34 | 55.3 | +| turbo4 | f16 | 3981 ± 28 | 55.7 | +| turbo3 | turbo3 | 3933 ± 25 | 55.2 | +| turbo3 | f16 | 3966 ± 17 | 55.9 | + +Any turbo operand lifts pp512 by ~53% over f16/f16 while token generation stays +flat. The likely cause is a different flash-attention kernel being selected once +an operand is a turbo type — unverified, worth a look with `GGML_CUDA_DEBUG`. + +KV cache at 131072 context, 8 attention layers (Qwen3.5 is hybrid — the other +24 layers are linear attention with a constant-size recurrent state): + +| KV Type | KV cache | Total VRAM | Savings | +|---------|----------|-----------|---------| +| f16 | 4096 MiB | 14.81 GB | — | +| turbo4 | 1152 MiB | 11.75 GB | 72% | +| turbo3 | 896 MiB | 11.47 GB | 78% | + +At 131072 context f16 needs 14.81 GB of the card's 16.3 GB, leaving almost no +headroom for the vision encoder. turbo3 brings that down to 11.47 GB. + +### Token generation degrades with context depth + +The pre-dequantize strategy converts the whole KV cache to f16 once per FA call. +That cost is O(n_kv) per generated token, so it grows with context depth while the +attention work per token stays flat. Measured on Qwen3.5-9B Q8_0, RX 9070 XT: + +| type_k | type_v | tg128 @ d0 | tg128 @ d16384 | tg128 @ d65536 | +|--------|--------|-----------|---------------|---------------| +| f16 | f16 | 56.8 | 54.1 | **46.9** | +| f16 | turbo4 | 55.9 | 48.0 | 32.3 | +| f16 | turbo3 | 56.1 | 47.3 | 31.5 | +| turbo4 | f16 | 55.7 | 48.0 | 32.5 | +| turbo3 | f16 | 56.4 | 48.2 | 31.9 | +| turbo3 | turbo4 | 55.9 | 43.1 | **24.5** | + +One quantized cache costs ~32 t/s at depth 65536, two cost ~24.5, f16 gets 46.9. +The penalty tracks the *number* of converted caches, not which one — exactly what +a per-cache bulk conversion predicts. + +Prompt processing is unaffected (turbo is in fact faster there, 1012 vs 730 t/s at +d65536) because the single conversion amortizes over the whole batch. + +### Fused kernel: implemented, correct, currently slower + +A fused vector kernel that decodes turbo blocks inside flash attention exists +(`fattn-vec-turbo.cuh`). It removes the bulk conversion entirely by working in +the Hadamard domain — Q is transformed once at kernel start, V is accumulated +untransformed and transformed once at the end, so two transforms per launch +replace two per KV token. + +It is correct (test-backend-ops FLASH_ATTN_EXT: turbo3 640/640, turbo4 640/640) +but **slower than the path it replaces**, measured A/B in the same build: + +| depth | fused (KQ=32) | fused (KQ=8) | bulk conversion | f16 | +|-------|--------------|-------------|-----------------|-----| +| 0 | 52.5 | 53.2 | 55.0 | 56.8 | +| 16384 | 35.7 | 39.2 | 42.5 | 54.1 | +| 65536 | 16.6 | **20.4** | **23.6** | **46.9** | + +Narrowing the KQ reduction from 32 lanes per token to 8 (one warp reduction is +then 3 stages instead of 5, and 4 tokens are dotted at once) bought 23% at depth +65536. It does not close the gap, because the remaining cause is structural: + +**GQA amortization.** For `Q->ne[1] == 1` with `gqa_opt_applies`, f16 does not +use the vector kernel at all — `fattn.cu` routes it to the tile kernel, which +processes all `gqa_ratio` query heads per KV read. The fused vector kernel +handles one query head per block and therefore reads K and V `gqa_ratio` times +as often. Per element and query head, with Qwen3.5's gqa_ratio = 4: + + tile + f16: 2 bytes / 4 heads = 0.50 bytes + fused + turbo3: 0.4375 bytes x 1 = 0.4375 bytes + +The 4.57x compression collapses to a 12% bandwidth edge, which the unpacking and +codebook lookups more than consume. No amount of tuning inside the vector kernel +recovers this; a fused kernel has to support GQA to beat the tile path. + +Therefore **disabled by default**. Enable with `GGML_CUDA_TURBO_FUSED_FA=1`. + +### Fused tile path: K and V decoded inside attention + +The tile kernel is what f16 uses for decode with GQA, so decoding there avoids +the vector kernel's GQA penalty. Both sides are now fused (`type_K`/`type_V` +threaded through `flash_attn_tile` / `_iter` / `_iter_KQ`, `need_f16_K/V = false`). + +The decisive detail is *where* the Hadamard transform happens. A first version +decoded each chunk as it was loaded — one inverse FWHT per KV token, the same +O(n_kv) cost as the bulk conversion, just relocated. It measured exactly as fast +as what it replaced. Working in the Hadamard domain instead reduces that to O(1): + + V: sum_j p_j H c_j = H (sum_j p_j c_j) accumulate raw, transform once at the end + K: = transform Q once at kernel start + +tg128 at depth 65536, Qwen3.5-9B Q8_0 on gfx1201: + +| stage | t/s | of f16 | +|-------|-----|--------| +| bulk conversion (baseline) | 23.58 | 0.50x | +| V fused | 26.62 | 0.57x | +| K and V fused | **29.59** | **0.63x** | +| f16 reference | 46.90 | 1.00x | + +0.63x matches what llama.cpp discussion #20969 reports for quantized KV at 110K +context. Quantized KV is slower than f16 at decode across methods and backends — +that is a property of the approach, not of this implementation. What was +recoverable has been recovered. + +Two bugs worth recording, both found by the test matrix rather than by reading: + +- The K tile loader assumed a slice narrower than one 128-element chunk. True for + `nbatch_K = 64`, false for 128 and 256, which the config table also selects. + The chunk (and its norm) is now derived per element. +- Staging all of Q for the transform needed `ncols*DKQ` floats in `KV_tmp`, which + is sized by `nbatch_fa*nbatch_K` — and those shrink as `ncols` grows. At + `ncols=64` the buffer was short by a factor of 1.8. Transforming one chunk at a + time needs 128 floats regardless of `ncols`. This is why only large batches + failed while `nb=1` stayed green. + ### Perplexity (lower is better) Measured on ~960KB C++ source code corpus, context=2048. @@ -46,6 +173,45 @@ Measured on ~960KB C++ source code corpus, context=2048. | turbo4 | 1.7134 | +0.010 | | turbo3 | 1.7544 | +0.051 | +Qwen3.5-9B Q8_0 (head_dim 256), 600KB C++ corpus, context=2048: + +| KV Type | PPL | Delta vs f16 | +|---------|-----|-------------| +| f16 | 1.5258 ± 0.0071 | — | +| turbo4 | 1.5281 ± 0.0071 | +0.002 | +| turbo3 | 1.5343 ± 0.0072 | +0.009 | + +Both deltas are markedly smaller than at head_dim 128 (+0.010 / +0.051). Splitting +a 256-dim head into two independently normalized halves stores two norms instead +of one, which tracks variance differences between the halves more closely — at no +extra cost, since the norm is replicated per block either way. + +#### The error does not accumulate with context depth + +KVarN (arXiv 2606.03458, llama.cpp #24139) argues that KV-cache quantization +error compounds during autoregressive decoding, so a short-context perplexity +number can understate the real cost. Measured against that concern: + +| context | f16 | turbo4 | delta | turbo3 | delta | +|---------|-----|--------|-------|--------|-------| +| 2048 | 1.5258 | 1.5276 | +0.0018 | 1.5338 | +0.0080 | +| 8192 | 1.4107 | 1.4124 | +0.0017 | 1.4172 | +0.0065 | +| 16384 | 1.3431 | not measured | — | 1.3497 | +0.0066 | + +The gap does not widen. For turbo3 it drops from +0.0080 at 2048 to +0.0065 at +8192 and then holds at +0.0066 through 16384 — 0.52%, 0.46%, 0.49% in relative +terms. Three depths across an 8x range, flat after the first step. + +(turbo4 at 16384 is missing; its delta is ~0.0017 at both measured depths, small +enough that the trend question hardly applies.) + +That is consistent with how TurboQuant stores data: the per-chunk L2 norm is kept +in fp16 and never quantized, so only the *direction* of the normalized vector is +approximated. KVarN attributes the accumulating error specifically to incorrect +per-token magnitude scaling — the mechanism does not exist here. The per-chunk +normalization that makes TurboQuant expensive at runtime is the same property +that makes it robust over long contexts. + ### KV Cache Memory | KV Type | Bytes per element | Savings vs f16 | @@ -67,6 +233,14 @@ where the Lloyd-Max codebook is optimal. The codebook centroids are precomputed for the Beta((d-1)/2, (d-1)/2) distribution that arises after FWHT of unit vectors in d=128 dimensions. +The transform is always applied in fixed 128-element chunks, walking the row. +A head larger than 128 is therefore split into independently normalized chunks: +a 256-dim head (Qwen3.5) becomes two halves, each with its own L2 norm and its +own FWHT. Because every normalized half is again a unit vector in R^128, the +same codebooks stay optimal — no per-head-size codebook is needed. The stored +norm is already replicated per 32-element block, so the extra half-norm costs +no additional bytes and the bit rate is unchanged. + Current implementation uses a pre-dequantize strategy: turbo KV data is bulk-converted to f16 before standard Flash Attention runs. This adds minimal overhead (~3% pp, ~8% tg) while avoiding the complexity of a fused FA kernel. @@ -74,12 +248,14 @@ while avoiding the complexity of a fused FA kernel. ## Requirements - Flash Attention enabled (`-fa 1`) -- head_dim = 128 (covers Llama, Qwen, Mistral, Gemma, and most current models) +- head_dim must be a multiple of 128 — 128 covers Llama, Qwen3, Mistral and + Gemma; 256 covers Qwen3.5 - AMD ROCm (gfx1201 tested) or CPU ## Limitations -- head_dim must be exactly 128 (FWHT requirement) +- head_dim must be a multiple of 128 (FWHT chunk size). Models with head_dim + 80, 96 or 112 are not supported. - Mixed turbo/quantized combinations (e.g. turbo4/q8_0) are not supported — use turbo/turbo or turbo/f16 - No CUDA support yet (HIP/ROCm only for GPU path) - Pre-dequantize strategy means full KV cache is converted to f16 per FA call (a fused kernel would eliminate this) diff --git a/ggml/src/ggml-cpu/ggml-cpu.c b/ggml/src/ggml-cpu/ggml-cpu.c index df17cc553..4d6ead497 100644 --- a/ggml/src/ggml-cpu/ggml-cpu.c +++ b/ggml/src/ggml-cpu/ggml-cpu.c @@ -393,6 +393,24 @@ static const struct ggml_type_traits_cpu type_traits_cpu[GGML_TYPE_COUNT] = { [GGML_TYPE_I32] = { .from_float = (ggml_from_float_t) ggml_cpu_fp32_to_i32, }, + // TurboQuant: KV-cache only. vec_dot takes f32 on the right-hand side and + // dequantizes the left-hand side internally (inverse FWHT + norm), so + // vec_dot_type is F32. Without these entries the CPU flash-attention path + // calls a null vec_dot pointer, which is what test-backend-ops hits when it + // computes its CPU reference. + // No .from_float: quantize_row_turbo*_ref asserts k % 128 == 0, which the + // generic per-block callers do not honor. Quantization goes through + // ggml_quantize_chunk instead. + [GGML_TYPE_TURBO3_0] = { + .vec_dot = ggml_vec_dot_turbo3_0, + .vec_dot_type = GGML_TYPE_F32, + .nrows = 1, + }, + [GGML_TYPE_TURBO4_0] = { + .vec_dot = ggml_vec_dot_turbo4_0, + .vec_dot_type = GGML_TYPE_F32, + .nrows = 1, + }, }; const struct ggml_type_traits_cpu * ggml_get_type_traits_cpu(enum ggml_type type) { diff --git a/ggml/src/ggml-cpu/quants.h b/ggml/src/ggml-cpu/quants.h index 3584aaa43..cb82824a9 100644 --- a/ggml/src/ggml-cpu/quants.h +++ b/ggml/src/ggml-cpu/quants.h @@ -42,6 +42,11 @@ void ggml_vec_dot_q5_0_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const voi void ggml_vec_dot_q5_1_q8_1(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc); void ggml_vec_dot_q8_0_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc); +// TurboQuant. Defined in ggml-quants.c (not arch-specific): dequantizes the +// left-hand side internally, so the right-hand side is plain f32. +void ggml_vec_dot_turbo3_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc); +void ggml_vec_dot_turbo4_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc); + void ggml_vec_dot_mxfp4_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc); void ggml_vec_dot_nvfp4_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc); diff --git a/ggml/src/ggml-cuda/CMakeLists.txt b/ggml/src/ggml-cuda/CMakeLists.txt index 1824afe4a..419862101 100644 --- a/ggml/src/ggml-cuda/CMakeLists.txt +++ b/ggml/src/ggml-cuda/CMakeLists.txt @@ -120,9 +120,7 @@ if (CUDAToolkit_FOUND) template-instances/fattn-vec-instance-f16-f16.cu template-instances/fattn-vec-instance-q4_0-q4_0.cu template-instances/fattn-vec-instance-q8_0-q8_0.cu - template-instances/fattn-vec-instance-bf16-bf16.cu - template-instances/fattn-vec-instance-turbo3_0-turbo3_0.cu - template-instances/fattn-vec-instance-turbo4_0-turbo4_0.cu) + template-instances/fattn-vec-instance-bf16-bf16.cu) endif() ggml_add_backend_library(ggml-cuda diff --git a/ggml/src/ggml-cuda/convert.cu b/ggml/src/ggml-cuda/convert.cu index 1aa64b65d..50cc56d1d 100644 --- a/ggml/src/ggml-cuda/convert.cu +++ b/ggml/src/ggml-cuda/convert.cu @@ -1,5 +1,6 @@ #include "convert.cuh" #include "dequantize.cuh" +#include "turbo.cuh" #include @@ -667,9 +668,6 @@ static void dequantize_row_nvfp4_cuda( // Step 3: normalize by 1/sqrt(128) and scale by stored norm // Step 4: write to output -#define TURBO_HEAD_DIM_GPU 128 -#define TURBO_BLOCKS_PER_CHUNK_GPU (TURBO_HEAD_DIM_GPU / 32) // 4 - template static __global__ void dequantize_block_turbo3_0_kernel(const void * __restrict__ vx, dst_t * __restrict__ y, const int64_t k) { __shared__ float smem[TURBO_HEAD_DIM_GPU]; @@ -680,49 +678,9 @@ static __global__ void dequantize_block_turbo3_0_kernel(const void * __restrict_ const int64_t output_offset = chunk_idx * TURBO_HEAD_DIM_GPU; if (output_offset + tid >= k) return; - // Which of the 4 blocks within this chunk does this thread belong to? - const int local_block = tid / TURBO3_BLOCK_SIZE; // 0..3 - const int elem_in_block = tid % TURBO3_BLOCK_SIZE; // 0..31 - - const int64_t global_block_idx = chunk_idx * TURBO_BLOCKS_PER_CHUNK_GPU + local_block; - - // Unpack 3-bit index and look up centroid - const block_turbo3_0 * x = (const block_turbo3_0 *)vx + global_block_idx; - const uint8_t * qs = x->qs; - - int bit_off = elem_in_block * 3; - int byte_idx = bit_off / 8; - int shift = bit_off % 8; - uint16_t raw = (uint16_t)qs[byte_idx] >> shift; - if (shift > 5 && byte_idx + 1 < 12) - raw |= (uint16_t)qs[byte_idx + 1] << (8 - shift); - uint8_t idx = (uint8_t)(raw & 0x07); + turbo_decode_chunk_smem( + (const block_turbo3_0 *)vx, chunk_idx * TURBO_BLOCKS_PER_CHUNK_GPU, smem, tid); - smem[tid] = dc_codebook_3bit[idx]; - __syncthreads(); - - // FWHT butterfly stages (7 stages for n=128) - for (int h = 1; h < TURBO_HEAD_DIM_GPU; h *= 2) { - if (tid < 64) { // 128/2 = 64 butterflies per stage - int group = tid / h; - int pos = tid % h; - int i = group * h * 2 + pos; - float a = smem[i]; - float b = smem[i + h]; - smem[i] = a + b; - smem[i + h] = a - b; - } - __syncthreads(); - } - - // Normalize by 1/sqrt(128) and scale by stored norm - const float fwht_scale = 0.08838834764831844f; // 1/sqrt(128) - const block_turbo3_0 * first_block = (const block_turbo3_0 *)vx + chunk_idx * TURBO_BLOCKS_PER_CHUNK_GPU; - float norm = __half2float(first_block->d); - smem[tid] *= fwht_scale * norm; - __syncthreads(); - - // Write to output y[output_offset + tid] = ggml_cuda_cast(smem[tid]); } @@ -743,51 +701,94 @@ static __global__ void dequantize_block_turbo4_0_kernel(const void * __restrict_ const int64_t output_offset = chunk_idx * TURBO_HEAD_DIM_GPU; if (output_offset + tid >= k) return; - // Which of the 4 blocks within this chunk does this thread belong to? - const int local_block = tid / TURBO4_BLOCK_SIZE; // 0..3 - const int elem_in_block = tid % TURBO4_BLOCK_SIZE; // 0..31 + turbo_decode_chunk_smem( + (const block_turbo4_0 *)vx, chunk_idx * TURBO_BLOCKS_PER_CHUNK_GPU, smem, tid); - const int64_t global_block_idx = chunk_idx * TURBO_BLOCKS_PER_CHUNK_GPU + local_block; + y[output_offset + tid] = ggml_cuda_cast(smem[tid]); +} - // Unpack 4-bit index and look up centroid - const block_turbo4_0 * x = (const block_turbo4_0 *)vx + global_block_idx; - int pair_idx = elem_in_block / 2; - uint8_t packed = x->qs[pair_idx]; - uint8_t idx = (elem_in_block & 1) ? ((packed >> 4) & 0x0F) : (packed & 0x0F); +template +static void dequantize_row_turbo4_0_cuda(const void * vx, dst_t * y, const int64_t k, cudaStream_t stream) { + GGML_ASSERT(k % TURBO_HEAD_DIM_GPU == 0); + const int num_chunks = (int)(k / TURBO_HEAD_DIM_GPU); + dequantize_block_turbo4_0_kernel<<>>(vx, y, k); +} - smem[tid] = dc_codebook_4bit[idx]; - __syncthreads(); +// ============================================================ +// TurboQuant strided (non-contiguous) dequantize +// ============================================================ +// +// Same decode as the bulk kernels above, but the source block index is computed +// from the row/slice strides instead of assuming a flat buffer. Needed because a +// KV-cache view is generally not contiguously allocated: nb[3] can skip over +// unused cells, so walking the source linearly reads past the end. +// +// s01/s02/s03 are in block units (nb / ggml_type_size), matching the convention +// of the other *_nc_ kernels. + +template +static __global__ void dequantize_block_turbo_nc_kernel( + const void * __restrict__ vx, dst_t * __restrict__ y, + const int64_t ne00, const int64_t ne01, + const int64_t ne0203, const uint3 ne02, + const int64_t s01, const int64_t s02, const int64_t s03) { + __shared__ float smem[TURBO_HEAD_DIM_GPU]; - // FWHT butterfly stages (7 stages for n=128) - for (int h = 1; h < TURBO_HEAD_DIM_GPU; h *= 2) { - if (tid < 64) { // 128/2 = 64 butterflies per stage - int group = tid / h; - int pos = tid % h; - int i = group * h * 2 + pos; - float a = smem[i]; - float b = smem[i + h]; - smem[i] = a + b; - smem[i + h] = a - b; + const int64_t chunk_in_row = blockIdx.x; + const int tid = threadIdx.x; // 0..127 + + const block_t * blocks = (const block_t *) vx; + + for (int64_t i01 = blockIdx.y; i01 < ne01; i01 += gridDim.y) { + for (int64_t i0203 = blockIdx.z; i0203 < ne0203; i0203 += gridDim.z) { + const uint2 dm = fast_div_modulo((uint32_t)i0203, ne02); + const int64_t i02 = dm.y; + const int64_t i03 = dm.x; + + const int64_t ibx0 = i03*s03 + i02*s02 + i01*s01; + + turbo_decode_chunk_smem( + blocks, ibx0 + chunk_in_row*TURBO_BLOCKS_PER_CHUNK_GPU, smem, tid); + + const int64_t iy = (i0203*ne01 + i01)*ne00 + chunk_in_row*TURBO_HEAD_DIM_GPU + tid; + y[iy] = ggml_cuda_cast(smem[tid]); + + // smem is reused by the next iteration + __syncthreads(); } - __syncthreads(); } +} - // Normalize by 1/sqrt(128) and scale by stored norm - const float fwht_scale = 0.08838834764831844f; // 1/sqrt(128) - const block_turbo4_0 * first_block = (const block_turbo4_0 *)vx + chunk_idx * TURBO_BLOCKS_PER_CHUNK_GPU; - float norm = __half2float(first_block->d); - smem[tid] *= fwht_scale * norm; - __syncthreads(); +template +static void dequantize_turbo_nc_cuda(const void * vx, dst_t * y, + const int64_t ne00, const int64_t ne01, const int64_t ne02, const int64_t ne03, + const int64_t s01, const int64_t s02, const int64_t s03, cudaStream_t stream) { + GGML_ASSERT(ne00 % TURBO_HEAD_DIM_GPU == 0); - // Write to output - y[output_offset + tid] = ggml_cuda_cast(smem[tid]); + const int64_t ne0203 = ne02*ne03; + const int64_t chunks_per_row = ne00 / TURBO_HEAD_DIM_GPU; + const uint3 ne02_fdv = init_fastdiv_values(ne02); + + const dim3 num_blocks((int)chunks_per_row, + (int)std::min(ne01, (int64_t)65535), + (int)std::min(ne0203, (int64_t)65535)); + + dequantize_block_turbo_nc_kernel<<>> + (vx, y, ne00, ne01, ne0203, ne02_fdv, s01, s02, s03); } template -static void dequantize_row_turbo4_0_cuda(const void * vx, dst_t * y, const int64_t k, cudaStream_t stream) { - GGML_ASSERT(k % TURBO_HEAD_DIM_GPU == 0); - const int num_chunks = (int)(k / TURBO_HEAD_DIM_GPU); - dequantize_block_turbo4_0_kernel<<>>(vx, y, k); +static void dequantize_row_turbo3_0_nc_cuda(const void * vx, dst_t * y, + const int64_t ne00, const int64_t ne01, const int64_t ne02, const int64_t ne03, + const int64_t s01, const int64_t s02, const int64_t s03, cudaStream_t stream) { + dequantize_turbo_nc_cuda(vx, y, ne00, ne01, ne02, ne03, s01, s02, s03, stream); +} + +template +static void dequantize_row_turbo4_0_nc_cuda(const void * vx, dst_t * y, + const int64_t ne00, const int64_t ne01, const int64_t ne02, const int64_t ne03, + const int64_t s01, const int64_t s02, const int64_t s03, cudaStream_t stream) { + dequantize_turbo_nc_cuda(vx, y, ne00, ne01, ne02, ne03, s01, s02, s03, stream); } template @@ -975,9 +976,11 @@ to_fp16_nc_cuda_t ggml_get_to_fp16_nc_cuda(ggml_type type) { case GGML_TYPE_Q8_0: return dequantize_block_cuda; case GGML_TYPE_TURBO3_0: - return dequantize_block_cuda; + // NOT dequantize_block_cuda: TurboQuant cannot be decoded element-wise, + // a value needs the inverse FWHT over its whole 128-element chunk. + return dequantize_row_turbo3_0_nc_cuda; case GGML_TYPE_TURBO4_0: - return dequantize_block_cuda; + return dequantize_row_turbo4_0_nc_cuda; case GGML_TYPE_BF16: return convert_unary_cuda; default: @@ -1000,9 +1003,11 @@ to_bf16_nc_cuda_t ggml_get_to_bf16_nc_cuda(ggml_type type) { case GGML_TYPE_Q8_0: return dequantize_block_cuda; case GGML_TYPE_TURBO3_0: - return dequantize_block_cuda; + // NOT dequantize_block_cuda: TurboQuant cannot be decoded element-wise, + // a value needs the inverse FWHT over its whole 128-element chunk. + return dequantize_row_turbo3_0_nc_cuda; case GGML_TYPE_TURBO4_0: - return dequantize_block_cuda; + return dequantize_row_turbo4_0_nc_cuda; case GGML_TYPE_F16: return convert_unary_cuda; default: @@ -1025,9 +1030,11 @@ to_fp32_nc_cuda_t ggml_get_to_fp32_nc_cuda(ggml_type type) { case GGML_TYPE_Q8_0: return dequantize_block_cuda; case GGML_TYPE_TURBO3_0: - return dequantize_block_cuda; + // NOT dequantize_block_cuda: TurboQuant cannot be decoded element-wise, + // a value needs the inverse FWHT over its whole 128-element chunk. + return dequantize_row_turbo3_0_nc_cuda; case GGML_TYPE_TURBO4_0: - return dequantize_block_cuda; + return dequantize_row_turbo4_0_nc_cuda; case GGML_TYPE_BF16: return convert_unary_cuda; default: diff --git a/ggml/src/ggml-cuda/dequantize.cuh b/ggml/src/ggml-cuda/dequantize.cuh index 1e4353f43..e6a845182 100644 --- a/ggml/src/ggml-cuda/dequantize.cuh +++ b/ggml/src/ggml-cuda/dequantize.cuh @@ -1,3 +1,5 @@ +#pragma once + #include "common.cuh" static __device__ __forceinline__ void dequantize_q4_0(const void * vx, const int64_t ib, const int iqs, float2 & v){ diff --git a/ggml/src/ggml-cuda/fattn-common.cuh b/ggml/src/ggml-cuda/fattn-common.cuh index c9eeba987..047b558fe 100644 --- a/ggml/src/ggml-cuda/fattn-common.cuh +++ b/ggml/src/ggml-cuda/fattn-common.cuh @@ -578,58 +578,6 @@ static __device__ __forceinline__ void dequantize_V_q8_0(const void * __restrict } } -// ============================================================ -// TurboQuant V-cache dequantize functions for flash attention -// ============================================================ - -template -static __device__ __forceinline__ void dequantize_V_turbo3_0(const void * __restrict__ vx, void * __restrict__ dst, const int64_t i0) { - const block_turbo3_0 * x = (const block_turbo3_0 *) vx; - - const int64_t ib = i0 / QK_TURBO3; - const int iqs = (int)(i0 % QK_TURBO3) / 2; - - static_assert(ne % 2 == 0, "bad ne"); - T * dst_t = (T *) dst; - -#pragma unroll - for (int l = 0; l < ne/2; ++l) { - float2 v; - dequantize_turbo3_0(vx, ib, iqs + l, v); - if constexpr (std::is_same_v) { - dst_t[2*l + 0] = __float2half(v.x); - dst_t[2*l + 1] = __float2half(v.y); - } else { - dst_t[2*l + 0] = (T)v.x; - dst_t[2*l + 1] = (T)v.y; - } - } -} - -template -static __device__ __forceinline__ void dequantize_V_turbo4_0(const void * __restrict__ vx, void * __restrict__ dst, const int64_t i0) { - const block_turbo4_0 * x = (const block_turbo4_0 *) vx; - - const int64_t ib = i0 / QK_TURBO4; - const int iqs = (int)(i0 % QK_TURBO4) / 2; - - static_assert(ne % 2 == 0, "bad ne"); - T * dst_t = (T *) dst; - -#pragma unroll - for (int l = 0; l < ne/2; ++l) { - float2 v; - dequantize_turbo4_0(vx, ib, iqs + l, v); - if constexpr (std::is_same_v) { - dst_t[2*l + 0] = __float2half(v.x); - dst_t[2*l + 1] = __float2half(v.y); - } else { - dst_t[2*l + 0] = (T)v.x; - dst_t[2*l + 1] = (T)v.y; - } - } - GGML_UNUSED(x); -} template constexpr __device__ vec_dot_KQ_t get_vec_dot_KQ() { @@ -647,12 +595,10 @@ constexpr __device__ vec_dot_KQ_t get_vec_dot_KQ() { return vec_dot_fattn_vec_KQ_q8_0; } else if constexpr (type_K == GGML_TYPE_BF16) { return vec_dot_fattn_vec_KQ_bf16; - // TurboQuant K-cache: Phase 1 - dequantize to FP16 before attention - // (use FP16 dot product after conversion in the dispatch layer) - } else if constexpr (type_K == GGML_TYPE_TURBO3_0) { - return vec_dot_fattn_vec_KQ_f16; - } else if constexpr (type_K == GGML_TYPE_TURBO4_0) { - return vec_dot_fattn_vec_KQ_f16; + // No turbo case on purpose. TurboQuant cannot be decoded element-wise — a + // value needs the inverse FWHT over its whole 128-element chunk — so there is + // no valid per-element dot product for it. Turbo KV reaches the vector kernel + // only through a dispatch bug; making that a compile error is the point. } else { static_assert(type_K == -1, "bad type"); return nullptr; @@ -675,10 +621,7 @@ constexpr __device__ dequantize_V_t get_dequantize_V() { return dequantize_V_q8_0; } else if constexpr (type_V == GGML_TYPE_BF16) { return dequantize_V_bf16; - } else if constexpr (type_V == GGML_TYPE_TURBO3_0) { - return dequantize_V_turbo3_0; - } else if constexpr (type_V == GGML_TYPE_TURBO4_0) { - return dequantize_V_turbo4_0; + // No turbo case — see get_vec_dot_KQ above. } else { static_assert(type_V == -1, "bad type"); return nullptr; diff --git a/ggml/src/ggml-cuda/fattn-tile.cu b/ggml/src/ggml-cuda/fattn-tile.cu index 3fcb09b7a..dd7182dde 100644 --- a/ggml/src/ggml-cuda/fattn-tile.cu +++ b/ggml/src/ggml-cuda/fattn-tile.cu @@ -5,6 +5,35 @@ void ggml_cuda_flash_attn_ext_tile(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { const ggml_tensor * K = dst->src[1]; const ggml_tensor * V = dst->src[2]; + + // TurboQuant V is decoded inside the kernel rather than expanded to f16 by + // launch_fattn. Only V matters here: the K tile holds nbatch_K < 128 columns + // at these head sizes, which is less than one FWHT chunk, so K still goes + // through the conversion path. + // Only the homogeneous case is instantiated; mixed turbo/f16 keeps using the + // conversion path, which stays correct because launch_fattn handles it. + if ((V->type == GGML_TYPE_TURBO3_0 || V->type == GGML_TYPE_TURBO4_0) && K->type == V->type) { + GGML_ASSERT(V->ne[0] == K->ne[0]); + switch (K->ne[0]) { + case 128: + if (V->type == GGML_TYPE_TURBO3_0) { + ggml_cuda_flash_attn_ext_tile_case<128, 128, GGML_TYPE_TURBO3_0, GGML_TYPE_TURBO3_0>(ctx, dst); + } else { + ggml_cuda_flash_attn_ext_tile_case<128, 128, GGML_TYPE_TURBO4_0, GGML_TYPE_TURBO4_0>(ctx, dst); + } + return; + case 256: + if (V->type == GGML_TYPE_TURBO3_0) { + ggml_cuda_flash_attn_ext_tile_case<256, 256, GGML_TYPE_TURBO3_0, GGML_TYPE_TURBO3_0>(ctx, dst); + } else { + ggml_cuda_flash_attn_ext_tile_case<256, 256, GGML_TYPE_TURBO4_0, GGML_TYPE_TURBO4_0>(ctx, dst); + } + return; + default: + GGML_ABORT("unsupported head size for turbo V"); + } + } + switch (K->ne[0]) { case 40: { GGML_ASSERT(V->ne[0] == K->ne[0]); diff --git a/ggml/src/ggml-cuda/fattn-tile.cuh b/ggml/src/ggml-cuda/fattn-tile.cuh index f3fa80ab2..d5e9e3e10 100644 --- a/ggml/src/ggml-cuda/fattn-tile.cuh +++ b/ggml/src/ggml-cuda/fattn-tile.cuh @@ -1,6 +1,7 @@ #include "common.cuh" #include "fattn-common.cuh" #include "fattn-wmma-f16.cuh" +#include "turbo.cuh" // nbatch_fa == number of KQ rows to process per iteration // nbatch_K == number of K columns to load in parallel for KQ calculation @@ -426,9 +427,131 @@ static __device__ __forceinline__ void flash_attn_tile_load_tile( ggml_cuda_unroll<5>{}(load); } +// TurboQuant variant: reads compressed blocks and decodes them straight into the +// tile, so the KV cache never has to be expanded to f16 beforehand. +// +// Rows are complete head_dim vectors here (the V tile is loaded with J = DV), so +// each row is J/128 whole FWHT chunks. The tile is filled with RAW codebook +// values scaled by norm/sqrt(128) — no rotation. Because H is linear and the +// softmax weights are scalars, sum_j p_j H c_j = H (sum_j p_j c_j): the kernel +// accumulates in the Hadamard domain and applies one inverse FWHT at the end, +// instead of one per KV token. The transform is a dependency chain; the +// accumulation is not, which is what makes this worth doing. +// +// KV is the raw block pointer; stride_KV_b is the row stride in bytes. +template +static __device__ __forceinline__ void flash_attn_tile_load_tile_turbo( + const char * const __restrict__ KV, tile_t * const __restrict__ tile_KV, + const int stride_KV_b, const int i_sup) { + static_assert(J % TURBO_HEAD_DIM_GPU == 0, "turbo tile rows must be whole chunks"); + constexpr int nchunks = J / TURBO_HEAD_DIM_GPU; + + const int warp = threadIdx.y; + const int lane = threadIdx.x; + +#pragma unroll + for (int i = 0; i < I; ++i) { + const bool oob = oob_check && i >= i_sup; + +#pragma unroll + for (int c0 = 0; c0 < nchunks; c0 += nwarps) { + const int c = c0 + warp; + if (nchunks % nwarps != 0 && c >= nchunks) { + continue; + } + + float v[4] = {0.0f, 0.0f, 0.0f, 0.0f}; + if (!oob) { + const block_t * row = (const block_t *) (KV + (size_t)i*stride_KV_b); + // Raw codebook values, no rotation: the accumulator stays in the + // Hadamard domain and is transformed once at the end of the kernel. + turbo_load_chunk_raw_warp(row, (int64_t)c*TURBO_BLOCKS_PER_CHUNK_GPU, v, lane); + } + + const int e0 = c*TURBO_HEAD_DIM_GPU + 4*lane; // first element this lane owns + if constexpr (std::is_same::value) { + half2 * dst = tile_KV + i*(J/2 + J_padding) + e0/2; + dst[0] = make_half2(v[0], v[1]); + dst[1] = make_half2(v[2], v[3]); + } else { + float * dst = tile_KV + i*(J + J_padding) + e0; +#pragma unroll + for (int r = 0; r < 4; ++r) { + dst[r] = v[r]; + } + } + } + } +} + +// TurboQuant K tile: raw codebook values scaled by norm/sqrt(128), no rotation. +// +// Unlike the V tile this loads a *slice* of each row — nbatch_K is 64 at the head +// sizes turbo supports, half a chunk. That is fine here: the KQ dot is +// element-wise, so it never needs a whole chunk at once. It only requires Q to +// have been transformed into the Hadamard domain beforehand, since +// = . +// +// A slice never straddles a chunk boundary (nbatch_K divides 128 and elem_off is +// a multiple of nbatch_K), so the norm is one load per row. +template +static __device__ __forceinline__ void flash_attn_tile_load_tile_turbo_k( + const char * const __restrict__ KV, tile_t * const __restrict__ tile_KV, + const int stride_KV_b, const int i_sup, const int elem_off) { + constexpr int bs = (sizeof(block_t) == sizeof(block_turbo3_0)) ? TURBO3_BLOCK_SIZE : TURBO4_BLOCK_SIZE; + constexpr int nthreads = warp_size*nwarps; + + // A slice may span several chunks (nbatch_K is 256 for some configs), so the + // chunk — and with it the norm — is derived per element rather than per row. + static_assert(J % 2 == 0 || J == 0, "turbo K slice must be even for half2 packing"); + + const int tid = warp_size*threadIdx.y + threadIdx.x; + + // Work in pairs so a single thread owns both halves of a half2 slot; letting + // two threads write the two halves would be a read-modify-write race. + constexpr int JP = std::is_same::value ? J/2 : J; + +#pragma unroll + for (int idx0 = 0; idx0 < I*JP; idx0 += nthreads) { + const int idx = idx0 + tid; + if (I*JP % nthreads != 0 && idx >= I*JP) { + continue; + } + const int i = idx / JP; + const int jp = idx % JP; + + const bool live = !oob_check || i < i_sup; + const block_t * row = (const block_t *) (KV + (size_t)i*stride_KV_b); + + auto fetch = [&] __device__ (const int j) -> float { + if (!live) { + return 0.0f; + } + const int e = elem_off + j; + const int chunk = e / TURBO_HEAD_DIM_GPU; + const int eic = e % TURBO_HEAD_DIM_GPU; + const int64_t ib0 = (int64_t)chunk * TURBO_BLOCKS_PER_CHUNK_GPU; + const float nrm = TURBO_FWHT_SCALE_GPU * __half2float(row[ib0].d); + if constexpr (sizeof(block_t) == sizeof(block_turbo3_0)) { + return dc_codebook_3bit[turbo3_unpack_index((const block_turbo3_0 *)(row + ib0 + eic/bs), eic % bs)] * nrm; + } else { + return dc_codebook_4bit[turbo4_unpack_index((const block_turbo4_0 *)(row + ib0 + eic/bs), eic % bs)] * nrm; + } + }; + + if constexpr (std::is_same::value) { + tile_KV[i*(J/2 + J_padding) + jp] = make_half2(fetch(2*jp), fetch(2*jp + 1)); + } else { + tile_KV[i*(J + J_padding) + jp] = fetch(jp); + } + } +} + // Function that performs a single iteration in for the KQ matrix multiplication: template + bool use_logit_softcap, bool oob_check, ggml_type type_K, typename T_vec_dot> static __device__ __forceinline__ void flash_attn_tile_iter_KQ( T_vec_dot * const Q_tmp, const half2 * const __restrict__ K_h2, @@ -445,8 +568,16 @@ static __device__ __forceinline__ void flash_attn_tile_iter_KQ( constexpr int cpw = ncols > nwarps ? ncols/nwarps : 1; // Q columns per warp constexpr int np = nwarps > ncols ? nwarps/ncols : 1; // number of parallel warps per Q column - flash_attn_tile_load_tile - (K_h2 + int64_t(k_VKQ_0)*stride_K2 + k_KQ_0/2, KV_tmp, stride_K2, k_VKQ_sup); + if constexpr (type_K == GGML_TYPE_TURBO3_0 || type_K == GGML_TYPE_TURBO4_0) { + using block_t = typename std::conditional::type; + const int stride_K_b = stride_K2 * (int) sizeof(half2); + flash_attn_tile_load_tile_turbo_k + ((const char *) K_h2 + (size_t)k_VKQ_0*stride_K_b, KV_tmp, stride_K_b, k_VKQ_sup, k_KQ_0); + } else { + flash_attn_tile_load_tile + (K_h2 + int64_t(k_VKQ_0)*stride_K2 + k_KQ_0/2, KV_tmp, stride_K2, k_VKQ_sup); + } __syncthreads(); #ifdef FAST_FP16_AVAILABLE @@ -503,7 +634,8 @@ static __device__ __forceinline__ void flash_attn_tile_iter_KQ( // Function that performs a single iteration of the main loop over up to nbatch_fa tokens. template + bool use_logit_softcap, bool oob_check, ggml_type type_K, ggml_type type_V, + typename T_vec_dot, typename T_KQ, typename T_acc> static __device__ __forceinline__ void flash_attn_tile_iter( T_vec_dot * const Q_tmp, const half2 * const __restrict__ K_h2, @@ -554,12 +686,12 @@ static __device__ __forceinline__ void flash_attn_tile_iter( constexpr int nbatch_K_last = DKQ % nbatch_K; #pragma unroll for (int k_KQ_0 = 0; k_KQ_0 < DKQ - nbatch_K_last; k_KQ_0 += nbatch_K) { - flash_attn_tile_iter_KQ( + flash_attn_tile_iter_KQ( Q_tmp, K_h2, KV_tmp, stride_K2, k_VKQ_0, k_VKQ_sup, k_KQ_0, KQ_acc); } if (nbatch_K_last > 0) { constexpr int k_KQ_0 = DKQ - nbatch_K_last; - flash_attn_tile_iter_KQ( + flash_attn_tile_iter_KQ( Q_tmp, K_h2, KV_tmp, stride_K2, k_VKQ_0, k_VKQ_sup, k_KQ_0, KQ_acc); } @@ -665,8 +797,19 @@ static __device__ __forceinline__ void flash_attn_tile_iter( static_assert(nbatch_V % np == 0, "bad nbatch_V"); #pragma unroll for (int k0 = 0; k0 < nbatch_fa; k0 += nbatch_V) { - flash_attn_tile_load_tile - (V_h2 + int64_t(k_VKQ_0 + k0)*stride_V2, KV_tmp, stride_V2, k_VKQ_sup - k0); + if constexpr (type_V == GGML_TYPE_TURBO3_0 || type_V == GGML_TYPE_TURBO4_0) { + // stride_V2 was derived as nb21/sizeof(half2); for turbo the row stride + // in bytes is what we need, and nb21 is always a multiple of 4 here + // (14 or 18 bytes per 32-element block, times a multiple of 4 blocks). + using block_t = typename std::conditional::type; + const int stride_V_b = stride_V2 * (int) sizeof(half2); + flash_attn_tile_load_tile_turbo + ((const char *) V_h2 + (size_t)(k_VKQ_0 + k0)*stride_V_b, KV_tmp, stride_V_b, k_VKQ_sup - k0); + } else { + flash_attn_tile_load_tile + (V_h2 + int64_t(k_VKQ_0 + k0)*stride_V2, KV_tmp, stride_V2, k_VKQ_sup - k0); + } __syncthreads(); #ifdef FAST_FP16_AVAILABLE @@ -735,7 +878,7 @@ static __device__ __forceinline__ void flash_attn_tile_iter( } } -template // D == head size +template // D == head size __launch_bounds__(ggml_cuda_fattn_tile_get_nthreads(DKQ, DV, ncols1*ncols2), ggml_cuda_fattn_tile_get_occupancy(DKQ, DV, ncols1*ncols2)) static __global__ void flash_attn_tile( const char * __restrict__ Q, @@ -892,6 +1035,68 @@ static __global__ void flash_attn_tile( __syncthreads(); + // TurboQuant K: move Q into the Hadamard domain once, so the KQ dot can run + // against raw codebook values. = , and H is symmetric. + // + // Done here rather than during the load above because the transform needs a + // whole 128-element chunk, while the load distributes elements across threads. + // KV_tmp is still unused at this point and serves as the float scratch — Q_tmp + // may be half, and rotating in half would cost accuracy the accumulator cannot + // recover. + if constexpr (type_K == GGML_TYPE_TURBO3_0 || type_K == GGML_TYPE_TURBO4_0) { + static_assert(DKQ % TURBO_HEAD_DIM_GPU == 0, "turbo needs DKQ to be a multiple of 128"); + constexpr int nthreads_blk = warp_size*nwarps; + constexpr int chunks_per_col = DKQ / TURBO_HEAD_DIM_GPU; + + // One chunk at a time: the scratch then needs 128 floats regardless of + // ncols. Staging all of Q at once would need ncols*DKQ floats, which + // outgrows KV_tmp as ncols rises (KV_tmp scales with nbatch_fa*nbatch_K, + // and those shrink as ncols grows). + float * Qs = (float *) KV_tmp; + const int tid_blk = warp_size*threadIdx.y + threadIdx.x; + +#pragma unroll + for (int jc = 0; jc < ncols; ++jc) { +#pragma unroll + for (int ch = 0; ch < chunks_per_col; ++ch) { + const int base = ch*TURBO_HEAD_DIM_GPU; + +#pragma unroll + for (int e0 = 0; e0 < TURBO_HEAD_DIM_GPU; e0 += nthreads_blk) { + const int e = e0 + tid_blk; + if (nthreads_blk > TURBO_HEAD_DIM_GPU && e >= TURBO_HEAD_DIM_GPU) { + continue; + } +#ifdef FAST_FP16_AVAILABLE + const half2 h = Q_tmp[jc*(DKQ/2) + (base + e)/2]; + Qs[e] = ((base + e) % 2 == 0) ? __half2float(__low2half(h)) : __half2float(__high2half(h)); +#else + Qs[e] = Q_tmp[jc*DKQ + base + e]; +#endif // FAST_FP16_AVAILABLE + } + __syncthreads(); + + turbo_fwht_smem_128_n(Qs, tid_blk); + +#pragma unroll + for (int e0 = 0; e0 < TURBO_HEAD_DIM_GPU; e0 += nthreads_blk) { + const int e = e0 + tid_blk; + if (nthreads_blk > TURBO_HEAD_DIM_GPU && e >= TURBO_HEAD_DIM_GPU) { + continue; + } +#ifdef FAST_FP16_AVAILABLE + if ((base + e) % 2 == 0) { + Q_tmp[jc*(DKQ/2) + (base + e)/2] = make_half2(Qs[e], Qs[e + 1]); + } +#else + Q_tmp[jc*DKQ + base + e] = Qs[e]; +#endif // FAST_FP16_AVAILABLE + } + __syncthreads(); + } + } + } + // Main loop over KV cache: const int k_VKQ_max = KV_max ? KV_max[sequence*gridDim.x + blockIdx.x] : ne11; if (ncols2 == 1) { @@ -899,14 +1104,14 @@ static __global__ void flash_attn_tile( int k_VKQ_0 = blockIdx.y*nbatch_fa; while (k_VKQ_0 < k_VKQ_max - nbatch_fa) { constexpr bool oob_check = false; - flash_attn_tile_iter + flash_attn_tile_iter (Q_tmp, K_h2, V_h2, maskh, ne01, logit_softcap, slope, KQ, KV_tmp, stride_K2, stride_V2, stride_mask, KQ_max, KQ_sum, VKQ, k_VKQ_0, k_VKQ_max, col_Q_0); k_VKQ_0 += gridDim.y*nbatch_fa; } if (k_VKQ_0 < k_VKQ_max) { constexpr bool oob_check = true; - flash_attn_tile_iter + flash_attn_tile_iter (Q_tmp, K_h2, V_h2, maskh, ne01, logit_softcap, slope, KQ, KV_tmp, stride_K2, stride_V2, stride_mask, KQ_max, KQ_sum, VKQ, k_VKQ_0, k_VKQ_max, col_Q_0); } @@ -914,7 +1119,7 @@ static __global__ void flash_attn_tile( // Branch without out-of-bounds checks. for (int k_VKQ_0 = blockIdx.y*nbatch_fa; k_VKQ_0 < k_VKQ_max; k_VKQ_0 += gridDim.y*nbatch_fa) { constexpr bool oob_check = false; - flash_attn_tile_iter + flash_attn_tile_iter (Q_tmp, K_h2, V_h2, maskh, ne01, logit_softcap, slope, KQ, KV_tmp, stride_K2, stride_V2, stride_mask, KQ_max, KQ_sum, VKQ, k_VKQ_0, k_VKQ_max, col_Q_0); } @@ -1021,6 +1226,60 @@ static __global__ void flash_attn_tile( } } + // TurboQuant: the accumulator is still in the Hadamard domain. One inverse + // FWHT per 128-element chunk brings it back — once per launch, not once per + // KV token. Valid because H is linear and every scaling applied along the way + // (softmax rescaling, the 1/KQ_sum below) is a scalar: H(a*v) = a*H(v). + if constexpr (type_V == GGML_TYPE_TURBO3_0 || type_V == GGML_TYPE_TURBO4_0) { + static_assert(DV % TURBO_HEAD_DIM_GPU == 0, "turbo needs DV to be a multiple of 128"); + static_assert(DVp == DV, "turbo assumes an unpadded V accumulator"); + + __syncthreads(); // KV_tmp is free from here on + float * fwht_buf = ((float *) KV_tmp) + threadIdx.y*DV; + +#pragma unroll + for (int jc0 = 0; jc0 < cpw; ++jc0) { +#pragma unroll + for (int i0 = 0; i0 < DV/2; i0 += warp_size) { + const int idx = i0 + threadIdx.x; +#ifdef FAST_FP16_AVAILABLE + const float2 t = __half22float2(VKQ[jc0*((DV/2)/warp_size) + i0/warp_size]); +#else + const float2 t = VKQ[jc0*((DV/2)/warp_size) + i0/warp_size]; +#endif // FAST_FP16_AVAILABLE + fwht_buf[2*idx] = t.x; + fwht_buf[2*idx + 1] = t.y; + } +#ifdef GGML_USE_HIP + __builtin_amdgcn_wave_barrier(); +#else + __syncwarp(); +#endif + +#pragma unroll + for (int c = 0; c < DV/TURBO_HEAD_DIM_GPU; ++c) { + turbo_fwht_smem_128_warp(fwht_buf + c*TURBO_HEAD_DIM_GPU, threadIdx.x); + } + +#pragma unroll + for (int i0 = 0; i0 < DV/2; i0 += warp_size) { + const int idx = i0 + threadIdx.x; +#ifdef FAST_FP16_AVAILABLE + VKQ[jc0*((DV/2)/warp_size) + i0/warp_size] = make_half2(fwht_buf[2*idx], fwht_buf[2*idx + 1]); +#else + VKQ[jc0*((DV/2)/warp_size) + i0/warp_size].x = fwht_buf[2*idx]; + VKQ[jc0*((DV/2)/warp_size) + i0/warp_size].y = fwht_buf[2*idx + 1]; +#endif // FAST_FP16_AVAILABLE + } +#ifdef GGML_USE_HIP + __builtin_amdgcn_wave_barrier(); +#else + __syncwarp(); +#endif + } + __syncthreads(); + } + // Write back results: #pragma unroll for (int jc0 = 0; jc0 < cpw; ++jc0) { @@ -1087,7 +1346,7 @@ static __global__ void flash_attn_tile( #endif // FLASH_ATTN_AVAILABLE } -template +template static void launch_fattn_tile_switch_ncols1(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { const ggml_tensor * Q = dst->src[0]; @@ -1097,15 +1356,19 @@ static void launch_fattn_tile_switch_ncols1(ggml_backend_cuda_context & ctx, ggm constexpr size_t nbytes_shared = 0; + // Turbo V is decoded inside the kernel, so launch_fattn must not expand it. + constexpr bool need_f16_V = !(type_V == GGML_TYPE_TURBO3_0 || type_V == GGML_TYPE_TURBO4_0); + constexpr bool need_f16_K = !(type_K == GGML_TYPE_TURBO3_0 || type_K == GGML_TYPE_TURBO4_0); + #ifdef GGML_USE_HIP if constexpr (DV <= 128) { if (Q->ne[1] > 32/ncols2) { constexpr int cols_per_block = 64; const int nwarps = ggml_cuda_fattn_tile_get_nthreads (DKQ, DV, cols_per_block, cc) / warp_size; const int nbatch_fa = ggml_cuda_fattn_tile_get_nbatch_fa(DKQ, DV, cols_per_block, cc); - fattn_kernel_t fattn_kernel = flash_attn_tile; + fattn_kernel_t fattn_kernel = flash_attn_tile; launch_fattn - (ctx, dst, fattn_kernel, nwarps, nbytes_shared, nbatch_fa, true, true, false, warp_size); + (ctx, dst, fattn_kernel, nwarps, nbytes_shared, nbatch_fa, need_f16_K, need_f16_V, false, warp_size); return; } } @@ -1119,9 +1382,9 @@ static void launch_fattn_tile_switch_ncols1(ggml_backend_cuda_context & ctx, ggm constexpr int cols_per_block = 32; const int nwarps = ggml_cuda_fattn_tile_get_nthreads (DKQ, DV, cols_per_block, cc) / warp_size; const int nbatch_fa = ggml_cuda_fattn_tile_get_nbatch_fa(DKQ, DV, cols_per_block, cc); - fattn_kernel_t fattn_kernel = flash_attn_tile; + fattn_kernel_t fattn_kernel = flash_attn_tile; launch_fattn - (ctx, dst, fattn_kernel, nwarps, nbytes_shared, nbatch_fa, true, true, false, warp_size); + (ctx, dst, fattn_kernel, nwarps, nbytes_shared, nbatch_fa, need_f16_K, need_f16_V, false, warp_size); return; } } @@ -1130,9 +1393,9 @@ static void launch_fattn_tile_switch_ncols1(ggml_backend_cuda_context & ctx, ggm constexpr int cols_per_block = 16; const int nwarps = ggml_cuda_fattn_tile_get_nthreads (DKQ, DV, cols_per_block, cc) / warp_size; const int nbatch_fa = ggml_cuda_fattn_tile_get_nbatch_fa(DKQ, DV, cols_per_block, cc); - fattn_kernel_t fattn_kernel = flash_attn_tile; + fattn_kernel_t fattn_kernel = flash_attn_tile; launch_fattn - (ctx, dst, fattn_kernel, nwarps, nbytes_shared, nbatch_fa, true, true, false, warp_size); + (ctx, dst, fattn_kernel, nwarps, nbytes_shared, nbatch_fa, need_f16_K, need_f16_V, false, warp_size); return; } @@ -1141,9 +1404,9 @@ static void launch_fattn_tile_switch_ncols1(ggml_backend_cuda_context & ctx, ggm constexpr int cols_per_block = 8; const int nwarps = ggml_cuda_fattn_tile_get_nthreads (DKQ, DV, cols_per_block, cc) / warp_size; const int nbatch_fa = ggml_cuda_fattn_tile_get_nbatch_fa(DKQ, DV, cols_per_block, cc); - fattn_kernel_t fattn_kernel = flash_attn_tile; + fattn_kernel_t fattn_kernel = flash_attn_tile; launch_fattn - (ctx, dst, fattn_kernel, nwarps, nbytes_shared, nbatch_fa, true, true, false, warp_size); + (ctx, dst, fattn_kernel, nwarps, nbytes_shared, nbatch_fa, need_f16_K, need_f16_V, false, warp_size); return; } } @@ -1153,9 +1416,9 @@ static void launch_fattn_tile_switch_ncols1(ggml_backend_cuda_context & ctx, ggm constexpr int cols_per_block = 4; const int nwarps = ggml_cuda_fattn_tile_get_nthreads (DKQ, DV, cols_per_block, cc) / warp_size; const int nbatch_fa = ggml_cuda_fattn_tile_get_nbatch_fa(DKQ, DV, cols_per_block, cc); - fattn_kernel_t fattn_kernel = flash_attn_tile; + fattn_kernel_t fattn_kernel = flash_attn_tile; launch_fattn - (ctx, dst, fattn_kernel, nwarps, nbytes_shared, nbatch_fa, true, true, false, warp_size); + (ctx, dst, fattn_kernel, nwarps, nbytes_shared, nbatch_fa, need_f16_K, need_f16_V, false, warp_size); return; } } @@ -1164,16 +1427,16 @@ static void launch_fattn_tile_switch_ncols1(ggml_backend_cuda_context & ctx, ggm constexpr int cols_per_block = 2; const int nwarps = ggml_cuda_fattn_tile_get_nthreads (DKQ, DV, cols_per_block, cc) / warp_size; const int nbatch_fa = ggml_cuda_fattn_tile_get_nbatch_fa(DKQ, DV, cols_per_block, cc); - fattn_kernel_t fattn_kernel = flash_attn_tile; + fattn_kernel_t fattn_kernel = flash_attn_tile; launch_fattn - (ctx, dst, fattn_kernel, nwarps, nbytes_shared, nbatch_fa, true, true, false, warp_size); + (ctx, dst, fattn_kernel, nwarps, nbytes_shared, nbatch_fa, need_f16_K, need_f16_V, false, warp_size); return; } GGML_ABORT("fatal error"); } -template +template static void launch_fattn_tile_switch_ncols2(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { const ggml_tensor * KQV = dst; const ggml_tensor * Q = dst->src[0]; @@ -1194,38 +1457,38 @@ static void launch_fattn_tile_switch_ncols2(ggml_backend_cuda_context & ctx, ggm if constexpr (DV == 512) { if (use_gqa_opt && gqa_ratio % 16 == 0) { - launch_fattn_tile_switch_ncols1(ctx, dst); + launch_fattn_tile_switch_ncols1(ctx, dst); return; } if (use_gqa_opt && gqa_ratio % 4 == 0) { - launch_fattn_tile_switch_ncols1(ctx, dst); + launch_fattn_tile_switch_ncols1(ctx, dst); return; } } if constexpr (DV <= 256) { if (use_gqa_opt && gqa_ratio % 8 == 0) { - launch_fattn_tile_switch_ncols1(ctx, dst); + launch_fattn_tile_switch_ncols1(ctx, dst); return; } if (use_gqa_opt && gqa_ratio % 4 == 0) { - launch_fattn_tile_switch_ncols1(ctx, dst); + launch_fattn_tile_switch_ncols1(ctx, dst); return; } if (use_gqa_opt && gqa_ratio % 2 == 0) { - launch_fattn_tile_switch_ncols1(ctx, dst); + launch_fattn_tile_switch_ncols1(ctx, dst); return; } - launch_fattn_tile_switch_ncols1(ctx, dst); + launch_fattn_tile_switch_ncols1(ctx, dst); return; } GGML_ABORT("fatal error"); } -template +template void ggml_cuda_flash_attn_ext_tile_case(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { const ggml_tensor * KQV = dst; @@ -1234,10 +1497,10 @@ void ggml_cuda_flash_attn_ext_tile_case(ggml_backend_cuda_context & ctx, ggml_te if (logit_softcap == 0.0f) { constexpr bool use_logit_softcap = false; - launch_fattn_tile_switch_ncols2(ctx, dst); + launch_fattn_tile_switch_ncols2(ctx, dst); } else { constexpr bool use_logit_softcap = true; - launch_fattn_tile_switch_ncols2(ctx, dst); + launch_fattn_tile_switch_ncols2(ctx, dst); } } @@ -1247,6 +1510,10 @@ void ggml_cuda_flash_attn_ext_tile(ggml_backend_cuda_context & ctx, ggml_tensor template void ggml_cuda_flash_attn_ext_tile_case \ (ggml_backend_cuda_context & ctx, ggml_tensor * dst) \ +#define DECL_FATTN_TILE_CASE_TURBO(DKQ, DV, TYPE_KV) \ + template void ggml_cuda_flash_attn_ext_tile_case \ + (ggml_backend_cuda_context & ctx, ggml_tensor * dst) \ + extern DECL_FATTN_TILE_CASE( 40, 40); extern DECL_FATTN_TILE_CASE( 64, 64); extern DECL_FATTN_TILE_CASE( 72, 72); @@ -1256,3 +1523,9 @@ extern DECL_FATTN_TILE_CASE(112, 112); extern DECL_FATTN_TILE_CASE(128, 128); extern DECL_FATTN_TILE_CASE(256, 256); extern DECL_FATTN_TILE_CASE(576, 512); + +// TurboQuant: decoded inside the kernel, only the head sizes turbo supports. +extern DECL_FATTN_TILE_CASE_TURBO(128, 128, GGML_TYPE_TURBO3_0); +extern DECL_FATTN_TILE_CASE_TURBO(128, 128, GGML_TYPE_TURBO4_0); +extern DECL_FATTN_TILE_CASE_TURBO(256, 256, GGML_TYPE_TURBO3_0); +extern DECL_FATTN_TILE_CASE_TURBO(256, 256, GGML_TYPE_TURBO4_0); diff --git a/ggml/src/ggml-cuda/fattn-vec-turbo.cuh b/ggml/src/ggml-cuda/fattn-vec-turbo.cuh new file mode 100644 index 000000000..ae6fa5f7a --- /dev/null +++ b/ggml/src/ggml-cuda/fattn-vec-turbo.cuh @@ -0,0 +1,465 @@ +#include "common.cuh" +#include "fattn-common.cuh" +#include "turbo.cuh" + +// ============================================================ +// Fused flash-attention vector kernel for TurboQuant KV +// ============================================================ +// +// The generic vector kernel cannot handle turbo: it decodes KV element-wise, +// while a TurboQuant value needs the inverse FWHT over its whole 128-element +// chunk. Decoding every KV token would cost O(n_kv) transforms and be slower +// than the bulk conversion this is meant to replace. +// +// Two properties of the Hadamard matrix H (symmetric, orthogonal) avoid that. +// With s = 1/sqrt(128), a stored chunk decodes as x = norm * s * H c: +// +// K side: = norm_j * s * = norm_j * s * +// Transform Q once at kernel start, then dot against raw codebook +// values — no transform in the inner loop. +// +// V side: sum_j p_j x_j = H ( sum_j p_j norm_j s c_j ) +// Accumulate in the Hadamard domain, transform once at the end. +// Valid because the online-softmax rescaling is a scalar factor and +// H is linear: H(a*v) = a*H(v). +// +// Two transforms per launch instead of two per KV token. +// +// Layout: one warp owns a full 128-element chunk, lane l holding elements +// [4l, 4l+4). That is exactly the register layout turbo_fwht_warp_128 expects, +// so the transforms need no shared memory. +// +// The accumulator is float, not half: codebook values are small (|c| <= 0.24) +// and the final transform amplifies rounding by sqrt(128). + +#ifdef __clang__ +#pragma clang diagnostic push +#pragma clang diagnostic ignored "-Wpass-failed" +#endif // __clang__ + +static constexpr __device__ int ggml_cuda_fattn_vec_turbo_nthreads() { + return 128; +} + +// Ordering fence within a warp. The lanes run in lockstep, so no data movement is +// needed — this only stops the compiler from sinking loads across the point. +static __device__ __forceinline__ void turbo_warp_fence() { +#ifdef GGML_USE_HIP + __builtin_amdgcn_wave_barrier(); +#else + __syncwarp(); +#endif +} + +template +__launch_bounds__(ggml_cuda_fattn_vec_turbo_nthreads(), 1) +static __global__ void flash_attn_ext_vec_turbo( + const char * __restrict__ Q, + const char * __restrict__ K, + const char * __restrict__ V, + const char * __restrict__ mask, + const char * __restrict__ sinks, + const int * __restrict__ KV_max, + float * __restrict__ dst, + float2 * __restrict__ dst_meta, + const float scale, + const float max_bias, + const float m0, + const float m1, + const uint32_t n_head_log2, + const float logit_softcap, + const int32_t ne00, const uint3 ne01, const int32_t ne02, const int32_t ne03, + const int32_t nb01, const int32_t nb02, const int32_t nb03, + const int32_t ne10, const int32_t ne11, const int32_t ne12, const int32_t ne13, + const int32_t nb11, const int32_t nb12, const int64_t nb13, + const int32_t nb21, const int32_t nb22, const int64_t nb23, + const int32_t ne31, const int32_t ne32, const int32_t ne33, + const int32_t nb31, const int32_t nb32, const int64_t nb33) { +#ifdef FLASH_ATTN_AVAILABLE + using block_t = typename std::conditional::type; + + static_assert(D % TURBO_HEAD_DIM_GPU == 0, "TurboQuant needs head_dim to be a multiple of 128"); + constexpr int nchunks = D / TURBO_HEAD_DIM_GPU; // 1 for D=128, 2 for D=256 + constexpr int nthreads = ggml_cuda_fattn_vec_turbo_nthreads(); + constexpr int nwarps = nthreads / WARP_SIZE; + + const int lane = threadIdx.x; // 0..31, owns elements [4*lane, 4*lane+4) + const int tid = WARP_SIZE*threadIdx.y + threadIdx.x; + + const int ic0 = blockIdx.x * ncols; + + const int sequence = blockIdx.z / ne02; + const int head = blockIdx.z - sequence*ne02; + const int gqa_ratio = ne02 / ne12; + + Q += nb03*sequence + nb02*head + nb01*ic0; + K += nb13*sequence + nb12*(head / gqa_ratio); + V += nb23*sequence + nb22*(head / gqa_ratio); + + const half * maskh = (const half *) (mask + nb33*(sequence % ne33) + nb31*ic0); + + const float slope = get_alibi_slope(max_bias, head, n_head_log2, m0, m1); + + // Shared memory: KQ logits during the loop, warp combination at the end. + constexpr int ne_KQ = ncols*nthreads; + constexpr int ne_combine = nwarps*D; + __shared__ float KQ[ne_KQ > ne_combine ? ne_KQ : ne_combine]; + + // --- Q into the Hadamard domain, once per launch --- + // + // Two different lane layouts are needed here. turbo_fwht_warp_128 requires 4 + // consecutive elements per lane (that is what makes the butterfly stages + // register-local). The KQ dot wants the opposite: as few lanes per token as + // possible, because each token costs one warp reduction. With nthreads_KQ + // lanes per token, a reduction is log2(nthreads_KQ) stages instead of 5, and + // nthreads/nthreads_KQ tokens are processed at once. + // + // So: transform in the 4-per-lane layout, then redistribute through shared + // memory into E = D/nthreads_KQ consecutive elements per lane. At E=16 the + // 3-bit indices of a turbo3 block also land on byte boundaries (16*3 = 48 + // bit = 6 byte), which keeps the unpacking cheap. + constexpr int nthreads_KQ = 8; + constexpr int E = D / nthreads_KQ; // elements per lane in the dot + constexpr int toks_per_it = WARP_SIZE / nthreads_KQ; // tokens handled simultaneously + + const int lane_KQ = lane % nthreads_KQ; // which slice of the head this lane dots + const int tok_KQ = lane / nthreads_KQ; // which of the concurrent tokens + + float Q_dot[ncols][E]; + { + float * Qs = KQ + threadIdx.y*D; // per-warp scratch, no cross-warp barrier +#pragma unroll + for (int j = 0; j < ncols; ++j) { + const float * Q_j = (const float *) (Q + j*nb01); +#pragma unroll + for (int c = 0; c < nchunks; ++c) { + float qh[4]; +#pragma unroll + for (int r = 0; r < 4; ++r) { + const int i = c*TURBO_HEAD_DIM_GPU + 4*lane + r; + qh[r] = (ncols == 1 || ic0 + j < int(ne01.z)) ? Q_j[i] : 0.0f; + } + turbo_fwht_warp_128(qh, lane); + // Fold in the attention scale and the decode-side 1/sqrt(128) here + // so the inner loop only multiplies by the per-chunk norm. +#pragma unroll + for (int r = 0; r < 4; ++r) { + Qs[c*TURBO_HEAD_DIM_GPU + 4*lane + r] = qh[r] * scale * TURBO_FWHT_SCALE_GPU; + } + } + turbo_warp_fence(); +#pragma unroll + for (int e = 0; e < E; ++e) { + Q_dot[j][e] = Qs[lane_KQ*E + e]; + } + turbo_warp_fence(); + } + } + + // VKQ accumulator, kept in the Hadamard domain until the very end. + float VKQ[ncols][nchunks][4] = {}; + float KQ_max[ncols]; + float KQ_sum[ncols]; +#pragma unroll + for (int j = 0; j < ncols; ++j) { + KQ_max[j] = -FLT_MAX/2.0f; + KQ_sum[j] = 0.0f; + } + + const int k_VKQ_max = KV_max ? KV_max[sequence*gridDim.x + blockIdx.x] : ne11; + K += blockIdx.y*nthreads * nb11; + V += blockIdx.y*nthreads * nb21; + maskh += blockIdx.y*nthreads; + + for (int k_VKQ_0 = blockIdx.y*nthreads; k_VKQ_0 < k_VKQ_max; k_VKQ_0 += gridDim.y*nthreads, + K += gridDim.y*nthreads*nb11, V += gridDim.y*nthreads*nb21, maskh += gridDim.y*nthreads) { + + // --- logits: warp y handles KV tokens [y*WARP_SIZE, (y+1)*WARP_SIZE) --- + // + // toks_per_it tokens at a time, nthreads_KQ lanes each. Each lane covers + // E consecutive elements, which is one chunk's worth or less, so the + // per-chunk norm is a single load per lane per token. + float KQ_max_new[ncols]; +#pragma unroll + for (int j = 0; j < ncols; ++j) { + KQ_max_new[j] = KQ_max[j]; + } + +#pragma unroll + for (int i_KQ_0 = 0; i_KQ_0 < WARP_SIZE; i_KQ_0 += toks_per_it) { + const int i_KQ = threadIdx.y*WARP_SIZE + i_KQ_0 + tok_KQ; + + const block_t * Kb = (const block_t *) (K + i_KQ*nb11); + + // E is a divisor of the 128-element chunk, so a lane's slice never + // straddles a chunk boundary: one norm, fetched once. + const int elem0 = lane_KQ*E; + const int chunk = elem0 / TURBO_HEAD_DIM_GPU; + const int64_t ib0 = (int64_t)chunk * TURBO_BLOCKS_PER_CHUNK_GPU; + const float norm = __half2float(Kb[ib0].d); + + float sum[ncols]; +#pragma unroll + for (int j = 0; j < ncols; ++j) { + sum[j] = 0.0f; + } +#pragma unroll + for (int e = 0; e < E; ++e) { + const int elem = elem0 + e; + const int eic = elem % TURBO_HEAD_DIM_GPU; // element within the chunk + const int blk = eic / 32; + const int eib = eic % 32; + float cb; + if constexpr (type_KV == GGML_TYPE_TURBO3_0) { + cb = dc_codebook_3bit[turbo3_unpack_index(Kb + ib0 + blk, eib)]; + } else { + cb = dc_codebook_4bit[turbo4_unpack_index(Kb + ib0 + blk, eib)]; + } +#pragma unroll + for (int j = 0; j < ncols; ++j) { + sum[j] += cb * Q_dot[j][e]; + } + } + +#pragma unroll + for (int j = 0; j < ncols; ++j) { + sum[j] = warp_reduce_sum(sum[j]) * norm; + + if (use_logit_softcap) { + sum[j] = logit_softcap * tanhf(sum[j]); + } + // maskh already advances with the loop, so index by i_KQ only, and + // the row stride is ne11 (the KV length), not ne31. + if (mask && (ncols == 1 || ic0 + j < int(ne01.z))) { + sum[j] += slope * __half2float(maskh[j*ne11 + i_KQ]); + } + + // The offset keeps the running maximum finite when a whole tile is + // masked out; without it expf(-inf - -inf) yields NaN. + KQ_max_new[j] = fmaxf(KQ_max_new[j], sum[j] + FATTN_KQ_MAX_OFFSET); + + if (lane_KQ == 0) { + KQ[j*nthreads + threadIdx.y*WARP_SIZE + i_KQ_0 + tok_KQ] = sum[j]; + } + } + } + + // The logits were written by lane_KQ == 0 of each group and are read back + // by every lane below; both happen inside one warp. + turbo_warp_fence(); + + // --- online softmax rescale --- +#pragma unroll + for (int j = 0; j < ncols; ++j) { + KQ_max_new[j] = warp_reduce_max(KQ_max_new[j]); + const float KQ_max_scale = expf(KQ_max[j] - KQ_max_new[j]); + KQ_max[j] = KQ_max_new[j]; + + // Every lane owns exactly one token's logit at this point. + const float p = expf(KQ[j*nthreads + tid] - KQ_max[j]); + KQ_sum[j] = KQ_sum[j]*KQ_max_scale + p; + KQ[j*nthreads + tid] = p; + + // Rescaling in the Hadamard domain is exact: H(a*v) = a*H(v). +#pragma unroll + for (int c = 0; c < nchunks; ++c) { +#pragma unroll + for (int r = 0; r < 4; ++r) { + VKQ[j][c][r] *= KQ_max_scale; + } + } + } + + __syncthreads(); + + // --- V accumulation, no transform in the loop --- +#pragma unroll + for (int k0 = 0; k0 < WARP_SIZE; ++k0) { + const int k = threadIdx.y*WARP_SIZE + k0; + + const block_t * Vb = (const block_t *) (V + k*nb21); + +#pragma unroll + for (int c = 0; c < nchunks; ++c) { + const int64_t ib = (int64_t)c * TURBO_BLOCKS_PER_CHUNK_GPU; + float cb[4]; +#pragma unroll + for (int r = 0; r < 4; ++r) { + const int elem = 4*lane + r; + const int blk = elem / 32; + const int eib = elem % 32; + if constexpr (type_KV == GGML_TYPE_TURBO3_0) { + cb[r] = dc_codebook_3bit[turbo3_unpack_index(Vb + ib + blk, eib)]; + } else { + cb[r] = dc_codebook_4bit[turbo4_unpack_index(Vb + ib + blk, eib)]; + } + } + const float vnorm = TURBO_FWHT_SCALE_GPU * __half2float(Vb[ib].d); +#pragma unroll + for (int j = 0; j < ncols; ++j) { + const float p = KQ[j*nthreads + threadIdx.y*WARP_SIZE + k0] * vnorm; +#pragma unroll + for (int r = 0; r < 4; ++r) { + VKQ[j][c][r] += p * cb[r]; + } + } + } + } + + __syncthreads(); + } + + // --- attention sinks --- + if (sinks && blockIdx.y == 0) { + const float sink = ((const float *) sinks)[head]; +#pragma unroll + for (int j = 0; j < ncols; ++j) { + const float kqmax_new = fmaxf(sink, KQ_max[j]); + const float KQ_max_scale = expf(KQ_max[j] - kqmax_new); + KQ_max[j] = kqmax_new; + + KQ_sum[j] = KQ_sum[j]*KQ_max_scale + (tid == 0 ? expf(sink - KQ_max[j]) : 0.0f); +#pragma unroll + for (int c = 0; c < nchunks; ++c) { +#pragma unroll + for (int r = 0; r < 4; ++r) { + VKQ[j][c][r] *= KQ_max_scale; + } + } + } + } + + // --- leave the Hadamard domain: one transform per chunk, per warp --- +#pragma unroll + for (int j = 0; j < ncols; ++j) { +#pragma unroll + for (int c = 0; c < nchunks; ++c) { + turbo_fwht_warp_128(VKQ[j][c], lane); + } + } + + __shared__ float KQ_max_shared[ncols][WARP_SIZE]; + __shared__ float KQ_sum_shared[ncols][WARP_SIZE]; +#pragma unroll + for (int j = 0; j < ncols; ++j) { + if (threadIdx.y == 0) { + KQ_max_shared[j][threadIdx.x] = -FLT_MAX/2.0f; + KQ_sum_shared[j][threadIdx.x] = 0.0f; + } + } + __syncthreads(); +#pragma unroll + for (int j = 0; j < ncols; ++j) { + if (threadIdx.x == 0) { + KQ_max_shared[j][threadIdx.y] = KQ_max[j]; + } + } + __syncthreads(); + +#pragma unroll + for (int j_VKQ = 0; j_VKQ < ncols; ++j_VKQ) { + if (ncols > 1 && ic0 + j_VKQ >= int(ne01.z)) { + break; + } + + float kqmax_new = KQ_max_shared[j_VKQ][threadIdx.x]; + kqmax_new = warp_reduce_max(kqmax_new); + const float kqmax_scale = expf(KQ_max[j_VKQ] - kqmax_new); + KQ_max[j_VKQ] = kqmax_new; + + // Stage this warp's contribution for the cross-warp sum. + float * VKQ_tmp = KQ + threadIdx.y*D; +#pragma unroll + for (int c = 0; c < nchunks; ++c) { +#pragma unroll + for (int r = 0; r < 4; ++r) { + VKQ_tmp[c*TURBO_HEAD_DIM_GPU + 4*lane + r] = VKQ[j_VKQ][c][r] * kqmax_scale; + } + } + + KQ_sum[j_VKQ] *= kqmax_scale; + KQ_sum[j_VKQ] = warp_reduce_sum(KQ_sum[j_VKQ]); + if (threadIdx.x == 0) { + KQ_sum_shared[j_VKQ][threadIdx.y] = KQ_sum[j_VKQ]; + } + + __syncthreads(); + + if (nthreads <= D || tid < D) { + KQ_sum[j_VKQ] = KQ_sum_shared[j_VKQ][threadIdx.x]; + KQ_sum[j_VKQ] = warp_reduce_sum(KQ_sum[j_VKQ]); + +#pragma unroll + for (int i0 = 0; i0 < D; i0 += nthreads) { + float dst_val = 0.0f; +#pragma unroll + for (int w = 0; w < nwarps; ++w) { + dst_val += KQ[w*D + i0 + tid]; + } + if (gridDim.y == 1) { + dst_val /= KQ_sum[j_VKQ]; + } + dst[(((sequence*int(ne01.z) + ic0 + j_VKQ)*ne02 + head)*gridDim.y + blockIdx.y)*D + i0 + tid] = dst_val; + } + } + + if (j_VKQ < ncols-1) { + __syncthreads(); + } + } + + if (gridDim.y != 1 && tid < ncols && (ncols == 1 || ic0 + tid < int(ne01.z))) { + dst_meta[((sequence*int(ne01.z) + ic0 + tid)*ne02 + head)*gridDim.y + blockIdx.y] = + make_float2(KQ_max[tid], KQ_sum[tid]); + } +#else + GGML_UNUSED_VARS(Q, K, V, mask, sinks, KV_max, dst, dst_meta, scale, + max_bias, m0, m1, n_head_log2, logit_softcap, + ne00, ne01, ne02, ne03, nb01, nb02, nb03, + ne10, ne11, ne12, ne13, nb11, nb12, nb13, + nb21, nb22, nb23, ne31, ne32, ne33, nb31, nb32, nb33); + NO_DEVICE_CODE; +#endif // FLASH_ATTN_AVAILABLE +} + +#ifdef __clang__ +#pragma clang diagnostic pop +#endif // __clang__ + +template +void ggml_cuda_flash_attn_ext_vec_turbo_case_impl(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { + constexpr int nthreads = 128; + constexpr int nwarps = nthreads / WARP_SIZE; + + fattn_kernel_t fattn_kernel = flash_attn_ext_vec_turbo; + + // No conversion: the kernel reads turbo blocks directly. That is the point. + launch_fattn(ctx, dst, fattn_kernel, nwarps, /*nbytes_shared*/ 0, + D, /*need_f16_K*/ false, /*need_f16_V*/ false, /*stream_k*/ false); +} + +template +void ggml_cuda_flash_attn_ext_vec_turbo_case(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { + const ggml_tensor * KQV = dst; + const ggml_tensor * Q = dst->src[0]; + + float logit_softcap; + memcpy(&logit_softcap, (const float *) KQV->op_params + 2, sizeof(float)); + + const int cols_per_block = Q->ne[1] > 1 ? 2 : 1; + + if (logit_softcap == 0.0f) { + if (cols_per_block == 1) { + ggml_cuda_flash_attn_ext_vec_turbo_case_impl(ctx, dst); + } else { + ggml_cuda_flash_attn_ext_vec_turbo_case_impl(ctx, dst); + } + } else { + if (cols_per_block == 1) { + ggml_cuda_flash_attn_ext_vec_turbo_case_impl(ctx, dst); + } else { + ggml_cuda_flash_attn_ext_vec_turbo_case_impl(ctx, dst); + } + } +} diff --git a/ggml/src/ggml-cuda/fattn-vec.cuh b/ggml/src/ggml-cuda/fattn-vec.cuh index 13981deb2..5b0246993 100644 --- a/ggml/src/ggml-cuda/fattn-vec.cuh +++ b/ggml/src/ggml-cuda/fattn-vec.cuh @@ -602,11 +602,8 @@ EXTERN_DECL_FATTN_VEC_CASES(256, GGML_TYPE_Q5_1) EXTERN_DECL_FATTN_VEC_CASES(256, GGML_TYPE_Q8_0) EXTERN_DECL_FATTN_VEC_CASES(256, GGML_TYPE_BF16) -// TurboQuant extern declarations (homogeneous K/V only) -extern DECL_FATTN_VEC_CASE( 64, GGML_TYPE_TURBO3_0, GGML_TYPE_TURBO3_0); -extern DECL_FATTN_VEC_CASE(128, GGML_TYPE_TURBO3_0, GGML_TYPE_TURBO3_0); -extern DECL_FATTN_VEC_CASE(256, GGML_TYPE_TURBO3_0, GGML_TYPE_TURBO3_0); - -extern DECL_FATTN_VEC_CASE( 64, GGML_TYPE_TURBO4_0, GGML_TYPE_TURBO4_0); -extern DECL_FATTN_VEC_CASE(128, GGML_TYPE_TURBO4_0, GGML_TYPE_TURBO4_0); -extern DECL_FATTN_VEC_CASE(256, GGML_TYPE_TURBO4_0, GGML_TYPE_TURBO4_0); +// No TurboQuant instances: the vector kernel does not convert its inputs, and +// TurboQuant cannot be decoded element-wise (a value needs the inverse FWHT over +// its whole 128-element chunk). Turbo KV is served by the tile/MMA kernels, which +// request f16 and are converted in launch_fattn. Leaving the instances out means a +// dispatch mistake is a link error rather than silently wrong numbers. diff --git a/ggml/src/ggml-cuda/fattn.cu b/ggml/src/ggml-cuda/fattn.cu index e4840d147..db1dac5cf 100644 --- a/ggml/src/ggml-cuda/fattn.cu +++ b/ggml/src/ggml-cuda/fattn.cu @@ -3,6 +3,7 @@ #include "fattn-mma-f16.cuh" #include "fattn-tile.cuh" #include "fattn-vec.cuh" +#include "fattn-vec-turbo.cuh" #include "fattn-wmma-f16.cuh" #include "fattn.cuh" #include "convert.cuh" @@ -396,8 +397,15 @@ static best_fattn_kernel ggml_cuda_get_best_fattn_kernel(const int device, const return BEST_FATTN_KERNEL_NONE; } + // The vector kernel is the only one that does not request f16 inputs + // (fattn-vec.cuh: need_f16_K = type_K == GGML_TYPE_F16), so it would be handed + // raw turbo blocks and reinterpret them as f16. Tile/MMA/WMMA always pass + // need_f16_K/V = true and are converted by launch_fattn, so they are safe. + // Keep turbo off the vector path until a fused turbo kernel exists. + const bool kv_is_turbo = ggml_type_is_turbo(K->type) || ggml_type_is_turbo(V->type); + // For small batch sizes the vector kernel may be preferable over the kernels optimized for large batch sizes: - const bool can_use_vector_kernel = Q->ne[0] <= 256 && Q->ne[0] % 64 == 0 && K->ne[1] % FATTN_KQ_STRIDE == 0; + const bool can_use_vector_kernel = !kv_is_turbo && Q->ne[0] <= 256 && Q->ne[0] % 64 == 0 && K->ne[1] % FATTN_KQ_STRIDE == 0; // If Turing tensor cores are available, use them: if (turing_mma_available(cc) && Q->ne[0] != 40 && Q->ne[0] != 72) { @@ -501,72 +509,80 @@ static best_fattn_kernel ggml_cuda_get_best_fattn_kernel(const int device, const return BEST_FATTN_KERNEL_TILE; } -// Pre-dequantize a turbo tensor to f16, returning a stack-allocated tensor copy. -// The caller must keep pool_buf alive until after FA completes. -static ggml_tensor turbo_pre_dequantize( - const ggml_tensor * src, - ggml_cuda_pool_alloc & pool_buf, - cudaStream_t stream) { - const int64_t n_elements = ggml_nelements(src); - - pool_buf.alloc(n_elements); - - to_fp16_cuda_t dequant = ggml_get_to_fp16_cuda(src->type); - GGML_ASSERT(dequant != nullptr); - dequant(src->data, pool_buf.ptr, n_elements, stream); - - // Scale existing strides from turbo block layout to f16 element layout. - // This preserves any permutation (e.g. ggml_permute swapping dims 1 and 2). - // The dequantized f16 data is in the same physical order as the turbo data, - // so the stride relationships must be preserved, just rescaled. - const size_t bs = ggml_blck_size(src->type); - const size_t ts = ggml_type_size(src->type); - - ggml_tensor tmp = *src; - tmp.type = GGML_TYPE_F16; - tmp.data = pool_buf.ptr; - tmp.nb[0] = sizeof(half); - tmp.nb[1] = src->nb[1] * bs * sizeof(half) / ts; - tmp.nb[2] = src->nb[2] * bs * sizeof(half) / ts; - tmp.nb[3] = src->nb[3] * bs * sizeof(half) / ts; - tmp.view_src = nullptr; - tmp.view_offs = 0; - - return tmp; -} - -void ggml_cuda_flash_attn_ext(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { - ggml_cuda_set_device(ctx.device); - +// The fused turbo kernel decodes TurboQuant blocks inside flash attention instead +// of expanding the whole KV cache to f16 first. It covers the decode case only +// (ncols <= 2); everything else falls back to the tile/MMA path, which converts +// via launch_fattn. +static bool ggml_cuda_fattn_turbo_fused_applies(const ggml_tensor * dst) { + const ggml_tensor * Q = dst->src[0]; const ggml_tensor * K = dst->src[1]; const ggml_tensor * V = dst->src[2]; - const bool k_is_turbo = ggml_type_is_turbo(K->type); - const bool v_is_turbo = V && ggml_type_is_turbo(V->type); - - // Pre-dequantize turbo KV to f16 so standard FA kernels can handle them. - // Pool buffers must outlive the FA dispatch (RAII frees on scope exit). - ggml_cuda_pool_alloc k_pool(ctx.pool()); - ggml_cuda_pool_alloc v_pool(ctx.pool()); - ggml_tensor k_f16, v_f16; + if (!K || !V) { + return false; + } + if (!ggml_type_is_turbo(K->type) || K->type != V->type) { + return false; // mixed turbo/f16 has no instance + } + if (Q->ne[0] != K->ne[0] || Q->ne[0] != V->ne[0]) { + return false; + } + if (Q->ne[0] != 128 && Q->ne[0] != 256) { + return false; + } + if (Q->ne[1] > 2) { + return false; // prompt processing: not a vector-kernel shape + } + // Like the generic vector kernel, this one consumes whole nthreads-sized + // blocks of KV without a tail check, so the KV length has to be padded. + if (K->ne[1] % FATTN_KQ_STRIDE != 0) { + return false; + } - cudaStream_t stream = ctx.stream(); + // Off by default: measured slower than the bulk-conversion path it replaces + // (16.6 vs 23.6 t/s at depth 65536 on gfx1201, see docs/turboquant.md). + // Correct, but not yet worth enabling. Opt in with GGML_CUDA_TURBO_FUSED_FA=1. + static const bool enabled = getenv("GGML_CUDA_TURBO_FUSED_FA") != nullptr; + return enabled; +} - // Save original src pointers - ggml_tensor * orig_k = dst->src[1]; - ggml_tensor * orig_v = dst->src[2]; +static void ggml_cuda_flash_attn_ext_vec_turbo(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { + const ggml_tensor * Q = dst->src[0]; + const ggml_tensor * K = dst->src[1]; - if (k_is_turbo) { - k_f16 = turbo_pre_dequantize(K, k_pool, stream); - dst->src[1] = &k_f16; - } - if (v_is_turbo) { - v_f16 = turbo_pre_dequantize(V, v_pool, stream); - dst->src[2] = &v_f16; + switch (Q->ne[0]) { + case 128: + if (K->type == GGML_TYPE_TURBO3_0) { + ggml_cuda_flash_attn_ext_vec_turbo_case<128, GGML_TYPE_TURBO3_0>(ctx, dst); + } else { + ggml_cuda_flash_attn_ext_vec_turbo_case<128, GGML_TYPE_TURBO4_0>(ctx, dst); + } + break; + case 256: + if (K->type == GGML_TYPE_TURBO3_0) { + ggml_cuda_flash_attn_ext_vec_turbo_case<256, GGML_TYPE_TURBO3_0>(ctx, dst); + } else { + ggml_cuda_flash_attn_ext_vec_turbo_case<256, GGML_TYPE_TURBO4_0>(ctx, dst); + } + break; + default: + GGML_ABORT("unsupported head size for fused turbo FA"); } +} + +void ggml_cuda_flash_attn_ext(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { + ggml_cuda_set_device(ctx.device); + if (ggml_cuda_fattn_turbo_fused_applies(dst)) { + ggml_cuda_flash_attn_ext_vec_turbo(ctx, dst); + return; + } - // Standard FA dispatch — now sees f16 tensors + // Turbo KV needs no special handling here: the tile/MMA/WMMA kernels all pass + // need_f16_K/V = true, so launch_fattn converts them via ggml_get_to_fp16_cuda + // (contiguous) or ggml_get_to_fp16_nc_cuda (strided) — both of which apply the + // inverse FWHT. The vector kernel, which would not convert, is excluded in + // ggml_cuda_get_best_fattn_kernel. switch (ggml_cuda_get_best_fattn_kernel(ggml_cuda_get_device(), dst)) { case BEST_FATTN_KERNEL_NONE: GGML_ABORT("fatal error"); @@ -583,10 +599,6 @@ void ggml_cuda_flash_attn_ext(ggml_backend_cuda_context & ctx, ggml_tensor * dst ggml_cuda_flash_attn_ext_mma_f16(ctx, dst); break; } - - // Restore original src pointers - dst->src[1] = orig_k; - dst->src[2] = orig_v; } bool ggml_cuda_flash_attn_ext_supported(int device, const ggml_tensor * dst) { diff --git a/ggml/src/ggml-cuda/template-instances/fattn-tile-instance-turbo.cu b/ggml/src/ggml-cuda/template-instances/fattn-tile-instance-turbo.cu new file mode 100644 index 000000000..f3ab4b148 --- /dev/null +++ b/ggml/src/ggml-cuda/template-instances/fattn-tile-instance-turbo.cu @@ -0,0 +1,12 @@ +// This file has been autogenerated by generate_cu_files.py, do not edit manually. + +#include "../fattn-tile.cuh" + +// TurboQuant instances: the tile kernel decodes the V cache itself instead of +// having launch_fattn expand it to f16 first. Only head sizes 128 and 256 exist +// because TurboQuant requires head_dim to be a multiple of 128. + +DECL_FATTN_TILE_CASE_TURBO(128, 128, GGML_TYPE_TURBO3_0); +DECL_FATTN_TILE_CASE_TURBO(128, 128, GGML_TYPE_TURBO4_0); +DECL_FATTN_TILE_CASE_TURBO(256, 256, GGML_TYPE_TURBO3_0); +DECL_FATTN_TILE_CASE_TURBO(256, 256, GGML_TYPE_TURBO4_0); diff --git a/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-turbo3_0-turbo3_0.cu b/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-turbo3_0-turbo3_0.cu deleted file mode 100644 index f8b44cee8..000000000 --- a/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-turbo3_0-turbo3_0.cu +++ /dev/null @@ -1,7 +0,0 @@ -// TurboQuant flash attention template instances - -#include "../fattn-vec.cuh" - -DECL_FATTN_VEC_CASE( 64, GGML_TYPE_TURBO3_0, GGML_TYPE_TURBO3_0); -DECL_FATTN_VEC_CASE(128, GGML_TYPE_TURBO3_0, GGML_TYPE_TURBO3_0); -DECL_FATTN_VEC_CASE(256, GGML_TYPE_TURBO3_0, GGML_TYPE_TURBO3_0); diff --git a/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-turbo4_0-turbo4_0.cu b/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-turbo4_0-turbo4_0.cu deleted file mode 100644 index 061000d03..000000000 --- a/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-turbo4_0-turbo4_0.cu +++ /dev/null @@ -1,7 +0,0 @@ -// TurboQuant flash attention template instances - -#include "../fattn-vec.cuh" - -DECL_FATTN_VEC_CASE( 64, GGML_TYPE_TURBO4_0, GGML_TYPE_TURBO4_0); -DECL_FATTN_VEC_CASE(128, GGML_TYPE_TURBO4_0, GGML_TYPE_TURBO4_0); -DECL_FATTN_VEC_CASE(256, GGML_TYPE_TURBO4_0, GGML_TYPE_TURBO4_0); diff --git a/ggml/src/ggml-cuda/turbo.cuh b/ggml/src/ggml-cuda/turbo.cuh new file mode 100644 index 000000000..ad7c44724 --- /dev/null +++ b/ggml/src/ggml-cuda/turbo.cuh @@ -0,0 +1,333 @@ +#pragma once + +#include "common.cuh" +#include "dequantize.cuh" + +// ============================================================ +// TurboQuant shared building blocks +// ============================================================ +// +// TurboQuant stores a KV row as fixed 128-element chunks. Each chunk is +// L2-normalized, rotated by a Walsh-Hadamard transform and scalar-quantized +// against a Lloyd-Max codebook: +// +// stored: c = quantize(H x / (||x|| sqrt(128))) (codebook indices + ||x||) +// recovered: x = ||x|| / sqrt(128) * H c +// +// H is the unnormalized +-1 Hadamard matrix, which is symmetric and orthogonal. +// A single element therefore cannot be decoded on its own — it needs the whole +// 128-element chunk. Everything below operates on complete chunks. + +#define TURBO_HEAD_DIM_GPU 128 +#define TURBO_BLOCKS_PER_CHUNK_GPU (TURBO_HEAD_DIM_GPU / 32) // 4 + +// 1/sqrt(128) — matches turbo_fwht_f32 in ggml-quants.c +#define TURBO_FWHT_SCALE_GPU 0.08838834764831844f + +// Unpack one 3-bit codebook index. Mirrors turbo_unpack3 in ggml-quants.c. +static __device__ __forceinline__ uint8_t turbo3_unpack_index(const block_turbo3_0 * x, int elem_in_block) { + const uint8_t * qs = x->qs; + + const int bit_off = elem_in_block * 3; + const int byte_idx = bit_off / 8; + const int shift = bit_off % 8; + + uint16_t raw = (uint16_t)qs[byte_idx] >> shift; + if (shift > 5 && byte_idx + 1 < 12) { + raw |= (uint16_t)qs[byte_idx + 1] << (8 - shift); + } + return (uint8_t)(raw & 0x07); +} + +static __device__ __forceinline__ uint8_t turbo4_unpack_index(const block_turbo4_0 * x, int elem_in_block) { + const uint8_t packed = x->qs[elem_in_block / 2]; + return (elem_in_block & 1) ? ((packed >> 4) & 0x0F) : (packed & 0x0F); +} + +// ------------------------------------------------------------ +// FWHT, shared-memory variant: 128 threads, one element each +// ------------------------------------------------------------ +// +// Stage order and pairing (i, i+h) are deliberately identical to +// turbo_fwht_f32 in ggml-quants.c so both produce the same rounding. +static __device__ __forceinline__ void turbo_fwht_smem_128(float * smem, int tid) { + for (int h = 1; h < TURBO_HEAD_DIM_GPU; h *= 2) { + if (tid < TURBO_HEAD_DIM_GPU/2) { + const int group = tid / h; + const int pos = tid % h; + const int i = group * h * 2 + pos; + const float a = smem[i]; + const float b = smem[i + h]; + smem[i] = a + b; + smem[i + h] = a - b; + } + __syncthreads(); + } +} + +// ------------------------------------------------------------ +// FWHT, warp variant: 32 lanes x 4 consecutive elements +// ------------------------------------------------------------ +// +// Lane l holds elements [4l, 4l+4) of the chunk, so the element index splits as +// i = 4*lane + r, r in [0,4) +// with bits 0..1 selecting the register and bits 2..6 selecting the lane. The +// first two butterfly stages (h = 1, 2) are therefore register-local, the +// remaining five (h = 4..64) are lane exchanges with XOR mask h/4. +// +// The arithmetic per butterfly is the same as in the shared-memory variant: the +// lower partner computes a+b, the upper computes a-b, in that order. Results are +// bit-identical to turbo_fwht_smem_128, which is what makes the fused kernel +// verifiable against the bulk dequantize path. +// +// Requires a full 32-lane warp with all lanes active. +static __device__ __forceinline__ void turbo_fwht_warp_128(float v[4], int lane) { + // h = 1: pairs (0,1) and (2,3) inside the register file + { + const float a0 = v[0], b0 = v[1]; + const float a1 = v[2], b1 = v[3]; + v[0] = a0 + b0; v[1] = a0 - b0; + v[2] = a1 + b1; v[3] = a1 - b1; + } + // h = 2: pairs (0,2) and (1,3) + { + const float a0 = v[0], b0 = v[2]; + const float a1 = v[1], b1 = v[3]; + v[0] = a0 + b0; v[2] = a0 - b0; + v[1] = a1 + b1; v[3] = a1 - b1; + } + // h = 4, 8, 16, 32, 64: cross-lane, XOR mask h/4 +#pragma unroll + for (int mask = 1; mask < 32; mask *= 2) { + const bool upper = (lane & mask) != 0; +#pragma unroll + for (int r = 0; r < 4; ++r) { + const float p = __shfl_xor_sync(0xffffffff, v[r], mask, 32); + v[r] = upper ? (p - v[r]) : (v[r] + p); + } + } +} + +// ------------------------------------------------------------ +// Chunk decode into shared memory (128 threads per chunk) +// ------------------------------------------------------------ + +template +static __device__ __forceinline__ void turbo_decode_chunk_smem( + const block_t * blocks, int64_t ib_chunk, float * smem, int tid); + +template <> +__device__ __forceinline__ void turbo_decode_chunk_smem( + const block_turbo3_0 * blocks, int64_t ib_chunk, float * smem, int tid) { + const int local_block = tid / TURBO3_BLOCK_SIZE; + const int elem_in_block = tid % TURBO3_BLOCK_SIZE; + + smem[tid] = dc_codebook_3bit[turbo3_unpack_index(blocks + ib_chunk + local_block, elem_in_block)]; + __syncthreads(); + + turbo_fwht_smem_128(smem, tid); + + // The norm is replicated in every block of the chunk; read it from the first. + const float norm = __half2float(blocks[ib_chunk].d); + smem[tid] *= TURBO_FWHT_SCALE_GPU * norm; + __syncthreads(); +} + +template <> +__device__ __forceinline__ void turbo_decode_chunk_smem( + const block_turbo4_0 * blocks, int64_t ib_chunk, float * smem, int tid) { + const int local_block = tid / TURBO4_BLOCK_SIZE; + const int elem_in_block = tid % TURBO4_BLOCK_SIZE; + + smem[tid] = dc_codebook_4bit[turbo4_unpack_index(blocks + ib_chunk + local_block, elem_in_block)]; + __syncthreads(); + + turbo_fwht_smem_128(smem, tid); + + const float norm = __half2float(blocks[ib_chunk].d); + smem[tid] *= TURBO_FWHT_SCALE_GPU * norm; + __syncthreads(); +} + +// ------------------------------------------------------------ +// Chunk decode with an arbitrary thread count +// ------------------------------------------------------------ +// +// The tile kernel runs with fewer threads than the 128 a chunk has elements +// (64 for head_dim 256 on RDNA), so the fixed one-thread-per-element form above +// does not fit. This variant strides over the chunk instead. The butterfly +// stages keep the same order and operand order, so results stay bit-identical +// to turbo_fwht_smem_128. +// +// All nthreads threads of the block must call this; it synchronizes internally. + +template +static __device__ __forceinline__ void turbo_fwht_smem_128_n(float * smem, int tid) { + for (int h = 1; h < TURBO_HEAD_DIM_GPU; h *= 2) { +#pragma unroll + for (int b0 = 0; b0 < TURBO_HEAD_DIM_GPU/2; b0 += nthreads) { + const int b = b0 + tid; + if (nthreads >= TURBO_HEAD_DIM_GPU/2 && b >= TURBO_HEAD_DIM_GPU/2) { + continue; + } + const int group = b / h; + const int pos = b % h; + const int i = group * h * 2 + pos; + const float a = smem[i]; + const float c = smem[i + h]; + smem[i] = a + c; + smem[i + h] = a - c; + } + __syncthreads(); + } +} + +template +static __device__ __forceinline__ void turbo_decode_chunk_smem_n( + const block_t * blocks, int64_t ib_chunk, float * smem, int tid) { + constexpr int bs = (sizeof(block_t) == sizeof(block_turbo3_0)) ? TURBO3_BLOCK_SIZE : TURBO4_BLOCK_SIZE; + +#pragma unroll + for (int e0 = 0; e0 < TURBO_HEAD_DIM_GPU; e0 += nthreads) { + const int e = e0 + tid; + if (nthreads >= TURBO_HEAD_DIM_GPU && e >= TURBO_HEAD_DIM_GPU) { + continue; + } + const int blk = e / bs; + const int eib = e % bs; + if constexpr (sizeof(block_t) == sizeof(block_turbo3_0)) { + smem[e] = dc_codebook_3bit[turbo3_unpack_index((const block_turbo3_0 *)(blocks + ib_chunk + blk), eib)]; + } else { + smem[e] = dc_codebook_4bit[turbo4_unpack_index((const block_turbo4_0 *)(blocks + ib_chunk + blk), eib)]; + } + } + __syncthreads(); + + turbo_fwht_smem_128_n(smem, tid); + + const float scale = TURBO_FWHT_SCALE_GPU * __half2float(blocks[ib_chunk].d); +#pragma unroll + for (int e0 = 0; e0 < TURBO_HEAD_DIM_GPU; e0 += nthreads) { + const int e = e0 + tid; + if (nthreads >= TURBO_HEAD_DIM_GPU && e >= TURBO_HEAD_DIM_GPU) { + continue; + } + smem[e] *= scale; + } + __syncthreads(); +} + +// ------------------------------------------------------------ +// Warp-local FWHT over a 128-float shared-memory buffer +// ------------------------------------------------------------ +// +// Like turbo_fwht_smem_128_n but synchronizing only within the warp, so each +// warp can transform its own buffer independently. Used once per kernel launch +// to bring an accumulator back out of the Hadamard domain. + +template +static __device__ __forceinline__ void turbo_fwht_smem_128_warp(float * smem, int lane) { + for (int h = 1; h < TURBO_HEAD_DIM_GPU; h *= 2) { +#pragma unroll + for (int b0 = 0; b0 < TURBO_HEAD_DIM_GPU/2; b0 += nthreads) { + const int b = b0 + lane; + if (nthreads > TURBO_HEAD_DIM_GPU/2 && b >= TURBO_HEAD_DIM_GPU/2) { + continue; + } + const int group = b / h; + const int pos = b % h; + const int i = group * h * 2 + pos; + const float a = smem[i]; + const float c = smem[i + h]; + smem[i] = a + c; + smem[i + h] = a - c; + } +#ifdef GGML_USE_HIP + __builtin_amdgcn_wave_barrier(); +#else + __syncwarp(); +#endif + } +} + +// ------------------------------------------------------------ +// Raw chunk load: codebook values only, no transform +// ------------------------------------------------------------ +// +// For accumulating in the Hadamard domain. The caller applies one inverse FWHT +// at the end instead of one per KV token — the whole point of the exercise, since +// the transform is a dependency chain and the accumulation is not. +// +// Returns the chunk's codebook values already scaled by norm/sqrt(128), so the +// downstream math is identical to decoded values apart from the missing rotation. + +template +static __device__ __forceinline__ void turbo_load_chunk_raw_warp( + const block_t * blocks, int64_t ib_chunk, float v[4], int lane) { + constexpr int bs = (sizeof(block_t) == sizeof(block_turbo3_0)) ? TURBO3_BLOCK_SIZE : TURBO4_BLOCK_SIZE; + + const float scale = TURBO_FWHT_SCALE_GPU * __half2float(blocks[ib_chunk].d); + +#pragma unroll + for (int r = 0; r < 4; ++r) { + const int elem = 4*lane + r; + const int blk = elem / bs; + const int eib = elem % bs; + if constexpr (sizeof(block_t) == sizeof(block_turbo3_0)) { + v[r] = dc_codebook_3bit[turbo3_unpack_index((const block_turbo3_0 *)(blocks + ib_chunk + blk), eib)] * scale; + } else { + v[r] = dc_codebook_4bit[turbo4_unpack_index((const block_turbo4_0 *)(blocks + ib_chunk + blk), eib)] * scale; + } + } +} + +// ------------------------------------------------------------ +// Chunk decode into registers (one warp per chunk) +// ------------------------------------------------------------ +// +// Lane l fills v[0..3] with elements [4l, 4l+4) of the chunk. Equivalent to +// turbo_decode_chunk_smem but without shared memory or block synchronization. + +template +static __device__ __forceinline__ void turbo_decode_chunk_warp( + const block_t * blocks, int64_t ib_chunk, float v[4], int lane); + +template <> +__device__ __forceinline__ void turbo_decode_chunk_warp( + const block_turbo3_0 * blocks, int64_t ib_chunk, float v[4], int lane) { +#pragma unroll + for (int r = 0; r < 4; ++r) { + const int elem = 4*lane + r; + const int local_block = elem / TURBO3_BLOCK_SIZE; + const int elem_in_block = elem % TURBO3_BLOCK_SIZE; + v[r] = dc_codebook_3bit[turbo3_unpack_index(blocks + ib_chunk + local_block, elem_in_block)]; + } + + turbo_fwht_warp_128(v, lane); + + const float scale = TURBO_FWHT_SCALE_GPU * __half2float(blocks[ib_chunk].d); +#pragma unroll + for (int r = 0; r < 4; ++r) { + v[r] *= scale; + } +} + +template <> +__device__ __forceinline__ void turbo_decode_chunk_warp( + const block_turbo4_0 * blocks, int64_t ib_chunk, float v[4], int lane) { +#pragma unroll + for (int r = 0; r < 4; ++r) { + const int elem = 4*lane + r; + const int local_block = elem / TURBO4_BLOCK_SIZE; + const int elem_in_block = elem % TURBO4_BLOCK_SIZE; + v[r] = dc_codebook_4bit[turbo4_unpack_index(blocks + ib_chunk + local_block, elem_in_block)]; + } + + turbo_fwht_warp_128(v, lane); + + const float scale = TURBO_FWHT_SCALE_GPU * __half2float(blocks[ib_chunk].d); +#pragma unroll + for (int r = 0; r < 4; ++r) { + v[r] *= scale; + } +} diff --git a/ggml/src/ggml-hip/CMakeLists.txt b/ggml/src/ggml-hip/CMakeLists.txt index 9a4a4252e..291b48374 100644 --- a/ggml/src/ggml-hip/CMakeLists.txt +++ b/ggml/src/ggml-hip/CMakeLists.txt @@ -75,9 +75,7 @@ else() ../ggml-cuda/template-instances/fattn-vec-instance-f16-f16.cu ../ggml-cuda/template-instances/fattn-vec-instance-q4_0-q4_0.cu ../ggml-cuda/template-instances/fattn-vec-instance-q8_0-q8_0.cu - ../ggml-cuda/template-instances/fattn-vec-instance-bf16-bf16.cu - ../ggml-cuda/template-instances/fattn-vec-instance-turbo3_0-turbo3_0.cu - ../ggml-cuda/template-instances/fattn-vec-instance-turbo4_0-turbo4_0.cu) + ../ggml-cuda/template-instances/fattn-vec-instance-bf16-bf16.cu) endif() ggml_add_backend_library(ggml-hip diff --git a/ggml/src/ggml-quants.c b/ggml/src/ggml-quants.c index 52d3c2dd4..1cf1b4d1d 100644 --- a/ggml/src/ggml-quants.c +++ b/ggml/src/ggml-quants.c @@ -748,6 +748,30 @@ void dequantize_row_turbo4_0(const block_turbo4_0 * GGML_RESTRICT x, float * GGM } } +// ============================================================ +// TurboQuant chunked quantization entry points +// ============================================================ +// +// Row-major wrappers used by ggml_quantize_chunk. The FWHT runs over fixed +// 128-element chunks, so a row must be a multiple of TURBO_HEAD_DIM — otherwise +// a chunk would straddle two rows and mix their values. + +size_t quantize_turbo3_0(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrow, int64_t n_per_row, const float * quant_weights) { + (void)quant_weights; // not used + GGML_ASSERT(n_per_row % TURBO_HEAD_DIM == 0); + const size_t row_size = ggml_row_size(GGML_TYPE_TURBO3_0, n_per_row); + quantize_row_turbo3_0_ref(src, (block_turbo3_0 *)dst, (int64_t)nrow*n_per_row); + return nrow * row_size; +} + +size_t quantize_turbo4_0(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrow, int64_t n_per_row, const float * quant_weights) { + (void)quant_weights; // not used + GGML_ASSERT(n_per_row % TURBO_HEAD_DIM == 0); + const size_t row_size = ggml_row_size(GGML_TYPE_TURBO4_0, n_per_row); + quantize_row_turbo4_0_ref(src, (block_turbo4_0 *)dst, (int64_t)nrow*n_per_row); + return nrow * row_size; +} + // ============================================================ // TurboQuant vec_dot (for flash attention compatibility) // ============================================================ diff --git a/ggml/src/ggml-quants.h b/ggml/src/ggml-quants.h index 7a796321c..e18561931 100644 --- a/ggml/src/ggml-quants.h +++ b/ggml/src/ggml-quants.h @@ -105,6 +105,13 @@ GGML_API size_t quantize_q8_0(const float * GGML_RESTRICT src, void * GGML_RESTR GGML_API size_t quantize_mxfp4(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix); GGML_API size_t quantize_nvfp4(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix); +// TurboQuant. n_per_row must be a multiple of 128 (the FWHT chunk size). +GGML_API size_t quantize_turbo3_0(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix); +GGML_API size_t quantize_turbo4_0(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix); + +GGML_API void ggml_vec_dot_turbo3_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc); +GGML_API void ggml_vec_dot_turbo4_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc); + GGML_API void iq2xs_init_impl(enum ggml_type type); GGML_API void iq2xs_free_impl(enum ggml_type type); GGML_API void iq3xs_init_impl(int grid_size); diff --git a/ggml/src/ggml.c b/ggml/src/ggml.c index 6fcd3e261..4b7786d3d 100644 --- a/ggml/src/ggml.c +++ b/ggml/src/ggml.c @@ -7691,6 +7691,8 @@ size_t ggml_quantize_chunk( case GGML_TYPE_IQ1_M: result = quantize_iq1_m (src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break; case GGML_TYPE_IQ4_NL: result = quantize_iq4_nl (src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break; case GGML_TYPE_IQ4_XS: result = quantize_iq4_xs (src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break; + case GGML_TYPE_TURBO3_0: result = quantize_turbo3_0(src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break; + case GGML_TYPE_TURBO4_0: result = quantize_turbo4_0(src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break; case GGML_TYPE_F16: { size_t elemsize = sizeof(ggml_fp16_t); diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp index 33493cb70..dd061faec 100644 --- a/src/llama-kv-cache.cpp +++ b/src/llama-kv-cache.cpp @@ -132,19 +132,22 @@ llama_kv_cache::llama_kv_cache( throw std::runtime_error("failed to create ggml context for kv cache"); } - // TurboQuant requires head_dim=128 for the FWHT transform + // TurboQuant applies the FWHT in fixed 128-element chunks, so head_dim + // must be a multiple of 128. A 256-dim head is quantized as two + // independently normalized halves — each half is a unit vector in R^128, + // which is the distribution the codebooks were fitted for. if (type_k == GGML_TYPE_TURBO3_0 || type_k == GGML_TYPE_TURBO4_0) { const uint32_t n_embd_head_k = hparams.n_embd_head_k(il); - if (n_embd_head_k != 128) { - LLAMA_LOG_ERROR("%s: TurboQuant requires head_dim=128, got %d (layer %d)\n", __func__, n_embd_head_k, il); - throw std::runtime_error("turbo types require head_dim=128"); + if (n_embd_head_k % 128 != 0) { + LLAMA_LOG_ERROR("%s: TurboQuant requires head_dim to be a multiple of 128, got %d (layer %d)\n", __func__, n_embd_head_k, il); + throw std::runtime_error("turbo types require head_dim to be a multiple of 128"); } } if (type_v == GGML_TYPE_TURBO3_0 || type_v == GGML_TYPE_TURBO4_0) { const uint32_t n_embd_head_v = hparams.n_embd_head_v(il); - if (n_embd_head_v != 128) { - LLAMA_LOG_ERROR("%s: TurboQuant requires head_dim=128, got %d (layer %d)\n", __func__, n_embd_head_v, il); - throw std::runtime_error("turbo types require head_dim=128"); + if (n_embd_head_v % 128 != 0) { + LLAMA_LOG_ERROR("%s: TurboQuant requires head_dim to be a multiple of 128, got %d (layer %d)\n", __func__, n_embd_head_v, il); + throw std::runtime_error("turbo types require head_dim to be a multiple of 128"); } } diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index 6a4f9b634..36862d412 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -109,6 +109,12 @@ static void init_tensor_uniform(ggml_tensor * tensor, float min = -1.0f, float m { // parallel quantization by block size_t blck_size = ggml_blck_size(tensor->type); + if (tensor->type == GGML_TYPE_TURBO3_0 || tensor->type == GGML_TYPE_TURBO4_0) { + // TurboQuant normalizes and rotates whole 128-element chunks, so a + // quantization call has to cover one — the 32-element block size would + // split a chunk across calls and produce garbage. + blck_size = 128; + } size_t n_blocks = nels / blck_size; auto quantize_thread = [&](size_t start, size_t end) { @@ -8610,8 +8616,13 @@ static std::vector> make_test_cases_eval() { for (int nb : { 1, 3, 32, 75, }) { for (ggml_prec prec : {GGML_PREC_F32, GGML_PREC_DEFAULT}) { if (hsk != 128 && prec == GGML_PREC_DEFAULT) continue; - for (ggml_type type_KV : {GGML_TYPE_F32, GGML_TYPE_F16, GGML_TYPE_BF16, GGML_TYPE_Q8_0, GGML_TYPE_Q4_0}) { - if (type_KV != GGML_TYPE_F16 && hsk != 64 && hsk != 72) continue; + for (ggml_type type_KV : {GGML_TYPE_F32, GGML_TYPE_F16, GGML_TYPE_BF16, GGML_TYPE_Q8_0, GGML_TYPE_Q4_0, GGML_TYPE_TURBO3_0, GGML_TYPE_TURBO4_0}) { + const bool is_turbo = type_KV == GGML_TYPE_TURBO3_0 || type_KV == GGML_TYPE_TURBO4_0; + // TurboQuant applies the FWHT over 128-element chunks, so head_dim + // must be a multiple of 128 — but it is not limited to 64/72 like + // the other quantized KV types. + if (is_turbo && (hsk % 128 != 0 || hsv % 128 != 0)) continue; + if (!is_turbo && type_KV != GGML_TYPE_F16 && hsk != 64 && hsk != 72) continue; test_cases.emplace_back(new test_flash_attn_ext( hsk, hsv, nh, {nr2, nr3}, kv, nb, mask, sinks, max_bias, logit_softcap, prec, type_KV)); // run fewer test cases permuted diff --git a/tests/test-turboquant.cpp b/tests/test-turboquant.cpp index 96320b032..671a81638 100644 --- a/tests/test-turboquant.cpp +++ b/tests/test-turboquant.cpp @@ -2,6 +2,12 @@ // // Validates that the quantize/dequantize pipeline in ggml-quants.c produces // MSE*d values consistent with the paper (Zandieh et al., ICLR 2026). +// +// head_dim is exercised at 128 and at multiples of it. TurboQuant processes a +// row in fixed 128-element chunks, so a 256-dim head is quantized as two +// independently normalized halves. Each half is a unit vector in R^128, which +// is exactly the distribution the codebooks were fitted for — the reconstruction +// error must therefore land in the same range as for a plain 128-dim head. #include "ggml.h" @@ -12,11 +18,13 @@ #include #include -static constexpr int HEAD_DIM = 128; -static constexpr int BLOCK_SIZE = 32; -static constexpr int BLOCKS_PER_CHUNK = HEAD_DIM / BLOCK_SIZE; -static constexpr int N_VECTORS = 10000; -static constexpr int N_ELEMENTS = N_VECTORS * HEAD_DIM; +static constexpr int CHUNK_DIM = 128; // TurboQuant FWHT chunk size (fixed) +static constexpr int BLOCK_SIZE = 32; +static constexpr int N_VECTORS = 10000; + +// head_dim values under test. Every one must be a multiple of CHUNK_DIM. +static constexpr int HEAD_DIMS[] = { 128, 256, 384 }; +static constexpr int N_HEAD_DIMS = (int)(sizeof(HEAD_DIMS) / sizeof(HEAD_DIMS[0])); // Expected MSE*d ranges (paper: TQ3 ~0.034, TQ4 ~0.009 for d=128) static constexpr float TQ3_MSE_D_MIN = 0.025f; @@ -46,25 +54,28 @@ static void fwht_f32(float * x, int n) { } // ============================================================ -// Test 1: FWHT self-inverse (FWHT(FWHT(x)) == x) +// Test 1: FWHT self-inverse (FWHT(FWHT(x)) == x), per 128-chunk // ============================================================ -static int test_fwht_self_inverse(void) { - printf(" FWHT self-inverse (d=%d)... ", HEAD_DIM); +static int test_fwht_self_inverse(int head_dim) { + printf(" FWHT self-inverse (d=%d, %d chunk(s))... ", head_dim, head_dim / CHUNK_DIM); - float orig[HEAD_DIM]; - float work[HEAD_DIM]; + std::vector orig(head_dim); + std::vector work(head_dim); - for (int i = 0; i < HEAD_DIM; i++) { + for (int i = 0; i < head_dim; i++) { orig[i] = sinf((float)(i + 1) * 0.7f) * 2.0f; } - memcpy(work, orig, sizeof(orig)); + work = orig; - fwht_f32(work, HEAD_DIM); - fwht_f32(work, HEAD_DIM); + // Applied chunk-wise, mirroring how ggml-quants.c walks the row + for (int off = 0; off < head_dim; off += CHUNK_DIM) { + fwht_f32(work.data() + off, CHUNK_DIM); + fwht_f32(work.data() + off, CHUNK_DIM); + } float max_err = 0.0f; - for (int i = 0; i < HEAD_DIM; i++) { + for (int i = 0; i < head_dim; i++) { float err = fabsf(work[i] - orig[i]); if (err > max_err) { max_err = err; } } @@ -86,33 +97,35 @@ static void generate_random_vectors(float * dst, int n_elements, unsigned int se } } -static int test_roundtrip_mse(ggml_type type, float mse_d_min, float mse_d_max) { +static int test_roundtrip_mse(ggml_type type, float mse_d_min, float mse_d_max, int head_dim) { const char * name = ggml_type_name(type); - printf(" %s roundtrip MSE*d (n=%d, d=%d)... ", name, N_VECTORS, HEAD_DIM); + printf(" %s roundtrip MSE*d (n=%d, d=%d)... ", name, N_VECTORS, head_dim); const ggml_type_traits * traits = ggml_get_type_traits(type); assert(traits->from_float_ref != nullptr); assert(traits->to_float != nullptr); - std::vector src(N_ELEMENTS); - std::vector dst(N_ELEMENTS); - size_t quant_size = (size_t)N_ELEMENTS / BLOCK_SIZE * ggml_type_size(type); + const int n_elements = N_VECTORS * head_dim; + + std::vector src(n_elements); + std::vector dst(n_elements); + size_t quant_size = (size_t)n_elements / BLOCK_SIZE * ggml_type_size(type); std::vector quant(quant_size); - generate_random_vectors(src.data(), N_ELEMENTS, 42); + generate_random_vectors(src.data(), n_elements, 42); - traits->from_float_ref(src.data(), quant.data(), N_ELEMENTS); - traits->to_float(quant.data(), dst.data(), N_ELEMENTS); + traits->from_float_ref(src.data(), quant.data(), n_elements); + traits->to_float(quant.data(), dst.data(), n_elements); // MSE*d = E[ ||x - x̃||² / ||x||² ] (normalized reconstruction error) double total_nmse = 0.0; for (int v = 0; v < N_VECTORS; v++) { - double err_sq = 0.0; + double err_sq = 0.0; double norm_sq = 0.0; - for (int i = 0; i < HEAD_DIM; i++) { - double diff = (double)src[v * HEAD_DIM + i] - (double)dst[v * HEAD_DIM + i]; - err_sq += diff * diff; - norm_sq += (double)src[v * HEAD_DIM + i] * (double)src[v * HEAD_DIM + i]; + for (int i = 0; i < head_dim; i++) { + double diff = (double)src[v * head_dim + i] - (double)dst[v * head_dim + i]; + err_sq += diff * diff; + norm_sq += (double)src[v * head_dim + i] * (double)src[v * head_dim + i]; } if (norm_sq > 1e-20) { total_nmse += err_sq / norm_sq; @@ -129,50 +142,50 @@ static int test_roundtrip_mse(ggml_type type, float mse_d_min, float mse_d_max) // Test 4: Bit-pack determinism and sanity // ============================================================ -static int test_bitpack_deterministic(ggml_type type) { +static int test_bitpack_deterministic(ggml_type type, int head_dim) { const char * name = ggml_type_name(type); - printf(" %s pack determinism... ", name); + printf(" %s pack determinism (d=%d)... ", name, head_dim); const ggml_type_traits * traits = ggml_get_type_traits(type); - float src[HEAD_DIM]; - for (int i = 0; i < HEAD_DIM; i++) { + std::vector src(head_dim); + for (int i = 0; i < head_dim; i++) { src[i] = cosf((float)i * 0.31415f); } - size_t qsize = BLOCKS_PER_CHUNK * ggml_type_size(type); + size_t qsize = (size_t)(head_dim / BLOCK_SIZE) * ggml_type_size(type); std::vector q1(qsize); std::vector q2(qsize); - traits->from_float_ref(src, q1.data(), HEAD_DIM); - traits->from_float_ref(src, q2.data(), HEAD_DIM); + traits->from_float_ref(src.data(), q1.data(), head_dim); + traits->from_float_ref(src.data(), q2.data(), head_dim); bool pass = memcmp(q1.data(), q2.data(), qsize) == 0; printf("%s\n", pass ? "ok" : "FAILED (non-deterministic)"); return pass ? 0 : 1; } -static int test_bitpack_sanity(ggml_type type) { +static int test_bitpack_sanity(ggml_type type, int head_dim) { const char * name = ggml_type_name(type); - printf(" %s dequantize sanity... ", name); + printf(" %s dequantize sanity (d=%d)... ", name, head_dim); const ggml_type_traits * traits = ggml_get_type_traits(type); - float src[HEAD_DIM]; - float dst[HEAD_DIM]; - for (int i = 0; i < HEAD_DIM; i++) { + std::vector src(head_dim); + std::vector dst(head_dim); + for (int i = 0; i < head_dim; i++) { src[i] = cosf((float)i * 0.31415f); } - size_t qsize = BLOCKS_PER_CHUNK * ggml_type_size(type); + size_t qsize = (size_t)(head_dim / BLOCK_SIZE) * ggml_type_size(type); std::vector q(qsize); - traits->from_float_ref(src, q.data(), HEAD_DIM); - traits->to_float(q.data(), dst, HEAD_DIM); + traits->from_float_ref(src.data(), q.data(), head_dim); + traits->to_float(q.data(), dst.data(), head_dim); - bool all_finite = true; + bool all_finite = true; bool any_nonzero = false; - for (int i = 0; i < HEAD_DIM; i++) { + for (int i = 0; i < head_dim; i++) { if (!std::isfinite(dst[i])) { all_finite = false; } if (fabsf(dst[i]) > 1e-10f) { any_nonzero = true; } } @@ -185,6 +198,38 @@ static int test_bitpack_sanity(ggml_type type) { return pass ? 0 : 1; } +// ============================================================ +// Test 5: chunk independence +// ============================================================ +// +// A 256-dim head must quantize to exactly the same bytes as two separate +// 128-dim heads holding the same values. This is what makes the existing +// codebooks valid for head_dim 256: the halves are processed independently. + +static int test_chunk_independence(ggml_type type) { + const char * name = ggml_type_name(type); + printf(" %s chunk independence (256 == 2x128)... ", name); + + const ggml_type_traits * traits = ggml_get_type_traits(type); + const size_t bytes_per_chunk = (size_t)(CHUNK_DIM / BLOCK_SIZE) * ggml_type_size(type); + + std::vector src(2 * CHUNK_DIM); + generate_random_vectors(src.data(), 2 * CHUNK_DIM, 7); + + // One pass over 256 elements + std::vector q_joint(2 * bytes_per_chunk); + traits->from_float_ref(src.data(), q_joint.data(), 2 * CHUNK_DIM); + + // Two separate passes over 128 elements each + std::vector q_split(2 * bytes_per_chunk); + traits->from_float_ref(src.data(), q_split.data(), CHUNK_DIM); + traits->from_float_ref(src.data() + CHUNK_DIM, q_split.data() + bytes_per_chunk, CHUNK_DIM); + + bool pass = memcmp(q_joint.data(), q_split.data(), 2 * bytes_per_chunk) == 0; + printf("%s\n", pass ? "ok" : "FAILED (chunks are not independent)"); + return pass ? 0 : 1; +} + // ============================================================ // Main // ============================================================ @@ -193,25 +238,43 @@ int main(void) { printf("TurboQuant CPU reference tests\n"); printf("==============================\n\n"); - int n_fail = 0; + int n_fail = 0; + int n_total = 0; printf("Test 1: FWHT self-inverse\n"); - n_fail += test_fwht_self_inverse(); + for (int i = 0; i < N_HEAD_DIMS; i++) { + n_fail += test_fwht_self_inverse(HEAD_DIMS[i]); + n_total++; + } printf("\nTest 2: TQ3 roundtrip MSE\n"); - n_fail += test_roundtrip_mse(GGML_TYPE_TURBO3_0, TQ3_MSE_D_MIN, TQ3_MSE_D_MAX); + for (int i = 0; i < N_HEAD_DIMS; i++) { + n_fail += test_roundtrip_mse(GGML_TYPE_TURBO3_0, TQ3_MSE_D_MIN, TQ3_MSE_D_MAX, HEAD_DIMS[i]); + n_total++; + } printf("\nTest 3: TQ4 roundtrip MSE\n"); - n_fail += test_roundtrip_mse(GGML_TYPE_TURBO4_0, TQ4_MSE_D_MIN, TQ4_MSE_D_MAX); + for (int i = 0; i < N_HEAD_DIMS; i++) { + n_fail += test_roundtrip_mse(GGML_TYPE_TURBO4_0, TQ4_MSE_D_MIN, TQ4_MSE_D_MAX, HEAD_DIMS[i]); + n_total++; + } printf("\nTest 4: Bit-pack tests\n"); - n_fail += test_bitpack_deterministic(GGML_TYPE_TURBO3_0); - n_fail += test_bitpack_deterministic(GGML_TYPE_TURBO4_0); - n_fail += test_bitpack_sanity(GGML_TYPE_TURBO3_0); - n_fail += test_bitpack_sanity(GGML_TYPE_TURBO4_0); + for (int i = 0; i < N_HEAD_DIMS; i++) { + n_fail += test_bitpack_deterministic(GGML_TYPE_TURBO3_0, HEAD_DIMS[i]); + n_fail += test_bitpack_deterministic(GGML_TYPE_TURBO4_0, HEAD_DIMS[i]); + n_fail += test_bitpack_sanity(GGML_TYPE_TURBO3_0, HEAD_DIMS[i]); + n_fail += test_bitpack_sanity(GGML_TYPE_TURBO4_0, HEAD_DIMS[i]); + n_total += 4; + } + + printf("\nTest 5: Chunk independence\n"); + n_fail += test_chunk_independence(GGML_TYPE_TURBO3_0); + n_fail += test_chunk_independence(GGML_TYPE_TURBO4_0); + n_total += 2; printf("\n==============================\n"); - printf("%d/%d tests passed\n", 7 - n_fail, 7); + printf("%d/%d tests passed\n", n_total - n_fail, n_total); return n_fail > 0 ? 1 : 0; }