upload model

__init__.py

@@ -0,0 +1,4 @@
+from .configuration_qwen2nsa import Qwen2NsaConfig
+from .modeling_qwen2nsa import Qwen2NsaForCausalLM
+
+__all__ = ["Qwen2NsaConfig", "Qwen2NsaForCausalLM"]

config.json

@@ -0,0 +1,90 @@
+{
+  "_attn_implementation_autoset": false,
+  "_name_or_path": "",
+  "add_cross_attention": false,
+  "architectures": [
+    "Qwen2NsaForCausalLM"
+  ],
+  "attention_dropout": 0.0,
+  "bad_words_ids": null,
+  "begin_suppress_tokens": null,
+  "bos_token_id": 151643,
+  "chunk_size_feed_forward": 0,
+  "cross_attention_hidden_size": null,
+  "decoder_start_token_id": null,
+  "diversity_penalty": 0.0,
+  "do_sample": true,
+  "early_stopping": false,
+  "encoder_no_repeat_ngram_size": 0,
+  "eos_token_id": 151643,
+  "exponential_decay_length_penalty": null,
+  "finetuning_task": null,
+  "forced_bos_token_id": null,
+  "forced_eos_token_id": null,
+  "gradient_checkpointing": false,
+  "hidden_act": "silu",
+  "hidden_size": 1536,
+  "id2label": {
+    "0": "LABEL_0",
+    "1": "LABEL_1"
+  },
+  "initializer_range": 0.02,
+  "intermediate_size": 8960,
+  "is_decoder": false,
+  "is_encoder_decoder": false,
+  "label2id": {
+    "LABEL_0": 0,
+    "LABEL_1": 1
+  },
+  "length_penalty": 1.0,
+  "max_length": 20,
+  "max_position_embeddings": 131072,
+  "max_window_layers": 21,
+  "min_length": 0,
+  "model_type": "qwen2nsa",
+  "auto_map": {
+    "AutoConfig": "configuration_qwen2nsa.Qwen2NsaConfig",
+    "AutoModelForCausalLM": "modeling_qwen2nsa.Qwen2NsaForCausalLM"
+  },
+  "no_repeat_ngram_size": 0,
+  "num_attention_heads": 12,
+  "num_beam_groups": 1,
+  "num_beams": 1,
+  "num_hidden_layers": 28,
+  "num_key_value_heads": 2,
+  "num_return_sequences": 1,
+  "output_attentions": false,
+  "output_hidden_states": false,
+  "output_scores": false,
+  "pad_token_id": 0,
+  "prefix": null,
+  "problem_type": null,
+  "pruned_heads": {},
+  "remove_invalid_values": false,
+  "repetition_penalty": 1.1,
+  "return_dict": true,
+  "return_dict_in_generate": false,
+  "rms_norm_eps": 1e-06,
+  "rope_scaling": null,
+  "rope_theta": 10000,
+  "sep_token_id": null,
+  "sliding_window": 4096,
+  "suppress_tokens": null,
+  "task_specific_params": null,
+  "temperature": 0.6,
+  "tf_legacy_loss": false,
+  "tie_encoder_decoder": false,
+  "tie_word_embeddings": false,
+  "tokenizer_class": null,
+  "top_k": 0,
+  "top_p": 0.95,
+  "torch_dtype": "float32",
+  "torchscript": false,
+  "transformers_version": "4.51.3",
+  "typical_p": 1.0,
+  "use_bfloat16": false,
+  "use_cache": true,
+  "use_mrope": false,
+  "use_sliding_window": false,
+  "vocab_size": 151936
+}

configuration_qwen2nsa.py

@@ -0,0 +1,74 @@
+# SPDX-License-Identifier: Apache-2.0
+"""
+Configuration class for Qwen-2-NSA
+"""
+
+import json
+from typing import Optional, Dict, Any
+
+from transformers import PretrainedConfig
+
+
+class Qwen2NsaConfig(PretrainedConfig):
+    model_type: str = "qwen2nsa"
+
+    def __init__(
+        self,
+        vocab_size: int = 151_936,
+        hidden_size: int = 2_048,
+        intermediate_size: int = 5_632,
+        num_hidden_layers: int = 24,
+        num_attention_heads: int = 16,
+        num_key_value_heads: int = 16,
+        hidden_act: str = "silu",
+        max_position_embeddings: int = 32_768,
+        initializer_range: float = 0.02,
+        rms_norm_eps: float = 1e-6,
+        tie_word_embeddings: bool = False,
+        pad_token_id: int = 0,
+        bos_token_id: int = 1,
+        eos_token_id: int = 2,
+        sliding_window: int = 128,
+        rope_theta: float = 10_000.0,
+        rope_scaling: Optional[Dict[str, Any]] = None,
+        attention_dropout: float = 0.0,
+        use_cache: bool = True,
+        use_sliding_window: bool = False,
+        gradient_checkpointing: bool = False,
+        **kwargs,
+    ):
+        super().__init__(
+            pad_token_id=pad_token_id,
+            bos_token_id=bos_token_id,
+            eos_token_id=eos_token_id,
+            tie_word_embeddings=tie_word_embeddings,
+            **kwargs,
+        )
+
+        self.vocab_size = vocab_size
+        self.hidden_size = hidden_size
+        self.intermediate_size = intermediate_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.num_key_value_heads = num_key_value_heads
+        self.hidden_act = hidden_act
+        self.max_position_embeddings = max_position_embeddings
+        self.initializer_range = initializer_range
+        self.rms_norm_eps = rms_norm_eps
+        self.sliding_window = sliding_window
+        self.rope_theta = rope_theta
+        self.rope_scaling = rope_scaling
+        self.attention_dropout = attention_dropout
+        self.use_cache = use_cache
+        self.use_sliding_window = use_sliding_window
+        self.gradient_checkpointing = gradient_checkpointing
+
+    @classmethod
+    def from_json(cls, json_str: str) -> "Qwen2NsaConfig":
+        return cls(**json.loads(json_str))
+
+    def to_json_string(self, use_diff: bool = True) -> str:
+        return json.dumps(self.to_dict(), indent=2, sort_keys=True)
+
+    def get(self, key, default=None):
+        return getattr(self, key, default)

generation_config.json

@@ -0,0 +1,7 @@
+{
+  "_from_model_config": true,
+  "bos_token_id": 151643,
+  "eos_token_id": 151643,
+  "pad_token_id": 0,
+  "transformers_version": "4.51.3"
+}

model-00001-of-00003.safetensors

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

model-00002-of-00003.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:49f17db38d254ce91a705fc833cdfcb5a07e633435082b9b917acd9979880bf7
+size 4953514664

model-00003-of-00003.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:499484d278654a89fb0dfefb771bf3b107e3183417637a697bb6364a257d2e3c
+size 1043614352

model.safetensors.index.json

@@ -0,0 +1,654 @@
+{
+  "metadata": {
+    "total_size": 10993103808
+  },
+  "weight_map": {
+    "lm_head.weight": "model-00003-of-00003.safetensors",
+    "model.embed_tokens.weight": "model-00001-of-00003.safetensors",
+    "model.layers.0.mlp.down_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.0.mlp.gate_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.0.mlp.up_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.0.post_attn_norm.weight": "model-00001-of-00003.safetensors",
+    "model.layers.0.pre_attn_norm.weight": "model-00001-of-00003.safetensors",
+    "model.layers.0.self_attn.combine_heads.weight": "model-00001-of-00003.safetensors",
+    "model.layers.0.self_attn.compress_mem_kv": "model-00001-of-00003.safetensors",
+    "model.layers.0.self_attn.k_compress.1.bias": "model-00001-of-00003.safetensors",
+    "model.layers.0.self_attn.k_compress.1.weight": "model-00001-of-00003.safetensors",
+    "model.layers.0.self_attn.k_compress.3.bias": "model-00001-of-00003.safetensors",
+    "model.layers.0.self_attn.k_compress.3.weight": "model-00001-of-00003.safetensors",
+    "model.layers.0.self_attn.k_intrablock_positions": "model-00001-of-00003.safetensors",
+    "model.layers.0.self_attn.norm.weight": "model-00001-of-00003.safetensors",
+    "model.layers.0.self_attn.rotary_emb.freqs": "model-00001-of-00003.safetensors",
+    "model.layers.0.self_attn.to_qkv.bias": "model-00001-of-00003.safetensors",
+    "model.layers.0.self_attn.to_qkv.weight": "model-00001-of-00003.safetensors",
+    "model.layers.0.self_attn.to_strategy_combine.0.bias": "model-00001-of-00003.safetensors",
+    "model.layers.0.self_attn.to_strategy_combine.0.weight": "model-00001-of-00003.safetensors",
+    "model.layers.0.self_attn.v_compress.1.bias": "model-00001-of-00003.safetensors",
+    "model.layers.0.self_attn.v_compress.1.weight": "model-00001-of-00003.safetensors",
+    "model.layers.0.self_attn.v_compress.3.bias": "model-00001-of-00003.safetensors",
+    "model.layers.0.self_attn.v_compress.3.weight": "model-00001-of-00003.safetensors",
+    "model.layers.0.self_attn.v_intrablock_positions": "model-00001-of-00003.safetensors",
+    "model.layers.1.mlp.down_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.1.mlp.gate_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.1.mlp.up_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.1.post_attn_norm.weight": "model-00001-of-00003.safetensors",
+    "model.layers.1.pre_attn_norm.weight": "model-00001-of-00003.safetensors",
+    "model.layers.1.self_attn.combine_heads.weight": "model-00001-of-00003.safetensors",
+    "model.layers.1.self_attn.compress_mem_kv": "model-00001-of-00003.safetensors",
+    "model.layers.1.self_attn.k_compress.1.bias": "model-00001-of-00003.safetensors",
+    "model.layers.1.self_attn.k_compress.1.weight": "model-00001-of-00003.safetensors",
+    "model.layers.1.self_attn.k_compress.3.bias": "model-00001-of-00003.safetensors",
+    "model.layers.1.self_attn.k_compress.3.weight": "model-00001-of-00003.safetensors",
+    "model.layers.1.self_attn.k_intrablock_positions": "model-00001-of-00003.safetensors",
+    "model.layers.1.self_attn.norm.weight": "model-00001-of-00003.safetensors",
+    "model.layers.1.self_attn.rotary_emb.freqs": "model-00001-of-00003.safetensors",
+    "model.layers.1.self_attn.to_qkv.bias": "model-00001-of-00003.safetensors",
+    "model.layers.1.self_attn.to_qkv.weight": "model-00001-of-00003.safetensors",
+    "model.layers.1.self_attn.to_strategy_combine.0.bias": "model-00001-of-00003.safetensors",
+    "model.layers.1.self_attn.to_strategy_combine.0.weight": "model-00001-of-00003.safetensors",
+    "model.layers.1.self_attn.v_compress.1.bias": "model-00001-of-00003.safetensors",
+    "model.layers.1.self_attn.v_compress.1.weight": "model-00001-of-00003.safetensors",
+    "model.layers.1.self_attn.v_compress.3.bias": "model-00001-of-00003.safetensors",
+    "model.layers.1.self_attn.v_compress.3.weight": "model-00001-of-00003.safetensors",
+    "model.layers.1.self_attn.v_intrablock_positions": "model-00001-of-00003.safetensors",
+    "model.layers.10.mlp.down_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.10.mlp.gate_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.10.mlp.up_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.10.post_attn_norm.weight": "model-00001-of-00003.safetensors",
+    "model.layers.10.pre_attn_norm.weight": "model-00001-of-00003.safetensors",
+    "model.layers.10.self_attn.combine_heads.weight": "model-00001-of-00003.safetensors",
+    "model.layers.10.self_attn.compress_mem_kv": "model-00001-of-00003.safetensors",
+    "model.layers.10.self_attn.k_compress.1.bias": "model-00001-of-00003.safetensors",
+    "model.layers.10.self_attn.k_compress.1.weight": "model-00001-of-00003.safetensors",
+    "model.layers.10.self_attn.k_compress.3.bias": "model-00001-of-00003.safetensors",
+    "model.layers.10.self_attn.k_compress.3.weight": "model-00001-of-00003.safetensors",
+    "model.layers.10.self_attn.k_intrablock_positions": "model-00001-of-00003.safetensors",
+    "model.layers.10.self_attn.norm.weight": "model-00001-of-00003.safetensors",
+    "model.layers.10.self_attn.rotary_emb.freqs": "model-00001-of-00003.safetensors",
+    "model.layers.10.self_attn.to_qkv.bias": "model-00001-of-00003.safetensors",
+    "model.layers.10.self_attn.to_qkv.weight": "model-00001-of-00003.safetensors",
+    "model.layers.10.self_attn.to_strategy_combine.0.bias": "model-00001-of-00003.safetensors",
+    "model.layers.10.self_attn.to_strategy_combine.0.weight": "model-00001-of-00003.safetensors",
+    "model.layers.10.self_attn.v_compress.1.bias": "model-00001-of-00003.safetensors",
+    "model.layers.10.self_attn.v_compress.1.weight": "model-00001-of-00003.safetensors",
+    "model.layers.10.self_attn.v_compress.3.bias": "model-00001-of-00003.safetensors",
+    "model.layers.10.self_attn.v_compress.3.weight": "model-00001-of-00003.safetensors",
+    "model.layers.10.self_attn.v_intrablock_positions": "model-00001-of-00003.safetensors",
+    "model.layers.11.mlp.down_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.11.mlp.gate_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.11.mlp.up_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.11.post_attn_norm.weight": "model-00001-of-00003.safetensors",
+    "model.layers.11.pre_attn_norm.weight": "model-00001-of-00003.safetensors",
+    "model.layers.11.self_attn.combine_heads.weight": "model-00001-of-00003.safetensors",
+    "model.layers.11.self_attn.compress_mem_kv": "model-00001-of-00003.safetensors",
+    "model.layers.11.self_attn.k_compress.1.bias": "model-00001-of-00003.safetensors",
+    "model.layers.11.self_attn.k_compress.1.weight": "model-00001-of-00003.safetensors",
+    "model.layers.11.self_attn.k_compress.3.bias": "model-00001-of-00003.safetensors",
+    "model.layers.11.self_attn.k_compress.3.weight": "model-00001-of-00003.safetensors",
+    "model.layers.11.self_attn.k_intrablock_positions": "model-00001-of-00003.safetensors",
+    "model.layers.11.self_attn.norm.weight": "model-00001-of-00003.safetensors",
+    "model.layers.11.self_attn.rotary_emb.freqs": "model-00001-of-00003.safetensors",
+    "model.layers.11.self_attn.to_qkv.bias": "model-00001-of-00003.safetensors",
+    "model.layers.11.self_attn.to_qkv.weight": "model-00001-of-00003.safetensors",
+    "model.layers.11.self_attn.to_strategy_combine.0.bias": "model-00001-of-00003.safetensors",
+    "model.layers.11.self_attn.to_strategy_combine.0.weight": "model-00001-of-00003.safetensors",
+    "model.layers.11.self_attn.v_compress.1.bias": "model-00001-of-00003.safetensors",
+    "model.layers.11.self_attn.v_compress.1.weight": "model-00001-of-00003.safetensors",
+    "model.layers.11.self_attn.v_compress.3.bias": "model-00001-of-00003.safetensors",
+    "model.layers.11.self_attn.v_compress.3.weight": "model-00001-of-00003.safetensors",
+    "model.layers.11.self_attn.v_intrablock_positions": "model-00001-of-00003.safetensors",
+    "model.layers.12.mlp.down_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.12.mlp.gate_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.12.mlp.up_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.12.post_attn_norm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.12.pre_attn_norm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.12.self_attn.combine_heads.weight": "model-00002-of-00003.safetensors",
+    "model.layers.12.self_attn.compress_mem_kv": "model-00001-of-00003.safetensors",
+    "model.layers.12.self_attn.k_compress.1.bias": "model-00001-of-00003.safetensors",
+    "model.layers.12.self_attn.k_compress.1.weight": "model-00001-of-00003.safetensors",
+    "model.layers.12.self_attn.k_compress.3.bias": "model-00001-of-00003.safetensors",
+    "model.layers.12.self_attn.k_compress.3.weight": "model-00001-of-00003.safetensors",
+    "model.layers.12.self_attn.k_intrablock_positions": "model-00001-of-00003.safetensors",
+    "model.layers.12.self_attn.norm.weight": "model-00001-of-00003.safetensors",
+    "model.layers.12.self_attn.rotary_emb.freqs": "model-00001-of-00003.safetensors",
+    "model.layers.12.self_attn.to_qkv.bias": "model-00001-of-00003.safetensors",
+    "model.layers.12.self_attn.to_qkv.weight": "model-00001-of-00003.safetensors",
+    "model.layers.12.self_attn.to_strategy_combine.0.bias": "model-00001-of-00003.safetensors",
+    "model.layers.12.self_attn.to_strategy_combine.0.weight": "model-00001-of-00003.safetensors",
+    "model.layers.12.self_attn.v_compress.1.bias": "model-00001-of-00003.safetensors",
+    "model.layers.12.self_attn.v_compress.1.weight": "model-00001-of-00003.safetensors",
+    "model.layers.12.self_attn.v_compress.3.bias": "model-00001-of-00003.safetensors",
+    "model.layers.12.self_attn.v_compress.3.weight": "model-00001-of-00003.safetensors",
+    "model.layers.12.self_attn.v_intrablock_positions": "model-00001-of-00003.safetensors",
+    "model.layers.13.mlp.down_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.13.mlp.gate_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.13.mlp.up_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.13.post_attn_norm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.13.pre_attn_norm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.13.self_attn.combine_heads.weight": "model-00002-of-00003.safetensors",
+    "model.layers.13.self_attn.compress_mem_kv": "model-00002-of-00003.safetensors",
+    "model.layers.13.self_attn.k_compress.1.bias": "model-00002-of-00003.safetensors",
+    "model.layers.13.self_attn.k_compress.1.weight": "model-00002-of-00003.safetensors",
+    "model.layers.13.self_attn.k_compress.3.bias": "model-00002-of-00003.safetensors",
+    "model.layers.13.self_attn.k_compress.3.weight": "model-00002-of-00003.safetensors",
+    "model.layers.13.self_attn.k_intrablock_positions": "model-00002-of-00003.safetensors",
+    "model.layers.13.self_attn.norm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.13.self_attn.rotary_emb.freqs": "model-00002-of-00003.safetensors",
+    "model.layers.13.self_attn.to_qkv.bias": "model-00002-of-00003.safetensors",
+    "model.layers.13.self_attn.to_qkv.weight": "model-00002-of-00003.safetensors",
+    "model.layers.13.self_attn.to_strategy_combine.0.bias": "model-00002-of-00003.safetensors",
+    "model.layers.13.self_attn.to_strategy_combine.0.weight": "model-00002-of-00003.safetensors",
+    "model.layers.13.self_attn.v_compress.1.bias": "model-00002-of-00003.safetensors",
+    "model.layers.13.self_attn.v_compress.1.weight": "model-00002-of-00003.safetensors",
+    "model.layers.13.self_attn.v_compress.3.bias": "model-00002-of-00003.safetensors",
+    "model.layers.13.self_attn.v_compress.3.weight": "model-00002-of-00003.safetensors",
+    "model.layers.13.self_attn.v_intrablock_positions": "model-00002-of-00003.safetensors",
+    "model.layers.14.mlp.down_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.14.mlp.gate_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.14.mlp.up_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.14.post_attn_norm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.14.pre_attn_norm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.14.self_attn.combine_heads.weight": "model-00002-of-00003.safetensors",
+    "model.layers.14.self_attn.compress_mem_kv": "model-00002-of-00003.safetensors",
+    "model.layers.14.self_attn.k_compress.1.bias": "model-00002-of-00003.safetensors",
+    "model.layers.14.self_attn.k_compress.1.weight": "model-00002-of-00003.safetensors",
+    "model.layers.14.self_attn.k_compress.3.bias": "model-00002-of-00003.safetensors",
+    "model.layers.14.self_attn.k_compress.3.weight": "model-00002-of-00003.safetensors",
+    "model.layers.14.self_attn.k_intrablock_positions": "model-00002-of-00003.safetensors",
+    "model.layers.14.self_attn.norm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.14.self_attn.rotary_emb.freqs": "model-00002-of-00003.safetensors",
+    "model.layers.14.self_attn.to_qkv.bias": "model-00002-of-00003.safetensors",
+    "model.layers.14.self_attn.to_qkv.weight": "model-00002-of-00003.safetensors",
+    "model.layers.14.self_attn.to_strategy_combine.0.bias": "model-00002-of-00003.safetensors",
+    "model.layers.14.self_attn.to_strategy_combine.0.weight": "model-00002-of-00003.safetensors",
+    "model.layers.14.self_attn.v_compress.1.bias": "model-00002-of-00003.safetensors",
+    "model.layers.14.self_attn.v_compress.1.weight": "model-00002-of-00003.safetensors",
+    "model.layers.14.self_attn.v_compress.3.bias": "model-00002-of-00003.safetensors",
+    "model.layers.14.self_attn.v_compress.3.weight": "model-00002-of-00003.safetensors",
+    "model.layers.14.self_attn.v_intrablock_positions": "model-00002-of-00003.safetensors",
+    "model.layers.15.mlp.down_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.15.mlp.gate_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.15.mlp.up_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.15.post_attn_norm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.15.pre_attn_norm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.15.self_attn.combine_heads.weight": "model-00002-of-00003.safetensors",
+    "model.layers.15.self_attn.compress_mem_kv": "model-00002-of-00003.safetensors",
+    "model.layers.15.self_attn.k_compress.1.bias": "model-00002-of-00003.safetensors",
+    "model.layers.15.self_attn.k_compress.1.weight": "model-00002-of-00003.safetensors",
+    "model.layers.15.self_attn.k_compress.3.bias": "model-00002-of-00003.safetensors",
+    "model.layers.15.self_attn.k_compress.3.weight": "model-00002-of-00003.safetensors",
+    "model.layers.15.self_attn.k_intrablock_positions": "model-00002-of-00003.safetensors",
+    "model.layers.15.self_attn.norm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.15.self_attn.rotary_emb.freqs": "model-00002-of-00003.safetensors",
+    "model.layers.15.self_attn.to_qkv.bias": "model-00002-of-00003.safetensors",
+    "model.layers.15.self_attn.to_qkv.weight": "model-00002-of-00003.safetensors",
+    "model.layers.15.self_attn.to_strategy_combine.0.bias": "model-00002-of-00003.safetensors",
+    "model.layers.15.self_attn.to_strategy_combine.0.weight": "model-00002-of-00003.safetensors",
+    "model.layers.15.self_attn.v_compress.1.bias": "model-00002-of-00003.safetensors",
+    "model.layers.15.self_attn.v_compress.1.weight": "model-00002-of-00003.safetensors",
+    "model.layers.15.self_attn.v_compress.3.bias": "model-00002-of-00003.safetensors",
+    "model.layers.15.self_attn.v_compress.3.weight": "model-00002-of-00003.safetensors",
+    "model.layers.15.self_attn.v_intrablock_positions": "model-00002-of-00003.safetensors",
+    "model.layers.16.mlp.down_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.16.mlp.gate_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.16.mlp.up_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.16.post_attn_norm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.16.pre_attn_norm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.16.self_attn.combine_heads.weight": "model-00002-of-00003.safetensors",
+    "model.layers.16.self_attn.compress_mem_kv": "model-00002-of-00003.safetensors",
+    "model.layers.16.self_attn.k_compress.1.bias": "model-00002-of-00003.safetensors",
+    "model.layers.16.self_attn.k_compress.1.weight": "model-00002-of-00003.safetensors",
+    "model.layers.16.self_attn.k_compress.3.bias": "model-00002-of-00003.safetensors",
+    "model.layers.16.self_attn.k_compress.3.weight": "model-00002-of-00003.safetensors",
+    "model.layers.16.self_attn.k_intrablock_positions": "model-00002-of-00003.safetensors",
+    "model.layers.16.self_attn.norm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.16.self_attn.rotary_emb.freqs": "model-00002-of-00003.safetensors",
+    "model.layers.16.self_attn.to_qkv.bias": "model-00002-of-00003.safetensors",
+    "model.layers.16.self_attn.to_qkv.weight": "model-00002-of-00003.safetensors",
+    "model.layers.16.self_attn.to_strategy_combine.0.bias": "model-00002-of-00003.safetensors",
+    "model.layers.16.self_attn.to_strategy_combine.0.weight": "model-00002-of-00003.safetensors",
+    "model.layers.16.self_attn.v_compress.1.bias": "model-00002-of-00003.safetensors",
+    "model.layers.16.self_attn.v_compress.1.weight": "model-00002-of-00003.safetensors",
+    "model.layers.16.self_attn.v_compress.3.bias": "model-00002-of-00003.safetensors",
+    "model.layers.16.self_attn.v_compress.3.weight": "model-00002-of-00003.safetensors",
+    "model.layers.16.self_attn.v_intrablock_positions": "model-00002-of-00003.safetensors",
+    "model.layers.17.mlp.down_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.17.mlp.gate_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.17.mlp.up_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.17.post_attn_norm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.17.pre_attn_norm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.17.self_attn.combine_heads.weight": "model-00002-of-00003.safetensors",
+    "model.layers.17.self_attn.compress_mem_kv": "model-00002-of-00003.safetensors",
+    "model.layers.17.self_attn.k_compress.1.bias": "model-00002-of-00003.safetensors",
+    "model.layers.17.self_attn.k_compress.1.weight": "model-00002-of-00003.safetensors",
+    "model.layers.17.self_attn.k_compress.3.bias": "model-00002-of-00003.safetensors",
+    "model.layers.17.self_attn.k_compress.3.weight": "model-00002-of-00003.safetensors",
+    "model.layers.17.self_attn.k_intrablock_positions": "model-00002-of-00003.safetensors",
+    "model.layers.17.self_attn.norm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.17.self_attn.rotary_emb.freqs": "model-00002-of-00003.safetensors",
+    "model.layers.17.self_attn.to_qkv.bias": "model-00002-of-00003.safetensors",
+    "model.layers.17.self_attn.to_qkv.weight": "model-00002-of-00003.safetensors",
+    "model.layers.17.self_attn.to_strategy_combine.0.bias": "model-00002-of-00003.safetensors",
+    "model.layers.17.self_attn.to_strategy_combine.0.weight": "model-00002-of-00003.safetensors",
+    "model.layers.17.self_attn.v_compress.1.bias": "model-00002-of-00003.safetensors",
+    "model.layers.17.self_attn.v_compress.1.weight": "model-00002-of-00003.safetensors",
+    "model.layers.17.self_attn.v_compress.3.bias": "model-00002-of-00003.safetensors",
+    "model.layers.17.self_attn.v_compress.3.weight": "model-00002-of-00003.safetensors",
+    "model.layers.17.self_attn.v_intrablock_positions": "model-00002-of-00003.safetensors",
+    "model.layers.18.mlp.down_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.18.mlp.gate_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.18.mlp.up_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.18.post_attn_norm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.18.pre_attn_norm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.18.self_attn.combine_heads.weight": "model-00002-of-00003.safetensors",
+    "model.layers.18.self_attn.compress_mem_kv": "model-00002-of-00003.safetensors",
+    "model.layers.18.self_attn.k_compress.1.bias": "model-00002-of-00003.safetensors",
+    "model.layers.18.self_attn.k_compress.1.weight": "model-00002-of-00003.safetensors",
+    "model.layers.18.self_attn.k_compress.3.bias": "model-00002-of-00003.safetensors",
+    "model.layers.18.self_attn.k_compress.3.weight": "model-00002-of-00003.safetensors",
+    "model.layers.18.self_attn.k_intrablock_positions": "model-00002-of-00003.safetensors",
+    "model.layers.18.self_attn.norm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.18.self_attn.rotary_emb.freqs": "model-00002-of-00003.safetensors",
+    "model.layers.18.self_attn.to_qkv.bias": "model-00002-of-00003.safetensors",
+    "model.layers.18.self_attn.to_qkv.weight": "model-00002-of-00003.safetensors",
+    "model.layers.18.self_attn.to_strategy_combine.0.bias": "model-00002-of-00003.safetensors",
+    "model.layers.18.self_attn.to_strategy_combine.0.weight": "model-00002-of-00003.safetensors",
+    "model.layers.18.self_attn.v_compress.1.bias": "model-00002-of-00003.safetensors",
+    "model.layers.18.self_attn.v_compress.1.weight": "model-00002-of-00003.safetensors",
+    "model.layers.18.self_attn.v_compress.3.bias": "model-00002-of-00003.safetensors",
+    "model.layers.18.self_attn.v_compress.3.weight": "model-00002-of-00003.safetensors",
+    "model.layers.18.self_attn.v_intrablock_positions": "model-00002-of-00003.safetensors",
+    "model.layers.19.mlp.down_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.19.mlp.gate_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.19.mlp.up_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.19.post_attn_norm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.19.pre_attn_norm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.19.self_attn.combine_heads.weight": "model-00002-of-00003.safetensors",
+    "model.layers.19.self_attn.compress_mem_kv": "model-00002-of-00003.safetensors",
+    "model.layers.19.self_attn.k_compress.1.bias": "model-00002-of-00003.safetensors",
+    "model.layers.19.self_attn.k_compress.1.weight": "model-00002-of-00003.safetensors",
+    "model.layers.19.self_attn.k_compress.3.bias": "model-00002-of-00003.safetensors",
+    "model.layers.19.self_attn.k_compress.3.weight": "model-00002-of-00003.safetensors",
+    "model.layers.19.self_attn.k_intrablock_positions": "model-00002-of-00003.safetensors",
+    "model.layers.19.self_attn.norm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.19.self_attn.rotary_emb.freqs": "model-00002-of-00003.safetensors",
+    "model.layers.19.self_attn.to_qkv.bias": "model-00002-of-00003.safetensors",
+    "model.layers.19.self_attn.to_qkv.weight": "model-00002-of-00003.safetensors",
+    "model.layers.19.self_attn.to_strategy_combine.0.bias": "model-00002-of-00003.safetensors",
+    "model.layers.19.self_attn.to_strategy_combine.0.weight": "model-00002-of-00003.safetensors",
+    "model.layers.19.self_attn.v_compress.1.bias": "model-00002-of-00003.safetensors",
+    "model.layers.19.self_attn.v_compress.1.weight": "model-00002-of-00003.safetensors",
+    "model.layers.19.self_attn.v_compress.3.bias": "model-00002-of-00003.safetensors",
+    "model.layers.19.self_attn.v_compress.3.weight": "model-00002-of-00003.safetensors",
+    "model.layers.19.self_attn.v_intrablock_positions": "model-00002-of-00003.safetensors",
+    "model.layers.2.mlp.down_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.2.mlp.gate_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.2.mlp.up_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.2.post_attn_norm.weight": "model-00001-of-00003.safetensors",
+    "model.layers.2.pre_attn_norm.weight": "model-00001-of-00003.safetensors",
+    "model.layers.2.self_attn.combine_heads.weight": "model-00001-of-00003.safetensors",
+    "model.layers.2.self_attn.compress_mem_kv": "model-00001-of-00003.safetensors",
+    "model.layers.2.self_attn.k_compress.1.bias": "model-00001-of-00003.safetensors",
+    "model.layers.2.self_attn.k_compress.1.weight": "model-00001-of-00003.safetensors",
+    "model.layers.2.self_attn.k_compress.3.bias": "model-00001-of-00003.safetensors",
+    "model.layers.2.self_attn.k_compress.3.weight": "model-00001-of-00003.safetensors",
+    "model.layers.2.self_attn.k_intrablock_positions": "model-00001-of-00003.safetensors",
+    "model.layers.2.self_attn.norm.weight": "model-00001-of-00003.safetensors",
+    "model.layers.2.self_attn.rotary_emb.freqs": "model-00001-of-00003.safetensors",
+    "model.layers.2.self_attn.to_qkv.bias": "model-00001-of-00003.safetensors",
+    "model.layers.2.self_attn.to_qkv.weight": "model-00001-of-00003.safetensors",
+    "model.layers.2.self_attn.to_strategy_combine.0.bias": "model-00001-of-00003.safetensors",
+    "model.layers.2.self_attn.to_strategy_combine.0.weight": "model-00001-of-00003.safetensors",
+    "model.layers.2.self_attn.v_compress.1.bias": "model-00001-of-00003.safetensors",
+    "model.layers.2.self_attn.v_compress.1.weight": "model-00001-of-00003.safetensors",
+    "model.layers.2.self_attn.v_compress.3.bias": "model-00001-of-00003.safetensors",
+    "model.layers.2.self_attn.v_compress.3.weight": "model-00001-of-00003.safetensors",
+    "model.layers.2.self_attn.v_intrablock_positions": "model-00001-of-00003.safetensors",
+    "model.layers.20.mlp.down_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.20.mlp.gate_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.20.mlp.up_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.20.post_attn_norm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.20.pre_attn_norm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.20.self_attn.combine_heads.weight": "model-00002-of-00003.safetensors",
+    "model.layers.20.self_attn.compress_mem_kv": "model-00002-of-00003.safetensors",
+    "model.layers.20.self_attn.k_compress.1.bias": "model-00002-of-00003.safetensors",
+    "model.layers.20.self_attn.k_compress.1.weight": "model-00002-of-00003.safetensors",
+    "model.layers.20.self_attn.k_compress.3.bias": "model-00002-of-00003.safetensors",
+    "model.layers.20.self_attn.k_compress.3.weight": "model-00002-of-00003.safetensors",
+    "model.layers.20.self_attn.k_intrablock_positions": "model-00002-of-00003.safetensors",
+    "model.layers.20.self_attn.norm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.20.self_attn.rotary_emb.freqs": "model-00002-of-00003.safetensors",
+    "model.layers.20.self_attn.to_qkv.bias": "model-00002-of-00003.safetensors",
+    "model.layers.20.self_attn.to_qkv.weight": "model-00002-of-00003.safetensors",
+    "model.layers.20.self_attn.to_strategy_combine.0.bias": "model-00002-of-00003.safetensors",
+    "model.layers.20.self_attn.to_strategy_combine.0.weight": "model-00002-of-00003.safetensors",
+    "model.layers.20.self_attn.v_compress.1.bias": "model-00002-of-00003.safetensors",
+    "model.layers.20.self_attn.v_compress.1.weight": "model-00002-of-00003.safetensors",
+    "model.layers.20.self_attn.v_compress.3.bias": "model-00002-of-00003.safetensors",
+    "model.layers.20.self_attn.v_compress.3.weight": "model-00002-of-00003.safetensors",
+    "model.layers.20.self_attn.v_intrablock_positions": "model-00002-of-00003.safetensors",
+    "model.layers.21.mlp.down_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.21.mlp.gate_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.21.mlp.up_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.21.post_attn_norm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.21.pre_attn_norm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.21.self_attn.combine_heads.weight": "model-00002-of-00003.safetensors",
+    "model.layers.21.self_attn.compress_mem_kv": "model-00002-of-00003.safetensors",
+    "model.layers.21.self_attn.k_compress.1.bias": "model-00002-of-00003.safetensors",
+    "model.layers.21.self_attn.k_compress.1.weight": "model-00002-of-00003.safetensors",
+    "model.layers.21.self_attn.k_compress.3.bias": "model-00002-of-00003.safetensors",
+    "model.layers.21.self_attn.k_compress.3.weight": "model-00002-of-00003.safetensors",
+    "model.layers.21.self_attn.k_intrablock_positions": "model-00002-of-00003.safetensors",
+    "model.layers.21.self_attn.norm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.21.self_attn.rotary_emb.freqs": "model-00002-of-00003.safetensors",
+    "model.layers.21.self_attn.to_qkv.bias": "model-00002-of-00003.safetensors",
+    "model.layers.21.self_attn.to_qkv.weight": "model-00002-of-00003.safetensors",
+    "model.layers.21.self_attn.to_strategy_combine.0.bias": "model-00002-of-00003.safetensors",
+    "model.layers.21.self_attn.to_strategy_combine.0.weight": "model-00002-of-00003.safetensors",
+    "model.layers.21.self_attn.v_compress.1.bias": "model-00002-of-00003.safetensors",
+    "model.layers.21.self_attn.v_compress.1.weight": "model-00002-of-00003.safetensors",
+    "model.layers.21.self_attn.v_compress.3.bias": "model-00002-of-00003.safetensors",
+    "model.layers.21.self_attn.v_compress.3.weight": "model-00002-of-00003.safetensors",
+    "model.layers.21.self_attn.v_intrablock_positions": "model-00002-of-00003.safetensors",
+    "model.layers.22.mlp.down_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.22.mlp.gate_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.22.mlp.up_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.22.post_attn_norm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.22.pre_attn_norm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.22.self_attn.combine_heads.weight": "model-00002-of-00003.safetensors",
+    "model.layers.22.self_attn.compress_mem_kv": "model-00002-of-00003.safetensors",
+    "model.layers.22.self_attn.k_compress.1.bias": "model-00002-of-00003.safetensors",
+    "model.layers.22.self_attn.k_compress.1.weight": "model-00002-of-00003.safetensors",
+    "model.layers.22.self_attn.k_compress.3.bias": "model-00002-of-00003.safetensors",
+    "model.layers.22.self_attn.k_compress.3.weight": "model-00002-of-00003.safetensors",
+    "model.layers.22.self_attn.k_intrablock_positions": "model-00002-of-00003.safetensors",
+    "model.layers.22.self_attn.norm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.22.self_attn.rotary_emb.freqs": "model-00002-of-00003.safetensors",
+    "model.layers.22.self_attn.to_qkv.bias": "model-00002-of-00003.safetensors",
+    "model.layers.22.self_attn.to_qkv.weight": "model-00002-of-00003.safetensors",
+    "model.layers.22.self_attn.to_strategy_combine.0.bias": "model-00002-of-00003.safetensors",
+    "model.layers.22.self_attn.to_strategy_combine.0.weight": "model-00002-of-00003.safetensors",
+    "model.layers.22.self_attn.v_compress.1.bias": "model-00002-of-00003.safetensors",
+    "model.layers.22.self_attn.v_compress.1.weight": "model-00002-of-00003.safetensors",
+    "model.layers.22.self_attn.v_compress.3.bias": "model-00002-of-00003.safetensors",
+    "model.layers.22.self_attn.v_compress.3.weight": "model-00002-of-00003.safetensors",
+    "model.layers.22.self_attn.v_intrablock_positions": "model-00002-of-00003.safetensors",
+    "model.layers.23.mlp.down_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.23.mlp.gate_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.23.mlp.up_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.23.post_attn_norm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.23.pre_attn_norm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.23.self_attn.combine_heads.weight": "model-00002-of-00003.safetensors",
+    "model.layers.23.self_attn.compress_mem_kv": "model-00002-of-00003.safetensors",
+    "model.layers.23.self_attn.k_compress.1.bias": "model-00002-of-00003.safetensors",
+    "model.layers.23.self_attn.k_compress.1.weight": "model-00002-of-00003.safetensors",
+    "model.layers.23.self_attn.k_compress.3.bias": "model-00002-of-00003.safetensors",
+    "model.layers.23.self_attn.k_compress.3.weight": "model-00002-of-00003.safetensors",
+    "model.layers.23.self_attn.k_intrablock_positions": "model-00002-of-00003.safetensors",
+    "model.layers.23.self_attn.norm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.23.self_attn.rotary_emb.freqs": "model-00002-of-00003.safetensors",
+    "model.layers.23.self_attn.to_qkv.bias": "model-00002-of-00003.safetensors",
+    "model.layers.23.self_attn.to_qkv.weight": "model-00002-of-00003.safetensors",
+    "model.layers.23.self_attn.to_strategy_combine.0.bias": "model-00002-of-00003.safetensors",
+    "model.layers.23.self_attn.to_strategy_combine.0.weight": "model-00002-of-00003.safetensors",
+    "model.layers.23.self_attn.v_compress.1.bias": "model-00002-of-00003.safetensors",
+    "model.layers.23.self_attn.v_compress.1.weight": "model-00002-of-00003.safetensors",
+    "model.layers.23.self_attn.v_compress.3.bias": "model-00002-of-00003.safetensors",
+    "model.layers.23.self_attn.v_compress.3.weight": "model-00002-of-00003.safetensors",
+    "model.layers.23.self_attn.v_intrablock_positions": "model-00002-of-00003.safetensors",
+    "model.layers.24.mlp.down_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.24.mlp.gate_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.24.mlp.up_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.24.post_attn_norm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.24.pre_attn_norm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.24.self_attn.combine_heads.weight": "model-00002-of-00003.safetensors",
+    "model.layers.24.self_attn.compress_mem_kv": "model-00002-of-00003.safetensors",
+    "model.layers.24.self_attn.k_compress.1.bias": "model-00002-of-00003.safetensors",
+    "model.layers.24.self_attn.k_compress.1.weight": "model-00002-of-00003.safetensors",
+    "model.layers.24.self_attn.k_compress.3.bias": "model-00002-of-00003.safetensors",
+    "model.layers.24.self_attn.k_compress.3.weight": "model-00002-of-00003.safetensors",
+    "model.layers.24.self_attn.k_intrablock_positions": "model-00002-of-00003.safetensors",
+    "model.layers.24.self_attn.norm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.24.self_attn.rotary_emb.freqs": "model-00002-of-00003.safetensors",
+    "model.layers.24.self_attn.to_qkv.bias": "model-00002-of-00003.safetensors",
+    "model.layers.24.self_attn.to_qkv.weight": "model-00002-of-00003.safetensors",
+    "model.layers.24.self_attn.to_strategy_combine.0.bias": "model-00002-of-00003.safetensors",
+    "model.layers.24.self_attn.to_strategy_combine.0.weight": "model-00002-of-00003.safetensors",
+    "model.layers.24.self_attn.v_compress.1.bias": "model-00002-of-00003.safetensors",
+    "model.layers.24.self_attn.v_compress.1.weight": "model-00002-of-00003.safetensors",
+    "model.layers.24.self_attn.v_compress.3.bias": "model-00002-of-00003.safetensors",
+    "model.layers.24.self_attn.v_compress.3.weight": "model-00002-of-00003.safetensors",
+    "model.layers.24.self_attn.v_intrablock_positions": "model-00002-of-00003.safetensors",
+    "model.layers.25.mlp.down_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.25.mlp.gate_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.25.mlp.up_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.25.post_attn_norm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.25.pre_attn_norm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.25.self_attn.combine_heads.weight": "model-00002-of-00003.safetensors",
+    "model.layers.25.self_attn.compress_mem_kv": "model-00002-of-00003.safetensors",
+    "model.layers.25.self_attn.k_compress.1.bias": "model-00002-of-00003.safetensors",
+    "model.layers.25.self_attn.k_compress.1.weight": "model-00002-of-00003.safetensors",
+    "model.layers.25.self_attn.k_compress.3.bias": "model-00002-of-00003.safetensors",
+    "model.layers.25.self_attn.k_compress.3.weight": "model-00002-of-00003.safetensors",
+    "model.layers.25.self_attn.k_intrablock_positions": "model-00002-of-00003.safetensors",
+    "model.layers.25.self_attn.norm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.25.self_attn.rotary_emb.freqs": "model-00002-of-00003.safetensors",
+    "model.layers.25.self_attn.to_qkv.bias": "model-00002-of-00003.safetensors",
+    "model.layers.25.self_attn.to_qkv.weight": "model-00002-of-00003.safetensors",
+    "model.layers.25.self_attn.to_strategy_combine.0.bias": "model-00002-of-00003.safetensors",
+    "model.layers.25.self_attn.to_strategy_combine.0.weight": "model-00002-of-00003.safetensors",
+    "model.layers.25.self_attn.v_compress.1.bias": "model-00002-of-00003.safetensors",
+    "model.layers.25.self_attn.v_compress.1.weight": "model-00002-of-00003.safetensors",
+    "model.layers.25.self_attn.v_compress.3.bias": "model-00002-of-00003.safetensors",
+    "model.layers.25.self_attn.v_compress.3.weight": "model-00002-of-00003.safetensors",
+    "model.layers.25.self_attn.v_intrablock_positions": "model-00002-of-00003.safetensors",
+    "model.layers.26.mlp.down_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.26.mlp.gate_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.26.mlp.up_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.26.post_attn_norm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.26.pre_attn_norm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.26.self_attn.combine_heads.weight": "model-00002-of-00003.safetensors",
+    "model.layers.26.self_attn.compress_mem_kv": "model-00002-of-00003.safetensors",
+    "model.layers.26.self_attn.k_compress.1.bias": "model-00002-of-00003.safetensors",
+    "model.layers.26.self_attn.k_compress.1.weight": "model-00002-of-00003.safetensors",
+    "model.layers.26.self_attn.k_compress.3.bias": "model-00002-of-00003.safetensors",
+    "model.layers.26.self_attn.k_compress.3.weight": "model-00002-of-00003.safetensors",
+    "model.layers.26.self_attn.k_intrablock_positions": "model-00002-of-00003.safetensors",
+    "model.layers.26.self_attn.norm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.26.self_attn.rotary_emb.freqs": "model-00002-of-00003.safetensors",
+    "model.layers.26.self_attn.to_qkv.bias": "model-00002-of-00003.safetensors",
+    "model.layers.26.self_attn.to_qkv.weight": "model-00002-of-00003.safetensors",
+    "model.layers.26.self_attn.to_strategy_combine.0.bias": "model-00002-of-00003.safetensors",
+    "model.layers.26.self_attn.to_strategy_combine.0.weight": "model-00002-of-00003.safetensors",
+    "model.layers.26.self_attn.v_compress.1.bias": "model-00002-of-00003.safetensors",
+    "model.layers.26.self_attn.v_compress.1.weight": "model-00002-of-00003.safetensors",
+    "model.layers.26.self_attn.v_compress.3.bias": "model-00002-of-00003.safetensors",
+    "model.layers.26.self_attn.v_compress.3.weight": "model-00002-of-00003.safetensors",
+    "model.layers.26.self_attn.v_intrablock_positions": "model-00002-of-00003.safetensors",
+    "model.layers.27.mlp.down_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.27.mlp.gate_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.27.mlp.up_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.27.post_attn_norm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.27.pre_attn_norm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.27.self_attn.combine_heads.weight": "model-00002-of-00003.safetensors",
+    "model.layers.27.self_attn.compress_mem_kv": "model-00002-of-00003.safetensors",
+    "model.layers.27.self_attn.k_compress.1.bias": "model-00002-of-00003.safetensors",
+    "model.layers.27.self_attn.k_compress.1.weight": "model-00002-of-00003.safetensors",
+    "model.layers.27.self_attn.k_compress.3.bias": "model-00002-of-00003.safetensors",
+    "model.layers.27.self_attn.k_compress.3.weight": "model-00002-of-00003.safetensors",
+    "model.layers.27.self_attn.k_intrablock_positions": "model-00002-of-00003.safetensors",
+    "model.layers.27.self_attn.norm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.27.self_attn.rotary_emb.freqs": "model-00002-of-00003.safetensors",
+    "model.layers.27.self_attn.to_qkv.bias": "model-00002-of-00003.safetensors",
+    "model.layers.27.self_attn.to_qkv.weight": "model-00002-of-00003.safetensors",
+    "model.layers.27.self_attn.to_strategy_combine.0.bias": "model-00002-of-00003.safetensors",
+    "model.layers.27.self_attn.to_strategy_combine.0.weight": "model-00002-of-00003.safetensors",
+    "model.layers.27.self_attn.v_compress.1.bias": "model-00002-of-00003.safetensors",
+    "model.layers.27.self_attn.v_compress.1.weight": "model-00002-of-00003.safetensors",
+    "model.layers.27.self_attn.v_compress.3.bias": "model-00002-of-00003.safetensors",
+    "model.layers.27.self_attn.v_compress.3.weight": "model-00002-of-00003.safetensors",
+    "model.layers.27.self_attn.v_intrablock_positions": "model-00002-of-00003.safetensors",
+    "model.layers.3.mlp.down_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.3.mlp.gate_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.3.mlp.up_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.3.post_attn_norm.weight": "model-00001-of-00003.safetensors",
+    "model.layers.3.pre_attn_norm.weight": "model-00001-of-00003.safetensors",
+    "model.layers.3.self_attn.combine_heads.weight": "model-00001-of-00003.safetensors",
+    "model.layers.3.self_attn.compress_mem_kv": "model-00001-of-00003.safetensors",
+    "model.layers.3.self_attn.k_compress.1.bias": "model-00001-of-00003.safetensors",
+    "model.layers.3.self_attn.k_compress.1.weight": "model-00001-of-00003.safetensors",
+    "model.layers.3.self_attn.k_compress.3.bias": "model-00001-of-00003.safetensors",
+    "model.layers.3.self_attn.k_compress.3.weight": "model-00001-of-00003.safetensors",
+    "model.layers.3.self_attn.k_intrablock_positions": "model-00001-of-00003.safetensors",
+    "model.layers.3.self_attn.norm.weight": "model-00001-of-00003.safetensors",
+    "model.layers.3.self_attn.rotary_emb.freqs": "model-00001-of-00003.safetensors",
+    "model.layers.3.self_attn.to_qkv.bias": "model-00001-of-00003.safetensors",
+    "model.layers.3.self_attn.to_qkv.weight": "model-00001-of-00003.safetensors",
+    "model.layers.3.self_attn.to_strategy_combine.0.bias": "model-00001-of-00003.safetensors",
+    "model.layers.3.self_attn.to_strategy_combine.0.weight": "model-00001-of-00003.safetensors",
+    "model.layers.3.self_attn.v_compress.1.bias": "model-00001-of-00003.safetensors",
+    "model.layers.3.self_attn.v_compress.1.weight": "model-00001-of-00003.safetensors",
+    "model.layers.3.self_attn.v_compress.3.bias": "model-00001-of-00003.safetensors",
+    "model.layers.3.self_attn.v_compress.3.weight": "model-00001-of-00003.safetensors",
+    "model.layers.3.self_attn.v_intrablock_positions": "model-00001-of-00003.safetensors",
+    "model.layers.4.mlp.down_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.4.mlp.gate_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.4.mlp.up_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.4.post_attn_norm.weight": "model-00001-of-00003.safetensors",
+    "model.layers.4.pre_attn_norm.weight": "model-00001-of-00003.safetensors",
+    "model.layers.4.self_attn.combine_heads.weight": "model-00001-of-00003.safetensors",
+    "model.layers.4.self_attn.compress_mem_kv": "model-00001-of-00003.safetensors",
+    "model.layers.4.self_attn.k_compress.1.bias": "model-00001-of-00003.safetensors",
+    "model.layers.4.self_attn.k_compress.1.weight": "model-00001-of-00003.safetensors",
+    "model.layers.4.self_attn.k_compress.3.bias": "model-00001-of-00003.safetensors",
+    "model.layers.4.self_attn.k_compress.3.weight": "model-00001-of-00003.safetensors",
+    "model.layers.4.self_attn.k_intrablock_positions": "model-00001-of-00003.safetensors",
+    "model.layers.4.self_attn.norm.weight": "model-00001-of-00003.safetensors",
+    "model.layers.4.self_attn.rotary_emb.freqs": "model-00001-of-00003.safetensors",
+    "model.layers.4.self_attn.to_qkv.bias": "model-00001-of-00003.safetensors",
+    "model.layers.4.self_attn.to_qkv.weight": "model-00001-of-00003.safetensors",
+    "model.layers.4.self_attn.to_strategy_combine.0.bias": "model-00001-of-00003.safetensors",
+    "model.layers.4.self_attn.to_strategy_combine.0.weight": "model-00001-of-00003.safetensors",
+    "model.layers.4.self_attn.v_compress.1.bias": "model-00001-of-00003.safetensors",
+    "model.layers.4.self_attn.v_compress.1.weight": "model-00001-of-00003.safetensors",
+    "model.layers.4.self_attn.v_compress.3.bias": "model-00001-of-00003.safetensors",
+    "model.layers.4.self_attn.v_compress.3.weight": "model-00001-of-00003.safetensors",
+    "model.layers.4.self_attn.v_intrablock_positions": "model-00001-of-00003.safetensors",
+    "model.layers.5.mlp.down_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.5.mlp.gate_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.5.mlp.up_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.5.post_attn_norm.weight": "model-00001-of-00003.safetensors",
+    "model.layers.5.pre_attn_norm.weight": "model-00001-of-00003.safetensors",
+    "model.layers.5.self_attn.combine_heads.weight": "model-00001-of-00003.safetensors",
+    "model.layers.5.self_attn.compress_mem_kv": "model-00001-of-00003.safetensors",
+    "model.layers.5.self_attn.k_compress.1.bias": "model-00001-of-00003.safetensors",
+    "model.layers.5.self_attn.k_compress.1.weight": "model-00001-of-00003.safetensors",
+    "model.layers.5.self_attn.k_compress.3.bias": "model-00001-of-00003.safetensors",
+    "model.layers.5.self_attn.k_compress.3.weight": "model-00001-of-00003.safetensors",
+    "model.layers.5.self_attn.k_intrablock_positions": "model-00001-of-00003.safetensors",
+    "model.layers.5.self_attn.norm.weight": "model-00001-of-00003.safetensors",
+    "model.layers.5.self_attn.rotary_emb.freqs": "model-00001-of-00003.safetensors",
+    "model.layers.5.self_attn.to_qkv.bias": "model-00001-of-00003.safetensors",
+    "model.layers.5.self_attn.to_qkv.weight": "model-00001-of-00003.safetensors",
+    "model.layers.5.self_attn.to_strategy_combine.0.bias": "model-00001-of-00003.safetensors",
+    "model.layers.5.self_attn.to_strategy_combine.0.weight": "model-00001-of-00003.safetensors",
+    "model.layers.5.self_attn.v_compress.1.bias": "model-00001-of-00003.safetensors",
+    "model.layers.5.self_attn.v_compress.1.weight": "model-00001-of-00003.safetensors",
+    "model.layers.5.self_attn.v_compress.3.bias": "model-00001-of-00003.safetensors",
+    "model.layers.5.self_attn.v_compress.3.weight": "model-00001-of-00003.safetensors",
+    "model.layers.5.self_attn.v_intrablock_positions": "model-00001-of-00003.safetensors",
+    "model.layers.6.mlp.down_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.6.mlp.gate_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.6.mlp.up_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.6.post_attn_norm.weight": "model-00001-of-00003.safetensors",
+    "model.layers.6.pre_attn_norm.weight": "model-00001-of-00003.safetensors",
+    "model.layers.6.self_attn.combine_heads.weight": "model-00001-of-00003.safetensors",
+    "model.layers.6.self_attn.compress_mem_kv": "model-00001-of-00003.safetensors",
+    "model.layers.6.self_attn.k_compress.1.bias": "model-00001-of-00003.safetensors",
+    "model.layers.6.self_attn.k_compress.1.weight": "model-00001-of-00003.safetensors",
+    "model.layers.6.self_attn.k_compress.3.bias": "model-00001-of-00003.safetensors",
+    "model.layers.6.self_attn.k_compress.3.weight": "model-00001-of-00003.safetensors",
+    "model.layers.6.self_attn.k_intrablock_positions": "model-00001-of-00003.safetensors",
+    "model.layers.6.self_attn.norm.weight": "model-00001-of-00003.safetensors",
+    "model.layers.6.self_attn.rotary_emb.freqs": "model-00001-of-00003.safetensors",
+    "model.layers.6.self_attn.to_qkv.bias": "model-00001-of-00003.safetensors",
+    "model.layers.6.self_attn.to_qkv.weight": "model-00001-of-00003.safetensors",
+    "model.layers.6.self_attn.to_strategy_combine.0.bias": "model-00001-of-00003.safetensors",
+    "model.layers.6.self_attn.to_strategy_combine.0.weight": "model-00001-of-00003.safetensors",
+    "model.layers.6.self_attn.v_compress.1.bias": "model-00001-of-00003.safetensors",
+    "model.layers.6.self_attn.v_compress.1.weight": "model-00001-of-00003.safetensors",
+    "model.layers.6.self_attn.v_compress.3.bias": "model-00001-of-00003.safetensors",
+    "model.layers.6.self_attn.v_compress.3.weight": "model-00001-of-00003.safetensors",
+    "model.layers.6.self_attn.v_intrablock_positions": "model-00001-of-00003.safetensors",
+    "model.layers.7.mlp.down_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.7.mlp.gate_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.7.mlp.up_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.7.post_attn_norm.weight": "model-00001-of-00003.safetensors",
+    "model.layers.7.pre_attn_norm.weight": "model-00001-of-00003.safetensors",
+    "model.layers.7.self_attn.combine_heads.weight": "model-00001-of-00003.safetensors",
+    "model.layers.7.self_attn.compress_mem_kv": "model-00001-of-00003.safetensors",
+    "model.layers.7.self_attn.k_compress.1.bias": "model-00001-of-00003.safetensors",
+    "model.layers.7.self_attn.k_compress.1.weight": "model-00001-of-00003.safetensors",
+    "model.layers.7.self_attn.k_compress.3.bias": "model-00001-of-00003.safetensors",
+    "model.layers.7.self_attn.k_compress.3.weight": "model-00001-of-00003.safetensors",
+    "model.layers.7.self_attn.k_intrablock_positions": "model-00001-of-00003.safetensors",
+    "model.layers.7.self_attn.norm.weight": "model-00001-of-00003.safetensors",
+    "model.layers.7.self_attn.rotary_emb.freqs": "model-00001-of-00003.safetensors",
+    "model.layers.7.self_attn.to_qkv.bias": "model-00001-of-00003.safetensors",
+    "model.layers.7.self_attn.to_qkv.weight": "model-00001-of-00003.safetensors",
+    "model.layers.7.self_attn.to_strategy_combine.0.bias": "model-00001-of-00003.safetensors",
+    "model.layers.7.self_attn.to_strategy_combine.0.weight": "model-00001-of-00003.safetensors",
+    "model.layers.7.self_attn.v_compress.1.bias": "model-00001-of-00003.safetensors",
+    "model.layers.7.self_attn.v_compress.1.weight": "model-00001-of-00003.safetensors",
+    "model.layers.7.self_attn.v_compress.3.bias": "model-00001-of-00003.safetensors",
+    "model.layers.7.self_attn.v_compress.3.weight": "model-00001-of-00003.safetensors",
+    "model.layers.7.self_attn.v_intrablock_positions": "model-00001-of-00003.safetensors",
+    "model.layers.8.mlp.down_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.8.mlp.gate_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.8.mlp.up_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.8.post_attn_norm.weight": "model-00001-of-00003.safetensors",
+    "model.layers.8.pre_attn_norm.weight": "model-00001-of-00003.safetensors",
+    "model.layers.8.self_attn.combine_heads.weight": "model-00001-of-00003.safetensors",
+    "model.layers.8.self_attn.compress_mem_kv": "model-00001-of-00003.safetensors",
+    "model.layers.8.self_attn.k_compress.1.bias": "model-00001-of-00003.safetensors",
+    "model.layers.8.self_attn.k_compress.1.weight": "model-00001-of-00003.safetensors",
+    "model.layers.8.self_attn.k_compress.3.bias": "model-00001-of-00003.safetensors",
+    "model.layers.8.self_attn.k_compress.3.weight": "model-00001-of-00003.safetensors",
+    "model.layers.8.self_attn.k_intrablock_positions": "model-00001-of-00003.safetensors",
+    "model.layers.8.self_attn.norm.weight": "model-00001-of-00003.safetensors",
+    "model.layers.8.self_attn.rotary_emb.freqs": "model-00001-of-00003.safetensors",
+    "model.layers.8.self_attn.to_qkv.bias": "model-00001-of-00003.safetensors",
+    "model.layers.8.self_attn.to_qkv.weight": "model-00001-of-00003.safetensors",
+    "model.layers.8.self_attn.to_strategy_combine.0.bias": "model-00001-of-00003.safetensors",
+    "model.layers.8.self_attn.to_strategy_combine.0.weight": "model-00001-of-00003.safetensors",
+    "model.layers.8.self_attn.v_compress.1.bias": "model-00001-of-00003.safetensors",
+    "model.layers.8.self_attn.v_compress.1.weight": "model-00001-of-00003.safetensors",
+    "model.layers.8.self_attn.v_compress.3.bias": "model-00001-of-00003.safetensors",
+    "model.layers.8.self_attn.v_compress.3.weight": "model-00001-of-00003.safetensors",
+    "model.layers.8.self_attn.v_intrablock_positions": "model-00001-of-00003.safetensors",
+    "model.layers.9.mlp.down_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.9.mlp.gate_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.9.mlp.up_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.9.post_attn_norm.weight": "model-00001-of-00003.safetensors",
+    "model.layers.9.pre_attn_norm.weight": "model-00001-of-00003.safetensors",
+    "model.layers.9.self_attn.combine_heads.weight": "model-00001-of-00003.safetensors",
+    "model.layers.9.self_attn.compress_mem_kv": "model-00001-of-00003.safetensors",
+    "model.layers.9.self_attn.k_compress.1.bias": "model-00001-of-00003.safetensors",
+    "model.layers.9.self_attn.k_compress.1.weight": "model-00001-of-00003.safetensors",
+    "model.layers.9.self_attn.k_compress.3.bias": "model-00001-of-00003.safetensors",
+    "model.layers.9.self_attn.k_compress.3.weight": "model-00001-of-00003.safetensors",
+    "model.layers.9.self_attn.k_intrablock_positions": "model-00001-of-00003.safetensors",
+    "model.layers.9.self_attn.norm.weight": "model-00001-of-00003.safetensors",
+    "model.layers.9.self_attn.rotary_emb.freqs": "model-00001-of-00003.safetensors",
+    "model.layers.9.self_attn.to_qkv.bias": "model-00001-of-00003.safetensors",
+    "model.layers.9.self_attn.to_qkv.weight": "model-00001-of-00003.safetensors",
+    "model.layers.9.self_attn.to_strategy_combine.0.bias": "model-00001-of-00003.safetensors",
+    "model.layers.9.self_attn.to_strategy_combine.0.weight": "model-00001-of-00003.safetensors",
+    "model.layers.9.self_attn.v_compress.1.bias": "model-00001-of-00003.safetensors",
+    "model.layers.9.self_attn.v_compress.1.weight": "model-00001-of-00003.safetensors",
+    "model.layers.9.self_attn.v_compress.3.bias": "model-00001-of-00003.safetensors",
+    "model.layers.9.self_attn.v_compress.3.weight": "model-00001-of-00003.safetensors",
+    "model.layers.9.self_attn.v_intrablock_positions": "model-00001-of-00003.safetensors",
+    "model.norm.weight": "model-00003-of-00003.safetensors"
+  }
+}

modeling_qwen2nsa.py

@@ -0,0 +1,346 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+from typing import Optional, Tuple, Dict, Any, List
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch import Tensor
+from torch.utils.checkpoint import checkpoint
+
+from transformers import (
+    PreTrainedModel,
+    GenerationMixin,
+    AutoConfig,
+    AutoModelForCausalLM,
+)
+from transformers.modeling_outputs import CausalLMOutputWithPast
+
+try:
+    from .configuration_qwen2nsa import Qwen2NsaConfig
+except ImportError:
+    from configuration_qwen2nsa import Qwen2NsaConfig
+
+from native_sparse_attention_pytorch import SparseAttention
+
+class Qwen2RMSNorm(nn.Module):
+    def __init__(self, hidden_size: int, eps: float = 1e-6):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+        self.eps = eps
+
+    def forward(self, x: Tensor) -> Tensor:
+        orig_dtype = x.dtype
+        x = x.float()
+        var = x.pow(2).mean(-1, keepdim=True)
+        x = x * torch.rsqrt(var + self.eps)
+        return (self.weight * x).to(orig_dtype)
+
+
+class Qwen2MLP(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.gate_proj = nn.Linear(
+            config.hidden_size, config.intermediate_size, bias=False
+        )
+        self.up_proj = nn.Linear(
+            config.hidden_size, config.intermediate_size, bias=False
+        )
+        self.down_proj = nn.Linear(
+            config.intermediate_size, config.hidden_size, bias=False
+        )
+        self.act = F.silu
+
+    def forward(self, x):
+        return self.down_proj(self.act(self.gate_proj(x)) * self.up_proj(x))
+
+
+class Qwen2DecoderLayer(nn.Module):
+    def __init__(self, config: Qwen2NsaConfig, layer_idx: int):
+        super().__init__()
+        self.self_attn = SparseAttention(
+            dim=config.hidden_size,
+            dim_head=config.hidden_size // config.num_attention_heads,
+            heads=config.num_attention_heads,
+            kv_heads=config.num_key_value_heads,
+            sliding_window_size=1024,
+            compress_block_size=32,
+            compress_block_sliding_stride=16,
+            selection_block_size=64,
+            num_selected_blocks=16,
+            causal=True,
+            norm=True,
+        )
+        self.mlp = Qwen2MLP(config)
+        self.pre_attn_norm = Qwen2RMSNorm(config.hidden_size, config.rms_norm_eps)
+        self.post_attn_norm = Qwen2RMSNorm(config.hidden_size, config.rms_norm_eps)
+
+    def forward(
+        self,
+        hidden_states: Tensor,
+        attn_mask: Optional[Tensor] = None,  # (unused)
+        position_embeddings: Optional[Tuple[Tensor, Tensor]] = None,
+        *,
+        output_attentions: bool = False,
+        output_all_states: bool = False,
+        return_cache: bool = False,
+    ):
+        residual = hidden_states
+        x = self.pre_attn_norm(hidden_states)
+
+        if return_cache:
+            attn_out, cache_obj, att_weights = self.self_attn(x, return_cache=True)
+        else:
+            attn_out, att_weights = checkpoint(
+                lambda y: self.self_attn(y), x, use_reentrant=False
+            )
+            cache_obj = None
+
+        post_attn = residual + attn_out
+        mlp_out = self.mlp(self.post_attn_norm(post_attn))
+        hidden = post_attn + mlp_out
+
+        extras = None
+        if output_all_states:
+            extras = dict(
+                pre_attn_states=x,
+                attn_output=attn_out,
+                post_attn_states=post_attn,
+                mlp_output=mlp_out,
+            )
+
+        if not output_attentions:
+            att_weights = None
+
+        return hidden, att_weights, extras, cache_obj
+
+    def forward_inference(
+        self,
+        hidden_states: Tensor,
+        cache: Any,
+        *,
+        output_attentions: bool = False,
+        output_all_states: bool = False,
+    ):
+        residual = hidden_states
+        x = self.pre_attn_norm(hidden_states)
+
+        attn_out, new_cache, att_weights = self.self_attn.forward_inference(
+            x, cache, return_cache=True
+        )
+
+        post_attn = residual + attn_out
+        mlp_out = self.mlp(self.post_attn_norm(post_attn))
+        hidden = post_attn + mlp_out
+
+        extras = None
+        if output_all_states:
+            extras = dict(
+                pre_attn_states=x,
+                attn_output=attn_out,
+                post_attn_states=post_attn,
+                mlp_output=mlp_out,
+            )
+
+        if not output_attentions:
+            att_weights = None
+
+        return hidden, new_cache, att_weights, extras
+
+
+class Qwen2Model(nn.Module):
+    def __init__(self, config: Qwen2NsaConfig):
+        super().__init__()
+        self.config = config
+        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
+        self.layers = nn.ModuleList(
+            Qwen2DecoderLayer(config, i) for i in range(config.num_hidden_layers)
+        )
+        self.norm = Qwen2RMSNorm(config.hidden_size, config.rms_norm_eps)
+
+        self.apply(self._init_weights)
+
+    def _init_weights(self, module: nn.Module):
+        if isinstance(module, (nn.Linear, nn.Embedding)):
+            module.weight.data.normal_(0.0, self.config.initializer_range)
+            if getattr(module, "bias", None) is not None:
+                module.bias.data.zero_()
+
+    # ------------------------------------------------------------------ #
+    def forward(
+        self,
+        input_ids: torch.LongTensor,
+        *,
+        attention_mask: Optional[Tensor] = None,
+        position_ids: Optional[Tensor] = None,
+        past_key_values: Optional[List[Any]] = None,
+        use_cache: bool = False,
+        output_hidden_states: bool = False,
+        output_attentions: bool = False,
+        output_all_layer_states: bool = False,
+    ) -> Dict[str, Any]:
+
+        hidden = self.embed_tokens(input_ids)
+
+        if past_key_values is None:
+            past_key_values = [None] * len(self.layers)
+
+        next_cache: List[Any] = [] if use_cache else None
+        all_hs = [] if output_hidden_states else None
+        all_att = [] if output_attentions else None
+        all_ex = [] if output_all_layer_states else None
+
+        for layer, past in zip(self.layers, past_key_values):
+            if output_hidden_states:
+                all_hs.append(hidden)
+
+            if use_cache and past is not None:
+                hidden, new_cache, att, ex = layer.forward_inference(
+                    hidden,
+                    past,
+                    output_attentions=output_attentions,
+                    output_all_states=output_all_layer_states,
+                )
+            else:
+                hidden, att, ex, new_cache = layer(
+                    hidden,
+                    output_attentions=output_attentions,
+                    output_all_states=output_all_layer_states,
+                    return_cache=use_cache,
+                )
+
+            if use_cache:
+                next_cache.append(new_cache)
+            if output_attentions:
+                all_att.append(att)
+            if output_all_layer_states:
+                all_ex.append(ex)
+
+        hidden = self.norm(hidden)
+        if output_hidden_states:
+            all_hs.append(hidden)
+
+        out = dict(last_hidden_state=hidden)
+        if output_hidden_states:
+            out["hidden_states"] = tuple(all_hs)
+        if output_attentions:
+            out["attentions"] = tuple(all_att)
+        if output_all_layer_states:
+            out["all_layer_states"] = all_ex
+        if use_cache:
+            out["past_key_values"] = next_cache
+        return out
+
+class Qwen2NsaForCausalLM(PreTrainedModel, GenerationMixin):
+    config_class = Qwen2NsaConfig
+    base_model_prefix = "model"
+    _tied_weights_keys = ["lm_head.weight"]
+
+    def __init__(self, config: Qwen2NsaConfig):
+        super().__init__(config)
+
+        self.model = Qwen2Model(config)
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+
+        if config.tie_word_embeddings:
+            self.lm_head.weight = self.model.embed_tokens.weight
+
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.model.embed_tokens
+
+    def set_input_embeddings(self, v):
+        self.model.embed_tokens = v
+
+    def get_output_embeddings(self):
+        return self.lm_head
+
+    def set_output_embeddings(self, v):
+        self.lm_head = v
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.Tensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[List[Any]] = None,
+        use_cache: Optional[bool] = None,
+        output_hidden_states: bool = False,
+        output_attentions: bool = False,
+        output_all_layer_states: bool = False,
+        **kwargs,
+    ) -> CausalLMOutputWithPast:
+
+        use_cache = self.config.use_cache if use_cache is None else use_cache
+
+        model_out = self.model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            use_cache=use_cache,
+            output_hidden_states=output_hidden_states,
+            output_attentions=output_attentions,
+            output_all_layer_states=output_all_layer_states,
+        )
+
+        logits = self.lm_head(model_out["last_hidden_state"])
+
+        loss = None
+        if labels is not None:
+            pad = self.config.pad_token_id
+            if pad is not None:
+                labels = labels.masked_fill(labels == pad, -100)
+            loss = F.cross_entropy(
+                logits[..., :-1, :].contiguous().view(-1, logits.size(-1)),
+                labels[..., 1:].contiguous().view(-1),
+            )
+
+        return CausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            past_key_values=model_out.get("past_key_values"),
+            hidden_states=model_out.get("hidden_states"),
+            attentions=model_out.get("attentions"),
+        )
+
+    def _reorder_cache(self, past_key_values: list, beam_idx: torch.LongTensor) -> list:
+        new = []
+        for layer_cache in past_key_values:
+            (kv, (ckv, rkv)) = layer_cache
+            k, v = kv
+            ck, cv = ckv
+            rk, rv = rkv
+            pick = lambda t: t.index_select(0, beam_idx) if t.numel() else t
+            new.append(
+                ((pick(k), pick(v)), ((pick(ck), pick(cv)), (pick(rk), pick(rv))))
+            )
+        return new
+
+    def prepare_inputs_for_generation(
+        self,
+        input_ids: torch.LongTensor,
+        past_key_values: Optional[List[Any]] = None,
+        attention_mask: Optional[torch.LongTensor] = None,
+        **kwargs,
+    ):
+        if past_key_values is not None:
+            input_ids = input_ids[:, -1:]
+            if attention_mask is not None:
+                attention_mask = attention_mask[:, -1:]
+        return dict(
+            input_ids=input_ids,
+            past_key_values=past_key_values,
+            attention_mask=attention_mask,
+            use_cache=True,
+        )
+
+
+AutoConfig.register(Qwen2NsaConfig.model_type, Qwen2NsaConfig)
+AutoModelForCausalLM.register(Qwen2NsaConfig, Qwen2NsaForCausalLM)
+
+__all__ = ["Qwen2NsaConfig", "Qwen2NsaForCausalLM"]