Loading video...

Video Failed to Load

Go Home

I implemented Google Research's TurboQuant as a CUDA-native compression engine on Blackwell B200. 5x KV cache compression on Qwen 2.5-1.5B, near-loseless attention scores, generating live from compressed memory. 5 custom cuTile CUDA kernels ft: - fused attention (with QJL corrections) - online softmax -on-chip cache decompression - pipelined TMA...

820,150 views • 7 months ago •via X (Twitter)

66 Comments

ani's profile picture
ani7 months ago

The numbers: - 5.02x compression (3 bits/coordinate) ~0.985 cosine similarity vs FP16 attention (layers 7–27) - 144.7 tok/s on B200 - Fused attention kernel beats FP16 matmul at 4K+ sequence length Random rotation spreads outliers into a uniform bell curve → Lloyd-Max quantization bins where data is dense → QJL corrects the bias with 1-bit signs so attention scores come out lossless. The value cache gets 3-bit MSE, keys get 2-bit MSE + 1-bit QJL correction.

ani's profile picture
ani7 months ago

Started here, last week. Ended with ~2,000 lines of cuTile. Started with pseudocode and basic kernels. Profiled and iterated to optimize them. Shaved 6% compute time with fused V cache decompression and GPU memory swizzling.

Marc A. Munoz's profile picture
Marc A. Munoz7 months ago

@GoogleResearch I’m so happy I’m on this side of twitter

ani's profile picture
ani7 months ago

@GoogleResearch :)

sensho's profile picture
sensho7 months ago

@GoogleResearch did u start this from just writing it out by hand? holy

ani's profile picture
ani7 months ago

@GoogleResearch was j putting thoughts done on how a basic version should work lol

AiDevCraft's profile picture
AiDevCraft7 months ago

5x KV cache compression with near-lossless attention scores is the kind of result that directly translates to longer context windows at the same memory budget. The B200's TMA hardware making this practical on-chip is what separates it from software-only approaches that eat into throughput.

ani's profile picture
ani7 months ago

@GoogleResearch Couldn't have said it better. The newer chips can handle larger throughput- that's how they're designed. It's just a matter of recognizing it and adjusting your kernels accordingly.

Sriram Krishnan's profile picture
Sriram Krishnan7 months ago

@GoogleResearch Very cool

ani's profile picture
ani7 months ago

@GoogleResearch Thanks for checking my project out Sriram!

CV.YH's profile picture
CV.YH7 months ago

@GoogleResearch I did that

ani's profile picture
ani7 months ago

@GoogleResearch Cool stuff!!!

Abhishek Ekaanth's profile picture
Abhishek Ekaanth7 months ago

@GoogleResearch I did the same implementation on one of my projects, TQ makes it scalable for edge device LLMs

ani's profile picture
ani7 months ago

@humblefool @GoogleResearch Cool project!!!

Abhishek Ekaanth's profile picture
Abhishek Ekaanth7 months ago

@GoogleResearch Thanks

Hubert Thieblot's profile picture
Hubert Thieblot7 months ago

@GoogleResearch Cracked

Gavin's profile picture
Gavin7 months ago

@GoogleResearch Great work! Sent you a dm

ani's profile picture
ani7 months ago

@GoogleResearch 👋😁

E.H.X.Z.E.L's profile picture
E.H.X.Z.E.L7 months ago

@GoogleResearch This is impressive bro, I had to feed the papers into Notebook LM to understand what was going on. I'm relatively new into AI research and Model engineering, and seeing the speed at which you implemented this is mind blowing and a huge inspiration. Let's connect 🖖💪

ani's profile picture
ani7 months ago

@N_Ehxzel7 @GoogleResearch Yeah the paper was a lot to take in. I spent an entire day breaking down the math, tryna put the ideas together. Once things clicked, I put my ideas down on paper and began the kernels :)

E.H.X.Z.E.L's profile picture
E.H.X.Z.E.L7 months ago

@GoogleResearch Now that's some Tony Stark type shii Bro 🤣🤣🤣 Impressive.

NOIR_BD's profile picture
NOIR_BD7 months ago

@GoogleResearch My implementation on an A10 got me till here (I did 4 bit tho)

ani's profile picture
ani7 months ago

@GoogleResearch That's sick wtf Send your implementation!

surya's profile picture
surya7 months ago

@GoogleResearch This is sick

ani's profile picture
ani7 months ago

@GoogleResearch 🙏 following in your footsteps

John D. Pope  🦒's profile picture
John D. Pope 🦒7 months ago

@GoogleResearch Can you 10x compression ?

ani's profile picture
ani7 months ago

@GoogleResearch Worth taking a look 👀

John D. Pope  🦒's profile picture
John D. Pope 🦒7 months ago

@GoogleResearch I didn’t get it to compile to 4 bit - Claude got stuck in a loop - but 3bit results looked stellar - johndpope/llama-cpp-turboquant branch feature/planarquant-kv-cache, commit 20efe75 -

Kenny Guo's profile picture
Kenny Guo7 months ago

Insane work!!!

InternationalOptions's profile picture
InternationalOptions7 months ago

@GoogleResearch Anirudh, how many offers by how to join AI labs and start ups? :-) I’d wager >10!!

ani's profile picture
ani7 months ago

@GoogleResearch LOL 🙏

vivek varikuti's profile picture
vivek varikuti7 months ago

@GoogleResearch I implemented this from scratch on H100. Got 44-59% KV-cache reduction with exact prefill fidelity across 5 model families. Full benchmarks in my thread i found something that no one noticed

Ashwin Gopinath's profile picture
Ashwin Gopinath7 months ago

@GoogleResearch Thank for implementation! It helped us as we figured out how to beat @googleresearch’s provably near-optimal 5.02× compression ratio :-)

ani's profile picture
ani7 months ago

@GoogleResearch 🙏 this is so cool!!!

Navjot Singh's profile picture
Navjot Singh7 months ago

@GoogleResearch yo are you at shopify right now ?

ani's profile picture
ani7 months ago

@GoogleResearch Bro we literally went to the Sidekick team burst dinner together last month LOL Tall guy, intern who was from Turner lmao

Navjot Singh's profile picture
Navjot Singh7 months ago

@GoogleResearch yes i just remembered. check slack

Nengjia Li's profile picture
Nengjia Li7 months ago

@GoogleResearch 🦫bv ce

ani's profile picture
ani7 months ago

@GoogleResearch 🫡 you alr know

Ilyas's profile picture
Ilyas7 months ago

@GoogleResearch impressive work ! but i noticed in your video that the cosine similarity is not 0.9 across the board as declared, seems to be very low on some layers and some weights especially the first layers where it dips to 0.29 the max absolute weight error of 1 is also concerning

b/acc, context platform engineer's profile picture
b/acc, context platform engineer7 months ago

@GoogleResearch Absolute 🔥 Anirudh Combining this with Dynamo KVBM and 1000x memory capacity will unlock enormous agent swarm scale and efficiency 💪

ani's profile picture
ani7 months ago

@GoogleResearch 🙏 he knows

Eshan's profile picture
Eshan7 months ago

5x KV cache compression with near-lossless attention scores means you're fitting roughly 5x the context into the same GPU memory without degrading output quality. for anyone not tracking this, the practical implication is that models that currently need 8 GPUs to serve could potentially run on 2. the cost reduction isn't incremental, it's structural. this is the kind of work that changes who can afford to run large models

Yashvardhan Verma's profile picture
Yashvardhan Verma7 months ago

@GoogleResearch gonna try this, seems awesome

ani's profile picture
ani7 months ago

@GoogleResearch It is! I'm looking for feedback and ways to improve. Pls try and lmk 👍

ruslan's profile picture
ruslan7 months ago

@GoogleResearch this is so fire

ani's profile picture
ani7 months ago

@GoogleResearch 🙏

levi's profile picture
levi7 months ago

@GoogleResearch so sick!! you're killing it man

Utkarsh's profile picture
Utkarsh7 months ago

@GoogleResearch Yoo solid work. How does the chip cache decompression work?

ani's profile picture
ani7 months ago

Great question :) Instead of decompressing in a separate kernel, we do it inside the attention loop, right before they hit the Tensor Cores. This pushes our kernel more into compute-bound instead of memory-bound. A couple of things: 1.) I load the Pi (rotation matrix) beforehand. It's constant across all streamed KV blocks. 2.) For every value block, I use the indices and norms for Llyod-Max centroid lookup. 3.) Then, I unrotate k_mse (centroid-based coords) and multiply with the norm to turn it back into the full 128 dimension vector (not normalized anymore). Because we reconstructed V, I just accumulate it to finish our attention score. FP16 decompressed values never touch HBM.

Pakhandi billa's profile picture
Pakhandi billa7 months ago

@GoogleResearch @MeetSiddhapura

adi's profile picture
adi7 months ago

@GoogleResearch super cool anirudh

ani's profile picture
ani7 months ago

@GoogleResearch Thanks big bro 🙏

chi's profile picture
chi7 months ago

@GoogleResearch insanely cool work

ani's profile picture
ani7 months ago

@GoogleResearch Thanks! 🙏

Armaan Sidhu's profile picture
Armaan Sidhu7 months ago

A Waterloo undergrad just built the first working implementation of Google's most disruptive AI paper. In one week. Memory stocks crashed. Investors got it exactly backwards. TurboQuant compresses the KV cache — the part of an AI model that eats your GPU memory during long conversations — from 16 bits down to 3 bits per value. 5x compression. Near-zero quality loss. No retraining needed. Drop it into any existing model and it just works. Micron, SK Hynix, Samsung, SanDisk all dumped 3-6% in a single day. Wall Street's logic: if AI needs less memory, sell memory stocks. That logic has been wrong every single time in the history of computing. This is Jevons paradox — the 160-year-old principle that says when you make a resource more efficient, you don't use less of it. You use more. Coal engines got efficient, coal consumption exploded. Bandwidth got cheaper, data consumption went vertical. Storage got dense, we filled it with 4K video. Every time. Here's what TurboQuant actually does to demand. A 70B parameter model that used to need a $10,000 workstation to run at long context? Now runs on a $1,500 consumer GPU. A model that choked at 32K tokens? Now handles 100K+. A smartphone that could barely run a 3B toy model? Now runs 7-13B with thousands of tokens of usable context. You didn't reduce demand. You just made AI accessible to a billion more devices. Every one of those devices now needs memory. Every user who couldn't run local AI before now can — and will. And every developer who hits the new ceiling will immediately ask for more. Anirudh BV — a student at UWaterloo — took Google's paper, wrote 2,000 lines of CUDA, and had it running on Nvidia's Blackwell B200 within a week. From pseudocode to live generation from compressed memory. If a college kid can ship this before most corporate labs, the adoption curve isn't years out. It's already here. The investors who sold memory stocks on this news will be buying them back within two quarters. Efficiency doesn't kill demand. It creates the demand that didn't exist yet.

Parsa Ahmadnezhad's profile picture
Parsa Ahmadnezhad7 months ago

@GoogleResearch @nvidia

ani's profile picture
ani7 months ago

@GoogleResearch @nvidia Appreciate the lob parsaxa

Satvik Garimella's profile picture
Satvik Garimella7 months ago

@GoogleResearch Fire asf g🔥

ani's profile picture
ani7 months ago

@GoogleResearch thanks big bro :)

Eyada's profile picture
Eyada7 months ago

@GoogleResearch Implementing a cutting-edge paper like TurboQuant end-to-end on new hardware (Blackwell) is the real test of a system's potential. Moving from theory to live generation is impressive.

ani's profile picture
ani7 months ago

@GoogleResearch I appreciate it! Had lots of fun picking this one out

luthira's profile picture
luthira7 months ago

@GoogleResearch 😼🤘

ani's profile picture
ani7 months ago

@GoogleResearch 🙏 the one and only

Udula Abeykoon's profile picture
Udula Abeykoon7 months ago

@GoogleResearch solved the turboquant problem

ani's profile picture
ani7 months ago

@GoogleResearch Had to solve it bc @nengjiali already solved the drift problem and you solved the seo problem. Seo and drift are old news.

Related Videos

Introducing The AI CUDA Engineer: An agentic AI system that automates the production of highly optimized CUDA kernels. The AI CUDA Engineer can produce highly optimized CUDA kernels, reaching 10-100x speedup over common machine learning operations in PyTorch. Our system is also able to produce highly optimized CUDA kernels that are much faster than existing CUDA kernels commonly used in production. We believe that fundamentally, AI systems can and should be as resource-efficient as the human brain, and that the best path to achieve this efficiency is to use AI to make AI more efficient! We are excited to publish our paper, The AI CUDA Engineer: Agentic CUDA Kernel Discovery, Optimization and Composition. We also release a dataset of over 17,000 verified CUDA kernels produced by The AI CUDA Engineer. Paper: Kernel Archive Webpage: HuggingFace Dataset: The AI CUDA Engineer utilizes evolutionary LLM-driven code optimization to autonomously improve the runtime of machine learning operations. Our system is not only able to convert PyTorch code into CUDA kernels, but through the use of evolution, it can also optimize the runtime performance of CUDA kernels, fuse multiple operations, and even discover novel solutions for writing efficient CUDA operations by learning from past innovations! We believe The AI CUDA Engineer opens a new era of AI-driven acceleration of AI and automated inference time optimization. We (Robert Lange, Aaditya Prasad 🇺🇸, sssss, Maxence Faldor, Yujin Tang, hardmaru) are excited to continue Sakana AI's mission of leveraging AI to improve AI.

Sakana AI

1,160,959 views • 1 year ago

A tricky LLM interview question: You're serving a reasoning model on vLLM, and it keeps running out of GPU memory on long traces. So you add KV cache compression and evict 90% of the cached tokens. VRAM usage stays as is and GPU still runs out of memory. Why? (answer below) Evicting 90% of the KV cache can free almost none of the memory it was using. This sounds counterintuitive, but it follows directly from how production servers store the cache today. The KV cache grows with every token a model generates. Each token appends its key and value vectors across every layer, and nothing is freed while generation continues. This is the dominant memory cost for reasoning models. If a 32K-token CoT caches ~32K tokens of KV vectors, a Qwen3-32B with 4-bit weights will run out-of-memory around 24K tokens on a 24GB GPU. One obvious solution is to keep the important tokens and drop the rest, since attention is sparse enough to allow it. But this does not solve the memory problem yet. The reason is paged attention, which is the memory manager behind vLLM and most production servers. Under the hood, it splits GPU memory into fixed physical blocks, each one holds the KV for about 16 tokens. This block returns to the allocator only when every slot inside it is empty. Since the eviction logic selects tokens by importance, and such tokens are scattered across blocks... ...so despite eviction, almost every block is left with at least some survivor tokens. For instance, if the logic evicts 14k of 16k tokens across 1,000 blocks, most likely every block will still have a token. This means the allocator frees almost nothing. Placing the new tokens into those freed slots is not ideal because it breaks the cache's layout. Say token 16,001 arrives, and it's placed in the slot the 40th token used to hold. The cache now reads position 38, then 16,001, then 41, so the cache is no longer in token order. Attention can still compute the right answer from that, but only if every slot now carries a separate note recording which position it actually holds. This introduces another bookkeeping cost that an in-order layout inherently avoids. So the cache is logically 90% smaller and still physically the same size. Many compression results miss this because they measure on pre-allocated contiguous tensors rather than a paged server. There's another problem. Eviction methods pick which tokens to keep by looking at the attention scores themselves (as expected). But fast attention kernels used in production, like FlashAttention, never save those scores. They compute attention in small pieces and throw the full score grid away as they go, which is also why they're fast. So the exact signal eviction methods need isn't available in memory. The workaround is to fall back to eager attention and build the full matrix, which gives up the speed FlashAttention was there to provide. NVIDIA published a method called TriAttention to solve both these problems. It never needs attention scores. Instead, it scores tokens from the geometry of the model's key and query vectors before RoPE is applied, where those vectors sit in stable clusters. For the memory problem, it runs a compaction pass every 128 decoded tokens. The surviving tokens slide forward to close the holes eviction creates, so whole blocks empty out and return to the allocator while the cache stays in token order. On long reasoning traces, the approach matches full-attention accuracy while decoding 2.5x faster and using 10.7x less KV memory. KV cache compression is a big infrastructure problem. The number that decides whether it works is the count of freed blocks, not the count of evicted tokens. You can find the NVIDIA write-up here: I wrote a first-principles breakdown of how the KV cache works. It walks through why the model stores keys and values at all, why the cache grows with every token, and a comparison of LLM generation speed with and without KV caching. Read it below.

Avi Chawla

273,899 views • 4 months ago

A tricky LLM interview question: You're serving a reasoning model on vLLM, and long traces keep pushing it out of GPU memory. You add KV cache compression and evict 90% of the cached tokens. VRAM usage doesn't move, and the GPU still runs out of memory. Why? (answer below) Evicting 90% of the KV cache can free close to zero bytes on a paged server. During decoding, every new token appends its key and value vectors at every layer. vLLM holds all of them until the request finishes. For reasoning models with long chains of thought, this cache is the main memory cost. A Qwen3-32B with 4-bit weights on a 24GB GPU needs KV for every token of a 32K-token trace, and it hits OOM around 24K tokens. Attention over long contexts is sparse, so you can keep the high-importance tokens and drop the rest with little accuracy loss. That reduces the number of cached tokens but it doesn't necessarily reduce allocated memory. vLLM manages the cache with PagedAttention. It splits GPU memory into fixed physical blocks of 16 token slots each. The allocator reclaims a block only when all 16 slots are empty. Importance-based eviction keeps tokens scattered across the full sequence, so the retained tokens are spread over nearly every block. Evict 14K of 16K tokens across 1,000 blocks, and almost every block still holds at least one token. So the allocator reclaims close to nothing. You could write new tokens into the empty slots. If token 16,001 goes into the slot token 40 held, the block now reads positions 38, 16,001, 41. The cache is no longer in position order. Attention still gives the correct result, but each slot now needs a stored position index, and every kernel has to read it. An in-order layout doesn't need that metadata. So the cache shrinks 90% logically and 0% physically. Many compression papers miss this because they benchmark on pre-allocated contiguous tensors instead of a paged server. The second problem is the eviction signal. Most eviction methods rank tokens by their attention scores. FlashAttention never materializes those scores. It computes attention tile by tile in on-chip SRAM and discards each score tile after use, which is where its speed comes from. To get the scores, you have to fall back to eager attention and build the full score matrix in HBM, which removes the speedup FlashAttention provides. NVIDIA's TriAttention handles both problems. It ranks tokens without attention scores. It uses the geometry of the query and key vectors before RoPE is applied, where they form stable clusters. Every 128 decoded tokens, it runs a compaction pass. Retained tokens move forward to fill the evicted slots, so trailing blocks empty out completely and return to the allocator, and the cache stays in position order. On long reasoning traces, it matches full-attention accuracy with 2.5x faster decoding and 10.7x less KV memory. For KV cache compression on a real server, the metric to track is blocks freed, not tokens evicted. NVIDIA's write-up: I also wrote a first-principles breakdown of the KV cache. It covers why the model caches keys and values, why the cache grows with every token, and how generation speed compares with and without caching. Read it below.

Avi Chawla

40,630 views • 6 days ago