From 49123f0f06e8c40af0d6139286e13a82a6910b6f Mon Sep 17 00:00:00 2001
From: Dr Daniel Wu <95159536+drdanielwuwu@users.noreply.github.com>
Date: Fri, 3 Apr 2026 11:40:27 +0800
Subject: [PATCH 01/16] Add files via upload

Implements the MalaAttention layer, providing a high-level interface for MALA with support for multi-head attention, MQA, LEPE, and output gating.
---
 fla/layers/mala.py | 173 +++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 173 insertions(+)
 create mode 100644 fla/layers/mala.py

diff --git a/fla/layers/mala.py b/fla/layers/mala.py
new file mode 100644
index 0000000000..5353fc2243
--- /dev/null
+++ b/fla/layers/mala.py
@@ -0,0 +1,173 @@
+# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang
+
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import torch
+import torch.nn as nn
+from einops import rearrange
+
+from fla.ops.mala import naive_mala_attn
+from fla.modules.rotary import RoPE
+
+if TYPE_CHECKING:
+    from transformers.processing_utils import Unpack
+
+    from fla.models.utils import Cache
+
+"""
+MALA (Magnitude-Aware Linear Attention) implementation.
+
+Based on the paper:
+"Rectifying Magnitude Neglect in Linear Attention"
+ICCV 2025 (highlight)
+https://arxiv.org/abs/2507.00698
+
+Original implementation:
+https://github.com/aldjalkdf/MAViT
+"""
+
+
+class MalaAttention(nn.Module):
+    r"""
+    The layer implementation for MALA (Magnitude-Aware Linear Attention).
+    
+    Based on the paper:
+    "Rectifying Magnitude Neglect in Linear Attention"
+    ICCV 2025 (highlight)
+    https://arxiv.org/abs/2507.00698
+    
+    Args:
+        hidden_size (int, Optional):
+            The hidden size of the input. Default: 1024.
+        expand_k (float, Optional):
+            The expansion ratio for the key dim. Default: 1.0.
+        expand_v (float, Optional):
+            The expansion ratio for the value dim. Default: 1.0.
+        num_heads (int, Optional):
+            The number of heads. Default: 4.
+        num_kv_heads (int, Optional):
+            The number of key/value heads, used for MQA. Default: None.
+        use_lepe (bool, Optional):
+            Whether to use local positional embedding. Default: `True`.
+        lepe_kernel_size (int, Optional):
+            The kernel size for local positional embedding. Default: 5.
+        layer_idx (int, Optional):
+            The index of the layer. Default: None.
+    """
+
+    def __init__(
+        self,
+        hidden_size: int = 1024,
+        expand_k: float = 1.0,
+        expand_v: float = 1.0,
+        num_heads: int = 4,
+        num_kv_heads: int | None = None,
+        use_lepe: bool = True,
+        lepe_kernel_size: int = 5,
+        layer_idx: int | None = None,
+    ) -> MalaAttention:
+        super().__init__()
+
+        self.hidden_size = hidden_size
+        self.expand_k = expand_k
+        self.expand_v = expand_v
+        self.num_heads = num_heads
+        self.num_kv_heads = num_kv_heads if num_kv_heads is not None else num_heads
+        self.num_kv_groups = self.num_heads // self.num_kv_heads
+        self.use_lepe = use_lepe
+        self.lepe_kernel_size = lepe_kernel_size
+        self.layer_idx = layer_idx
+
+        self.key_dim = int(hidden_size * expand_k)
+        self.value_dim = int(hidden_size * expand_v)
+        self.key_dim_per_group = self.key_dim // self.num_kv_groups
+        self.value_dim_per_group = self.value_dim // self.num_kv_groups
+
+        assert self.key_dim % num_heads == 0, f"key dim must be divisible by num_heads of {num_heads}"
+        assert self.value_dim % num_heads == 0, f"value dim must be divisible by num_heads of {num_heads}"
+
+        self.head_k_dim = self.key_dim // num_heads
+        self.head_v_dim = self.value_dim // num_heads
+
+        # Projection layers
+        self.q_proj = nn.Linear(hidden_size, self.key_dim, bias=False)
+        self.k_proj = nn.Linear(hidden_size, self.key_dim_per_group, bias=False)
+        self.v_proj = nn.Linear(hidden_size, self.value_dim_per_group, bias=False)
+        self.o_proj = nn.Linear(self.value_dim, hidden_size, bias=False)
+        
+        # Local positional embedding
+        if use_lepe:
+            self.lepe = nn.Conv1d(
+                in_channels=self.value_dim,
+                out_channels=self.value_dim,
+                kernel_size=lepe_kernel_size,
+                padding=lepe_kernel_size // 2,
+                groups=self.value_dim
+            )
+
+        # Output gate projection
+        self.o_gate_proj = nn.Linear(hidden_size, self.value_dim, bias=False)
+
+        # RoPE embedding
+        self.rope = RoPE(self.head_k_dim)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+        past_key_values: Cache | None = None,
+        use_cache: bool | None = False,
+        output_attentions: bool | None = False,
+        **kwargs: Unpack[dict],
+    ) -> tuple[torch.Tensor, torch.Tensor | None, Cache | None]:
+        batch_size, seq_len, _ = hidden_states.shape
+
+        # Projection
+        q = self.q_proj(hidden_states)
+        k = self.k_proj(hidden_states)
+        v = self.v_proj(hidden_states)
+        o_gate = self.o_gate_proj(hidden_states)
+
+        # Reshape tensors
+        q = rearrange(q, 'b s (h d) -> b s h d', d=self.head_k_dim)
+        if self.num_kv_groups > 1:
+            k = rearrange(k, 'b s (h d) -> b s h d', d=self.head_k_dim)
+            k = k.repeat_interleave(self.num_kv_groups, dim=2)
+            v = rearrange(v, 'b s (h d) -> b s h d', d=self.head_v_dim)
+            v = v.repeat_interleave(self.num_kv_groups, dim=2)
+        else:
+            k = rearrange(k, 'b s (h d) -> b s h d', d=self.head_k_dim)
+            v = rearrange(v, 'b s (h d) -> b s h d', d=self.head_v_dim)
+
+        # Compute RoPE embeddings
+        sin, cos = self.rope(seq_len)
+        sin = sin.unsqueeze(0).unsqueeze(2).expand(batch_size, seq_len, self.num_heads, self.head_k_dim)
+        cos = cos.unsqueeze(0).unsqueeze(2).expand(batch_size, seq_len, self.num_heads, self.head_k_dim)
+
+        # Compute attention
+        o = naive_mala_attn(
+            q=q,
+            k=k,
+            v=v,
+            sin=sin,
+            cos=cos,
+        )
+
+        # Apply local positional embedding if enabled
+        if self.use_lepe:
+            lepe_input = rearrange(v, 'b s h d -> b (h d) s')
+            lepe_output = self.lepe(lepe_input)
+            lepe_output = rearrange(lepe_output, 'b (h d) s -> b s h d', h=self.num_heads)
+            o = o + lepe_output
+
+        # Apply output gate
+        o = rearrange(o, 'b s h d -> b s (h d)')
+        o = o * o_gate
+
+        # Final projection
+        o = self.o_proj(o)
+
+        return o, None, past_key_values

From 6ddece0058b01240d1e8942d045aa4d072cb56d1 Mon Sep 17 00:00:00 2001
From: Dr Daniel Wu <95159536+drdanielwuwu@users.noreply.github.com>
Date: Fri, 3 Apr 2026 11:41:31 +0800
Subject: [PATCH 02/16] Create /mala

---
 fla/ops/mala | 1 +
 1 file changed, 1 insertion(+)
 create mode 100644 fla/ops/mala

diff --git a/fla/ops/mala b/fla/ops/mala
new file mode 100644
index 0000000000..8b13789179
--- /dev/null
+++ b/fla/ops/mala
@@ -0,0 +1 @@
+

From e93f46a1f0d22c466619dc841305b32529435a2b Mon Sep 17 00:00:00 2001
From: Dr Daniel Wu <95159536+drdanielwuwu@users.noreply.github.com>
Date: Fri, 3 Apr 2026 11:42:41 +0800
Subject: [PATCH 03/16] Delete fla/ops/mala

---
 fla/ops/mala | 1 -
 1 file changed, 1 deletion(-)
 delete mode 100644 fla/ops/mala

diff --git a/fla/ops/mala b/fla/ops/mala
deleted file mode 100644
index 8b13789179..0000000000
--- a/fla/ops/mala
+++ /dev/null
@@ -1 +0,0 @@
-

From 25bf0787f70f2516e92e759e3acac2504f77d7b1 Mon Sep 17 00:00:00 2001
From: Dr Daniel Wu <95159536+drdanielwuwu@users.noreply.github.com>
Date: Fri, 3 Apr 2026 11:44:53 +0800
Subject: [PATCH 04/16] Create naive.py

---
 fla/ops/mala/naive.py | 1 +
 1 file changed, 1 insertion(+)
 create mode 100644 fla/ops/mala/naive.py

diff --git a/fla/ops/mala/naive.py b/fla/ops/mala/naive.py
new file mode 100644
index 0000000000..8b13789179
--- /dev/null
+++ b/fla/ops/mala/naive.py
@@ -0,0 +1 @@
+

From d325dc25d47ccb82bba3316efe15c1cbc141d7ea Mon Sep 17 00:00:00 2001
From: Dr Daniel Wu <95159536+drdanielwuwu@users.noreply.github.com>
Date: Fri, 3 Apr 2026 11:45:50 +0800
Subject: [PATCH 05/16] Create 10.py

---
 fla/models/mala/10.py | 1 +
 1 file changed, 1 insertion(+)
 create mode 100644 fla/models/mala/10.py

diff --git a/fla/models/mala/10.py b/fla/models/mala/10.py
new file mode 100644
index 0000000000..8b13789179
--- /dev/null
+++ b/fla/models/mala/10.py
@@ -0,0 +1 @@
+

From 0b1a8b0336f7cbd76b60711d32ce2a5bcb089062 Mon Sep 17 00:00:00 2001
From: Dr Daniel Wu <95159536+drdanielwuwu@users.noreply.github.com>
Date: Fri, 3 Apr 2026 11:48:04 +0800
Subject: [PATCH 06/16] Add files via upload

Implements the MALA model architecture, including MalaModel, MalaLayer, MalaMLP, and MalaForCausalLM classes for language modeling.
---
 fla/models/mala/modeling_mala.py | 272 +++++++++++++++++++++++++++++++
 1 file changed, 272 insertions(+)
 create mode 100644 fla/models/mala/modeling_mala.py

diff --git a/fla/models/mala/modeling_mala.py b/fla/models/mala/modeling_mala.py
new file mode 100644
index 0000000000..c1e3f77cec
--- /dev/null
+++ b/fla/models/mala/modeling_mala.py
@@ -0,0 +1,272 @@
+# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang
+
+
+from typing import Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+from transformers import PreTrainedModel
+from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
+
+from fla.layers import MalaAttention
+from fla.models.utils import cache_to_cpu, recycle_cache
+
+"""
+MALA (Magnitude-Aware Linear Attention) model implementation.
+
+Based on the paper:
+"Rectifying Magnitude Neglect in Linear Attention"
+ICCV 2025 (highlight)
+https://arxiv.org/abs/2507.00698
+
+Original implementation:
+https://github.com/aldjalkdf/MAViT
+"""
+
+
+class MalaPreTrainedModel(PreTrainedModel):
+    r"""
+    An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained models.
+    """
+
+    config_class = None
+    base_model_prefix = "model"
+
+    def _init_weights(self, module):
+        """Initialize the weights."""
+        if isinstance(module, nn.Linear):
+            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
+            if module.bias is not None:
+                module.bias.data.zero_()
+        elif isinstance(module, nn.Embedding):
+            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
+            if module.padding_idx is not None:
+                module.weight.data[module.padding_idx].zero_()
+
+
+class MalaModel(MalaPreTrainedModel):
+    r"""
+    The MALA (Magnitude-Aware Linear Attention) model.
+    """
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.config = config
+
+        # Embedding layer
+        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
+        
+        # Layers
+        self.layers = nn.ModuleList()
+        for i in range(config.num_hidden_layers):
+            layer = MalaLayer(config, layer_idx=i)
+            self.layers.append(layer)
+
+        # Final layer norm
+        self.norm = nn.RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+
+        # Initialize weights
+        self.apply(self._init_weights)
+
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        past_key_values: Optional[Tuple[Tuple[torch.Tensor]]] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[Tuple[torch.Tensor], BaseModelOutputWithPast]:
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        use_cache = use_cache if use_cache is not None else self.config.use_cache
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        # Get embeddings
+        hidden_states = self.embed_tokens(input_ids)
+
+        # Process layers
+        all_hidden_states = () if output_hidden_states else None
+        all_attentions = () if output_attentions else None
+        next_decoder_cache = () if use_cache else None
+
+        for i, layer in enumerate(self.layers):
+            if output_hidden_states:
+                all_hidden_states = all_hidden_states + (hidden_states,)
+
+            layer_past = past_key_values[i] if past_key_values is not None else None
+
+            layer_outputs = layer(
+                hidden_states,
+                attention_mask=attention_mask,
+                past_key_values=layer_past,
+                use_cache=use_cache,
+                output_attentions=output_attentions,
+            )
+
+            hidden_states = layer_outputs[0]
+
+            if use_cache:
+                next_decoder_cache += (layer_outputs[2],)
+
+            if output_attentions:
+                all_attentions = all_attentions + (layer_outputs[1],)
+
+        # Apply final layer norm
+        hidden_states = self.norm(hidden_states)
+
+        if output_hidden_states:
+            all_hidden_states = all_hidden_states + (hidden_states,)
+
+        if not return_dict:
+            return tuple(v for v in [hidden_states, next_decoder_cache, all_hidden_states, all_attentions] if v is not None)
+
+        return BaseModelOutputWithPast(
+            last_hidden_state=hidden_states,
+            past_key_values=next_decoder_cache,
+            hidden_states=all_hidden_states,
+            attentions=all_attentions,
+        )
+
+
+class MalaLayer(nn.Module):
+    r"""
+    A MALA layer.
+    """
+
+    def __init__(self, config, layer_idx):
+        super().__init__()
+        self.config = config
+        self.layer_idx = layer_idx
+
+        # Self-attention
+        self.self_attn = MalaAttention(
+            hidden_size=config.hidden_size,
+            expand_k=config.expand_k,
+            expand_v=config.expand_v,
+            num_heads=config.num_attention_heads,
+            num_kv_heads=config.num_key_value_heads,
+            use_lepe=config.use_lepe,
+            lepe_kernel_size=config.lepe_kernel_size,
+            layer_idx=layer_idx,
+        )
+
+        # MLP
+        self.mlp = MalaMLP(config)
+
+        # Layer norms
+        self.input_layernorm = nn.RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.post_attention_layernorm = nn.RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        past_key_values: Optional[Tuple[torch.Tensor]] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
+        # Self-attention
+        residual = hidden_states
+        hidden_states = self.input_layernorm(hidden_states)
+        attn_output, attn_weights, past_key_values = self.self_attn(
+            hidden_states,
+            attention_mask=attention_mask,
+            past_key_values=past_key_values,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+        )
+        hidden_states = residual + attn_output
+
+        # MLP
+        residual = hidden_states
+        hidden_states = self.post_attention_layernorm(hidden_states)
+        hidden_states = self.mlp(hidden_states)
+        hidden_states = residual + hidden_states
+
+        return hidden_states, attn_weights, past_key_values
+
+
+class MalaMLP(nn.Module):
+    r"""
+    MLP for MALA.
+    """
+
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        self.gate_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
+        self.up_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
+        self.down_proj = nn.Linear(config.intermediate_size, config.hidden_size, bias=False)
+        self.act_fn = nn.SiLU()
+
+    def forward(self, x):
+        x = self.gate_proj(x) * self.act_fn(self.up_proj(x))
+        x = self.down_proj(x)
+        return x
+
+
+class MalaForCausalLM(MalaPreTrainedModel):
+    r"""
+    MALA model for causal language modeling.
+    """
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.model = MalaModel(config)
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+
+        # Initialize weights
+        self.apply(self._init_weights)
+
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        past_key_values: Optional[Tuple[Tuple[torch.Tensor]]] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        labels: Optional[torch.LongTensor] = None,
+    ) -> Union[Tuple[torch.Tensor], CausalLMOutputWithPast]:
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        # Get model outputs
+        outputs = self.model(
+            input_ids,
+            attention_mask=attention_mask,
+            past_key_values=past_key_values,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+
+        # Get logits
+        logits = self.lm_head(outputs.last_hidden_state)
+
+        loss = None
+        if labels is not None:
+            # Shift logits and labels
+            shift_logits = logits[..., :-1, :].contiguous()
+            shift_labels = labels[..., 1:].contiguous()
+            # Compute loss
+            loss = nn.functional.cross_entropy(
+                shift_logits.view(-1, shift_logits.size(-1)),
+                shift_labels.view(-1),
+                ignore_index=-100,
+            )
+
+        if not return_dict:
+            output = (logits,) + outputs[1:]
+            return ((loss,) + output) if loss is not None else output
+
+        return CausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )

From ccc3991092bc21b90bfe2c404f50b9f996c4565c Mon Sep 17 00:00:00 2001
From: Dr Daniel Wu <95159536+drdanielwuwu@users.noreply.github.com>
Date: Fri, 3 Apr 2026 11:49:27 +0800
Subject: [PATCH 07/16] Add files via upload

Defines the MalaConfig class, which configures MALA model parameters such as hidden size, number of heads, and position embedding settings.
---
 fla/models/mala/configuration_mala.py | 76 +++++++++++++++++++++++++++
 1 file changed, 76 insertions(+)
 create mode 100644 fla/models/mala/configuration_mala.py

diff --git a/fla/models/mala/configuration_mala.py b/fla/models/mala/configuration_mala.py
new file mode 100644
index 0000000000..da50e7d8a3
--- /dev/null
+++ b/fla/models/mala/configuration_mala.py
@@ -0,0 +1,76 @@
+# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang
+
+
+from transformers import PretrainedConfig
+from transformers.utils import logging
+
+logger = logging.get_logger(__name__)
+
+"""
+MALA (Magnitude-Aware Linear Attention) configuration.
+
+Based on the paper:
+"Rectifying Magnitude Neglect in Linear Attention"
+ICCV 2025 (highlight)
+https://arxiv.org/abs/2507.00698
+
+Original implementation:
+https://github.com/aldjalkdf/MAViT
+"""
+
+
+class MalaConfig(PretrainedConfig):
+    r"""
+    Configuration class for MALA (Magnitude-Aware Linear Attention).
+    
+    Based on the paper:
+    "Rectifying Magnitude Neglect in Linear Attention"
+    ICCV 2025 (highlight)
+    https://arxiv.org/abs/2507.00698
+    """
+
+    model_type = "mala"
+
+    def __init__(
+        self,
+        vocab_size=32000,
+        hidden_size=1024,
+        intermediate_size=None,
+        num_hidden_layers=24,
+        num_attention_heads=16,
+        num_key_value_heads=None,
+        hidden_act="silu",
+        max_position_embeddings=8192,
+        initializer_range=0.02,
+        rms_norm_eps=1e-5,
+        use_cache=True,
+        tie_word_embeddings=False,
+        rope_theta=10000.0,
+        attention_dropout=0.0,
+        mlp_dropout=0.0,
+        expand_k=1.0,
+        expand_v=1.0,
+        use_lepe=True,
+        lepe_kernel_size=5,
+        **kwargs,
+    ):
+        super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs)
+
+        self.vocab_size = vocab_size
+        self.hidden_size = hidden_size
+        self.intermediate_size = intermediate_size if intermediate_size is not None else 4 * hidden_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.num_key_value_heads = num_key_value_heads if num_key_value_heads is not None else num_attention_heads
+        self.hidden_act = hidden_act
+        self.max_position_embeddings = max_position_embeddings
+        self.initializer_range = initializer_range
+        self.rms_norm_eps = rms_norm_eps
+        self.use_cache = use_cache
+        self.rope_theta = rope_theta
+        self.attention_dropout = attention_dropout
+        self.mlp_dropout = mlp_dropout
+        self.expand_k = expand_k
+        self.expand_v = expand_v
+        self.use_lepe = use_lepe
+        self.lepe_kernel_size = lepe_kernel_size

From 225dbc1adf2057d13546727919bed6b28b1babde Mon Sep 17 00:00:00 2001
From: Dr Daniel Wu <95159536+drdanielwuwu@users.noreply.github.com>
Date: Fri, 3 Apr 2026 11:50:17 +0800
Subject: [PATCH 08/16] Add files via upload

Initialization file for the MALA model module, exporting configuration and model classes.
---
 fla/models/mala/__init__.py | 11 +++++++++++
 1 file changed, 11 insertions(+)
 create mode 100644 fla/models/mala/__init__.py

diff --git a/fla/models/mala/__init__.py b/fla/models/mala/__init__.py
new file mode 100644
index 0000000000..b12630a868
--- /dev/null
+++ b/fla/models/mala/__init__.py
@@ -0,0 +1,11 @@
+# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang
+
+
+from fla.models.mala.configuration_mala import MalaConfig
+from fla.models.mala.modeling_mala import MalaModel, MalaForCausalLM
+
+__all__ = [
+    "MalaConfig",
+    "MalaModel",
+    "MalaForCausalLM",
+]

From a2a44cff4d9de6ebaa4ec075396cba639506ef86 Mon Sep 17 00:00:00 2001
From: Dr Daniel Wu <95159536+drdanielwuwu@users.noreply.github.com>
Date: Fri, 3 Apr 2026 11:50:55 +0800
Subject: [PATCH 09/16] Delete fla/models/mala/10.py

---
 fla/models/mala/10.py | 1 -
 1 file changed, 1 deletion(-)
 delete mode 100644 fla/models/mala/10.py

diff --git a/fla/models/mala/10.py b/fla/models/mala/10.py
deleted file mode 100644
index 8b13789179..0000000000
--- a/fla/models/mala/10.py
+++ /dev/null
@@ -1 +0,0 @@
-

From ecb40102091d9b217364f9027a5490c8b6ab7c82 Mon Sep 17 00:00:00 2001
From: Dr Daniel Wu <95159536+drdanielwuwu@users.noreply.github.com>
Date: Fri, 3 Apr 2026 11:52:26 +0800
Subject: [PATCH 10/16] Add files via upload

Initialization file for the MALA operations module, exporting the core functions.
---
 fla/ops/mala/__init__.py | 8 ++++++++
 1 file changed, 8 insertions(+)
 create mode 100644 fla/ops/mala/__init__.py

diff --git a/fla/ops/mala/__init__.py b/fla/ops/mala/__init__.py
new file mode 100644
index 0000000000..d4de052fb6
--- /dev/null
+++ b/fla/ops/mala/__init__.py
@@ -0,0 +1,8 @@
+# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang
+
+
+from fla.ops.mala.naive import naive_mala_attn
+
+__all__ = [
+    "naive_mala_attn",
+]

From f1bd705cf61c41b3181b79f8df1d98256a122eae Mon Sep 17 00:00:00 2001
From: Dr Daniel Wu <95159536+drdanielwuwu@users.noreply.github.com>
Date: Fri, 3 Apr 2026 11:52:50 +0800
Subject: [PATCH 11/16] Delete fla/ops/mala/naive.py

---
 fla/ops/mala/naive.py | 1 -
 1 file changed, 1 deletion(-)
 delete mode 100644 fla/ops/mala/naive.py

diff --git a/fla/ops/mala/naive.py b/fla/ops/mala/naive.py
deleted file mode 100644
index 8b13789179..0000000000
--- a/fla/ops/mala/naive.py
+++ /dev/null
@@ -1 +0,0 @@
-

From 7949e9e5226e8b5f4618104d524354e0df8f6e27 Mon Sep 17 00:00:00 2001
From: Dr Daniel Wu <95159536+drdanielwuwu@users.noreply.github.com>
Date: Fri, 3 Apr 2026 11:54:34 +0800
Subject: [PATCH 12/16] Add files via upload

Implements the core computational logic for MALA (Magnitude-Aware Linear Attention), including the forward pass with magnitude normalization and RoPE support.
---
 fla/ops/mala/naive.py | 82 +++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 82 insertions(+)
 create mode 100644 fla/ops/mala/naive.py

diff --git a/fla/ops/mala/naive.py b/fla/ops/mala/naive.py
new file mode 100644
index 0000000000..8382d7b1e0
--- /dev/null
+++ b/fla/ops/mala/naive.py
@@ -0,0 +1,82 @@
+# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang
+
+
+import torch
+import torch.nn.functional as F
+
+"""
+MALA (Magnitude-Aware Linear Attention) implementation.
+
+Based on the paper:
+"Rectifying Magnitude Neglect in Linear Attention"
+ICCV 2025 (highlight)
+https://arxiv.org/abs/2507.00698
+
+Original implementation:
+https://github.com/aldjalkdf/MAViT
+"""
+
+def rotate_every_two(x):
+    x1 = x[..., ::2]
+    x2 = x[..., 1::2]
+    x = torch.stack([-x2, x1], dim=-1)
+    return x.flatten(-2)
+
+
+def theta_shift(x, sin, cos):
+    return (x * cos) + (rotate_every_two(x) * sin)
+
+
+def naive_mala_attn(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    sin: torch.Tensor | None = None,
+    cos: torch.Tensor | None = None,
+    scale: float | None = None,
+):
+    """
+    Naive implementation of MALA (Magnitude-Aware Linear Attention).
+    
+    Based on the paper:
+    "Rectifying Magnitude Neglect in Linear Attention"
+    ICCV 2025 (highlight)
+    https://arxiv.org/abs/2507.00698
+    
+    Args:
+        q (torch.Tensor): Queries of shape [B, T, H, K].
+        k (torch.Tensor): Keys of shape [B, T, H, K].
+        v (torch.Tensor): Values of shape [B, T, H, V].
+        sin (torch.Tensor, optional): Sinusoidal position embeddings of shape [T, K].
+        cos (torch.Tensor, optional): Cosine position embeddings of shape [T, K].
+        scale (float, optional): Scale factor for attention scores.
+        
+    Returns:
+        torch.Tensor: Output of shape [B, T, H, V].
+    """
+    dtype = q.dtype
+    q, k, v = map(lambda x: x.float(), (q, k, v))
+    B, T, H, K, V = *q.shape, v.shape[-1]
+    
+    if scale is None:
+        scale = K ** -0.5
+    
+    # Apply ELU activation and add 1
+    q = F.elu(q) + 1
+    k = F.elu(k) + 1
+    
+    # Compute normalization factor
+    z = q @ k.mean(dim=1, keepdim=True).transpose(-1, -2) * scale
+    
+    # Apply RoPE if provided
+    if sin is not None and cos is not None:
+        q = theta_shift(q, sin, cos)
+        k = theta_shift(k, sin, cos)
+    
+    # Compute key-value product
+    kv = (k.transpose(-2, -1) * (scale / T ** 0.5)) @ (v * (scale / T ** 0.5))
+    
+    # Compute attention output
+    res = q @ kv * (1 + 1/(z + 1e-6)) - z * v.mean(dim=1, keepdim=True)
+    
+    return res.to(dtype)

From 077b84f74df05b5da38964fe12b52c756259d5f3 Mon Sep 17 00:00:00 2001
From: Dr Daniel Wu <95159536+drdanielwuwu@users.noreply.github.com>
Date: Fri, 3 Apr 2026 12:41:19 +0800
Subject: [PATCH 13/16] Refactor naive MALA attention implementation

Refactor naive MALA implementation to apply RoPE before normalization and adjust key-value product calculation.
---
 fla/ops/mala/naive.py | 14 +++++++-------
 1 file changed, 7 insertions(+), 7 deletions(-)

diff --git a/fla/ops/mala/naive.py b/fla/ops/mala/naive.py
index 8382d7b1e0..11126abe10 100644
--- a/fla/ops/mala/naive.py
+++ b/fla/ops/mala/naive.py
@@ -61,6 +61,11 @@ def naive_mala_attn(
     if scale is None:
         scale = K ** -0.5
     
+    # Apply RoPE if provided (before computing normalization factor)
+    if sin is not None and cos is not None:
+        q = theta_shift(q, sin, cos)
+        k = theta_shift(k, sin, cos)
+    
     # Apply ELU activation and add 1
     q = F.elu(q) + 1
     k = F.elu(k) + 1
@@ -68,13 +73,8 @@ def naive_mala_attn(
     # Compute normalization factor
     z = q @ k.mean(dim=1, keepdim=True).transpose(-1, -2) * scale
     
-    # Apply RoPE if provided
-    if sin is not None and cos is not None:
-        q = theta_shift(q, sin, cos)
-        k = theta_shift(k, sin, cos)
-    
-    # Compute key-value product
-    kv = (k.transpose(-2, -1) * (scale / T ** 0.5)) @ (v * (scale / T ** 0.5))
+    # Compute key-value product (sum over sequence dimension)
+    kv = (k.transpose(-2, -1) @ v) * (scale ** 2 / T)
     
     # Compute attention output
     res = q @ kv * (1 + 1/(z + 1e-6)) - z * v.mean(dim=1, keepdim=True)

From 7fcf1e63eb525546658e5b5d01328ba1a32df5e5 Mon Sep 17 00:00:00 2001
From: Dr Daniel Wu <95159536+drdanielwuwu@users.noreply.github.com>
Date: Fri, 3 Apr 2026 12:42:25 +0800
Subject: [PATCH 14/16] Modify constructor return type and add assertion

Change return type of the constructor from MalaAttention to None and add an assertion for num_kv_heads.
---
 fla/layers/mala.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/fla/layers/mala.py b/fla/layers/mala.py
index 5353fc2243..4d71569fc9 100644
--- a/fla/layers/mala.py
+++ b/fla/layers/mala.py
@@ -68,7 +68,7 @@ def __init__(
         use_lepe: bool = True,
         lepe_kernel_size: int = 5,
         layer_idx: int | None = None,
-    ) -> MalaAttention:
+    ) -> None:
         super().__init__()
 
         self.hidden_size = hidden_size
@@ -76,6 +76,7 @@ def __init__(
         self.expand_v = expand_v
         self.num_heads = num_heads
         self.num_kv_heads = num_kv_heads if num_kv_heads is not None else num_heads
+        assert self.num_heads % self.num_kv_heads == 0, f"num_heads must be divisible by num_kv_heads, got {self.num_heads} and {self.num_kv_heads}"
         self.num_kv_groups = self.num_heads // self.num_kv_heads
         self.use_lepe = use_lepe
         self.lepe_kernel_size = lepe_kernel_size

From 8d4e79c1a0e06354beb09f4e02b3a5fd806d93d6 Mon Sep 17 00:00:00 2001
From: Dr Daniel Wu <95159536+drdanielwuwu@users.noreply.github.com>
Date: Fri, 3 Apr 2026 12:43:31 +0800
Subject: [PATCH 15/16] Restore copyright and update MALA model comments

Restored copyright notice and updated comments in the MALA model implementation. Adjusted the handling of hidden states and attentions in the forward methods.
---
 fla/models/mala/modeling_mala.py | 22 ++++++++++++++--------
 1 file changed, 14 insertions(+), 8 deletions(-)

diff --git a/fla/models/mala/modeling_mala.py b/fla/models/mala/modeling_mala.py
index c1e3f77cec..9afeea9b0d 100644
--- a/fla/models/mala/modeling_mala.py
+++ b/fla/models/mala/modeling_mala.py
@@ -9,7 +9,6 @@
 from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
 
 from fla.layers import MalaAttention
-from fla.models.utils import cache_to_cpu, recycle_cache
 
 """
 MALA (Magnitude-Aware Linear Attention) model implementation.
@@ -87,13 +86,13 @@ def forward(
         hidden_states = self.embed_tokens(input_ids)
 
         # Process layers
-        all_hidden_states = () if output_hidden_states else None
-        all_attentions = () if output_attentions else None
-        next_decoder_cache = () if use_cache else None
+        all_hidden_states = [] if output_hidden_states else None
+        all_attentions = [] if output_attentions else None
+        next_decoder_cache = [] if use_cache else None
 
         for i, layer in enumerate(self.layers):
             if output_hidden_states:
-                all_hidden_states = all_hidden_states + (hidden_states,)
+                all_hidden_states.append(hidden_states)
 
             layer_past = past_key_values[i] if past_key_values is not None else None
 
@@ -108,16 +107,23 @@ def forward(
             hidden_states = layer_outputs[0]
 
             if use_cache:
-                next_decoder_cache += (layer_outputs[2],)
+                next_decoder_cache.append(layer_outputs[2])
 
             if output_attentions:
-                all_attentions = all_attentions + (layer_outputs[1],)
+                all_attentions.append(layer_outputs[1])
 
         # Apply final layer norm
         hidden_states = self.norm(hidden_states)
 
         if output_hidden_states:
-            all_hidden_states = all_hidden_states + (hidden_states,)
+            all_hidden_states.append(hidden_states)
+            all_hidden_states = tuple(all_hidden_states)
+
+        if use_cache:
+            next_decoder_cache = tuple(next_decoder_cache)
+
+        if output_attentions:
+            all_attentions = tuple(all_attentions)
 
         if not return_dict:
             return tuple(v for v in [hidden_states, next_decoder_cache, all_hidden_states, all_attentions] if v is not None)

From 4c8c5611920474805812d1774aa7843c87218b4e Mon Sep 17 00:00:00 2001
From: Dr Daniel Wu <95159536+drdanielwuwu@users.noreply.github.com>
Date: Fri, 3 Apr 2026 12:52:57 +0800
Subject: [PATCH 16/16] Restore copyright and documentation comments

---
 fla/ops/mala/naive.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/fla/ops/mala/naive.py b/fla/ops/mala/naive.py
index 11126abe10..39dfd13a6a 100644
--- a/fla/ops/mala/naive.py
+++ b/fla/ops/mala/naive.py
@@ -56,7 +56,7 @@ def naive_mala_attn(
     """
     dtype = q.dtype
     q, k, v = map(lambda x: x.float(), (q, k, v))
-    B, T, H, K, V = *q.shape, v.shape[-1]
+    _, T, _, K, _ = *q.shape, v.shape[-1]
     
     if scale is None:
         scale = K ** -0.5