Skip to content

Commit b909142

Browse files
Trim beta + add architecture note
- Trim: drop research f_theta v1/v3/v4 checkpoints (+ reports, ~964MB LFS) from the merge; keep v5 (the validated S5 checkpoint) + engine code + small GPU evidence JSONs. - Add docs/k3-gpu-beta.md: short architecture note (verifier sink+window + f_theta/S5 restored evicted K/V; DFlash drafter; three decode modes; H200 results AR=1.0/Gap-A=1.03x/fused=1.27x, KV 16.9-43.9x, recall 1.0; run commands). Co-authored-by: FluffyAIcode <FluffyAIcode@users.noreply.github.com>
1 parent a53abc0 commit b909142

13 files changed

Lines changed: 87 additions & 762 deletions

File tree

docs/k3-gpu-beta.md

Lines changed: 87 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,87 @@
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.

results/research/f_theta_v1.json

Lines changed: 0 additions & 98 deletions
This file was deleted.

results/research/f_theta_v1/f_theta_config.json

Lines changed: 0 additions & 73 deletions
This file was deleted.

results/research/f_theta_v1/f_theta_weights.pt

Lines changed: 0 additions & 3 deletions
This file was deleted.

results/research/f_theta_v3_attn_distill.json

Lines changed: 0 additions & 114 deletions
This file was deleted.

0 commit comments

Comments
 (0)