CoMem Explained — From Paper to Working Code in 10 Minutes

Short summary

CoMem (Comprehension Memory) caches intermediate transformer layer states instead of full KV pairs, cutting VRAM from 89 GB to 18 GB at 128k tokens with a 7.83x prefill speedup. The method exploits the observation that lower transformer layers handle semantic understanding while upper layers handle prediction, so you only need to cache and retrieve at the split boundary. The article includes a working PyTorch implementation of the cache, retrieve, and recompute pipeline tested on Qwen3-8B.

• •CoMem caches residual stream states at an intermediate split layer, reducing 128k-token VRAM from 89 GB to 18 GB
• •Achieves 7.83x prefill speedup and RULER score 97.05, outperforming KV-Direct by 10.6% on LoCoMo
• •Includes working PyTorch code for write, retrieve, and recompute operations

Generated with AI, which can make mistakes.