Skip to main content

Optimized CUDAgraph-enabled kernels and attention backend for vLLM, SGLang and more based on TurboQuant near-lossless KV cache compression. SOTA performance with Gemma 4, Qwen 3.6 and other modern LLMs.

Project description

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: MPL-2.0

# same flags on all three supported engines (arbi-serve / vLLM-fork / SGLang-fork):
<engine> serve <model> --kv-cache-dtype tkv --attention-backend turbo-attn

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 arbi split-D mainloop goes further, beating the FlashAttention-4 mainloop it grew out of 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 the FA4 / FA2 comparisons are measured 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. TurboKVCodec 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 TurboKVCodec

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

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

See examples/ 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. TurboKVCodec is a pure, framework-agnostic compressor: compress with TKV, decompress to bf16 / fp16, hand the result to vanilla flash_attn_varlen_func, FlashInfer, SGLang attention, anything that takes raw KV. See examples/06_tkv_codec_with_third_party_attention.py.

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

Run with Docker

Three inference servers are supported: vLLM, SGLang, and arbi-serve. 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

vLLM

Uses our vllm-fork rebased onto upstream v0.20.1 (small overlay — CacheDType Literal relaxation, per-group block-pool bookkeeping, named TURBO_ATTN slot in AttentionBackendEnum; full layout in docker/PATCHES.md).

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

Calibration

Pre-built calibration bundles for common models live in HuggingFace at arbi-dev/turbo-attn-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

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 the FlashAttention-4 and FlashAttention-2 comparisons live. FA4 and FA2 consume raw bf16 — and so does the arbi mainloop in bypass mode. The reported 1.2–1.4× over 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 arbi split-D mainloop

tkv/kernels/cuda/prefill/arbi_prefill.py, the unconditional production prefill path at every head-dim. An original schedule on the same CuTeDSL base as FlashAttention-4 (our tkv/kernels/cute/_fa/ fork of Dao-AILab's flash_attn.cute), not a tuning of it.

  • Split-D (head-dim-in-SMEM) sequence-parallel walk — beats the FA4 mainloop bit-identically by 1.2–1.4× geomean across Ada/Blackwell/Ampere (up to 1.6× at head_dim=256), and matches or beats mature stock FA2 on bf16. Because it shares FA4's loader/softmax/epilogue via one cute_prefill entry, the win is pure schedule at zero accuracy cost.
  • Inline TQ dequant in the MMA pipeline — K dequants into n-major SMEM as e4m3; V dequants transposed (d-major sVt) so P·V reads it directly with ldmatrix, avoiding a separate transpose pass, and the warps that own the dequant STS also own the pipelining.
  • fp8 (e4m3) split-D path (TKV_FA4_SPLITD_FP8) — an fp8 producer for the compute-bound long-context prefill regime.
  • 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.
  • Fallbacks: the FA4 mainloop (TKV_PREFILL_KERNEL=fa4) owns head_dim > 256; a decompress + stock FlashAttention path handles any shape both decline. See Prefill performance for numbers.

Decode — the unified split-K kernel

tkv/kernels/_cuda_decode_unified_splitk.cu, the sole production decode path.

  • 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 chat_baseline.sh/serve_bench.sh harness — 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 arbi prefill mainloop is an original schedule on the FlashAttention-4 CuTeDSL base. Since it and the FA4 mainloop feed the same loader / softmax / epilogue via one cute_prefill entry, they produce bit-identical output — so the wall-clock gap is a pure mainloop win at zero accuracy cost. These comparisons run in bypass (raw bf16) mode — FA4 and stock FA2 see the exact same uncompressed inputs the arbi mainloop does, so the delta isolates the schedule and nothing else (no codec confound). Measured with benchmarks/bench_prefill_splitd_vs_fa4.py (eager, torch.cuda.Event, median of 50, dense causal prefill across head_dim ∈ {64,128,256}, GQA ratios {1,4,7,8}, and both symmetric and asymmetric chunked-prefill shapes 512…131072):

GPU vs FA4 (bf16 / tkv codec), 42 shapes, Δ=0 vs stock FA2 (bf16)
RTX 4090 (Ada, sm89) 1.29× / 1.21× geomean, 42/42 net-positive (1.00× geomean; FA2 edges asym only)
RTX 5090 (Blackwell, sm120) 1.27× / 1.25× geomean, 42/42 ahead (1.07× geomean, 29/42)
RTX A5500 (Ampere, sm86) 1.44× / 1.18× geomean, 42/42 ahead (1.03× geomean, 27/42)

Two things to read off this: the arbi mainloop decisively and bit-identically beats the FA4 mainloop it replaced across all three GPU generations (up to 1.6× at head_dim=256), and on raw bf16 it is at or above parity with the mature stock FlashAttention-2 — net-positive on every GPU, with the sole soft spot being asymmetric long-prefill on Ada — while being the only one of the two that can serve the compressed TKV codec path at all (FA2 has no inline dequant). Reproduce:

python benchmarks/bench_prefill_splitd_vs_fa4.py --dtypes bf16,tq4 --fa2 --flashinfer

Configuration

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

Bit width and calibration

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

Engine selection

Variable Default Description
TKV_ENGINE "" (auto) Decode engine: native_tq, flash_attn, or bypass.
TKV_PREFILL_ENGINE fa4 Prefill family (dispatcher), not the mainloop: fa4 = the CuTeDSL cute_prefill path (default), adaptive = add the decompress+FA fallback for layouts CuTeDSL declines, decomp_fa_main_only = bench-only. Despite the name, fa4 here does not force the FA4 mainloop — see TKV_PREFILL_KERNEL.
TKV_PREFILL_KERNEL "" (=arbi) The CuTeDSL mainloop: arbi/unset = the arbi prefill mainloop (production default, every head-dim); fa4 = the FA4 mainloop fallback (also auto-selected at head_dim > 256).
TKV_PREFILL_BYPASS 1 First-chunk prefill bypass — skip codec on prompt-prefill, then re-rotate to TQ basis for decode.
TKV_FUSE_QROT "" (auto) Fused Q-rotation prologue. Decode-only.
TKV_O_PROJ_FOLD on Fold rotate_output into o_proj weights.
TKV_MTP_SPLITK 1 Use split-K decode kernel for MTP layers.
TKV_DECODE_SPLITS "" (autotune) Force decode-kernel split count.

Backend behaviour

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

MLA (DeepSeek V2/V3/V4)

Variable Default Description
TKV_MLA_ENABLE 0 Master switch for the MLA backend.
TKV_MLA_ROPE_HEAD_DIM 64 RoPE head dimension for MLA latent + RoPE split.

Why we ship engine forks (for now)

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. As upstream gains runtime registration, these forks shrink toward zero.

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

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

Project details


Release history Release notifications | RSS feed

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.26.0.tar.gz (822.2 kB 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.26.0-py3-none-any.whl (895.8 kB view details)

Uploaded Python 3

File details

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

File metadata

  • Download URL: turbo_attn-0.26.0.tar.gz
  • Upload date:
  • Size: 822.2 kB
  • Tags: Source
  • Uploaded using Trusted Publishing? No
  • Uploaded via: uv/0.11.32 {"installer":{"name":"uv","version":"0.11.32","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.26.0.tar.gz
Algorithm Hash digest
SHA256 34155434498bd84d18e7e4b0377bca4cc7de1a08a040859235bbd06f086b6744
MD5 6289a65d13ddf130137d1ac63223164e
BLAKE2b-256 d0695f41194bc0b1391960130d4e96b3cd4afcd6d294fec806318363e4016fc6

See more details on using hashes here.

File details

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

File metadata

  • Download URL: turbo_attn-0.26.0-py3-none-any.whl
  • Upload date:
  • Size: 895.8 kB
  • Tags: Python 3
  • Uploaded using Trusted Publishing? No
  • Uploaded via: uv/0.11.32 {"installer":{"name":"uv","version":"0.11.32","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.26.0-py3-none-any.whl
Algorithm Hash digest
SHA256 24bffec2f06a941a178975a84b2ac0576f4a685dbfdd310d68f74e25111d426e
MD5 31c881f0174881d3231f3bd6f94ccf37
BLAKE2b-256 ec404447bdbb8337c0b6f928e5d0896cf8e7f56ac5f979a6c40fad212de21397

See more details on using hashes here.

Supported by

AWS Cloud computing and Security Sponsor Datadog Monitoring Depot Continuous Integration Fastly CDN Google Download Analytics Pingdom Monitoring Sentry Error logging StatusPage Status page