Metadata-Version: 2.4
Name: turbo-attn
Version: 0.54.0
Summary: Optimized CUDAgraph-enabled kernels and attention backend for vLLM, SGLang and more based on TurboQuant near-lossless KV cache compression.
Author-email: "Dmitri Evseev (Arbi City)" <dmitri.evseev@arbi.city>
Maintainer-email: "Dmitri Evseev (Arbi City)" <dmitri.evseev@arbi.city>
License: Apache-2.0
Project-URL: Homepage, https://github.com/arbi-dev/turbo-attn
Project-URL: Repository, https://github.com/arbi-dev/turbo-attn
Project-URL: Issues, https://github.com/arbi-dev/turbo-attn/issues
Project-URL: Changelog, https://github.com/arbi-dev/turbo-attn/blob/main/CHANGELOG.md
Keywords: kv-cache,quantization,vllm,sglang,flash-attention,llm-inference,transformers,cuda
Classifier: Development Status :: 4 - Beta
Classifier: Intended Audience :: Developers
Classifier: Intended Audience :: Science/Research
Classifier: License :: OSI Approved :: Apache Software License
Classifier: Operating System :: POSIX :: Linux
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Requires-Python: >=3.10
Description-Content-Type: text/markdown
License-File: LICENSE
License-File: NOTICE
Requires-Dist: torch>=2.1
Provides-Extra: codebook
Requires-Dist: scipy>=1.10; extra == "codebook"
Provides-Extra: triton
Requires-Dist: triton>=3.0; extra == "triton"
Provides-Extra: cute
Requires-Dist: nvidia-cutlass-dsl[cu13]>=4.6.0; extra == "cute"
Requires-Dist: quack-kernels>=0.6; extra == "cute"
Requires-Dist: apache-tvm-ffi>=0.1.11; extra == "cute"
Provides-Extra: vllm
Requires-Dist: vllm>=0.19; extra == "vllm"
Provides-Extra: flash-attn
Requires-Dist: flash-attn>=2.5; extra == "flash-attn"
Provides-Extra: fa4
Requires-Dist: flash-attn-4==4.0.0b24; extra == "fa4"
Provides-Extra: eval
Requires-Dist: lm-eval>=0.4.5; extra == "eval"
Requires-Dist: ray; extra == "eval"
Requires-Dist: datasets; extra == "eval"
Requires-Dist: langdetect; extra == "eval"
Requires-Dist: immutabledict; extra == "eval"
Requires-Dist: nltk; extra == "eval"
Requires-Dist: sacrebleu; extra == "eval"
Requires-Dist: absl-py; extra == "eval"
Provides-Extra: dev
Requires-Dist: pytest>=7.0; extra == "dev"
Requires-Dist: pytest-xdist>=3.0; extra == "dev"
Requires-Dist: pytest-timeout>=2.1; extra == "dev"
Requires-Dist: scipy>=1.10; extra == "dev"
Requires-Dist: triton>=3.0; extra == "dev"
Provides-Extra: all
Requires-Dist: turbo-attn[codebook,eval,flash-attn,triton,vllm]; extra == "all"
Dynamic: license-file

# Turbo Attention

[![CI](https://github.com/arbi-dev/turbo-attn/actions/workflows/ci.yml/badge.svg)](https://github.com/arbi-dev/turbo-attn/actions/workflows/ci.yml)
[![PyPI](https://img.shields.io/pypi/v/turbo-attn.svg)](https://pypi.org/project/turbo-attn/)
[![License](https://img.shields.io/badge/license-Apache--2.0-blue.svg)](LICENSE)

**The KV-cache backend for VRAM-constrained inference.** Fit ~3.75× more
context into the same GPU, at decode-parity speed. Serve it on
**[arbi-serve](https://github.com/arbi-dev/arbi-serve)** — our own
OpenAI-compatible inference engine, where TKV is a native first-class
backend — or on our thin vLLM / SGLang forks.

PyPI: `turbo-attn` · Import: `tkv` · License: Apache-2.0

---

## Why turbo-attn

Attention is memory-bound and the KV cache is what eats your VRAM. turbo-attn
attacks that directly: a near-lossless 2/4/8-bit KV codec (TurboQuant) welded
to custom prefill + decode kernels that dequantize **inline**, so you pay for
the compression in *bytes stored*, not in tokens/sec.

### 1 · Fit far more context in the same VRAM

The headline. Cross-engine chat-serving baseline on **Qwen3.6-27B-AWQ-INT4**,
TP=2 on 2× RTX 4090 (locked clocks, single-tenant, prefix-caching OFF,
cudagraph, real prompts):

| engine / KV codec | c1 tok/s | c8 tok/s | KV-cache tokens | density |
|---|---|---|---|---|
| SGLang / bf16     | 76.5 | 317.9 | 340,539       | 1× |
| SGLang / **tkv**  | 73.1 | 307.5 | **1,282,029** | **3.76×** |
| vLLM / bf16       | 75.0 | 338.6 | 317,290       | 1× |
| vLLM / **tkv**    | 75.1 | 330.3 | **1,190,272** | **3.75×** |

Same card, ~3.75× the resident context — **at decode parity** (the dequant is
amortized under cudagraph; ~13 ms TPOT is the model's memory-bandwidth
roofline either way). That is longer prompts, more concurrent sequences, or a
bigger model on the hardware you already have.

- **SmartMix per-layer allocation.** `TKV_BITS` is an *average* bits-per-element
  target; a calibrated solver gives each layer its own `(k_bits, v_bits)` over
  the full 2–8-bit range, and the kernels read that **heterogeneous cache in a
  single forward pass** — closer to bf16 at equal memory. See
  [SmartMix](#smartmix-per-layer-optimal-kv-read-in-one-pass).
- **Near-lossless, verified.** Correctness is gated on PPL / needle /
  task-accuracy against a bf16+FlashAttention control arm — not vibes.
- **Big-model shapes covered.** GQA, sliding-window, MTP (multi-token
  predict), and MLA (DeepSeek V2/V3/V4) are all first-class.

### 2 · Fast, bit-identical kernels — no speed tax

Compressing the KV cache only pays off if the kernel can *read* it compressed.
Decompress to a bf16 buffer before attention and **the memory savings never
materialize** — you write the decompressed cache back through HBM and hand
back, in bandwidth, every byte you just saved. turbo-attn's kernels instead
**dequantize inline**, inside the MMA pipeline, so the *compressed* bytes are
the only thing that crosses HBM and the dequant hides under the math. That is precisely why ~3.75× smaller KV comes at
**decode parity** rather than a speed penalty — the codec's memory win and the
kernel's speed are the *same* win.

On prefill the custom **Turbo split-D prefill mainloop** — an original CuTeDSL
tensor-core kernel, not a fork of anyone's — goes further, beating **upstream FlashAttention-4** by
**1.2–1.4× geomean, bit-identically**, across Ada / Blackwell / Ampere (and at
or above stock FlashAttention-2 on raw bf16). Those kernels also run
**uncompressed bf16 / fp16** (BypassLoader) — so turbo-attn is a fast general
attention backend on its own, and that raw-bf16 path is where we measure
*against* upstream FA4 / FA2 like-for-like (no codec confound). Everything runs
under full CUDAGraph capture.

See **[The kernels](#the-kernels)** for the per-kernel innovation breakdown.

### 3 · Built for fast iteration & developer-friendliness

turbo-attn is designed to be *adopted and extended* without forking the world:

- **First-class in arbi-serve; a thin overlay elsewhere.** On
  **arbi-serve**, our own inference engine, TKV is a native backend — nothing
  to fork. For vLLM and SGLang we ship *small* forks, because neither upstream
  can yet register a new KV-cache dtype / attention backend at runtime; once
  the fork is installed, enabling TKV is a single `--kv-cache-dtype tkv` flag
  via an auto-registered plugin entry-point. (HuggingFace Transformers is used
  only as a correctness *oracle* — a bf16 reference arm — not a serving path.)
- **Two pieces you can consume separately** (see below) — take the pure-PyTorch
  codec with *any* attention backend, or take the kernels with *any* KV
  format.
- **One clean extension point — the Loader Protocol.** Adding a new KV format
  (fp8, int8, nvfp4, GPTQ, …) is a *sibling module*, not mainloop surgery. You
  write the per-tile SMEM-fill body; the scheduler, softmax, and epilogue stay
  ours. See [`docs/writing_a_loader.md`](docs/writing_a_loader.md) for a
  complete worked fp8 example.
- **Hackable codec.** `TkvCodec` is pure PyTorch with a Triton/CUDA fast
  path and a readable fallback — easy to read, ablate, and unit-test on CPU.
- **Apples-to-apples by construction.** `BypassLoader` runs raw bf16/fp16
  through the *byte-identical* kernel, so codec-vs-baseline ablations isolate
  exactly one variable.
- **Turn-key ops.** A Dockerfile per engine (`docker compose up`),
  auto-calibration on first init, and a "just-works" config surface with
  optimal-by-default settings and loud, named, validated caps.

---

## Install

```bash
pip install turbo-attn                  # codec + CUDA/Triton kernels
pip install "turbo-attn[vllm]"          # + vLLM attention backend
pip install "turbo-attn[all]"           # + SGLang, FlashInfer, flash-attn, eval harness
```

## Quickstart

```python
import torch
from tkv import TkvCodec

codec = TkvCodec(head_dim=128, bit_width=4, device="cuda")
keys = torch.randn(8, 128, device="cuda")

packed, norms = codec.compress_k(keys)
recon = codec.decompress_k(packed, norms)
```

See [`examples/`](examples/) for runnable snippets and
[ARCHITECTURE.md](ARCHITECTURE.md) for a codebase tour.

## Two independently-usable pieces

`turbo-attn` ships two pieces that are sold as a stack but designed to be consumed separately:

1. **Codec → any attention backend.** `TkvCodec` is a pure, framework-agnostic compressor: compress with TKV, decompress to bf16 / fp16, hand the result to vanilla `flash_attn_varlen_func`, FlashInfer, SGLang attention, anything that takes raw KV. See [`examples/06_tkv_codec_with_third_party_attention.py`](examples/06_tkv_codec_with_third_party_attention.py).

2. **Kernels → any KV format.** The cute-DSL prefill and split-K paged decode kernels are policy-parametric on the K/V format via the **Loader extension point**. The bundled set is `{TkvLoader, BypassLoader}`:
   - **`TkvLoader`** — TKV centroid-based codec dequant (the production path).
   - **`BypassLoader`** — raw bf16 / fp16 KV, no codec. Useful for apples-to-apples ablations under an otherwise-byte-identical kernel.
   Third-party formats (fp8, int8, nvfp4, …) are **not** shipped — write a sibling Loader for your format. The Loader is the public extension surface; mainloop / scheduler / softmax / epilogue stay turbo-attn's. See [`docs/writing_a_loader.md`](docs/writing_a_loader.md) for a worked fp8 example.

## Repo layout

- `tkv/` — the package (codec, kernels, runtime, vLLM/SGLang plugins, calibration pipeline).
- `tkv/kernels/loaders/` — bundled cute-DSL prefill Loaders (`tkv`, `bypass`).
- `tkv/kernels/_decode_loader_*.cuh` — bundled decode Loaders (`TkvDecodeLoader`, `BypassDecodeLoader`), over the shared staging in `_decode_loader_common.cuh`.
- `docs/`, `docker/`, `scripts/`, and `examples/` — public docs, deploy recipes, helper scripts, and runnable examples.

## Run with Docker

Three inference servers are supported: vLLM, SGLang, and [arbi-serve](https://github.com/arbi-dev/arbi-serve). Each ships a turn-key Dockerfile. Calibration files for the bit-width / model combo go in a host directory; the `TKV_CALIBRATION_FILE` env var inside the container points to one. Examples below use `Qwen3.5-0.8B` + a K4V4 calibration.

### Layout assumed

```
/path/to/models/Qwen3.5-0.8B/...                       # HF snapshot
/path/to/calibrations/qwen3.5-0.8b_k4v4.json           # calibration bundle
```

All three accept the same CLI flags: `--kv-cache-dtype tkv --attention-backend turbo-attn` plus `TKV_BITS=<float>` and `TKV_CALIBRATION_FILE=<path>`. `TKV_BITS` is the *average* bits-per-element across K and V (e.g. `4.0`, `5.0`, `6.0`); a per-layer Lagrangian solver turns that target into a per-layer `(k_bits, v_bits)` allocation that lives in the calibration bundle. `TKV_BITS=4.0` does *not* mean "K and V both at 4 bits" — it means "average 4 bits-per-element under the smart per-layer allocation".

### Sibling-checkout layout

All three Dockerfiles `COPY` from a sibling-repo layout. Clone the relevant repos as siblings of `turbo-attn/`:

```
GIT/
├── turbo-attn/        # this repo
├── vllm-fork/         # arbicity/vllm-turbo        (only needed for vLLM image)
├── sglang-fork/       # arbicity/sglang-turbo      (only needed for SGLang image)
└── arbi-serve/        # arbi-dev/arbi-serve  (only needed for arbi-serve image)
```

```bash
mkdir -p ~/GIT && cd ~/GIT
git clone https://github.com/arbi-dev/turbo-attn
git clone https://github.com/arbicity/vllm-turbo       vllm-fork    # for vLLM
git clone https://github.com/arbicity/sglang-turbo     sglang-fork  # for SGLang
git clone https://github.com/arbi-dev/arbi-serve              # for arbi-serve
```

### First boot on a new machine

The first boot on any given machine (fresh image, empty `TKV_CACHE_DIR`) pays two real, one-time costs before the server is ready. **This is expected — not a hang:**

1. **CUDA kernel JIT compile** (nvcc/cutlass) — every decode/prefill kernel variant your deployment needs gets compiled to a `.so` and cached under `TKV_CACHE_DIR` (default `/root/.cache/torch_extensions`). Real cost scales with the number of `(head_dim, GQA-ratio, bit-width)` combos your model + calibration bundle need.
2. **Decode-kernel autotune sweep** — times every candidate kernel config (splits × tile size × min-blocks-per-SM) per *distinct KV shape* and caches the winning pick. A uniform bit-width bundle has one shape, so it sweeps once; a SmartMix (per-layer heterogeneous bit-width) bundle sweeps once per distinct per-layer shape. Real measured cost, live RTX 5090, 2026-08-17: the default (trimmed) bucket ladder swept in 66.5s at a small batch grid, vs. 273.4s untrimmed (4.1× slower) — and a real production boot log at the *full* batch grid (up to 256) showed **1413.99s (~23.5 min) for the two largest buckets alone**, on a single KV shape. ([PR #836](https://github.com/arbicity/turbo-attn/pull/836) trims those buckets out by default; a fingerprint that misses that trim still pays the untrimmed cost.)

**Both phases print real progress as they run** — watch `docker compose logs -f` for `[TKV precompile] ... OK (Ns)` lines (one per compiled kernel) and `[TKV autotune] ... PICK/HELD ...` lines (one per swept shape/batch cell). If the log genuinely stops advancing for several minutes with no new line, *that* — not the wait itself — is the real signal something's wrong.

**Subsequent boots on the same machine, with the same model / bit-width / TP config, are fast** (typically seconds): both the compiled-kernel cache and the autotune table persist across container restarts as long as `TKV_CACHE_DIR` is a mounted volume (the default compose files do this).

**To skip the JIT-compile cost entirely**, build a pre-baked image — bakes a fixed grid of kernel `.so`s at `docker build` time for a specific GPU arch, so a `docker run` from that image boots with no runtime nvcc:

```bash
docker build --build-arg TKV_BAKE=1 \
              --build-arg TKV_CUDA_ARCH=8.9 \
              --build-arg TKV_BAKE_SHAPES=256:2:8 \
              --build-arg TKV_BAKE_BITS=4 \
              -f docker/Dockerfile .
```

`TKV_CUDA_ARCH` must match the GPU you serve on (`8.9` = Ada / RTX 4090, `12.0` = Blackwell / RTX 5090, `9.0` = Hopper). This is *not* the default build (every CI image build would otherwise pay a multi-minute from-scratch nvcc compile with no layer cache) — opt in explicitly when you want a fast-boot demo image. Note it only removes the kernel-compile phase above; the autotune sweep (phase 2) still runs on first boot regardless, since it depends on the actual serving shapes.

### vLLM

Uses our [`vllm-fork`](https://github.com/arbicity/vllm-turbo) rebased onto upstream `v0.28.0` (small overlay — `CacheDType` support, per-group block-pool bookkeeping, named `TURBO_ATTN` slot in the attention backend registry; full layout in [`docker/PATCHES.md`](docker/PATCHES.md)).

```bash
cd ~/GIT/turbo-attn/docker

# build + run
TKV_MODELS_ROOT=/path/to/models \
  docker compose -f compose.vllm.yaml up -d --build

docker compose -f compose.vllm.yaml logs -f

# serve a request once "Application startup complete" appears:
curl http://localhost:8000/v1/completions \
  -H "Content-Type: application/json" \
  -d '{"model": "/models/Qwen3.5-0.8B", "prompt": "The capital of France is", "max_tokens": 12, "temperature": 0}'
```

Optional env (override on the `docker compose` command line):

| Variable | Default | Purpose |
|---|---|---|
| `TKV_MODEL` | `/models/Qwen3.5-0.8B` | Container-side model path |
| `TKV_BITS` | `4.0` | Average bits-per-element target |
| `TKV_CALIBRATION_FILE` | unset | Path to a calibration bundle. Unset → uniform K4V4 fallback (tests/CI only; production needs a bundle) |
| `TKV_PORT` | `8000` | Host port |
| `TKV_GPU_DEVICE` | `0` | `NVIDIA_VISIBLE_DEVICES` |
| `TKV_MAX_MODEL_LEN` | `2048` | Max context length |

### SGLang

Uses our [`sglang-fork`](https://github.com/arbicity/sglang-turbo) rebased onto upstream `v0.5.18` (small overlay — plugin registries for KV-cache dtypes and attention backends; full layout in [`docker/PATCHES.md`](docker/PATCHES.md)).

```bash
cd ~/GIT/turbo-attn/docker

TKV_MODELS_ROOT=/path/to/models \
  docker compose -f compose.sglang.yaml up -d --build

docker compose -f compose.sglang.yaml logs -f

curl http://localhost:30000/v1/completions \
  -H "Content-Type: application/json" \
  -d '{"model": "/models/Qwen3.5-0.8B", "prompt": "The capital of France is", "max_tokens": 12, "temperature": 0}'
```

Same env-var contract as vLLM (drop `TKV_MAX_MODEL_LEN`; SGLang uses `TKV_CONTEXT_LEN` and `TKV_MEM_FRAC` for `--mem-fraction-static`, default `0.45`).

### arbi-serve

Standalone OpenAI-compatible server with TKV backends as a first-class citizen. Lives in [arbi-dev/arbi-serve](https://github.com/arbi-dev/arbi-serve).

```bash
cd ~/GIT/arbi-serve

ARBI_MODELS_ROOT=/path/to/models \
  docker compose up -d --build

docker compose logs -f

curl http://localhost:8000/v1/completions \
  -H "Content-Type: application/json" \
  -d '{"model": "/models/Qwen3.5-0.8B", "prompt": "The capital of France is", "max_tokens": 12, "temperature": 0}'
```

### MLA models

For DeepSeek V2/V3/V4, additionally set `-e TKV_MLA_ENABLE=1` on whichever container.
`--kv-cache-dtype tkv-bypass` serves the same models on their own KV — in whatever
format the model itself produced — and needs neither `TKV_MLA_ENABLE` nor `TKV_BITS`.
A model that quantized its own KV declares that format through
`tkv.runtime.declare_bypass_slot`; the cache is then allocated and read in it rather
than inflated to bf16. Decode is served; prefill over a model-quantized latent is
refused by name.

### Calibration

Calibration bundles are versioned under [`calibrations/`](calibrations/). To
roll your own:

```bash
python -m tkv.calibration.calibrate_centroids \
    --model Qwen/Qwen3.5-0.8B \
    --output qwen3.5-0.8b_k4v4.json \
    --bit-width 4
```

## How it works

1. **Rotate** each KV vector with a fast Walsh–Hadamard transform.
2. **Normalize** — store the magnitude as a single BF16 value.
3. **Quantize** each rotated coordinate to a shared codebook.

Attention scores on rotated KV are bit-identical to attention on unrotated KV when the query is rotated by the same matrix; we pre-rotate Q once per request and compute everything in the rotated space.

## The kernels

turbo-attn ships **two** original attention kernels, and they are different pieces of machinery — do not collapse them into one, and neither is FlashAttention:

- **Prefill — the Turbo prefill kernel.** `TurboPrefillForward` (+ its Bypass / Hybrid siblings), a tensor-core **CuTeDSL** kernel entered via `turbo_prefill`. Tuning knobs are `TKV_PREFILL_*`; the engine value is `TKV_PREFILL_ENGINE=turbo` and its row class is `ROW_PREFILL`.
- **Decode — the split-K / unified decode kernel.** An original **SIMT CUDA C++** kernel (`cuda_decode_splitk.py` / `_cuda_decode_unified_splitk.cu`). Tuning knobs are `TKV_DECODE_*` / `TKV_MTP_*` and its row class is `ROW_DECODE`.

Both are turbo-attn's own work, and neither should be confused with the **TKV (TurboQuant) codec** — the codec is the compression format these kernels *read*; the kernels are what read it. The kernels are where the codec turns from a storage trick into a serving win. Each is a substantial piece of engineering in its own right; the throughline is **inline dequant**.

### Inline dequant — why the codec's savings actually survive

A naive integration decompresses the packed KV to a bf16 tensor before attention. That round-trip writes the *decompressed* cache back through HBM — paying back every byte the codec saved, and erasing the bandwidth win that makes KV compression worth doing on a memory-bound kernel.

turbo-attn's kernels **never materialize a decompressed KV tensor.** Both the prefill mainloop and the decode kernel unpack + dequantize *inside* the MMA pipeline, tile by tile, straight into SMEM/registers. Only the compressed bytes cross HBM; the dequant math hides under the matmul. This is the single most important property in the repo — it is what makes "3.75× smaller cache" and "decode-parity throughput" the *same* fact instead of a trade-off.

The seam that keeps this reusable is the **Loader Protocol** (see [Two independently-usable pieces](#two-independently-usable-pieces)): the per-tile SMEM-fill body is a **compile-time-specialized** functor, so the dequant loop is unrolled and any format axes (bit width, scale presence) are dead-code-eliminated against the mainloop. Runtime branches or function pointers in the hot loop would collapse this by an order of magnitude — so there are none.

### Also a full bf16 backend — the confound-free baseline

The very same prefill and decode kernels run a **`BypassLoader`**: raw bf16 / fp16 KV, no codec, through a byte-identical mainloop / softmax / epilogue. This matters twice:

- turbo-attn is a **general attention backend even with the codec off** — you can adopt the faster prefill/decode schedules on ordinary uncompressed KV, no calibration, no compression.
- It is **where we benchmark against upstream FlashAttention-4 and FlashAttention-2.** Those are the baselines we compare *against*, never something turbo-attn is built from. Upstream FA4 and FA2 consume raw bf16 — and so does the Turbo prefill mainloop in bypass mode. The reported 1.2–1.4× over upstream FA4 and the FA2 parity are therefore **confound-free, like-for-like schedule wins on identical bf16 inputs**, not "our dequant path vs their plain path." The codec-path numbers are measured separately against the *same kernel in bypass mode*, so the codec's own overhead is isolated on its own axis too.

### Prefill — the Turbo prefill kernel (CuTeDSL split-D mainloop)

`tkv/kernels/cuda/prefill/turbo_prefill.py`, the unconditional production prefill path at every head-dim. `TurboPrefillForward` is an **original CuTeDSL kernel** — it owns its mainloop, epilogue, base class and parameter plumbing, and is not a fork or subclass of any FlashAttention forward kernel. What it shares with upstream is a vendored *support substrate* (`tkv/kernels/cute/_fa/`: masks, softmax, seqlen/block info, pack-GQA, named barriers, tile schedulers — BSD-3-Clause, see [ATTRIBUTION.md §4](ATTRIBUTION.md) and `tkv/kernels/cute/_fa/UPSTREAM.md`), not an attention mainloop.

- **Split-D (head-dim-in-SMEM) sequence-parallel walk** — **1.15× geomean over upstream FlashAttention-4** with **bit-identical output** on Ada, and at or above mature stock FA2 on bf16. Zero accuracy cost: the codec, not the schedule, is the only accuracy axis.
- **Inline TQ dequant in the MMA pipeline** — K and V dequant straight into the MMA operand SMEM slabs, with no staging pass; the warps that own the dequant STS also own the pipelining.
- **First-chunk bypass** (`TKV_PREFILL_BYPASS`, default on) — prompt prefill runs *uncompressed*, then re-rotates to the TQ basis for decode, so prefill quality is never codec-limited.
- **Strict dispatch:** every prefill request routes through this kernel; a shape it declines raises rather than falling back to a decompress + third-party-attention path (there is no such path left). See [Prefill performance](#prefill-performance) for numbers.

### Decode — the unified split-K kernel

`tkv/kernels/_cuda_decode_unified_splitk.cu`, the sole production decode path — an **original SIMT CUDA C++** kernel, a different codebase from the CuTeDSL prefill mainloop above.

- **One fused split-K SIMT kernel** — unpack → dequant → Q·K → online-softmax → P·V in a single pass, no decompress buffer.
- **One template-policy body** serves **full attention and sliding-window** via a compile-time `RangePolicy` (autotune branches on `sliding_window`), *plus* the **MTP-verify `BLOCK_M=N`** path — the earlier standalone `BLOCK_M=1` kernel was retired into this one.
- **Rotation fusion** — `TKV_FUSE_QROT` folds the Q-rotation prologue in, and `rotate_output` is folded into the `o_proj` weights (`TKV_O_PROJ_FOLD`), so the rotation cost leaves the token loop entirely.
- **MLA-shaped decode** variant for DeepSeek V2/V3/V4.
- **Per-shape autotuner** (`tkv/runtime/autotune.py`) picks split count / MTP variant / `fuse_qr` per geometry.

### Compress-on-the-wire + CUDAGraph

- **Fused compress + store** (`fused_compress_store`) encodes K/V and writes the packed cache in a single kernel — compression happens on the *write* path, never as a separate materialize-then-compress pass.
- **Full CUDAGraph capture** end-to-end with CG-safe per-step metadata precomputed host-side, so the entire compressed path is invisible at the token level.

## SmartMix: per-layer optimal KV, read in one pass

Uniform KV quantization spends the same bits on every layer — but layers are not equally sensitive: some tolerate 3-bit keys with 5-bit values, others the reverse. **SmartMix** is turbo-attn's per-layer bit-allocation. Given an *average* bits-per-element budget, a calibrated solver assigns each layer its own `(k_bits, v_bits)` over the full 2–8-bit range (odd widths included), spending bits where the model's own decode drift says they matter.

Two things make this more than a calibration knob:

- **A single forward pass reads a heterogeneous cache.** The cache holds a *different* `(k_bits, v_bits)` per layer — e.g. layer 15 at K4V3, layer 19 at K5V3, layer 23 at K4V5 — and each layer's attention dispatches to a kernel **compile-time-specialized to that layer's widths** through the Loader Protocol's geometry-keyed dispatch. No per-token branching, no runtime format check in the hot loop; every width is its own specialized PTX, with CG-safe metadata precomputed host-side so the whole heterogeneous forward still runs under one CUDA graph. We are not aware of another production serving stack that reads per-layer mixed-precision KV natively — most KV-quant implementations are a single width for the entire cache.

- **It buys fidelity at equal memory** — equivalently, a target fidelity at *fewer* bits (i.e. a smaller cache). Measured as KL drift from bf16 (lower = closer), SmartMix beats the **best-global** uniform allocation at the *same* bytes/token, and the margin **grows as compression gets more aggressive** (−12.6% at 6 bits, −19.5% at 4 bits, reaching ~−36% at 3.5 bits on the raw-Lloyd bundle):

  > ⚠️ **Preliminary numbers** — Qwen3.5-0.8B, 2026-07-12, correctly baselined vs best-global uniform. A refreshed run is pending; treat the exact values as provisional.

  | KV format (equal bytes/tok) | avg bits | KL drift vs bf16 ↓ |
  |---|---|---|
  | bf16 self-jitter floor | 16 | 0.831 |
  | uniform K8V8 (calibrated) | 8 | 1.050 |
  | fp8 e4m3 | 8 | 1.466 |
  | uniform K6V6 | 6 | 2.122 |
  | **SmartMix @ 6.0** | 6 | **1.855  (−12.6%)** |
  | uniform K4V4 | 4 | 14.409 |
  | **SmartMix @ 4.0** | 4 | **11.595  (−19.5%)** |

  The same curve shows **calibrated KV beating fp8** — 1.05 vs 1.47 mbits/tok at 8 real bits — so spending the bit-budget on a fitted codebook beats spending it on an fp8 mantissa.

The fidelity win is unambiguous and monotone in the bit-budget. Turning it into a downstream **task-accuracy** number is gated on running a benchmark hard enough to be KV-limited: at 4 bits the tasks measured so far (NIAH to 131K, GSM8K, math500) put SmartMix *and* uniform at the bf16 ceiling, so the discriminating benchmark is the open eval item — not a sign the headroom isn't there. Allocations are produced by the solver (`python -m tkv.auto_calibrate --solve`) and stored in the calibration bundle's `byte_budget_table`.

## Performance

Cross-engine chat-serving baseline on **Qwen3.6-27B-AWQ-INT4**, TP=2 on 2× RTX 4090, via the canonical `benchmarks/serve_mtp/` harness (retired `chat_baseline.sh`/`serve_bench.sh`; numbers are historical) — locked clocks, single-tenant, **prefix-caching OFF**, cudagraph, real prompts, `vllm bench serve`. The decode/density table is in [Why turbo-attn](#1--fit-far-more-context-in-the-same-vram) above.

**MTP (native head, thinking mode, c1):** vLLM ~145–147 tok/s, SGLang ~128 — each ~1.7–2× its own no-MTP baseline. The residual ~13% is SGLang's NEXTN per-step overhead at equal draft acceptance (~3.5), not draft quality.

### Prefill performance

The Turbo prefill mainloop is an original CuTeDSL kernel; upstream FlashAttention-2 and FlashAttention-4 are the baselines it is measured **against**. **These comparisons run in bypass (raw bf16) mode** — stock FA2 sees the exact same uncompressed inputs the Turbo prefill mainloop does, so the delta isolates the schedule and nothing else (no codec confound). Measured with `benchmarks/bench_prefill_vs_reference.py` (eager, `torch.cuda.Event`, median of 50, dense causal prefill across `head_dim ∈ {64,128,256}`, GQA ratios `{1,4,7,8}`, and both symmetric and asymmetric chunked-prefill shapes `512…131072`):

| GPU | vs stock **FA2** (bf16) |
|---|---|
| RTX 4090 (Ada, sm89) | net-positive (1.00× geomean; FA2 edges asym only) |
| RTX 5090 (Blackwell, sm120) | ahead (1.07× geomean, 29/42) |
| RTX A5500 (Ampere, sm86) | ahead (1.03× geomean, 27/42) |

#### vs upstream FlashAttention-4 (2026-08-20)

Against the real `flash-attn-4==4.0.0b24` PyPI package through its public
`flash_attn.cute.flash_attn_func` API — the upstream kernel, not turbo-attn's
vendored `_fa/` support substrate. `flash-attn-4` is a benchmark-only extra;
no turbo-attn runtime path imports it. Bypass mode (raw bf16 K/V, no codec), so the delta isolates the
mainloop schedule. RTX 4090 (Ada, sm89), `nvidia-cutlass-dsl>=4.6.0`, CUDA
events, median of 20 after 5 warmup.

| shape | S=1024 | S=4096 | S=16384 |
|---|---|---|---|
| `qwen_gqa4_d128` (B2, 32q/8kv) | **1.18×** | **1.14×** | **1.13×** |
| `gemma_mqa_d128` (B2, 16q/2kv) | **1.18×** | **1.14×** | **1.14×** |

**1.15× geomean, 6/6 cells.** Every cell is numerically exact against
upstream: `max_abs_diff = 0.0`, `cosine = 1.000000`, and `lse_max_abs_diff
= 0.0` on the softmax LSE both kernels compute internally. The schedule
differs; the arithmetic does not.

The three `head_dim=256` cells are **unmeasured, and no result is claimed
for them.** Upstream sizes its tiles for SM80's 163 KB of shared memory and
ships a dedicated `flash_fwd_sm120.py` subclass for Blackwell GeForce's
99 KB, but `flash_fwd.py` still carries `# TODO: sm86 and sm89` — so Ada
inherits SM80 tiles it cannot fit, asking 131072 bytes against a 101376 byte
cap, and CUDA rejects the launch (`CUDA_LAUNCH_INVALID_CONFIG`). This is an
Ada arch-support gap upstream has not closed yet, **not a head-dim limit**.
Measuring `d256` and the Gemma-4-style 256/512 layouts against upstream
needs a card whose shared-memory budget upstream targets.

Single-GPU measurement; clocks not locked (the 13–18% margins sit well
clear of clock drift). Reproduce with `pip install "turbo-attn[cute,fa4]"`
— installing `fa4` *after* `cute` silently downgrades cutlass-dsl — then:

```bash
python benchmarks/bench_turbo_vs_stock_fa4.py
```

On raw bf16 the Turbo prefill mainloop is **at or above parity with the mature stock FlashAttention-2** — net-positive on every GPU, with the sole soft spot being asymmetric long-prefill on Ada — while being the only one of the two that can serve the compressed TKV codec path at all (FA2 has no inline dequant). For a head-to-head against the upstream `flash-attn-4` package, see `benchmarks/bench_turbo_vs_stock_fa4.py`. Reproduce:

```bash
python benchmarks/bench_prefill_vs_reference.py --dtypes bf16,tq4 --fa2 --flashinfer
```

## Configuration

All runtime configuration is via `TKV_`-prefixed environment variables. The supported surface is below; anything unlisted is internal and may change.

### Bit width and calibration

| Variable | Default | Description |
|---|---|---|
| `TKV_BITS` | `4.0` | *Average* bits-per-element target across K and V (float in `[2.0, 8.0]`). The runtime looks up the calibration bundle's `byte_budget_table[<TKV_BITS>]` for the per-layer `(k_bits, v_bits)` allocation. Hard error if the entry is missing — no silent fallback. |
| `TKV_CALIBRATION_FILE` | `""` | Path to a calibration bundle (centroids + per-channel scales + `byte_budget_table`). Required for production; bundle generation: `python -m tkv.auto_calibrate`. When unset, the plugin falls back to uniform K4V4 (tests/CI only). |
| `TKV_AUTO_CALIBRATE_MODEL` | `""` | Model path for plugin-side auto-calibration when `TKV_CALIBRATION_FILE` doesn't exist on first init. |

### Engine selection

| Variable | Default | Description |
|---|---|---|
| `TKV_ENGINE` | `""` (auto) | Decode engine. `""` (auto), `native_tq` and `cuda_standalone` all select the fused CUDA decode + Turbo split-D prefill path; any other value raises. |
| `TKV_PREFILL_ENGINE` | `turbo` | Prefill dispatcher family. `turbo` (the CuTeDSL `turbo_prefill` path) is the only accepted value; every prefill request routes through it and a shape it declines raises. Removed values — including the old spelling `fa4` — raise a `ValueError` listing the valid set. |
| `TKV_PREFILL_BYPASS` | `1` | First-chunk prefill bypass — skip codec on prompt-prefill, then re-rotate to TQ basis for decode. |
| `TKV_FUSE_QROT` | `""` (auto) | Fused Q-rotation prologue. Decode-only. |
| `TKV_O_PROJ_FOLD` | `on` | Fold `rotate_output` into `o_proj` weights. |
| `TKV_MTP_SPLITK` | `1` | Use split-K decode kernel for MTP layers. |
| `TKV_DECODE_SPLITS` | `""` (autotune) | Force decode-kernel split count. |

### Backend behaviour

| Variable | Default | Description |
|---|---|---|
| `TKV_NO_JIT` | `0` | Fail if a kernel variant is not pre-compiled. |
| `TKV_K_NC` | `1` | Apply norm-correction to K reads in the dequant path. |
| `TKV_DISABLE_PRESCALE` | `0` | Disable per-channel pre-scaling on compress upload. |
| `TKV_STRICT_NO_SDPA` | `0` | Raise instead of taking the `head_dim>256` SDPA fallback. Recommended for `head_dim>256` deployments. |

### MLA (DeepSeek V2/V3/V4)

| Variable | Default | Description |
|---|---|---|
| `TKV_MLA_ENABLE` | `0` | Master switch for the MLA codec. Not read by `--kv-cache-dtype tkv-bypass`, which serves MLA layers on the model's own latent. |
| `TKV_MLA_ROPE_HEAD_DIM` | `64` | RoPE head dimension for MLA latent + RoPE split. |

## Why we ship engine forks

turbo-attn needs each host engine to let a plugin register a new KV-cache dtype **and** attention backend at runtime. Neither vLLM nor SGLang allows that upstream yet, so we ship thin overlay forks — [`vllm-turbo`](https://github.com/arbicity/vllm-turbo) and [`sglang-turbo`](https://github.com/arbicity/sglang-turbo):

- **vLLM** — `CacheDType` in `vllm/config/cache.py` is a Pydantic `Literal` validated at class-definition time, which blocks runtime registration of new KV-cache dtypes. Our overlay relaxes it, adds a named `TURBO_ATTN` backend slot, and per-group block-pool bookkeeping.
- **SGLang** — needs plugin registries for KV-cache dtypes and attention backends wired into the engine.

Both are thin overlays rebased on tagged upstream releases; full layout in [`docker/PATCHES.md`](docker/PATCHES.md). **arbi-serve** needs no overlay — TKV is a native backend there.

## Citation

If Turbo Attention helps your work, please cite both the underlying TurboQuant paper and this implementation:

```bibtex
@misc{turbo_attention2026,
  title = {Turbo Attention: Production attention backend for TurboQuant KV cache compression},
  author = {Evseev, Dmitri},
  year = {2026},
  url = {https://github.com/arbi-dev/turbo-attn}
}

@inproceedings{zandieh2026turboquant,
  title = {TurboQuant: Near-optimal KV Cache Quantization for LLM Inference},
  author = {Zandieh, Amir and others},
  booktitle = {ICLR},
  year = {2026}
}
```

## License

Apache License, Version 2.0. See `LICENSE` and `NOTICE`.

## Measured KV-cache fidelity

Teacher-forced decode-path divergence from an fp32 oracle (mean |Δ log-prob| of the target token, long-decode positions), Qwen3.5-0.8B, n=8 prompts @ 16k prefill + 2k decode. Lower = higher fidelity. Reproduce: `benchmarks/eval/paper_kv_ablation/` (committed harness, pinned target + bundle sha).

| model / codec | bits | Δ from fp32 (↓ = better) |
|---|---|---|
| BF16 (reference) | 16 | 0.0059 |
| **tkv K6V6 — ours** | 6 | 0.0089 |
| FP8 | 8 | 0.0091 |
| **tkv K4V4 — ours** | 4 | 0.0237 |
| turboquant — upstream | 4 | 0.0353 |

- Our **6-bit KV matches FP8's 8-bit fidelity** at ¾ the bits.
- Our **4-bit KV beats upstream turboquant 4-bit** (and 3.1× better on the first token via the fresh-diagonal).
