--- library_name: transformers base_model: poolside/Laguna-XS.2 tags: - laguna - duo-attention - custom-code --- # 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. Paper reference: [DuoAttention: Efficient Long-Context LLM Inference with Retrieval and Streaming Heads](https://arxiv.org/abs/2410.10819). ## 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: ```bash pip install sentencepiece tiktoken ``` Load the tokenizer from the base Laguna model and the patched model from this adapter repository: ```python 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.
