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

819,707 views • 6 months ago •via X (Twitter)

66 Comments

ani's profile picture
ani6 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
ani6 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. Munoz6 months ago

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

ani's profile picture
ani6 months ago

@GoogleResearch :)

sensho's profile picture
sensho6 months ago

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

ani's profile picture
ani6 months ago

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

AiDevCraft's profile picture
AiDevCraft6 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
ani6 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 Krishnan6 months ago

@GoogleResearch Very cool

ani's profile picture
ani6 months ago

@GoogleResearch Thanks for checking my project out Sriram!

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

@GoogleResearch I did that

ani's profile picture
ani6 months ago

@GoogleResearch Cool stuff!!!

Abhishek Ekaanth's profile picture
Abhishek Ekaanth6 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
ani6 months ago

@humblefool @GoogleResearch Cool project!!!

Abhishek Ekaanth's profile picture
Abhishek Ekaanth6 months ago

@GoogleResearch Thanks

Hubert Thieblot's profile picture
Hubert Thieblot6 months ago

@GoogleResearch Cracked

Gavin's profile picture
Gavin6 months ago

@GoogleResearch Great work! Sent you a dm

ani's profile picture
ani6 months ago

@GoogleResearch 👋😁

E.H.X.Z.E.L's profile picture
E.H.X.Z.E.L6 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
ani6 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.L6 months ago

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

NOIR_BD's profile picture
NOIR_BD6 months ago

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

ani's profile picture
ani6 months ago

@GoogleResearch That's sick wtf Send your implementation!

surya's profile picture
surya6 months ago

@GoogleResearch This is sick

ani's profile picture
ani6 months ago

@GoogleResearch 🙏 following in your footsteps

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

@GoogleResearch Can you 10x compression ?

ani's profile picture
ani6 months ago

@GoogleResearch Worth taking a look 👀

John D. Pope  🦒's profile picture
John D. Pope 🦒6 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 Guo6 months ago

Insane work!!!

InternationalOptions's profile picture
InternationalOptions6 months ago

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

ani's profile picture
ani6 months ago

@GoogleResearch LOL 🙏

vivek varikuti's profile picture
vivek varikuti6 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 Gopinath6 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
ani6 months ago

@GoogleResearch 🙏 this is so cool!!!

Navjot Singh's profile picture
Navjot Singh6 months ago

@GoogleResearch yo are you at shopify right now ?

ani's profile picture
ani6 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 Singh6 months ago

@GoogleResearch yes i just remembered. check slack

Nengjia Li's profile picture
Nengjia Li6 months ago

@GoogleResearch 🦫bv ce

ani's profile picture
ani6 months ago

@GoogleResearch 🫡 you alr know

Ilyas's profile picture
Ilyas6 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 engineer6 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
ani6 months ago

@GoogleResearch 🙏 he knows

Eshan's profile picture
Eshan6 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 Verma6 months ago

@GoogleResearch gonna try this, seems awesome

ani's profile picture
ani6 months ago

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

ruslan's profile picture
ruslan6 months ago

@GoogleResearch this is so fire

ani's profile picture
ani6 months ago

@GoogleResearch 🙏

levi's profile picture
levi6 months ago

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

Utkarsh's profile picture
Utkarsh6 months ago

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

ani's profile picture
ani6 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 billa6 months ago

@GoogleResearch @MeetSiddhapura

adi's profile picture
adi6 months ago

@GoogleResearch super cool anirudh

ani's profile picture
ani6 months ago

@GoogleResearch Thanks big bro 🙏

chi's profile picture
chi6 months ago

@GoogleResearch insanely cool work

ani's profile picture
ani6 months ago

@GoogleResearch Thanks! 🙏

Armaan Sidhu's profile picture
Armaan Sidhu6 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 Ahmadnezhad6 months ago

@GoogleResearch @nvidia

ani's profile picture
ani6 months ago

@GoogleResearch @nvidia Appreciate the lob parsaxa

Satvik Garimella's profile picture
Satvik Garimella6 months ago

@GoogleResearch Fire asf g🔥

ani's profile picture
ani6 months ago

@GoogleResearch thanks big bro :)

Eyada's profile picture
Eyada6 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
ani6 months ago

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

luthira's profile picture
luthira6 months ago

@GoogleResearch 😼🤘

ani's profile picture
ani6 months ago

@GoogleResearch 🙏 the one and only

Udula Abeykoon's profile picture
Udula Abeykoon6 months ago

@GoogleResearch solved the turboquant problem

ani's profile picture
ani6 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