diff --git a/README.md b/README.md index 97752a0dc..ca81897f8 100644 --- a/README.md +++ b/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) diff --git a/dflash/README.md b/dflash/README.md index eb8aa43d1..9d6d17c41 100644 --- a/dflash/README.md +++ b/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 diff --git a/dflash/deps/llama.cpp b/dflash/deps/llama.cpp index b16de6590..182346026 160000 --- a/dflash/deps/llama.cpp +++ b/dflash/deps/llama.cpp @@ -1 +1 @@ -Subproject commit b16de65904ed7e468397f5417ad130f092cba8f4 +Subproject commit 1823460262950917e7ddcf65040295b41a423c55 diff --git a/dflash/scripts/bench_daemon.py b/dflash/scripts/bench_daemon.py index a39be924a..314af3325 100644 --- a/dflash/scripts/bench_daemon.py +++ b/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: diff --git a/dflash/scripts/run.py b/dflash/scripts/run.py index 7bac16420..a3fe991cc 100644 --- a/dflash/scripts/run.py +++ b/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") diff --git a/dflash/scripts/server.py b/dflash/scripts/server.py index 4d136f06d..bcc50f1ea 100644 --- a/dflash/scripts/server.py +++ b/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}") diff --git a/dflash/scripts/server_tools.py b/dflash/scripts/server_tools.py index 53ba7ee5c..c5a273f35 100644 --- a/dflash/scripts/server_tools.py +++ b/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}") diff --git a/dflash/src/internal.h b/dflash/src/internal.h index 294531a76..f56718c67 100644 --- a/dflash/src/internal.h +++ b/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 attn_k; // size = n_full_attn_layers (16) diff --git a/dflash/src/qwen35_target_graph.cpp b/dflash/src/qwen35_target_graph.cpp index d19846bfe..9f1cb1557 100644 --- a/dflash/src/qwen35_target_graph.cpp +++ b/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; diff --git a/dflash/test/test_dflash.cpp b/dflash/test/test_dflash.cpp index a2a875259..50a938008 100644 --- a/dflash/test/test_dflash.cpp +++ b/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;