Skip to main content

Turbo Attention

CI PyPI 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 — 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
SGLang / tkv 73.1 307.5 1,282,029 3.76×
vLLM / bf16 75.0 338.6 317,290
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.
  • 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 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 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

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:

  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.

  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 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. 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:

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

  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. TurboAttnCute (+ 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 turbo_attn_simt kernel. An original SIMT CUDA C++ kernel (turbo_attn_simt.py / _turbo_attn_simt.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): 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 walk1.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 on sliding_window), plus the MTP-verify BLOCK_M=N path — the earlier standalone BLOCK_M=1 kernel was retired into this one.
  • Rotation fusionTKV_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 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:

  • vLLMCacheDType 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. 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

turbo_attn-0.55.2.tar.gz (3.7 MB view details)

Uploaded Source

Built Distribution

If you're not sure about the file name format, learn more about wheel file names.

turbo_attn-0.55.2-py3-none-any.whl (3.7 MB view details)

Uploaded Python 3

File details

Details for the file turbo_attn-0.55.2.tar.gz.

File metadata

  • Download URL: turbo_attn-0.55.2.tar.gz
  • Upload date:
  • Size: 3.7 MB
  • Tags: Source
  • Uploaded using Trusted Publishing? No
  • Uploaded via: uv/0.12.12 {"installer":{"name":"uv","version":"0.12.12","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

Hashes for turbo_attn-0.55.2.tar.gz
Algorithm Hash digest
SHA256 9cb9793d072527e103daf11f6f4d74bb6f811fc505d8ad3bb6f581617b146149
MD5 f4952e050c7b21c9c79a78dcd882aa5c
BLAKE2b-256 e3411db9e4c269c2b43e489b781b28f15c4945a7f5a18e5ed395c1696c12aaf4

See more details on using hashes here.

File details

Details for the file turbo_attn-0.55.2-py3-none-any.whl.

File metadata

  • Download URL: turbo_attn-0.55.2-py3-none-any.whl
  • Upload date:
  • Size: 3.7 MB
  • Tags: Python 3
  • Uploaded using Trusted Publishing? No
  • Uploaded via: uv/0.12.12 {"installer":{"name":"uv","version":"0.12.12","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

Hashes for turbo_attn-0.55.2-py3-none-any.whl
Algorithm Hash digest
SHA256 828ff6a3fdb891e07abba73d6f7fac2f504d13d28f17f5b9bf646d05f263b645
MD5 36bb1e66b7f5f211a200d66db2d5a139
BLAKE2b-256 37a4f9233b58cb93f35f55dd0e19cf594b7262138c00c9f1bb94597f991c7db0

See more details on using hashes here.

Release history Release notifications | RSS feed

This release

0.55.2 This release

2 files

0.55.1

2 files

0.55.0

2 files

0.54.0

2 files

0.53.0

2 files

0.52.0

2 files

0.51.0

2 files

0.50.0

2 files

0.49.0

2 files

0.48.0

2 files

0.47.0

2 files

0.46.0

2 files

0.44.0

2 files

0.43.1

2 files

0.43.0

2 files

0.42.0

2 files

0.41.2

2 files

0.41.1

2 files

0.41.0

2 files

0.40.1

2 files

0.40.0

2 files

0.39.0

2 files

0.38.6

2 files

0.38.5

2 files

0.38.4

2 files

0.38.3

2 files

0.38.2

2 files

0.38.1

2 files

0.38.0

2 files

0.37.8

2 files

0.37.7

2 files

0.37.6

2 files

0.37.5

2 files

0.37.4

2 files

0.37.3

2 files

0.37.2

2 files

0.37.1

2 files

0.37.0

2 files

0.36.1

2 files

0.36.0

2 files

0.35.4

2 files

0.35.3

2 files

0.35.2

2 files

0.35.1

2 files

0.35.0

2 files

0.34.0

2 files

0.33.3

2 files

0.33.2

2 files

0.33.1

2 files

0.33.0

2 files

0.32.6

2 files

0.32.5

2 files

0.32.4

2 files

0.32.3

2 files

0.32.2

2 files

0.32.1

2 files

0.32.0

2 files

0.31.1

2 files

0.31.0

2 files

0.30.2

2 files

0.30.1

2 files

0.30.0

2 files

0.29.0

2 files

0.28.4

2 files

0.28.3

2 files

0.28.2

2 files

0.28.1

2 files

0.28.0

2 files

0.27.0

2 files

0.26.0

2 files

0.25.0

2 files

0.24.4

2 files

0.24.3

2 files

0.24.0

2 files

0.23.0

2 files

0.22.0

2 files

0.21.0

2 files

0.20.3

2 files

0.20.2

2 files

0.20.1

2 files

0.20.0

2 files

0.19.1

2 files

0.19.0

2 files

0.18.1

2 files

0.18.0

2 files

0.17.1

2 files

0.17.0

2 files

0.16.7

2 files

0.16.6

2 files

0.16.4

2 files

0.16.3

2 files

0.16.2

2 files

0.16.1

2 files

0.16.0

2 files

0.15.1

2 files

0.15.0

2 files

0.14.3

2 files

0.14.2

2 files

0.14.1

2 files

0.14.0

2 files

0.13.0

2 files

0.12.4

2 files

0.12.3

2 files

0.12.2

2 files

0.12.1

2 files

0.12.0

2 files

0.11.2

2 files

0.11.1

2 files

0.11.0

2 files

0.10.0

2 files

0.9.2

2 files

0.9.1

2 files

0.9.0

2 files

0.8.9

2 files

0.8.8

2 files

0.8.7

2 files

0.8.6

2 files

0.8.5

2 files

0.8.4

2 files

0.8.3

2 files

0.8.2

1 file

0.8.1

2 files

0.8.0

2 files

0.7.0

2 files

0.6.5

2 files

0.6.4

2 files

0.6.3

2 files

0.6.2

2 files

0.6.1

2 files

0.6.0

2 files

0.5.1

2 files

0.5.0

2 files

0.4.1

2 files

0.4.0

2 files

0.3.2

2 files

0.3.1

2 files

0.3.0

1 file

0.2.0

2 files

0.1.2

2 files

0.1.1

2 files

0.1.0

2 files

Anthropic, PBC Visionary sponsor Bloomberg Visionary sponsor Hudson River Trading Visionary sponsor Meta Visionary sponsor NVIDIA Visionary sponsor Microsoft Sustainability sponsor Depot Continuous Integration AWS Cloud computing and Security Sponsor Datadog Monitoring Fastly CDN Google Download Analytics Sentry Error logging StatusPage Status page