Fitting 4.4x More Context on Jetson: Implementing TurboQuant on NVIDIA’s AGX Thor

Six iterations of KV-cache compression: from a naive PyTorch baseline to fused vectorized Triton kernels running real LLM inference on edge hardware.

The Problem

A 7B-parameter model like Qwen2.5 needs 14 GB of memory just for weights. But during inference, the KV cache (the key and value tensors stored for every token the model has recently seen) often exceeds the size of the model if we have multiple users. For example, at 32k context with 28 attention layers, the KV cache alone consumes more memory than the model weights.

In this blog we are going to explore this problem with the Nvidia Jetson AGX Thor. On the NVIDIA Jetson AGX Thor (a Blackwell-based edge AI computer with 128 GB of unified CPU/GPU memory (122 GB usable) and 20 streaming multiprocessors) every gigabyte of KV cache is a gigabyte unavailable for model weights or additional users.

TurboQuant (arXiv 2504.19874, Zandieh et al.) aims to solve this problem by compressing the KV cache to ~3 bits per element while preserving attention quality. The paper proves near-optimal distortion bounds and unbiased inner-product estimation, but does not provide GPU kernel implementations. This post describes an implementation on the Jetson AGX Thor, covering correctness validation and six iterations of performance optimization.

Note that we compare our Triton kernel against SDPA because (at the time of experimenting) vLLM does not natively support TurboQuant-compressed KV caches. Standard SDPA implementations expect uncompressed FP16/BF16 tensors and cannot process bit-packed TurboQuant data. We implemented this custom Triton kernel to bridge that gap, and we benchmark against SDPA as a performance baseline to track how closely our implementation approaches the hardware’s theoretical throughput limits.

Proving It Works: The Paper’s Guarantees

Before optimizing and attempting to solve this problem, the implementation must be validated against the paper’s formal guarantees. TurboQuant makes three claims; all three must hold, or the compression degrades attention quality rather than preserving it.

Theorem 1: MSE Distortion Bounds

The paper proves that for bit-width b, the MSE between original and reconstructed vectors is bounded:

D_mse <= (sqrt(3) * pi / 2) * (1/4^b)

For b = 1, 2, 3, 4, the specific bounds are 0.36, 0.117, 0.03, 0.009.

Our measurement on Thor (Section 2 of demo_jetson.py, 8192 unit vectors, D=128):

The measured distortions match the theoretical bounds closely. At b=1, we hit the bound exactly (ratio 1.00x), confirming the Lloyd-Max codebook is optimal. At b=3, the measured MSE slightly exceeds the theoretical expected value (1.14x). This is expected at d=128, where the Gaussian approximation to the true Beta distribution isn’t yet fully tight and finite-sample variance contributes.

What this proves: Our TurboQuantMSE implements Algorithm 1 correctly: random rotation via QR decomposition, Lloyd-Max optimal centroids for the Beta distribution f_X(x), and nearest-centroid quantization via searchsorted.

Theorem 2: Unbiased Inner Products

For KV-cache use, the central guarantee is that TurboQuant_prod provides an unbiased inner product estimator:

E[<y, x_tilde>] = <y, x>

A biased estimator would cause attention weights to systematically favor or penalize certain tokens, degrading generation quality.

Our measurement on Thor (Section 3 of demo_jetson.py, 512 keys, 32 queries, b=3):

Mean bias:      0.000047  (should be ~ 0)
Mean |error|:   0.02989
Relative error: 42.09%

The measured bias of 4.7e-05 is consistent with zero. This confirms that TurboQuantProd implements Algorithm 2 correctly: MSE quantization at (b-1) bits for the first stage, QJL sign projection on the residual for the second stage, and the sqrt(pi/2)/d scaling constant for unbiased dequantization.

Relevance to v6: The fused Triton kernel computes the same inner product as the paper’s Algorithm 2. The correspondence is:

1. Paper: <y, x_tilde> = <y, x_mse> + ||r|| * sqrt(pi/2)/d * <S * y, qjl>
2. Code: score = mse_score + qjl_score * res_norm * QJL_SCALE
3. Kernel: Same formula, but the loop over coordinates is vectorized into (BLOCK_N, PACKED_D) tile operations instead of scalar iterations.

Vectorization does not change the arithmetic — the same bit-unpacking, centroid gather, multiply-accumulate, and QJL sign extraction occur. The tiles are processed in parallel instead of serially.

Theorem 3: Lower Bounds

The paper proves that no quantizer can achieve MSE distortion better than 1/4^b, and no inner-product quantizer can beat ||y||^2 / (d * 4^b) (which reduces to 1/(d * 4^b) for unit-norm queries). TurboQuant is within a sqrt(3) * pi / 2 ≈ 2.7x constant factor of the MSE lower bound.

The gap between the upper and lower bounds is a small constant factor, leaving limited room for any quantizer to improve upon.

Triton Kernel Correctness

The fused decode kernel (kernel 3, which kernel 4 extends) is validated against the PyTorch reference path:

Fused decode vs PyTorch reference:
Max  |diff| = 0.00000
Mean |diff| = 0.00000

The Triton kernel produces bit-identical results to the PyTorch dequant-then-matmul path. All subsequent optimizations — vectorized tiles, fused hybrid, pre-unpacked values — preserve this exactness.

Compression Ratios

At long contexts (where compression matters most), TurboQuant achieves 5.1x per-layer compression. On Thor with a 16 GB model, fp16 KV can handle ~875k tokens. TurboQuant extends this to ~3.86 million tokens, a 4.4x increase in context capacity. (The capacity multiplier is lower than the 5.1x raw compression because per-token metadata, such as norms, residual norms, QJL signs, and value scales, adds overhead across all 28+ layers.)

Real Model Validation

Running Qwen2.5-7B-Instruct through vLLM with TurboQuant hooks on all 28 attention layers:

Prompt:    "What is KV cache compression? Explain in two sentences."
Baseline:  "KV cache compression is a technique used to reduce the
            size of key-value caches, making them more efficient..."
TurboQuant: [identical output]

In our evaluation, the experiments are run more like a smoke test, rather than a thorough evaluation for every checkpoint hit. The paper validates quality more rigorously: Needle-In-A-Haystack retrieval across 4k–104k tokens (perfect recall at 4x compression), and LongBench-E scores within ~1 point of uncompressed baselines at 3.5 effective bits. The paper also uses an outlier channel strategy, allocating higher bit precision to 32 outlier channels and lower to the remaining 96, which we don’t implement here. Our uniform 3-bit keys + 2-bit values is a simpler configuration.

On Thor, the 17.5 GB of paged KV cache was freed after inference, dropping VRAM from 31.4 GB to 14.6 GB.

The Optimization Journey: v1 to v6

v1: Baseline: 35x Slower Than SDPA (1.29 ms at ctx=1k)

The naive implementation follows the paper literally:

1. Dequantize all compressed keys: unpack indices, lookup centroids, rotate back via D x D matmul, scale by norms
2. torch.matmul(query, keys.T) for scores
3. Separate softmax, value dequant, value matmul

Bottleneck: The rotation back is O(N * D^2) per head. For N=16k, D=128, that is 268M FLOPs per head per decode step. The paper’s optimization (Section 3.1) is to rotate the query forward once instead; v1 does not use it.

v2: Triton Score Kernels: 14x Slower (0.58 ms)

Implemented the paper’s “rotate query forward” insight as Triton kernels. Three kernels:

1. MSE score: Unpack bit-packed indices, lookup centroids, dot with rotated query — all streaming through packed data, no key materialization
2. QJL score: Unpack sign bits, dot with sketched query, scale by residual norms
3. Fused decode: Online softmax over compressed KV with on-the-fly value dequantization (flash-attention style)

2.2x faster than v1. But the _fused_hybrid path (for caches with both compressed and buffer tokens) still used 19 separate CUDA kernel launches: 2 Triton + 3 cuBLAS + 1 softmax + 4 value unpack + 5 dtype casts + 2 concats

• 1 scalar multiply + 1 bmm.

v3: Combined Transform: ~Same Speed (0.56 ms)

Precomputed [Pi^T | S^T] as a single matrix, replacing two separate query matmuls with one. Added fp16 value aggregation. Run-to-run variance dominated at these timescales — no measurable improvement over v2.

v4: SDPA (Dead End): 36x Slower (1.51 ms)

Hypothesis: dequantize KV to fp16, then use PyTorch’s F.scaled_dot_product_ attention (FlashAttention). SDPA runs at 0.04 ms for ctx=1k. If decompress + SDPA < Triton, this path wins.

Result: Slower. The ~20 small PyTorch operations needed to prepare data for SDPA (unpack indices, gather centroids, multiply norms, unpack signs, cast dtypes, concatenate) each launch a separate CUDA kernel. At ~15 us per launch, 20 launches = 300 us of dispatch overhead alone.

Lesson: On this hardware, kernel launch count dominates over algorithmic complexity. The Triton kernel’s advantage is that it performs all work in a single launch.

v5: Extended 2D SDPA (Dead End): 32x Slower (1.34 ms)

Derived from the paper’s Theorem 2. The inner product decomposes as:

<y, x_tilde> = <[yPi^T | yS^T], [a | b]> in R^{2d}

By extending Q and K to 2D dimensions, a single SDPA call gives exact TurboQuant_prod scores with zero dequantization matmuls. Buffer keys use [k*Pi^T | 0] which preserves exact inner products since Pi is orthogonal.

Same kernel launch overhead problem as v4.

v6: Fused Vectorized Hybrid: 3x Slower (0.31 ms)

Three optimizations that stack multiplicatively:

1. Fused hybrid Triton kernel. A single kernel handles both compressed and buffer tokens with online softmax. Phase 1 iterates over compressed blocks (TQ score + value dequant + accumulate). Phase 2 iterates over buffer blocks (simple dot product + accumulate). The softmax state carries seamlessly across phases. This reduced kernel launches from 19 to 1.

2. Vectorized 2D tile score computation. Instead of iterating per-coordinate in a scalar loop (256 iterations for D=128), load entire (BLOCK_N, PACKED_D) tiles and process with Triton’s vectorized operations:

# import triton.language as tl
# Old: 256 serial iterations per token
for byte_idx in range(PACKED_D_MSE):
    packed = tl.load(MSE_ptr + byte_idx * stride)
    for sub in range(VALS_PER_BYTE):
        coord_idx = byte_idx * VALS_PER_BYTE + sub
        idx = (packed >> (sub * BITS)) & BIT_MASK
        centroid_val = tl.load(CENTROIDS_ptr + idx)
        q_val = tl.load(q_rot_base + coord_idx * stride)
        mse_scores += q_val * centroid_val

# New: VALS_PER_BYTE iterations per stage
coord_offsets = tl.arange(0, BLOCK_D)
all_packed = tl.load(MSE_ptr + ...)  # (BLOCK_N, PACKED_D_MSE) tile
for sub in range(VALS_PER_BYTE):
    idx = (all_packed >> (sub * BITS)) & BIT_MASK
    centroid_vals = tl.load(CENTROIDS_ptr + idx)  # L1-cached gather
    q_vals = tl.load(q_rot_base + coord_offsets * stride)
    mse_scores += tl.sum(centroid_vals * q_vals[None, :], axis=1)

12 tile-parallel iterations (representing 4 iterations for each of 3 sub-stages) instead of 256 scalar iterations. These 12 total iterations span both MSE and QJL stages. Each iteration processes a (64, 32) tile, which is 2048 elements in parallel.

3. Pre-unpacked values + combined Pi_S_T matmul. Pre-unpack value data to uint8 at prefill time (once) instead of at every decode step (saves 2 kernel launches). Combined [Pi^T | S^T] matmul for query transform (1 matmul instead of float cast + 2 matmuls).

Dead End: Byte-Level Lookup Table

A further attempt: precompute all 256 possible query-centroid products per packed byte position into a (PACKED_D, 256) lookup table. This reduces inner-loop ops by ~5x in theory. In practice, this was slower. Writing 48KB to global memory then performing scattered gathers has higher latency than the vectorized tile approach, which stays in registers.

The Numbers

All measurements on NVIDIA Jetson AGX Thor (SM 11.0 Blackwell, 20 SMs, 122.8 GB unified memory), 8 KV heads, head_dim=128, 3-bit keys + 2-bit values.

v1 -> v6: 4.2x faster at ctx=1k, 8.7x faster at ctx=16k.

The optimization progression within v6:

Can TurboQuant Match Scaled Dot-Product Attention?

The remaining 3-7x gap comes from:

• Centroid gather (~40%): Data-dependent indexed load per coordinate. Even with vectorized tiles, 2048 indexed loads from a 4-entry table have latency.
• Value dequantization (~25%): Loading uint8 + scales + zeros, multiply-add per tile.
• Query transform (~15%): One matmul outside the kernel.
• Softmax + accumulate (~10%): Irreducible overhead.
• Kernel launch (~10%): 2-3 remaining launches.

TQ reads 5x less KV data than fp16. At long contexts (>64k tokens) where fp16 KV exceeds L2 cache, SDPA becomes memory-bandwidth bound. TQ’s 5x bandwidth reduction should offset the decompression compute at sufficient context length.

In a full LLM inference pipeline, attention is 10-20% of decode time. A 3-7x attention overhead translates to a smaller end-to-end impact, while fitting 4.4x more context on the same hardware.

Real Model Running

./run_jetson.sh vllm

Loads Qwen2.5-7B-Instruct, hooks 28 attention layers with TurboQuant, generates text side-by-side:

Model: Qwen/Qwen2.5-7B-Instruct (14.25 GB, bf16)
Layers hooked: 28 attention layers

Output comparison:
  Baseline:   "KV cache compression is a technique used to reduce
               the size of key-value caches..."
  TurboQuant: "KV cache compression is a technique used to reduce
               the size of key-value caches..."
  [identical]

KV cache freed: 17,544 MB (31.4 GB -> 14.6 GB VRAM)

The model produces identical text. The KV cache was compressed from 17.5 GB (fp16 paged cache) to ~3 MB (TurboQuant compressed store), and the paged cache memory was freed.

Observations

1. Kernel launch overhead is the primary bottleneck. On the Jetson Thor has 20 SMs,each kernel launch incurs ~15 µs of CPU dispatch time. Consolidating 19 kernels into one provided a greater performance gain than any single algorithmic change.
2. Vectorized tiling is more efficient than scalar processing. Processing (64, 32) tiles in parallel (12 iterations) was 1.5x faster than per-coordinate scalar loops (256 iterations), at similar total FLOPs.
3. Kernel count matters more than the raw FLOPs. The extended 2D SDPA approach (v5) removed the dequantization matmuls, but the 20+ PyTorch operations needed to prepare the extended tensors made it 4x slower than the fused Triton kernel.
4. Move work to prefill time. Unpacking 2-bit values at prefill (once per token) instead of at each decode step (once per decode per token) eliminated 2 kernel launches per decode.
5. Register parallelism was better than memory lookups. Precomputing 256 byte-centroid products into a lookup table reduced theoretical operations by 5x, but the overhead of writing to and reading from the global memory made it slower in practice. It was better to keep the data in registers for tile-based operations, avoiding LUTs for KV dequantization.

The implementation validates the paper’s theorems. All three theorems (MSE bounds, unbiased inner products, near-optimal lower bounds) were validated on real hardware, with measurements closely matching the paper’s predictions (within expected finite-dimension variance at d=128).
Try It

# Synthetic benchmark (no model needed)
./run_jetson.sh demo

# Real LLM inference with Qwen2.5-7B
./run_jetson.sh vllm

# Interactive shell for development
./run_jetson.sh shell

Next Steps

If you need to compress KV caches and run large language models efficiently on edge hardware, ZEDEDA can help. Our Edge Intelligence Platform combines an AI-driven approach to build and orchestrate edge agents, models, applications, and infrastructure. This lets you build, test, and deploy AI agents and models on any edge hardware, using the same edge orchestration platform trusted by the world’s largest enterprises. Get in touch at https://zededa.com/contact-us/.

Built on NVIDIA Jetson AGX Thor (SM 11.0 Blackwell, 122.8 GB unified memory, JetPack 7.1). PyTorch 2.11, Triton 3.6, vLLM 0.17.1.

FAQ

What exactly is the KV cache, and why does it grow so large?

During inference, a transformer re-reads all previous tokens at every generation step. The KV cache stores the key and value matrices computed for those tokens so they don’t have to be recomputed. The problem is that the cache grows linearly with context length: for a model with 28 layers and 128-dimensional heads, every new token adds 28 × 2 × 128 floats (one key, one value, per layer). At 32k tokens in fp16, that’s over 14 GB, more than the model weights themselves.

Why does TurboQuant compress to “approximately 3 bits” rather than a clean power of two like 4 bits or 2 bits?

TurboQuant splits the key vector into two stages: a (b-1)-bit MSE quantizer for the bulk of the signal, plus a 1-bit QJL sign sketch on the residual. The paper’s sweet spot uses 2-bit MSE + 1-bit QJL = 3 bits per key element, and 2 bits per value element. That asymmetry (keys at 3, values at 2) comes from attention scores being more sensitive to key precision than to value precision.

What does “unbiased inner product estimator” mean in plain English?

When you compress a key vector and then compute its dot product with a query, you usually get a slightly wrong answer. “Unbiased” means that if you ran the compression many times with different random rotations, the average answer would be exactly right. No systematic skew means the attention softmax doesn’t consistently over- or under-weight certain tokens, so the model’s output distribution stays correct on average, even though any individual compressed inner product has some random error.

Why is kernel launch overhead so significant on Jetson? Can’t the GPU handle many small operations efficiently?

Each CUDA kernel launch incurs a fixed CPU-side dispatch cost, roughly 15 microseconds on Jetson Thor, before any GPU work begins. The Jetson AGX Thor has only 20 streaming multiprocessors, so small operations finish very quickly; the 15 µs dispatch overhead then dominates total time. v4 (dequantize to fp16, then use SDPA) required ~20 separate kernel launches just to prepare the data, adding ~300 µs of pure overhead before the attention kernel even started. Fusing everything into one Triton kernel eliminates 18 of those 19 launches and is the single largest contributor to v6’s speedup.

How is compressing the KV cache different from quantizing the model weights?

Weight quantization (as in GPTQ or AWQ) reduces the precision of the model’s parameters, e.g., the numbers that define the model’s behavior, and is done once at “model preparation” time. KV cache quantization compresses the dynamic, per-request, per-token cache that grows during inference. The two are independent and complementary: you can run a 4-bit quantized model with a TurboQuant-compressed KV cache at the same time.

If attention is only 10–20% of decode time, does a 3x slowdown in attention really matter?

At short contexts (1k-4k tokens) the answer is: not much. A 3x attention slowdown translates to roughly 1.3–1.6x end-to-end slowdown. But the value proposition is memory, not speed: TurboQuant lets you fit 4.4x more context on the same hardware, which unlocks workloads that would simply fail with fp16 KV, not just run slower. At very long contexts (>64k tokens), SDPA becomes memory-bandwidth bound and TurboQuant’s 5x reduction in KV data read may close most of the throughput gap.

What is “online softmax” and why does the fused kernel need it?

Standard softmax requires two passes over the data: one to find the maximum value (for numerical stability) and one to compute the exponentials and normalize. Online softmax (introduced by the FlashAttention paper) computes a running maximum and running sum in a single pass, which allows the attention computation to be fused with KV dequantization into one kernel sweep. Without it, the fused kernel would have to materialize all attention scores before normalizing, requiring a separate kernel launch and more memory.

Can I use TurboQuant with grouped-query attention (GQA) models, or only with standard multi-head attention?

The TurboQuant algorithm operates per key/value head independently, it doesn’t assume any relationship between heads. GQA models share KV heads across multiple query heads, so the same compressed key/value block is reused for several queries. TurboQuant should work with GQA with no algorithmic changes; you simply compress fewer KV heads. The implementation in this post hooks into vLLM’s attention layer, so GQA support depends on how vLLM exposes the KV tensors to the hook.