Upload KV Tenants DuoAttention Laguna adapter

.gitattributes

@@ -33,3 +33,8 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+figures/efficiency_decoding.jpg filter=lfs diff=lfs merge=lfs -text
+figures/efficiency_prefilling.jpg filter=lfs diff=lfs merge=lfs -text
+figures/kv_capacity.jpg filter=lfs diff=lfs merge=lfs -text
+figures/method1.jpg filter=lfs diff=lfs merge=lfs -text
+figures/method2.jpg filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,53 @@
+---
+library_name: transformers
+base_model: poolside/Laguna-XS.2
+tags:
+- laguna
+- duo-attention
+- custom-code
+---
+
+# DuoAttention Laguna Adapter
+
+This repository contains the 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.
+
+Install the optional tokenizer dependencies if they are not already present:
+
+```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
+
+repo_id = "<this-repo-id>"
+base_model = "poolside/Laguna-XS.2"
+
+tokenizer = AutoTokenizer.from_pretrained(
+    base_model,
+    trust_remote_code=True,
+    token=True,
+)
+model = AutoModelForCausalLM.from_pretrained(
+    repo_id,
+    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.

config.json

@@ -0,0 +1,218 @@
+{
+  "_name_or_path": "poolside/Laguna-XS.2",
+  "architectures": [
+    "DuoLagunaForCausalLM"
+  ],
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "auto_map": {
+    "AutoModelForCausalLM": "modeling_duo_laguna.DuoLagunaForCausalLM"
+  },
+  "bos_token_id": 2,
+  "chunk_size_feed_forward": 0,
+  "dtype": "bfloat16",
+  "duo_attention": {
+    "architecture": "laguna",
+    "base_config_name_or_path": null,
+    "base_model_name_or_path": "poolside/Laguna-XS.2",
+    "enabled": true,
+    "format_version": 1,
+    "full_attention_heads_file": "duo_attention/full_attention_heads.pt",
+    "full_attention_heads_tsv_file": "duo_attention/full_attention_heads.tsv",
+    "patch_mode": "eval",
+    "recent_size": 256,
+    "sink_size": 64,
+    "training_recent_size": 256,
+    "training_sink_size": 64
+  },
+  "eos_token_id": [
+    2,
+    24
+  ],
+  "gating": true,
+  "head_dim": 128,
+  "hidden_act": "silu",
+  "hidden_size": 2048,
+  "id2label": {
+    "0": "LABEL_0",
+    "1": "LABEL_1"
+  },
+  "initializer_range": 0.02,
+  "intermediate_size": 8192,
+  "is_encoder_decoder": false,
+  "label2id": {
+    "LABEL_0": 0,
+    "LABEL_1": 1
+  },
+  "layer_types": [
+    "full_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "full_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "full_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "full_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "full_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "full_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "full_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "full_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "full_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "full_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention"
+  ],
+  "max_position_embeddings": 262144,
+  "mlp_layer_types": [
+    "dense",
+    "sparse",
+    "sparse",
+    "sparse",
+    "sparse",
+    "sparse",
+    "sparse",
+    "sparse",
+    "sparse",
+    "sparse",
+    "sparse",
+    "sparse",
+    "sparse",
+    "sparse",
+    "sparse",
+    "sparse",
+    "sparse",
+    "sparse",
+    "sparse",
+    "sparse",
+    "sparse",
+    "sparse",
+    "sparse",
+    "sparse",
+    "sparse",
+    "sparse",
+    "sparse",
+    "sparse",
+    "sparse",
+    "sparse",
+    "sparse",
+    "sparse",
+    "sparse",
+    "sparse",
+    "sparse",
+    "sparse",
+    "sparse",
+    "sparse",
+    "sparse",
+    "sparse"
+  ],
+  "model_type": "laguna",
+  "moe_apply_router_weight_on_input": false,
+  "moe_intermediate_size": 512,
+  "moe_routed_scaling_factor": 2.5,
+  "moe_router_logit_softcapping": 0.0,
+  "num_attention_heads": 48,
+  "num_attention_heads_per_layer": [
+    48,
+    64,
+    64,
+    64,
+    48,
+    64,
+    64,
+    64,
+    48,
+    64,
+    64,
+    64,
+    48,
+    64,
+    64,
+    64,
+    48,
+    64,
+    64,
+    64,
+    48,
+    64,
+    64,
+    64,
+    48,
+    64,
+    64,
+    64,
+    48,
+    64,
+    64,
+    64,
+    48,
+    64,
+    64,
+    64,
+    48,
+    64,
+    64,
+    64
+  ],
+  "num_experts": 256,
+  "num_experts_per_tok": 8,
+  "num_hidden_layers": 40,
+  "num_key_value_heads": 8,
+  "output_attentions": false,
+  "output_hidden_states": false,
+  "output_router_logits": false,
+  "pad_token_id": 9,
+  "partial_rotary_factor": 0.5,
+  "problem_type": null,
+  "return_dict": true,
+  "rms_norm_eps": 1e-06,
+  "rope_parameters": {
+    "full_attention": {
+      "attention_factor": 1.0,
+      "beta_fast": 64.0,
+      "beta_slow": 1.0,
+      "factor": 64.0,
+      "original_max_position_embeddings": 4096,
+      "partial_rotary_factor": 0.5,
+      "rope_theta": 500000.0,
+      "rope_type": "yarn"
+    },
+    "original_max_position_embeddings": 4096,
+    "rope_type": "default",
+    "sliding_attention": {
+      "partial_rotary_factor": 1.0,
+      "rope_theta": 10000.0,
+      "rope_type": "default"
+    }
+  },
+  "router_aux_loss_coef": 0.0,
+  "shared_expert_intermediate_size": 512,
+  "sliding_window": 512,
+  "tie_word_embeddings": false,
+  "transformers_version": "5.9.0",
+  "use_cache": true,
+  "vocab_size": 100352
+}

duo_attention/config.json

@@ -0,0 +1,14 @@
+{
+  "architecture": "laguna",
+  "base_config_name_or_path": null,
+  "base_model_name_or_path": "poolside/Laguna-XS.2",
+  "enabled": true,
+  "format_version": 1,
+  "full_attention_heads_file": "duo_attention/full_attention_heads.pt",
+  "full_attention_heads_tsv_file": "duo_attention/full_attention_heads.tsv",
+  "patch_mode": "eval",
+  "recent_size": 256,
+  "sink_size": 64,
+  "training_recent_size": 256,
+  "training_sink_size": 64
+}

duo_attention/full_attention_heads.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7bf485ba3f39e94d04660d87faddabebdf6a70820f7f472c4f5d2171ec32b94b
+size 2525

duo_attention/full_attention_heads.tsv

@@ -0,0 +1,40 @@
+8.984375000000000000e-01	5.468750000000000000e-01	9.882812500000000000e-01	1.000000000000000000e+00	8.906250000000000000e-01	1.000000000000000000e+00	9.921875000000000000e-01	9.882812500000000000e-01
+1.000000000000000000e+00	1.000000000000000000e+00	1.000000000000000000e+00	1.000000000000000000e+00	1.000000000000000000e+00	9.687500000000000000e-01	9.960937500000000000e-01	1.000000000000000000e+00
+1.000000000000000000e+00	1.000000000000000000e+00	1.000000000000000000e+00	1.000000000000000000e+00	9.179687500000000000e-01	7.070312500000000000e-01	8.554687500000000000e-01	9.765625000000000000e-01
+1.000000000000000000e+00	1.000000000000000000e+00	1.000000000000000000e+00	9.570312500000000000e-01	1.000000000000000000e+00	9.726562500000000000e-01	1.000000000000000000e+00	9.648437500000000000e-01
+1.000000000000000000e+00	1.000000000000000000e+00	1.000000000000000000e+00	1.000000000000000000e+00	1.000000000000000000e+00	8.242187500000000000e-01	1.000000000000000000e+00	1.000000000000000000e+00
+1.000000000000000000e+00	1.000000000000000000e+00	8.437500000000000000e-01	9.960937500000000000e-01	1.000000000000000000e+00	9.101562500000000000e-01	1.000000000000000000e+00	8.945312500000000000e-01
+8.085937500000000000e-01	9.765625000000000000e-01	3.652343750000000000e-01	1.000000000000000000e+00	9.765625000000000000e-01	1.000000000000000000e+00	3.359375000000000000e-01	1.000000000000000000e+00
+1.000000000000000000e+00	1.000000000000000000e+00	1.007812500000000000e+00	9.609375000000000000e-01	1.000000000000000000e+00	1.000000000000000000e+00	1.000000000000000000e+00	9.609375000000000000e-01
+1.000000000000000000e+00	9.375000000000000000e-01	1.000000000000000000e+00	8.593750000000000000e-01	9.570312500000000000e-01	9.843750000000000000e-01	9.882812500000000000e-01	7.187500000000000000e-01
+8.632812500000000000e-01	7.187500000000000000e-01	1.000000000000000000e+00	7.226562500000000000e-01	1.000000000000000000e+00	1.000000000000000000e+00	1.000000000000000000e+00	1.000000000000000000e+00
+7.109375000000000000e-01	1.000000000000000000e+00	8.750000000000000000e-01	1.000000000000000000e+00	9.179687500000000000e-01	8.906250000000000000e-01	9.414062500000000000e-01	1.000000000000000000e+00
+1.000000000000000000e+00	7.187500000000000000e-01	6.132812500000000000e-01	1.000000000000000000e+00	1.000000000000000000e+00	9.023437500000000000e-01	1.000000000000000000e+00	9.882812500000000000e-01
+1.000000000000000000e+00	8.945312500000000000e-01	9.960937500000000000e-01	1.000000000000000000e+00	1.000000000000000000e+00	1.000000000000000000e+00	1.000000000000000000e+00	6.875000000000000000e-01
+4.140625000000000000e-01	7.812500000000000000e-01	1.007812500000000000e+00	9.492187500000000000e-01	9.062500000000000000e-01	5.703125000000000000e-01	5.351562500000000000e-01	8.945312500000000000e-01
+1.000000000000000000e+00	9.609375000000000000e-01	1.000000000000000000e+00	3.906250000000000000e-01	6.718750000000000000e-01	7.617187500000000000e-02	5.585937500000000000e-01	1.000000000000000000e+00
+1.000000000000000000e+00	8.085937500000000000e-01	1.031494140625000000e-02	1.000000000000000000e+00	9.921875000000000000e-01	3.984375000000000000e-01	1.640625000000000000e-01	1.867675781250000000e-02
+9.921875000000000000e-01	1.000000000000000000e+00	2.558593750000000000e-01	1.000000000000000000e+00	1.000000000000000000e+00	9.882812500000000000e-01	8.242187500000000000e-01	9.375000000000000000e-01
+1.000000000000000000e+00	3.398437500000000000e-01	4.570312500000000000e-01	2.578125000000000000e-01	9.609375000000000000e-01	9.375000000000000000e-01	6.757812500000000000e-01	8.906250000000000000e-01
+1.000000000000000000e+00	7.421875000000000000e-01	4.824218750000000000e-01	2.792968750000000000e-01	4.467773437500000000e-02	3.574218750000000000e-01	8.945312500000000000e-01	8.320312500000000000e-01
+-1.564025878906250000e-04	4.218750000000000000e-01	6.054687500000000000e-01	1.000000000000000000e+00	9.375000000000000000e-01	9.882812500000000000e-01	4.746093750000000000e-01	2.880859375000000000e-02
+9.882812500000000000e-01	1.000000000000000000e+00	1.000000000000000000e+00	1.000000000000000000e+00	1.000000000000000000e+00	1.000000000000000000e+00	8.437500000000000000e-01	2.197265625000000000e-01
+2.558593750000000000e-01	8.320312500000000000e-01	7.812500000000000000e-01	4.414062500000000000e-01	1.000000000000000000e+00	6.640625000000000000e-01	1.000000000000000000e+00	4.570312500000000000e-01
+7.421875000000000000e-01	5.625000000000000000e-01	9.375000000000000000e-01	3.417968750000000000e-01	9.335937500000000000e-01	6.445312500000000000e-01	1.708984375000000000e-01	4.863281250000000000e-01
+5.468750000000000000e-01	4.062500000000000000e-01	3.496093750000000000e-01	5.419921875000000000e-02	4.843750000000000000e-01	1.060485839843750000e-03	-1.953125000000000000e-03	3.593750000000000000e-01
+1.000000000000000000e+00	9.140625000000000000e-01	1.007812500000000000e+00	1.007812500000000000e+00	1.007812500000000000e+00	3.710937500000000000e-01	1.007812500000000000e+00	9.257812500000000000e-01
+2.431640625000000000e-01	9.179687500000000000e-01	9.296875000000000000e-01	1.000000000000000000e+00	6.796875000000000000e-01	4.726562500000000000e-01	3.847656250000000000e-01	1.000000000000000000e+00
+9.609375000000000000e-01	9.921875000000000000e-01	5.664062500000000000e-01	1.000000000000000000e+00	7.031250000000000000e-01	3.125000000000000000e-01	1.000000000000000000e+00	1.000000000000000000e+00
+1.079101562500000000e-01	5.507812500000000000e-01	2.138671875000000000e-01	-1.899719238281250000e-03	9.216308593750000000e-03	8.867187500000000000e-01	2.187500000000000000e-01	2.324218750000000000e-01
+1.007812500000000000e+00	8.515625000000000000e-01	1.007812500000000000e+00	1.000000000000000000e+00	9.023437500000000000e-01	1.000000000000000000e+00	8.007812500000000000e-01	1.000000000000000000e+00
+5.507812500000000000e-01	4.414062500000000000e-01	9.687500000000000000e-01	8.593750000000000000e-01	8.437500000000000000e-01	9.335937500000000000e-01	1.000000000000000000e+00	1.000000000000000000e+00
+1.000000000000000000e+00	1.777343750000000000e-01	8.476562500000000000e-01	9.257812500000000000e-01	1.000000000000000000e+00	6.640625000000000000e-01	7.617187500000000000e-01	1.000000000000000000e+00
+9.179687500000000000e-02	3.769531250000000000e-01	7.304687500000000000e-01	7.070312500000000000e-01	9.843750000000000000e-01	8.398437500000000000e-01	1.007812500000000000e+00	9.453125000000000000e-01
+1.007812500000000000e+00	9.335937500000000000e-01	9.023437500000000000e-01	1.000000000000000000e+00	1.000000000000000000e+00	-3.967285156250000000e-03	1.000000000000000000e+00	1.007812500000000000e+00
+1.000000000000000000e+00	1.000000000000000000e+00	9.843750000000000000e-01	9.726562500000000000e-01	9.335937500000000000e-01	5.000000000000000000e-01	5.781250000000000000e-01	9.804687500000000000e-01
+1.007812500000000000e+00	9.804687500000000000e-01	9.453125000000000000e-01	9.726562500000000000e-01	4.667968750000000000e-01	9.960937500000000000e-01	8.085937500000000000e-01	1.000000000000000000e+00
+7.617187500000000000e-01	3.847656250000000000e-01	8.710937500000000000e-01	8.593750000000000000e-01	5.117187500000000000e-01	7.500000000000000000e-01	1.007812500000000000e+00	5.625000000000000000e-01
+9.023437500000000000e-01	9.179687500000000000e-01	8.632812500000000000e-01	1.007812500000000000e+00	9.765625000000000000e-01	6.914062500000000000e-01	5.351562500000000000e-01	9.765625000000000000e-01
+8.125000000000000000e-01	3.847656250000000000e-01	9.453125000000000000e-01	8.906250000000000000e-01	7.890625000000000000e-01	8.203125000000000000e-01	9.140625000000000000e-01	8.593750000000000000e-01
+7.148437500000000000e-01	9.179687500000000000e-01	9.414062500000000000e-01	1.000000000000000000e+00	6.367187500000000000e-01	8.164062500000000000e-01	1.007812500000000000e+00	1.000000000000000000e+00
+1.000000000000000000e+00	5.312500000000000000e-01	8.476562500000000000e-01	6.250000000000000000e-01	4.785156250000000000e-01	1.007812500000000000e+00	9.375000000000000000e-01	1.007812500000000000e+00

duo_laguna_remote.py

@@ -0,0 +1,448 @@
+import types
+from typing import Optional, Tuple
+
+import torch
+import torch.nn.functional as F
+from transformers.modeling_outputs import MoeCausalLMOutputWithPast, MoeModelOutputWithPast
+from transformers.models.laguna.modeling_laguna import apply_rotary_pos_emb
+
+
+try:
+    from flash_attn import flash_attn_func
+except ImportError:
+
+    def flash_attn_func(
+        query_states,
+        key_states,
+        value_states,
+        dropout_p=0.0,
+        softmax_scale=None,
+        causal=False,
+        **kwargs,
+    ):
+        if key_states.shape[2] != query_states.shape[2]:
+            repeat = query_states.shape[2] // key_states.shape[2]
+            key_states = key_states.repeat_interleave(repeat, dim=2)
+            value_states = value_states.repeat_interleave(repeat, dim=2)
+
+        scale = softmax_scale or query_states.shape[-1] ** -0.5
+        attn_weights = torch.einsum("bqhd,bkhd->bhqk", query_states, key_states) * scale
+        if causal:
+            q_len = query_states.shape[1]
+            kv_len = key_states.shape[1]
+            causal_mask = torch.ones(
+                q_len, kv_len, dtype=torch.bool, device=query_states.device
+            ).triu(kv_len - q_len + 1)
+            attn_weights = attn_weights.masked_fill(
+                causal_mask[None, None], float("-inf")
+            )
+        attn_weights = torch.softmax(attn_weights, dim=-1, dtype=torch.float32).to(
+            query_states.dtype
+        )
+        if dropout_p:
+            attn_weights = torch.nn.functional.dropout(attn_weights, p=dropout_p)
+        return torch.einsum("bhqk,bkhd->bqhd", attn_weights, value_states)
+
+
+@torch.no_grad()
+def reorder_linear_weights(linear_module, full_attention_heads, repeat_num, reorder_channel):
+    full_attention_heads = torch.repeat_interleave(
+        full_attention_heads, repeats=repeat_num
+    ).to(linear_module.weight.device)
+    full_attn_mask = full_attention_heads > 0.5
+    if reorder_channel == "in":
+        reordered_weight = torch.cat(
+            [
+                linear_module.weight.data[:, full_attn_mask],
+                linear_module.weight.data[:, ~full_attn_mask],
+            ],
+            dim=1,
+        )
+    else:
+        reordered_weight = torch.cat(
+            [
+                linear_module.weight.data[full_attn_mask, :],
+                linear_module.weight.data[~full_attn_mask, :],
+            ],
+            dim=0,
+        )
+    linear_module.weight.data = reordered_weight
+    if linear_module.bias is not None:
+        linear_module.bias.data = torch.cat(
+            [
+                linear_module.bias.data[full_attn_mask],
+                linear_module.bias.data[~full_attn_mask],
+            ],
+            dim=0,
+        )
+    return linear_module
+
+
+@torch.no_grad()
+def reorder_full_attn_heads(full_attention_heads):
+    full_attn_mask = full_attention_heads > 0.5
+    num_full_attn_heads = full_attn_mask.sum().item()
+    full_attention_heads[:num_full_attn_heads] = 1
+    full_attention_heads[num_full_attn_heads:] = 0
+    return full_attention_heads
+
+
+def _num_key_value_heads(module):
+    return getattr(module, "num_key_value_heads", module.config.num_key_value_heads)
+
+
+def _shape_qkv(module, hidden_states):
+    bsz, q_len, _ = hidden_states.size()
+    num_key_value_heads = _num_key_value_heads(module)
+    query_states = module.q_proj(hidden_states).view(
+        bsz, q_len, module.num_heads, module.head_dim
+    )
+    key_states = module.k_proj(hidden_states).view(
+        bsz, q_len, num_key_value_heads, module.head_dim
+    )
+    value_states = module.v_proj(hidden_states).view(
+        bsz, q_len, num_key_value_heads, module.head_dim
+    )
+    return module.q_norm(query_states), module.k_norm(key_states), value_states
+
+
+def _apply_gate(module, attn_output, hidden_states):
+    input_shape = hidden_states.shape[:-1]
+    gate = F.softplus(module.g_proj(hidden_states).float()).to(attn_output.dtype)
+    attn_output = attn_output.reshape(*input_shape, module.num_heads, module.head_dim)
+    attn_output = attn_output * gate.unsqueeze(-1)
+    return attn_output.reshape(*input_shape, module.num_heads * module.head_dim)
+
+
+def laguna_duo_attention_forward(
+    self,
+    hidden_states: torch.Tensor,
+    attention_mask: Optional[torch.Tensor] = None,
+    position_ids: Optional[torch.LongTensor] = None,
+    position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+    past_key_value: Optional[Tuple[torch.Tensor]] = None,
+    output_attentions: bool = False,
+    use_cache: bool = False,
+    **kwargs,
+):
+    bsz, q_len, _ = hidden_states.size()
+    query_states, key_states, value_states = _shape_qkv(self, hidden_states)
+
+    kv_seq_len = key_states.shape[1]
+    if past_key_value is not None:
+        kv_seq_len += past_key_value[0].shape[2]
+
+    if position_embeddings is None:
+        raise ValueError("Duo Laguna requires position_embeddings")
+    query_states, key_states = apply_rotary_pos_emb(
+        query_states,
+        key_states,
+        *position_embeddings,
+        unsqueeze_dim=2,
+    )
+
+    if not hasattr(self, "full_attn_head_mask") or self.full_attn_head_mask is None:
+        self.full_attn_head_mask = self.full_attention_heads > 0.5
+        self.num_full_attn_head = self.full_attn_head_mask.sum().item()
+        self.num_streaming_attn_head = _num_key_value_heads(self) - self.num_full_attn_head
+        self.num_full_query_head = self.num_full_attn_head * self.num_key_value_groups
+        self.num_streaming_query_head = self.num_heads - self.num_full_query_head
+
+    full_key_states = key_states[:, :, : self.num_full_attn_head, :]
+    full_value_states = value_states[:, :, : self.num_full_attn_head, :]
+    streaming_key_states = key_states[:, :, self.num_full_attn_head :, :]
+    streaming_value_states = value_states[:, :, self.num_full_attn_head :, :]
+
+    if past_key_value is not None:
+        past_full_kv = past_key_value[0].transpose(1, 2)
+        past_streaming_kv = past_key_value[1].transpose(1, 2)
+        full_key_states = torch.cat([past_full_kv[:bsz], full_key_states], dim=1)
+        full_value_states = torch.cat([past_full_kv[bsz:], full_value_states], dim=1)
+        streaming_key_states = torch.cat(
+            [past_streaming_kv[:bsz], streaming_key_states], dim=1
+        )
+        streaming_value_states = torch.cat(
+            [past_streaming_kv[bsz:], streaming_value_states], dim=1
+        )
+
+    if q_len == kv_seq_len:
+        attn_output = flash_attn_func(
+            query_states, key_states, value_states, causal=True, dropout_p=0.0
+        )
+    else:
+        full_attn_output = None
+        streaming_attn_output = None
+        if self.num_full_attn_head > 0:
+            full_attn_output = flash_attn_func(
+                query_states[:, :, : self.num_full_query_head, :],
+                full_key_states,
+                full_value_states,
+                causal=True,
+                dropout_p=0.0,
+            )
+        if self.num_streaming_attn_head > 0:
+            streaming_attn_output = flash_attn_func(
+                query_states[:, :, self.num_full_query_head :, :],
+                streaming_key_states,
+                streaming_value_states,
+                causal=True,
+                dropout_p=0.0,
+            )
+        if full_attn_output is None:
+            attn_output = streaming_attn_output
+        elif streaming_attn_output is None:
+            attn_output = full_attn_output
+        else:
+            attn_output = torch.cat([full_attn_output, streaming_attn_output], dim=2)
+
+    attn_output = self.o_proj(_apply_gate(self, attn_output, hidden_states))
+
+    if streaming_key_states.shape[1] > self.recent_size + self.sink_size:
+        recent_key_states = streaming_key_states[:, -self.recent_size :, :, :].clone()
+        streaming_key_states[:, self.sink_size : self.sink_size + self.recent_size].copy_(
+            recent_key_states
+        )
+        streaming_key_states = streaming_key_states[:, : self.sink_size + self.recent_size]
+
+        recent_value_states = streaming_value_states[:, -self.recent_size :, :, :].clone()
+        streaming_value_states[
+            :, self.sink_size : self.sink_size + self.recent_size
+        ].copy_(recent_value_states)
+        streaming_value_states = streaming_value_states[
+            :, : self.sink_size + self.recent_size
+        ]
+
+    past_key_value = (
+        (
+            torch.cat([full_key_states, full_value_states], dim=0).transpose(1, 2),
+            torch.cat([streaming_key_states, streaming_value_states], dim=0).transpose(
+                1, 2
+            ),
+        )
+        if use_cache
+        else None
+    )
+    return attn_output, None, past_key_value
+
+
+def laguna_for_causal_lm_forward(
+    self,
+    input_ids=None,
+    attention_mask=None,
+    position_ids=None,
+    past_key_values=None,
+    inputs_embeds=None,
+    labels=None,
+    use_cache=None,
+    output_attentions=None,
+    output_hidden_states=None,
+    return_dict=None,
+    logits_to_keep=0,
+    **kwargs,
+):
+    output_hidden_states = (
+        output_hidden_states
+        if output_hidden_states is not None
+        else self.config.output_hidden_states
+    )
+    return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+    outputs = self.model(
+        input_ids=input_ids,
+        attention_mask=attention_mask,
+        position_ids=position_ids,
+        past_key_values=past_key_values,
+        inputs_embeds=inputs_embeds,
+        use_cache=use_cache,
+        output_attentions=output_attentions,
+        output_hidden_states=output_hidden_states,
+        return_dict=return_dict,
+        **kwargs,
+    )
+    hidden_states = outputs[0]
+    slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
+    logits = self.lm_head(hidden_states[:, slice_indices if self.training else slice(-1, None), :])
+    loss = None
+    if labels is not None:
+        loss = self.loss_function(logits, labels, self.vocab_size, **kwargs)
+    if not return_dict:
+        output = (logits,) + outputs[1:]
+        return (loss,) + output if loss is not None else output
+    return MoeCausalLMOutputWithPast(
+        loss=loss,
+        logits=logits,
+        past_key_values=outputs.past_key_values,
+        hidden_states=outputs.hidden_states,
+        attentions=outputs.attentions,
+        router_logits=None,
+    )
+
+
+def laguna_model_forward(
+    self,
+    input_ids=None,
+    attention_mask=None,
+    position_ids=None,
+    past_key_values=None,
+    inputs_embeds=None,
+    use_cache=None,
+    output_attentions=None,
+    output_hidden_states=None,
+    return_dict=None,
+    **kwargs,
+):
+    output_attentions = (
+        output_attentions
+        if output_attentions is not None
+        else self.config.output_attentions
+    )
+    output_hidden_states = (
+        output_hidden_states
+        if output_hidden_states is not None
+        else self.config.output_hidden_states
+    )
+    use_cache = use_cache if use_cache is not None else self.config.use_cache
+    return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+    if input_ids is not None and inputs_embeds is not None:
+        raise ValueError("You cannot specify both input_ids and inputs_embeds")
+    if inputs_embeds is None:
+        inputs_embeds = self.embed_tokens(input_ids)
+    batch_size, seq_length, _ = inputs_embeds.shape
+
+    past_key_values_length = past_key_values[0][0].shape[2] if past_key_values is not None else 0
+    if position_ids is None:
+        position_ids = torch.arange(
+            past_key_values_length,
+            seq_length + past_key_values_length,
+            dtype=torch.long,
+            device=inputs_embeds.device,
+        ).unsqueeze(0)
+    else:
+        position_ids = position_ids.view(-1, seq_length).long()
+
+    hidden_states = inputs_embeds
+    all_hidden_states = () if output_hidden_states else None
+    all_self_attns = () if output_attentions else None
+    next_decoder_cache = () if use_cache else None
+
+    position_embeddings = {}
+    for layer_type in set(self.config.layer_types):
+        position_embeddings[layer_type] = self.rotary_emb(hidden_states, position_ids, layer_type)
+
+    for idx, decoder_layer in enumerate(self.layers):
+        if output_hidden_states:
+            all_hidden_states += (hidden_states,)
+        past_key_value = past_key_values[idx] if past_key_values is not None else None
+        layer_outputs = decoder_layer(
+            hidden_states,
+            attention_mask=None,
+            position_ids=position_ids,
+            past_key_value=past_key_value,
+            output_attentions=output_attentions,
+            use_cache=use_cache,
+            position_embeddings=position_embeddings[self.config.layer_types[idx]],
+        )
+        hidden_states = layer_outputs[0]
+        if use_cache:
+            next_decoder_cache += (layer_outputs[2 if output_attentions else 1],)
+        if output_attentions:
+            all_self_attns += (layer_outputs[1],)
+
+    hidden_states = self.norm(hidden_states)
+    if output_hidden_states:
+        all_hidden_states += (hidden_states,)
+
+    next_cache = next_decoder_cache if use_cache else None
+    if not return_dict:
+        return tuple(
+            v
+            for v in [hidden_states, next_cache, all_hidden_states, all_self_attns]
+            if v is not None
+        )
+    return MoeModelOutputWithPast(
+        last_hidden_state=hidden_states,
+        past_key_values=next_cache,
+        hidden_states=all_hidden_states,
+        attentions=all_self_attns,
+        router_logits=None,
+    )
+
+
+def laguna_decoder_layer_forward(
+    self,
+    hidden_states,
+    attention_mask=None,
+    position_ids=None,
+    past_key_value=None,
+    output_attentions=False,
+    use_cache=False,
+    position_embeddings=None,
+    **kwargs,
+):
+    residual = hidden_states
+    hidden_states = self.input_layernorm(hidden_states)
+    hidden_states, self_attn_weights, present_key_value = self.self_attn(
+        hidden_states=hidden_states,
+        attention_mask=attention_mask,
+        position_ids=position_ids,
+        position_embeddings=position_embeddings,
+        past_key_value=past_key_value,
+        output_attentions=output_attentions,
+        use_cache=use_cache,
+    )
+    hidden_states = residual + hidden_states
+    residual = hidden_states
+    hidden_states = self.post_attention_layernorm(hidden_states)
+    hidden_states = self.mlp(hidden_states)
+    hidden_states = residual + hidden_states
+
+    outputs = (hidden_states,)
+    if output_attentions:
+        outputs += (self_attn_weights,)
+    if use_cache:
+        outputs += (present_key_value,)
+    return outputs
+
+
+def enable_duo_laguna_eval(model, full_attention_heads, sink_size, recent_size):
+    model.model.forward = types.MethodType(laguna_model_forward, model.model)
+    for layer in model.model.layers:
+        layer.forward = types.MethodType(laguna_decoder_layer_forward, layer)
+
+    device = next(model.parameters()).device
+    dtype = next(model.parameters()).dtype
+    for idx, layer in enumerate(model.model.layers):
+        module = layer.self_attn
+        module.num_key_value_heads = _num_key_value_heads(module)
+        layer_full_attention_heads = torch.as_tensor(
+            full_attention_heads[idx], device=device, dtype=dtype
+        )
+        module.forward = types.MethodType(laguna_duo_attention_forward, module)
+        module.q_proj = reorder_linear_weights(
+            module.q_proj,
+            layer_full_attention_heads,
+            module.num_key_value_groups * module.head_dim,
+            "out",
+        )
+        module.k_proj = reorder_linear_weights(
+            module.k_proj, layer_full_attention_heads, module.head_dim, "out"
+        )
+        module.v_proj = reorder_linear_weights(
+            module.v_proj, layer_full_attention_heads, module.head_dim, "out"
+        )
+        module.g_proj = reorder_linear_weights(
+            module.g_proj, layer_full_attention_heads, module.num_key_value_groups, "out"
+        )
+        module.o_proj = reorder_linear_weights(
+            module.o_proj,
+            layer_full_attention_heads,
+            module.num_key_value_groups * module.head_dim,
+            "in",
+        )
+        layer_full_attention_heads = reorder_full_attn_heads(layer_full_attention_heads)
+        module.sink_size = sink_size
+        module.recent_size = recent_size
+        module.register_buffer("full_attention_heads", layer_full_attention_heads)
+
+    model.forward = types.MethodType(laguna_for_causal_lm_forward, model)
+    return model

modeling_duo_laguna.py

@@ -0,0 +1,103 @@
+import os
+import json
+
+import torch
+from huggingface_hub import hf_hub_download
+from transformers.models.laguna.modeling_laguna import LagunaForCausalLM
+
+from .duo_laguna_remote import enable_duo_laguna_eval
+
+
+def _load_repo_json(pretrained_model_name_or_path, filename, revision=None):
+    if os.path.isdir(pretrained_model_name_or_path):
+        path = os.path.join(pretrained_model_name_or_path, filename)
+    else:
+        path = hf_hub_download(
+            repo_id=pretrained_model_name_or_path,
+            filename=filename,
+            revision=revision,
+        )
+    with open(path) as f:
+        return json.load(f)
+
+
+def _load_duo_tensor(pretrained_model_name_or_path, filename, revision=None):
+    if os.path.isdir(pretrained_model_name_or_path):
+        path = os.path.join(pretrained_model_name_or_path, filename)
+    else:
+        path = hf_hub_download(
+            repo_id=pretrained_model_name_or_path,
+            filename=filename,
+            revision=revision,
+        )
+    return torch.load(path, map_location="cpu", weights_only=True)
+
+
+class DuoLagunaForCausalLM(LagunaForCausalLM):
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs):
+        duo_attention = kwargs.pop("duo_attention", True)
+        duo_sink_size = kwargs.pop("duo_sink_size", None)
+        duo_recent_size = kwargs.pop("duo_recent_size", None)
+        duo_heads_file = kwargs.pop("duo_heads_file", None)
+        revision = kwargs.get("revision")
+        config_dict = _load_repo_json(
+            pretrained_model_name_or_path,
+            "config.json",
+            revision=revision,
+        )
+        duo_config = config_dict.get("duo_attention")
+        base_model_name_or_path = kwargs.pop("duo_base_model_name_or_path", None)
+        if duo_config is not None:
+            base_model_name_or_path = (
+                base_model_name_or_path
+                or duo_config.get("base_model_name_or_path")
+                or duo_config.get("base_model")
+            )
+        base_revision = kwargs.pop("duo_base_revision", None)
+        if duo_config is not None:
+            base_revision = base_revision or duo_config.get("base_model_revision")
+
+        load_path = base_model_name_or_path or pretrained_model_name_or_path
+        if base_revision is not None:
+            kwargs["revision"] = base_revision
+
+        model = super().from_pretrained(
+            load_path,
+            *model_args,
+            **kwargs,
+        )
+        if base_revision is not None and revision is not None:
+            kwargs["revision"] = revision
+
+        if not duo_attention:
+            return model
+
+        if duo_config is None:
+            duo_config = getattr(model.config, "duo_attention", None)
+        if duo_config is None:
+            raise ValueError(
+                "This repository does not define config.duo_attention; "
+                "reload with duo_attention=False to use the unpatched Laguna model."
+            )
+
+        full_attention_heads = _load_duo_tensor(
+            pretrained_model_name_or_path,
+            duo_heads_file or duo_config["full_attention_heads_file"],
+            revision=revision,
+        )
+        sink_size = duo_sink_size or duo_config["sink_size"]
+        recent_size = duo_recent_size or duo_config["recent_size"]
+
+        enable_duo_laguna_eval(
+            model,
+            full_attention_heads,
+            sink_size=sink_size,
+            recent_size=recent_size,
+        )
+        model.duo_attention_config = {
+            **duo_config,
+            "sink_size": sink_size,
+            "recent_size": recent_size,
+        }
+        return model

requirements.txt

@@ -0,0 +1,4 @@
+torch
+transformers>=5.9.0
+huggingface_hub
+numpy