duo-laguna-adapter

DuoAttention Laguna Adapter

This repository contains adapter-only DuoAttention head weights for poolside/Laguna-XS.2. It does not include the Laguna base weights or tokenizer.

DuoAttention reduces long-context KV-cache growth by learning which KV heads need full history and letting the remaining heads keep only a sink window plus recent tokens. This Laguna adapter loads the base model, applies the learned head mask from duo_attention/full_attention_heads.pt, and enables DuoAttention with sink size 64 and recent size 256.

Why Use It

• Smaller KV cache for long prompts and generation.
• Adapter-only distribution, so the base Laguna model remains separate.
• trust_remote_code=True loading applies the Laguna DuoAttention patch.
• Laguna-specific support for gated attention projections and KV-head reordering.

Figures From The DuoAttention Paper

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

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

Laguna Adapter Figure

This figure visualizes the optimized Laguna DuoAttention gating values produced for this adapter.

<img src="figures/laguna_optimized_gate_values_booksum.png" alt="Laguna optimized DuoAttention gating values" width="420">

Laguna Results

We ran a mixed-precision KV benchmark as a Hugging Face Job on poolside/Laguna-XS.2. The base line accounts for dense Laguna FP8 KV cache; the DuoAttention path stores retrieval heads as FP8 and streaming heads as packed INT4 with per-group scale/zero-point metadata.

<img src="figures/laguna_mixed_kv_reduction_pct.png" alt="Laguna DuoAttention mixed KV cache reduction" width="820">

Job: 6a1ab49e5c8d10ffa11088c0

W&B run: ox2c0m6s

Laguna-Specific Changes

• Ported DuoAttention from Llama/Mistral-style attention modules to Laguna's gated attention structure.
• Preserved Laguna's g_proj gated output path when splitting full-context and streaming heads.
• Reordered Laguna Q/K/V/gating/output projections so full and streaming KV heads remain aligned after patching.
• Added adapter-only loading that fetches the base Laguna model separately and applies the learned full_attention_heads tensor at load time.
• Kept decode compatible with the patched tuple KV cache path used by the current Laguna remote code.

Usage

Install optional tokenizer dependencies if needed:

pip install sentencepiece tiktoken

Load the base tokenizer and compare the base Laguna cache with the DuoAttention cache on the same non-trivial prompt:

import gc
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

adapter_repo = "dogeplusplus/duo-laguna-adapter"
base_model_id = "poolside/Laguna-XS.2"

tokenizer = AutoTokenizer.from_pretrained(
    base_model_id,
    trust_remote_code=True,
    token=True,
)
model_kwargs = {
    "trust_remote_code": True,
    "token": True,
}
if torch.cuda.is_available():
    model_kwargs["dtype"] = torch.bfloat16
    model_kwargs["device_map"] = {"": "cuda:0"}
else:
    model_kwargs["torch_dtype"] = "auto"
    model_kwargs["device_map"] = "auto"


def cache_nbytes(value):
    if value is None:
        return 0
    if torch.is_tensor(value):
        return value.numel() * value.element_size()
    if hasattr(value, "key_cache") and hasattr(value, "value_cache"):
        return cache_nbytes(value.key_cache) + cache_nbytes(value.value_cache)
    if hasattr(value, "to_legacy_cache"):
        try:
            return cache_nbytes(value.to_legacy_cache())
        except Exception:
            pass
    if isinstance(value, dict):
        return sum(cache_nbytes(v) for v in value.values())
    if isinstance(value, (list, tuple)):
        return sum(cache_nbytes(v) for v in value)
    return 0


def first_parameter_device(model):
    return next(model.parameters()).device


def dense_kv_cache_nbytes(config, tokens, dtype):
    num_layers = config.num_hidden_layers
    num_key_value_heads = config.num_key_value_heads
    head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)
    bytes_per_value = torch.empty((), dtype=dtype).element_size()
    return num_layers * 2 * num_key_value_heads * tokens * head_dim * bytes_per_value


def clear_cuda():
    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
        torch.cuda.ipc_collect()


def greedy_decode_from_prefill(model, prefill, input_ids, max_new_tokens):
    past_key_values = prefill.past_key_values
    next_token = prefill.logits[:, -1, :].argmax(dim=-1, keepdim=True)
    generated = [input_ids, next_token]
    for _ in range(max_new_tokens - 1):
        out = model(
            input_ids=next_token,
            past_key_values=past_key_values,
            use_cache=True,
        )
        past_key_values = out.past_key_values
        next_token = out.logits[:, -1, :].argmax(dim=-1, keepdim=True)
        generated.append(next_token)
    return torch.cat(generated, dim=-1)


prompt = (
    "Remember this retrieval key: RIVER-4821. "
    + "The notebook contains many irrelevant meeting notes. " * 180
    + "Question: what is the retrieval key?"
)

base_model = AutoModelForCausalLM.from_pretrained(
    base_model_id,
    **model_kwargs,
).eval()
inputs = tokenizer(prompt, return_tensors="pt").to(first_parameter_device(base_model))
with torch.no_grad():
    base_out = base_model(**inputs, use_cache=True)
base_cache_bytes = cache_nbytes(base_out.past_key_values)
if base_cache_bytes == 0:
    base_cache_bytes = dense_kv_cache_nbytes(
        base_model.config,
        inputs["input_ids"].shape[-1],
        next(base_model.parameters()).dtype,
    )
base_cache_mib = base_cache_bytes / 2**20
del base_out, inputs, base_model
clear_cuda()

duo_model = AutoModelForCausalLM.from_pretrained(
    adapter_repo,
    **model_kwargs,
).eval()
duo_inputs = tokenizer(prompt, return_tensors="pt").to(first_parameter_device(duo_model))
with torch.no_grad():
    duo_out = duo_model(**duo_inputs, use_cache=True)
    generated = greedy_decode_from_prefill(duo_model, duo_out, duo_inputs["input_ids"], 64)
duo_cache_mib = cache_nbytes(duo_out.past_key_values) / 2**20

print(f"Base KV cache: {base_cache_mib:.2f} MiB")
print(f"Duo KV cache:  {cache_nbytes(duo_out.past_key_values) / 2**20:.2f} MiB")
print(f"KV reduction:   {100 * (1 - duo_cache_mib / base_cache_mib):.1f}%")

print(tokenizer.decode(generated[0], skip_special_tokens=True))

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