Add files using upload-large-folder tool

.gitattributes

@@ -33,3 +33,4 @@ 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
+tokenizer.json filter=lfs diff=lfs merge=lfs -text

LICENSE

@@ -0,0 +1,202 @@
+
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright 2026 Alibaba Cloud
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

Qwen3_5_wrapper.py

@@ -0,0 +1,305 @@
+from .utils import DecoderLayerWrapperForMetis, CausalLMWrapperForMetis
+from transformers.models.qwen3_5.modeling_qwen3_5 import apply_rotary_pos_emb, eager_attention_forward, Qwen3_5TextModel
+import torch
+from torch.utils.checkpoint import checkpoint as _ckpt
+from transformers import Cache
+from transformers.utils import TransformersKwargs
+from typing_extensions import Unpack
+from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
+from transformers.masking_utils import create_causal_mask
+from typing import Callable
+import torch.nn as nn
+from transformers.models.qwen3_5.modeling_qwen3_5 import Qwen3_5ModelOutputWithPast
+
+# Qwen3_5DynamicCache was introduced in a later transformers version.
+# Fall back to DynamicCache when unavailable (training uses use_cache=False,
+# so the custom cache class is only needed for autoregressive generation).
+try:
+    from transformers.models.qwen3_5.modeling_qwen3_5 import Qwen3_5DynamicCache
+except ImportError:
+    from transformers import DynamicCache as Qwen3_5DynamicCache 
+from transformers.utils.generic import merge_with_config_defaults
+from transformers.utils.output_capturing import capture_outputs
+from transformers.modeling_outputs import BaseModelOutputWithPast
+
+
+class Qwen3_5DecoderLayerForMetis(DecoderLayerWrapperForMetis):
+    def __init__(self, config, raw_decoder):
+        super().__init__(config, raw_decoder)
+
+    # Modified from Qwen3.5
+    def before_mixin(self,
+        hidden_states: torch.Tensor,
+        position_embeddings: tuple[torch.Tensor, torch.Tensor],
+        attention_mask: torch.Tensor | None = None,
+        position_ids: torch.LongTensor | None = None,
+        past_key_values: Cache | None = None,
+        cache_position: torch.LongTensor | None = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ):
+        if self.raw_decoder.layer_type == "linear_attention":
+            return None, None, None, None
+
+        self_attn = self.raw_decoder.self_attn
+        use_ckpt = getattr(self, "_use_gradient_checkpointing", False) and self.training
+
+        if not use_ckpt:
+            residual = hidden_states
+            hidden_states = self.raw_decoder.input_layernorm(hidden_states)
+
+            input_shape = hidden_states.shape[:-1]
+            hidden_shape = (*input_shape, -1, self_attn.head_dim)
+
+            query_states, gate = torch.chunk(
+                self_attn.q_proj(hidden_states).view(*input_shape, -1, self_attn.head_dim * 2), 2, dim=-1
+            )
+            gate = gate.reshape(*input_shape, -1)
+
+            query_states = self_attn.q_norm(query_states.view(hidden_shape)).transpose(1, 2)
+            key_states = self_attn.k_norm(self_attn.k_proj(hidden_states).view(hidden_shape)).transpose(1, 2)
+            value_states = self_attn.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+
+            memory_for_query = query_states.clone()
+
+            cos, sin = position_embeddings
+            query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
+
+            if past_key_values is not None:
+                cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}
+                key_states, value_states = past_key_values.update(
+                    key_states, value_states, self_attn.layer_idx, cache_kwargs
+                )
+
+            attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface(
+                self_attn.config._attn_implementation, eager_attention_forward
+            )
+
+            attn_output, attn_weights = attention_interface(
+                self_attn,
+                query_states,
+                key_states,
+                value_states,
+                attention_mask,
+                dropout=0.0 if not self.training else self_attn.attention_dropout,
+                scaling=self_attn.scaling,
+                **kwargs,
+            )
+
+            attn_output = attn_output.reshape(*input_shape, -1).contiguous()
+            attn_output = attn_output * torch.sigmoid(gate)
+            attn_output = self_attn.o_proj(attn_output)
+            return memory_for_query, attn_output, {'residual': residual}, self_attn.o_proj
+
+        def _attn(hidden_states):
+            residual = hidden_states
+            normed = self.raw_decoder.input_layernorm(hidden_states)
+
+            input_shape = normed.shape[:-1]
+            hidden_shape = (*input_shape, -1, self_attn.head_dim)
+
+            query_states, gate = torch.chunk(
+                self_attn.q_proj(normed).view(*input_shape, -1, self_attn.head_dim * 2), 2, dim=-1
+            )
+            gate = gate.reshape(*input_shape, -1)
+
+            query_states = self_attn.q_norm(query_states.view(hidden_shape)).transpose(1, 2)
+            key_states = self_attn.k_norm(self_attn.k_proj(normed).view(hidden_shape)).transpose(1, 2)
+            value_states = self_attn.v_proj(normed).view(hidden_shape).transpose(1, 2)
+
+            memory_for_query = query_states.clone()
+
+            cos, sin = position_embeddings
+            query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
+
+            if past_key_values is not None:
+                cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}
+                key_states, value_states = past_key_values.update(
+                    key_states, value_states, self_attn.layer_idx, cache_kwargs
+                )
+
+            attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface(
+                self_attn.config._attn_implementation, eager_attention_forward
+            )
+
+            attn_output, attn_weights = attention_interface(
+                self_attn,
+                query_states,
+                key_states,
+                value_states,
+                attention_mask,
+                dropout=0.0 if not self.training else self_attn.attention_dropout,
+                scaling=self_attn.scaling,
+                **kwargs,
+            )
+
+            attn_output = attn_output.reshape(*input_shape, -1).contiguous()
+            attn_output = attn_output * torch.sigmoid(gate)
+            attn_output = self_attn.o_proj(attn_output)
+            return memory_for_query, attn_output, residual
+
+        memory_for_query, attn_output, residual = _ckpt(_attn, hidden_states, use_reentrant=False)
+
+        return memory_for_query, attn_output, {'residual': residual}, self_attn.o_proj
+    
+    def after_mixin(self, memory_carrier,
+        cache_dict,
+        hidden_states: torch.Tensor,
+        position_embeddings: tuple[torch.Tensor, torch.Tensor],
+        attention_mask: torch.Tensor | None = None,
+        position_ids: torch.LongTensor | None = None,
+        past_key_values: Cache | None = None,
+        cache_position: torch.LongTensor | None = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> torch.Tensor:
+        use_ckpt = getattr(self, "_use_gradient_checkpointing", False) and self.training
+
+        if not use_ckpt:
+            if self.raw_decoder.layer_type == "linear_attention":
+                residual = hidden_states
+                hidden_states = self.raw_decoder.input_layernorm(hidden_states)
+                hidden_states = self.raw_decoder.linear_attn(
+                    hidden_states=hidden_states,
+                    cache_params=past_key_values,
+                    attention_mask=attention_mask,
+                )
+            elif self.raw_decoder.layer_type == "full_attention":
+                residual = cache_dict['residual']
+                hidden_states = memory_carrier
+
+            hidden_states = residual + hidden_states
+            residual = hidden_states
+            hidden_states = self.raw_decoder.post_attention_layernorm(hidden_states)
+            hidden_states = self.raw_decoder.mlp(hidden_states)
+            return residual + hidden_states
+
+        if self.raw_decoder.layer_type == "linear_attention":
+            def _mix(hidden_states):
+                residual = hidden_states
+                normed = self.raw_decoder.input_layernorm(hidden_states)
+                mixed = self.raw_decoder.linear_attn(
+                    hidden_states=normed,
+                    cache_params=past_key_values,
+                    attention_mask=attention_mask,
+                )
+                return residual + mixed
+
+            hidden_states = _ckpt(_mix, hidden_states, use_reentrant=False)
+
+        elif self.raw_decoder.layer_type == "full_attention":
+            residual = cache_dict['residual']
+            hidden_states = residual + memory_carrier
+
+        def _mlp(hidden_states):
+            residual = hidden_states
+            normed = self.raw_decoder.post_attention_layernorm(hidden_states)
+            return residual + self.raw_decoder.mlp(normed)
+
+        hidden_states = _ckpt(_mlp, hidden_states, use_reentrant=False)
+
+        return hidden_states
+    
+
+class Qwen3_5CausalLMForMetis(CausalLMWrapperForMetis):
+    def __init__(self, config):
+        super().__init__(config)
+
+        self.model = Qwen3_5TextModel(config.backbone_configs.text_config)
+
+        self.vocab_size = config.backbone_configs.text_config.vocab_size
+        self.lm_head = nn.Linear(config.backbone_configs.text_config.hidden_size, config.backbone_configs.text_config.vocab_size, bias=False)
+
+        # Initialize weights and apply final processing
+        self.model.post_init()
+
+    def get_decoder_layer_by_id(self, layer_id: int):
+        return self.model.layers[layer_id]
+    
+    @merge_with_config_defaults
+    @capture_outputs
+    def forward_with_memory(
+        self,
+        input_ids: torch.LongTensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        position_ids: torch.LongTensor | None = None,
+        past_key_values: Cache | None = None,
+        inputs_embeds: torch.FloatTensor | None = None,
+        use_cache: bool | None = None,
+        cache_position: torch.LongTensor | None = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> BaseModelOutputWithPast:
+        if (input_ids is None) ^ (inputs_embeds is not None):
+            raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
+
+        if inputs_embeds is None:
+            inputs_embeds = self.model.embed_tokens(input_ids)
+
+        # generate() in transformers 5.x pre-creates a standard DynamicCache.
+        # Qwen3.5's hybrid (full + linear attention) model requires Qwen3_5DynamicCache,
+        # which manages recurrent states for linear attention layers.
+        # On the first call the incoming cache is always empty, so replacement is safe.
+        # On subsequent calls the cache is already Qwen3_5DynamicCache and is left as-is.
+        if use_cache and not isinstance(past_key_values, Qwen3_5DynamicCache):
+            past_key_values = Qwen3_5DynamicCache(config=self.model.config)
+
+        if cache_position is None:
+            past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
+            cache_position = torch.arange(
+                past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device
+            )
+
+        # mrope: the hard coded `4` is for text, temporal, height and width.
+        if position_ids is None:
+            position_ids = cache_position.view(1, 1, -1).expand(4, inputs_embeds.shape[0], -1)
+        elif position_ids.ndim == 2:
+            position_ids = position_ids[None, ...].expand(4, position_ids.shape[0], -1)
+
+        if position_ids.ndim == 3 and position_ids.shape[0] == 4:
+            text_position_ids = position_ids[0]
+            position_ids = position_ids[1:]
+        else:
+            text_position_ids = None
+
+        causal_mask = create_causal_mask(
+            config=self.config,
+            inputs_embeds=inputs_embeds,
+            attention_mask=attention_mask,
+            cache_position=cache_position,
+            past_key_values=past_key_values,
+            position_ids=text_position_ids,
+        )
+        linear_attn_mask = self.model._update_linear_attn_mask(attention_mask, past_key_values)
+
+        output_hidden_states = kwargs.get("output_hidden_states", self.config.output_hidden_states)
+        all_hidden_states = () if output_hidden_states else None
+
+        hidden_states = inputs_embeds
+        position_embeddings = self.model.rotary_emb(hidden_states, position_ids)
+
+        for layer_idx, decoder_layer in enumerate(self.model.layers[: self.model.config.num_hidden_layers]):
+            layer_mask = linear_attn_mask if decoder_layer.layer_type == "linear_attention" else causal_mask
+
+            if output_hidden_states:
+                all_hidden_states += (hidden_states,)
+
+            # Use Metis Block.
+            hidden_states = self._metis_blocks_ref[layer_idx](
+                hidden_states,
+                position_embeddings=position_embeddings,
+                attention_mask=layer_mask,
+                position_ids=position_ids,
+                past_key_values=past_key_values,
+                use_cache=use_cache,
+                cache_position=cache_position,
+                **kwargs,
+            )
+
+        hidden_states = self.model.norm(hidden_states)
+
+        if output_hidden_states:
+            all_hidden_states += (hidden_states,)
+
+        return Qwen3_5ModelOutputWithPast(
+            last_hidden_state=hidden_states,
+            past_key_values=past_key_values,
+            hidden_states=all_hidden_states
+        )

README.md

@@ -0,0 +1,32 @@
+---
+license: apache-2.0
+library_name: transformers
+base_model: Qwen/Qwen3.5-9B
+tags:
+- metis
+- memory
+- custom-code
+---
+
+# Metis-9B
+
+Metis-9B is a Metis persistent-memory model built on `Qwen/Qwen3.5-9B`.
+The repository contains the complete merged model weights rather than a delta-only checkpoint.
+
+## Load
+
+```python
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+model_id = "IAAR-Shanghai/Metis-9B"
+tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
+model = AutoModelForCausalLM.from_pretrained(
+    model_id,
+    trust_remote_code=True,
+    dtype=torch.bfloat16,
+    device_map="auto",
+)
+```
+
+Metis uses custom Transformers code included in this repository. Use Transformers 5.4.0 or newer.

chat_template.jinja

@@ -0,0 +1,154 @@
+{%- set image_count = namespace(value=0) %}
+{%- set video_count = namespace(value=0) %}
+{%- macro render_content(content, do_vision_count, is_system_content=false) %}
+    {%- if content is string %}
+        {{- content }}
+    {%- elif content is iterable and content is not mapping %}
+        {%- for item in content %}
+            {%- if 'image' in item or 'image_url' in item or item.type == 'image' %}
+                {%- if is_system_content %}
+                    {{- raise_exception('System message cannot contain images.') }}
+                {%- endif %}
+                {%- if do_vision_count %}
+                    {%- set image_count.value = image_count.value + 1 %}
+                {%- endif %}
+                {%- if add_vision_id %}
+                    {{- 'Picture ' ~ image_count.value ~ ': ' }}
+                {%- endif %}
+                {{- '<|vision_start|><|image_pad|><|vision_end|>' }}
+            {%- elif 'video' in item or item.type == 'video' %}
+                {%- if is_system_content %}
+                    {{- raise_exception('System message cannot contain videos.') }}
+                {%- endif %}
+                {%- if do_vision_count %}
+                    {%- set video_count.value = video_count.value + 1 %}
+                {%- endif %}
+                {%- if add_vision_id %}
+                    {{- 'Video ' ~ video_count.value ~ ': ' }}
+                {%- endif %}
+                {{- '<|vision_start|><|video_pad|><|vision_end|>' }}
+            {%- elif 'text' in item %}
+                {{- item.text }}
+            {%- else %}
+                {{- raise_exception('Unexpected item type in content.') }}
+            {%- endif %}
+        {%- endfor %}
+    {%- elif content is none or content is undefined %}
+        {{- '' }}
+    {%- else %}
+        {{- raise_exception('Unexpected content type.') }}
+    {%- endif %}
+{%- endmacro %}
+{%- if not messages %}
+    {{- raise_exception('No messages provided.') }}
+{%- endif %}
+{%- if tools and tools is iterable and tools is not mapping %}
+    {{- '<|im_start|>system\n' }}
+    {{- "# Tools\n\nYou have access to the following functions:\n\n<tools>" }}
+    {%- for tool in tools %}
+        {{- "\n" }}
+        {{- tool | tojson }}
+    {%- endfor %}
+    {{- "\n</tools>" }}
+    {{- '\n\nIf you choose to call a function ONLY reply in the following format with NO suffix:\n\n<tool_call>\n<function=example_function_name>\n<parameter=example_parameter_1>\nvalue_1\n</parameter>\n<parameter=example_parameter_2>\nThis is the value for the second parameter\nthat can span\nmultiple lines\n</parameter>\n</function>\n</tool_call>\n\n<IMPORTANT>\nReminder:\n- Function calls MUST follow the specified format: an inner <function=...></function> block must be nested within <tool_call></tool_call> XML tags\n- Required parameters MUST be specified\n- You may provide optional reasoning for your function call in natural language BEFORE the function call, but NOT after\n- If there is no function call available, answer the question like normal with your current knowledge and do not tell the user about function calls\n</IMPORTANT>' }}
+    {%- if messages[0].role == 'system' %}
+        {%- set content = render_content(messages[0].content, false, true)|trim %}
+        {%- if content %}
+            {{- '\n\n' + content }}
+        {%- endif %}
+    {%- endif %}
+    {{- '<|im_end|>\n' }}
+{%- else %}
+    {%- if messages[0].role == 'system' %}
+        {%- set content = render_content(messages[0].content, false, true)|trim %}
+        {{- '<|im_start|>system\n' + content + '<|im_end|>\n' }}
+    {%- endif %}
+{%- endif %}
+{%- set ns = namespace(multi_step_tool=true, last_query_index=messages|length - 1) %}
+{%- for message in messages[::-1] %}
+    {%- set index = (messages|length - 1) - loop.index0 %}
+    {%- if ns.multi_step_tool and message.role == "user" %}
+        {%- set content = render_content(message.content, false)|trim %}
+        {%- if not(content.startswith('<tool_response>') and content.endswith('</tool_response>')) %}
+            {%- set ns.multi_step_tool = false %}
+            {%- set ns.last_query_index = index %}
+        {%- endif %}
+    {%- endif %}
+{%- endfor %}
+{%- if ns.multi_step_tool %}
+    {{- raise_exception('No user query found in messages.') }}
+{%- endif %}
+{%- for message in messages %}
+    {%- set content = render_content(message.content, true)|trim %}
+    {%- if message.role == "system" %}
+        {%- if not loop.first %}
+            {{- raise_exception('System message must be at the beginning.') }}
+        {%- endif %}
+    {%- elif message.role == "user" %}
+        {{- '<|im_start|>' + message.role + '\n' + content + '<|im_end|>' + '\n' }}
+    {%- elif message.role == "assistant" %}
+        {%- set reasoning_content = '' %}
+        {%- if message.reasoning_content is string %}
+            {%- set reasoning_content = message.reasoning_content %}
+        {%- else %}
+            {%- if '</think>' in content %}
+                {%- set reasoning_content = content.split('</think>')[0].rstrip('\n').split('<think>')[-1].lstrip('\n') %}
+                {%- set content = content.split('</think>')[-1].lstrip('\n') %}
+            {%- endif %}
+        {%- endif %}
+        {%- set reasoning_content = reasoning_content|trim %}
+        {%- if loop.index0 > ns.last_query_index %}
+            {{- '<|im_start|>' + message.role + '\n<think>\n' + reasoning_content + '\n</think>\n\n' + content }}
+        {%- else %}
+            {{- '<|im_start|>' + message.role + '\n' + content }}
+        {%- endif %}
+        {%- if message.tool_calls and message.tool_calls is iterable and message.tool_calls is not mapping %}
+            {%- for tool_call in message.tool_calls %}
+                {%- if tool_call.function is defined %}
+                    {%- set tool_call = tool_call.function %}
+                {%- endif %}
+                {%- if loop.first %}
+                    {%- if content|trim %}
+                        {{- '\n\n<tool_call>\n<function=' + tool_call.name + '>\n' }}
+                    {%- else %}
+                        {{- '<tool_call>\n<function=' + tool_call.name + '>\n' }}
+                    {%- endif %}
+                {%- else %}
+                    {{- '\n<tool_call>\n<function=' + tool_call.name + '>\n' }}
+                {%- endif %}
+                {%- if tool_call.arguments is defined %}
+                    {%- for args_name, args_value in tool_call.arguments|items %}
+                        {{- '<parameter=' + args_name + '>\n' }}
+                        {%- set args_value = args_value | tojson | safe if args_value is mapping or (args_value is sequence and args_value is not string) else args_value | string %}
+                        {{- args_value }}
+                        {{- '\n</parameter>\n' }}
+                    {%- endfor %}
+                {%- endif %}
+                {{- '</function>\n</tool_call>' }}
+            {%- endfor %}
+        {%- endif %}
+        {{- '<|im_end|>\n' }}
+    {%- elif message.role == "tool" %}
+        {%- if loop.previtem and loop.previtem.role != "tool" %}
+            {{- '<|im_start|>user' }}
+        {%- endif %}
+        {{- '\n<tool_response>\n' }}
+        {{- content }}
+        {{- '\n</tool_response>' }}
+        {%- if not loop.last and loop.nextitem.role != "tool" %}
+            {{- '<|im_end|>\n' }}
+        {%- elif loop.last %}
+            {{- '<|im_end|>\n' }}
+        {%- endif %}
+    {%- else %}
+        {{- raise_exception('Unexpected message role.') }}
+    {%- endif %}
+{%- endfor %}
+{%- if add_generation_prompt %}
+    {{- '<|im_start|>assistant\n' }}
+    {%- if enable_thinking is defined and enable_thinking is false %}
+        {{- '<think>\n\n</think>\n\n' }}
+    {%- else %}
+        {{- '<think>\n' }}
+    {%- endif %}
+{%- endif %}

config.json

@@ -0,0 +1,198 @@
+{
+  "architectures": [
+    "MetisForCausalLM"
+  ],
+  "auto_map": {
+    "AutoConfig": "configuration_metis.MetisConfig",
+    "AutoModelForCausalLM": "modeling_metis.MetisForCausalLM"
+  },
+  "backbone_configs": {
+    "_name_or_path": "/mnt/afs/models/Qwen3.5-9B",
+    "architectures": [
+      "Qwen3_5ForConditionalGeneration"
+    ],
+    "chunk_size_feed_forward": 0,
+    "dtype": null,
+    "id2label": {
+      "0": "LABEL_0",
+      "1": "LABEL_1"
+    },
+    "image_token_id": 248056,
+    "is_encoder_decoder": false,
+    "label2id": {
+      "LABEL_0": 0,
+      "LABEL_1": 1
+    },
+    "model_type": "qwen3_5",
+    "output_attentions": false,
+    "output_hidden_states": false,
+    "problem_type": null,
+    "return_dict": true,
+    "text_config": {
+      "_name_or_path": "",
+      "architectures": null,
+      "attention_bias": false,
+      "attention_dropout": 0.0,
+      "attn_output_gate": true,
+      "bos_token_id": null,
+      "chunk_size_feed_forward": 0,
+      "dtype": "bfloat16",
+      "eos_token_id": 248044,
+      "full_attention_interval": 4,
+      "head_dim": 256,
+      "hidden_act": "silu",
+      "hidden_size": 4096,
+      "id2label": {
+        "0": "LABEL_0",
+        "1": "LABEL_1"
+      },
+      "initializer_range": 0.02,
+      "intermediate_size": 12288,
+      "is_encoder_decoder": false,
+      "label2id": {
+        "LABEL_0": 0,
+        "LABEL_1": 1
+      },
+      "layer_types": [
+        "linear_attention",
+        "linear_attention",
+        "linear_attention",
+        "full_attention",
+        "linear_attention",
+        "linear_attention",
+        "linear_attention",
+        "full_attention",
+        "linear_attention",
+        "linear_attention",
+        "linear_attention",
+        "full_attention",
+        "linear_attention",
+        "linear_attention",
+        "linear_attention",
+        "full_attention",
+        "linear_attention",
+        "linear_attention",
+        "linear_attention",
+        "full_attention",
+        "linear_attention",
+        "linear_attention",
+        "linear_attention",
+        "full_attention",
+        "linear_attention",
+        "linear_attention",
+        "linear_attention",
+        "full_attention",
+        "linear_attention",
+        "linear_attention",
+        "linear_attention",
+        "full_attention"
+      ],
+      "linear_conv_kernel_dim": 4,
+      "linear_key_head_dim": 128,
+      "linear_num_key_heads": 16,
+      "linear_num_value_heads": 32,
+      "linear_value_head_dim": 128,
+      "mamba_ssm_dtype": "float32",
+      "max_position_embeddings": 262144,
+      "mlp_only_layers": [],
+      "model_type": "qwen3_5_text",
+      "mtp_num_hidden_layers": 1,
+      "mtp_use_dedicated_embeddings": false,
+      "num_attention_heads": 16,
+      "num_hidden_layers": 32,
+      "num_key_value_heads": 4,
+      "output_attentions": false,
+      "output_hidden_states": false,
+      "pad_token_id": null,
+      "partial_rotary_factor": 0.25,
+      "problem_type": null,
+      "return_dict": true,
+      "rms_norm_eps": 1e-06,
+      "rope_parameters": {
+        "mrope_interleaved": true,
+        "mrope_section": [
+          11,
+          11,
+          10
+        ],
+        "partial_rotary_factor": 0.25,
+        "rope_theta": 10000000,
+        "rope_type": "default"
+      },
+      "tie_word_embeddings": false,
+      "use_cache": true,
+      "vocab_size": 248320
+    },
+    "tie_word_embeddings": false,
+    "video_token_id": 248057,
+    "vision_config": {
+      "_name_or_path": "",
+      "architectures": null,
+      "chunk_size_feed_forward": 0,
+      "deepstack_visual_indexes": [],
+      "depth": 27,
+      "dtype": null,
+      "hidden_act": "gelu_pytorch_tanh",
+      "hidden_size": 1152,
+      "id2label": {
+        "0": "LABEL_0",
+        "1": "LABEL_1"
+      },
+      "in_channels": 3,
+      "initializer_range": 0.02,
+      "intermediate_size": 4304,
+      "is_encoder_decoder": false,
+      "label2id": {
+        "LABEL_0": 0,
+        "LABEL_1": 1
+      },
+      "model_type": "qwen3_5",
+      "num_heads": 16,
+      "num_position_embeddings": 2304,
+      "out_hidden_size": 4096,
+      "output_attentions": false,
+      "output_hidden_states": false,
+      "patch_size": 16,
+      "problem_type": null,
+      "return_dict": true,
+      "spatial_merge_size": 2,
+      "temporal_patch_size": 2
+    },
+    "vision_end_token_id": 248054,
+    "vision_start_token_id": 248053
+  },
+  "backbone_meta": {
+    "backbone_path": "Qwen/Qwen3.5-9B",
+    "backbone_type": "Qwen3_5"
+  },
+  "bos_token_id": null,
+  "dtype": "bfloat16",
+  "eos_token_id": 248044,
+  "hidden_size": 4096,
+  "memory_configs": {
+    "alpha_max_fraction": 0.0,
+    "alpha_max_tokens": 0,
+    "alpha_min_tokens": 1,
+    "alpha_top_p": 0.9,
+    "commit_hidden_offset": 0,
+    "forget_ratio": 1.0,
+    "gated_delta_alpha_init": 1.0,
+    "gated_delta_beta_init": 1.0,
+    "gumbel_topk_noise": true,
+    "mem_norm_init": 1.0,
+    "metis_block_type": "NormedReweightLearnedQueryMetisBlock",
+    "metis_hyper_memory_type": "StraightThroughAlphaTopPGatedDeltaRuleMetisHyperMemory",
+    "metis_local_memory_type": "NormalizedDeltaNetMetisLocalMemory",
+    "metis_reweight_gamma": 0.9,
+    "pool_temperature": 1.0,
+    "qk_kernel_type": "elu_plus_one",
+    "stride_interval": 8,
+    "uniform_num_selected": 16,
+    "update_ratio": 0.9
+  },
+  "model_type": "metis",
+  "num_attention_heads": 16,
+  "num_hidden_layers": 32,
+  "pad_token_id": null,
+  "transformers_version": "5.4.0"
+}

configuration_metis.py

@@ -0,0 +1,63 @@
+from typing import Optional
+
+from transformers import PretrainedConfig, AutoConfig
+
+class MetisConfig(PretrainedConfig):
+    model_type = "metis"
+
+    def __init__(
+        self,
+        backbone_meta: dict = {
+            'backbone_type': "Qwen3.5",
+            'backbone_path': "",
+        },
+        backbone_configs=None,
+        memory_configs=None,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.backbone_meta = backbone_meta
+
+        if backbone_configs is not None:
+            # When loading from saved config JSON, backbone_configs arrives
+            # as a plain dict.  Convert it back to the proper config class.
+            if isinstance(backbone_configs, dict) and "model_type" in backbone_configs:
+                cfg = dict(backbone_configs)
+                model_type = cfg.pop("model_type")
+                backbone_configs = AutoConfig.for_model(model_type, **cfg)
+            self.backbone_configs = backbone_configs
+        else:
+            # Eagerly load only when an explicit path is provided.
+            # During default construction (e.g. HuggingFace internal
+            # diff/serialization), backbone_path is "" so this is skipped.
+            backbone_path = backbone_meta.get('backbone_path', '')
+            if backbone_path:
+                self.backbone_configs = AutoConfig.from_pretrained(backbone_path)
+            else:
+                self.backbone_configs = None
+
+        self.memory_configs = memory_configs
+
+        # Expose backbone fields that HuggingFace internals (e.g. cache init,
+        # generation utils) expect to find directly on the top-level config.
+        if self.backbone_configs is not None:
+            text_cfg = getattr(self.backbone_configs, 'text_config', self.backbone_configs)
+            self.num_hidden_layers = text_cfg.num_hidden_layers
+            self.num_attention_heads = text_cfg.num_attention_heads
+            self.hidden_size = text_cfg.hidden_size
+            self.bos_token_id = text_cfg.bos_token_id
+            self.eos_token_id = text_cfg.eos_token_id
+            self.pad_token_id = text_cfg.pad_token_id
+
+        # ----- [Delete this block when releasing] -----
+        if self.memory_configs is None:
+            self.memory_configs = {
+                'metis_block_type': 'ReweightMetisBlock',
+                'metis_hyper_memory_type': 'FullTokensKeyNormMetisHyperMemory',
+                'metis_local_memory_type': 'NormalizedDeltaNetMetisLocalMemory',
+                'metis_reweight_gamma': 0.9,
+                'commit_hidden_offset': 0,
+            }
+        # ----- [Delete this block when releasing] -----
+
+        

generation_config.json

@@ -0,0 +1,7 @@
+{
+  "_from_model_config": true,
+  "eos_token_id": 248044,
+  "output_attentions": false,
+  "output_hidden_states": false,
+  "transformers_version": "5.4.0"
+}

metis_block.py

@@ -0,0 +1,474 @@
+from transformers import GradientCheckpointingLayer
+import torch
+import torch.nn as nn
+from .metis_hyper_memory import create_metis_hyper_memory
+from .metis_local_memory import create_metis_local_memory
+from .utils import create_metis_decoder_layer
+from abc import ABC
+
+
+def create_metis_block(config, layer_idx, raw_decoder):
+    if getattr(raw_decoder, "layer_type", "full_attention") == "linear_attention":
+        return NonMemoryMetisBlock(config, layer_idx, raw_decoder)
+    return eval(config.memory_configs['metis_block_type'])(config, layer_idx, raw_decoder)
+
+
+def _is_boundary_memory_layer(config, layer_idx: int) -> bool:
+    text_cfg = getattr(config.backbone_configs, 'text_config', config.backbone_configs)
+    num_layers = getattr(text_cfg, "num_hidden_layers", getattr(config, "num_hidden_layers", 0))
+    layer_types = getattr(text_cfg, "layer_types", None)
+    if layer_types is None:
+        return layer_idx in {0, num_layers - 1}
+    memory_layers = [
+        i for i, lt in enumerate(layer_types)
+        if lt != "linear_attention"
+    ]
+    return bool(memory_layers) and layer_idx in {memory_layers[0], memory_layers[-1]}
+
+class NonMemoryMetisBlock(GradientCheckpointingLayer):
+    """Block with no memory modules — used for debugging / ablation.
+    Passes through the backbone decoder with no memory read or write."""
+
+    def __init__(self, config, layer_idx, raw_decoder):
+        super().__init__()
+        self.config = config
+        self.layer_idx = layer_idx
+        self.local_memory = None
+        self.hyper_memory = None
+        self.monitor_branch_norms = False
+        self.last_memory_branch_norm = None
+        self.last_attention_branch_norm = None
+        self.last_memory_attention_norm_ratio = None
+        self._backbone_decoder_ref = [create_metis_decoder_layer(config, raw_decoder)]
+
+    @property
+    def backbone_decoder(self):
+        return self._backbone_decoder_ref[0]
+
+    def forward(self, hidden_states, **kwargs) -> torch.Tensor:
+        _, memory_carrier_before_mixin, cache_dict, _ = \
+            self.backbone_decoder.before_mixin(hidden_states, **kwargs)
+        return self.backbone_decoder.after_mixin(
+            memory_carrier_before_mixin, cache_dict, hidden_states, **kwargs,
+        )
+
+
+# ── Basic blocks ────────────────────────────────────────────────────────────
+
+class MetisBlockBase(GradientCheckpointingLayer, ABC):
+    def __init__(self, config, layer_idx: int, raw_decoder):
+        super().__init__()
+        self.config = config
+        self.layer_idx = layer_idx
+
+        self.local_memory = create_metis_local_memory(config)
+        self.hyper_memory = create_metis_hyper_memory(config)
+        self._backbone_decoder_ref = [create_metis_decoder_layer(config, raw_decoder)]
+        self.hyper_memory.register_raw_decoder(self._backbone_decoder_ref)
+
+        # Branch norm monitoring (on boundary full-attention layers only).
+        self.monitor_branch_norms = _is_boundary_memory_layer(config, layer_idx)
+        self.last_memory_branch_norm = None
+        self.last_attention_branch_norm = None
+        self.last_memory_attention_norm_ratio = None
+        self.last_memory_query_norm = None
+
+    @property
+    def backbone_decoder(self):
+        return self._backbone_decoder_ref[0]
+
+    def memory_integrate(self, memory_branch, memory_carrier_before_mixin):
+        raise NotImplementedError
+
+    @torch.no_grad()
+    def _record_branch_norms(self, memory_branch, memory_carrier_before_mixin):
+        if not self.monitor_branch_norms:
+            return
+        mem_norm = memory_branch.detach().float().norm()
+        attn_norm = memory_carrier_before_mixin.detach().float().norm()
+        self.last_memory_branch_norm = mem_norm.item()
+        self.last_attention_branch_norm = attn_norm.item()
+        self.last_memory_attention_norm_ratio = (
+            mem_norm / attn_norm.clamp_min(1e-12)
+        ).item()
+
+    @torch.no_grad()
+    def _record_query_norm(self, query_for_memory):
+        self.last_memory_query_norm = (
+            query_for_memory.detach().float().norm(dim=-1).mean().item()
+        )
+
+    def _init_learned_query(self, raw_decoder):
+        text_cfg = getattr(self.config.backbone_configs, 'text_config', self.config.backbone_configs)
+        hidden_dim = text_cfg.hidden_size
+        self.query_num_heads = text_cfg.num_attention_heads
+        self.query_head_dim = getattr(text_cfg, "head_dim", hidden_dim // self.query_num_heads)
+        mem_dim = self.query_num_heads * self.query_head_dim
+
+        self_attn = getattr(raw_decoder, "self_attn", None)
+        q_proj = getattr(self_attn, "q_proj", None) if self_attn is not None else None
+        has_q_bias = q_proj is not None and getattr(q_proj, "bias", None) is not None
+        self.query_proj = nn.Linear(hidden_dim, mem_dim, bias=has_q_bias)
+        nn.init.normal_(self.query_proj.weight, std=hidden_dim ** -0.5)
+        if self.query_proj.bias is not None:
+            nn.init.zeros_(self.query_proj.bias)
+
+        ln = getattr(self_attn, "q_norm", None) or raw_decoder.input_layernorm
+        norm_cls = type(ln)
+        eps = getattr(ln, "variance_epsilon", getattr(ln, "eps", 1e-6))
+        self.query_norm = norm_cls(self.query_head_dim, eps=eps)
+
+    def _make_learned_query(self, hidden_states):
+        input_shape = hidden_states.shape[:-1]
+        query_for_memory = self.query_proj(hidden_states).view(
+            *input_shape, self.query_num_heads, self.query_head_dim,
+        )
+        query_for_memory = self.query_norm(query_for_memory).transpose(1, 2).contiguous()
+        self._record_query_norm(query_for_memory)
+        return query_for_memory
+
+    def forward(self, hidden_states, **kwargs) -> torch.Tensor:
+        query_for_memory, memory_carrier_before_mixin, cache_dict, o_proj_weight = \
+            self.backbone_decoder.before_mixin(hidden_states, **kwargs)
+        if query_for_memory is None:
+            return self.backbone_decoder.after_mixin(
+                memory_carrier_before_mixin, cache_dict, hidden_states, **kwargs,
+            )
+        self._record_query_norm(query_for_memory)
+        memory_branch = o_proj_weight(self.local_memory.read(query_for_memory))
+        self._record_branch_norms(memory_branch, memory_carrier_before_mixin)
+        memory_carrier_after_mixin = self.memory_integrate(memory_branch, memory_carrier_before_mixin)
+        return self.backbone_decoder.after_mixin(
+            memory_carrier_after_mixin, cache_dict, hidden_states, **kwargs,
+        )
+
+class NaiveMetisBlock(MetisBlockBase):
+    def memory_integrate(self, memory_branch, memory_carrier_before_mixin):
+        return memory_branch + memory_carrier_before_mixin
+
+# ── Normed blocks ───────────────────────────────────────────────────────────
+
+class NormedNaiveMetisBlock(NaiveMetisBlock):
+    """NaiveMetisBlock with RMSNorm applied to the memory readout.
+
+    Inserts mem_norm between local_memory.read() and o_proj:
+        before: memory_branch = o_proj( M.read(q) )
+        after:  memory_branch = o_proj( mem_norm( M.read(q) ) )
+
+    mem_norm_init defaults to 1.0 (safe backbone-style init).  Set
+    memory_configs['mem_norm_init'] to a float to override (e.g. 0.2).
+    """
+
+    def __init__(self, config, layer_idx: int, raw_decoder):
+        super().__init__(config, layer_idx, raw_decoder)
+        text_cfg = getattr(config.backbone_configs, 'text_config', config.backbone_configs)
+        mem_dim = text_cfg.num_attention_heads * text_cfg.head_dim
+
+        ln = raw_decoder.input_layernorm
+        norm_cls = type(ln)
+        eps = getattr(ln, "variance_epsilon", getattr(ln, "eps", 1e-6))
+        self.mem_norm = norm_cls(mem_dim, eps=eps)
+
+        init_val = float(config.memory_configs.get('mem_norm_init', 1.0))
+        if hasattr(self.mem_norm, "weight") and init_val != 1.0:
+            with torch.no_grad():
+                self.mem_norm.weight.fill_(init_val)
+
+    def forward(self, hidden_states, **kwargs) -> torch.Tensor:
+        query_for_memory, memory_carrier_before_mixin, cache_dict, o_proj_weight = \
+            self.backbone_decoder.before_mixin(hidden_states, **kwargs)
+        if query_for_memory is None:
+            return self.backbone_decoder.after_mixin(
+                memory_carrier_before_mixin, cache_dict, hidden_states, **kwargs,
+            )
+        self._record_query_norm(query_for_memory)
+        normed_mem = self.mem_norm(self.local_memory.read(query_for_memory))
+        memory_branch = o_proj_weight(normed_mem)
+        self._record_branch_norms(memory_branch, memory_carrier_before_mixin)
+        memory_carrier_after_mixin = self.memory_integrate(memory_branch, memory_carrier_before_mixin)
+        return self.backbone_decoder.after_mixin(
+            memory_carrier_after_mixin, cache_dict, hidden_states, **kwargs,
+        )
+
+
+# ── Reweight blocks ─────────────────────────────────────────────────────────
+
+# class ReweightMetisBlock(NaiveMetisBlock):
+#     """Attention = gamma * Original Attention + (1 - gamma) * Memory Attention."""
+
+#     def __init__(self, config, layer_idx, raw_decoder):
+#         super().__init__(config, layer_idx, raw_decoder)
+#         self.gamma = config.memory_configs.get('metis_reweight_gamma', 0.5)
+
+#     def memory_integrate(self, memory_branch, memory_carrier_before_mixin):
+#         return self.gamma * memory_carrier_before_mixin + (1 - self.gamma) * memory_branch
+
+
+class NormedReweightMetisBlock(NormedNaiveMetisBlock):
+    """Normed memory + reweighted integration.
+    Attention = gamma * Original Attention + (1 - gamma) * Memory Attention."""
+
+    def __init__(self, config, layer_idx, raw_decoder):
+        super().__init__(config, layer_idx, raw_decoder)
+        self.gamma = config.memory_configs.get('metis_reweight_gamma', 0.5)
+
+    def memory_integrate(self, memory_branch, memory_carrier_before_mixin):
+        return self.gamma * memory_carrier_before_mixin + (1 - self.gamma) * memory_branch
+
+
+class NormedReweightLearnedQueryMetisBlock(NormedReweightMetisBlock):
+    """Normed reweight block with a dedicated trainable memory-read query."""
+
+    def __init__(self, config, layer_idx: int, raw_decoder):
+        super().__init__(config, layer_idx, raw_decoder)
+        self._init_learned_query(raw_decoder)
+
+    def forward(self, hidden_states, **kwargs) -> torch.Tensor:
+        query_for_memory, memory_carrier_before_mixin, cache_dict, o_proj_weight = \
+            self.backbone_decoder.before_mixin(hidden_states, **kwargs)
+        if query_for_memory is None:
+            return self.backbone_decoder.after_mixin(
+                memory_carrier_before_mixin, cache_dict, hidden_states, **kwargs,
+            )
+        query_for_memory = self._make_learned_query(hidden_states)
+        normed_mem = self.mem_norm(self.local_memory.read(query_for_memory))
+        memory_branch = o_proj_weight(normed_mem)
+        self._record_branch_norms(memory_branch, memory_carrier_before_mixin)
+        memory_carrier_after_mixin = self.memory_integrate(memory_branch, memory_carrier_before_mixin)
+        return self.backbone_decoder.after_mixin(
+            memory_carrier_after_mixin, cache_dict, hidden_states, **kwargs,
+        )
+
+
+class NormedGatedMetisBlock(NormedNaiveMetisBlock):
+    """NormedNaiveMetisBlock with a query-conditioned single-layer linear gate.
+
+    Replaces the fixed ``gamma`` scalar with a single-layer MLP (no bias) that
+    takes the block input (``hidden_states``) and produces a per-token blending
+    coefficient:
+
+        g      = sigmoid( gamma_gate_proj(hidden_states) )   # [B, T, 1]
+        output = g * attn_out + (1 - g) * memory_branch
+
+    Weights are initialised with a small normal std (0.01) so ``sigmoid`` output
+    starts near 0.5 for all positions.
+
+    The parent's ``mem_norm`` / ``mem_norm_init`` logic is fully preserved.
+    """
+
+    def __init__(self, config, layer_idx: int, raw_decoder):
+        super().__init__(config, layer_idx, raw_decoder)
+        text_cfg = getattr(config.backbone_configs, 'text_config', config.backbone_configs)
+        hidden_dim = text_cfg.hidden_size
+
+        self.gamma_gate_proj = nn.Linear(hidden_dim, 1, bias=False)
+        # Small init so sigmoid output starts near 0.5 (≈ uniform gate).
+        nn.init.normal_(self.gamma_gate_proj.weight, mean=0.0, std=0.01)
+
+    def forward(self, hidden_states, **kwargs) -> torch.Tensor:
+        query_for_memory, memory_carrier_before_mixin, cache_dict, o_proj_weight = \
+            self.backbone_decoder.before_mixin(hidden_states, **kwargs)
+        if query_for_memory is None:
+            return self.backbone_decoder.after_mixin(
+                memory_carrier_before_mixin, cache_dict, hidden_states, **kwargs,
+            )
+        self._record_query_norm(query_for_memory)
+        normed_mem = self.mem_norm(self.local_memory.read(query_for_memory))
+        memory_branch = o_proj_weight(normed_mem)
+        self._record_branch_norms(memory_branch, memory_carrier_before_mixin)
+        # Per-token gate conditioned on the block input hidden state.
+        g = torch.sigmoid(self.gamma_gate_proj(hidden_states))   # [B, T, 1]
+        memory_carrier_after_mixin = g * memory_carrier_before_mixin + (1.0 - g) * memory_branch
+        return self.backbone_decoder.after_mixin(
+            memory_carrier_after_mixin, cache_dict, hidden_states, **kwargs,
+        )
+
+
+class NormedGatedLearnedQueryMetisBlock(NormedGatedMetisBlock):
+    """NormedGatedMetisBlock where ``query_for_memory`` is derived from a
+    dedicated learnable projection of ``hidden_states`` rather than the
+    backbone attention Q output.
+
+    Architecture:
+        query_for_memory = query_norm( query_proj(hidden_states) )   # [B, H, T, D]
+        normed_mem       = mem_norm( local_memory.read(query_for_memory) )
+        memory_branch    = o_proj(normed_mem)
+        g                = sigmoid( gamma_gate_proj(hidden_states) ) # [B, T, 1]
+        output           = g * attn_out + (1 - g) * memory_branch
+
+    ``query_proj`` maps ``hidden_dim → num_heads × head_dim``; its bias setting
+    follows the backbone q_proj.
+    ``query_norm`` is the same RMSNorm class as the backbone's attention Q norm,
+    applied in head space so the memory read is scale-normalised.
+
+    Initialisation: ``query_proj`` uses a 1/√hidden_dim normal std, matching
+    standard attention Q-projection scale; ``query_norm`` starts at weight 1.
+    """
+
+    def __init__(self, config, layer_idx: int, raw_decoder):
+        super().__init__(config, layer_idx, raw_decoder)
+        self._init_learned_query(raw_decoder)
+
+    def forward(self, hidden_states, **kwargs) -> torch.Tensor:
+        query_for_memory, memory_carrier_before_mixin, cache_dict, o_proj_weight = \
+            self.backbone_decoder.before_mixin(hidden_states, **kwargs)
+        if query_for_memory is None:
+            return self.backbone_decoder.after_mixin(
+                memory_carrier_before_mixin, cache_dict, hidden_states, **kwargs,
+            )
+        # Replace backbone Q with a learned projection from hidden_states.
+        query_for_memory = self._make_learned_query(hidden_states)
+        normed_mem = self.mem_norm(self.local_memory.read(query_for_memory))
+        memory_branch = o_proj_weight(normed_mem)
+        self._record_branch_norms(memory_branch, memory_carrier_before_mixin)
+        g = torch.sigmoid(self.gamma_gate_proj(hidden_states))        # [B, T, 1]
+        memory_carrier_after_mixin = g * memory_carrier_before_mixin + (1.0 - g) * memory_branch
+        return self.backbone_decoder.after_mixin(
+            memory_carrier_after_mixin, cache_dict, hidden_states, **kwargs,
+        )
+
+
+class NormedSwiGLUGatedMetisBlock(NormedNaiveMetisBlock):
+    """NormedNaiveMetisBlock with a two-layer SwiGLU gate MLP.
+
+    The blending gate is computed by a SwiGLU MLP over the block input:
+
+        h = SiLU( gate_proj(hidden_states) ) * up_proj(hidden_states)   # [B, T, D_gate]
+        g = sigmoid( down_proj(h) )                                      # [B, T, 1]
+        output = g * attn_out + (1 − g) * memory_branch
+
+    Intermediate dimension ``D_gate`` defaults to ``hidden_dim // 4`` and can
+    be overridden via ``memory_configs['swiglu_gate_hidden_dim']``.
+
+    Initialization: ``gate_proj`` / ``up_proj`` use small-std normal so the
+    SwiGLU activations start near zero; ``down_proj`` is zeroed so
+    ``sigmoid(0) = 0.5`` at the first step (balanced blending).
+    """
+
+    def __init__(self, config, layer_idx: int, raw_decoder):
+        super().__init__(config, layer_idx, raw_decoder)
+        text_cfg = getattr(config.backbone_configs, 'text_config', config.backbone_configs)
+        hidden_dim = text_cfg.hidden_size
+        gate_dim = int(config.memory_configs.get('swiglu_gate_hidden_dim', hidden_dim // 4))
+
+        self.gamma_gate_proj = nn.Linear(hidden_dim, gate_dim, bias=False)
+        self.gamma_up_proj   = nn.Linear(hidden_dim, gate_dim, bias=False)
+        self.gamma_down_proj = nn.Linear(gate_dim,   1,        bias=False)
+
+        nn.init.normal_(self.gamma_gate_proj.weight, std=0.01)
+        nn.init.normal_(self.gamma_up_proj.weight,   std=0.01)
+        nn.init.zeros_(self.gamma_down_proj.weight)
+
+    def forward(self, hidden_states, **kwargs) -> torch.Tensor:
+        query_for_memory, memory_carrier_before_mixin, cache_dict, o_proj_weight = \
+            self.backbone_decoder.before_mixin(hidden_states, **kwargs)
+        if query_for_memory is None:
+            return self.backbone_decoder.after_mixin(
+                memory_carrier_before_mixin, cache_dict, hidden_states, **kwargs,
+            )
+        self._record_query_norm(query_for_memory)
+        normed_mem = self.mem_norm(self.local_memory.read(query_for_memory))
+        memory_branch = o_proj_weight(normed_mem)
+        self._record_branch_norms(memory_branch, memory_carrier_before_mixin)
+        h = nn.functional.silu(self.gamma_gate_proj(hidden_states)) * self.gamma_up_proj(hidden_states)
+        g = torch.sigmoid(self.gamma_down_proj(h))                   # [B, T, 1]
+        memory_carrier_after_mixin = g * memory_carrier_before_mixin + (1.0 - g) * memory_branch
+        return self.backbone_decoder.after_mixin(
+            memory_carrier_after_mixin, cache_dict, hidden_states, **kwargs,
+        )
+
+
+# # ── Branch Norm Matching blocks ─────────────────────────────────────────────
+
+# class BNMMetisBlock(NaiveMetisBlock):
+#     """Branch Norm Matching: scale mem_branch so its sequence-average norm matches attn_out.
+
+#     Per-token alignment is noisy, so one scale factor per sequence is used:
+#         mean_norm_attn = mean_t( ||attn_out_t|| )
+#         mean_norm_mem  = mean_t( ||mem_branch_t|| )
+#         scale          = mean_norm_attn / (mean_norm_mem + eps)
+#         output         = mem_branch * scale + attn_out
+#     """
+
+#     BNM_EPS: float = 1e-6
+
+#     def memory_integrate(self, memory_branch, memory_carrier_before_mixin):
+#         norm_attn = memory_carrier_before_mixin.norm(dim=-1)
+#         norm_mem = memory_branch.norm(dim=-1)
+#         mean_norm_attn = norm_attn.mean(dim=1)
+#         mean_norm_mem = norm_mem.mean(dim=1)
+#         scale = (mean_norm_attn / (mean_norm_mem + self.BNM_EPS)).view(-1, 1, 1)
+#         return memory_branch * scale + memory_carrier_before_mixin
+
+
+# class BNMReweightMetisBlock(BNMMetisBlock):
+#     """BNM with an additional scalar gate gamma.
+#     output = gamma * attn_out + (1 - gamma) * mem_branch_scaled"""
+
+#     def __init__(self, config, layer_idx, raw_decoder):
+#         super().__init__(config, layer_idx, raw_decoder)
+#         self.gamma = config.memory_configs.get('metis_reweight_gamma', 0.5)
+
+#     def memory_integrate(self, memory_branch, memory_carrier_before_mixin):
+#         norm_attn = memory_carrier_before_mixin.norm(dim=-1)
+#         norm_mem = memory_branch.norm(dim=-1)
+#         mean_norm_attn = norm_attn.mean(dim=1)
+#         mean_norm_mem = norm_mem.mean(dim=1)
+#         scale = (mean_norm_attn / (mean_norm_mem + BNMMetisBlock.BNM_EPS)).view(-1, 1, 1)
+#         mem_scaled = memory_branch * scale
+#         return self.gamma * memory_carrier_before_mixin + (1 - self.gamma) * mem_scaled
+
+
+# ── Gated blocks (from operation branch) ─────────────────────────────────────
+
+# class GatedNormedNaiveMetisBlock(NaiveMetisBlock):
+#     """NaiveMetisBlock + RMSNorm + learnable per-dimension read gate.
+
+#     Architecture:
+#         raw_mem   = local_memory.read(q)
+#         gated_mem = sigmoid(read_gate) * mem_norm(raw_mem)
+#         memory_branch = o_proj(gated_mem)
+#         output = memory_branch + attn_output
+
+#     read_gate is init to -1.5 (sigmoid(-1.5) ≈ 0.18) so memory starts small.
+#     """
+
+#     def __init__(self, config, layer_idx: int, raw_decoder):
+#         super().__init__(config, layer_idx, raw_decoder)
+#         text_cfg = getattr(config.backbone_configs, 'text_config', config.backbone_configs)
+#         mem_dim = text_cfg.num_attention_heads * text_cfg.head_dim
+
+#         ln = raw_decoder.input_layernorm
+#         norm_cls = type(ln)
+#         eps = getattr(ln, "variance_epsilon", getattr(ln, "eps", 1e-6))
+#         self.mem_norm = norm_cls(mem_dim, eps=eps)
+#         if hasattr(self.mem_norm, "weight"):
+#             with torch.no_grad():
+#                 self.mem_norm.weight.fill_(0.2)
+#         self.read_gate = nn.Parameter(torch.full((mem_dim,), -1.5))
+
+#     def forward(self, hidden_states, **kwargs) -> torch.Tensor:
+#         query_for_memory, memory_carrier_before_mixin, cache_dict, o_proj_weight = \
+#             self.backbone_decoder.before_mixin(hidden_states, **kwargs)
+#         if query_for_memory is None:
+#             return self.backbone_decoder.after_mixin(
+#                 memory_carrier_before_mixin, cache_dict, hidden_states, **kwargs,
+#             )
+#         raw_mem = self.local_memory.read(query_for_memory)
+#         gated_mem = torch.sigmoid(self.read_gate) * self.mem_norm(raw_mem)
+#         memory_branch = o_proj_weight(gated_mem)
+#         self._record_branch_norms(memory_branch, memory_carrier_before_mixin)
+#         memory_carrier_after_mixin = self.memory_integrate(memory_branch, memory_carrier_before_mixin)
+#         return self.backbone_decoder.after_mixin(
+#             memory_carrier_after_mixin, cache_dict, hidden_states, **kwargs,
+#         )
+
+
+# class LegacyGatedNormedNaiveMetisBlock(GatedNormedNaiveMetisBlock):
+#     """Legacy variant: train read_gate, keep mem_norm at backbone-style init (1.0)."""
+
+#     def __init__(self, config, layer_idx: int, raw_decoder):
+#         super().__init__(config, layer_idx, raw_decoder)
+#         if hasattr(self.mem_norm, "weight"):
+#             with torch.no_grad():
+#                 self.mem_norm.weight.fill_(1.0)

metis_hyper_memory.py

@@ -0,0 +1,1561 @@
+import math
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from abc import ABC
+
+def _qk_kernel(x: torch.Tensor, kernel_type: str = "elu_plus_one") -> torch.Tensor:
+    if kernel_type == "elu_plus_one":
+        return F.elu(x) + 1.0
+    if kernel_type == "relu_square":
+        return F.relu(x).square()
+    if kernel_type == "softplus":
+        return F.softplus(x)
+    raise ValueError(f"Unsupported qk kernel type: {kernel_type}")
+
+def create_metis_hyper_memory(config):
+    return eval(config.memory_configs['metis_hyper_memory_type'])(config)
+
+class MetisHyperMemoryBase(nn.Module, ABC):
+    def __init__(self, config) -> None:
+        super().__init__()
+        self.config = config
+        # Qwen 3.5 has text config, but Qwen 3 does not.
+        self.text_cfg = getattr(config.backbone_configs, 'text_config', config.backbone_configs)
+        # Reference set later by MetisBlock via register_raw_decoder().
+        # Stored as a list to avoid registering the backbone decoder as a submodule.
+        self._backbone_decoder_ref: list | None = None
+
+    def register_raw_decoder(self, _backbone_decoder_ref: list) -> None:
+        """Called by MetisBlock to give HyperMemory access to the backbone decoder.
+
+        Enables hyper memory variants to apply backbone-style normalisations
+        (e.g. input_layernorm) on hidden states before computing W_k/W_v.
+        """
+        self._backbone_decoder_ref = _backbone_decoder_ref
+
+    @property
+    def backbone_decoder(self):
+        if self._backbone_decoder_ref is None:
+            raise RuntimeError("backbone_decoder not registered; call register_raw_decoder first")
+        return self._backbone_decoder_ref[0]
+
+    def update_local_memory(self, raw_info, local_memory) -> None:
+        local_memory.write(self.get_new_info_for_local_memory(raw_info))
+
+    def get_new_info_for_local_memory(self, raw_info):
+        raise NotImplementedError
+
+
+class LinearLastMetisHyperMemory(MetisHyperMemoryBase):
+    """Additive memory update using the last token's hidden state.
+
+    M_new = (1 - update_ratio) * M_old + update_ratio * (W_k(h_norm)^T @ W_v(h_norm))
+
+    where h_norm = backbone.input_layernorm(h_last).  Applying the backbone's
+    RMSNorm before W_k / W_v bounds pre-projection magnitudes (mirroring how
+    the backbone's own attention consumes its input).
+
+    All token-selection subclasses (Uniform / Stride / AllTokens) follow the
+    same layernorm-then-project pattern.  The exception is
+    ``NormalizedLinearLastMetisHyperMemory``, which keeps the legacy
+    L2-normalize-on-output behaviour.
+
+    update_ratio is read from memory_configs (default: 1.0).
+    """
+
+    def __init__(self, config) -> None:
+        super().__init__(config)
+        hidden_size = self.text_cfg.hidden_size
+
+        num_q_heads  = self.text_cfg.num_attention_heads
+        num_kv_heads = getattr(self.text_cfg, "num_key_value_heads", num_q_heads)
+        head_dim     = getattr(self.text_cfg, "head_dim", hidden_size // num_q_heads)
+
+        # W_k / W_v output dim must equal the local memory matrix's kv_dim so
+        # that write vectors align with read queries.  Default = GQA layout
+        # (num_kv_heads * head_dim).  Switch to MHA layout (num_q_heads * head_dim)
+        # when the user chose the legacy MHA local memory.
+        local_mem_type = config.memory_configs.get('metis_local_memory_type', '')
+        if local_mem_type.startswith('MHA'):
+            self.kv_dim = num_q_heads * head_dim
+        else:
+            self.kv_dim = num_kv_heads * head_dim
+
+        self.update_ratio = config.memory_configs.get('update_ratio', 1.0)
+
+        self.W_k = nn.Linear(hidden_size, self.kv_dim, bias=False)
+        self.W_v = nn.Linear(hidden_size, self.kv_dim, bias=False)
+
+    def get_new_info_for_local_memory(self, raw_info: torch.Tensor,
+                                       attention_mask: torch.Tensor | None = None) -> torch.Tensor:
+        """Compute the memory delta from the last *real* token of each sample.
+
+        Args:
+            raw_info:       hidden states  (b, s, hidden_size)
+            attention_mask: binary mask    (b, s)  — 1 for real tokens, 0 for pad.
+                            When None, the last position is used (safe for unbatched
+                            or already-trimmed sequences).
+
+        Returns:
+            delta: (b, D, D) — outer product of write key and write value.
+        """
+        if attention_mask is not None:
+            # last real token index per sample: sum of 1s minus 1
+            last_idx = attention_mask.sum(dim=1) - 1          # (b,)
+            b = raw_info.size(0)
+            h_last = raw_info[torch.arange(b, device=raw_info.device),
+                               last_idx, :].unsqueeze(1)       # (b, 1, hidden_size)
+        else:
+            h_last = raw_info[:, -1:, :]                      # (b, 1, hidden_size)
+        h_last = self.backbone_decoder.raw_decoder.input_layernorm(h_last)
+        write_key = self.W_k(h_last)                          # (b, 1, kv_dim)
+        write_value = self.W_v(h_last)                        # (b, 1, kv_dim)
+        # (b, kv_dim, 1) @ (b, 1, kv_dim) -> (b, kv_dim, kv_dim)
+        return torch.matmul(write_key.transpose(-1, -2), write_value)
+
+    def update_local_memory(self, raw_info: torch.Tensor, local_memory,
+                            attention_mask: torch.Tensor | None = None) -> None:
+        """Blend the existing memory with the new additive update."""
+        delta = self.get_new_info_for_local_memory(raw_info, attention_mask)
+        if local_memory.state is not None:
+            new_state = (1.0 - self.update_ratio) * local_memory.state + self.update_ratio * delta
+        else:
+            new_state = self.update_ratio * delta
+        local_memory.write(new_state)
+
+
+class NormalizedLinearLastMetisHyperMemory(LinearLastMetisHyperMemory):
+    """Exception class: legacy L2-normalized W_k/W_v output, no input_layernorm.
+
+    Unlike all other subclasses (which apply backbone.input_layernorm to
+    hidden states before W_k / W_v), this class operates on raw hidden
+    states and L2-normalises the projection *outputs*:
+
+        write_key   = F.normalize(W_k(h_last), dim=-1)   # ‖·‖ = 1
+        write_value = F.normalize(W_v(h_last), dim=-1)   # ‖·‖ = 1
+        ‖delta‖_F   = ‖write_key‖ · ‖write_value‖ = 1
+
+    The per-step memory increment is bounded by update_ratio.  Kept
+    primarily for reproducing earlier experiments.
+    """
+
+    def get_new_info_for_local_memory(self, raw_info: torch.Tensor,
+                                       attention_mask: torch.Tensor | None = None) -> torch.Tensor:
+        if attention_mask is not None:
+            last_idx = attention_mask.sum(dim=1) - 1          # (b,)
+            b = raw_info.size(0)
+            h_last = raw_info[torch.arange(b, device=raw_info.device),
+                               last_idx, :].unsqueeze(1)       # (b, 1, hidden_size)
+        else:
+            h_last = raw_info[:, -1:, :]                      # (b, 1, hidden_size)
+        write_key   = F.normalize(self.W_k(h_last), dim=-1)   # (b, 1, D), ‖·‖=1
+        write_value = F.normalize(self.W_v(h_last), dim=-1)   # (b, 1, D), ‖·‖=1
+        # (b, D, 1) @ (b, 1, D) -> (b, D, D),  ‖delta‖_F ≤ 1
+        return torch.matmul(write_key.transpose(-1, -2), write_value)
+
+
+class UniformNormalizedMetisHyperMemory(LinearLastMetisHyperMemory):
+    """Memory update using uniformly sampled tokens with backbone-norm pre-projection.
+
+    This class selects N = ``uniform_num_selected`` tokens evenly spaced
+    across the real sequence (always including the last real token),
+    applies the backbone's input_layernorm to the selected hidden states,
+    then projects with W_k / W_v.  Each selected token contributes one
+    rank-1 outer product to the memory delta:
+
+        step  = L / N                 (L = real sequence length)
+        idx_j = round(j * step)  for j in 0..N-1
+        idx_{N-1} = L - 1             (force-include last)
+
+        h_normed   = input_layernorm(h[idx])         # (b, N, hidden)
+        write_key  = W_k(h_normed)                   # (b, N, kv_dim)
+        write_val  = W_v(h_normed)                   # (b, N, kv_dim)
+        delta      = write_key.T @ write_val         # (b, kv_dim, kv_dim)
+
+    Configurable via ``memory_configs``:
+        - ``uniform_num_selected``  (int, default 16): number of tokens N
+
+    When the real sequence is shorter than N, all real tokens are used
+    and the last one is repeated to fill the remaining slots.
+    """
+
+    DEFAULT_NUM_SELECTED: int = 16
+
+    def __init__(self, config) -> None:
+        super().__init__(config)
+        self.num_selected = int(
+            config.memory_configs.get('uniform_num_selected', self.DEFAULT_NUM_SELECTED)
+        )
+
+    def _select_tokens(
+        self,
+        hidden_states: torch.Tensor,            # (b, s, hidden_size)
+        attention_mask: torch.Tensor | None,
+    ) -> torch.Tensor:                           # (b, N, hidden_size)
+        b, S, hidden_size = hidden_states.shape
+        N = self.num_selected
+        device = hidden_states.device
+
+        # Real sequence length per sample.
+        if attention_mask is not None:
+            lengths = attention_mask.sum(dim=1).long()          # (b,)
+        else:
+            lengths = torch.full((b,), S, dtype=torch.long, device=device)
+
+        # Build per-sample index tensors (b, N).
+        indices_list = []
+        for bi in range(b):
+            L = lengths[bi].item()
+            if L <= N:
+                # Fewer real tokens than slots: use all, repeat last to pad.
+                idx = list(range(L)) + [L - 1] * (N - L)
+            else:
+                # Uniformly spaced: step = L/N, always land last on L-1.
+                step = L / N
+                idx = [min(int(i * step), L - 1) for i in range(N)]
+                idx[-1] = L - 1
+            indices_list.append(idx)
+
+        indices = torch.tensor(indices_list, dtype=torch.long, device=device)
+        idx_exp = indices.unsqueeze(-1).expand(b, N, hidden_size)  # (b, N, hidden)
+        return hidden_states.gather(1, idx_exp)                    # (b, N, hidden)
+
+    def get_new_info_for_local_memory(
+        self,
+        raw_info: torch.Tensor,                 # (b, s, hidden_size)
+        attention_mask: torch.Tensor | None = None,
+    ) -> torch.Tensor:                           # (b, kv_dim, kv_dim)
+        h_sel = self._select_tokens(raw_info, attention_mask)         # (b, N, hidden)
+        # Apply backbone's RMSNorm (same one that gates the layer's attention).
+        h_sel = self.backbone_decoder.raw_decoder.input_layernorm(h_sel)
+        write_key   = self.W_k(h_sel)                                 # (b, N, kv_dim)
+        write_value = self.W_v(h_sel)                                 # (b, N, kv_dim)
+        # (b, kv_dim, N) @ (b, N, kv_dim) -> (b, kv_dim, kv_dim)
+        return torch.matmul(write_key.transpose(-1, -2), write_value)
+
+
+class StrideNormalizedMetisHyperMemory(LinearLastMetisHyperMemory):
+    """Memory update using stride-based token selection with backbone-norm pre-projection.
+
+    Unlike ``UniformNormalizedMetisHyperMemory`` (fixed-N evenly-spaced),
+    this class selects **every K-th real token** from each sample, plus
+    the last real token.  K is configurable; the number of selected tokens
+    per sample varies with sequence length:
+
+        L  = real (non-padding) sequence length
+        K  = ``stride_interval``  (config, default 16)
+        idx = [0, K, 2K, ...]  intersected with [0, L-1],  union {L-1}
+
+    For mixed-length batches the per-sample selection counts differ; padded
+    slots in the resulting (b, N_max, hidden) tensor are masked to zero so
+    they contribute nothing to the rank-1 outer products.
+
+    Configurable via ``memory_configs``:
+        - ``stride_interval``  (int, default 16): K, the spacing between picks
+    """
+
+    DEFAULT_STRIDE: int = 16
+
+    def __init__(self, config) -> None:
+        super().__init__(config)
+        self.stride = int(config.memory_configs.get('stride_interval', self.DEFAULT_STRIDE))
+        if self.stride <= 0:
+            raise ValueError(f"stride_interval must be > 0, got {self.stride}")
+
+    def _select_tokens_with_mask(
+        self,
+        hidden_states: torch.Tensor,            # (b, s, hidden_size)
+        attention_mask: torch.Tensor | None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:    # (b, N_max, hidden), (b, N_max)
+        b, S, hidden_size = hidden_states.shape
+        device = hidden_states.device
+        K = self.stride
+
+        # Real length per sample — only non-padding tokens are eligible.
+        if attention_mask is not None:
+            lengths = attention_mask.sum(dim=1).long().tolist()
+        else:
+            lengths = [S] * b
+
+        # Per-sample stride-K indices, always force-including the last real token.
+        per_sample_idx: list[list[int]] = []
+        for L in lengths:
+            if L <= 0:
+                # Edge case: empty sample.  Use index 0 (will be masked out).
+                per_sample_idx.append([0])
+                continue
+            idx = list(range(0, L, K))
+            if idx[-1] != L - 1:
+                idx.append(L - 1)
+            per_sample_idx.append(idx)
+
+        N_max = max(len(idx) for idx in per_sample_idx)
+
+        # Right-pad each sample's index list with 0 (a real position) and
+        # record a 0/1 mask so padded slots contribute zero to the outer product.
+        indices_padded: list[list[int]] = []
+        masks: list[list[float]] = []
+        for idx, L in zip(per_sample_idx, lengths):
+            n_valid = len(idx) if L > 0 else 0
+            pad_n = N_max - len(idx)
+            indices_padded.append(idx + [0] * pad_n)
+            masks.append([1.0] * n_valid + [0.0] * (N_max - n_valid))
+
+        indices = torch.tensor(indices_padded, dtype=torch.long, device=device)
+        mask    = torch.tensor(masks, dtype=hidden_states.dtype, device=device)
+
+        idx_exp = indices.unsqueeze(-1).expand(b, N_max, hidden_size)  # (b, N_max, hidden)
+        h_sel   = hidden_states.gather(1, idx_exp)                     # (b, N_max, hidden)
+        return h_sel, mask
+
+    def get_new_info_for_local_memory(
+        self,
+        raw_info: torch.Tensor,                 # (b, s, hidden_size)
+        attention_mask: torch.Tensor | None = None,
+    ) -> torch.Tensor:                           # (b, kv_dim, kv_dim)
+        h_sel, mask = self._select_tokens_with_mask(raw_info, attention_mask)
+        # Apply backbone's RMSNorm before the W_k / W_v projections.
+        h_sel = self.backbone_decoder.raw_decoder.input_layernorm(h_sel)
+        write_key   = self.W_k(h_sel)                                   # (b, N, kv_dim)
+        write_value = self.W_v(h_sel)                                   # (b, N, kv_dim)
+        # Zero-out padded slots so they contribute nothing to the matmul.
+        mask = mask.unsqueeze(-1)                                       # (b, N, 1)
+        write_key   = write_key   * mask
+        write_value = write_value * mask
+        # (b, kv_dim, N) @ (b, N, kv_dim) -> (b, kv_dim, kv_dim)
+        return torch.matmul(write_key.transpose(-1, -2), write_value)
+
+
+class FullTokensNormalizedv3MetisHyperMemory(LinearLastMetisHyperMemory):
+    """Memory update using all real tokens with v3-style normalization.
+
+    Every non-padding token contributes one rank-1 outer product:
+
+        h_normed   = input_layernorm(h)             # (b, s, hidden)
+        write_key  = W_k(h_normed)                  # (b, s, kv_dim)
+        write_val  = W_v(h_normed)                  # (b, s, kv_dim)
+        delta      = write_key.T @ write_val / (L * sqrt(D))
+
+    This matches ``StrideNormalizedv3MetisHyperMemory``'s normalization while
+    selecting the full real-token sequence instead of stride-sampled tokens.
+    """
+
+    def get_new_info_for_local_memory(
+        self,
+        raw_info: torch.Tensor,                 # (b, s, hidden_size)
+        attention_mask: torch.Tensor | None = None,
+    ) -> torch.Tensor:                           # (b, kv_dim, kv_dim)
+        h = self.backbone_decoder.raw_decoder.input_layernorm(raw_info)
+        write_key   = self.W_k(h)                                # (b, s, kv_dim)
+        write_value = self.W_v(h)                                # (b, s, kv_dim)
+
+        if attention_mask is not None:
+            # Broadcast mask over hidden dim so pad positions contribute 0.
+            mask = attention_mask.unsqueeze(-1).to(write_key.dtype)  # (b, s, 1)
+            L_prime = attention_mask.sum(dim=1).clamp(min=1)         # (b,)
+            write_key   = write_key   * mask
+            write_value = write_value * mask
+        else:
+            L_prime = torch.full(
+                (raw_info.size(0),),
+                raw_info.size(1),
+                dtype=write_key.dtype,
+                device=raw_info.device,
+            ).clamp(min=1)
+
+        # (b, kv_dim, s) @ (b, s, kv_dim) -> (b, kv_dim, kv_dim)
+        delta = torch.matmul(write_key.transpose(-1, -2), write_value)
+        scale = L_prime.to(delta.dtype) * (self.kv_dim ** 0.5)
+        return delta / scale.view(-1, 1, 1)
+
+
+class KeyNormTokenAggMetisHyperMemory(LinearLastMetisHyperMemory):
+    """Shared write path for token-aggregation experiments.
+
+    Subclasses choose or pool hidden states into ``(h_tokens, mask)``.  This
+    base class then applies the same key-normalized DeltaNet write protocol as
+    ``FullTokensKeyNormMetisHyperMemory``:
+
+        k = normalize(W_k(input_layernorm(h))) / sqrt(D)
+        v = W_v(input_layernorm(h))
+        state = mean_t(k_t^T @ v_t)
+        key_state = mean_t(k_t)
+    """
+
+    def _delta_from_normed_tokens(
+        self,
+        h_normed: torch.Tensor,
+        mask: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        write_key = F.normalize(self.W_k(h_normed), dim=-1) / (self.kv_dim ** 0.5)
+        write_value = self.W_v(h_normed)
+
+        if mask is not None:
+            mask_exp = mask.unsqueeze(-1).to(write_key.dtype)
+            lengths = mask.sum(dim=1).clamp(min=1).to(write_key.dtype)
+            write_key = write_key * mask_exp
+            write_value = write_value * mask_exp
+        else:
+            lengths = torch.full(
+                (h_normed.size(0),),
+                h_normed.size(1),
+                dtype=write_key.dtype,
+                device=h_normed.device,
+            ).clamp(min=1)
+
+        delta_state = torch.matmul(write_key.transpose(-1, -2), write_value)
+        ones = torch.ones(
+            write_key.size(0),
+            write_key.size(1),
+            1,
+            device=write_key.device,
+            dtype=write_key.dtype,
+        )
+        delta_key_state = torch.matmul(write_key.transpose(-1, -2), ones)
+        scale = lengths.view(-1, 1, 1)
+        return delta_state / scale, delta_key_state / scale
+
+    def _delta_from_raw_tokens(
+        self,
+        h_tokens: torch.Tensor,
+        mask: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        h_normed = self.backbone_decoder.raw_decoder.input_layernorm(h_tokens)
+        return self._delta_from_normed_tokens(h_normed, mask)
+
+    def _write_keynorm_update(
+        self,
+        raw_info: torch.Tensor,
+        local_memory,
+        attention_mask: torch.Tensor | None = None,
+    ) -> None:
+        delta_state, delta_key_state = self.get_new_info_for_local_memory(raw_info, attention_mask)
+        if getattr(local_memory, "key_state", None) is None:
+            new_state = self.update_ratio * delta_state
+            new_key_state = self.update_ratio * delta_key_state
+        else:
+            new_state = (1.0 - self.update_ratio) * local_memory.state + self.update_ratio * delta_state
+            new_key_state = (
+                (1.0 - self.update_ratio) * local_memory.key_state
+                + self.update_ratio * delta_key_state
+            )
+        local_memory.write(new_state, new_key_state)
+
+    def update_local_memory(
+        self,
+        raw_info: torch.Tensor,
+        local_memory,
+        attention_mask: torch.Tensor | None = None,
+    ) -> None:
+        self._write_keynorm_update(raw_info, local_memory, attention_mask)
+
+
+class MeanPoolKeyNormMetisHyperMemory(KeyNormTokenAggMetisHyperMemory):
+    """Mean-pool all real hidden states into one write token."""
+
+    def get_new_info_for_local_memory(
+        self,
+        raw_info: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        if attention_mask is not None:
+            mask = attention_mask.unsqueeze(-1).to(raw_info.dtype)
+            lengths = attention_mask.sum(dim=1).clamp(min=1).to(raw_info.dtype)
+            h_pool = (raw_info * mask).sum(dim=1, keepdim=True) / lengths.view(-1, 1, 1)
+        else:
+            h_pool = raw_info.mean(dim=1, keepdim=True)
+        return self._delta_from_raw_tokens(h_pool)
+
+
+class StridePoolKeyNormMetisHyperMemory(KeyNormTokenAggMetisHyperMemory):
+    """Mean-pool every stride-sized chunk into one write token per chunk."""
+
+    DEFAULT_STRIDE: int = 8
+
+    def __init__(self, config) -> None:
+        super().__init__(config)
+        self.stride = int(config.memory_configs.get("stride_interval", self.DEFAULT_STRIDE))
+        if self.stride <= 0:
+            raise ValueError(f"stride_interval must be > 0, got {self.stride}")
+
+    def _pool_stride_windows(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: torch.Tensor | None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        b, S, hidden_size = hidden_states.shape
+        device = hidden_states.device
+        dtype = hidden_states.dtype
+        if attention_mask is not None:
+            lengths = attention_mask.sum(dim=1).long().tolist()
+        else:
+            lengths = [S] * b
+
+        per_sample: list[torch.Tensor] = []
+        masks: list[list[float]] = []
+        max_chunks = 1
+        for bi, L in enumerate(lengths):
+            L = max(int(L), 1)
+            chunks = []
+            for start in range(0, L, self.stride):
+                end = min(start + self.stride, L)
+                chunks.append(hidden_states[bi, start:end].mean(dim=0))
+            sample = torch.stack(chunks, dim=0)
+            per_sample.append(sample)
+            max_chunks = max(max_chunks, sample.size(0))
+
+        padded = []
+        for sample in per_sample:
+            pad_n = max_chunks - sample.size(0)
+            if pad_n > 0:
+                pad = torch.zeros(pad_n, hidden_size, device=device, dtype=dtype)
+                sample = torch.cat([sample, pad], dim=0)
+            padded.append(sample)
+            masks.append([1.0] * (sample.size(0) - pad_n) + [0.0] * pad_n)
+
+        return torch.stack(padded, dim=0), torch.tensor(masks, device=device, dtype=dtype)
+
+    def get_new_info_for_local_memory(
+        self,
+        raw_info: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        h_pool, mask = self._pool_stride_windows(raw_info, attention_mask)
+        return self._delta_from_raw_tokens(h_pool, mask)
+
+
+class AttentionPoolKeyNormMetisHyperMemory(KeyNormTokenAggMetisHyperMemory):
+    """Learn a global attention pooling query and write one pooled token."""
+
+    def __init__(self, config) -> None:
+        super().__init__(config)
+        self.pool_score = nn.Linear(self.text_cfg.hidden_size, 1, bias=False)
+        self.pool_temperature = float(config.memory_configs.get("pool_temperature", 1.0))
+        if self.pool_temperature <= 0:
+            raise ValueError(f"pool_temperature must be > 0, got {self.pool_temperature}")
+
+    def get_new_info_for_local_memory(
+        self,
+        raw_info: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        h_normed = self.backbone_decoder.raw_decoder.input_layernorm(raw_info)
+        scores = self.pool_score(h_normed).squeeze(-1)
+        if attention_mask is not None:
+            scores = scores.masked_fill(attention_mask == 0, torch.finfo(scores.dtype).min)
+        weights = torch.softmax(scores / self.pool_temperature, dim=1).unsqueeze(1)
+        h_pool = torch.matmul(weights, h_normed)
+        return self._delta_from_normed_tokens(h_pool)
+
+
+class WindowAttentionPoolKeyNormMetisHyperMemory(AttentionPoolKeyNormMetisHyperMemory):
+    """Soft-select one pooled write token per stride-sized window.
+
+    This is the differentiable replacement for hard top-k token selection used
+    by the token-aggregation experiments.  For stride ``R`` it writes roughly
+    ``ceil(L / R)`` tokens, matching stride/top-k compression, but the scorer
+    receives gradients from every real token in each window during both
+    training and inference.
+    """
+
+    DEFAULT_STRIDE: int = 8
+
+    def __init__(self, config) -> None:
+        super().__init__(config)
+        self.stride = int(config.memory_configs.get("stride_interval", self.DEFAULT_STRIDE))
+        if self.stride <= 0:
+            raise ValueError(f"stride_interval must be > 0, got {self.stride}")
+
+    def _pool_attention_windows(
+        self,
+        h_normed: torch.Tensor,
+        attention_mask: torch.Tensor | None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        b, S, hidden_size = h_normed.shape
+        if attention_mask is None:
+            attention_mask = torch.ones(b, S, device=h_normed.device, dtype=torch.long)
+        pad_n = (-S) % self.stride
+        if pad_n > 0:
+            h_normed = F.pad(h_normed, (0, 0, 0, pad_n))
+            attention_mask = F.pad(attention_mask, (0, pad_n))
+
+        W = h_normed.size(1) // self.stride
+        h_win = h_normed.view(b, W, self.stride, hidden_size)
+        scores = self.pool_score(h_win).squeeze(-1)
+        mask_win = attention_mask.view(b, W, self.stride).bool()
+
+        valid_window = mask_win.any(dim=2)
+        scores = scores.masked_fill(~mask_win, torch.finfo(scores.dtype).min)
+        weights = torch.softmax(scores / self.pool_temperature, dim=2)
+        weights = weights.masked_fill(~mask_win, 0.0)
+
+        h_pool = (weights.unsqueeze(-1).to(h_win.dtype) * h_win).sum(dim=2)
+        return h_pool, valid_window.to(h_normed.dtype)
+
+    def get_new_info_for_local_memory(
+        self,
+        raw_info: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        h_normed = self.backbone_decoder.raw_decoder.input_layernorm(raw_info)
+        h_pool, mask = self._pool_attention_windows(h_normed, attention_mask)
+        return self._delta_from_normed_tokens(h_pool, mask)
+
+
+class TopKKeyNormMetisHyperMemory(KeyNormTokenAggMetisHyperMemory):
+    """Learn token scores, then write the top ceil(L / stride_interval) tokens."""
+
+    DEFAULT_STRIDE: int = 8
+
+    def __init__(self, config) -> None:
+        super().__init__(config)
+        self.stride = int(config.memory_configs.get("stride_interval", self.DEFAULT_STRIDE))
+        if self.stride <= 0:
+            raise ValueError(f"stride_interval must be > 0, got {self.stride}")
+        self.pool_score = nn.Linear(self.text_cfg.hidden_size, 1, bias=False)
+        self.pool_temperature = float(config.memory_configs.get("pool_temperature", 1.0))
+        if self.pool_temperature <= 0:
+            raise ValueError(f"pool_temperature must be > 0, got {self.pool_temperature}")
+
+    def _select_topk(
+        self,
+        raw_info: torch.Tensor,
+        attention_mask: torch.Tensor | None,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        b, S, hidden_size = raw_info.shape
+        device = raw_info.device
+        dtype = raw_info.dtype
+        h_normed = self.backbone_decoder.raw_decoder.input_layernorm(raw_info)
+        scores = self.pool_score(h_normed).squeeze(-1)
+        if attention_mask is not None:
+            lengths_t = attention_mask.sum(dim=1).long().clamp(min=1)
+            scores = scores.masked_fill(attention_mask == 0, torch.finfo(scores.dtype).min)
+        else:
+            lengths_t = torch.full((b,), S, dtype=torch.long, device=device)
+
+        k_per_sample = torch.div(lengths_t + self.stride - 1, self.stride, rounding_mode="floor")
+        k_max = int(k_per_sample.max().item())
+        selected = []
+        selected_scores = []
+        masks = []
+        for bi in range(b):
+            k = int(k_per_sample[bi].item())
+            idx = torch.topk(scores[bi], k=k, dim=0).indices.sort().values
+            h_sel = raw_info[bi].index_select(0, idx)
+            score_sel = scores[bi].index_select(0, idx)
+            pad_n = k_max - k
+            if pad_n > 0:
+                h_pad = torch.zeros(pad_n, hidden_size, device=device, dtype=dtype)
+                s_pad = torch.zeros(pad_n, device=device, dtype=scores.dtype)
+                h_sel = torch.cat([h_sel, h_pad], dim=0)
+                score_sel = torch.cat([score_sel, s_pad], dim=0)
+            selected.append(h_sel)
+            selected_scores.append(score_sel)
+            masks.append([1.0] * k + [0.0] * pad_n)
+
+        return (
+            torch.stack(selected, dim=0),
+            torch.tensor(masks, device=device, dtype=dtype),
+            torch.stack(selected_scores, dim=0),
+        )
+
+    def get_new_info_for_local_memory(
+        self,
+        raw_info: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        h_sel, mask, _scores = self._select_topk(raw_info, attention_mask)
+        return self._delta_from_raw_tokens(h_sel, mask)
+
+
+class SoftTopKKeyNormMetisHyperMemory(TopKKeyNormMetisHyperMemory):
+    """All-token soft select used for both training and inference.
+
+    This is the no-hard-selection counterpart of top-k writes:
+
+        soft = softmax(scores / tau) * K
+        state = sum_i soft_i * outer(k_i, v_i) / K
+
+    Forward and backward are both soft.  The total gate mass is K, so after the
+    final divide-by-K this is a convex weighted sum of per-token outer products
+    and stays on the same scale as K hard selected tokens averaged by K.
+    """
+
+    def get_new_info_for_local_memory(
+        self,
+        raw_info: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        b, S, _hidden_size = raw_info.shape
+        device = raw_info.device
+        h_normed = self.backbone_decoder.raw_decoder.input_layernorm(raw_info)
+        scores = self.pool_score(h_normed).squeeze(-1)
+
+        if attention_mask is not None:
+            lengths_t = attention_mask.sum(dim=1).long().clamp(min=1)
+            valid_mask = attention_mask.bool()
+            scores = scores.masked_fill(~valid_mask, torch.finfo(scores.dtype).min)
+        else:
+            lengths_t = torch.full((b,), S, dtype=torch.long, device=device)
+            valid_mask = torch.ones(b, S, device=device, dtype=torch.bool)
+
+        k_per_sample = torch.div(lengths_t + self.stride - 1, self.stride, rounding_mode="floor")
+        gate = torch.softmax(scores / self.pool_temperature, dim=1)
+        gate = gate.masked_fill(~valid_mask, 0.0)
+        gate = gate * k_per_sample.to(gate.dtype).unsqueeze(1)
+
+        write_key = F.normalize(self.W_k(h_normed), dim=-1) / (self.kv_dim ** 0.5)
+        write_value = self.W_v(h_normed)
+        gate_exp = gate.unsqueeze(-1).to(write_key.dtype)
+        write_key = write_key * gate_exp
+
+        delta_state = torch.matmul(write_key.transpose(-1, -2), write_value)
+        ones = torch.ones(b, S, 1, device=device, dtype=write_key.dtype)
+        delta_key_state = torch.matmul(write_key.transpose(-1, -2), ones)
+        scale = k_per_sample.clamp(min=1).to(write_key.dtype).view(-1, 1, 1)
+        return delta_state / scale, delta_key_state / scale
+
+
+class StraightThroughTopKKeyNormMetisHyperMemory(TopKKeyNormMetisHyperMemory):
+    """Hard top-k forward with softmax surrogate gradients for the scorer."""
+
+    def get_new_info_for_local_memory(
+        self,
+        raw_info: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        b, S, _hidden_size = raw_info.shape
+        device = raw_info.device
+        h_normed = self.backbone_decoder.raw_decoder.input_layernorm(raw_info)
+        scores = self.pool_score(h_normed).squeeze(-1)
+
+        if attention_mask is not None:
+            lengths_t = attention_mask.sum(dim=1).long().clamp(min=1)
+            valid_mask = attention_mask.bool()
+            scores = scores.masked_fill(~valid_mask, torch.finfo(scores.dtype).min)
+        else:
+            lengths_t = torch.full((b,), S, dtype=torch.long, device=device)
+            valid_mask = torch.ones(b, S, device=device, dtype=torch.bool)
+
+        k_per_sample = torch.div(lengths_t + self.stride - 1, self.stride, rounding_mode="floor")
+        soft = torch.softmax(scores / self.pool_temperature, dim=1) * k_per_sample.to(scores.dtype).unsqueeze(1)
+        soft = soft.masked_fill(~valid_mask, 0.0)
+
+        hard = torch.zeros_like(scores)
+        for bi in range(b):
+            k = int(k_per_sample[bi].item())
+            idx = torch.topk(scores[bi], k=k, dim=0).indices
+            hard[bi].scatter_(0, idx, 1.0)
+
+        gate = hard.detach() - soft.detach() + soft
+        write_key = F.normalize(self.W_k(h_normed), dim=-1) / (self.kv_dim ** 0.5)
+        write_value = self.W_v(h_normed)
+        gate_exp = gate.unsqueeze(-1).to(write_key.dtype)
+        write_key = write_key * gate_exp
+
+        delta_state = torch.matmul(write_key.transpose(-1, -2), write_value)
+        ones = torch.ones(b, S, 1, device=device, dtype=write_key.dtype)
+        delta_key_state = torch.matmul(write_key.transpose(-1, -2), ones)
+        scale = k_per_sample.clamp(min=1).to(write_key.dtype).view(-1, 1, 1)
+        return delta_state / scale, delta_key_state / scale
+
+
+class GumbelTopKKeyNormMetisHyperMemory(TopKKeyNormMetisHyperMemory):
+    """Continuous Gumbel-TopK approximation that writes K soft-selected tokens."""
+
+    def __init__(self, config) -> None:
+        super().__init__(config)
+        self.gumbel_topk_noise = bool(config.memory_configs.get("gumbel_topk_noise", True))
+        self.gumbel_eps = float(config.memory_configs.get("gumbel_eps", 1e-6))
+
+    def _sample_gumbel(self, scores: torch.Tensor) -> torch.Tensor:
+        uniform = torch.rand_like(scores).clamp_(self.gumbel_eps, 1.0 - self.gumbel_eps)
+        return -torch.log(-torch.log(uniform))
+
+    def get_new_info_for_local_memory(
+        self,
+        raw_info: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        b, S, _hidden_size = raw_info.shape
+        device = raw_info.device
+        h_normed = self.backbone_decoder.raw_decoder.input_layernorm(raw_info)
+        scores = self.pool_score(h_normed).squeeze(-1)
+
+        if attention_mask is not None:
+            lengths_t = attention_mask.sum(dim=1).long().clamp(min=1)
+            valid_mask = attention_mask.bool()
+            scores = scores.masked_fill(~valid_mask, torch.finfo(scores.dtype).min)
+        else:
+            lengths_t = torch.full((b,), S, dtype=torch.long, device=device)
+            valid_mask = torch.ones(b, S, device=device, dtype=torch.bool)
+
+        k_per_sample = torch.div(lengths_t + self.stride - 1, self.stride, rounding_mode="floor")
+        k_max = int(k_per_sample.max().item())
+        logits = scores
+        if self.training and self.gumbel_topk_noise:
+            logits = logits + self._sample_gumbel(scores)
+
+        remaining = valid_mask.to(scores.dtype)
+        selections = []
+        for _ in range(k_max):
+            masked_logits = logits + torch.log(remaining.clamp(min=self.gumbel_eps))
+            weights = torch.softmax(masked_logits / self.pool_temperature, dim=1)
+            weights = weights.masked_fill(~valid_mask, 0.0)
+            weights = weights / weights.sum(dim=1, keepdim=True).clamp(min=self.gumbel_eps)
+            selections.append(weights)
+            remaining = remaining * (1.0 - weights).clamp(min=0.0)
+
+        selection = torch.stack(selections, dim=1)
+        row_mask = (
+            torch.arange(k_max, device=device).unsqueeze(0)
+            < k_per_sample.unsqueeze(1)
+        ).to(h_normed.dtype)
+        h_pool = torch.matmul(selection.to(h_normed.dtype), h_normed)
+        return self._delta_from_normed_tokens(h_pool, row_mask)
+
+
+class AlphaTopPKeyNormMetisHyperMemory(TopKKeyNormMetisHyperMemory):
+    """Adaptive top-p/nucleus token selection with selected soft weights.
+
+    Select the smallest set whose scorer probability mass reaches
+    ``alpha_top_p``.  The selected tokens are written as a convex weighted sum
+    of per-token outer products:
+
+        weights_i = p_i / sum_{j in S_alpha} p_j
+        state = sum_{i in S_alpha} weights_i * outer(k_i, v_i)
+    """
+
+    def __init__(self, config) -> None:
+        super().__init__(config)
+        self.alpha_top_p = float(config.memory_configs.get("alpha_top_p", 0.9))
+        if not 0.0 < self.alpha_top_p <= 1.0:
+            raise ValueError(f"alpha_top_p must be in (0, 1], got {self.alpha_top_p}")
+        self.alpha_min_tokens = int(config.memory_configs.get("alpha_min_tokens", 1))
+        if self.alpha_min_tokens <= 0:
+            raise ValueError(f"alpha_min_tokens must be > 0, got {self.alpha_min_tokens}")
+        self.alpha_max_tokens = int(config.memory_configs.get("alpha_max_tokens", 0))
+        if self.alpha_max_tokens < 0:
+            raise ValueError(f"alpha_max_tokens must be >= 0, got {self.alpha_max_tokens}")
+        self.alpha_max_fraction = float(config.memory_configs.get("alpha_max_fraction", 0.0))
+        if not 0.0 <= self.alpha_max_fraction <= 1.0:
+            raise ValueError(f"alpha_max_fraction must be in [0, 1], got {self.alpha_max_fraction}")
+        self.last_alpha_stats: dict[str, float] = {}
+
+    def _alpha_top_p_mask(
+        self,
+        probs: torch.Tensor,
+        valid_mask: torch.Tensor,
+        lengths_t: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        b, S = probs.shape
+        sorted_probs, sorted_idx = torch.sort(probs, descending=True, dim=1)
+        cum = sorted_probs.cumsum(dim=1)
+        k_raw = (cum <= self.alpha_top_p).sum(dim=1) + 1
+        k_raw = torch.minimum(k_raw, lengths_t)
+
+        k_min = torch.minimum(
+            torch.full_like(lengths_t, self.alpha_min_tokens),
+            lengths_t,
+        )
+        k_max = lengths_t.clone()
+        if self.alpha_max_fraction > 0.0:
+            frac_cap = torch.ceil(lengths_t.to(probs.dtype) * self.alpha_max_fraction).long()
+            k_max = torch.minimum(k_max, frac_cap.clamp(min=1))
+        if self.alpha_max_tokens > 0:
+            fixed_cap = torch.full_like(lengths_t, self.alpha_max_tokens)
+            k_max = torch.minimum(k_max, fixed_cap.clamp(min=1))
+        k_max = torch.maximum(k_max, k_min)
+        k_alpha = torch.minimum(torch.maximum(k_raw, k_min), k_max)
+
+        rank = torch.arange(S, device=probs.device).unsqueeze(0)
+        keep_sorted = rank < k_alpha.unsqueeze(1)
+        hard = torch.zeros_like(probs)
+        hard.scatter_(1, sorted_idx, keep_sorted.to(probs.dtype))
+        hard = hard.masked_fill(~valid_mask, 0.0)
+
+        selected_mass = (probs * hard).sum(dim=1, keepdim=True).clamp(min=1e-6)
+        return hard, k_alpha, selected_mass
+
+    def _record_alpha_stats(
+        self,
+        probs: torch.Tensor,
+        hard: torch.Tensor,
+        k_alpha: torch.Tensor,
+        selected_mass: torch.Tensor,
+        lengths_t: torch.Tensor,
+    ) -> None:
+        with torch.no_grad():
+            probs_f = probs.detach().float()
+            k_f = k_alpha.detach().float()
+            lengths_f = lengths_t.detach().float().clamp(min=1)
+            entropy = -(probs_f * probs_f.clamp(min=1e-12).log()).sum(dim=1)
+            self.last_alpha_stats = {
+                "k_mean": float(k_f.mean().item()),
+                "k_min": float(k_f.min().item()),
+                "k_max": float(k_f.max().item()),
+                "k_ratio": float((k_f / lengths_f).mean().item()),
+                "score_entropy": float(entropy.mean().item()),
+                "p_max": float(probs_f.max(dim=1).values.mean().item()),
+                "selected_mass": float(selected_mass.detach().float().mean().item()),
+                "alpha": self.alpha_top_p,
+            }
+
+    def _alpha_weights(
+        self,
+        scores: torch.Tensor,
+        valid_mask: torch.Tensor,
+        lengths_t: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        probs = torch.softmax(scores / self.pool_temperature, dim=1)
+        probs = probs.masked_fill(~valid_mask, 0.0)
+        probs = probs / probs.sum(dim=1, keepdim=True).clamp(min=1e-6)
+        hard, k_alpha, selected_mass = self._alpha_top_p_mask(probs, valid_mask, lengths_t)
+        self._record_alpha_stats(probs, hard, k_alpha, selected_mass, lengths_t)
+        weights = probs * hard / selected_mass
+        return weights, probs
+
+    def _delta_from_weights(
+        self,
+        h_normed: torch.Tensor,
+        weights: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        write_key = F.normalize(self.W_k(h_normed), dim=-1) / (self.kv_dim ** 0.5)
+        write_value = self.W_v(h_normed)
+        weight_exp = weights.unsqueeze(-1).to(write_key.dtype)
+        write_key = write_key * weight_exp
+        delta_state = torch.matmul(write_key.transpose(-1, -2), write_value)
+        ones = torch.ones(
+            write_key.size(0),
+            write_key.size(1),
+            1,
+            device=write_key.device,
+            dtype=write_key.dtype,
+        )
+        delta_key_state = torch.matmul(write_key.transpose(-1, -2), ones)
+        return delta_state, delta_key_state
+
+    def get_new_info_for_local_memory(
+        self,
+        raw_info: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        b, S, _hidden_size = raw_info.shape
+        device = raw_info.device
+        h_normed = self.backbone_decoder.raw_decoder.input_layernorm(raw_info)
+        scores = self.pool_score(h_normed).squeeze(-1)
+        if attention_mask is not None:
+            lengths_t = attention_mask.sum(dim=1).long().clamp(min=1)
+            valid_mask = attention_mask.bool()
+            scores = scores.masked_fill(~valid_mask, torch.finfo(scores.dtype).min)
+        else:
+            lengths_t = torch.full((b,), S, dtype=torch.long, device=device)
+            valid_mask = torch.ones(b, S, device=device, dtype=torch.bool)
+
+        weights, _probs = self._alpha_weights(scores, valid_mask, lengths_t)
+        return self._delta_from_weights(h_normed, weights)
+
+
+class StraightThroughAlphaTopPKeyNormMetisHyperMemory(AlphaTopPKeyNormMetisHyperMemory):
+    """Alpha top-p forward with full-softmax surrogate gradients."""
+
+    def get_new_info_for_local_memory(
+        self,
+        raw_info: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        b, S, _hidden_size = raw_info.shape
+        device = raw_info.device
+        h_normed = self.backbone_decoder.raw_decoder.input_layernorm(raw_info)
+        scores = self.pool_score(h_normed).squeeze(-1)
+        if attention_mask is not None:
+            lengths_t = attention_mask.sum(dim=1).long().clamp(min=1)
+            valid_mask = attention_mask.bool()
+            scores = scores.masked_fill(~valid_mask, torch.finfo(scores.dtype).min)
+        else:
+            lengths_t = torch.full((b,), S, dtype=torch.long, device=device)
+            valid_mask = torch.ones(b, S, device=device, dtype=torch.bool)
+
+        hard_weights, probs = self._alpha_weights(scores, valid_mask, lengths_t)
+        soft_weights = probs.masked_fill(~valid_mask, 0.0)
+        soft_weights = soft_weights / soft_weights.sum(dim=1, keepdim=True).clamp(min=1e-6)
+        weights = hard_weights.detach() - soft_weights.detach() + soft_weights
+        return self._delta_from_weights(h_normed, weights)
+
+
+class GatedDeltaRuleMixin:
+    """Mixin implementing the gated delta rule memory recurrence.
+
+    The paper formula is written for column-vector reads:
+
+        S_t = S_{t-1}(alpha_t (I - beta_t k_t k_t^T)) + beta_t v_t k_t^T
+
+    Metis stores row-vector memories read as ``q @ M``.  The equivalent
+    single-token update is:
+
+        M_t = alpha_t (I - beta_t k_t k_t^T) M_{t-1} + beta_t k_t v_t^T
+
+    For a selected token set, this implementation applies the batched parallel
+    approximation ``sum_t beta_t k_t k_t^T`` / ``sum_t beta_t k_t v_t^T`` in one
+    write.  The same erase/write rule is applied to ``key_state`` when the
+    paired local memory keeps one for key-normalized reads.
+    """
+
+    @staticmethod
+    def _logit_clamped(value: float) -> float:
+        eps = 1e-4
+        p = min(max(float(value), eps), 1.0 - eps)
+        return math.log(p / (1.0 - p))
+
+    def _init_gated_delta_rule(self) -> None:
+        hidden_size = self.text_cfg.hidden_size
+        self.gated_delta_alpha = nn.Linear(hidden_size, 1, bias=True)
+        self.gated_delta_beta = nn.Linear(hidden_size, 1, bias=True)
+
+        nn.init.zeros_(self.gated_delta_alpha.weight)
+        nn.init.zeros_(self.gated_delta_beta.weight)
+
+        alpha_init = self.config.memory_configs.get("gated_delta_alpha_init", 1.0)
+        beta_init = self.config.memory_configs.get("gated_delta_beta_init", 1.0)
+        nn.init.constant_(self.gated_delta_alpha.bias, self._logit_clamped(alpha_init))
+        nn.init.constant_(self.gated_delta_beta.bias, self._logit_clamped(beta_init))
+
+    def _alpha_top_p_normed_weights(
+        self,
+        raw_info: torch.Tensor,
+        attention_mask: torch.Tensor | None,
+        straight_through: bool,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        b, S, _hidden_size = raw_info.shape
+        device = raw_info.device
+        h_normed = self.backbone_decoder.raw_decoder.input_layernorm(raw_info)
+        scores = self.pool_score(h_normed).squeeze(-1)
+        if attention_mask is not None:
+            lengths_t = attention_mask.sum(dim=1).long().clamp(min=1)
+            valid_mask = attention_mask.bool()
+            scores = scores.masked_fill(~valid_mask, torch.finfo(scores.dtype).min)
+        else:
+            lengths_t = torch.full((b,), S, dtype=torch.long, device=device)
+            valid_mask = torch.ones(b, S, device=device, dtype=torch.bool)
+
+        hard_weights, probs = self._alpha_weights(scores, valid_mask, lengths_t)
+        if not straight_through:
+            return h_normed, hard_weights
+
+        soft_weights = probs.masked_fill(~valid_mask, 0.0)
+        soft_weights = soft_weights / soft_weights.sum(dim=1, keepdim=True).clamp(min=1e-6)
+        weights = hard_weights.detach() - soft_weights.detach() + soft_weights
+        return h_normed, weights
+
+    def _apply_gated_delta_rule_update(
+        self,
+        h_normed: torch.Tensor,
+        weights: torch.Tensor,
+        local_memory,
+    ) -> None:
+        write_key = F.normalize(self.W_k(h_normed), dim=-1) / (self.kv_dim ** 0.5)
+        write_value = self.W_v(h_normed)
+
+        weights = weights.to(write_key.dtype)
+        weight_mass = weights.sum(dim=1, keepdim=True).clamp(min=1e-6)
+        alpha_gate = torch.sigmoid(self.gated_delta_alpha(h_normed).squeeze(-1))
+        beta_gate = torch.sigmoid(self.gated_delta_beta(h_normed).squeeze(-1))
+        alpha = (weights * alpha_gate).sum(dim=1) / weight_mass.squeeze(1)
+        beta = weights * (self.update_ratio * beta_gate)
+        beta_exp = beta.unsqueeze(-1)
+
+        bsz = write_key.size(0)
+        state = local_memory.state
+        if state is None:
+            state = torch.zeros(
+                bsz, self.kv_dim, self.kv_dim,
+                device=write_key.device,
+                dtype=write_key.dtype,
+            )
+
+        key_state = getattr(local_memory, "key_state", None)
+
+        # M_t = alpha * (M - K^T beta (K M)) + K^T beta V
+        key_memory = torch.matmul(write_key, state)
+        erase_state = torch.matmul(write_key.transpose(-1, -2), beta_exp * key_memory)
+        add_state = torch.matmul(write_key.transpose(-1, -2), beta_exp * write_value)
+        new_state = alpha.view(bsz, 1, 1) * (state - erase_state) + add_state
+
+        has_key_state = hasattr(local_memory, "key_state")
+        if key_state is None and has_key_state:
+            key_state = torch.zeros(
+                bsz, self.kv_dim, 1,
+                device=write_key.device,
+                dtype=write_key.dtype,
+            )
+
+        if key_state is not None:
+            key_memory_mass = torch.matmul(write_key, key_state)
+            erase_key_state = torch.matmul(
+                write_key.transpose(-1, -2),
+                beta_exp * key_memory_mass,
+            )
+            add_key_state = torch.matmul(write_key.transpose(-1, -2), beta_exp)
+            new_key_state = (
+                alpha.view(bsz, 1, 1) * (key_state - erase_key_state)
+                + add_key_state
+            )
+            local_memory.write(new_state, new_key_state)
+        else:
+            local_memory.write(new_state)
+
+
+class AlphaTopPGatedDeltaRuleMetisHyperMemory(
+    GatedDeltaRuleMixin,
+    AlphaTopPKeyNormMetisHyperMemory,
+):
+    """AlphaTopP token selection with gated-delta local-memory writes."""
+
+    def __init__(self, config) -> None:
+        super().__init__(config)
+        self._init_gated_delta_rule()
+
+    def update_local_memory(
+        self,
+        raw_info: torch.Tensor,
+        local_memory,
+        attention_mask: torch.Tensor | None = None,
+    ) -> None:
+        h_normed, weights = self._alpha_top_p_normed_weights(
+            raw_info, attention_mask, straight_through=False,
+        )
+        self._apply_gated_delta_rule_update(h_normed, weights, local_memory)
+
+
+class StraightThroughAlphaTopPGatedDeltaRuleMetisHyperMemory(
+    AlphaTopPGatedDeltaRuleMetisHyperMemory,
+):
+    """Straight-through AlphaTopP selection with gated-delta writes."""
+
+    def update_local_memory(
+        self,
+        raw_info: torch.Tensor,
+        local_memory,
+        attention_mask: torch.Tensor | None = None,
+    ) -> None:
+        h_normed, weights = self._alpha_top_p_normed_weights(
+            raw_info, attention_mask, straight_through=True,
+        )
+        self._apply_gated_delta_rule_update(h_normed, weights, local_memory)
+
+
+class WeightedTopKKeyNormMetisHyperMemory(TopKKeyNormMetisHyperMemory):
+    """Top-k token write with learned softmax weights over selected tokens."""
+
+    def get_new_info_for_local_memory(
+        self,
+        raw_info: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        h_sel, mask, scores = self._select_topk(raw_info, attention_mask)
+        h_normed = self.backbone_decoder.raw_decoder.input_layernorm(h_sel)
+        write_key = F.normalize(self.W_k(h_normed), dim=-1) / (self.kv_dim ** 0.5)
+        write_value = self.W_v(h_normed)
+
+        scores = scores.masked_fill(mask == 0, torch.finfo(scores.dtype).min)
+        weights = torch.softmax(scores, dim=1).unsqueeze(-1).to(write_key.dtype)
+        mask_exp = mask.unsqueeze(-1).to(write_key.dtype)
+        write_key = write_key * mask_exp
+        write_value = write_value * weights * mask_exp
+
+        delta_state = torch.matmul(write_key.transpose(-1, -2), write_value)
+        delta_key_state = torch.matmul(write_key.transpose(-1, -2), weights * mask_exp)
+        return delta_state, delta_key_state
+
+
+class Conv1dPoolKeyNormMetisHyperMemory(StridePoolKeyNormMetisHyperMemory):
+    """Depthwise conv1d pooling with kernel=stride=stride_interval."""
+
+    def __init__(self, config) -> None:
+        super().__init__(config)
+        hidden_size = self.text_cfg.hidden_size
+        self.pool_conv = nn.Conv1d(
+            hidden_size,
+            hidden_size,
+            kernel_size=self.stride,
+            stride=self.stride,
+            groups=hidden_size,
+            bias=False,
+        )
+        nn.init.constant_(self.pool_conv.weight, 1.0 / self.stride)
+
+    def get_new_info_for_local_memory(
+        self,
+        raw_info: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        b, S, hidden_size = raw_info.shape
+        pad_n = (-S) % self.stride
+        if pad_n > 0:
+            raw_info = F.pad(raw_info, (0, 0, 0, pad_n))
+            if attention_mask is not None:
+                attention_mask = F.pad(attention_mask, (0, pad_n))
+
+        if attention_mask is not None:
+            mask_exp = attention_mask.unsqueeze(-1).to(raw_info.dtype)
+            raw_info = raw_info * mask_exp
+            denom = F.avg_pool1d(
+                attention_mask.unsqueeze(1).to(raw_info.dtype),
+                kernel_size=self.stride,
+                stride=self.stride,
+                count_include_pad=False,
+            ).squeeze(1) * self.stride
+        else:
+            denom = torch.full(
+                (b, raw_info.size(1) // self.stride),
+                self.stride,
+                device=raw_info.device,
+                dtype=raw_info.dtype,
+            )
+
+        h_pool = self.pool_conv(raw_info.transpose(1, 2)).transpose(1, 2)
+        h_pool = h_pool * (self.stride / denom.clamp(min=1).unsqueeze(-1))
+        mask = (denom > 0).to(raw_info.dtype)
+        return self._delta_from_raw_tokens(h_pool, mask)
+
+
+class MixedKeyNormMetisHyperMemory(StridePoolKeyNormMetisHyperMemory):
+    """Stride-window pooled writes mixed with one global mean-pooled write token.
+
+    This uses the same stride-window pooling as ``StridePoolKeyNormMetisHyperMemory``
+    and appends one all-sequence mean token:
+
+        tokens = [mean(h[0:K]), mean(h[K:2K]), ..., mean(h[0:L])]
+    """
+
+    def get_new_info_for_local_memory(
+        self,
+        raw_info: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        h_stride, stride_mask = self._pool_stride_windows(raw_info, attention_mask)
+
+        if attention_mask is not None:
+            mask = attention_mask.unsqueeze(-1).to(raw_info.dtype)
+            lengths = attention_mask.sum(dim=1).clamp(min=1).to(raw_info.dtype)
+            h_mean = (raw_info * mask).sum(dim=1, keepdim=True) / lengths.view(-1, 1, 1)
+        else:
+            h_mean = raw_info.mean(dim=1, keepdim=True)
+
+        h_mix = torch.cat([h_stride, h_mean], dim=1)
+        mean_mask = torch.ones(
+            stride_mask.size(0),
+            1,
+            device=stride_mask.device,
+            dtype=stride_mask.dtype,
+        )
+        mix_mask = torch.cat([stride_mask, mean_mask], dim=1)
+        return self._delta_from_raw_tokens(h_mix, mix_mask)
+
+
+class FullTokensKeyNormMetisHyperMemory(LinearLastMetisHyperMemory):
+    """Full-token write path for key-normalized DeltaNet memory.
+
+    This mirrors the recent metis_modular normalization scheme while keeping it
+    opt-in as a separate dev_beta class:
+
+        h_normed    = input_layernorm(h)
+        write_key   = normalize(W_k(h_normed)) / sqrt(D)
+        write_value = W_v(h_normed)
+        state       = mean_t(write_key_t^T @ write_value_t)
+        key_state   = mean_t(write_key_t)
+
+    ``NormalizedDeltaNetMetisLocalMemory`` uses ``key_state`` at read time to
+    divide memory outputs by ``q @ key_state + 1``.
+    """
+
+    def get_new_info_for_local_memory(
+        self,
+        raw_info: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        h = self.backbone_decoder.raw_decoder.input_layernorm(raw_info)
+        write_key = F.normalize(self.W_k(h), dim=-1) / (self.kv_dim ** 0.5)
+        write_value = self.W_v(h)
+
+        if attention_mask is not None:
+            mask = attention_mask.unsqueeze(-1).to(write_key.dtype)
+            lengths = attention_mask.sum(dim=1).clamp(min=1).to(write_key.dtype)
+            write_key = write_key * mask
+            write_value = write_value * mask
+        else:
+            lengths = torch.full(
+                (raw_info.size(0),),
+                raw_info.size(1),
+                dtype=write_key.dtype,
+                device=raw_info.device,
+            ).clamp(min=1)
+
+        delta_state = torch.matmul(write_key.transpose(-1, -2), write_value)
+        ones = torch.ones(
+            write_key.size(0),
+            write_key.size(1),
+            1,
+            device=write_key.device,
+            dtype=write_key.dtype,
+        )
+        delta_key_state = torch.matmul(write_key.transpose(-1, -2), ones)
+        scale = lengths.view(-1, 1, 1)
+        return delta_state / scale, delta_key_state / scale
+
+    def update_local_memory(
+        self,
+        raw_info: torch.Tensor,
+        local_memory,
+        attention_mask: torch.Tensor | None = None,
+    ) -> None:
+        delta_state, delta_key_state = self.get_new_info_for_local_memory(raw_info, attention_mask)
+        if getattr(local_memory, "key_state", None) is None:
+            new_state = self.update_ratio * delta_state
+            new_key_state = self.update_ratio * delta_key_state
+        else:
+            new_state = (1.0 - self.update_ratio) * local_memory.state + self.update_ratio * delta_state
+            new_key_state = (
+                (1.0 - self.update_ratio) * local_memory.key_state
+                + self.update_ratio * delta_key_state
+            )
+        local_memory.write(new_state, new_key_state)
+
+
+class StrideKeyNormMetisHyperMemory(StrideNormalizedMetisHyperMemory):
+    """Stride-token write path for key-normalized DeltaNet memory.
+
+    This is the stride-sampled counterpart of
+    ``FullTokensKeyNormMetisHyperMemory``:
+
+        h_sel       = input_layernorm(h[stride_indices])
+        write_key   = normalize(W_k(h_sel)) / sqrt(D)
+        write_value = W_v(h_sel)
+        state       = mean_selected(write_key_t^T @ write_value_t)
+        key_state   = mean_selected(write_key_t)
+
+    It should be paired with ``NormalizedDeltaNetMetisLocalMemory`` so reads
+    can use ``key_state`` for the q @ key_state + 1 normalization factor.
+    """
+
+    def get_new_info_for_local_memory(
+        self,
+        raw_info: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        h_sel, mask = self._select_tokens_with_mask(raw_info, attention_mask)
+        lengths = mask.sum(dim=1).clamp(min=1).to(raw_info.dtype)
+
+        h_sel = self.backbone_decoder.raw_decoder.input_layernorm(h_sel)
+        write_key = F.normalize(self.W_k(h_sel), dim=-1) / (self.kv_dim ** 0.5)
+        write_value = self.W_v(h_sel)
+
+        mask_exp = mask.unsqueeze(-1).to(write_key.dtype)
+        write_key = write_key * mask_exp
+        write_value = write_value * mask_exp
+
+        delta_state = torch.matmul(write_key.transpose(-1, -2), write_value)
+        ones = torch.ones(
+            write_key.size(0),
+            write_key.size(1),
+            1,
+            device=write_key.device,
+            dtype=write_key.dtype,
+        )
+        delta_key_state = torch.matmul(write_key.transpose(-1, -2), ones)
+        scale = lengths.to(delta_state.dtype).view(-1, 1, 1)
+        return delta_state / scale, delta_key_state / scale
+
+    def update_local_memory(
+        self,
+        raw_info: torch.Tensor,
+        local_memory,
+        attention_mask: torch.Tensor | None = None,
+    ) -> None:
+        delta_state, delta_key_state = self.get_new_info_for_local_memory(raw_info, attention_mask)
+        if getattr(local_memory, "key_state", None) is None:
+            new_state = self.update_ratio * delta_state
+            new_key_state = self.update_ratio * delta_key_state
+        else:
+            new_state = (1.0 - self.update_ratio) * local_memory.state + self.update_ratio * delta_state
+            new_key_state = (
+                (1.0 - self.update_ratio) * local_memory.key_state
+                + self.update_ratio * delta_key_state
+            )
+        local_memory.write(new_state, new_key_state)
+
+
+class StrideKernelKeyNormMetisHyperMemory(StrideKeyNormMetisHyperMemory):
+    """Stride keynorm write path with a kernel feature map on write keys.
+
+    Paired with ``KernelizedDeltaNetMetisLocalMemory``.  The same q/k feature
+    map should be used on both sides:
+
+        phi(k)     = kernel(W_k(input_layernorm(h_sel)))
+        state      = mean_selected(phi(k)_t^T @ v_t)
+        key_state  = mean_selected(phi(k)_t)
+
+    Configurable via ``memory_configs['qk_kernel_type']``:
+        - ``elu_plus_one`` (default)
+        - ``relu_square``
+        - ``softplus``
+    """
+
+    def __init__(self, config) -> None:
+        super().__init__(config)
+        self.qk_kernel_type = config.memory_configs.get("qk_kernel_type", "elu_plus_one")
+
+    def get_new_info_for_local_memory(
+        self,
+        raw_info: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        h_sel, mask = self._select_tokens_with_mask(raw_info, attention_mask)
+        lengths = mask.sum(dim=1).clamp(min=1).to(raw_info.dtype)
+
+        h_sel = self.backbone_decoder.raw_decoder.input_layernorm(h_sel)
+        write_key = _qk_kernel(self.W_k(h_sel), self.qk_kernel_type)
+        write_value = self.W_v(h_sel)
+
+        mask_exp = mask.unsqueeze(-1).to(write_key.dtype)
+        write_key = write_key * mask_exp
+        write_value = write_value * mask_exp
+
+        delta_state = torch.matmul(write_key.transpose(-1, -2), write_value)
+        ones = torch.ones(
+            write_key.size(0),
+            write_key.size(1),
+            1,
+            device=write_key.device,
+            dtype=write_key.dtype,
+        )
+        delta_key_state = torch.matmul(write_key.transpose(-1, -2), ones)
+        scale = lengths.to(delta_state.dtype).view(-1, 1, 1)
+        return delta_state / scale, delta_key_state / scale
+
+
+class StrideL2NormMetisHyperMemory(StrideNormalizedMetisHyperMemory):
+    """Stride-token write path with L2-normalized keys scaled by sqrt(D).
+
+    Pair this with ``L2NormalizedDeltaNetMetisLocalMemory``.  Unlike
+    ``StrideKeyNormMetisHyperMemory``, this class does not apply a kernel,
+    does not produce ``key_state``, and therefore has no key-state denominator
+    at read time:
+
+        k      = normalize(W_k(input_layernorm(h_sel))) / sqrt(D)
+        q      = normalize(q)  # in the paired local memory
+        state  = mean_selected(k_t^T @ v_t)
+    """
+
+    def get_new_info_for_local_memory(
+        self,
+        raw_info: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        h_sel, mask = self._select_tokens_with_mask(raw_info, attention_mask)
+        lengths = mask.sum(dim=1).clamp(min=1).to(raw_info.dtype)
+
+        h_sel = self.backbone_decoder.raw_decoder.input_layernorm(h_sel)
+        write_key = F.normalize(self.W_k(h_sel), dim=-1) / (self.kv_dim ** 0.5)
+        write_value = self.W_v(h_sel)
+
+        mask_exp = mask.unsqueeze(-1).to(write_key.dtype)
+        write_key = write_key * mask_exp
+        write_value = write_value * mask_exp
+
+        delta = torch.matmul(write_key.transpose(-1, -2), write_value)
+        return delta / lengths.to(delta.dtype).view(-1, 1, 1)
+
+
+class StrideNormalizedv3MetisHyperMemory(StrideNormalizedMetisHyperMemory):
+    """Stride-based memory update scaled by 1 / (L' * sqrt(D)).
+    """
+
+    def get_new_info_for_local_memory(
+        self,
+        raw_info: torch.Tensor,                 # (b, s, hidden_size)
+        attention_mask: torch.Tensor | None = None,
+    ) -> torch.Tensor:                           # (b, kv_dim, kv_dim)
+        h_sel, mask = self._select_tokens_with_mask(raw_info, attention_mask)
+        # mask: (b, N_max)  — 1.0 for valid tokens, 0.0 for padding
+        L_prime = mask.sum(dim=1).clamp(min=1)              # (b,)  actual token count
+        h_sel = self.backbone_decoder.raw_decoder.input_layernorm(h_sel)
+        write_key   = self.W_k(h_sel)                       # (b, N_max, kv_dim)
+        write_value = self.W_v(h_sel)                       # (b, N_max, kv_dim)
+        # Zero-out padded slots so they contribute nothing to the outer product.
+        mask_exp    = mask.unsqueeze(-1)                    # (b, N_max, 1)
+        write_key   = write_key   * mask_exp
+        write_value = write_value * mask_exp
+        # (b, kv_dim, N_max) @ (b, N_max, kv_dim) -> (b, kv_dim, kv_dim)
+        delta = torch.matmul(write_key.transpose(-1, -2), write_value)
+        # Scale per sample: divide by L' * sqrt(D)
+        scale = L_prime * (self.kv_dim ** 0.5)              # (b,)
+        scale = scale.view(-1, 1, 1)                        # (b, 1, 1)  broadcast
+        return delta / scale
+
+class StrideNormalizedv4MetisHyperMemory(StrideNormalizedMetisHyperMemory):
+    """Stride-based memory update scaled by 1 / (L' * sqrt(D)).
+    """
+
+    def get_new_info_for_local_memory(
+        self,
+        raw_info: torch.Tensor,                 # (b, s, hidden_size)
+        attention_mask: torch.Tensor | None = None,
+    ) -> torch.Tensor:                           # (b, kv_dim, kv_dim)
+        h_sel, mask = self._select_tokens_with_mask(raw_info, attention_mask)
+        # mask: (b, N_max)  — 1.0 for valid tokens, 0.0 for padding
+        L_prime = mask.sum(dim=1).clamp(min=1)              # (b,)  actual token count
+        h_sel = self.backbone_decoder.raw_decoder.input_layernorm(h_sel)
+        write_key   = self.W_k(h_sel)                       # (b, N_max, kv_dim)
+        write_value = self.W_v(h_sel)                       # (b, N_max, kv_dim)
+        # Zero-out padded slots so they contribute nothing to the outer product.
+        mask_exp    = mask.unsqueeze(-1)                    # (b, N_max, 1)
+        write_key   = write_key   * mask_exp
+        write_value = write_value * mask_exp
+        # (b, kv_dim, N_max) @ (b, N_max, kv_dim) -> (b, kv_dim, kv_dim)
+        delta = torch.matmul(write_key.transpose(-1, -2), write_value)
+        # Scale per sample: divide by L' * sqrt(D)
+        scale = L_prime * (self.kv_dim)              # (b,)
+        scale = scale.view(-1, 1, 1)                        # (b, 1, 1)  broadcast
+        return delta / scale
+
+
+class StrideNormalizedv5MetisHyperMemory(StrideNormalizedMetisHyperMemory):
+    """Stride-based memory update: L2-normalize write vectors, then divide by L'.
+
+        delta = F.normalize(W_k H, dim=-1).T  @  F.normalize(W_v H, dim=-1)  /  L'
+
+    Differences vs v3 (which divides raw projections by L' * sqrt(D)):
+      - Each token's write_key / write_value is L2-normalised to unit norm before
+        the outer product, so every rank-1 contribution has ||·||_F = 1 exactly.
+      - Dividing by L' averages the L' unit outer products.
+      - Result: ||delta||_F <= 1 always, independent of D, sequence length,
+        and weight magnitudes.
+
+    Note: mask is applied AFTER F.normalize so that batch-padding slots (filled
+    with a copy of position 0) are first given unit norm and then zeroed out.
+    Applying mask before normalize would produce 0/0 for zero vectors.
+    """
+
+    def get_new_info_for_local_memory(
+        self,
+        raw_info: torch.Tensor,                 # (b, s, hidden_size)
+        attention_mask: torch.Tensor | None = None,
+    ) -> torch.Tensor:                           # (b, kv_dim, kv_dim)
+        h_sel, mask = self._select_tokens_with_mask(raw_info, attention_mask)
+        # mask: (b, N_max) — 1.0 for valid tokens, 0.0 for batch-padding slots
+        L_prime = mask.sum(dim=1).clamp(min=1)              # (b,)  actual token count
+        h_sel = self.backbone_decoder.raw_decoder.input_layernorm(h_sel)
+        # L2-normalize each token's projection to unit norm along kv_dim axis
+        write_key   = F.normalize(self.W_k(h_sel), dim=-1)  # (b, N_max, kv_dim), ‖·‖=1
+        write_value = F.normalize(self.W_v(h_sel), dim=-1)  # (b, N_max, kv_dim), ‖·‖=1
+        # Zero-out batch-padding slots after normalization to avoid 0/0 issues
+        mask_exp    = mask.unsqueeze(-1)                    # (b, N_max, 1)
+        write_key   = write_key   * mask_exp
+        write_value = write_value * mask_exp
+        # (b, kv_dim, N_max) @ (b, N_max, kv_dim) -> (b, kv_dim, kv_dim)
+        # ||delta||_F <= L' (sum of L' unit outer products), divide by L' to average
+        delta = torch.matmul(write_key.transpose(-1, -2), write_value)
+        return delta / L_prime.view(-1, 1, 1)               # (b, 1, 1) broadcast

metis_local_memory.py

@@ -0,0 +1,287 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from abc import ABC
+
+def _qk_kernel(x: torch.Tensor, kernel_type: str = "elu_plus_one") -> torch.Tensor:
+    if kernel_type == "elu_plus_one":
+        return F.elu(x) + 1.0
+    if kernel_type == "relu_square":
+        return F.relu(x).square()
+    if kernel_type == "softplus":
+        return F.softplus(x)
+    raise ValueError(f"Unsupported qk kernel type: {kernel_type}")
+
+def create_metis_local_memory(config):
+    return eval(config.memory_configs['metis_local_memory_type'])(config)
+
+class MetisLocalMemoryBase(nn.Module, ABC):
+    def __init__(self, config) -> None:
+        super().__init__()
+        self.config = config
+        # Qwen 3.5 has text config, but Qwen 3 does not.
+        self.text_cfg = getattr(config.backbone_configs, 'text_config', config.backbone_configs)
+
+    def initialize(self) -> None:
+        raise NotImplementedError
+
+    def reset(self) -> None:
+        raise NotImplementedError
+
+    def read(self, query_for_memory):
+        raise NotImplementedError
+
+    def write(self, new_info) -> None:
+        raise NotImplementedError
+
+    @property
+    def state(self):
+        raise NotImplementedError
+
+
+class DeltaNetMetisLocalMemory(MetisLocalMemoryBase):
+    """Linear (DeltaNet-style) memory matrix of shape (b, D, D).
+
+    Read:  output = Q_flat @ M,  where Q_flat = (b, s, D)
+    Write: M = new_state  (forget + additive update computed by HyperMemory)
+    """
+
+    def __init__(self, config) -> None:
+        super().__init__(config)
+
+        num_q_heads = self.text_cfg.num_attention_heads
+        # If num_key_value_heads is not set, use num_attention_heads (MHA).
+        num_kv_heads = getattr(self.text_cfg, "num_key_value_heads", num_q_heads)
+        head_dim = getattr(self.text_cfg, "head_dim", self.text_cfg.hidden_size // num_q_heads)
+
+        self.q_dim = num_q_heads * head_dim
+        self.kv_dim = self._compute_kv_dim(num_q_heads, num_kv_heads, head_dim)
+        self.num_kv_groups = self.q_dim // self.kv_dim
+
+        self._state: torch.Tensor | None = None
+
+    @staticmethod
+    def _compute_kv_dim(num_q_heads: int, num_kv_heads: int, head_dim: int) -> int:
+        """GQA layout: kv_dim = num_kv_heads * head_dim."""
+        return num_kv_heads * head_dim
+
+    def initialize(self) -> None:
+        self._state = None
+
+    def reset(self) -> None:
+        self._state = None
+
+    def _ensure_ready(
+        self,
+        batch_size: int,
+        device: torch.device,
+        dtype: torch.dtype,
+    ) -> None:
+        if self._state is None or self._state.shape[0] != batch_size:
+            self._state = torch.zeros(
+                batch_size, self.kv_dim, self.kv_dim, device=device, dtype=dtype,
+            )
+
+    def read(self, query_for_memory: torch.Tensor) -> torch.Tensor:
+        """Linear memory read: output = Q_flat @ M.
+
+        Args:
+            query_for_memory: (b, h, s, d)
+
+        Returns:
+            (b, s, D) — memory readout, ready to be fused with attention output.
+        """
+        bsz, _h, seq_len, _d = query_for_memory.shape
+        self._ensure_ready(bsz, query_for_memory.device, query_for_memory.dtype)
+        # (b, h, s, d) → (b, s, q_dim)
+        q_flat = query_for_memory.transpose(1, 2).reshape(bsz, seq_len, -1)
+        if self.num_kv_groups > 1:
+            # GQA mode.
+            q_2d = q_flat.view(bsz, seq_len * self.num_kv_groups, self.kv_dim)
+            out_2d = torch.matmul(q_2d, self._state)
+            return out_2d.view(bsz, seq_len, self.q_dim).contiguous()
+        else:
+            # For MHA mode
+            # (b, s, q_dim) @ (b, q_dim, q_dim) → (b, s, q_dim)
+            return torch.matmul(q_flat, self._state).contiguous()
+
+    def write(self, new_state: torch.Tensor) -> None:
+        self._state = new_state  # no detach here, so gradients flow through W_k / W_v
+
+    @property
+    def state(self) -> torch.Tensor | None:
+        return self._state
+
+    @property
+    def is_initialized(self) -> bool:
+        return self._state is not None
+
+    def norm(self) -> float:
+        return self._state.norm().item() if self._state is not None else 0.0
+
+class MHADeltaNetMetisLocalMemory(DeltaNetMetisLocalMemory):
+    """Legacy MHA-style memory: kv_dim = num_q_heads * head_dim (no GQA grouping).
+
+    Memory matrix is (b, q_dim, q_dim) — for Qwen3.5-4B that's 4096×4096.
+    Read collapses to a single MHA matmul: (b, s, q_dim) @ (b, q_dim, q_dim).
+
+    Use this for loading checkpoints trained before the GQA refactor
+    (e.g. experiments/4.17-* and 4.18-*).
+    """
+
+    @staticmethod
+    def _compute_kv_dim(num_q_heads: int, num_kv_heads: int, head_dim: int) -> int:
+        return num_q_heads * head_dim
+
+
+class NormalizedDeltaNetMetisLocalMemory(DeltaNetMetisLocalMemory):
+    """DeltaNet memory with metis_modular-style key normalization.
+
+    Read path:
+        q = normalize(q)
+        y = q @ state
+        y = y / (q @ key_state + 1)
+
+    The paired hyper-memory class ``FullTokensKeyNormMetisHyperMemory`` writes
+    both ``state`` and ``key_state``.
+    """
+
+    def __init__(self, config) -> None:
+        super().__init__(config)
+        self._key_state: torch.Tensor | None = None
+
+    def initialize(self) -> None:
+        self._state = None
+        self._key_state = None
+
+    def reset(self) -> None:
+        self._state = None
+        self._key_state = None
+
+    def _ensure_ready(
+        self,
+        batch_size: int,
+        device: torch.device,
+        dtype: torch.dtype,
+    ) -> None:
+        if self._state is None or self._state.shape[0] != batch_size:
+            self._state = torch.zeros(
+                batch_size, self.kv_dim, self.kv_dim, device=device, dtype=dtype,
+            )
+            self._key_state = torch.zeros(
+                batch_size, self.kv_dim, 1, device=device, dtype=dtype,
+            )
+
+    def read(self, query_for_memory: torch.Tensor) -> torch.Tensor:
+        bsz, _h, seq_len, _d = query_for_memory.shape
+        self._ensure_ready(bsz, query_for_memory.device, query_for_memory.dtype)
+        query_for_memory = F.normalize(query_for_memory, dim=-1)
+        q_flat = query_for_memory.transpose(1, 2).reshape(bsz, seq_len, -1)
+
+        if self.num_kv_groups > 1:
+            q_2d = q_flat.view(bsz, seq_len * self.num_kv_groups, self.kv_dim)
+            out_2d = torch.matmul(q_2d, self._state)
+            # if self._key_state is not None:
+            norm_factor = torch.matmul(q_2d, self._key_state)
+            # print(norm_factor[0])
+            out_2d = out_2d / (norm_factor + 1.0)
+            return out_2d.view(bsz, seq_len, self.q_dim).contiguous()
+
+        out = torch.matmul(q_flat, self._state)
+        # if self._key_state is not None:
+        norm_factor = torch.matmul(q_flat, self._key_state)
+        # print(norm_factor.shape)
+        out = out / (norm_factor + 1.0)
+        return out.contiguous()
+
+    def write(self, new_state: torch.Tensor, key_state: torch.Tensor) -> None:
+        self._state = new_state
+        self._key_state = key_state
+
+    @property
+    def key_state(self) -> torch.Tensor | None:
+        return self._key_state
+
+    @property
+    def is_initialized(self) -> bool:
+        return self._state is not None and self._key_state is not None
+
+
+class KernelizedDeltaNetMetisLocalMemory(NormalizedDeltaNetMetisLocalMemory):
+    """DeltaNet memory read path with a kernel feature map on queries.
+
+    Pair this with ``StrideKernelKeyNormMetisHyperMemory`` so the same feature
+    map is applied to q and k before the key-state normalization:
+
+        phi(q) = kernel(q)
+        y      = phi(q) @ state
+        y      = y / (phi(q) @ key_state + 1)
+    """
+
+    def __init__(self, config) -> None:
+        super().__init__(config)
+        self.qk_kernel_type = config.memory_configs.get("qk_kernel_type", "elu_plus_one")
+
+    def read(self, query_for_memory: torch.Tensor) -> torch.Tensor:
+        bsz, _h, seq_len, _d = query_for_memory.shape
+        self._ensure_ready(bsz, query_for_memory.device, query_for_memory.dtype)
+        q_flat = query_for_memory.transpose(1, 2).reshape(bsz, seq_len, -1)
+
+        if self.num_kv_groups > 1:
+            q_2d = q_flat.view(bsz, seq_len * self.num_kv_groups, self.kv_dim)
+            q_2d = _qk_kernel(q_2d, self.qk_kernel_type)
+            out_2d = torch.matmul(q_2d, self._state)
+            if self._key_state is not None:
+                norm_factor = torch.matmul(q_2d, self._key_state)
+                out_2d = out_2d / norm_factor
+            return out_2d.view(bsz, seq_len, self.q_dim).contiguous()
+
+        q_flat = _qk_kernel(q_flat, self.qk_kernel_type)
+        out = torch.matmul(q_flat, self._state)
+        if self._key_state is not None:
+            norm_factor = torch.matmul(q_flat, self._key_state)
+            out = out / norm_factor
+        return out.contiguous()
+
+
+class L2NormalizedDeltaNetMetisLocalMemory(DeltaNetMetisLocalMemory):
+    """DeltaNet memory with L2-normalized queries and no key-state denominator.
+
+    Pair this with ``StrideL2NormMetisHyperMemory``:
+
+        q = normalize(q)
+        y = q @ state
+
+    This keeps the read-side q scale controlled by L2 normalization while the
+    write-side hyper memory normalizes k.
+    """
+
+    def read(self, query_for_memory: torch.Tensor) -> torch.Tensor:
+        bsz, _h, seq_len, _d = query_for_memory.shape
+        self._ensure_ready(bsz, query_for_memory.device, query_for_memory.dtype)
+        query_for_memory = F.normalize(query_for_memory, dim=-1)
+        q_flat = query_for_memory.transpose(1, 2).reshape(bsz, seq_len, -1)
+
+        if self.num_kv_groups > 1:
+            q_2d = q_flat.view(bsz, seq_len * self.num_kv_groups, self.kv_dim)
+            out_2d = torch.matmul(q_2d, self._state)
+            return out_2d.view(bsz, seq_len, self.q_dim).contiguous()
+
+        return torch.matmul(q_flat, self._state).contiguous()
+
+
+class OneStepAblationMetisLocalMemory(MetisLocalMemoryBase):
+    def __init__(self, config) -> None:
+        super().__init__(config)
+    
+    def initialize(self) -> None:
+        self.memory_state = None
+
+    def reset(self) -> None:
+        self.initialize()
+
+    def read(self, query_for_memory):
+        return self.memory_state
+
+    def write(self, new_info) -> None:
+        self.memory_state = new_info

model-00001-of-00004.safetensors

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

model-00002-of-00004.safetensors

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

model-00003-of-00004.safetensors

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

model-00004-of-00004.safetensors

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

model.safetensors.index.json

@@ -0,0 +1,515 @@
+{
+  "metadata": {
+    "total_parameters": 9155262992,
+    "total_size": 18310525984
+  },
+  "weight_map": {
+    "model.language_model.metis_backbone.lm_head.weight": "model-00001-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.embed_tokens.weight": "model-00001-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.0.input_layernorm.weight": "model-00001-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.0.linear_attn.A_log": "model-00001-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.0.linear_attn.conv1d.weight": "model-00001-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.0.linear_attn.dt_bias": "model-00001-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.0.linear_attn.in_proj_a.weight": "model-00001-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.0.linear_attn.in_proj_b.weight": "model-00001-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.0.linear_attn.in_proj_qkv.weight": "model-00001-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.0.linear_attn.in_proj_z.weight": "model-00001-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.0.linear_attn.norm.weight": "model-00001-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.0.linear_attn.out_proj.weight": "model-00001-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.0.mlp.down_proj.weight": "model-00001-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.0.mlp.gate_proj.weight": "model-00001-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.0.mlp.up_proj.weight": "model-00001-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.0.post_attention_layernorm.weight": "model-00001-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.1.input_layernorm.weight": "model-00001-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.1.linear_attn.A_log": "model-00001-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.1.linear_attn.conv1d.weight": "model-00001-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.1.linear_attn.dt_bias": "model-00001-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.1.linear_attn.in_proj_a.weight": "model-00001-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.1.linear_attn.in_proj_b.weight": "model-00001-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.1.linear_attn.in_proj_qkv.weight": "model-00001-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.1.linear_attn.in_proj_z.weight": "model-00001-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.1.linear_attn.norm.weight": "model-00001-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.1.linear_attn.out_proj.weight": "model-00001-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.1.mlp.down_proj.weight": "model-00001-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.1.mlp.gate_proj.weight": "model-00001-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.1.mlp.up_proj.weight": "model-00001-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.1.post_attention_layernorm.weight": "model-00001-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.10.input_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.10.linear_attn.A_log": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.10.linear_attn.conv1d.weight": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.10.linear_attn.dt_bias": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.10.linear_attn.in_proj_a.weight": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.10.linear_attn.in_proj_b.weight": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.10.linear_attn.in_proj_qkv.weight": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.10.linear_attn.in_proj_z.weight": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.10.linear_attn.norm.weight": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.10.linear_attn.out_proj.weight": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.10.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.10.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.10.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.10.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.11.input_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.11.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.11.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.11.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.11.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.11.self_attn.k_norm.weight": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.11.self_attn.k_proj.weight": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.11.self_attn.o_proj.weight": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.11.self_attn.q_norm.weight": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.11.self_attn.q_proj.weight": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.11.self_attn.v_proj.weight": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.12.input_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.12.linear_attn.A_log": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.12.linear_attn.conv1d.weight": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.12.linear_attn.dt_bias": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.12.linear_attn.in_proj_a.weight": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.12.linear_attn.in_proj_b.weight": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.12.linear_attn.in_proj_qkv.weight": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.12.linear_attn.in_proj_z.weight": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.12.linear_attn.norm.weight": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.12.linear_attn.out_proj.weight": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.12.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.12.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.12.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.12.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.13.input_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.13.linear_attn.A_log": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.13.linear_attn.conv1d.weight": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.13.linear_attn.dt_bias": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.13.linear_attn.in_proj_a.weight": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.13.linear_attn.in_proj_b.weight": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.13.linear_attn.in_proj_qkv.weight": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.13.linear_attn.in_proj_z.weight": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.13.linear_attn.norm.weight": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.13.linear_attn.out_proj.weight": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.13.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.13.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.13.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.13.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.14.input_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.14.linear_attn.A_log": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.14.linear_attn.conv1d.weight": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.14.linear_attn.dt_bias": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.14.linear_attn.in_proj_a.weight": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.14.linear_attn.in_proj_b.weight": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.14.linear_attn.in_proj_qkv.weight": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.14.linear_attn.in_proj_z.weight": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.14.linear_attn.norm.weight": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.14.linear_attn.out_proj.weight": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.14.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.14.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.14.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.14.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.15.input_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.15.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.15.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.15.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.15.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.15.self_attn.k_norm.weight": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.15.self_attn.k_proj.weight": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.15.self_attn.o_proj.weight": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.15.self_attn.q_norm.weight": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.15.self_attn.q_proj.weight": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.15.self_attn.v_proj.weight": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.16.input_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.16.linear_attn.A_log": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.16.linear_attn.conv1d.weight": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.16.linear_attn.dt_bias": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.16.linear_attn.in_proj_a.weight": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.16.linear_attn.in_proj_b.weight": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.16.linear_attn.in_proj_qkv.weight": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.16.linear_attn.in_proj_z.weight": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.16.linear_attn.norm.weight": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.16.linear_attn.out_proj.weight": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.16.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.16.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.16.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.16.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.17.input_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.17.linear_attn.A_log": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.17.linear_attn.conv1d.weight": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.17.linear_attn.dt_bias": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.17.linear_attn.in_proj_a.weight": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.17.linear_attn.in_proj_b.weight": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.17.linear_attn.in_proj_qkv.weight": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.17.linear_attn.in_proj_z.weight": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.17.linear_attn.norm.weight": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.17.linear_attn.out_proj.weight": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.17.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.17.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.17.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.17.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.18.input_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.18.linear_attn.A_log": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.18.linear_attn.conv1d.weight": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.18.linear_attn.dt_bias": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.18.linear_attn.in_proj_a.weight": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.18.linear_attn.in_proj_b.weight": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.18.linear_attn.in_proj_qkv.weight": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.18.linear_attn.in_proj_z.weight": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.18.linear_attn.norm.weight": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.18.linear_attn.out_proj.weight": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.18.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.18.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.18.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.18.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.19.input_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.19.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.19.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.19.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.19.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.19.self_attn.k_norm.weight": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.19.self_attn.k_proj.weight": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.19.self_attn.o_proj.weight": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.19.self_attn.q_norm.weight": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.19.self_attn.q_proj.weight": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.19.self_attn.v_proj.weight": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.2.input_layernorm.weight": "model-00001-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.2.linear_attn.A_log": "model-00001-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.2.linear_attn.conv1d.weight": "model-00001-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.2.linear_attn.dt_bias": "model-00001-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.2.linear_attn.in_proj_a.weight": "model-00001-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.2.linear_attn.in_proj_b.weight": "model-00001-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.2.linear_attn.in_proj_qkv.weight": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.2.linear_attn.in_proj_z.weight": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.2.linear_attn.norm.weight": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.2.linear_attn.out_proj.weight": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.2.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.2.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.2.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.2.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.20.input_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.20.linear_attn.A_log": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.20.linear_attn.conv1d.weight": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.20.linear_attn.dt_bias": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.20.linear_attn.in_proj_a.weight": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.20.linear_attn.in_proj_b.weight": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.20.linear_attn.in_proj_qkv.weight": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.20.linear_attn.in_proj_z.weight": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.20.linear_attn.norm.weight": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.20.linear_attn.out_proj.weight": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.20.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.20.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.20.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.20.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.21.input_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.21.linear_attn.A_log": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.21.linear_attn.conv1d.weight": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.21.linear_attn.dt_bias": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.21.linear_attn.in_proj_a.weight": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.21.linear_attn.in_proj_b.weight": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.21.linear_attn.in_proj_qkv.weight": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.21.linear_attn.in_proj_z.weight": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.21.linear_attn.norm.weight": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.21.linear_attn.out_proj.weight": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.21.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.21.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.21.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.21.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.22.input_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.22.linear_attn.A_log": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.22.linear_attn.conv1d.weight": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.22.linear_attn.dt_bias": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.22.linear_attn.in_proj_a.weight": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.22.linear_attn.in_proj_b.weight": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.22.linear_attn.in_proj_qkv.weight": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.22.linear_attn.in_proj_z.weight": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.22.linear_attn.norm.weight": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.22.linear_attn.out_proj.weight": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.22.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.22.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.22.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.22.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.23.input_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.23.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.23.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.23.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.23.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.23.self_attn.k_norm.weight": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.23.self_attn.k_proj.weight": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.23.self_attn.o_proj.weight": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.23.self_attn.q_norm.weight": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.23.self_attn.q_proj.weight": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.23.self_attn.v_proj.weight": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.24.input_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.24.linear_attn.A_log": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.24.linear_attn.conv1d.weight": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.24.linear_attn.dt_bias": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.24.linear_attn.in_proj_a.weight": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.24.linear_attn.in_proj_b.weight": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.24.linear_attn.in_proj_qkv.weight": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.24.linear_attn.in_proj_z.weight": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.24.linear_attn.norm.weight": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.24.linear_attn.out_proj.weight": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.24.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.24.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.24.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.24.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.25.input_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.25.linear_attn.A_log": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.25.linear_attn.conv1d.weight": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.25.linear_attn.dt_bias": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.25.linear_attn.in_proj_a.weight": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.25.linear_attn.in_proj_b.weight": "model-00003-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.25.linear_attn.in_proj_qkv.weight": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.25.linear_attn.in_proj_z.weight": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.25.linear_attn.norm.weight": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.25.linear_attn.out_proj.weight": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.25.mlp.down_proj.weight": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.25.mlp.gate_proj.weight": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.25.mlp.up_proj.weight": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.25.post_attention_layernorm.weight": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.26.input_layernorm.weight": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.26.linear_attn.A_log": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.26.linear_attn.conv1d.weight": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.26.linear_attn.dt_bias": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.26.linear_attn.in_proj_a.weight": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.26.linear_attn.in_proj_b.weight": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.26.linear_attn.in_proj_qkv.weight": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.26.linear_attn.in_proj_z.weight": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.26.linear_attn.norm.weight": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.26.linear_attn.out_proj.weight": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.26.mlp.down_proj.weight": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.26.mlp.gate_proj.weight": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.26.mlp.up_proj.weight": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.26.post_attention_layernorm.weight": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.27.input_layernorm.weight": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.27.mlp.down_proj.weight": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.27.mlp.gate_proj.weight": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.27.mlp.up_proj.weight": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.27.post_attention_layernorm.weight": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.27.self_attn.k_norm.weight": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.27.self_attn.k_proj.weight": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.27.self_attn.o_proj.weight": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.27.self_attn.q_norm.weight": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.27.self_attn.q_proj.weight": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.27.self_attn.v_proj.weight": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.28.input_layernorm.weight": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.28.linear_attn.A_log": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.28.linear_attn.conv1d.weight": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.28.linear_attn.dt_bias": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.28.linear_attn.in_proj_a.weight": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.28.linear_attn.in_proj_b.weight": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.28.linear_attn.in_proj_qkv.weight": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.28.linear_attn.in_proj_z.weight": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.28.linear_attn.norm.weight": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.28.linear_attn.out_proj.weight": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.28.mlp.down_proj.weight": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.28.mlp.gate_proj.weight": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.28.mlp.up_proj.weight": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.28.post_attention_layernorm.weight": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.29.input_layernorm.weight": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.29.linear_attn.A_log": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.29.linear_attn.conv1d.weight": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.29.linear_attn.dt_bias": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.29.linear_attn.in_proj_a.weight": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.29.linear_attn.in_proj_b.weight": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.29.linear_attn.in_proj_qkv.weight": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.29.linear_attn.in_proj_z.weight": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.29.linear_attn.norm.weight": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.29.linear_attn.out_proj.weight": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.29.mlp.down_proj.weight": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.29.mlp.gate_proj.weight": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.29.mlp.up_proj.weight": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.29.post_attention_layernorm.weight": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.3.input_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.3.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.3.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.3.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.3.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.3.self_attn.k_norm.weight": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.3.self_attn.k_proj.weight": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.3.self_attn.o_proj.weight": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.3.self_attn.q_norm.weight": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.3.self_attn.q_proj.weight": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.3.self_attn.v_proj.weight": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.30.input_layernorm.weight": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.30.linear_attn.A_log": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.30.linear_attn.conv1d.weight": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.30.linear_attn.dt_bias": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.30.linear_attn.in_proj_a.weight": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.30.linear_attn.in_proj_b.weight": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.30.linear_attn.in_proj_qkv.weight": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.30.linear_attn.in_proj_z.weight": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.30.linear_attn.norm.weight": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.30.linear_attn.out_proj.weight": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.30.mlp.down_proj.weight": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.30.mlp.gate_proj.weight": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.30.mlp.up_proj.weight": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.30.post_attention_layernorm.weight": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.31.input_layernorm.weight": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.31.mlp.down_proj.weight": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.31.mlp.gate_proj.weight": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.31.mlp.up_proj.weight": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.31.post_attention_layernorm.weight": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.31.self_attn.k_norm.weight": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.31.self_attn.k_proj.weight": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.31.self_attn.o_proj.weight": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.31.self_attn.q_norm.weight": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.31.self_attn.q_proj.weight": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.31.self_attn.v_proj.weight": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.4.input_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.4.linear_attn.A_log": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.4.linear_attn.conv1d.weight": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.4.linear_attn.dt_bias": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.4.linear_attn.in_proj_a.weight": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.4.linear_attn.in_proj_b.weight": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.4.linear_attn.in_proj_qkv.weight": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.4.linear_attn.in_proj_z.weight": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.4.linear_attn.norm.weight": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.4.linear_attn.out_proj.weight": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.4.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.4.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.4.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.4.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.5.input_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.5.linear_attn.A_log": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.5.linear_attn.conv1d.weight": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.5.linear_attn.dt_bias": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.5.linear_attn.in_proj_a.weight": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.5.linear_attn.in_proj_b.weight": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.5.linear_attn.in_proj_qkv.weight": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.5.linear_attn.in_proj_z.weight": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.5.linear_attn.norm.weight": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.5.linear_attn.out_proj.weight": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.5.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.5.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.5.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.5.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.6.input_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.6.linear_attn.A_log": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.6.linear_attn.conv1d.weight": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.6.linear_attn.dt_bias": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.6.linear_attn.in_proj_a.weight": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.6.linear_attn.in_proj_b.weight": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.6.linear_attn.in_proj_qkv.weight": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.6.linear_attn.in_proj_z.weight": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.6.linear_attn.norm.weight": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.6.linear_attn.out_proj.weight": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.6.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.6.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.6.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.6.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.7.input_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.7.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.7.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.7.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.7.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.7.self_attn.k_norm.weight": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.7.self_attn.k_proj.weight": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.7.self_attn.o_proj.weight": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.7.self_attn.q_norm.weight": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.7.self_attn.q_proj.weight": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.7.self_attn.v_proj.weight": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.8.input_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.8.linear_attn.A_log": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.8.linear_attn.conv1d.weight": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.8.linear_attn.dt_bias": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.8.linear_attn.in_proj_a.weight": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.8.linear_attn.in_proj_b.weight": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.8.linear_attn.in_proj_qkv.weight": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.8.linear_attn.in_proj_z.weight": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.8.linear_attn.norm.weight": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.8.linear_attn.out_proj.weight": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.8.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.8.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.8.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.8.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.9.input_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.9.linear_attn.A_log": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.9.linear_attn.conv1d.weight": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.9.linear_attn.dt_bias": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.9.linear_attn.in_proj_a.weight": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.9.linear_attn.in_proj_b.weight": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.9.linear_attn.in_proj_qkv.weight": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.9.linear_attn.in_proj_z.weight": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.9.linear_attn.norm.weight": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.9.linear_attn.out_proj.weight": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.9.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.9.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.9.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.layers.9.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.language_model.metis_backbone.model.language_model.norm.weight": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_blocks.11.hyper_memory.W_k.weight": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_blocks.11.hyper_memory.W_v.weight": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_blocks.11.hyper_memory.gated_delta_alpha.bias": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_blocks.11.hyper_memory.gated_delta_alpha.weight": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_blocks.11.hyper_memory.gated_delta_beta.bias": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_blocks.11.hyper_memory.gated_delta_beta.weight": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_blocks.11.hyper_memory.pool_score.weight": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_blocks.11.mem_norm.weight": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_blocks.11.query_norm.weight": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_blocks.11.query_proj.weight": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_blocks.15.hyper_memory.W_k.weight": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_blocks.15.hyper_memory.W_v.weight": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_blocks.15.hyper_memory.gated_delta_alpha.bias": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_blocks.15.hyper_memory.gated_delta_alpha.weight": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_blocks.15.hyper_memory.gated_delta_beta.bias": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_blocks.15.hyper_memory.gated_delta_beta.weight": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_blocks.15.hyper_memory.pool_score.weight": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_blocks.15.mem_norm.weight": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_blocks.15.query_norm.weight": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_blocks.15.query_proj.weight": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_blocks.19.hyper_memory.W_k.weight": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_blocks.19.hyper_memory.W_v.weight": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_blocks.19.hyper_memory.gated_delta_alpha.bias": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_blocks.19.hyper_memory.gated_delta_alpha.weight": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_blocks.19.hyper_memory.gated_delta_beta.bias": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_blocks.19.hyper_memory.gated_delta_beta.weight": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_blocks.19.hyper_memory.pool_score.weight": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_blocks.19.mem_norm.weight": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_blocks.19.query_norm.weight": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_blocks.19.query_proj.weight": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_blocks.23.hyper_memory.W_k.weight": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_blocks.23.hyper_memory.W_v.weight": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_blocks.23.hyper_memory.gated_delta_alpha.bias": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_blocks.23.hyper_memory.gated_delta_alpha.weight": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_blocks.23.hyper_memory.gated_delta_beta.bias": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_blocks.23.hyper_memory.gated_delta_beta.weight": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_blocks.23.hyper_memory.pool_score.weight": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_blocks.23.mem_norm.weight": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_blocks.23.query_norm.weight": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_blocks.23.query_proj.weight": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_blocks.27.hyper_memory.W_k.weight": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_blocks.27.hyper_memory.W_v.weight": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_blocks.27.hyper_memory.gated_delta_alpha.bias": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_blocks.27.hyper_memory.gated_delta_alpha.weight": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_blocks.27.hyper_memory.gated_delta_beta.bias": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_blocks.27.hyper_memory.gated_delta_beta.weight": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_blocks.27.hyper_memory.pool_score.weight": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_blocks.27.mem_norm.weight": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_blocks.27.query_norm.weight": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_blocks.27.query_proj.weight": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_blocks.3.hyper_memory.W_k.weight": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_blocks.3.hyper_memory.W_v.weight": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_blocks.3.hyper_memory.gated_delta_alpha.bias": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_blocks.3.hyper_memory.gated_delta_alpha.weight": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_blocks.3.hyper_memory.gated_delta_beta.bias": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_blocks.3.hyper_memory.gated_delta_beta.weight": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_blocks.3.hyper_memory.pool_score.weight": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_blocks.3.mem_norm.weight": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_blocks.3.query_norm.weight": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_blocks.3.query_proj.weight": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_blocks.31.hyper_memory.W_k.weight": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_blocks.31.hyper_memory.W_v.weight": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_blocks.31.hyper_memory.gated_delta_alpha.bias": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_blocks.31.hyper_memory.gated_delta_alpha.weight": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_blocks.31.hyper_memory.gated_delta_beta.bias": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_blocks.31.hyper_memory.gated_delta_beta.weight": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_blocks.31.hyper_memory.pool_score.weight": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_blocks.31.mem_norm.weight": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_blocks.31.query_norm.weight": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_blocks.31.query_proj.weight": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_blocks.7.hyper_memory.W_k.weight": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_blocks.7.hyper_memory.W_v.weight": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_blocks.7.hyper_memory.gated_delta_alpha.bias": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_blocks.7.hyper_memory.gated_delta_alpha.weight": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_blocks.7.hyper_memory.gated_delta_beta.bias": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_blocks.7.hyper_memory.gated_delta_beta.weight": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_blocks.7.hyper_memory.pool_score.weight": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_blocks.7.mem_norm.weight": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_blocks.7.query_norm.weight": "model-00004-of-00004.safetensors",
+    "model.language_model.metis_blocks.7.query_proj.weight": "model-00004-of-00004.safetensors"
+  }
+}

modeling_metis.py

@@ -0,0 +1,147 @@
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers import PreTrainedModel
+from transformers.modeling_outputs import (
+    BaseModelOutputWithPast,
+    CausalLMOutputWithPast,
+)
+
+from .configuration_metis import MetisConfig
+from .utils import TrajectoryGenerationMixin, create_metis_causallm
+from .metis_block import create_metis_block
+
+
+class MetisPreTrainedModel(PreTrainedModel):
+    config_class = MetisConfig
+    base_model_prefix = "model"
+    supports_gradient_checkpointing = True
+    _no_split_modules = ["MetisBlock"]
+    _skip_keys_device_placement = "past_key_values"
+    _tied_weights_keys = []
+
+    @torch.no_grad()
+    def _init_weights(self, module: nn.Module) -> None:
+        super()._init_weights(module)
+
+
+class MetisModel(MetisPreTrainedModel):
+    def __init__(self, config: MetisConfig):
+        super().__init__(config)
+
+        self.metis_backbone = create_metis_causallm(config)
+
+        self.metis_blocks = nn.ModuleList(
+            [create_metis_block(config, i, self.metis_backbone.get_decoder_layer_by_id(i))
+             for i in range(self.metis_backbone.model.config.num_hidden_layers)]
+        )
+        self.metis_backbone.register_metis_blocks(self.metis_blocks)
+
+        self.post_init()
+
+    def forward(self, **kwargs) -> BaseModelOutputWithPast:
+        return self.metis_backbone.forward_with_memory(**kwargs)
+
+
+class MetisForCausalLM(MetisPreTrainedModel, TrajectoryGenerationMixin):
+    def __init__(self, config: MetisConfig):
+        super().__init__(config)
+        self.model = MetisModel(config)
+        self.post_init()
+
+    @torch.no_grad()
+    def reset(self):
+        for layer in self.model.metis_blocks:
+            if layer.local_memory is None:
+                continue
+            layer.local_memory.reset()
+
+    # Backward-compat alias.
+    reset_memory = reset
+
+    def _commit_memory(self, outputs, attention_mask=None):
+        """Write per-layer hidden states into local memory with optional mask.
+
+        attention_mask is passed through to hyper_memory.update_local_memory
+        so mask-aware variants (e.g. LinearLastMetisHyperMemory) can select
+        the last *real* token rather than the last position.
+        """
+        offset = self.config.memory_configs.get('commit_hidden_offset', 0)
+        if offset not in (0, 1):
+            raise ValueError(f"commit_hidden_offset must be 0 or 1, got {offset!r}")
+        all_hidden = outputs.hidden_states
+        for k, layer in enumerate(self.model.metis_blocks):
+            if layer.local_memory is None:
+                continue
+            layer_h = all_hidden[k + offset]
+            layer.hyper_memory.update_local_memory(
+                layer_h, layer.local_memory, attention_mask=attention_mask,
+            )
+
+    def commit(self, outputs):
+        """Public commit wrapper.  For mask-aware writes use _commit_memory directly."""
+        self._commit_memory(outputs)
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        position_ids: torch.LongTensor | None = None,
+        past_key_values=None,
+        inputs_embeds: torch.FloatTensor | None = None,
+        labels: torch.LongTensor | None = None,
+        use_cache: bool | None = None,
+        cache_position: torch.LongTensor | None = None,
+        logits_to_keep: int | torch.Tensor = 0,
+        commit_memory: bool = False,
+        attention_mask_1d: torch.Tensor | None = None,
+        **kwargs,
+    ) -> CausalLMOutputWithPast:
+        if commit_memory:
+            kwargs['output_hidden_states'] = True
+
+        outputs: BaseModelOutputWithPast = 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,
+            cache_position=cache_position,
+            **kwargs,
+        )
+
+        hidden_states = outputs.last_hidden_state
+        slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
+        logits = self.model.metis_backbone.lm_head(hidden_states[:, slice_indices, :])
+
+        loss = None
+        if labels is not None:
+            shift_logits = logits[..., :-1, :].contiguous()
+            shift_labels = labels[..., 1:].contiguous()
+            loss = F.cross_entropy(
+                shift_logits.view(-1, shift_logits.size(-1)),
+                shift_labels.view(-1),
+                ignore_index=-100,
+            )
+
+        if commit_memory:
+            mask_1d = attention_mask_1d
+            if mask_1d is None and attention_mask is not None and attention_mask.ndim == 2:
+                mask_1d = attention_mask
+            self._commit_memory(outputs, attention_mask=mask_1d)
+
+        return CausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+if False:  # pragma: no cover - dependency markers for HF dynamic modules
+    from .metis_hyper_memory import create_metis_hyper_memory
+    from .metis_local_memory import create_metis_local_memory
+    from .Qwen3_5_wrapper import Qwen3_5CausalLMForMetis

preprocessor_config.json

@@ -0,0 +1,21 @@
+{
+    "size": {
+        "longest_edge": 16777216,
+        "shortest_edge": 65536
+    },
+    "patch_size": 16,
+    "temporal_patch_size": 2,
+    "merge_size": 2,
+    "image_mean": [
+        0.5,
+        0.5,
+        0.5
+    ],
+    "image_std": [
+        0.5,
+        0.5,
+        0.5
+    ],
+    "processor_class": "Qwen3VLProcessor",
+    "image_processor_type": "Qwen2VLImageProcessorFast"
+}

tokenizer.json

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

tokenizer_config.json

@@ -0,0 +1,31 @@
+{
+  "add_prefix_space": false,
+  "audio_bos_token": "<|audio_start|>",
+  "audio_eos_token": "<|audio_end|>",
+  "audio_token": "<|audio_pad|>",
+  "backend": "tokenizers",
+  "bos_token": null,
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<|im_end|>",
+  "errors": "replace",
+  "image_token": "<|image_pad|>",
+  "is_local": true,
+  "model_max_length": 262144,
+  "model_specific_special_tokens": {
+    "audio_bos_token": "<|audio_start|>",
+    "audio_eos_token": "<|audio_end|>",
+    "audio_token": "<|audio_pad|>",
+    "image_token": "<|image_pad|>",
+    "video_token": "<|video_pad|>",
+    "vision_bos_token": "<|vision_start|>",
+    "vision_eos_token": "<|vision_end|>"
+  },
+  "pad_token": "<|endoftext|>",
+  "pretokenize_regex": "(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\\r\\n\\p{L}\\p{N}]?[\\p{L}\\p{M}]+|\\p{N}| ?[^\\s\\p{L}\\p{M}\\p{N}]+[\\r\\n]*|\\s*[\\r\\n]+|\\s+(?!\\S)|\\s+",
+  "split_special_tokens": false,
+  "tokenizer_class": "TokenizersBackend",
+  "unk_token": null,
+  "video_token": "<|video_pad|>",
+  "vision_bos_token": "<|vision_start|>",
+  "vision_eos_token": "<|vision_end|>"
+}

utils.py

@@ -0,0 +1,83 @@
+# Define some utils for constructing Metis.
+from transformers import GenerationMixin
+
+class TrajectoryGenerationMixin(GenerationMixin):
+    def reset(self):
+        raise NotImplementedError
+
+    def commit(self):
+        raise NotImplementedError
+
+    def step_generate(self, input_ids, **kwargs):
+        outputs = self.generate(
+            input_ids=input_ids,
+            return_dict_in_generate=True,
+            output_hidden_states=False,
+            **kwargs
+        )
+
+        trajectory_ids = outputs.sequences
+
+        # Re-forward the trajectory to get the final hidden states.
+        final_outputs = self.model.forward(
+            input_ids=trajectory_ids,
+            output_hidden_states=True, 
+            use_cache=False 
+        )
+
+        self.commit(final_outputs)
+        
+        return trajectory_ids
+    
+import torch
+import torch.nn as nn
+import importlib
+
+def create_metis_decoder_layer(config, raw_decoder):
+    module = importlib.import_module(f"{__package__}.{config.backbone_meta['backbone_type']}_wrapper")
+    decoder_layer_class = getattr(module, '%sDecoderLayerForMetis' % config.backbone_meta['backbone_type'])
+
+    return decoder_layer_class(config, raw_decoder)
+
+def create_metis_causallm(config):
+    module = importlib.import_module(f"{__package__}.{config.backbone_meta['backbone_type']}_wrapper")
+    model_class = getattr(module, '%sCausalLMForMetis' % config.backbone_meta['backbone_type'])
+    
+    return model_class(config)
+
+class DecoderLayerWrapperForMetis(nn.Module):
+    def __init__(self, config, raw_decoder):
+        super().__init__()
+        self.config = config
+        self._raw_decoder_ref = [raw_decoder]
+    
+    @property
+    def raw_decoder(self):
+        return self._raw_decoder_ref[0]
+
+    def before_mixin(self, **kwargs):
+        raise NotImplementedError
+    
+    def after_mixin(self, memory_carrier, cache_dict, **kwargs) -> torch.Tensor:
+        raise NotImplementedError
+
+class CausalLMWrapperForMetis(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+    
+    def register_metis_blocks(self, metis_blocks):
+        # Store as a plain list so PyTorch does NOT register these as submodules
+        # of this wrapper. The MetisModel already owns the ModuleList; registering
+        # it here too would create duplicate state-dict keys (shared-tensor error).
+        self._metis_blocks_ref = list(metis_blocks)
+
+    def get_decoder_layer_by_id(self, layer_id: int):
+        raise NotImplementedError
+    
+    def forward_with_memory(self, **kwargs):
+        raise NotImplementedError
+
+
+if False:  # pragma: no cover - dependency marker for HF dynamic modules
+    from .Qwen3_5_wrapper import Qwen3_5CausalLMForMetis

video_preprocessor_config.json

@@ -0,0 +1,21 @@
+{
+    "size": {
+        "longest_edge": 25165824,
+        "shortest_edge": 4096
+    },
+    "patch_size": 16,
+    "temporal_patch_size": 2,
+    "merge_size": 2,
+    "image_mean": [
+        0.5,
+        0.5,
+        0.5
+    ],
+    "image_std": [
+        0.5,
+        0.5,
+        0.5
+    ],
+    "processor_class": "Qwen3VLProcessor",
+    "video_processor_type": "Qwen3VLVideoProcessor"
+}