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_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)
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 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 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.27.1 (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 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 rebased onto upstream v0.5.17 (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 bundles are versioned under calibrations/. To
roll your own:
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
- 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.
TurboPrefillForward(+ 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 split-K / unified decode kernel. An original SIMT CUDA C++ kernel (
cuda_decode_splitk.py/_cuda_decode_unified_splitk.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_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 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/_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 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]). 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 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).
Download files
Download the file for your platform. If you're not sure which to choose, learn more about installing packages.
Source Distribution
Built Distribution
Filter files by name, interpreter, ABI, and platform.
If you're not sure about the file name format, learn more about wheel file names.
Copy a direct link to the current filters
File details
Details for the file turbo_attn-0.53.0.tar.gz.
File metadata
- Download URL: turbo_attn-0.53.0.tar.gz
- Upload date:
- Size: 3.6 MB
- Tags: Source
- Uploaded using Trusted Publishing? No
- Uploaded via:
uv/0.12.9 {"installer":{"name":"uv","version":"0.12.9","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}
File hashes
| Algorithm | Hash digest | |
|---|---|---|
| SHA256 |
7ac429bd841a1208f8f11207ec4ebcaad4ece98da1960f072670a2912f042c2f
|
|
| MD5 |
2d89b17ad0e7bd167f457aea1b3edbde
|
|
| BLAKE2b-256 |
6098d7e628b6ebebc7faefb34b1c58cd6b683ad907612c90da6e7e2c82596121
|
File details
Details for the file turbo_attn-0.53.0-py3-none-any.whl.
File metadata
- Download URL: turbo_attn-0.53.0-py3-none-any.whl
- Upload date:
- Size: 3.6 MB
- Tags: Python 3
- Uploaded using Trusted Publishing? No
- Uploaded via:
uv/0.12.9 {"installer":{"name":"uv","version":"0.12.9","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}
File hashes
| Algorithm | Hash digest | |
|---|---|---|
| SHA256 |
1353fd0c660b96f493e37d0bb508d8104e0dd300c284098fafd2811e87b556b4
|
|
| MD5 |
9ed7f6c888232b8d68c91ac0f97e9050
|
|
| BLAKE2b-256 |
fe62fe39aa758d0b2a036c85f9ae41b5cd5feab6bce7b7d258547d4a31aa5365
|