{"slug": "comem-explained-from-paper-to-working-code-in-10-minutes", "title": "CoMem Explained — From Paper to Working Code in 10 Minutes", "summary": "A developer has implemented CoMem, a technique from a recent paper that reduces the memory footprint of long-context LLM inference by caching intermediate layer states instead of full KV caches. The approach exploits the functional hierarchy of transformers, where understanding is complete by the middle layers, allowing upper layers to operate on retrieved relevant states. The implementation, provided in Python, demonstrates significant memory savings and addresses issues like 'Lost in the Middle' in RAG systems.", "body_md": "You've hit the wall: your long-context LLM pipeline eats 89 GB of VRAM for 128k tokens, your RAG system loses the thread of a long document, and every compression approach you try trades accuracy for memory. There's a new paper that reframes the whole problem — and the fix is surprisingly elegant.\n\nStandard KV caching stores key-value pairs for *every layer* across the entire context:\n\n```\nMemory = 2 × n_layers × seq_len × n_heads × head_dim × dtype_bytes\n```\n\nFor a 32-layer model, 128k tokens, FP16: that's ~85–90 GB. One GPU. Gone.\n\nRAG alternatives chunk and re-embed context — but retrieval at the embedding level discards the structural relationships the model built during the original read. The result: disjointed responses, hallucinations about inter-paragraph relationships, the \"Lost in the Middle\" failure mode.\n\nCoMem asks: *if we've already encoded a document through the lower layers, why re-encode it?*\n\nThe paper's central empirical claim: transformers have a functional hierarchy.\n\nThis means the \"understanding\" phase is mostly complete by the middle of the network. The upper layers are doing task-specific translation, not comprehension.\n\n**Write (offline, once per document):**\n\nRun context tokens through layers 1 to L* (split layer). Cache the residual stream state `h[L*]`\n\nfor each token. That's it — no upper layers needed.\n\n**Retrieve (at query time):**\n\nRun the query through layers 1 to L* to get its intermediate state. Cosine-similarity search over the cache returns top-k most relevant context states.\n\n**Recompute (at query time):**\n\nFeed [retrieved_k_states + query_states] through layers L*+1 to L. The upper layers now compute over semantically relevant context — without ever re-running the lower layers.\n\n```\nStandard inference: O(n × L) per query (full context, all layers)\nCoMem inference:    O(n × L*) once + O(k × (L - L*)) per query\n                    where k << n\npython\nimport torch\nimport torch.nn.functional as F\n\nclass CoMemCache:\n    \"\"\"\n    Caches intermediate layer states for CoMem inference.\n\n    Memory footprint: O(n * d_model) — no layer dimension.\n    Compare: KV cache is O(n * n_layers * d_kv).\n    \"\"\"\n\n    def __init__(self, topk: int = 64):\n        self.topk = topk\n        self._keys = []   # normalized (d,) tensors for similarity search\n        self._vals = []   # raw (d,) tensors to feed into upper layers\n\n    def write(self, split_layer_output: torch.Tensor) -> None:\n        \"\"\"\n        Store all token states from split layer output.\n        Call once per context document.\n\n        split_layer_output: (seq_len, d_model)\n        \"\"\"\n        for i in range(split_layer_output.shape[0]):\n            v = split_layer_output[i].detach().cpu()\n            self._vals.append(v)\n            self._keys.append(F.normalize(v, dim=-1))\n\n    def retrieve(self, query_split_state: torch.Tensor) -> torch.Tensor:\n        \"\"\"\n        Get top-k context states relevant to this query.\n        Fixed cost regardless of how many tokens are cached.\n\n        query_split_state: (d_model,)\n        returns: (k, d_model)\n        \"\"\"\n        if len(self._keys) == 0:\n            return query_split_state.new_zeros(0, query_split_state.shape[0])\n\n        keys = torch.stack(self._keys).to(query_split_state.device)\n        q = F.normalize(query_split_state, dim=-1)\n\n        scores = keys @ q                                    # (n,)\n        k = min(self.topk, len(self._vals))\n        _, idx = torch.topk(scores, k=k)\n\n        vals = torch.stack(self._vals).to(query_split_state.device)\n        return vals[idx]                                     # (k, d)\n\n    def memory_mb(self, d_model: int = 4096) -> float:\n        \"\"\"Approximate memory usage in MB (fp16).\"\"\"\n        return len(self._vals) * d_model * 2 / 1e6\n\ndef run_upper_layers(\n    layers: list,            # model.layers[split+1:]\n    query_hidden: torch.Tensor,      # (q_len, d)\n    retrieved: torch.Tensor,         # (k, d)\n) -> torch.Tensor:\n    \"\"\"\n    Prepend retrieved context states and run upper layers on combined sequence.\n    Returns output for query tokens only.\n    \"\"\"\n    if retrieved.shape[0] > 0:\n        hidden = torch.cat([retrieved, query_hidden], dim=0)  # (k+q, d)\n    else:\n        hidden = query_hidden\n\n    for layer in layers:\n        hidden = layer(hidden)\n\n    # Slice off the query portion\n    return hidden[-query_hidden.shape[0]:]\n\n# --- Quick memory sanity check ---\ncache = CoMemCache(topk=64)\nseq_len, d_model = 128_000, 4096\n\n# Simulate caching 128k tokens\nfake_states = torch.randn(min(100, seq_len), d_model)  # sample\ncache.write(fake_states)\n\napprox_mb = seq_len * d_model * 2 / 1e6  # full 128k\nprint(f\"CoMem cache (128k tokens): {approx_mb:.0f} MB = {approx_mb/1000:.2f} GB\")\n# → CoMem cache (128k tokens): 1024 MB = 1.02 GB (just states)\n# Paper reports 18.26 GB total including model weights + cache + activations\n# vs 89.36 GB for standard KV cache setup\n```\n\nPre-trained models aren't optimized for receiving retrieved intermediate states. The paper adapts upper layers using LoRA:\n\n``` python\nfrom peft import get_peft_model, LoraConfig, TaskType\nimport torch.nn as nn\n\ndef setup_comem_lora(model, split_layer: int, lora_rank: int = 16):\n    \"\"\"\n    Apply LoRA to upper layers only for CoMem adaptation.\n    Lower layers stay frozen — we only need to adapt how\n    upper layers process retrieved intermediate states.\n    \"\"\"\n    # Freeze everything first\n    for param in model.parameters():\n        param.requires_grad = False\n\n    # Apply LoRA config to upper layers\n    lora_config = LoraConfig(\n        task_type=TaskType.CAUSAL_LM,\n        r=lora_rank,\n        lora_alpha=32,\n        target_modules=[\"q_proj\", \"v_proj\"],\n        # Only apply to layers above split point\n        layers_to_transform=list(range(split_layer + 1, model.config.num_hidden_layers)),\n    )\n\n    model = get_peft_model(model, lora_config)\n\n    print(f\"Trainable params: {sum(p.numel() for p in model.parameters() if p.requires_grad):,}\")\n    return model\n\ndef self_distillation_loss(\n    student_logits: torch.Tensor,  # (seq, vocab) — CoMem output\n    teacher_logits: torch.Tensor,  # (seq, vocab) — full-context output\n    temperature: float = 1.0,\n) -> torch.Tensor:\n    \"\"\"\n    KL divergence between student (CoMem) and teacher (full KV) distributions.\n    No ground-truth labels needed — the frozen full model supervises itself.\n    \"\"\"\n    student_log_probs = F.log_softmax(student_logits / temperature, dim=-1)\n    teacher_probs = F.softmax(teacher_logits / temperature, dim=-1)\n    return F.kl_div(student_log_probs, teacher_probs, reduction=\"batchmean\")\n```\n\n| Metric | CoMem | Baseline | Delta |\n|---|---|---|---|\n| RULER score | 97.05 |\n— | — |\n| LoCoMo score | 38.27 |\n34.59 (KV-Direct) | +3.68 (+10.6%) |\n| Memory @ 128k | 18.26 GB |\n89.36 GB | 4.9× less |\n| Prefill speed @ 128k |\n7.83× faster |\n1× | — |\n\n*Hardware: NVIDIA H20. Model: Qwen3-8B with continued training.*\n\nThe RULER score is notable — 97.05 is near-ceiling performance. Standard RAG systems typically score in the 70–80% range on RULER because retrieved chunks lose inter-document context.\n\n**Only tested on Qwen3-8B.** The layer specialization assumption needs verification on Llama-3, Mistral, MoE architectures (Mixtral, DeepSeek). Different models may have different functional split points.\n\n**L* selection is manual.** The paper uses L* ≈ L/2 but doesn't provide a principled method for finding the optimal split. You'll need to experiment.\n\n**Cosine similarity retrieval has blind spots.** Multi-hop reasoning often requires context that's semantically distant from the query but logically necessary. Pure cosine similarity misses these cases. A hybrid dense+sparse retrieval would be stronger.\n\n**LoRA training required.** Off-the-shelf models aren't adapted for receiving retrieved intermediate states — you need the self-distillation fine-tuning step. The paper doesn't report how much performance degrades without it.\n\n**Generation quality not evaluated.** RULER and LoCoMo are retrieval-heavy. Open-ended generation quality, hallucination rate, and coherence in long documents aren't assessed.\n\nThe paper is at [arxiv:2607.28263](https://arxiv.org/abs/2607.28263). While the official code isn't out yet, you can start experimenting with the concept:\n\n**Profile your model's layer specialization**: Measure representational similarity (CKA) across layers on your target model — find where semantics stabilize.\n\n**Prototype the cache**: The `CoMemCache`\n\nclass above is a drop-in starting point. Hook it into your model's forward pass at the split layer.\n\n**Test without LoRA first**: Grab the upper-layer outputs without adaptation. Measure quality degradation — this tells you whether LoRA training is strictly necessary for your use case.\n\n**Benchmark your target task**: Before fine-tuning, verify that RULER/LoCoMo improvements generalize to your domain.\n\n```\n# One-liner to estimate whether CoMem is worth it for your model\ndef should_use_comem(current_vram_gb: float, context_tokens: int) -> str:\n    # CoMem approximation: ~d_model * seq_len * 2 bytes + model weights\n    # vs full KV: n_layers * d_kv * seq_len * 2 * 2 bytes + model weights\n    kv_overhead_gb = 32 * 128 * context_tokens * 2 * 2 / 1e9  # 32L, 128 d_kv\n    comem_overhead_gb = 4096 * context_tokens * 2 / 1e9        # just d_model\n\n    if current_vram_gb + kv_overhead_gb > 80:  # A100/H100/H20 limit\n        return f\"YES — KV cache needs {kv_overhead_gb:.1f}GB, CoMem needs {comem_overhead_gb:.1f}GB\"\n    return f\"Optional — KV cache needs {kv_overhead_gb:.1f}GB (within 80GB budget)\"\n\nprint(should_use_comem(current_vram_gb=20, context_tokens=128_000))\n# → YES — KV cache needs 85.9GB, CoMem needs 1.0GB\n```\n\n*What's your experience with long-context inference? Drop a comment — especially if you've tried alternative approaches like StreamingLLM or KV compression in production.*", "url": "https://wpnews.pro/news/comem-explained-from-paper-to-working-code-in-10-minutes", "canonical_source": "https://dev.to/cofldus/comem-explained-from-paper-to-working-code-in-10-minutes-476j", "published_at": "2026-08-04 02:14:59+00:00", "updated_at": "2026-08-04 03:10:57.723860+00:00", "lang": "en", "topics": ["large-language-models", "machine-learning", "ai-research", "ai-infrastructure", "developer-tools"], "entities": ["CoMem"], "alternates": {"html": "https://wpnews.pro/news/comem-explained-from-paper-to-working-code-in-10-minutes", "markdown": "https://wpnews.pro/news/comem-explained-from-paper-to-working-code-in-10-minutes.md", "text": "https://wpnews.pro/news/comem-explained-from-paper-to-working-code-in-10-minutes.txt", "jsonld": "https://wpnews.pro/news/comem-explained-from-paper-to-working-code-in-10-minutes.jsonld"}}