Skip to content

Commit 69751d6

Browse files
cursoragentFluffyAIcode
cursoragent
and
FluffyAIcode
committed
release(v0.5-cuda): KakeyaVLLM entrypoint — Kakeya Attention on vLLM (fused-MoE+CUDA-graph+scheduler)
Package KIE-v2 as the product entrypoint inference_engine.engine.KakeyaVLLM: Kakeya Attention's bounded-window (S5) KV management running ON vLLM's Apache-2.0 runtime (fused-MoE Triton kernel + CUDA graphs + continuous-batching scheduler inherited unchanged). - engine/kakeya_vllm.py: KakeyaVLLM wrapper + KakeyaVLLMConfig + pure kakeya_hf_overrides/kakeya_window_total helpers. Model-aware text_config nesting (gemma-4 nested / Qwen-Llama flat) — fixes a crash where unconditional text_config injection broke text-only models. - validated end-to-end on H200 (vLLM 0.23.0): CUDA graphs captured, Kakeya window reaches vLLM config, coherent generation, 777 tok/s (Qwen3-4B). - 13 torch/vllm-free config unit tests. - docs/reports/kakeya-inference-engine-v0.5-cuda.md: consolidated scorecard (token throughput >= vLLM 1.15-1.23x, parallelism N=70/75, honest memory saving: gemma-4 ~7%, full-attn ~6x via v0.6 restoration). - README v0.5-for-CUDA section + release badge; ADR 0015 milestone row. Co-authored-by: FluffyAIcode <FluffyAIcode@users.noreply.github.com>
1 parent e2cf137 commit 69751d6

6 files changed

Lines changed: 469 additions & 1 deletion

File tree

‎README.md‎

Lines changed: 44 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
# Kakeya Inference engine — memory bounded local inference engine
22

33
[![CI](https://github.com/FluffyAIcode/Kakeya-LLM-Inference-engine/actions/workflows/ci.yaml/badge.svg?branch=main)](https://github.com/FluffyAIcode/Kakeya-LLM-Inference-engine/actions/workflows/ci.yaml)
4-
[![Release](https://img.shields.io/badge/release-v0.4-blue)](https://github.com/FluffyAIcode/Kakeya-LLM-Inference-engine/tags)
4+
[![Release](https://img.shields.io/badge/release-v0.5--cuda%20%7C%20v0.4-blue)](https://github.com/FluffyAIcode/Kakeya-LLM-Inference-engine/tags)
55
[![Platform](https://img.shields.io/badge/platform-Apple%20Silicon%20(MLX)%20%7C%20CUDA%20%7C%20Linux--CPU-lightgrey)](docs/quickstart.md)
66
[![Architecture](https://img.shields.io/badge/architecture-ADR%200008%20%7C%200014-green)](docs/adr/0014-agent-connection-capacity-and-cross-host-topology-tests.md)
77
[![License](https://img.shields.io/badge/license-MIT-lightblue)](LICENSE)
@@ -355,6 +355,49 @@ control + tool plane** (the Mac bridge) and **LAN = co-located data plane**. See
355355
[ADR 0014](docs/adr/0014-agent-connection-capacity-and-cross-host-topology-tests.md)
356356
for the full plan, evidence, and the served-MLX-gemma gap found during testing.
357357

358+
## v0.5 for CUDA — Kakeya Attention on the vLLM runtime
359+
360+
**v0.5-cuda** ships Kakeya Attention's bounded-window (S5) KV management **on top
361+
of the vLLM runtime**, so the three runtime components the engine needs are
362+
inherited unchanged — all **Apache-2.0**:
363+
364+
| component | owner in v0.5-cuda | role |
365+
| --- | --- | --- |
366+
| **Fused MoE Triton kernel** | **vLLM** | grouped-GEMM expert kernel — the dominant ~90 % of the gemma-4-26B-A4B decode forward |
367+
| **CUDA graphs** | **vLLM** | fixed-shape decode capture (`enforce_eager=False`) — removes per-token launch overhead |
368+
| **Continuous-batching scheduler** | **vLLM** | request scheduler + paged KV-manager — drives multi-tenant throughput |
369+
| **Kakeya Attention (bounded window / KV)** | **Kakeya** | bounds resident sliding-layer KV to `sink + window` (S5 = 68) |
370+
371+
This is the **KIE-v2** strategy ([ADR 0015](docs/adr/0015-kakeya-attention-and-engine-substrate.md),
372+
[feasibility](docs/design/kakeya-vllm-backend-feasibility.md)): rather than rebuild
373+
vLLM's fused-MoE + graphs + scheduler (shown blocked in the eager KIE-v1.1.z
374+
attempt), Kakeya Attention runs *on* vLLM and contributes the bounded-KV layer.
375+
376+
```python
377+
from inference_engine.engine import KakeyaVLLM
378+
from vllm import SamplingParams
379+
380+
engine = KakeyaVLLM("google/gemma-4-26b-a4b-it", sink=4, window=64, max_model_len=16384)
381+
out = engine.generate(prompts, SamplingParams(temperature=0.0, max_tokens=128))
382+
```
383+
384+
**Scorecard** ([full report](docs/reports/kakeya-inference-engine-v0.5-cuda.md)),
385+
H200, gemma-4-26B-A4B, recall **1.0**:
386+
387+
| axis | result |
388+
| --- | --- |
389+
| **Token throughput (decode)** | **≥ vLLM**: **1.15–1.23×** vs vLLM default at ctx 16k, N=1..70 (e.g. N=70 **1079 vs 894.9 tok/s**) |
390+
| **Parallel inference** | bounded window on vLLM measured to **N=70** @16k (recall 1.0); eager research engine reached **N=75 @62k** (≈4.8× vLLM concurrency) |
391+
| **Memory saving** | gemma-4 hybrid: **~7 % @ 62k** (vLLM already bounds 25/30 layers; the 5 full layers dominate both); **~6×** edge needs a **full-attention** model + the v0.6 restoration backend |
392+
393+
> **Honest scope.** v0.5-cuda is the **gemma-4 bounded-window** instantiation
394+
> (gemma-4's hybrid needs **no per-token restoration** — the S5 "free lunch" —
395+
> delivered via vLLM `hf_overrides`). The `KakeyaVLLM` wrapper itself was validated
396+
> end-to-end on an H200 (CUDA graphs captured, window applied, coherent generation,
397+
> 777 tok/s on Qwen3-4B — the model that fit the box). The **restoration backend**
398+
> (f_θ + dLLM-proposer at prefill + quantized-exact attention) for **full-attention**
399+
> models — the large ~6× memory differentiator — is the **v0.6** roadmap item.
400+
358401
## v0.4 for Mac — MLX speculative-decode port (the journey to parity)
359402

360403
After the **CUDA** path (f_θ + S5 K/V-restoration verifier, **fused DFlash

‎docs/adr/0015-kakeya-attention-and-engine-substrate.md‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -146,6 +146,7 @@ PR, so development context stays per-task:
146146
| **KIE-v1.1.z** | throughput + N=75 | **N=75 MET** (recall 1.0, 126.7 GB, ≈4.8× vLLM; int8+quant-attn). kakeyalattice v1.6.1 gemma-4 fix verified (recall 1.0). **decode ≥ vLLM NOT met**: ~31 tok/s aggregate (eager 26B-MoE forward dominates; torch.compile-attn 6.6× but 0% e2e; static-cache auto-compile segfaults). | #139 |
147147
| **KIE-v1.1.z2** | fused-MoE + graph-captured full-forward decode (rebuild vLLM's runtime) | **abandoned — superseded by KIE-v2** (run on vLLM instead of rebuilding it) | — |
148148
| **KIE-v2** | **Kakeya Attention on the vLLM runtime** (bounded window + restoration + quantized-exact as a vLLM backend) | **decode-parity MET + EXCEEDED**: gemma-4 bounded-window (sw=68) on vLLM = **195.6/231.9/539/1079 tok/s vs vLLM-default 159.3/198.6/467.5/894.9 @N=1/4/8/70, ~1.15–1.23× faster, recall 1.0** (ctx16k). Full backend (restoration + quantized-exact for full-attn models) = next. | #140 |
149+
| **v0.5-cuda** | **release**: package KIE-v2 as `inference_engine.engine.KakeyaVLLM` (Kakeya window on vLLM's Apache-2.0 fused-MoE + CUDA-graph + scheduler), consolidate reports, README | **done** — entrypoint validated end-to-end on H200 (CUDA graphs captured, window→vLLM config, coherent generation, 777 tok/s on Qwen3-4B); model-aware `text_config` nesting (gemma nested / Qwen flat); 13 config unit tests. Scorecard: `docs/reports/kakeya-inference-engine-v0.5-cuda.md`. | this PR |
149150

150151
**KIE-v1.1.z2 port attempt (2026-06-16) — all clean paths blocked:**
151152
- **HF `kernels` hub** (lowest-surgery kernelize): `kernels` is **version-incompatible

‎docs/reports/kakeya-inference-engine-v0.5-cuda.md‎

Lines changed: 144 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,144 @@
1+
# Kakeya Inference Engine v0.5 for CUDA — release scorecard
2+
3+
**What v0.5-cuda ships:** Kakeya Attention's bounded-window (S5) KV management
4+
running **on the vLLM runtime** — so the three runtime components the engine
5+
needs are inherited unchanged (all Apache-2.0):
6+
7+
| component | owner in v0.5-cuda | role |
8+
| --- | --- | --- |
9+
| **Fused MoE Triton kernel** | **vLLM** (Apache-2.0) | grouped-GEMM expert kernel — the dominant ~90 % of the gemma-4-26B-A4B decode forward |
10+
| **CUDA graphs** | **vLLM** (Apache-2.0) | fixed-shape decode capture (`enforce_eager=False`) — removes per-token launch overhead |
11+
| **Continuous-batching scheduler** | **vLLM** (Apache-2.0) | request scheduler + paged KV-manager — drives multi-tenant throughput |
12+
| **Kakeya Attention (bounded window / KV)** | **Kakeya** | bounds the resident sliding-layer KV to `sink + window` (S5 = 68) |
13+
14+
Entrypoint: `inference_engine.engine.KakeyaVLLM` (see *Usage* below). This is the
15+
**KIE-v2** strategy from [ADR 0015](../adr/0015-kakeya-attention-and-engine-substrate.md)
16+
and the [feasibility note](../design/kakeya-vllm-backend-feasibility.md): rather
17+
than rebuild vLLM's fused-MoE + graphs + scheduler (a multi-component kernel
18+
project that was shown blocked in KIE-v1.1.z), Kakeya Attention runs *on* vLLM
19+
and contributes the bounded-KV attention layer.
20+
21+
> **Honest scope.** v0.5-cuda is the **gemma-4 bounded-window** instantiation:
22+
> gemma-4's hybrid (5 full + 25 sliding) needs **no per-token restoration** (the
23+
> S5 "free lunch"), so the bounded window is delivered via vLLM `hf_overrides`.
24+
> The **restoration backend** (f_θ + dLLM-proposer at prefill) for **full-attention**
25+
> models (Qwen/Llama) — the large ~6× memory differentiator — is the **v0.6**
26+
> roadmap item, not in this release.
27+
28+
---
29+
30+
## 1. Token throughput — decode tok/s ≥ vLLM (H200, gemma-4-26B-A4B, recall 1.0)
31+
32+
Kakeya bounded window (`sliding_window=68`) **on vLLM** vs vLLM default
33+
(`sliding_window=1024`), same H200, ctx 16k, gen 128, recall **1.0** throughout:
34+
35+
| N (concurrent) | **Kakeya-on-vLLM** decode tok/s | vLLM default | ratio |
36+
| --- | --- | --- | --- |
37+
| 1 | **195.6** | 159.3 | **1.23×** |
38+
| 4 | **231.9** | 198.6 | 1.17× |
39+
| 8 | **539.0** | 467.5 | 1.15× |
40+
| 70 | **1079.0** | 894.9 | **1.21×** |
41+
42+
**Decode throughput exceeds vLLM by ~1.15–1.23×** at recall 1.0 — it inherits
43+
vLLM's fused-MoE + CUDA-graph + scheduler, and Kakeya's tighter sliding window
44+
makes the sliding-layer attention cheaper than vLLM's default. This is the axis
45+
the eager research engine (KIE-v1.1.y/z, ~25–31 tok/s aggregate) could not reach;
46+
on vLLM it is 195–1079 tok/s.
47+
48+
**Long context (ctx 62k, N=70) — the edge shrinks on gemma-4:**
49+
50+
| metric (62k, N=70) | Kakeya-on-vLLM sw=68 | vLLM default sw=1024 | ratio |
51+
| --- | --- | --- | --- |
52+
| decode tok/s | 20.38 | 19.03 | 1.07× |
53+
| vLLM max concurrency (66k req) | 16.15× | 15.51× | 1.04× |
54+
55+
The edge falls from ~1.2× (16k) to ~1.07× (62k) because gemma-4's **5
56+
full-attention layers hold full-ctx KV in both configs** and dominate the
57+
footprint — shrinking the sliding window is only a ~4–7 % saving. This is the
58+
gemma-4 caveat (see §3); the large win needs a full-attention model.
59+
60+
## 2. Parallel inference (concurrency)
61+
62+
- **On vLLM (v0.5-cuda path):** the bounded window is measured to **N=70**
63+
concurrent sessions at ctx 16k, recall 1.0, while *increasing* decode tok/s vs
64+
vLLM default (§1) — vLLM's continuous-batching scheduler + the smaller resident
65+
window scale cleanly.
66+
- **Research engine ceiling (KIE-v1.1.y/z, eager, demonstrator):** the bounded-KV
67+
+ int8 + quantized-attention path reached **N=75 @ 62k, recall 1.0** (≈4.8×
68+
vLLM's 15.5 concurrency ceiling) — proving the *memory/concurrency* axis even
69+
before moving to the vLLM runtime. Decode speed on that eager path was the weak
70+
axis (~31 tok/s aggregate), which is exactly what running on vLLM fixes.
71+
72+
So v0.5-cuda gets **both** axes: vLLM's decode speed **and** the bounded-window
73+
concurrency, with recall 1.0.
74+
75+
## 3. Memory-saving efficiency (honest, model-dependent)
76+
77+
The bounded-KV win scales with the model's **full-attention fraction**:
78+
79+
| model class | full-attn layers | Kakeya resident-KV edge vs vLLM | shipped in |
80+
| --- | --- | --- | --- |
81+
| **gemma-4-26B-A4B** (hybrid 5 full / 25 sliding) | 5 / 30 | **~7 % @ 62k** (vLLM already hybrid-bounds 25/30; the 5 full layers dominate both) | **v0.5-cuda** |
82+
| **full-attention** (Qwen/Llama, all layers full) | all | **~6×** (vLLM keeps all full; Kakeya keeps exact + window + restoration) | **v0.6** (restoration backend) |
83+
84+
- The bounded-KV cost model (`inference_engine.engine.admission`, 9 unit tests):
85+
per-session resident **2.56 GB @ 62k** vs full-KV **15.2 GB** — a **~6×** edge —
86+
but that edge is only realized end-to-end on a **full-attention** model. On
87+
gemma-4 the native hybrid means vLLM already bounds the sliding layers, so the
88+
*measured* long-context saving over vLLM is the honest **~7 %**.
89+
- **Takeaway:** v0.5-cuda's memory win on gemma-4 is modest-but-real; the engine's
90+
large memory differentiation is a **full-attention-model** property delivered by
91+
the v0.6 restoration backend. We report this rather than overclaim gemma-4.
92+
93+
## 4. Usage
94+
95+
```python
96+
from inference_engine.engine import KakeyaVLLM
97+
from vllm import SamplingParams
98+
99+
# Kakeya Attention (S5 bounded window) on vLLM's fused-MoE + CUDA-graph + scheduler.
100+
engine = KakeyaVLLM(
101+
"google/gemma-4-26b-a4b-it",
102+
sink=4, window=64, # Kakeya S5 window (total resident = 68)
103+
max_model_len=16384,
104+
)
105+
out = engine.generate(prompts, SamplingParams(temperature=0.0, max_tokens=128))
106+
```
107+
108+
`KakeyaVLLM` builds a `vllm.LLM` with `hf_overrides={"sliding_window": 68,
109+
"text_config": {"sliding_window": 68}}` and `enforce_eager=False` (CUDA graphs +
110+
fused-MoE on). The pure config layer (`kakeya_hf_overrides`, `KakeyaVLLMConfig`)
111+
is torch/vllm-free and unit-tested (`tests/inference_engine/engine/test_kakeya_vllm.py`).
112+
113+
## 5. Verification status
114+
115+
- ✅ **Throughput / concurrency / recall numbers** (gemma-4-26B) above were measured
116+
on H200 (Vast.ai) and committed in the KIE-v2 integration
117+
(`scripts/research/vllm_multitenant_parallel_bench.py --sliding-window 68`;
118+
commits `7ec3a03`, `48ded1e`, `e2cf137`).
119+
- ✅ **`KakeyaVLLM` entrypoint validated end-to-end on H200** (vLLM 0.23.0): it builds
120+
the vLLM engine with **CUDA graphs captured** (PIECEWISE + FULL), the Kakeya window
121+
(68) **reaches vLLM's model config** (`hf_config.sliding_window == 68`), and
122+
generation is coherent (Paris / 2+2=4 / story) at **777 tok/s** (batch 3). The
123+
validation used **Qwen/Qwen3-4B** — the 26B model does not fit this box's 4 GB free
124+
disk, so the wrapper *mechanism* is validated here and the 26B *performance* is the
125+
committed measured path above. Evidence:
126+
`kakeya_vllm_v05_h200_validation.log`.
127+
- ✅ The wrapper auto-detects `text_config` nesting (multimodal gemma-4 → nested;
128+
text-only Qwen/Llama → flat), fixing a crash where unconditional `text_config`
129+
injection broke text-only models. Config layer unit-tested (13 tests).
130+
131+
## 6. What's next (v0.6)
132+
133+
The **restoration backend** for full-attention models (Qwen/Llama): inject f_θ +
134+
dLLM-proposer restoration at vLLM prefill and a graph-capturable quantized-exact
135+
attention kernel, to realize the **~6×** resident-KV edge end-to-end on vLLM. That
136+
is the genuine custom-backend work (ADR 0015 §KIE-v2 caveats); v0.5-cuda proves
137+
the throughput axis and ships the gemma-4 bounded-window engine.
138+
139+
## Evidence
140+
141+
- Throughput tables: [`docs/design/kakeya-vllm-backend-feasibility.md`](../design/kakeya-vllm-backend-feasibility.md) §5b
142+
- Concurrency / memory journey: [`docs/reports/kakeya-engine-vs-vllm-h200.md`](kakeya-engine-vs-vllm-h200.md)
143+
- Architecture / milestones: [`docs/adr/0015-kakeya-attention-and-engine-substrate.md`](../adr/0015-kakeya-attention-and-engine-substrate.md)
144+
- Entrypoint: `inference_engine/engine/kakeya_vllm.py`; tests: `tests/inference_engine/engine/test_kakeya_vllm.py`

‎inference_engine/engine/__init__.py‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,11 +19,23 @@
1919
# kakeya_engine imports torch only lazily (inside methods), so this is safe on
2020
# torch-less hosts; the class is the engine runtime entry point.
2121
from inference_engine.engine.kakeya_engine import KakeyaEngine
22+
# kakeya_vllm defers the vllm import to construction, so importing these names is
23+
# safe on torch/vllm-free hosts; KakeyaVLLM is the v0.5-CUDA product entrypoint.
24+
from inference_engine.engine.kakeya_vllm import (
25+
KakeyaVLLM,
26+
KakeyaVLLMConfig,
27+
kakeya_hf_overrides,
28+
kakeya_window_total,
29+
)
2230

2331
__all__ = [
2432
"BoundedKVModel",
2533
"resident_kv_bytes_per_session",
2634
"full_kv_bytes_per_session",
2735
"max_concurrent_sessions",
2836
"KakeyaEngine",
37+
"KakeyaVLLM",
38+
"KakeyaVLLMConfig",
39+
"kakeya_hf_overrides",
40+
"kakeya_window_total",
2941
]

0 commit comments

Comments
 (0)