duo-laguna-adapter-rosy-energy-5-preview

DuoAttention Laguna Adapter

This repository contains learned DuoAttention attention-head weights and custom loading code for poolside/Laguna-XS.2. It intentionally does not include the full Laguna base-model weights or tokenizer files.

Summary

DuoAttention is a KV-cache reduction method for long-context decoder models. It learns which KV heads need full-context memory for retrieval-heavy behavior and lets the remaining heads use a streaming cache made from a fixed sink window and a recent-token window. The idea is simple: not every attention head needs to carry the whole past sequence. Keeping full history only for the heads that use it preserves the parts of the model most responsible for long-range retrieval, while reducing KV-cache growth for the rest.

For this adapter, the learned head mask is packaged separately from the base Laguna weights. At load time, the custom model code reads duo_attention/full_attention_heads.pt, patches the Laguna attention modules, and enables DuoAttention inference with sink size 64 and recent size 256.

Why It Works

• Transformer KV-cache memory scales with sequence length, layer count, KV head count, head dimension, and dtype.
• DuoAttention reduces the effective cache footprint by splitting KV heads into full-context heads and streaming heads.
• Full-context heads retain all prior tokens for retrieval and global-memory behavior.
• Streaming heads keep only sink tokens plus the most recent tokens, limiting cache growth while preserving local continuity and stable prefix anchoring.

Benefits

• Smaller KV-cache footprint for long prompts and long decode workloads.
• A deployable adapter artifact that avoids republishing the full Laguna base model.
• Custom trust_remote_code loading that applies the Laguna patch automatically.
• Benchmark tooling for base-vs-Duo latency and KV-cache utilization, including static Duo KV-cache measurements.

Figures

Selected figures from the DuoAttention paper show the method, cache tradeoff, and latency motivation. The final figure is Laguna-specific and visualizes the optimized gating values learned for this adapter.

<img src="figures/method1.jpg" alt="DuoAttention retrieval and streaming head split" width="820"> <img src="figures/method2.jpg" alt="DuoAttention full and streaming KV-cache pattern" width="820"> <img src="figures/kv_capacity.jpg" alt="DuoAttention KV-cache capacity comparison" width="620"> <img src="figures/efficiency_prefilling.jpg" alt="DuoAttention prefilling efficiency" width="820"> <img src="figures/efficiency_decoding.jpg" alt="DuoAttention decoding efficiency" width="820"> <img src="figures/laguna_optimized_gate_values_booksum.png" alt="Laguna optimized DuoAttention gating values" width="420">

Paper reference: DuoAttention: Efficient Long-Context LLM Inference with Retrieval and Streaming Heads.

Training

The adapter is produced by training per-layer, per-KV-head DuoAttention scores on long-context retrieval-style data. During training, the model learns which KV heads should remain full-context. The final artifact stores the learned full_attention_heads tensor and a small duo_attention config block with the base model id, sink/recent sizes, and custom-code metadata.

This package contains adapter state only:

• duo_attention/full_attention_heads.pt
• config.json with the DuoAttention metadata and remote-code mapping
• custom Laguna loading/patch code
• requirements.txt

Changes In This Submission

Engineering changes:

• Added Hugging Face Hub packaging for adapter-only DuoAttention artifacts.
• Added custom AutoModelForCausalLM loading via trust_remote_code=True.
• Added automatic loading of DuoAttention head weights from the Hub repo.
• Added a Laguna-specific benchmark suite for prefill latency, decode latency, KV-cache size, active KV-cache utilization, plots, JSONL/CSV output, and W&B artifact logging.
• Added Hugging Face Jobs scripts for training and evaluation runs.
• Added local and Hub dataset loading support for training datasets.
• Added Mac/local smoke-test paths and PyTorch fallback paths where practical.

Laguna-specific / novel changes:

• Ported DuoAttention patching from Llama/Mistral-style modules to Laguna's attention structure.
• Accounted for Laguna's gated attention output projection path.
• Added Laguna tuple KV-cache compatibility for ordinary generation.
• Added Laguna static Duo KV-cache support for measuring cache allocation and utilization.
• Added head reordering for Laguna Q/K/V/gating/output projections so full and streaming KV heads can be separated consistently.
• Preserved adapter-only distribution so users load the gated/base Laguna model separately from the learned DuoAttention heads.

Install optional tokenizer dependencies if needed:

pip install sentencepiece tiktoken

Load the tokenizer from the base Laguna model and the patched model from this adapter repository:

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

adapter_repo = "dogeplusplus/duo-laguna-adapter-rosy-energy-5-preview"
base_model = "poolside/Laguna-XS.2"

tokenizer = AutoTokenizer.from_pretrained(
    base_model,
    trust_remote_code=True,
    token=True,
)

model = AutoModelForCausalLM.from_pretrained(
    adapter_repo,
    trust_remote_code=True,
    token=True,
    torch_dtype="auto",
    device_map="auto",
)

prompt = "The capital of France is"
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
with torch.no_grad():
    generated = model.generate(**inputs, max_new_tokens=32)
print(tokenizer.decode(generated[0], skip_special_tokens=True))

Use token=True after running hf auth login, or pass a token string directly when loading private or gated repositories.