|
| 1 | +# K3 GPU beta — Kakeya inference (f_θ + S5 K/V-Restoration) |
| 2 | + |
| 3 | +Status: beta, GPU-validated on NVIDIA H200 with `google/gemma-4-26B-A4B-it` |
| 4 | +(verifier) + `z-lab/gemma-4-26B-A4B-it-DFlash` (drafter) + the trained f_θ v5 |
| 5 | +checkpoint (`results/research/f_theta_v5_s5_sliding/`). Recall 1.0 throughout. |
| 6 | + |
| 7 | +## What it is |
| 8 | + |
| 9 | +The verifier keeps only a **sink+window** local KV cache; at every *evicted* |
| 10 | +position its attention reads **reconstructed** K/V, so it attends over the full |
| 11 | +context while holding `O(sink+window)` resident KV (ADR 0008 §11). |
| 12 | + |
| 13 | + verifier (Gemma 4 26B-A4B): sink+window resident KV |
| 14 | + ├─ sliding layers → evicted K/V restored via f_θ(drafter K/V) |
| 15 | + └─ full-attn layers (S5: [5,11,17,23,29]) → verifier's OWN exact K/V |
| 16 | + (recall-critical; f_θ cannot reconstruct these — |
| 17 | + proven by the α-sweep, eval rel_mse floor ~1.4) |
| 18 | + |
| 19 | + drafter (DFlash 0.4B): no KV cache; constant-memory K/V reconstruction |
| 20 | + source (its K/V are projected into verifier space by f_θ). |
| 21 | + |
| 22 | +## Components (this branch) |
| 23 | + |
| 24 | +| piece | file | |
| 25 | +|---|---| |
| 26 | +| DFlash drafter (block diffusion, faithful to z-lab `qwen3_dflash`) | `inference_engine/v04/dflash_drafter.py` | |
| 27 | +| f_θ projection (drafter K/V → verifier K/V) | `inference_engine/v04/f_theta.py` | |
| 28 | +| Cross-model restored verifier (CUDA) + S5 | `inference_engine/v04/cross_model_dlm_verifier.py` | |
| 29 | +| Cross-model restored verifier (MLX / Apple Silicon) | `inference_engine/backends/mlx/cross_model_dlm_verifier.py` | |
| 30 | +| Incremental restored verifier (`SinkWindowVerifier` API) | `inference_engine/v04/restored_sink_window_verifier.py` | |
| 31 | +| Served-path factories + gRPC `--backend restored` | `inference_engine/v04/build_restored.py`, `scripts/start_grpc_runtime_server.py` | |
| 32 | + |
| 33 | +## Three engines (decode modes) |
| 34 | + |
| 35 | +* **Re-forward** (`incremental=False`) — memory-optimal, eval-grade; recomputes |
| 36 | + restoration each step (O(T)/step). Bit-equivalent reference for the gate. |
| 37 | +* **Gap-A incremental** (`incremental=True`) — capture restored K/V into a |
| 38 | + `DynamicCache` at prefill, decode natively (O(L)/block). **= AR decode speed**, |
| 39 | + KV 16.9×–43.9× smaller, recall 1.0. |
| 40 | +* **Fused spec-decode** (`restored_specdecode_fused`) — DFlash block draft + |
| 41 | + incremental verify, with three prefill-built, incrementally-extended caches: |
| 42 | + (A) verifier aux hidden captured from the verify forward, (B) drafter context |
| 43 | + K/V cache, (C) Gap-A restored KV. Per-block O(L). **> AR** (see below). |
| 44 | + |
| 45 | +## Validated results (H200, ctx 1238, gemma-4-26B-A4B) |
| 46 | + |
| 47 | +| path | decode tok/s | vs AR | recall | |
| 48 | +|---|---|---|---| |
| 49 | +| standalone AR | 21.1 | 1.0× | 1.0 | |
| 50 | +| Gap-A incremental restored | 21.7 | 1.03× | 1.0 | |
| 51 | +| fused DFlash spec-decode (aggregate) | 26.8 | **1.27×** | 1.0 | |
| 52 | + |
| 53 | +KV memory: restored resident KV constant **16.71 MB** vs AR 282 MB @1238 tok → |
| 54 | +733 MB @3238 tok (**16.9× → 43.9×**, grows with context). DFlash acceptance on |
| 55 | +HumanEval ≈ official gemma-4-26B parity (length ~3.9 ≈ official 3.3× speedup). |
| 56 | + |
| 57 | +## Run |
| 58 | + |
| 59 | +```bash |
| 60 | +# Incremental restored decode vs AR (memory + tok/s + recall) |
| 61 | +PYTHONPATH=.:sdks/python python scripts/research/k3_e2e_gpu_bench.py \ |
| 62 | + --verifier-id google/gemma-4-26B-A4B-it \ |
| 63 | + --drafter-id z-lab/gemma-4-26B-A4B-it-DFlash \ |
| 64 | + --f-theta-dir results/research/f_theta_v5_s5_sliding \ |
| 65 | + --incremental --haystack-lines 60,160 |
| 66 | + |
| 67 | +# Fused DFlash spec-decode vs AR |
| 68 | +PYTHONPATH=.:sdks/python python scripts/research/k3_specdecode_gpu_bench.py \ |
| 69 | + --drafter-id z-lab/gemma-4-26B-A4B-it-DFlash --skip-unfused |
| 70 | + |
| 71 | +# gRPC server with the restored backend |
| 72 | +PYTHONPATH=.:sdks/python python scripts/start_grpc_runtime_server.py \ |
| 73 | + --backend restored --device cuda \ |
| 74 | + --verifier-id google/gemma-4-26B-A4B-it \ |
| 75 | + --drafter-id z-lab/gemma-4-26B-A4B-it-DFlash \ |
| 76 | + --f-theta-dir results/research/f_theta_v5_s5_sliding --sink 4 --window 64 |
| 77 | +``` |
| 78 | + |
| 79 | +## Notes / scope |
| 80 | + |
| 81 | +* Drafting conditions on the restored verifier hidden for committed decode tokens |
| 82 | + (clean aux for the prompt) — resolves the bounded-KV vs clean-aux tension |
| 83 | + natively; no SGLang/vLLM dependency. |
| 84 | +* Stable decode requires loading the verifier without `device_map` (no accelerate |
| 85 | + per-forward hooks; the 26B-A4B fits on one H200) + a full-length warmup. |
| 86 | +* f_θ v5 restores the sliding layers; recall is carried by the S5 exact |
| 87 | + full-attention layers, so f_θ fidelity is not the recall bottleneck. |
0 commit comments