Turbo Attention
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 — 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_BITSis 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. - 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 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 tkvflag 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.mdfor a complete worked fp8 example. - Hackable codec.
TkvCodecis 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.
BypassLoaderruns 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
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
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/ for runnable snippets and
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:
-
Codec → any attention backend.
TkvCodecis a pure, framework-agnostic compressor: compress with TKV, decompress to bf16 / fp16, hand the result to vanillaflash_attn_varlen_func, FlashInfer, SGLang attention, anything that takes raw KV. Seeexamples/06_tkv_codec_with_third_party_attention.py. -
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. Seedocs/writing_a_loader.mdfor 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/, andexamples/— public docs, deploy recipes, helper scripts, and runnable examples.
Run with Docker
Three inference servers are supported: vLLM, SGLang, and 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.json # the model's calibration store
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 solved at load from the calibration store. 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)
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:
- CUDA kernel JIT compile (nvcc/cutlass) — every decode/prefill kernel variant your deployment needs gets compiled to a
.soand cached underTKV_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. - 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 deployment has one shape, so it sweeps once; a SmartMix (per-layer heterogeneous bit-width) allocation 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 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 .sos at docker build time for a specific GPU arch, so a docker run from that image boots with no runtime nvcc:
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 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).
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 the model's calibration store. Unset → uniform K4V4 fallback (tests/CI only; production needs a store) |
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 rebased onto upstream v0.5.18 (small overlay — plugin registries for KV-cache dtypes and attention backends; full layout in docker/PATCHES.md).
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.
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 stores (one file per model, schema 56) live under calibrations/. To
roll your own:
python -m tkv.auto_calibrate --model Qwen/Qwen3.5-0.8B --output qwen3.5-0.8b.json
# add --smartmix to also measure the marginals a fractional TKV_BITS needs
How it works
- Rotate each KV vector with a fast Walsh–Hadamard transform.
- Normalize — store the magnitude as a single BF16 value.
- 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.
TurboAttnCute(+ its Bypass / Hybrid siblings), a tensor-core CuTeDSL kernel entered viaturbo_prefill. Tuning knobs areTKV_PREFILL_*; the engine value isTKV_PREFILL_ENGINE=turboand its row class isROW_PREFILL. - Decode — the
turbo_attn_simtkernel. An original SIMT CUDA C++ kernel (turbo_attn_simt.py/_turbo_attn_simt.cu). Tuning knobs areTKV_DECODE_*/TKV_MTP_*and its row class isROW_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): 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_attn_cute.py, the unconditional production prefill path at every head-dim. TurboAttnCute 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 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 for numbers.
Decode — the unified split-K kernel
tkv/kernels/_turbo_attn_simt.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 onsliding_window), plus the MTP-verifyBLOCK_M=Npath — the earlier standaloneBLOCK_M=1kernel was retired into this one. - Rotation fusion —
TKV_FUSE_QROTfolds the Q-rotation prologue in, androtate_outputis folded into theo_projweights (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_qrper 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 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:
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:
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]). Any value in the store's validated range: the per-layer (k_bits, v_bits) allocation is solved once at load from the store's measured_marginals. Hard error if the store has no marginals or the target is outside the validated range — no silent fallback. An int is uniform; k<K>v<V> is uniform asymmetric. |
TKV_CALIBRATION_FILE |
"" |
Path to the model's calibration store (centroids + per-channel scales, plus measured_marginals for a fractional TKV_BITS, written only by python -m tkv.auto_calibrate --smartmix). Required for production; store 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 and sglang-turbo:
- vLLM —
CacheDTypeinvllm/config/cache.pyis a PydanticLiteralvalidated at class-definition time, which blocks runtime registration of new KV-cache dtypes. Our overlay relaxes it, adds a namedTURBO_ATTNbackend 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. 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:
@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).
Release files for turbo-attn 0.56.1
For a detailed explanation of source distributions (sdists) and built distributions (wheels), please see the package formats documentation.
Source distribution (sdist)
| File | Size | Uploaded | |
|---|---|---|---|
| turbo_attn-0.56.1.tar.gz | 3.8 MB | Details |
Built distribution (wheel)
| File | Interpreter | ABI | Platform | Reset |
|---|---|---|---|---|
| turbo_attn-0.56.1-py3-none-any.whl | Python 3 | none | any | Details |
Total release size: 7.5 MB
Release files / turbo_attn-0.56.1.tar.gz
| Download URL | turbo_attn-0.56.1.tar.gz |
|---|---|
| Size | 3.8 MB |
| Tags | Source |
|
SHA-256 checksum How to use checksums |
ea8ccf08c8d293443cdaf12cfae01fc90d238e0136ccff7d6915777fb43afbc4
|
|
BLAKE2b-256 checksum How to use checksums |
c9203cd94bcdaffc632df7ed8889608c609591c020b2560224fe5ed0ae4e5712
|
| Upload date | |
|
Uploaded using Trusted Publishing? What is trusted publishing? |
No |
| Uploaded via |
uv/0.12.18 {"installer":{"name":"uv","version":"0.12.18","subcommand":["publish"]},"python":null,"implementation":{"name":null,"version":null},"distro":{"name":"Ubuntu","version":"24.04","id":"noble","libc":null},"system":{"name":null,"release":null},"cpu":null,"openssl_version":null,"setuptools_version":null,"rustc_version":null,"ci":true}
|
Release files / turbo_attn-0.56.1-py3-none-any.whl
| Download URL | turbo_attn-0.56.1-py3-none-any.whl |
|---|---|
| Size | 3.8 MB |
| Tags | Python 3 |
|
SHA-256 checksum How to use checksums |
f05a3dd5fd08a7c97aea90e4e0acb98badfac092fb40dee8ba357ef4edc1bb04
|
|
BLAKE2b-256 checksum How to use checksums |
00bb0e468c4ac3a1367917477b8d176b70ee001c1281ead075e113043f71ebe9
|
| Upload date | |
|
Uploaded using Trusted Publishing? What is trusted publishing? |
No |
| Uploaded via |
uv/0.12.18 {"installer":{"name":"uv","version":"0.12.18","subcommand":["publish"]},"python":null,"implementation":{"name":null,"version":null},"distro":{"name":"Ubuntu","version":"24.04","id":"noble","libc":null},"system":{"name":null,"release":null},"cpu":null,"openssl_version":null,"setuptools_version":null,"rustc_version":null,"ci":true}
|