feat: TQ3_0 (TurboQuant 3.5bpv) KV cache — 22% memory savings

README.md

@@ -106,7 +106,7 @@ What we ported and tuned:
 - C++/CUDA decode engine on top of ggml (no libllama, no Python runtime, Q4_K_M target path).
 - Three custom CUDA kernels for tree-aware SSM state rollback: `ggml_ssm_conv_tree`, `ggml_gated_delta_net_tree`, `ggml_gated_delta_net_tree_persist`.
 - DDTree budget swept for RTX 3090 + Q4_K_M target: **budget=22** is the sweet spot.
-- Q4_0 KV cache + sliding `target_feat` ring to fit 128K context in 24 GB with ~3% AL hit.
+- TQ3_0 KV cache (TurboQuant 3.5 bpv, default) + sliding `target_feat` ring to fit up to 256K context in 24 GB (Q4_0 available as legacy, tops out near 128K).
 
 [Full writeup →](dflash/README.md) · [Benchmarks →](dflash/RESULTS.md) · [Blog post →](https://lucebox.com/blog/dflash27b)
 

dflash/README.md

@@ -57,7 +57,7 @@ Qwen3.5-27B Q4_K_M, concurrency=1, n_gen=256, 10 prompts/dataset:
 
 AR = autoregressive (`test_generate`). DFlash+DDTree = tree verify at budget=22 with fast rollback (`test_dflash`). AL = Acceptance Length, average committed tokens per draft/verify step. Reproduce via `python3 scripts/bench_llm.py`.
 
-**128K context on 24 GB** via Q4_0 KV cache + sliding `target_feat` ring (4096 slots): ~3% AL hit vs F16 KV, 8× memory saving.
+**Up to 256K context on 24 GB** via TQ3_0 KV cache (3.5 bpv, default; Q4_0 legacy path tops out near 128K) + sliding `target_feat` ring (4096 slots). TQ3 = ~9.7× memory saving vs F16; Q4_0 = 8×.
 
 | Prompt length | KV     | Prefill time | Decode tok/s |
 |:-------------:|:------:|:------------:|:------------:|
@@ -70,7 +70,7 @@ Prefill numbers assume `--max-ctx` sized to the prompt (auto-fit in `run.py` / `
 
 HE 10-prompt bench mean in 128K mode (ctx=131072, ddtree-budget=16): **134.78 tok/s** at AL 8.33.
 
-Set `DFLASH27B_KV_Q4=1` to enable. Full sweep in [RESULTS.md](RESULTS.md).
+Set `DFLASH27B_KV_TQ3=1` (TQ3_0, 3.5 bpv, default) or `DFLASH27B_KV_Q4=1` (Q4_0, 4.5 bpv, legacy) to enable. Full sweep in [RESULTS.md](RESULTS.md).
 
 ## Qwen3.6-27B target (experimental)
 
@@ -134,7 +134,7 @@ python3 scripts/bench_he.py --n-gen 256 --ddtree-budget 22   # minimal HE bench
 
 **128K context mode:**
 ```bash
-DFLASH27B_KV_Q4=1 DFLASH27B_PREFILL_UBATCH=16 \
+DFLASH27B_KV_TQ3=1 DFLASH27B_PREFILL_UBATCH=16 \
   build/test_dflash models/Qwen3.5-27B-Q4_K_M.gguf \
   models/draft/model.safetensors /tmp/long_prompt.bin 64 /tmp/out.bin \
   --fast-rollback --ddtree --ddtree-budget=16 --max-ctx=N   # N = align_up(prompt + n_gen + 64, 256); up to 131072
@@ -194,7 +194,6 @@ Open an issue or PR against `Luce-Org/lucebox-hub`. Good first picks:
 
 - **Temperature / top-k sampling** in the verify path
 - **Q5_K_M / Q6_K target** support
-- **TurboQuant KV cache** — K=Q8_0 for attention accuracy, V=TQ (`turbo3`) for denser long-context at Q4-ish memory cost
 - **Full llama.cpp integration**: new arch, `llama-speculative-dflash.cpp`, `llama-cli` / `llama-server` wiring
 
 ## Citation

dflash/scripts/bench_daemon.py

@@ -11,7 +11,7 @@
 fast under the daemon as under a one-shot test_dflash invocation.
 
 Start the server first (same config the published numbers use):
-    DFLASH27B_KV_Q4=1 python3 scripts/server_tools.py \\
+    DFLASH27B_KV_TQ3=1 python3 scripts/server_tools.py \\
         --budget 22 --max-ctx 16384 --port 8000
 
 Then:

dflash/scripts/run.py

@@ -64,6 +64,8 @@ def main():
     ap.add_argument("--system", type=str, default=None)
     ap.add_argument("--kv-q4", action="store_true",
                     help="Q4_0 KV cache (required for max_ctx=131072)")
+    ap.add_argument("--kv-tq3", action="store_true",
+                    help="TQ3_0 KV cache (3.5 bpv, near-lossless)")
     ap.add_argument("--max-ctx", type=int, default=0,
                     help="Override max KV context (default: auto-fit "
                          "prompt+n_gen+block, aligned to 256). Passing a "
@@ -102,6 +104,8 @@ def main():
         env["PATH"] = dll_dir + os.pathsep + bin_dir + os.pathsep + env.get("PATH", "")
     if args.kv_q4:
         env["DFLASH27B_KV_Q4"] = "1"
+    if args.kv_tq3:
+        env["DFLASH27B_KV_TQ3"] = "1"
 
     with tempfile.TemporaryDirectory() as tmp:
         in_bin  = os.path.join(tmp, "prompt.bin")

dflash/scripts/server.py

@@ -446,19 +446,19 @@ def main():
                          "can slow attention 20×+ until issue #10 is fixed)")
     ap.add_argument("--kv-f16", action="store_true",
                     help="Force F16 KV cache. When --max-ctx > 6144 the server "
-                         "auto-enables Q4 KV to fit; pass --kv-f16 to opt out.")
+                         "auto-enables TQ3_0 KV to fit; pass --kv-f16 to opt out.")
     ap.add_argument("--tokenizer", type=str, default=None,
                     help="HuggingFace tokenizer repo ID (default: auto-detect "
                          "from target GGUF basename; falls back to Qwen/Qwen3.5-27B)")
     ap.add_argument("--daemon", action="store_true", help="Run with persistent model daemon (now default)")
     args = ap.parse_args()
 
-    # Auto-enable Q4 KV cache when the requested context exceeds what F16 fits.
+    # Auto-enable TQ3_0 KV cache when the requested context exceeds what F16 fits.
     # Clients like Claude Code routinely send 10k+ token system prompts, so
     # 6144 is too tight for real-world use. setdefault so an explicit user
-    # DFLASH27B_KV_Q4=0 still wins.
+    # DFLASH27B_KV_TQ3=0 still wins.
     if args.max_ctx > 6144 and not args.kv_f16:
-        os.environ.setdefault("DFLASH27B_KV_Q4", "1")
+        os.environ.setdefault("DFLASH27B_KV_TQ3", "1")
 
     if not args.bin.is_file():
         raise SystemExit(f"binary not found at {args.bin}")

dflash/scripts/server_tools.py

@@ -806,15 +806,15 @@ def main():
                          "attention 20×+ until issue #10 is fixed)")
     ap.add_argument("--kv-f16", action="store_true",
                     help="Force F16 KV cache. When --max-ctx > 6144 the server "
-                         "auto-enables Q4 KV to fit; pass --kv-f16 to opt out.")
+                         "auto-enables TQ3_0 KV to fit; pass --kv-f16 to opt out.")
     ap.add_argument("--tokenizer", default="Qwen/Qwen3.5-27B",
                     help="HF tokenizer id; Qwen3.6 shares this tokenizer.")
     args = ap.parse_args()
 
-    # Auto-enable Q4 KV cache when the requested context exceeds what F16 fits.
-    # setdefault so an explicit user DFLASH27B_KV_Q4=0 still wins.
+    # Auto-enable TQ3_0 KV cache when the requested context exceeds what F16 fits.
+    # setdefault so an explicit user DFLASH27B_KV_TQ3=0 still wins.
     if args.max_ctx > 6144 and not args.kv_f16:
-        os.environ.setdefault("DFLASH27B_KV_Q4", "1")
+        os.environ.setdefault("DFLASH27B_KV_TQ3", "1")
 
     if not args.bin.is_file():
         raise SystemExit(f"binary not found at {args.bin}")

dflash/src/internal.h

@@ -183,6 +183,8 @@ struct TargetCache {
     int max_ctx  = 0;         // max tokens in the KV cache
     int cur_pos  = 0;         // number of tokens already committed
 
+    ggml_type kv_k_type = GGML_TYPE_Q8_0;
+
     // Full-attention KV cache: one K and one V per full-attention layer.
     // Layout: [head_dim, max_ctx, n_head_kv] f16, contiguous per layer.
     std::vector<ggml_tensor *> attn_k;   // size = n_full_attn_layers (16)

dflash/src/qwen35_target_graph.cpp

@@ -110,6 +110,7 @@ bool create_target_cache(const TargetWeights & w,
     // Env overrides (checked in order; last wins):
     //   DFLASH27B_KV_F16=1  → f16 (regression baseline)
     //   DFLASH27B_KV_Q4=1   → Q4_0 (8× vs f16, required for 128K on 24 GB, ~3% AL hit)
+    //   DFLASH27B_KV_TQ3=1  → TQ3_0 (3.5 bpv, near-lossless, enables 128K on 24 GB)
     //
     // Default: Q8_0 — best quality/memory tradeoff at short context.
     ggml_type kv_k_type = GGML_TYPE_Q8_0;
@@ -120,15 +121,22 @@ bool create_target_cache(const TargetWeights & w,
     if (const char * s = std::getenv("DFLASH27B_KV_Q4")) {
         if (std::atoi(s) != 0) { kv_k_type = GGML_TYPE_Q4_0; kv_v_type = GGML_TYPE_Q4_0; }
     }
+    if (const char * s = std::getenv("DFLASH27B_KV_TQ3")) {
+        if (std::atoi(s) != 0) { kv_k_type = GGML_TYPE_TQ3_0; kv_v_type = GGML_TYPE_TQ3_0; }
+    }
+    out.kv_k_type = kv_k_type;
+    const int max_ctx_alloc = (kv_k_type == GGML_TYPE_TQ3_0)
+        ? ((max_ctx + 255) / 256) * 256
+        : max_ctx;
     int fa_idx = 0, dn_idx = 0;
     for (int il = 0; il < w.n_layer; il++) {
         const bool is_attn = (((il + 1) % w.full_attention_interval) == 0);
         if (is_attn) {
-            // [head_dim, max_ctx, n_head_kv]
+            // [head_dim, max_ctx_alloc, n_head_kv]
             ggml_tensor * K = ggml_new_tensor_3d(out.ctx, kv_k_type,
-                                                 q35::HEAD_DIM, max_ctx, q35::N_HEAD_KV);
+                                                 q35::HEAD_DIM, max_ctx_alloc, q35::N_HEAD_KV);
             ggml_tensor * V = ggml_new_tensor_3d(out.ctx, kv_v_type,
-                                                 q35::HEAD_DIM, max_ctx, q35::N_HEAD_KV);
+                                                 q35::HEAD_DIM, max_ctx_alloc, q35::N_HEAD_KV);
             char name[64];
             std::snprintf(name, sizeof(name), "cache_k_%d", il);
             ggml_set_name(K, name);
@@ -282,14 +290,15 @@ static ggml_tensor * build_full_attn_block(
     ggml_context * ctx,
     ggml_cgraph * gf,
     const TargetLayer & L,
-    ggml_tensor * cur,              // [hidden, n_tokens]
-    ggml_tensor * positions,        // [n_tokens] i32
+    ggml_tensor * cur,
+    ggml_tensor * positions,
     const int * rope_sections,
-    ggml_tensor * cache_k,          // [head_dim, max_ctx, n_head_kv]
-    ggml_tensor * cache_v,          // [head_dim, max_ctx, n_head_kv]
-    ggml_tensor * attn_mask,        // [kv_len, n_tokens] f32 or nullptr
+    ggml_tensor * cache_k,
+    ggml_tensor * cache_v,
+    ggml_tensor * attn_mask,
     int kv_start,
-    int n_tokens
+    int n_tokens,
+    ggml_type kv_k_type
 ) {
     // ── Q projection (packed Q || gate), shape [2*q_dim, n_tokens]
     ggml_tensor * QG = ggml_mul_mat(ctx, L.wq, cur);
@@ -369,13 +378,20 @@ static ggml_tensor * build_full_attn_block(
     // FATTN_KQ_STRIDE=256 (see fattn.cu:get_best_fattn_kernel). Round up
     // for TBQ cache types; the caller's attn_mask is built with the same
     // padded length so positions beyond the real kv_len get -inf.
-    const int fattn_stride  = 1;
+    const int fattn_stride  = (kv_k_type == GGML_TYPE_TQ3_0) ? 256 : 1;
     const int kv_len_padded = ((kv_len + fattn_stride - 1) / fattn_stride) * fattn_stride;
 
     // Q needs to be [head_dim, n_tokens, n_head] for flash_attn_ext
     ggml_tensor * Qfa = ggml_permute(ctx, Q, 0, 2, 1, 3);   // [head_dim, n_tokens, n_head]
     Qfa = ggml_cont(ctx, Qfa);
 
+    // For TQ3_0 KV cache, K/V are stored in FWHT-rotated space.
+    // Rotate Q to match before computing KQ dot product.
+    const bool needs_rotation = (kv_k_type == GGML_TYPE_TQ3_0);
+    if (needs_rotation) {
+        Qfa = ggml_turbo_wht(ctx, Qfa, 0);
+    }
+
     // K and V from cache: a view into the first kv_len_padded slots. For
     // non-TBQ paths kv_len_padded == kv_len so this is identical to the
     // old behaviour.
@@ -394,6 +410,13 @@ static ggml_tensor * build_full_attn_block(
     ggml_tensor * attn = ggml_flash_attn_ext(ctx, Qfa, Kfa, Vfa, attn_mask,
                                              kq_scale, 0.0f, 0.0f);
     // attn: [head_dim, n_head, n_tokens] (permuted)
+
+    // Un-rotate the FA output from FWHT-rotated V space.
+    if (needs_rotation) {
+        attn = ggml_cont(ctx, attn);
+        attn = ggml_turbo_wht(ctx, attn, 1);
+    }
+
     attn = ggml_reshape_2d(ctx, attn, q35::Q_DIM, n_tokens);
 
     // ── Apply the sigmoid gate from the packed Q
@@ -709,7 +732,8 @@ QwenGraphOutputs build_qwen35_graph(
         if (is_attn) {
             cur = build_full_attn_block(ctx, gf, L, cur, in.positions, w.rope_sections,
                                         cache.attn_k[fa_idx], cache.attn_v[fa_idx],
-                                        in.attn_mask, in.kv_start, n_tokens);
+                                        in.attn_mask, in.kv_start, n_tokens,
+                                        cache.kv_k_type);
             fa_idx++;
         } else {
             DeltaNetCapture * cap_ptr = nullptr;

dflash/test/test_dflash.cpp

@@ -721,6 +721,9 @@ int main(int argc, char ** argv) {
     if (const char * s = std::getenv("DFLASH27B_KV_TBQ")) {
         if (std::atoi(s) != 0) g_kq_stride_pad = 256;
     }
+    if (const char * s = std::getenv("DFLASH27B_KV_TQ3")) {
+        if (std::atoi(s) != 0) g_kq_stride_pad = 256;
+    }
     const char * target_path = argv[1];
     const char * draft_path  = argv[2];
     const char * prompt_path = (argc >= 6 && argv[3][0] != '-') ? argv[3] : nullptr;