feat(engine): Kakeya Inference Engine v1 core (chunked restoration prefill + bounded-KV decode + peak-window admission)

Architecture design doc (docs/design/kakeya-inference-engine-architecture.md,
product terms only) + implementation:
- inference_engine/engine/admission.py: bounded-KV cost model + peak-window
  admission (pure; 9 unit tests).
- inference_engine/engine/kakeya_engine.py: KakeyaEngine runtime, NativeHybridBounded
  policy (chunked restoration prefill via prefill_chunk_size + sliding-window-bounded
  cache; bounded-KV decode), per ADR 0015.
- scripts/eval/kakeya_engine_throughput_eval.py: long-context concurrency/throughput
  evaluation driver.

Co-authored-by: FluffyAIcode <FluffyAIcode@users.noreply.github.com>

Author: cursoragent

docs/design/kakeya-inference-engine-architecture.md

@@ -0,0 +1,168 @@
+# Kakeya Inference Engine Architecture Design
+
+Status: Accepted (v1 core); Date: 2026-06-15. Governing decision: [ADR 0015](../adr/0015-kakeya-attention-and-engine-substrate.md).
+
+## 1. Purpose
+
+The Kakeya Inference Engine is a **product-grade LLM inference engine whose goal
+is to replace vLLM**. Its native, first-class attention algorithm is **Kakeya
+Attention** — sink+window bound + f_θ KV-projection + dLLM-proposer restoration,
+as one primitive. The engine is designed **bounded-KV-native**: the full token
+history is never resident; evicted context is reconstructed on demand. Every
+subsystem (prefill, KV layout, admission, decode) is built around that invariant,
+the way vLLM is built around full-KV PagedAttention.
+
+The engine **replaces** vLLM rather than extending it. Of vLLM's three
+prefill-engineering pieces:
+
+- **Chunked prefill** → adopted, reinterpreted as **chunked restoration**.
+- **FlashAttention** → wrapped and called as a kernel (table stakes).
+- **Paged KV** → **not used** — paging manages a growing full KV that this
+  engine never holds.
+
+## 2. Core invariant
+
+For a session of logical length `T`, the **resident KV is bounded** and
+independent of `T` beyond the exact-layer term:
+
+```
+resident_KV(session) = Σ_exact_layers (full T) + Σ_other_layers (sink + window)
+```
+
+- **Exact layers** keep full-context KV (the recall-critical layers).
+- **All other layers** keep only `sink + window` resident positions; positions
+  outside that window are **evicted** and, when a query needs them, **restored on
+  demand** by the restoration policy (§5).
+- The full per-layer KV for evicted positions is **never materialized or stored**.
+
+Memory is provisioned for the **peak resident window**, not the conversation
+length. This is the structural source of the engine's concurrency advantage over
+full-KV engines.
+
+## 3. Subsystems
+
+```
+        prompt stream                          token stream
+             │                                      ▲
+             ▼                                      │
+  ┌───────────────────────┐   bounded KV   ┌────────────────────┐
+  │ Chunked Restoration   │──────────────▶ │ Bounded-KV Decode   │
+  │ Prefill (§4)          │  + restore idx │ Engine (§6)         │
+  └───────────────────────┘                └────────────────────┘
+             ▲                                      ▲
+             │            ┌──────────────────┐      │
+             └────────────│ Peak-Window      │──────┘
+            admit/reject  │ Admission (§7)   │ schedule cohort
+                          └──────────────────┘
+                                   │
+                          ┌──────────────────┐
+                          │ Restoration       │  (§5: native-hybrid | f_θ)
+                          │ Policy            │
+                          └──────────────────┘
+```
+
+## 4. Chunked Restoration Prefill
+
+**Interface**
+
+```
+prefill(prompt_ids: int[T], policy: RestorationPolicy, chunk: int = 2048)
+    -> BoundedKVState
+```
+
+**Behavior** — consume the prompt in fixed `chunk`-sized blocks. For each block:
+
+1. Run the verifier forward over the block against the in-progress
+   `BoundedKVState` (the resident KV so far).
+2. Emit/update each layer's resident KV per the core invariant (§2): exact layers
+   accumulate full; other layers retain only `sink + window`.
+3. The restoration policy (§5) supplies the K/V for any evicted positions a block
+   needs to attend to (for full-attention models); on hybrid models with native
+   sliding the block simply does not attend beyond its window.
+
+**Invariants**
+
+- Per-block working memory is **O(N · chunk · cache_len)** for the mask/activation
+  — never **O(N · T²)**. Memory is decoupled from total prompt length.
+- The LM head is evaluated only where needed (last position for next-token), never
+  as a full `[N, T, vocab]` tensor.
+- No `[·, ·, T, T]` attention mask is ever materialized — attention is computed by
+  the wrapped flash kernel (§6) with window/causal as kernel parameters.
+
+## 5. Restoration Policy
+
+A pluggable policy supplies the K/V for evicted positions, keeping the rest of the
+engine model-agnostic:
+
+| Policy | Model class | Mechanism | Restoration cost |
+| --- | --- | --- | --- |
+| **NativeHybridBounded** | hybrid-attention (e.g. Gemma-4: full + sliding layers) | exact full-attn layers carry recall; other layers are natively local → no reconstruction needed | none (free) |
+| **FThetaRestored** | full-attention (e.g. Qwen/Llama) | dLLM proposer produces transient K/V over history; f_θ projects to verifier K/V at evicted positions | one proposer forward |
+
+The policy decides **where recall comes from** and **whether reconstruction
+runs**. The bounded-KV layout, prefill chunking, admission, and decode are
+identical across policies.
+
+## 6. Bounded-KV Decode Engine
+
+**Interface**
+
+```
+decode(cohort: list[BoundedKVState], max_new_tokens) -> list[token[]]
+```
+
+**Behavior** — decode an admitted cohort in one batched step per token. Each
+session is one batch row over its bounded KV. Attention is the **wrapped flash
+kernel** with the Kakeya window / exact-layer-full as kernel parameters; the
+decode step is **graph-capturable** (static shapes per cohort). The resident KV
+grows only within the bound (sink+window for non-exact layers).
+
+**Invariants**
+
+- Decode-step memory is bounded by the cohort's resident KV (§2), not by total
+  generated length beyond the exact-layer term.
+- No full-KV cache and no paged store: the resident set is the only KV that exists.
+
+## 7. Peak-Window Admission
+
+**Interface**
+
+```
+admit(memory_budget_bytes, model_bytes, sessions) -> admitted_cohort
+```
+
+**Behavior** — a session's cost is its **bounded resident KV** (§2), computed from
+the model's layer layout and the engine's `(sink, window)` — **not** its token
+count. Max concurrency:
+
+```
+max_concurrent = (memory_budget - model_bytes) // resident_KV_per_session
+```
+
+Admission is by **peak window**, so concurrency does **not** degrade as
+conversations lengthen — the defining difference from full-KV admission, where a
+session's cost grows with its history and concurrency collapses at long context.
+
+## 8. Where the engine wins vs vLLM
+
+The advantage scales with the model's **full-attention fraction** (the exact-layer
+term in §2):
+
+- **Hybrid models** (Gemma-4, 25/30 natively sliding): vLLM already bounds the
+  sliding layers and the 5 full-attention layers dominate KV in both engines →
+  the engine is **competitive** (it removes the O(N·T²) prefill cost and provisions
+  by peak window) but has no large structural KV edge.
+- **Full-attention models** (Qwen/Llama, no native sliding): vLLM must keep **all**
+  layers' full KV; the Kakeya engine keeps only `exact + sink + window` and
+  restores the rest via f_θ+proposer → a **large** resident-KV edge → it admits
+  many more concurrent long-context sessions. **This is the engine's target
+  regime.**
+
+## 9. v1 scope and sequencing
+
+- **v1 (this design):** Chunked Restoration Prefill + Bounded-KV Decode +
+  Peak-Window Admission, with the **NativeHybridBounded** policy (Gemma-4) and the
+  **FThetaRestored** policy interface. Flash kernel wrapped via SDPA/FlashAttention.
+- **v1.1:** graph-captured decode + fused-MoE kernels.
+- **v1.2:** FThetaRestored policy fully wired for a full-attention verifier — the
+  configuration that demonstrates the decisive concurrency win over vLLM.

inference_engine/engine/__init__.py

@@ -0,0 +1,25 @@
+"""Kakeya Inference Engine (product runtime).
+
+Bounded-KV-native LLM inference engine — the product-grade vLLM replacement
+defined in ADR 0015 and `docs/design/kakeya-inference-engine-architecture.md`.
+
+Public surface:
+  * :mod:`inference_engine.engine.admission` — peak-window admission + the
+    bounded-KV memory model (pure stdlib; the concurrency math).
+  * :mod:`inference_engine.engine.kakeya_engine` — the engine runtime
+    (chunked restoration prefill + bounded-KV decode). Imports torch lazily.
+"""
+
+from inference_engine.engine.admission import (
+    BoundedKVModel,
+    full_kv_bytes_per_session,
+    max_concurrent_sessions,
+    resident_kv_bytes_per_session,
+)
+
+__all__ = [
+    "BoundedKVModel",
+    "resident_kv_bytes_per_session",
+    "full_kv_bytes_per_session",
+    "max_concurrent_sessions",
+]

inference_engine/engine/admission.py

@@ -0,0 +1,100 @@
+"""Peak-window admission + bounded-KV memory model (§2, §7).
+
+Pure stdlib — this is the concurrency math the engine admits sessions by, and
+the model that quantifies the bounded-KV advantage over a full-KV engine. A
+session's cost is its **bounded resident KV** (exact layers full + other layers
+sink+window), NOT its token count — so concurrency does not degrade as
+conversations lengthen.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Sequence
+
+
+@dataclass(frozen=True)
+class BoundedKVModel:
+    """Per-session bounded-KV cost model for a given model + engine config.
+
+    Attributes
+    ----------
+    num_layers, num_kv_heads, head_dim
+        Verifier layer layout (KV side).
+    n_exact_layers
+        Number of full-context (recall-critical) layers kept exact.
+    sink, window
+        Resident sink + sliding window for the non-exact layers.
+    dtype_bytes
+        Bytes per KV element (bf16 → 2).
+    """
+
+    num_layers: int
+    num_kv_heads: int
+    head_dim: int
+    n_exact_layers: int
+    sink: int
+    window: int
+    dtype_bytes: int = 2
+
+    def __post_init__(self) -> None:
+        if self.n_exact_layers > self.num_layers:
+            raise ValueError("n_exact_layers cannot exceed num_layers")
+        for name in ("num_layers", "num_kv_heads", "head_dim"):
+            if getattr(self, name) <= 0:
+                raise ValueError(f"{name} must be positive")
+        if self.sink < 0 or self.window <= 0:
+            raise ValueError("sink must be >=0 and window must be >0")
+
+    @property
+    def per_token_per_layer_bytes(self) -> int:
+        """K and V for one token at one layer."""
+        return 2 * self.num_kv_heads * self.head_dim * self.dtype_bytes
+
+    def resident_bytes(self, context_len: int) -> int:
+        """Bounded resident KV for a session at logical length ``context_len``."""
+        if context_len < 0:
+            raise ValueError("context_len must be >= 0")
+        n_other = self.num_layers - self.n_exact_layers
+        exact = self.n_exact_layers * context_len
+        other = n_other * min(context_len, self.sink + self.window)
+        return (exact + other) * self.per_token_per_layer_bytes
+
+    def full_kv_bytes(self, context_len: int) -> int:
+        """Full-KV cost (every layer keeps full context) — the full-attention
+        engine's per-session cost, for the advantage ratio."""
+        if context_len < 0:
+            raise ValueError("context_len must be >= 0")
+        return self.num_layers * context_len * self.per_token_per_layer_bytes
+
+    def advantage_ratio(self, context_len: int) -> float:
+        """full-KV / bounded-KV per-session bytes at ``context_len``."""
+        b = self.resident_bytes(context_len)
+        return self.full_kv_bytes(context_len) / b if b else float("inf")
+
+
+def resident_kv_bytes_per_session(model: BoundedKVModel, context_len: int) -> int:
+    return model.resident_bytes(context_len)
+
+
+def full_kv_bytes_per_session(model: BoundedKVModel, context_len: int) -> int:
+    return model.full_kv_bytes(context_len)
+
+
+def max_concurrent_sessions(
+    *, memory_budget_bytes: int, model_weight_bytes: int,
+    per_session_bytes: int,
+) -> int:
+    """Max sessions that fit: (budget − weights) // per-session KV."""
+    if per_session_bytes <= 0:
+        raise ValueError("per_session_bytes must be positive")
+    free = memory_budget_bytes - model_weight_bytes
+    if free <= 0:
+        return 0
+    return int(free // per_session_bytes)
+
+
+def exact_layer_indices_for_layer_types(layer_types: Sequence[str]) -> list:
+    """Indices of full-attention layers given a hybrid model's layer_types."""
+    return [i for i, t in enumerate(layer_types)
+            if "full" in str(t).lower()]