Metadata-Version: 2.4
Name: turbo-attn
Version: 0.23.0
Summary: Optimized CUDAgraph-enabled kernels and attention backend for vLLM, SGLang and more based on TurboQuant near-lossless KV cache compression. SOTA performance with Gemma 4, Qwen 3.6 and other modern LLMs.
Author-email: "Dmitri Evseev (Arbi City)" <dmitri.evseev@arbi.city>
Maintainer-email: "Dmitri Evseev (Arbi City)" <dmitri.evseev@arbi.city>
License: MPL-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 :: Mozilla Public License 2.0 (MPL 2.0)
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: flashinfer
Requires-Dist: flashinfer>=0.6; extra == "flashinfer"
Provides-Extra: fa4
Requires-Dist: nvidia-cutlass-dsl>=4.4.2; extra == "fa4"
Requires-Dist: quack-kernels>=0.3.7; extra == "fa4"
Requires-Dist: apache-tvm-ffi>=0.1.9; extra == "fa4"
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: 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,fa4,flash-attn,flashinfer,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-MPL--2.0-blue.svg)](LICENSE)

A modular attention backend for vLLM, SGLang, and HuggingFace Transformers. Custom CUDA + Triton kernels with full CUDAGraph capture, asymmetric K/V quantization, hybrid-model support. Built on FlashAttention; based on TurboQuant near-lossless KV cache compression.

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

## 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 TurboKVCodec

codec = TurboKVCodec(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.** `TurboKVCodec` 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`).
- `docs/`, `docker/`, `scripts/`, `examples/`, `experiments/` — public docs, deploy recipes, helper scripts, runnable examples, research notes.
- The top-level `internal/` directory is engineering-only and unsupported — design notes, internal compose files, dev harnesses. It is kept in-tree for development history but excluded from the wheel (`MANIFEST.in prune internal`), so it never ships.

## 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
```

### vLLM

Uses our [`vllm-fork`](https://github.com/arbicity/vllm-turbo) rebased onto upstream `v0.20.1` (small overlay — `CacheDType` Literal relaxation, per-group block-pool bookkeeping, named `TURBO_ATTN` slot in `AttentionBackendEnum`; 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.11` (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.

### Calibration

Pre-built calibration bundles for common models live in HuggingFace at [`arbi-dev/turbo-attn-calibrations`](https://huggingface.co/arbi-dev/turbo-attn-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 decode path is one fused CUDA kernel that unpacks, dequantizes, and runs Q·K, online softmax, and P·V in a single pass — no decompress buffer.

Prefill runs on the **arbi prefill mainloop** (`tkv/kernels/cuda/prefill/arbi_prefill.py`) — turbo-attn's own CuTeDSL mainloop, and the unconditional production prefill path at every head-dim. It is built on the same CuTeDSL base as FlashAttention-4 (our `tkv/kernels/cute/_fa/` fork of Dao-AILab's `flash_attn.cute`) but replaces the mainloop schedule — internally a split-D (head-dim-in-smem) sequence-parallel walk — and it dequantizes the TKV codec inline during the MMA pipeline, no decompress buffer. The FA4 mainloop it grew out of is kept as a fallback: it owns `head_dim > 256` and the explicit `TKV_PREFILL_KERNEL=fa4` debug path; a decompress + stock FlashAttention path handles any shape both decline. See [Prefill performance](#prefill-performance) for the numbers.

## Performance

Cross-engine chat-serving baseline on **Qwen3.6-27B-AWQ-INT4**, TP=2 on 2× RTX 4090, via the canonical `chat_baseline.sh`/`serve_bench.sh` harness — locked clocks, single-tenant, **prefix-caching OFF**, cudagraph, real prompts, `vllm bench serve`.

**Decode throughput + KV density (no-MTP, 32k ctx):**

| 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×** |

tkv is **decode-neutral** (cudagraph amortizes the dequant) and packs **~3.75× more KV** into the same VRAM on both engines — the headline win at 27B. The two engines are at decode parity (~13 ms TPOT = the model's memory-bandwidth roofline).

**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 arbi prefill mainloop is an original schedule on the FlashAttention-4 CuTeDSL base. Since it and the FA4 mainloop feed the *same* loader / softmax / epilogue via one `cute_prefill` entry, they produce **bit-identical** output — so the wall-clock gap is a pure mainloop win at zero accuracy cost. Measured with `benchmarks/bench_prefill_splitd_vs_fa4.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 **FA4** (bf16 / tkv codec), 42 shapes, Δ=0 | vs stock **FA2** (bf16) |
|---|---|---|
| RTX 4090 (Ada, sm89) | **1.29× / 1.21×** geomean, 42/42 | net-positive (1.00× geomean; FA2 edges asym only) |
| RTX 5090 (Blackwell, sm120) | **1.27× / 1.25×** geomean, 42/42 | ahead (1.07× geomean, 29/42) |
| RTX A5500 (Ampere, sm86) | **1.44× / 1.18×** geomean, 42/42 | ahead (1.03× geomean, 27/42) |

Two things to read off this: the arbi mainloop **decisively and bit-identically beats the FA4 mainloop it replaced** across all three GPU generations (up to 1.6× at `head_dim=256`), and on raw bf16 it 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). Reproduce:

```bash
python benchmarks/bench_prefill_splitd_vs_fa4.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: `native_tq`, `flash_attn`, or `bypass`. |
| `TKV_PREFILL_ENGINE` | `fa4` | Prefill *family* (dispatcher), not the mainloop: `fa4` = the CuTeDSL `cute_prefill` path (default), `adaptive` = add the decompress+FA fallback for layouts CuTeDSL declines, `decomp_fa_main_only` = bench-only. Despite the name, `fa4` here does **not** force the FA4 mainloop — see `TKV_PREFILL_KERNEL`. |
| `TKV_PREFILL_KERNEL` | `""` (=`arbi`) | The CuTeDSL **mainloop**: `arbi`/unset = the **arbi prefill mainloop** (production default, every head-dim); `fa4` = the FA4 mainloop fallback (also auto-selected at `head_dim > 256`). |
| `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 backend. |
| `TKV_MLA_ROPE_HEAD_DIM` | `64` | RoPE head dimension for MLA latent + RoPE split. |

## Why a vLLM fork (for now)

`CacheDType` in `vllm/config/cache.py` is a Pydantic `Literal` validated at class-definition time, which blocks runtime registration of new KV-cache dtypes. Until that's relaxed upstream, we ship a fork. The fork is a thin overlay; full layout in [`docker/PATCHES.md`](docker/PATCHES.md). SGLang does not need a fork.

## 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

Mozilla Public License 2.0 (MPL-2.0). See `LICENSE` and `NOTICE`.
