Add fixed FLA folder for Quasar generation

README.md

@@ -2,7 +2,7 @@
 
 Repo-local Raven dependency and Quasar-Preview runner used with `silx-ai/Quasar-Preview` revision `48d9dc1d2dcdb2eeb087c7b79c89300d3e46a70b`.
 
-This packages the Raven implementation from `goombalab/raven` under `raven/` plus a simple script that wires it for Quasar greedy generation.
+This packages the Raven implementation from `goombalab/raven` under `raven/`, the fixed Quasar FLA folder under `fla/`, plus a simple script that wires them for Quasar greedy generation.
 
 ## Run
 
@@ -20,8 +20,7 @@ The runner handles:
 - restoring `experts_w3` tensors that Transformers may leave as meta tensors
 - running direct `model.generate()`
 
-The Raven/FLA generation-cache fix lives in `silx-ai/Quasar-Preview` under
-`fla/ops/gsa/fused_recurrent.py`.
+The Raven/FLA generation-cache fix is included here under `fla/ops/gsa/fused_recurrent.py` and has also been pushed to `silx-ai/Quasar-Preview`.
 
 ## Simpler script
 

fla/__init__.py

@@ -0,0 +1,8 @@
+# Lightweight package initializer for the vendored FLA subset used by Quasar.
+#
+# The upstream file eagerly imports every layer and model, which is slow and can
+# hang on fresh training containers while optional kernels are being resolved.
+# Import concrete modules directly, e.g. `from fla.layers.quasar import ...`.
+
+__version__ = "0.1.0"
+__all__ = []

fla/distributed_compat.py

@@ -0,0 +1,57 @@
+# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang
+"""
+Centralized compatibility module for torch.distributed imports.
+All distributed-related imports should go through here to handle environments
+where distributed tensor APIs are not available.
+"""
+
+import torch
+
+# DeviceMesh
+try:
+    from torch.distributed import DeviceMesh
+except ImportError:
+    try:
+        from torch.distributed.device_mesh import DeviceMesh
+    except ImportError:
+        DeviceMesh = None
+
+# DTensor
+try:
+    from torch.distributed.tensor import DTensor
+except (ImportError, AttributeError):
+    DTensor = None
+
+# Replicate, Shard, distribute_module, Placement
+try:
+    from torch.distributed.tensor import Placement, Replicate, Shard, distribute_module
+except (ImportError, AttributeError):
+    Placement = Replicate = Shard = distribute_module = None
+
+# ParallelStyle
+try:
+    from torch.distributed.tensor.parallel import ParallelStyle
+except (ImportError, AttributeError):
+    ParallelStyle = None
+
+# Convenience flag
+HAS_DISTRIBUTED = all([
+    DeviceMesh is not None,
+    DTensor is not None,
+    Placement is not None,
+    Replicate is not None,
+    Shard is not None,
+    distribute_module is not None,
+    ParallelStyle is not None,
+])
+
+__all__ = [
+    'DeviceMesh',
+    'DTensor',
+    'Placement',
+    'Replicate',
+    'Shard',
+    'distribute_module',
+    'ParallelStyle',
+    'HAS_DISTRIBUTED',
+]

fla/layers/__init__.py

@@ -0,0 +1,3 @@
+# Keep layer package imports lazy. Import specific layer modules directly.
+
+__all__ = []

fla/layers/abc.py

@@ -0,0 +1,231 @@
+# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang
+
+from __future__ import annotations
+
+import warnings
+from typing import TYPE_CHECKING
+
+import torch
+import torch.nn as nn
+from einops import rearrange
+
+from fla.layers.utils import get_layer_cache, update_layer_cache
+from fla.modules import FusedRMSNormGated, RMSNorm, RotaryEmbedding, ShortConvolution
+from fla.modules.activations import swiglu, swish
+from fla.ops.abc.chunk import chunk_abc
+
+if TYPE_CHECKING:
+    from fla.models.utils import Cache
+
+
+class ABCAttention(nn.Module):
+
+    def __init__(
+        self,
+        hidden_size: int = 1024,
+        expand_k: float = 0.5,
+        expand_v: float = 1.0,
+        num_heads: int = 4,
+        use_short_conv: bool = False,
+        conv_size: int = 4,
+        conv_bias: bool = False,
+        num_slots: int | None = None,
+        elementwise_affine: bool | None = True,
+        norm_eps: float = 1e-5,
+        gate_low_rank_dim: int = 16,
+        gate_logit_normalizer: int = 16,
+        use_rope: bool = True,
+        use_input_gate: bool = False,
+        use_output_gate: bool = True,
+        use_norm: bool = True,
+        clamp_min: float | None = -32,
+        clamp_max: float | None = 32,
+        layer_idx: int | None = None,
+        **kwargs,
+    ) -> ABCAttention:
+        super().__init__()
+
+        self.hidden_size = hidden_size
+        self.expand_k = expand_k
+        self.expand_v = expand_v
+        self.num_heads = num_heads
+        self.key_dim = int(self.hidden_size * self.expand_k)
+        self.value_dim = int(self.hidden_size * self.expand_v)
+        self.head_k_dim = self.key_dim // self.num_heads
+        self.head_v_dim = self.value_dim // self.num_heads
+
+        self.use_short_conv = use_short_conv
+        self.conv_size = conv_size
+        self.conv_bias = conv_bias
+
+        self.gate_low_rank_dim = gate_low_rank_dim
+        self.gate_logit_normalizer = gate_logit_normalizer
+
+        self.use_rope = use_rope
+        self.use_input_gate = use_input_gate
+        self.use_output_gate = use_output_gate
+        self.use_norm = use_norm
+
+        if num_slots is None:
+            num_slots = self.head_k_dim
+        self.num_slots = num_slots
+
+        self.norm_eps = norm_eps
+
+        self.clamp_min = clamp_min
+        self.clamp_max = clamp_max
+        self.layer_idx = layer_idx
+
+        if layer_idx is None:
+            warnings.warn(
+                f"Instantiating {self.__class__.__name__} without passing `layer_idx` is not recommended and will "
+                "to errors during the forward call, if caching is used. Please make sure to provide a `layer_idx` "
+                "when creating this class.",
+            )
+
+        self.q_proj = nn.Linear(self.hidden_size, self.key_dim, bias=False)
+        self.k_proj = nn.Linear(self.hidden_size, self.key_dim, bias=False)
+        self.v_proj = nn.Linear(self.hidden_size, self.value_dim, bias=False)
+
+        if use_output_gate:
+            self.g_proj = nn.Linear(self.hidden_size, self.value_dim, bias=False)
+        self.s_proj = nn.Linear(self.hidden_size, self.num_heads * self.num_slots, bias=False)
+        self.o_proj = nn.Linear(self.value_dim, self.hidden_size, bias=False)
+
+        if use_short_conv:
+            self.conv_size = conv_size
+            self.q_conv1d = ShortConvolution(
+                hidden_size=self.key_dim,
+                kernel_size=conv_size,
+                bias=conv_bias,
+                activation='silu',
+            )
+            self.k_conv1d = ShortConvolution(
+                hidden_size=self.key_dim,
+                kernel_size=conv_size,
+                bias=conv_bias,
+                activation='silu',
+            )
+            self.v_conv1d = ShortConvolution(
+                hidden_size=self.value_dim,
+                kernel_size=conv_size,
+                bias=conv_bias,
+                activation='silu',
+            )
+
+        if self.use_norm:
+            if self.use_output_gate:
+                self.g_norm = FusedRMSNormGated(
+                    hidden_size=self.head_v_dim,
+                    elementwise_affine=elementwise_affine,
+                    eps=norm_eps,
+                )
+            else:
+                self.g_norm = RMSNorm(
+                    hidden_size=self.head_v_dim,
+                    elementwise_affine=elementwise_affine,
+                    eps=norm_eps,
+                    dtype=torch.float32,
+                )
+
+        if self.use_rope:
+            self.rotary = RotaryEmbedding(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,
+    ) -> tuple[torch.Tensor, torch.Tensor | None, Cache | None]:
+        if attention_mask is not None:
+            assert len(attention_mask.shape) == 2, (
+                "Expected attention_mask as a 0-1 matrix with shape [batch_size, seq_len] "
+                "for padding purposes (0 indicating padding). "
+                "Arbitrary attention masks of shape [batch_size, seq_len, seq_len] are not allowed."
+            )
+
+        last_state = get_layer_cache(self, past_key_values)
+
+        cu_seqlens = kwargs.get('cu_seqlens')
+        if cu_seqlens is not None:
+            raise NotImplementedError("Training with cu_seqlens is not supported yet for ABCAttention")
+        if self.use_short_conv:
+            conv_state_q, conv_state_k, conv_state_v = None, None, None
+            if last_state is not None:
+                conv_state_q, conv_state_k, conv_state_v = last_state['conv_state']
+            conv_mask = attention_mask[:, -hidden_states.shape[1]:] if attention_mask is not None else None
+            q, conv_state_q = self.q_conv1d(
+                x=self.q_proj(hidden_states),
+                mask=conv_mask,
+                cache=conv_state_q,
+                output_final_state=use_cache,
+                cu_seqlens=cu_seqlens,
+            )
+            k, conv_state_k = self.k_conv1d(
+                x=self.k_proj(hidden_states),
+                mask=conv_mask,
+                cache=conv_state_k,
+                output_final_state=use_cache,
+                cu_seqlens=cu_seqlens,
+            )
+            v, conv_state_v = self.v_conv1d(
+                x=self.v_proj(hidden_states),
+                mask=conv_mask,
+                cache=conv_state_v,
+                output_final_state=use_cache,
+                cu_seqlens=cu_seqlens,
+            )
+        else:
+            q = self.q_proj(hidden_states)
+            k = self.k_proj(hidden_states)
+            v = self.v_proj(hidden_states)
+
+        if self.use_input_gate:
+            q, k, v = map(lambda x: swish(x), (q, k, v))
+        # dealing with left-padding
+        if attention_mask is not None:
+            v = v.mul_(attention_mask[:, -v.shape[-2]:, None])
+
+        q, k = map(lambda x: rearrange(x, '... (h d) -> ... h d', d=self.head_k_dim), (q, k))
+        v = rearrange(v, '... (h d) -> ... h d', d=self.head_v_dim)
+        if self.use_rope:
+            seqlen_offset = 0
+            if past_key_values is not None:
+                seqlen_offset = past_key_values.get_seq_length(self.layer_idx)
+            q, k = self.rotary(q, k, seqlen_offset=seqlen_offset)
+
+        s = rearrange(self.s_proj(hidden_states), '... (h m) -> ... h m', m=self.num_slots)
+        s = s.clamp_(self.clamp_min, self.clamp_max)
+
+        recurrent_state = last_state['recurrent_state'] if last_state is not None else None
+        o, recurrent_state = chunk_abc(
+            q=q,
+            k=k,
+            v=v,
+            s=s,
+            initial_state=recurrent_state,
+            output_final_state=use_cache,
+        )
+        update_layer_cache(
+            self,
+            past_key_values,
+            recurrent_state=recurrent_state,
+            conv_state=(conv_state_q, conv_state_k, conv_state_v) if self.use_short_conv else None,
+            offset=q.shape[1],
+        )
+
+        if self.use_norm and not self.use_output_gate:
+            o = self.g_norm(o)
+        elif self.use_output_gate:
+            g = rearrange(self.g_proj(hidden_states), '... (h d) -> ... h d', d=self.head_v_dim)
+            o = self.g_norm(o, g) if self.use_norm else swiglu(g, o)
+        o = rearrange(o, '... h d -> ... (h d)')
+        o = self.o_proj(o)
+
+        return o, None, past_key_values
+
+    def state_size(self, seq_len: int = 2048):
+        return 2 * self.num_slots * self.hidden_size

fla/layers/attn.py

@@ -0,0 +1,176 @@
+# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang
+
+from __future__ import annotations
+
+import warnings
+from typing import TYPE_CHECKING
+
+import torch
+import torch.nn as nn
+from einops import rearrange
+from transformers.utils import logging
+
+from fla.layers.utils import pad_input, unpad_input
+from fla.modules import RMSNorm, RotaryEmbedding
+from fla.ops.utils.index import prepare_lens_from_mask
+
+if TYPE_CHECKING:
+    from fla.models.utils import Cache
+
+try:
+    from flash_attn import flash_attn_func, flash_attn_varlen_func
+except ImportError:
+    warnings.warn(
+        "Flash Attention is not installed. Please install it via `pip install flash-attn --no-build-isolation`",
+        category=ImportWarning,
+    )
+    flash_attn_func = None
+
+logger = logging.get_logger(__name__)
+
+
+class Attention(nn.Module):
+
+    def __init__(
+        self,
+        hidden_size: int = 2048,
+        num_heads: int = 32,
+        num_kv_heads: int | None = None,
+        qkv_bias: bool = False,
+        qk_norm: bool = False,
+        window_size: int | None = None,
+        rope_theta: float | None = 10000.,
+        max_position_embeddings: int | None = None,
+        use_nope: bool = False,
+        layer_idx: int = None,
+    ):
+        super().__init__()
+
+        self.hidden_size = hidden_size
+        self.num_heads = num_heads
+        if num_kv_heads is None:
+            self.num_kv_heads = self.num_heads
+        else:
+            self.num_kv_heads = num_kv_heads
+        self.num_kv_groups = num_heads // self.num_kv_heads
+        self.head_dim = self.hidden_size // self.num_heads
+        self.kv_dim = self.num_kv_heads * self.head_dim
+        self.qkv_bias = qkv_bias
+        self.qk_norm = qk_norm
+
+        self.window_size = window_size
+        self.rope_theta = rope_theta
+        self.max_position_embeddings = max_position_embeddings
+        self.use_nope = use_nope
+        self.layer_idx = layer_idx
+
+        if flash_attn_func is None:
+            raise ImportError("Please install Flash Attention via `pip install flash-attn --no-build-isolation` first")
+
+        self.q_proj = nn.Linear(self.hidden_size, self.hidden_size, bias=self.qkv_bias)
+        self.k_proj = nn.Linear(self.hidden_size, self.kv_dim, bias=self.qkv_bias)
+        self.v_proj = nn.Linear(self.hidden_size, self.kv_dim, bias=self.qkv_bias)
+        self.o_proj = nn.Linear(self.hidden_size, self.hidden_size, bias=False)
+
+        if qk_norm:
+            self.q_norm = RMSNorm(self.head_dim, dtype=torch.float32)
+            self.k_norm = RMSNorm(self.head_dim, dtype=torch.float32)
+
+        self.rotary = RotaryEmbedding(dim=self.head_dim, base=self.rope_theta)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: torch.LongTensor | None = None,
+        past_key_values: Cache | None = None,
+        output_attentions: bool = False,
+        use_cache: bool = False,
+        **kwargs,
+    ) -> tuple[torch.Tensor, torch.Tensor | None, tuple[torch.Tensor] | None]:
+        if attention_mask is not None:
+            assert len(attention_mask.shape) == 2, (
+                "Expected attention_mask as a 0-1 matrix with shape [batch_size, seq_len] "
+                "for padding purposes (0 indicating padding). "
+                "Arbitrary attention masks of shape [batch_size, seq_len, seq_len] are not allowed."
+            )
+
+        batch_size, q_len, _ = hidden_states.size()
+
+        q = rearrange(self.q_proj(hidden_states), '... (h d) -> ... h d', d=self.head_dim)
+        k = rearrange(self.k_proj(hidden_states), '... (h d) -> ... h d', d=self.head_dim)
+        v = rearrange(self.v_proj(hidden_states), '... (h d) -> ... h d', d=self.head_dim)
+
+        if self.qk_norm:
+            q, k = self.q_norm(q), self.k_norm(k)
+
+        # equivalent to cu_seqlens in `flash_attn`
+        cu_seqlens = kwargs.get('cu_seqlens')
+
+        seqlen_offset, max_seqlen = 0, q_len
+        if past_key_values is not None:
+            seqlen_offset = past_key_values.get_seq_length(self.layer_idx)
+            max_seqlen = q.shape[1] + seqlen_offset
+
+            if attention_mask is not None:
+                # to deliminate the offsets of padding tokens
+                seqlen_offset = seqlen_offset + prepare_lens_from_mask(attention_mask) - attention_mask.shape[-1]
+                max_seqlen = q.shape[1] + max(seqlen_offset)
+
+        if self.max_position_embeddings is not None:
+            max_seqlen = max(max_seqlen, self.max_position_embeddings)
+        if not self.use_nope:
+            q, k = self.rotary(q, k, seqlen_offset=seqlen_offset, max_seqlen=max_seqlen, cu_seqlens=cu_seqlens)
+
+        if past_key_values is not None:
+            cache_has_content = past_key_values.get_seq_length(self.layer_idx) > 0
+            k_cached, v_cached = past_key_values.update(
+                attn_state=(k.flatten(-2, -1), v.flatten(-2, -1)),
+                layer_idx=self.layer_idx,
+                offset=q_len,
+                cache_kwargs=dict(window_size=self.window_size),
+            )['attn_state']
+            if cache_has_content:
+                k, v = k_cached, v_cached
+                k = rearrange(k, '... (h d) -> ... h d', d=self.head_dim)
+                v = rearrange(v, '... (h d) -> ... h d', d=self.head_dim)
+
+        # Contains at least one padding token in the sequence
+        if attention_mask is not None:
+            if q.shape[1] == 1 and self.window_size is not None:
+                attention_mask = attention_mask[:, -self.window_size:]
+            q, (k, v), indices_q, cu_seqlens, max_seq_lens = unpad_input(q, (k, v), attention_mask, q_len)
+            cu_seqlens_q, cu_seqlens_k = cu_seqlens
+            max_seqlen_q, max_seqlen_k = max_seq_lens
+            o = flash_attn_varlen_func(
+                q, k, v,
+                cu_seqlens_q=cu_seqlens_q,
+                cu_seqlens_k=cu_seqlens_k,
+                max_seqlen_q=max_seqlen_q,
+                max_seqlen_k=max_seqlen_k,
+                causal=True,
+                window_size=(-1, -1) if self.window_size is None else (self.window_size-1, 0),
+            )
+            o = pad_input(o, indices_q, batch_size, q_len)
+        elif cu_seqlens is not None:
+            o = flash_attn_varlen_func(
+                q.squeeze(0), k.squeeze(0), v.squeeze(0),
+                cu_seqlens_q=cu_seqlens,
+                cu_seqlens_k=cu_seqlens,
+                max_seqlen_q=max_seqlen,
+                max_seqlen_k=max_seqlen,
+                causal=True,
+                window_size=(-1, -1) if self.window_size is None else (self.window_size-1, 0),
+            ).unsqueeze(0)
+        else:
+            o = flash_attn_func(
+                q, k, v,
+                causal=True,
+                window_size=(-1, -1) if self.window_size is None else (self.window_size-1, 0),
+            )
+        o = o.reshape(batch_size, q_len, -1)
+        o = self.o_proj(o)
+
+        if not output_attentions:
+            attentions = None
+
+        return o, attentions, past_key_values

fla/layers/based.py

@@ -0,0 +1,93 @@
+# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang
+
+"""
+Linear attention in Based.
+https://github.com/HazyResearch/zoology/blob/main/zoology/mixers/based.py
+"""
+
+import torch
+import torch.nn as nn
+from einops import rearrange
+
+from fla.modules.feature_map import TaylorFeatureMap
+from fla.ops.based import parallel_based
+from fla.ops.linear_attn import chunk_linear_attn, fused_chunk_linear_attn
+
+
+class BasedLinearAttention(nn.Module):
+
+    def __init__(
+        self,
+        hidden_size: int,
+        feature_dim: int = 16,
+        num_key_value_heads: int = 12,
+        num_heads: int = 12,
+        feature_name: str = "taylor_exp",
+        eps: float = 1e-12,
+        causal: bool = True,
+        mode: str = "parallel",
+    ):
+        super().__init__()
+
+        self.hidden_size = hidden_size
+        self.mode = mode
+        self.feature_name = feature_name
+        self.feature_dim = feature_dim
+        self.num_key_value_heads = num_key_value_heads
+        self.num_heads = num_heads
+        self.head_dim = self.hidden_size // self.num_key_value_heads
+        assert self.hidden_size % self.head_dim == 0
+        self.causal = causal
+
+        self.q_proj = nn.Linear(self.hidden_size, self.feature_dim * self.num_heads, bias=False)
+        self.k_proj = nn.Linear(self.hidden_size, self.feature_dim * self.num_heads, bias=False)
+        self.v_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=False)
+        self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=False)
+        self.dropout = nn.Identity()
+        self.feature_map = TaylorFeatureMap(feature_dim)
+        self.eps = eps
+
+    def forward(self, hidden_states: torch.Tensor, **kwargs):
+        mode = self.mode
+        q, k, v = self.q_proj(hidden_states), self.k_proj(hidden_states), self.v_proj(hidden_states)
+        q, k, v = map(lambda x: rearrange(x, "... (h d) -> ... h d", d=self.head_dim), [q, k, v])
+        if mode == "fused_chunk":
+            q, k = self.feature_map(q), self.feature_map(k)
+            o, _ = fused_chunk_linear_attn(q, k, v, normalize=True, scale=1)
+        elif mode == 'chunk':
+            q, k = self.feature_map(q), self.feature_map(k)
+            o, _ = chunk_linear_attn(q, k, v, normalize=True, scale=1)
+        elif mode == 'parallel':
+            assert q.shape[-1] <= 128
+            o = parallel_based(q, k, v, scale=1, use_norm=True)
+        o = rearrange(o, 'b t h d -> b t (h d)')
+        o = self.o_proj(o)
+        o = self.dropout(o)
+        return o
+
+    def forward_reference(self, hidden_states: torch.Tensor, **kwargs):
+        """
+        x (torch.Tensor): tensor of shape (b, d, t)
+        y (torch.Tensor): tensor of shape (b, d, t)
+        """
+        # hidden_states = hidden_states.transpose(1, 2)
+        b, t, _ = hidden_states.size()
+        q, k, v = self.q_proj(hidden_states), self.k_proj(hidden_states), self.v_proj(hidden_states)
+
+        q = q.view(b, t, self.num_heads, self.feature_dim).transpose(1, 2)
+        k = k.view(b, t, self.num_key_value_heads, self.feature_dim).transpose(1, 2)
+        v = v.view(b, t, self.num_key_value_heads, self.head_dim).transpose(1, 2)
+
+        # Linear attention
+        q, k = self.feature_map(q), self.feature_map(k)
+        q, k, v = q.unsqueeze(-2), k.unsqueeze(-2), v.unsqueeze(-1)
+
+        # Compute attention
+        if self.causal:
+            y = ((q * (k * v).cumsum(2)).sum(-1) / ((q * k.cumsum(2)).sum(-1) + self.eps))
+        else:
+            y = ((q * (k * v).sum(2, True)).sum(-1) / ((q * k.sum(2, True)).sum(-1) + self.eps))
+        y = rearrange(y, 'b h t d -> b t (h d)')
+        y = self.o_proj(y.to(hidden_states.dtype))
+        y = self.dropout(y)
+        return y.to(hidden_states.dtype)

fla/layers/bitattn.py

@@ -0,0 +1,162 @@
+# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang
+
+from __future__ import annotations
+
+import warnings
+from typing import TYPE_CHECKING
+
+import torch
+import torch.nn as nn
+from einops import rearrange
+from transformers.utils import logging
+
+from fla.layers.utils import pad_input, unpad_input
+from fla.modules import RotaryEmbedding
+from fla.modules.fused_bitlinear import FusedBitLinear
+from fla.ops.utils.index import prepare_lens_from_mask
+
+if TYPE_CHECKING:
+    from fla.models.utils import Cache
+
+try:
+    from flash_attn import flash_attn_func, flash_attn_varlen_func
+except ImportError:
+    warnings.warn(
+        "Flash Attention is not installed. Please install it via `pip install flash-attn --no-build-isolation`",
+        category=ImportWarning,
+    )
+    flash_attn_func = None
+
+logger = logging.get_logger(__name__)
+
+
+class BitAttention(nn.Module):
+
+    def __init__(
+        self,
+        hidden_size: int = 2048,
+        num_heads: int = 32,
+        num_kv_heads: int | None = None,
+        window_size: int | None = None,
+        rope_theta: float | None = 10000.,
+        max_position_embeddings: int | None = None,
+        norm_eps: float = 1e-5,
+        layer_idx: int = None,
+    ):
+        super().__init__()
+
+        self.num_heads = num_heads
+        if num_kv_heads is None:
+            self.num_kv_heads = self.num_heads
+        else:
+            self.num_kv_heads = num_kv_heads
+        self.num_kv_groups = num_heads // self.num_kv_heads
+        self.hidden_size = hidden_size
+        self.head_dim = self.hidden_size // self.num_heads
+        self.kv_dim = self.num_kv_heads * self.head_dim
+        self.kv_dim = self.num_kv_heads * self.head_dim
+        self.window_size = window_size
+        self.rope_theta = rope_theta
+        self.max_position_embeddings = max_position_embeddings
+        self.layer_idx = layer_idx
+
+        self.q_proj = FusedBitLinear(self.hidden_size, self.hidden_size, bias=False)
+        self.k_proj = FusedBitLinear(self.hidden_size, self.kv_dim, bias=False)
+        self.v_proj = FusedBitLinear(self.hidden_size, self.kv_dim, bias=False)
+        self.o_proj = FusedBitLinear(self.hidden_size, self.hidden_size, bias=False)
+
+        self.rotary = RotaryEmbedding(dim=self.head_dim, base=self.rope_theta)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: torch.LongTensor | None = None,
+        past_key_values: Cache | None = None,
+        output_attentions: bool = False,
+        use_cache: bool = False,
+        **kwargs,
+    ) -> tuple[torch.Tensor, torch.Tensor | None, tuple[torch.Tensor] | None]:
+        if attention_mask is not None:
+            assert len(attention_mask.shape) == 2, (
+                "Expected attention_mask as a 0-1 matrix with shape [batch_size, seq_len] "
+                "for padding purposes (0 indicating padding). "
+                "Arbitrary attention masks of shape [batch_size, seq_len, seq_len] are not allowed."
+            )
+
+        batch_size, q_len, _ = hidden_states.size()
+
+        q = rearrange(self.q_proj(hidden_states), '... (h d) -> ... h d', d=self.head_dim)
+        k = rearrange(self.k_proj(hidden_states), '... (h d) -> ... h d', d=self.head_dim)
+        v = rearrange(self.v_proj(hidden_states), '... (h d) -> ... h d', d=self.head_dim)
+
+        # equivalent to cu_seqlens in `flash_attn`
+        cu_seqlens = kwargs.get('cu_seqlens')
+
+        seqlen_offset, max_seqlen = 0, q_len
+        if past_key_values is not None:
+            seqlen_offset = past_key_values.get_seq_length(self.layer_idx)
+            max_seqlen = q.shape[1] + seqlen_offset
+
+            if attention_mask is not None:
+                # to deliminate the offsets of padding tokens
+                seqlen_offset = seqlen_offset + prepare_lens_from_mask(attention_mask) - attention_mask.shape[-1]
+                max_seqlen = q.shape[1] + max(seqlen_offset)
+
+        if self.max_position_embeddings is not None:
+            max_seqlen = max(max_seqlen, self.max_position_embeddings)
+        q, k = self.rotary(q, k, seqlen_offset=seqlen_offset, max_seqlen=max_seqlen, cu_seqlens=cu_seqlens)
+
+        if past_key_values is not None:
+            cache_has_content = past_key_values.get_seq_length(self.layer_idx) > 0
+            k_cached, v_cached = past_key_values.update(
+                attn_state=(k.flatten(-2, -1), v.flatten(-2, -1)),
+                layer_idx=self.layer_idx,
+                offset=q_len,
+                cache_kwargs=dict(window_size=self.window_size),
+            )['attn_state']
+            if cache_has_content:
+                k, v = k_cached, v_cached
+                k = rearrange(k, '... (h d) -> ... h d', d=self.head_dim)
+                v = rearrange(v, '... (h d) -> ... h d', d=self.head_dim)
+
+        if flash_attn_func is None:
+            raise ImportError("Please install Flash Attention via `pip install flash-attn --no-build-isolation` first")
+
+        # Contains at least one padding token in the sequence
+        if attention_mask is not None:
+            q, (k, v), indices_q, cu_seqlens, max_seq_lens = unpad_input(q, (k, v), attention_mask, q_len)
+            cu_seqlens_q, cu_seqlens_k = cu_seqlens
+            max_seqlen_q, max_seqlen_k = max_seq_lens
+            o = flash_attn_varlen_func(
+                q, k, v,
+                cu_seqlens_q=cu_seqlens_q,
+                cu_seqlens_k=cu_seqlens_k,
+                max_seqlen_q=max_seqlen_q,
+                max_seqlen_k=max_seqlen_k,
+                causal=True,
+                window_size=(-1, -1) if self.window_size is None else (self.window_size-1, 0),
+            )
+            o = pad_input(o, indices_q, batch_size, q_len)
+        elif cu_seqlens is not None:
+            o = flash_attn_varlen_func(
+                q.squeeze(0), k.squeeze(0), v.squeeze(0),
+                cu_seqlens_q=cu_seqlens,
+                cu_seqlens_k=cu_seqlens,
+                max_seqlen_q=max_seqlen,
+                max_seqlen_k=max_seqlen,
+                causal=True,
+                window_size=(-1, -1) if self.window_size is None else (self.window_size-1, 0),
+            ).unsqueeze(0)
+        else:
+            o = flash_attn_func(
+                q, k, v,
+                causal=True,
+                window_size=(-1, -1) if self.window_size is None else (self.window_size-1, 0),
+            )
+        o = o.reshape(batch_size, q_len, -1)
+        o = self.o_proj(o)
+
+        if not output_attentions:
+            attentions = None
+
+        return o, attentions, past_key_values

fla/layers/comba.py

@@ -0,0 +1,328 @@
+# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang
+
+from __future__ import annotations
+
+import math
+import warnings
+from typing import TYPE_CHECKING
+
+import torch
+import torch.nn as nn
+from einops import rearrange, repeat
+from torch.nn import functional as F
+
+from fla.layers.utils import get_layer_cache, get_unpad_data, index_first_axis, pad_input, update_layer_cache
+from fla.modules import FusedRMSNormGated, RMSNorm, ShortConvolution
+from fla.ops.comba import chunk_comba, fused_recurrent_comba
+
+if TYPE_CHECKING:
+    from transformers.processing_utils import Unpack
+
+    from fla.models.utils import Cache
+
+
+class Comba(nn.Module):
+    """
+    The layer implementaion for [Comba: Improving Bilinear RNNs with Closed-loop Control](https://arxiv.org/abs/2506.02475).
+
+    Similar to Mamba2 and Gated-DeltaNet, each layer contains around 6*hidden_size*hidden_size parameters.
+
+    Parameter alloation when use_output_gate=True:
+        - 0.75 * hidden_size * hidden_size for the q_proj and k_proj each
+        - 1.5 * hidden_size * hidden_size for the v_proj, g_proj and o_proj each
+        - Others are ignorably small.
+        - In total = 0.75 * 2 + 1.5 * 3 = 6 * hidden_size * hidden_size
+    NOTE: num_heads * head_dim = 0.75 * hidden_size, please make sure to set the correct num_heads and head_dim.
+
+    Parameter allocation when use_output_gate=False:
+        - 1 * hidden_size * hidden_size for the q_proj and k_proj each
+        - 2 * hidden_size * hidden_size for the v_proj and o_proj each
+        - Others are ignorably small.
+        - In total = 1 * 2 + 2 * 2 = 6 * hidden_size * hidden_size
+
+    Args:
+        hidden_size (int, Optional):
+            The hidden size of the input. Default: 2048.
+        expand_v (float, Optional):
+            The expansion ratio for the value dim. Default: 2.0.
+        head_dim (int, Optional):
+            The dimension of each head. Default: 256.
+        num_heads (int, Optional):
+            The number of heads. Default: 4.
+        num_v_heads (int, Optional):
+            The number of heads for the value projection, equal to `num_heads` if `None`.
+            GVA is applied if `num_v_heads` > `num_heads`. Default: `None`.
+        mode (str, Optional):
+            Which Gated DeltaNet kernel to use.
+            Currently available: `chunk` and `fused_recurrent`.
+            Default: `chunk`.
+        use_beta (bool, Optional):
+            Whether to use beta. Default: `True`.
+        use_output_gate (bool, Optional):
+            Whether to use output gate. Default: `True`.
+        use_output_correction (bool, Optional):
+            Whether to use <q-dk>. Default: `True`.
+        use_short_conv (bool, Optional):
+            Whether to use short convolutions. Default: `True`.
+        conv_size (int, Optional):
+            The kernel size of the short convolution, only used when `use_short_conv` is `True`. Default: 4.
+        conv_bias (bool, Optional):
+            Whether to use bias in the short convolution, only used when `use_short_conv` is `True`. Default: `False`.
+        layer_idx (int, Optional):
+            The index of the layer. Default: None.
+        norm_eps (float, Optional):
+            The epsilon value for the normalization layer. Default: 1e-5.
+    """
+
+    def __init__(
+        self,
+        hidden_size: int = 2048,
+        expand_v: float = 2,
+        head_dim: int = 256,
+        num_heads: int = 6,
+        num_v_heads: int = None,
+        mode: str = 'chunk',
+        use_short_conv: bool = True,
+        use_output_gate: bool = True,
+        use_output_correction: bool = True,
+        use_inner_decay: bool = True,
+        correction_factor: float = 1.,
+        conv_size: int = 4,
+        conv_bias: bool = False,
+        layer_idx: int = None,
+        norm_eps: float = 1e-5,
+        **kwargs,
+    ) -> Comba:
+        super().__init__()
+
+        self.mode = mode
+
+        self.hidden_size = hidden_size
+        self.expand_v = expand_v
+
+        self.use_short_conv = use_short_conv
+        self.use_output_gate = use_output_gate
+        self.use_output_correction = use_output_correction
+        self.use_inner_decay = use_inner_decay
+        self.conv_size = conv_size
+        self.conv_bias = conv_bias
+
+        self.head_dim = head_dim
+        self.num_heads = num_heads
+        self.num_v_heads = num_v_heads if num_v_heads is not None else num_heads
+
+        self.head_k_dim = head_dim
+        self.head_v_dim = int(self.head_dim * self.expand_v)
+        self.key_dim = int(self.num_heads * self.head_k_dim)
+        self.value_dim = int(self.num_v_heads * self.head_v_dim)
+        self.layer_idx = layer_idx
+
+        # Consistency check: Ensure expand_v produces integer values
+        if not math.isclose(self.num_v_heads * self.head_dim * expand_v, self.value_dim, rel_tol=1e-5):
+            raise ValueError(
+                f"expand_v={expand_v} does not produce an integer value when multiplied by key_dim={self.key_dim}. "
+                f"Resulting value_dim would be {self.num_v_heads * self.head_dim * expand_v}, which is invalid for nn.Linear.",
+            )
+        if self.num_v_heads > self.num_heads and self.num_v_heads % self.num_heads != 0:
+            raise ValueError(
+                f"num_v_heads={self.num_v_heads} must be divisible by num_heads={self.num_heads}.",
+            )
+
+        if not math.isclose(head_dim * expand_v, self.head_v_dim, rel_tol=1e-5):
+            raise ValueError(
+                f"expand_v={expand_v} does not produce an integer value when multiplied by head_dim={head_dim}. "
+                f"Resulting head_v_dim would be {head_dim * expand_v}, which is invalid for FusedRMSNormGated.",
+            )
+        assert mode in ['chunk', 'fused_recurrent'], f"Not supported mode `{mode}`."
+
+        self.q_proj = nn.Linear(hidden_size, self.key_dim, bias=False)
+        self.k_proj = nn.Linear(hidden_size, self.key_dim, bias=False)
+        self.v_proj = nn.Linear(hidden_size, self.value_dim, bias=False)
+        self.a_proj = nn.Linear(hidden_size, self.num_v_heads, bias=False)
+        self.b_proj = nn.Linear(hidden_size, self.num_v_heads, bias=False)
+
+        if use_inner_decay:
+            self.decay = nn.Parameter(torch.ones(self.num_heads))
+
+        if use_output_correction:
+            warnings.warn(
+                "The correction_factor is set to 1 by default similar to Mamba2. "
+                "However, we find that sometimes correction_factor = 0.02 works better for small-scale models. "
+                "In practice, we recommend trying both settings. ",
+            )
+            self.D = nn.Parameter(torch.ones(self.num_heads) * correction_factor)
+            self.D._no_weight_decay = True
+
+        A = torch.empty(self.num_v_heads, dtype=torch.float32).uniform_(0, 16)
+        self.A_log = nn.Parameter(torch.log(A))
+        self.A_log._no_weight_decay = True
+        # hard coded for now
+        dt_min = 0.001
+        dt_max = 0.1
+        dt_init_floor = 1e-4
+        dt = torch.exp(
+            torch.rand(self.num_v_heads) * (math.log(dt_max) - math.log(dt_min))
+            + math.log(dt_min),
+        )
+        dt = torch.clamp(dt, min=dt_init_floor)
+        # Inverse of softplus: https://github.com/pytorch/pytorch/issues/72759
+        inv_dt = dt + torch.log(-torch.expm1(-dt))
+        self.dt_bias = nn.Parameter(inv_dt)
+        # Just to be explicit. Without this we already don't put wd on dt_bias because of the check
+        # name.endswith("bias") in param_grouping.py
+        self.dt_bias._no_weight_decay = True
+
+        if use_short_conv:
+            self.conv_size = conv_size
+            self.q_conv1d = ShortConvolution(
+                hidden_size=self.key_dim,
+                kernel_size=conv_size,
+                bias=conv_bias,
+                activation='silu',
+            )
+            self.k_conv1d = ShortConvolution(
+                hidden_size=self.key_dim,
+                kernel_size=conv_size,
+                bias=conv_bias,
+                activation='silu',
+            )
+            self.v_conv1d = ShortConvolution(
+                hidden_size=self.value_dim,
+                kernel_size=conv_size,
+                bias=conv_bias,
+                activation='silu',
+            )
+        else:
+            warnings.warn(
+                "ShortConvolution is crucial to the performance. "
+                "Do not turn it off, i.e., setting `use_short_conv=False` unless you know what you are doing.",
+            )
+        if use_output_gate:
+            self.g_proj = nn.Linear(hidden_size, self.value_dim, bias=False)
+            self.o_norm = FusedRMSNormGated(self.head_v_dim, activation='sigmoid', eps=norm_eps)
+        else:
+            self.o_norm = RMSNorm(self.head_v_dim, eps=norm_eps, dtype=torch.float32)
+        self.o_proj = nn.Linear(self.value_dim, hidden_size, bias=False)
+
+    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]:
+        if attention_mask is not None:
+            assert len(attention_mask.shape) == 2, (
+                "Expected attention_mask as a 0-1 matrix with shape [batch_size, seq_len] "
+                "for padding purposes (0 indicating padding). "
+                "Arbitrary attention masks of shape [batch_size, seq_len, seq_len] are not allowed."
+            )
+
+        batch_size, q_len, _ = hidden_states.shape
+        # change to inference mode.
+        mode = 'fused_recurrent' if (q_len <= 64 and not self.training) else self.mode
+        if self.training:
+            assert mode == 'chunk', "Only chunk mode is supported in training."
+        last_state = get_layer_cache(self, past_key_values)
+
+        cu_seqlens = kwargs.get('cu_seqlens')
+        if attention_mask is not None:
+            indices, cu_seqlens, _ = get_unpad_data(attention_mask[:, -q_len:])
+            hidden_states = index_first_axis(rearrange(hidden_states, "b s ... -> (b s) ..."), indices).unsqueeze(0)
+
+        if self.use_short_conv:
+            conv_state_q, conv_state_k, conv_state_v = None, None, None
+            if last_state is not None:
+                conv_state_q, conv_state_k, conv_state_v = last_state['conv_state']
+            q, conv_state_q = self.q_conv1d(
+                x=self.q_proj(hidden_states),
+                cache=conv_state_q,
+                output_final_state=use_cache,
+                cu_seqlens=cu_seqlens,
+            )
+            k, conv_state_k = self.k_conv1d(
+                x=self.k_proj(hidden_states),
+                cache=conv_state_k,
+                output_final_state=use_cache,
+                cu_seqlens=cu_seqlens,
+            )
+            v, conv_state_v = self.v_conv1d(
+                x=self.v_proj(hidden_states),
+                cache=conv_state_v,
+                output_final_state=use_cache,
+                cu_seqlens=cu_seqlens,
+            )
+        else:
+            q = F.silu(self.q_proj(hidden_states))
+            k = F.silu(self.k_proj(hidden_states))
+            v = F.silu(self.v_proj(hidden_states))
+
+        q, k = map(lambda x: rearrange(x, '... (h d) -> ... h d', d=self.head_k_dim), (q, k))
+
+        if self.use_inner_decay:
+            p = k * self.decay[None, None, :, None].sigmoid()
+        else:
+            p = k
+
+        if self.use_output_correction:
+            q = q - self.D[None, None, :, None] * p
+
+        v = rearrange(v, '... (h d) -> ... h d', d=self.head_v_dim)
+
+        if self.num_v_heads > self.num_heads:
+            q, k = map(lambda x: repeat(x, '... h d -> ... (h g) d', g=self.num_v_heads // self.num_heads), (q, k))
+
+        beta = self.b_proj(hidden_states).sigmoid()
+        g = -self.A_log.float().exp() * F.softplus(self.a_proj(hidden_states).float() + self.dt_bias)
+
+        recurrent_state = last_state['recurrent_state'] if last_state is not None else None
+        if mode == 'chunk':
+            o, recurrent_state = chunk_comba(
+                q=q,
+                k=k,
+                v=v,
+                p=p,
+                g=g,
+                beta=beta,
+                initial_state=recurrent_state,
+                output_final_state=use_cache,
+                cu_seqlens=cu_seqlens,
+                use_qk_l2norm_in_kernel=True,
+            )
+        elif mode == 'fused_recurrent':
+            o, recurrent_state = fused_recurrent_comba(
+                q=q,
+                k=k,
+                v=v,
+                p=p,
+                g=g,
+                beta=beta,
+                initial_state=recurrent_state,
+                output_final_state=use_cache,
+                cu_seqlens=cu_seqlens,
+                use_qk_l2norm_in_kernel=True,
+            )
+        else:
+            raise NotImplementedError(f"Not supported mode `{mode}`.")
+
+        update_layer_cache(
+            self,
+            past_key_values,
+            recurrent_state=recurrent_state,
+            conv_state=(conv_state_q, conv_state_k, conv_state_v) if self.use_short_conv else None,
+            offset=q_len,
+        )
+
+        if self.use_output_gate:
+            g = rearrange(self.g_proj(hidden_states), '... (h d) -> ... h d', d=self.head_v_dim)
+            o = self.o_norm(o, g)
+        else:
+            o = self.o_norm(o)
+        o = rearrange(o, 'b t h d -> b t (h d)')
+        o = self.o_proj(o)
+        if attention_mask is not None:
+            o = pad_input(o.squeeze(0), indices, batch_size, q_len)
+
+        return o, None, past_key_values

fla/layers/delta_net.py

@@ -0,0 +1,287 @@
+# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang
+
+from __future__ import annotations
+
+import warnings
+from typing import TYPE_CHECKING
+
+import torch
+import torch.nn as nn
+from einops import rearrange
+from torch.nn import functional as F
+
+from fla.layers.utils import get_layer_cache, get_unpad_data, index_first_axis, pad_input, update_layer_cache
+from fla.modules import FusedRMSNormGated, RMSNorm, ShortConvolution
+from fla.ops.delta_rule import chunk_delta_rule, fused_recurrent_delta_rule
+
+if TYPE_CHECKING:
+    from transformers.processing_utils import Unpack
+
+    from fla.models.utils import Cache
+
+
+def elu_p1(x):
+    return (F.elu(x, 1., False) + 1.).to(x)
+
+
+def sum_norm(x):
+    return (x / x.sum(-1, keepdim=True)).to(x)
+
+
+class DeltaNet(nn.Module):
+    r"""
+    The layer implementaion for [Parallelizing Linear Transformers with the Delta Rule over Sequence Length](https://arxiv.org/abs/2406.06484).  # noqa:
+    DeltaNet was originally proposed in [Linear Transformers Are Secretly Fast Weight Programmers](https://arxiv.org/abs/2102.11174). # noqa
+
+    Args:
+        mode (str, Optional):
+            Which DeltaNet kernel to use.
+            Currently available: `chunk`, `fused_recurrent`, and `fused_chunk`.
+            Default: `chunk`.
+        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.
+        use_beta (bool, Optional):
+            Whether to use beta. Default: `True`.
+        use_gate (bool, Optional):
+            Whether to use output gate. Default: `False`.
+        use_short_conv (bool, Optional):
+            Whether to use short convolutions. Default: `True`.
+        conv_size (int, Optional):
+            The kernel size of the short convolution, only used when `use_short_conv` is `True`. Default: 4.
+        conv_bias (bool, Optional):
+            Whether to use bias in the short convolution, only used when `use_short_conv` is `True`. Default: `False`.
+        allow_neg_eigval (bool, Optional):
+            Allow negative eigenvalues. Default: `False`. If set to `True`, the beta will be multiplied by 2.
+            See reference: [Unlocking State-Tracking in Linear RNNs Through Negative Eigenvalues](https://arxiv.org/abs/2411.12537)
+        layer_idx (int, Optional):
+            The index of the layer. Default: None.
+        norm_eps (float, Optional):
+            The epsilon value for the layernorm/rmsnorm layer. Default: 1e-5.
+        qk_activation (str, Optional):
+            The activation function for the query and key. Default: `silu`.
+        qk_norm (str, Optional):
+            The normalization method for the query and key. Default: `l2`.
+    """
+
+    def __init__(
+        self,
+        mode: str = 'chunk',
+        d_model: int = None,
+        hidden_size: int = 1024,
+        expand_k: float = 1.0,
+        expand_v: float = 1.0,
+        num_heads: int = 4,
+        use_beta: bool = True,
+        use_gate: bool = False,
+        use_short_conv: bool = True,
+        conv_size: int = 4,
+        conv_bias: bool = False,
+        allow_neg_eigval: bool = False,
+        layer_idx: int = None,
+        qk_activation: str = 'silu',
+        qk_norm: str = 'l2',
+        norm_eps: float = 1e-5,
+        **kwargs,
+    ) -> DeltaNet:
+        super().__init__()
+
+        self.mode = mode
+        self.qk_activation = qk_activation
+        self.qk_norm = qk_norm
+
+        assert self.qk_activation in ['silu', 'relu', 'elu', 'identity']
+        assert self.qk_norm in ['l2', 'sum']
+
+        if d_model is not None:
+            hidden_size = d_model
+        self.hidden_size = hidden_size
+        self.expand_k = expand_k
+        self.expand_v = expand_v
+        self.num_heads = num_heads
+        self.use_gate = use_gate
+        self.use_short_conv = use_short_conv
+        self.conv_size = conv_size
+        self.conv_bias = conv_bias
+        self.allow_neg_eigval = allow_neg_eigval
+
+        self.key_dim = int(hidden_size * expand_k)
+        self.value_dim = int(hidden_size * expand_v)
+        self.head_k_dim = self.key_dim // num_heads
+        self.head_v_dim = self.value_dim // num_heads
+        self.layer_idx = layer_idx
+
+        if mode == 'fused_chunk':
+            raise NotImplementedError("fused_chunk_delta_rule is now deprecated. Please use `chunk_delta_rule` instead.")
+        assert mode in ['chunk', 'fused_recurrent'], f"Not supported mode `{mode}`."
+        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.q_proj = nn.Linear(hidden_size, self.key_dim, bias=False)
+        self.k_proj = nn.Linear(hidden_size, self.key_dim, bias=False)
+        self.v_proj = nn.Linear(hidden_size, self.value_dim, bias=False)
+
+        self.use_beta = use_beta
+        if self.use_beta:
+            self.b_proj = nn.Linear(hidden_size, self.num_heads, bias=False)
+        if use_short_conv:
+            self.conv_size = conv_size
+            self.q_conv1d = ShortConvolution(
+                hidden_size=self.key_dim,
+                kernel_size=conv_size,
+                bias=conv_bias,
+                activation='silu' if qk_activation == 'silu' else None,
+            )
+            self.k_conv1d = ShortConvolution(
+                hidden_size=self.key_dim,
+                kernel_size=conv_size,
+                bias=conv_bias,
+                activation='silu' if qk_activation == 'silu' else None,
+            )
+            self.v_conv1d = ShortConvolution(
+                hidden_size=self.value_dim,
+                kernel_size=conv_size,
+                bias=conv_bias,
+                activation='silu',
+            )
+        else:
+            warnings.warn(
+                "ShortConvolution is crucial to the performance. "
+                "Do not turn it off, i.e., setting `use_short_conv=False` unless you know what you are doing.",
+            )
+        if use_gate:
+            self.g_proj = nn.Linear(hidden_size, self.value_dim, bias=False)
+            self.o_norm = FusedRMSNormGated(self.head_v_dim, eps=norm_eps)
+        else:
+            self.o_norm = RMSNorm(self.head_v_dim, eps=norm_eps, dtype=torch.float32)
+
+        self.o_proj = nn.Linear(self.value_dim, hidden_size, bias=False)
+
+    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]:
+        if attention_mask is not None:
+            assert len(attention_mask.shape) == 2, (
+                "Expected attention_mask as a 0-1 matrix with shape [batch_size, seq_len] "
+                "for padding purposes (0 indicating padding). "
+                "Arbitrary attention masks of shape [batch_size, seq_len, seq_len] are not allowed."
+            )
+
+        batch_size, q_len, _ = hidden_states.shape
+        # change to inference mode.
+        mode = 'fused_recurrent' if q_len <= 64 else self.mode
+
+        last_state = get_layer_cache(self, past_key_values)
+
+        cu_seqlens = kwargs.get('cu_seqlens')
+        if attention_mask is not None:
+            indices, cu_seqlens, _ = get_unpad_data(attention_mask[:, -q_len:])
+            hidden_states = index_first_axis(rearrange(hidden_states, "b s ... -> (b s) ..."), indices).unsqueeze(0)
+
+        if self.use_short_conv:
+            conv_state_q, conv_state_k, conv_state_v = None, None, None
+            if last_state is not None:
+                conv_state_q, conv_state_k, conv_state_v = last_state['conv_state']
+            q, conv_state_q = self.q_conv1d(
+                x=self.q_proj(hidden_states),
+                cache=conv_state_q,
+                output_final_state=use_cache,
+                cu_seqlens=cu_seqlens,
+            )
+            k, conv_state_k = self.k_conv1d(
+                x=self.k_proj(hidden_states),
+                cache=conv_state_k,
+                output_final_state=use_cache,
+                cu_seqlens=cu_seqlens,
+            )
+            v, conv_state_v = self.v_conv1d(
+                x=self.v_proj(hidden_states),
+                cache=conv_state_v,
+                output_final_state=use_cache,
+                cu_seqlens=cu_seqlens,
+            )
+        else:
+            q = self.q_proj(hidden_states)
+            k = self.k_proj(hidden_states)
+            if self.qk_activation == 'silu':
+                q, k = F.silu(q), F.silu(k)
+            v = F.silu(self.v_proj(hidden_states))
+
+        q, k = map(lambda x: rearrange(x, '... (h d) -> ... h d', d=self.head_k_dim), (q, k))
+        v = rearrange(v, '... (h d) -> ... h d', d=self.head_v_dim)
+        if self.qk_activation != 'silu':
+            if self.qk_activation == 'relu':
+                q, k = q.relu(), k.relu()
+            elif self.qk_activation == 'elu':
+                q, k = elu_p1(q), elu_p1(k)
+            elif self.qk_activation != 'identity':
+                raise NotImplementedError
+
+        if self.qk_norm == 'sum':
+            q = sum_norm(q).to(q)
+            k = sum_norm(k).to(k)
+
+        if self.use_beta:
+            beta = self.b_proj(hidden_states).sigmoid()
+        else:
+            beta = torch.ones_like(q[..., 0])
+
+        if self.allow_neg_eigval:
+            beta = beta * 2.
+
+        recurrent_state = last_state['recurrent_state'] if last_state is not None else None
+        if mode == 'fused_recurrent':
+            o, recurrent_state = fused_recurrent_delta_rule(
+                q=q,
+                k=k,
+                v=v,
+                beta=beta,
+                initial_state=recurrent_state,
+                output_final_state=use_cache,
+                cu_seqlens=cu_seqlens,
+                use_qk_l2norm_in_kernel=(self.qk_norm == 'l2'),
+            )
+        elif mode == 'chunk':
+            o, recurrent_state = chunk_delta_rule(
+                q=q,
+                k=k,
+                v=v,
+                beta=beta,
+                initial_state=recurrent_state,
+                output_final_state=use_cache,
+                cu_seqlens=cu_seqlens,
+                use_qk_l2norm_in_kernel=(self.qk_norm == 'l2'),
+            )
+        else:
+            raise NotImplementedError(f"Not supported mode `{mode}`.")
+
+        update_layer_cache(
+            self,
+            past_key_values,
+            recurrent_state=recurrent_state,
+            conv_state=(conv_state_q, conv_state_k, conv_state_v) if self.use_short_conv else None,
+            offset=q_len,
+        )
+
+        if self.use_gate:
+            g = rearrange(self.g_proj(hidden_states), '... (h d) -> ... h d', d=self.head_v_dim)
+            o = self.o_norm(o, g)
+        else:
+            o = self.o_norm(o)
+        o = rearrange(o, 'b t h d -> b t (h d)')
+        o = self.o_proj(o)
+        if attention_mask is not None:
+            o = pad_input(o.squeeze(0), indices, batch_size, q_len)
+
+        return o, None, past_key_values

fla/layers/deltaformer.py

@@ -0,0 +1,152 @@
+# 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 transformers.utils import logging
+
+from fla.modules import RMSNorm, RotaryEmbedding
+from fla.ops.deltaformer import deltaformer_attn
+from fla.ops.utils.index import prepare_lens_from_mask
+
+if TYPE_CHECKING:
+    from fla.models.utils import Cache
+
+logger = logging.get_logger(__name__)
+
+
+class DeltaFormerAttention(nn.Module):
+
+    r"""
+    The layer implementation for DeltaFormer,
+    [Understanding Transformer from the Perspective of Associative Memory]
+    (https://arxiv.org/pdf/2505.19488).
+
+    Notes
+        - DeltaFormer attention is implemented with Triton kernels in `fla.ops.deltaformer` and is tuned
+          for typical head dimensions (e.g., 64/128). It currently supports fixed-length inputs.
+        - For variable-length inputs (padding masks), the deltaformer computation falls back to using the
+          fixed-length path, while the second stage (softmax attention over U) uses FlashAttention's
+          varlen path when an attention mask is provided.
+        - K/V grouping (GQA) is supported natively by FlashAttention via `num_kv_heads`.
+        - Uses K-K similarity in deltaformer computation instead of Q-K similarity for better performance.
+
+    Args:
+        hidden_size (int, Optional):
+            The hidden size of the input. Default: 2048.
+        num_heads (int, Optional):
+            The number of attention heads. Default: 32.
+        num_kv_heads (int, Optional):
+            The number of key/value heads for grouped-query attention. If None, equals `num_heads`.
+            Default: None.
+        qkv_bias (bool, Optional):
+            Whether to use bias for Q/K/V projections. Default: False.
+        qk_norm (bool, Optional):
+            Whether to apply per-head RMSNorm to Q and K before attention. Default: False.
+        rope_theta (float, Optional):
+            The base frequency for rotary position embedding. Default: 10000.
+        max_position_embeddings (int, Optional):
+            The maximum position embeddings. Default: None.
+        layer_idx (int, Optional):
+            The index of the layer (used for cache compatibility). Default: None.
+    """
+
+    def __init__(
+        self,
+        hidden_size: int = 2048,
+        num_heads: int = 32,
+        num_kv_heads: int | None = None,
+        qkv_bias: bool = False,
+        qk_norm: bool = False,
+        rope_theta: float = 10000.,
+        max_position_embeddings: int | None = None,
+        layer_idx: int | None = None,
+    ):
+        super().__init__()
+
+        self.hidden_size = hidden_size
+        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 = num_heads // self.num_kv_heads
+        self.head_dim = self.hidden_size // self.num_heads
+        self.kv_dim = self.num_kv_heads * self.head_dim
+        self.qkv_bias = qkv_bias
+        self.qk_norm = qk_norm
+        self.rope_theta = rope_theta
+        self.max_position_embeddings = max_position_embeddings
+        self.layer_idx = layer_idx
+
+        self.q_proj = nn.Linear(self.hidden_size, self.hidden_size, bias=self.qkv_bias)
+        self.k_proj = nn.Linear(self.hidden_size, self.kv_dim, bias=self.qkv_bias)
+        self.v_proj = nn.Linear(self.hidden_size, self.kv_dim, bias=self.qkv_bias)
+        self.b_proj = nn.Linear(self.hidden_size, self.num_heads, bias=True)
+        self.o_proj = nn.Linear(self.hidden_size, self.hidden_size, bias=False)
+
+        if qk_norm:
+            self.q_norm = RMSNorm(self.head_dim, dtype=torch.float32)
+            self.k_norm = RMSNorm(self.head_dim, dtype=torch.float32)
+
+        self.rotary = RotaryEmbedding(dim=self.head_dim, base=self.rope_theta)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: torch.LongTensor | None = None,
+        past_key_values: Cache | None = None,
+        output_attentions: bool = False,
+        use_cache: bool = False,
+        **kwargs,
+    ) -> tuple[torch.Tensor, torch.Tensor | None, tuple[torch.Tensor] | None]:
+        attentions = None
+        if attention_mask is not None:
+            assert len(attention_mask.shape) == 2, (
+                "Expected attention_mask as a 0-1 matrix with shape [batch_size, seq_len] "
+                "for padding purposes (0 indicating padding). "
+                "Arbitrary attention masks of shape [batch_size, seq_len, seq_len] are not allowed."
+            )
+
+        batch_size, q_len, _ = hidden_states.size()
+
+        q = rearrange(self.q_proj(hidden_states), '... (h d) -> ... h d', d=self.head_dim)
+        k = rearrange(self.k_proj(hidden_states), '... (h d) -> ... h d', d=self.head_dim)
+        v = rearrange(self.v_proj(hidden_states), '... (h d) -> ... h d', d=self.head_dim)
+        beta = self.b_proj(hidden_states)
+
+        if self.qk_norm:
+            q, k = self.q_norm(q), self.k_norm(k)
+
+        cu_seqlens_kw = kwargs.get('cu_seqlens')
+        seqlen_offset, max_seqlen = 0, q_len
+        if past_key_values is not None:
+            seqlen_offset = past_key_values.get_seq_length(self.layer_idx)
+            max_seqlen = q_len + seqlen_offset
+
+            if attention_mask is not None:
+                seqlen_offset = seqlen_offset + prepare_lens_from_mask(attention_mask) - attention_mask.shape[-1]
+                max_seqlen = q_len + max(seqlen_offset)
+
+        if self.max_position_embeddings is not None:
+            max_seqlen = max(max_seqlen, self.max_position_embeddings)
+
+        q, k = self.rotary(q, k, seqlen_offset=seqlen_offset, max_seqlen=max_seqlen, cu_seqlens=cu_seqlens_kw)
+
+        o = deltaformer_attn(
+            q=q,
+            k=k,
+            v=v,
+            beta=beta,
+            attention_mask=attention_mask,
+            cu_seqlens=cu_seqlens_kw,
+        )
+
+        o = o.reshape(batch_size, q_len, -1)
+        o = self.o_proj(o)
+
+        if not output_attentions:
+            attentions = None
+
+        return o, attentions, past_key_values

fla/layers/forgetting_attn.py

@@ -0,0 +1,133 @@
+# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.utils.checkpoint
+from einops import rearrange
+from transformers.utils import logging
+
+from fla.layers.utils import pad_input, unpad_input
+from fla.modules import GroupNorm
+from fla.ops.attn.decoding import attn_decoding_one_step
+from fla.ops.forgetting_attn.parallel import parallel_forgetting_attn
+
+if TYPE_CHECKING:
+    from fla.models.utils import Cache
+
+logger = logging.get_logger(__name__)
+
+
+class ForgettingAttention(nn.Module):
+
+    def __init__(
+        self,
+        hidden_size: int = 2048,
+        num_heads: int = 32,
+        num_kv_heads: int | None = None,
+        qkv_bias: bool = False,
+        qk_norm: bool = False,
+        window_size: int | None = None,
+        use_output_gate: bool = False,
+        layer_idx: int = None,
+    ):
+        super().__init__()
+
+        self.hidden_size = hidden_size
+        self.num_heads = num_heads
+        if num_kv_heads is None:
+            self.num_kv_heads = self.num_heads
+        else:
+            self.num_kv_heads = num_kv_heads
+        self.num_kv_groups = num_heads // self.num_kv_heads
+        self.head_dim = self.hidden_size // self.num_heads
+        self.kv_dim = self.num_kv_heads * self.head_dim
+        self.qkv_bias = qkv_bias
+        self.qk_norm = qk_norm
+
+        self.window_size = window_size
+        self.use_output_gate = use_output_gate
+        self.layer_idx = layer_idx
+
+        self.q_proj = nn.Linear(self.hidden_size, self.hidden_size, bias=self.qkv_bias)
+        self.k_proj = nn.Linear(self.hidden_size, self.kv_dim, bias=self.qkv_bias)
+        self.v_proj = nn.Linear(self.hidden_size, self.kv_dim, bias=self.qkv_bias)
+        self.f_proj = nn.Linear(self.hidden_size, self.num_heads, bias=True)
+
+        if use_output_gate:
+            self.g_proj = nn.Linear(self.hidden_size, self.hidden_size, bias=False)
+        self.o_proj = nn.Linear(self.hidden_size, self.hidden_size, bias=False)
+
+        if qk_norm:
+            self.q_norm = GroupNorm(
+                num_groups=self.num_heads,
+                hidden_size=self.hidden_size,
+                is_rms_norm=True,
+            )
+            self.k_norm = GroupNorm(
+                num_groups=self.num_kv_heads,
+                hidden_size=self.kv_dim,
+                is_rms_norm=True,
+            )
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: torch.LongTensor | None = None,
+        past_key_values: Cache | None = None,
+        output_attentions: bool = False,
+        use_cache: bool = False,
+        **kwargs,
+    ) -> tuple[torch.Tensor, torch.Tensor | None, tuple[torch.Tensor] | None]:
+        if attention_mask is not None:
+            assert len(attention_mask.shape) == 2, (
+                "Expected attention_mask as a 0-1 matrix with shape [batch_size, seq_len] "
+                "for padding purposes (0 indicating padding). "
+                "Arbitrary attention masks of shape [batch_size, seq_len, seq_len] are not allowed."
+            )
+
+        batch_size, q_len, _ = hidden_states.size()
+
+        q, k, v = self.q_proj(hidden_states), self.k_proj(hidden_states), self.v_proj(hidden_states)
+        f = F.logsigmoid(self.f_proj(hidden_states).float())
+        if self.qk_norm:
+            q, k = self.q_norm(q), self.k_norm(k)
+
+        cu_seqlens = kwargs.get('cu_seqlens')
+        if past_key_values is not None:
+            assert cu_seqlens is None, "cu_seqlens should not be provided when past_key_values is not None"
+            state = past_key_values.update(
+                attn_state=(k, v, f),
+                layer_idx=self.layer_idx,
+                offset=q_len,
+                cache_kwargs=dict(window_size=self.window_size),
+            )
+            k, v, f = state['attn_state']
+
+        q = rearrange(q, '... (h d) -> ... h d', d=self.head_dim)
+        k = rearrange(k, '... (h d) -> ... h d', d=self.head_dim)
+        v = rearrange(v, '... (h d) -> ... h d', d=self.head_dim)
+
+        if attention_mask is not None:
+            q, (k, v, f), indices_q, cu_seqlens, max_seq_lens = unpad_input(q, (k, v, f), attention_mask, q_len, keepdim=True)
+            _, cu_seqlens_k = cu_seqlens
+            cu_seqlens = cu_seqlens_k
+            max_seqlen_q, max_seqlen_k = max_seq_lens
+            if max_seqlen_q != max_seqlen_k:
+                assert max_seqlen_q == 1, "only support q_len == 1 for decoding"
+                o = attn_decoding_one_step(q, k, v, f, cu_seqlens=cu_seqlens)
+            else:
+                o = parallel_forgetting_attn(q, k, v, f, cu_seqlens=cu_seqlens)
+        else:
+            o = parallel_forgetting_attn(q, k, v, f, cu_seqlens=cu_seqlens)
+        if attention_mask is not None:
+            o = pad_input(o.squeeze(0), indices_q, batch_size, q_len)
+        o = rearrange(o, '... h d -> ... (h d)')
+        if self.use_output_gate:
+            o = self.g_proj(hidden_states).sigmoid() * o
+        o = self.o_proj(o)
+        return o, None, past_key_values

fla/layers/gated_deltanet.py

@@ -0,0 +1,316 @@
+# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang
+
+from __future__ import annotations
+
+import math
+import warnings
+from typing import TYPE_CHECKING
+
+import torch
+import torch.nn as nn
+from einops import rearrange, repeat
+from torch.nn import functional as F
+
+from fla.layers.utils import get_layer_cache, get_unpad_data, index_first_axis, pad_input, update_layer_cache
+from fla.modules import FusedRMSNormGated, RMSNorm, ShortConvolution
+from fla.ops.gated_delta_rule import chunk_gated_delta_rule, fused_recurrent_gated_delta_rule
+
+if TYPE_CHECKING:
+    from transformers.processing_utils import Unpack
+
+    from fla.models.utils import Cache
+
+
+@torch.compile
+def elu_p1(x):
+    return (F.elu(x, 1., False) + 1.).to(x)
+
+
+@torch.compile
+def sum_norm(x):
+    return (x / x.sum(-1, keepdim=True)).to(x)
+
+
+class GatedDeltaNet(nn.Module):
+    """
+    The layer implementaion for [Gated Delta Networks: Improving Mamba2 with Delta Rule](https://arxiv.org/abs/2412.06464).  # noqa
+
+    Similar to Mamba2, each layer contains around 6*hidden_size*hidden_size parameters.
+
+    Parameter alloation when use_gate=True:
+        - 0.75 * hidden_size * hidden_size for the q_proj and k_proj each
+        - 1.5 * hidden_size * hidden_size for the v_proj, g_proj and o_proj each
+        - Others are ignorably small.
+        - In total = 0.75 * 2 + 1.5 * 3 = 6 * hidden_size * hidden_size
+    NOTE: num_heads * head_dim = 0.75 * hidden_size, please make sure to set the correct num_heads and head_dim.
+
+    Parameter allocation when use_gate=False:
+        - 1 * hidden_size * hidden_size for the q_proj and k_proj each
+        - 2 * hidden_size * hidden_size for the v_proj and o_proj each
+        - Others are ignorably small.
+        - In total = 1 * 2 + 2 * 2 = 6 * hidden_size * hidden_size
+
+    Args:
+        hidden_size (int, Optional):
+            The hidden size of the input. Default: 2048.
+        expand_v (float, Optional):
+            The expansion ratio for the value dim. Default: 2.0.
+        head_dim (int, Optional):
+            The dimension of each head. Default: 256.
+        num_heads (int, Optional):
+            The number of heads. Default: 4.
+        num_v_heads (int, Optional):
+            The number of heads for the value projection, equal to `num_heads` if `None`.
+            GVA is applied if `num_v_heads` > `num_heads`. Default: `None`.
+        mode (str, Optional):
+            Which Gated DeltaNet kernel to use.
+            Currently available: `chunk` and `fused_recurrent`.
+            Default: `chunk`.
+        use_beta (bool, Optional):
+            Whether to use beta. Default: `True`.
+        use_gate (bool, Optional):
+            Whether to use output gate. Default: `True`.
+        use_short_conv (bool, Optional):
+            Whether to use short convolutions. Default: `True`.
+        allow_neg_eigval (bool, Optional):
+            Allow negative eigenvalues. Default: `False`. If set to `True`, the beta will be multiplied by 2.
+            See reference: [Unlocking State-Tracking in Linear RNNs Through Negative Eigenvalues](https://arxiv.org/abs/2411.12537)
+        conv_size (int, Optional):
+            The kernel size of the short convolution, only used when `use_short_conv` is `True`. Default: 4.
+        conv_bias (bool, Optional):
+            Whether to use bias in the short convolution, only used when `use_short_conv` is `True`. Default: `False`.
+        layer_idx (int, Optional):
+            The index of the layer. Default: None.
+        norm_eps (float, Optional):
+            The epsilon value for the normalization layer. Default: 1e-5.
+    """
+
+    def __init__(
+        self,
+        hidden_size: int = 2048,
+        expand_v: float = 2,
+        head_dim: int = 256,
+        num_heads: int = 6,
+        num_v_heads: int = None,
+        mode: str = 'chunk',
+        use_gate: bool = True,
+        use_short_conv: bool = True,
+        allow_neg_eigval: bool = False,
+        conv_size: int = 4,
+        conv_bias: bool = False,
+        layer_idx: int = None,
+        norm_eps: float = 1e-5,
+        **kwargs,
+    ) -> GatedDeltaNet:
+        super().__init__()
+
+        self.mode = mode
+        self.allow_neg_eigval = allow_neg_eigval
+        self.hidden_size = hidden_size
+        self.expand_v = expand_v
+
+        self.use_gate = use_gate
+        self.use_short_conv = use_short_conv
+        self.conv_size = conv_size
+        self.conv_bias = conv_bias
+
+        self.head_dim = head_dim
+        self.num_heads = num_heads
+        self.num_v_heads = num_v_heads if num_v_heads is not None else num_heads
+
+        self.head_k_dim = head_dim
+        self.head_v_dim = int(self.head_dim * self.expand_v)
+        self.key_dim = int(self.num_heads * self.head_k_dim)
+        self.value_dim = int(self.num_v_heads * self.head_v_dim)
+        self.layer_idx = layer_idx
+
+        # Consistency check: Ensure expand_v produces integer values
+        if not math.isclose(self.num_v_heads * self.head_dim * expand_v, self.value_dim, rel_tol=1e-5):
+            raise ValueError(
+                f"expand_v={expand_v} does not produce an integer value when multiplied by key_dim={self.key_dim}. "
+                f"Resulting value_dim would be {self.num_v_heads * self.head_dim * expand_v}, which is invalid for nn.Linear.",
+            )
+        if self.num_v_heads > self.num_heads and self.num_v_heads % self.num_heads != 0:
+            raise ValueError(
+                f"num_v_heads={self.num_v_heads} must be divisible by num_heads={self.num_heads}.",
+            )
+
+        if not math.isclose(head_dim * expand_v, self.head_v_dim, rel_tol=1e-5):
+            raise ValueError(
+                f"expand_v={expand_v} does not produce an integer value when multiplied by head_dim={head_dim}. "
+                f"Resulting head_v_dim would be {head_dim * expand_v}, which is invalid for FusedRMSNormGated.",
+            )
+        assert mode in ['chunk', 'fused_recurrent'], f"Not supported mode `{mode}`."
+
+        self.q_proj = nn.Linear(hidden_size, self.key_dim, bias=False)
+        self.k_proj = nn.Linear(hidden_size, self.key_dim, bias=False)
+        self.v_proj = nn.Linear(hidden_size, self.value_dim, bias=False)
+        self.a_proj = nn.Linear(hidden_size, self.num_v_heads, bias=False)
+        self.b_proj = nn.Linear(hidden_size, self.num_v_heads, bias=False)
+
+        A = torch.empty(self.num_v_heads, dtype=torch.float32).uniform_(0, 16)
+        self.A_log = nn.Parameter(torch.log(A))
+        self.A_log._no_weight_decay = True
+        # hard coded for now
+        dt_min = 0.001
+        dt_max = 0.1
+        dt_init_floor = 1e-4
+        dt = torch.exp(
+            torch.rand(self.num_v_heads) * (math.log(dt_max) - math.log(dt_min))
+            + math.log(dt_min),
+        )
+        dt = torch.clamp(dt, min=dt_init_floor)
+        # Inverse of softplus: https://github.com/pytorch/pytorch/issues/72759
+        inv_dt = dt + torch.log(-torch.expm1(-dt))
+        self.dt_bias = nn.Parameter(inv_dt)
+        # Just to be explicit. Without this we already don't put wd on dt_bias because of the check
+        # name.endswith("bias") in param_grouping.py
+        self.dt_bias._no_weight_decay = True
+
+        if use_short_conv:
+            self.conv_size = conv_size
+            self.q_conv1d = ShortConvolution(
+                hidden_size=self.key_dim,
+                kernel_size=conv_size,
+                bias=conv_bias,
+                activation='silu',
+            )
+            self.k_conv1d = ShortConvolution(
+                hidden_size=self.key_dim,
+                kernel_size=conv_size,
+                bias=conv_bias,
+                activation='silu',
+            )
+            self.v_conv1d = ShortConvolution(
+                hidden_size=self.value_dim,
+                kernel_size=conv_size,
+                bias=conv_bias,
+                activation='silu',
+            )
+        else:
+            warnings.warn(
+                "ShortConvolution is crucial to the performance. "
+                "Do not turn it off, i.e., setting `use_short_conv=False` unless you know what you are doing.",
+            )
+        if use_gate:
+            self.g_proj = nn.Linear(hidden_size, self.value_dim, bias=False)
+            self.o_norm = FusedRMSNormGated(self.head_v_dim, eps=norm_eps)
+        else:
+            self.o_norm = RMSNorm(self.head_v_dim, eps=norm_eps, dtype=torch.float32)
+        self.o_proj = nn.Linear(self.value_dim, hidden_size, bias=False)
+
+    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]:
+        if attention_mask is not None:
+            assert len(attention_mask.shape) == 2, (
+                "Expected attention_mask as a 0-1 matrix with shape [batch_size, seq_len] "
+                "for padding purposes (0 indicating padding). "
+                "Arbitrary attention masks of shape [batch_size, seq_len, seq_len] are not allowed."
+            )
+
+        batch_size, q_len, _ = hidden_states.shape
+        # change to inference mode.
+        mode = 'fused_recurrent' if (q_len <= 64 and not self.training) else self.mode
+        if self.training:
+            assert mode == 'chunk', "Only chunk mode is supported in training."
+
+        last_state = get_layer_cache(self, past_key_values)
+
+        cu_seqlens = kwargs.get('cu_seqlens')
+        if attention_mask is not None:
+            indices, cu_seqlens, _ = get_unpad_data(attention_mask[:, -q_len:])
+            hidden_states = index_first_axis(rearrange(hidden_states, "b s ... -> (b s) ..."), indices).unsqueeze(0)
+
+        if self.use_short_conv:
+            conv_state_q, conv_state_k, conv_state_v = None, None, None
+            if last_state is not None:
+                conv_state_q, conv_state_k, conv_state_v = last_state['conv_state']
+            q, conv_state_q = self.q_conv1d(
+                x=self.q_proj(hidden_states),
+                cache=conv_state_q,
+                output_final_state=use_cache,
+                cu_seqlens=cu_seqlens,
+            )
+            k, conv_state_k = self.k_conv1d(
+                x=self.k_proj(hidden_states),
+                cache=conv_state_k,
+                output_final_state=use_cache,
+                cu_seqlens=cu_seqlens,
+            )
+            v, conv_state_v = self.v_conv1d(
+                x=self.v_proj(hidden_states),
+                cache=conv_state_v,
+                output_final_state=use_cache,
+                cu_seqlens=cu_seqlens,
+            )
+        else:
+            q = F.silu(self.q_proj(hidden_states))
+            k = F.silu(self.k_proj(hidden_states))
+            v = F.silu(self.v_proj(hidden_states))
+
+        q, k = map(lambda x: rearrange(x, '... (h d) -> ... h d', d=self.head_k_dim), (q, k))
+        v = rearrange(v, '... (h d) -> ... h d', d=self.head_v_dim)
+
+        if self.num_v_heads > self.num_heads:
+            q, k = map(lambda x: repeat(x, '... h d -> ... (h g) d', g=self.num_v_heads // self.num_heads), (q, k))
+
+        beta = self.b_proj(hidden_states).sigmoid()
+        if self.allow_neg_eigval:
+            beta = beta * 2.
+
+        g = -self.A_log.float().exp() * F.softplus(self.a_proj(hidden_states).float() + self.dt_bias)
+
+        recurrent_state = last_state['recurrent_state'] if last_state is not None else None
+        if mode == 'chunk':
+            o, recurrent_state = chunk_gated_delta_rule(
+                q=q,
+                k=k,
+                v=v,
+                g=g,
+                beta=beta,
+                initial_state=recurrent_state,
+                output_final_state=use_cache,
+                cu_seqlens=cu_seqlens,
+                use_qk_l2norm_in_kernel=True,
+            )
+        elif mode == 'fused_recurrent':
+            o, recurrent_state = fused_recurrent_gated_delta_rule(
+                q=q,
+                k=k,
+                v=v,
+                g=g,
+                beta=beta,
+                initial_state=recurrent_state,
+                output_final_state=use_cache,
+                cu_seqlens=cu_seqlens,
+                use_qk_l2norm_in_kernel=True,
+            )
+        else:
+            raise NotImplementedError(f"Not supported mode `{mode}`.")
+
+        update_layer_cache(
+            self,
+            past_key_values,
+            recurrent_state=recurrent_state,
+            conv_state=(conv_state_q, conv_state_k, conv_state_v) if self.use_short_conv else None,
+            offset=q_len,
+        )
+
+        if self.use_gate:
+            g = rearrange(self.g_proj(hidden_states), '... (h d) -> ... h d', d=self.head_v_dim)
+            o = self.o_norm(o, g)
+        else:
+            o = self.o_norm(o)
+        o = rearrange(o, 'b t h d -> b t (h d)')
+        o = self.o_proj(o)
+        if attention_mask is not None:
+            o = pad_input(o.squeeze(0), indices, batch_size, q_len)
+
+        return o, None, past_key_values

fla/layers/gated_deltaproduct.py

@@ -0,0 +1,287 @@
+# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang
+
+from __future__ import annotations
+
+import math
+import warnings
+from typing import TYPE_CHECKING
+
+import torch
+import torch.nn as nn
+from einops import rearrange, repeat
+from torch.nn import functional as F
+
+from fla.layers.utils import get_layer_cache, get_unpad_data, index_first_axis, pad_input, update_layer_cache
+from fla.modules import FusedRMSNormGated, RMSNorm, ShortConvolution
+from fla.ops.gated_delta_product import chunk_gated_delta_product
+from fla.ops.gated_delta_rule import fused_recurrent_gated_delta_rule
+
+if TYPE_CHECKING:
+    from transformers.processing_utils import Unpack
+
+    from fla.models.utils import Cache
+
+
+class GatedDeltaProduct(nn.Module):
+    """
+    Generalized version of GatedDoubleDeltaNet that supports arbitrary number of householder transformations.
+    """
+
+    def __init__(
+        self,
+        hidden_size: int = 2048,
+        expand_v: float = 2,
+        head_dim: int = 256,
+        num_heads: int = 6,
+        num_v_heads: int = None,
+        mode: str = 'chunk',
+        use_output_gate: bool = True,
+        use_short_conv: bool = True,
+        conv_size: int = 4,
+        conv_bias: bool = False,
+        layer_idx: int = None,
+        norm_eps: float = 1e-5,
+        use_forget_gate: bool = True,
+        allow_neg_eigval: bool = True,
+        num_householder: int = 2,
+        **kwargs,
+    ) -> GatedDeltaProduct:
+        super().__init__()
+
+        self.mode = mode
+
+        self.hidden_size = hidden_size
+        self.expand_v = expand_v
+
+        self.use_forget_gate = use_forget_gate
+        self.allow_neg_eigval = allow_neg_eigval
+        self.num_householder = num_householder
+        self.use_output_gate = use_output_gate
+        self.use_short_conv = use_short_conv
+        self.conv_size = conv_size
+        self.conv_bias = conv_bias
+
+        self.head_dim = head_dim
+        self.num_heads = num_heads
+        self.num_v_heads = num_v_heads if num_v_heads is not None else num_heads
+
+        self.head_k_dim = head_dim
+        self.head_v_dim = int(self.head_dim * self.expand_v)
+        self.key_dim = int(self.num_heads * self.head_k_dim)
+        self.value_dim = int(self.num_v_heads * self.head_v_dim)
+        self.layer_idx = layer_idx
+
+        # Consistency check: Ensure expand_v produces integer values
+        if not math.isclose(self.num_v_heads * self.head_dim * expand_v, self.value_dim, rel_tol=1e-5):
+            raise ValueError(
+                f"expand_v={expand_v} does not produce an integer value when multiplied by key_dim={self.key_dim}. "
+                f"Resulting value_dim would be {self.num_v_heads * self.head_dim * expand_v}, which is invalid for nn.Linear.",
+            )
+        if self.num_v_heads > self.num_heads and self.num_v_heads % self.num_heads != 0:
+            raise ValueError(
+                f"num_v_heads={self.num_v_heads} must be divisible by num_heads={self.num_heads}.",
+            )
+
+        if not math.isclose(head_dim * expand_v, self.head_v_dim, rel_tol=1e-5):
+            raise ValueError(
+                f"expand_v={expand_v} does not produce an integer value when multiplied by head_dim={head_dim}. "
+                f"Resulting head_v_dim would be {head_dim * expand_v}, which is invalid for FusedRMSNormGated.",
+            )
+        assert mode in ['chunk', 'fused_recurrent'], f"Not supported mode `{mode}`."
+
+        self.q_proj = nn.Linear(hidden_size, self.key_dim, bias=False)
+        self.k_proj = nn.Linear(hidden_size, self.key_dim * num_householder, bias=False)
+        self.v_proj = nn.Linear(hidden_size, self.value_dim * num_householder, bias=False)
+        self.b_proj = nn.Linear(hidden_size, self.num_v_heads * num_householder, bias=False)
+
+        if self.use_forget_gate:
+            self.a_proj = nn.Linear(hidden_size, self.num_v_heads, bias=False)
+            A = torch.empty(self.num_v_heads, dtype=torch.float32).uniform_(0, 16)
+            self.A_log = nn.Parameter(torch.log(A))
+            self.A_log._no_weight_decay = True
+            # hard coded for now
+            dt_min = 0.001
+            dt_max = 0.1
+            dt_init_floor = 1e-4
+            dt = torch.exp(
+                torch.rand(self.num_v_heads) * (math.log(dt_max) - math.log(dt_min))
+                + math.log(dt_min),
+            )
+            dt = torch.clamp(dt, min=dt_init_floor)
+            # Inverse of softplus: https://github.com/pytorch/pytorch/issues/72759
+            inv_dt = dt + torch.log(-torch.expm1(-dt))
+            self.dt_bias = nn.Parameter(inv_dt)
+            # Just to be explicit. Without this we already don't put wd on dt_bias because of the check
+            # name.endswith("bias") in param_grouping.py
+            self.dt_bias._no_weight_decay = True
+
+        if use_short_conv:
+            self.conv_size = conv_size
+            self.q_conv1d = ShortConvolution(
+                hidden_size=self.key_dim,
+                kernel_size=conv_size,
+                bias=conv_bias,
+                activation='silu',
+            )
+            self.k_conv1d = ShortConvolution(
+                hidden_size=self.key_dim * num_householder,
+                kernel_size=conv_size,
+                bias=conv_bias,
+                activation='silu',
+            )
+            self.v_conv1d = ShortConvolution(
+                hidden_size=self.value_dim * num_householder,
+                kernel_size=conv_size,
+                bias=conv_bias,
+                activation='silu',
+            )
+        else:
+            warnings.warn(
+                "ShortConvolution is crucial to the performance. "
+                "Do not turn it off, i.e., setting `use_short_conv=False` unless you know what you are doing.",
+            )
+        if use_output_gate:
+            self.g_proj = nn.Linear(hidden_size, self.value_dim, bias=False)
+            self.o_norm = FusedRMSNormGated(self.head_v_dim, eps=norm_eps)
+        else:
+            self.o_norm = RMSNorm(self.head_v_dim, eps=norm_eps, dtype=torch.float32)
+        self.o_proj = nn.Linear(self.value_dim, hidden_size, bias=False)
+
+    def _initialize_weights(self, module: nn.Module):
+        if getattr(module, "_is_hf_initialized", False):
+            return
+        if isinstance(module, nn.Linear):
+            nn.init.xavier_uniform_(module.weight, gain=2 ** -2.5)
+            if module.bias is not None:
+                nn.init.zeros_(module.bias)
+        module._is_hf_initialized = True
+
+    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]:
+        if attention_mask is not None:
+            assert len(attention_mask.shape) == 2, (
+                "Expected attention_mask as a 0-1 matrix with shape [batch_size, seq_len] "
+                "for padding purposes (0 indicating padding). "
+                "Arbitrary attention masks of shape [batch_size, seq_len, seq_len] are not allowed."
+            )
+
+        batch_size, q_len, _ = hidden_states.shape
+        # change to inference mode.
+        mode = 'fused_recurrent' if (q_len <= 64 and not self.training) else self.mode
+        if self.training:
+            assert mode == 'chunk', "Only chunk mode is supported in training."
+
+        last_state = get_layer_cache(self, past_key_values)
+
+        cu_seqlens = kwargs.get('cu_seqlens')
+        if attention_mask is not None:
+            indices, cu_seqlens, _ = get_unpad_data(attention_mask[:, -q_len:])
+            hidden_states = index_first_axis(rearrange(hidden_states, "b s ... -> (b s) ..."), indices).unsqueeze(0)
+
+        if self.use_short_conv:
+            conv_state_q, conv_state_k, conv_state_v = None, None, None
+            if last_state is not None:
+                conv_state_q, conv_state_k, conv_state_v = last_state['conv_state']
+            q, conv_state_q = self.q_conv1d(
+                x=self.q_proj(hidden_states),
+                cache=conv_state_q,
+                output_final_state=use_cache,
+                cu_seqlens=cu_seqlens,
+            )
+            k, conv_state_k = self.k_conv1d(
+                x=self.k_proj(hidden_states),
+                cache=conv_state_k,
+                output_final_state=use_cache,
+                cu_seqlens=cu_seqlens,
+            )
+            v, conv_state_v = self.v_conv1d(
+                x=self.v_proj(hidden_states),
+                cache=conv_state_v,
+                output_final_state=use_cache,
+                cu_seqlens=cu_seqlens,
+            )
+        else:
+            q = F.silu(self.q_proj(hidden_states))
+            k = F.silu(self.k_proj(hidden_states))
+            v = F.silu(self.v_proj(hidden_states))
+
+        q = rearrange(q, '... (h d) -> ... h d', d=self.head_k_dim)
+        k = rearrange(k, '... t (n h d) -> ... (t n) h d', n=self.num_householder, d=self.head_k_dim)
+        v = rearrange(v, '... t (n h d) -> ... (t n) h d', n=self.num_householder, d=self.head_v_dim)
+
+        if self.num_v_heads > self.num_heads:
+            q, k = map(lambda x: repeat(x, '... h d -> ... (h g) d', g=self.num_v_heads // self.num_heads), (q, k))
+
+        beta = self.b_proj(hidden_states).sigmoid()
+        if self.allow_neg_eigval:
+            beta = beta * 2.
+
+        beta = rearrange(beta, '... t (n h) -> ... (t n) h', n=self.num_householder)
+        if self.use_forget_gate:
+            g = -self.A_log.float().exp() * F.softplus(self.a_proj(hidden_states).float() + self.dt_bias)
+        else:
+            g = None
+
+        recurrent_state = last_state['recurrent_state'] if last_state is not None else None
+        if mode == 'chunk':
+            o, recurrent_state = chunk_gated_delta_product(
+                q=q,
+                k=k,
+                v=v,
+                g=g,
+                beta=beta,
+                initial_state=recurrent_state,
+                output_final_state=use_cache,
+                cu_seqlens=cu_seqlens,
+                num_householder=self.num_householder,
+                use_qk_l2norm_in_kernel=True,
+            )
+
+        elif mode == 'fused_recurrent':
+            if self.use_forget_gate:
+                g_new = g.new_zeros(g.shape[0], g.shape[1], self.num_householder, g.shape[2])
+                g_new[:, :, 0] = g
+                g = rearrange(g_new, '... t n h -> ... (t n) h')
+
+            q_new = q.new_zeros(q.shape[0], q.shape[1], self.num_householder, q.shape[2], q.shape[3])
+            q_new[:, :, -1] = q
+            q = rearrange(q_new, '... t n h d-> ... (t n) h d')
+            o, recurrent_state = fused_recurrent_gated_delta_rule(
+                q=q,
+                k=k,
+                v=v,
+                g=g,
+                beta=beta,
+                initial_state=recurrent_state,
+                output_final_state=use_cache,
+                cu_seqlens=cu_seqlens * self.num_householder if cu_seqlens is not None else None,
+                use_qk_l2norm_in_kernel=True,
+            )
+            o = rearrange(o, '... (t n) h d -> ... t n h d', n=self.num_householder)[..., -1, :, :].contiguous()
+
+        update_layer_cache(
+            self,
+            past_key_values,
+            recurrent_state=recurrent_state,
+            conv_state=(conv_state_q, conv_state_k, conv_state_v) if self.use_short_conv else None,
+            offset=q_len,
+        )
+
+        if self.use_output_gate:
+            g = rearrange(self.g_proj(hidden_states), '... (h d) -> ... h d', d=self.head_v_dim)
+            o = self.o_norm(o, g)
+        else:
+            o = self.o_norm(o)
+        o = rearrange(o, 'b t h d -> b t (h d)')
+        o = self.o_proj(o)
+        if attention_mask is not None:
+            o = pad_input(o.squeeze(0), indices, batch_size, q_len)
+
+        return o, None, past_key_values

fla/layers/gla.py

@@ -0,0 +1,304 @@
+# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang
+
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange, repeat
+
+from fla.layers.utils import get_layer_cache, get_unpad_data, index_first_axis, pad_input, update_layer_cache
+from fla.modules import FusedRMSNormGated, RMSNorm, ShortConvolution
+from fla.modules.activations import ACT2FN
+from fla.ops.gla import chunk_gla, fused_chunk_gla, fused_recurrent_gla
+
+if TYPE_CHECKING:
+    from transformers.processing_utils import Unpack
+
+    from fla.models.utils import Cache
+
+
+class GatedLinearAttention(nn.Module):
+    r"""
+    The layer implementaion for [Gated Linear Attention Transformers with Hardware-Efficient Training](https://arxiv.org/abs/2312.06635).  # noqa
+
+    Args:
+        mode (str, Optional):
+            Which GLA kernel to use.
+            Currently available: `chunk`, `fused_recurrent`, and `fused_chunk`.
+            Default: `chunk`.
+        hidden_size (int, Optional):
+            The hidden size of the input. Default: 1024.
+        expand_k (float, Optional):
+            The expansion ratio for the key dim. Default: 0.5.
+        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.
+        feature_map (str, Optional):
+            Feature map function applied to queries/keys. Default: None.
+        use_short_conv (bool, Optional):
+            Whether to use short convolutions. Default: `False`.
+        conv_size (int, Optional):
+            The kernel size of the short convolution, only used when `use_short_conv` is `True`. Default: 4.
+        conv_bias (bool, Optional):
+            Whether to use bias in the short convolution, only used when `use_short_conv` is `True`. Default: `False`.
+        use_output_gate (bool, Optional):
+            Whether to use output gate. Default: `True`.
+        gate_fn (str, Optional):
+            The activation function for the output gate. Default: `swish`.
+        elementwise_affine (bool, Optional):
+            If `True`, applies elementwise affine to LayerNorm with learnable parameters. Default: `True`.
+        norm_eps (float, Optional):
+            The epsilon value for the layernorm/rmsnorm layer. Default: 1e-5.
+        gate_logit_normalizer (int, Optional):
+            The normalizer for the gate logits, appied after `logsigmoid`. Default: 16.
+        gate_low_rank_dim (int, Optional):
+            The low rank dim for the gate projection. Default: 16.
+        clamp_min (float, Optional):
+            The minimum value for the gate logits. Default: None.
+        fuse_norm (bool, Optional):
+            Whether to fuse the norm and the output gate for better memory footprint. Default: `True`.
+        layer_idx (int, Optional):
+            The index of the layer. Default: None.
+    """
+
+    def __init__(
+        self,
+        mode: str = 'chunk',
+        hidden_size: int = 1024,
+        expand_k: float = 0.5,
+        expand_v: float = 1.0,
+        num_heads: int = 4,
+        num_kv_heads: int | None = None,
+        feature_map: str | None = None,
+        use_short_conv: bool = False,
+        conv_size: int = 4,
+        conv_bias: bool = False,
+        use_output_gate: bool = True,
+        gate_fn: str = 'swish',
+        elementwise_affine: bool | None = True,
+        norm_eps: float = 1e-5,
+        gate_logit_normalizer: int = 16,
+        gate_low_rank_dim: int = 16,
+        clamp_min: float | None = None,
+        fuse_norm: bool = True,
+        layer_idx: int = None,
+    ) -> GatedLinearAttention:
+        super().__init__()
+
+        self.mode = mode
+        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.feature_map_fn = ACT2FN[feature_map] if feature_map is not None else None
+
+        self.use_short_conv = use_short_conv
+        self.conv_size = conv_size
+        self.conv_bias = conv_bias
+        self.use_output_gate = use_output_gate
+
+        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
+        self.clamp_min = clamp_min
+        self.layer_idx = layer_idx
+
+        assert mode in ['chunk', 'fused_recurrent', 'fused_chunk'], f"Not supported mode `{mode}`."
+        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
+
+        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)
+        if self.use_output_gate:
+            self.g_proj = nn.Linear(hidden_size, self.value_dim, bias=False)
+
+        if use_short_conv:
+            self.conv_size = conv_size
+            self.q_conv1d = ShortConvolution(
+                hidden_size=self.key_dim,
+                kernel_size=conv_size,
+                bias=conv_bias,
+                activation='silu',
+            )
+            self.k_conv1d = ShortConvolution(
+                hidden_size=self.key_dim_per_group,
+                kernel_size=conv_size,
+                bias=conv_bias,
+                activation='silu',
+            )
+            self.v_conv1d = ShortConvolution(
+                hidden_size=self.value_dim_per_group,
+                kernel_size=conv_size,
+                bias=conv_bias,
+                activation='silu',
+            )
+
+        self.gk_proj = nn.Sequential(nn.Linear(hidden_size, gate_low_rank_dim, bias=False),
+                                     nn.Linear(gate_low_rank_dim, self.key_dim_per_group, bias=True))
+        self.o_proj = nn.Linear(self.value_dim, hidden_size, bias=False)
+
+        if gate_fn == 'swish' and fuse_norm and use_output_gate:
+            self.g_norm_swish_gate = FusedRMSNormGated(
+                hidden_size=self.head_v_dim,
+                elementwise_affine=elementwise_affine,
+                eps=norm_eps,
+            )
+            self.fuse_norm_and_gate = True
+        else:
+            self.fuse_norm_and_gate = False
+            self.g_norm = RMSNorm(
+                hidden_size=self.head_v_dim,
+                elementwise_affine=elementwise_affine,
+                eps=norm_eps,
+                dtype=torch.float32
+            )
+            self.gate_fn = ACT2FN[gate_fn]
+
+        self.gate_logit_normalizer = gate_logit_normalizer
+
+    def reset_parameters(self) -> None:
+        for module in self.children():
+            reset = getattr(module, "reset_parameters", None)
+            if callable(reset):
+                reset()
+
+    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]:
+        if attention_mask is not None:
+            assert len(attention_mask.shape) == 2, (
+                "Expected attention_mask as a 0-1 matrix with shape [batch_size, seq_len] "
+                "for padding purposes (0 indicating padding). "
+                "Arbitrary attention masks of shape [batch_size, seq_len, seq_len] are not allowed."
+            )
+
+        batch_size, q_len, _ = hidden_states.shape
+        mode = 'fused_recurrent' if hidden_states.shape[1] <= 64 else self.mode
+
+        last_state = get_layer_cache(self, past_key_values)
+
+        cu_seqlens = kwargs.get('cu_seqlens')
+        if attention_mask is not None:
+            indices, cu_seqlens, _ = get_unpad_data(attention_mask[:, -q_len:])
+            hidden_states = index_first_axis(rearrange(hidden_states, "b s ... -> (b s) ..."), indices).unsqueeze(0)
+
+        if self.use_short_conv:
+            conv_state_q, conv_state_k, conv_state_v = None, None, None
+            if last_state is not None:
+                conv_state_q, conv_state_k, conv_state_v = last_state['conv_state']
+            q, conv_state_q = self.q_conv1d(
+                x=self.q_proj(hidden_states),
+                cache=conv_state_q,
+                output_final_state=use_cache,
+                cu_seqlens=cu_seqlens,
+            )
+            k, conv_state_k = self.k_conv1d(
+                x=self.k_proj(hidden_states),
+                cache=conv_state_k,
+                output_final_state=use_cache,
+                cu_seqlens=cu_seqlens,
+            )
+            v, conv_state_v = self.v_conv1d(
+                x=self.v_proj(hidden_states),
+                cache=conv_state_v,
+                output_final_state=use_cache,
+                cu_seqlens=cu_seqlens,
+            )
+        else:
+            q = self.q_proj(hidden_states)
+            k = self.k_proj(hidden_states)
+            v = self.v_proj(hidden_states)
+        gk = self.gk_proj(hidden_states)
+
+        q = rearrange(q, '... (h d) -> ... h d', d=self.head_k_dim)
+        if self.num_kv_groups > 1:
+            k, gk = (repeat(x, '... (h d) -> ... (h g) d', g=self.num_kv_groups, d=self.head_k_dim) for x in (k, gk))
+            v = repeat(v, '... (h d) -> ... (h g) d', g=self.num_kv_groups, d=self.head_v_dim)
+        else:
+            k, gk = (rearrange(x, '... (h d) -> ... h d', d=self.head_k_dim) for x in (k, gk))
+            v = rearrange(v, '... (h d) -> ... h d', d=self.head_v_dim)
+
+        gk = F.logsigmoid(gk) / self.gate_logit_normalizer
+        if self.clamp_min is not None:
+            gk = torch.clamp_min(gk, self.clamp_min)
+
+        if self.feature_map_fn is not None:
+            q, k = map(self.feature_map_fn, (q, k))
+
+        recurrent_state = last_state['recurrent_state'] if last_state is not None else None
+        if mode == 'fused_recurrent':
+            o, recurrent_state = fused_recurrent_gla(
+                q=q,
+                k=k,
+                v=v,
+                gk=gk,
+                initial_state=recurrent_state,
+                output_final_state=use_cache,
+                cu_seqlens=cu_seqlens,
+            )
+        elif mode == 'fused_chunk':
+            o, recurrent_state = fused_chunk_gla(
+                q=q,
+                k=k,
+                v=v,
+                g=gk,
+                initial_state=recurrent_state,
+                output_final_state=use_cache,
+            )
+        elif mode == 'chunk':
+            o, recurrent_state = chunk_gla(
+                q=q,
+                k=k,
+                v=v,
+                g=gk,
+                initial_state=recurrent_state,
+                output_final_state=use_cache,
+                cu_seqlens=cu_seqlens,
+            )
+        else:
+            raise NotImplementedError(f"Not supported mode `{mode}`.")
+
+        update_layer_cache(
+            self,
+            past_key_values,
+            recurrent_state=recurrent_state,
+            conv_state=(conv_state_q, conv_state_k, conv_state_v) if self.use_short_conv else None,
+            offset=q_len,
+        )
+
+        if self.use_output_gate:
+            g = self.g_proj(hidden_states)
+            if self.fuse_norm_and_gate:
+                g = rearrange(g, '... (h d) -> ... h d', d=self.head_v_dim)
+                o = self.g_norm_swish_gate(o, g)
+                o = rearrange(o, '... h d -> ... (h d)')
+            else:
+                o = rearrange(self.g_norm(o), '... h d -> ... (h d)')
+                o = o * self.gate_fn(g)
+        else:
+            o = rearrange(self.g_norm(o), '... h d -> ... (h d)')
+        o = self.o_proj(o)
+        if attention_mask is not None:
+            o = pad_input(o.squeeze(0), indices, batch_size, q_len)
+
+        return o, None, past_key_values

fla/layers/gsa.py

@@ -0,0 +1,237 @@
+# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang
+
+from __future__ import annotations
+
+import warnings
+from typing import TYPE_CHECKING
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange, repeat
+
+from fla.layers.utils import get_layer_cache, get_unpad_data, index_first_axis, pad_input, update_layer_cache
+from fla.modules import RMSNorm, ShortConvolution
+from fla.modules.feature_map import ReLUFeatureMap, SwishFeatureMap, T2RFeatureMap
+from fla.modules.layernorm import rms_norm_linear
+from fla.ops.gsa import chunk_gsa, fused_recurrent_gsa
+
+if TYPE_CHECKING:
+    from transformers.processing_utils import Unpack
+
+    from fla.models.utils import Cache
+
+
+class GatedSlotAttention(nn.Module):
+
+    def __init__(
+        self,
+        mode: str = 'chunk',
+        hidden_size: int = 1024,
+        expand_k: float = 1.,
+        expand_v: float = 1.,
+        num_heads: int = 4,
+        num_kv_heads: int | None = None,
+        use_short_conv: bool = False,
+        conv_size: int = 4,
+        conv_bias: bool = False,
+        num_slots: int | None = None,
+        elementwise_affine: bool | None = True,
+        norm_eps: float = 1e-5,
+        gate_logit_normalizer: int = 8,
+        feature_map: str = 'swish',
+        use_output_gate: bool = False,
+        use_norm: bool = True,
+        layer_idx: int | None = None,
+        scale: float | None = 1.,
+        **kwargs,
+    ) -> GatedSlotAttention:
+        super().__init__()
+
+        self.mode = mode
+        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_heads if num_kv_heads is None else num_kv_heads
+        self.num_kv_groups = self.num_heads // self.num_kv_heads
+        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
+        self.head_k_dim = self.key_dim // self.num_heads
+        self.head_v_dim = self.value_dim // self.num_heads
+
+        self.use_short_conv = use_short_conv
+        self.conv_size = conv_size
+        self.conv_bias = conv_bias
+
+        self.gate_logit_normalizer = gate_logit_normalizer
+
+        self.use_output_gate = use_output_gate
+        self.use_norm = use_norm
+        self.scale = scale
+
+        if num_slots is None:
+            num_slots = self.head_k_dim
+        self.num_slots = num_slots
+
+        self.layer_idx = layer_idx
+
+        if layer_idx is None:
+            warnings.warn(
+                f"Instantiating {self.__class__.__name__} without passing `layer_idx` is not recommended and will "
+                "to errors during the forward call, if caching is used. Please make sure to provide a `layer_idx` "
+                "when creating this class.",
+            )
+
+        self.register_module('feature_map', None)
+        if feature_map == 'swish':
+            self.feature_map = SwishFeatureMap()
+        elif feature_map == 'relu':
+            self.feature_map = ReLUFeatureMap()
+        elif feature_map == 't2r':
+            self.feature_map = T2RFeatureMap(self.head_k_dim, self.head_k_dim)
+        else:
+            raise NotImplementedError(f"Feature map `{feature_map}` is not supported now.")
+
+        self.q_proj = nn.Linear(self.hidden_size, self.key_dim, bias=False)
+        self.k_proj = nn.Linear(self.hidden_size, self.key_dim_per_group, bias=False)
+        self.v_proj = nn.Linear(self.hidden_size, self.value_dim_per_group, bias=False)
+        self.f_proj = nn.Linear(self.hidden_size, self.num_kv_heads * self.num_slots, bias=False)
+
+        if use_short_conv:
+            self.conv_size = conv_size
+            self.q_conv1d = ShortConvolution(
+                hidden_size=self.key_dim,
+                kernel_size=conv_size,
+                bias=conv_bias,
+                activation='silu',
+            )
+            self.k_conv1d = ShortConvolution(
+                hidden_size=self.key_dim_per_group,
+                kernel_size=conv_size,
+                bias=conv_bias,
+                activation='silu',
+            )
+            self.v_conv1d = ShortConvolution(
+                hidden_size=self.value_dim_per_group,
+                kernel_size=conv_size,
+                bias=conv_bias,
+                activation='silu',
+            )
+
+        self.g_norm = RMSNorm(self.hidden_size, elementwise_affine, eps=norm_eps, dtype=torch.float32)
+        self.o_proj = nn.Linear(self.value_dim, self.hidden_size, bias=False)
+
+    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]:
+        if attention_mask is not None:
+            assert len(attention_mask.shape) == 2, (
+                "Expected attention_mask as a 0-1 matrix with shape [batch_size, seq_len] "
+                "for padding purposes (0 indicating padding). "
+                "Arbitrary attention masks of shape [batch_size, seq_len, seq_len] are not allowed."
+            )
+
+        batch_size, q_len, _ = hidden_states.shape
+        mode = 'fused_recurrent' if hidden_states.shape[1] <= 64 else self.mode
+
+        last_state = get_layer_cache(self, past_key_values)
+
+        cu_seqlens = kwargs.get('cu_seqlens')
+        if attention_mask is not None:
+            indices, cu_seqlens, _ = get_unpad_data(attention_mask[:, -q_len:])
+            hidden_states = index_first_axis(rearrange(hidden_states, "b s ... -> (b s) ..."), indices).unsqueeze(0)
+
+        if self.use_short_conv:
+            conv_state_q, conv_state_k, conv_state_v = None, None, None
+            if last_state is not None:
+                conv_state_q, conv_state_k, conv_state_v = last_state['conv_state']
+            q, conv_state_q = self.q_conv1d(
+                x=self.q_proj(hidden_states),
+                cache=conv_state_q,
+                output_final_state=use_cache,
+                cu_seqlens=cu_seqlens,
+            )
+            k, conv_state_k = self.k_conv1d(
+                x=self.k_proj(hidden_states),
+                cache=conv_state_k,
+                output_final_state=use_cache,
+                cu_seqlens=cu_seqlens,
+            )
+            v, conv_state_v = self.v_conv1d(
+                x=self.v_proj(hidden_states),
+                cache=conv_state_v,
+                output_final_state=use_cache,
+                cu_seqlens=cu_seqlens,
+            )
+        else:
+            q = self.q_proj(hidden_states)
+            k = self.k_proj(hidden_states)
+            v = self.v_proj(hidden_states)
+        f = self.f_proj(hidden_states)
+
+        q = rearrange(q, '... (h d) -> ... h d', d=self.head_k_dim)
+        k = rearrange(k, '... (h d) -> ... h d', d=self.head_k_dim)
+        v = rearrange(v, '... (h d) -> ... h d', d=self.head_v_dim)
+        f = rearrange(f, '... (h m) -> ... h m', m=self.num_slots)
+
+        if self.feature_map is not None:
+            q, k = map(lambda x: self.feature_map(x), (q, k))
+        v = F.silu(v)
+
+        f = F.logsigmoid(f) / self.gate_logit_normalizer
+        s = (1 - f.exp()).to(f.dtype)
+
+        if self.num_kv_groups > 1:
+            k, v, f, s = map(lambda x: repeat(x, '... h d -> ... (h g) d', g=self.num_kv_groups), (k, v, f, s))
+
+        recurrent_state = last_state['recurrent_state'] if last_state is not None else None
+        if mode == 'fused_recurrent':
+            o, recurrent_state = fused_recurrent_gsa(
+                q=q,
+                k=k,
+                v=v,
+                s=s,
+                g=f,
+                initial_state=recurrent_state,
+                output_final_state=use_cache,
+                scale=self.scale,
+                cu_seqlens=cu_seqlens,
+            )
+        elif mode == 'chunk':
+            o, recurrent_state = chunk_gsa(
+                q=q,
+                k=k,
+                v=v,
+                s=s,
+                g=f,
+                initial_state=recurrent_state,
+                output_final_state=use_cache,
+                scale=self.scale,
+                cu_seqlens=cu_seqlens,
+            )
+        else:
+            raise NotImplementedError(f"Not supported mode `{mode}`.")
+
+        update_layer_cache(
+            self,
+            past_key_values,
+            recurrent_state=recurrent_state,
+            conv_state=(conv_state_q, conv_state_k, conv_state_v) if self.use_short_conv else None,
+            offset=q_len,
+        )
+
+        o = rearrange(o, '... h d -> ... (h d)')
+        o = rms_norm_linear(F.silu(o), self.g_norm.weight, self.g_norm.bias, self.o_proj.weight, self.o_proj.bias)
+        if attention_mask is not None:
+            o = pad_input(o.squeeze(0), indices, batch_size, q_len)
+
+        return o, None, past_key_values

fla/layers/hgrn.py

@@ -0,0 +1,174 @@
+# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang
+
+# "Hierarchically Gated Recurrent Neural Network for Sequence Modeling" [https://arxiv.org/abs/2311.04823]
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from fla.layers.utils import get_layer_cache, update_layer_cache
+from fla.modules import FusedRMSNormGated, ShortConvolution
+from fla.modules.activations import swiglu
+from fla.ops.hgrn import chunk_hgrn, fused_recurrent_hgrn
+
+if TYPE_CHECKING:
+    from transformers.processing_utils import Unpack
+
+    from fla.models.utils import Cache
+
+
+class HGRNAttention(nn.Module):
+
+    def __init__(
+        self,
+        mode: str = 'chunk',
+        hidden_size: int = 1024,
+        expand_ratio: int | None = 1,
+        use_short_conv: bool = False,
+        conv_size: int = 4,
+        conv_bias: bool = False,
+        elementwise_affine: bool | None = True,
+        norm_eps: float = 1e-5,
+        layer_idx: int = None,
+    ) -> HGRNAttention:
+        super().__init__()
+
+        self.mode = mode
+        self.hidden_size = hidden_size
+        self.expand_ratio = expand_ratio
+        self.input_dim = int(hidden_size * expand_ratio)
+
+        self.use_short_conv = use_short_conv
+        self.conv_size = conv_size
+        self.conv_bias = conv_bias
+
+        self.layer_idx = layer_idx
+
+        assert mode in ['chunk', 'fused_recurrent'], f"Not supported mode `{mode}`."
+
+        self.i_proj = nn.Linear(hidden_size, self.input_dim, bias=False)
+        self.f_proj = nn.Linear(hidden_size, self.input_dim, bias=False)
+        self.g_proj = nn.Linear(hidden_size, self.input_dim, bias=False)
+
+        if use_short_conv:
+            self.conv_size = conv_size
+            self.f_conv1d = ShortConvolution(
+                hidden_size=self.input_dim,
+                kernel_size=conv_size,
+                bias=conv_bias,
+                activation=None,
+            )
+            self.i_conv1d = ShortConvolution(
+                hidden_size=self.input_dim,
+                kernel_size=conv_size,
+                bias=conv_bias,
+                activation=None,
+            )
+
+        self.g_norm = FusedRMSNormGated(
+            hidden_size=self.input_dim,
+            elementwise_affine=elementwise_affine,
+            eps=norm_eps,
+        )
+        self.o_proj = nn.Linear(self.input_dim, hidden_size, bias=False)
+
+    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,
+        lower_bound: torch.Tensor | None = None,
+        **kwargs: Unpack[dict],
+    ) -> tuple[torch.Tensor, torch.Tensor | None, Cache | None]:
+        if attention_mask is not None:
+            assert len(attention_mask.shape) == 2, (
+                "Expected attention_mask as a 0-1 matrix with shape [batch_size, seq_len] "
+                "for padding purposes (0 indicating padding). "
+                "Arbitrary attention masks of shape [batch_size, seq_len, seq_len] are not allowed."
+            )
+
+        # launching the triton kernel for just one token will actually be slower
+        mode = 'fused_recurrent' if not self.training and hidden_states.shape[1] <= 64 else self.mode
+
+        last_state = get_layer_cache(self, past_key_values)
+
+        cu_seqlens = kwargs.get('cu_seqlens')
+        if self.use_short_conv:
+            conv_state_i, conv_state_f = None, None
+            if last_state is not None:
+                conv_state_i, conv_state_f = last_state['conv_state']
+            conv_mask = attention_mask[:, -hidden_states.shape[1]:] if attention_mask is not None else None
+            i, conv_state_i = self.i_conv1d(
+                x=self.i_proj(hidden_states),
+                mask=conv_mask,
+                cache=conv_state_i,
+                output_final_state=use_cache,
+                cu_seqlens=cu_seqlens,
+            )
+            f, conv_state_f = self.f_conv1d(
+                x=self.f_proj(hidden_states),
+                mask=conv_mask,
+                cache=conv_state_f,
+                output_final_state=use_cache,
+                cu_seqlens=cu_seqlens,
+            )
+        else:
+            i = self.i_proj(hidden_states)
+            f = self.f_proj(hidden_states)
+
+        f = F.logsigmoid(f)
+        # the lower bound for the first layer is zero
+        if lower_bound is not None and self.layer_idx > 0:
+            f = torch.logaddexp(lower_bound.log(), torch.log1p(-lower_bound) + f).to(f)
+        i = swiglu(i, 1 - f.exp())
+
+        # dealing with left-padding
+        if attention_mask is not None:
+            i = i.mul(attention_mask[:, -i.shape[-2]:, None])
+
+        recurrent_state = last_state['recurrent_state'] if last_state is not None else None
+        if mode == 'chunk':
+            if cu_seqlens is not None:
+                raise NotImplementedError("Chunk mode does not support variable-length sequences.")
+            o, recurrent_state = chunk_hgrn(
+                x=i,
+                g=f,
+                initial_state=recurrent_state,
+                output_final_state=use_cache,
+            )
+        elif mode == 'fused_recurrent':
+            o, recurrent_state = fused_recurrent_hgrn(
+                x=i,
+                g=f,
+                initial_state=recurrent_state,
+                output_final_state=use_cache,
+                cu_seqlens=cu_seqlens,
+            )
+        else:
+            raise NotImplementedError(f"Not supported mode `{mode}`.")
+
+        update_layer_cache(
+            self,
+            past_key_values,
+            recurrent_state=recurrent_state,
+            conv_state=(conv_state_i, conv_state_f) if self.use_short_conv else None,
+            offset=i.shape[1],
+        )
+
+        o = self.g_norm(o, self.g_proj(hidden_states))
+        o = self.o_proj(o)
+
+        return o, None, past_key_values
+
+    def state_size(self, **kwargs) -> int:
+        state_size = self.hidden_size
+        for module in self.children():
+            if isinstance(module, ShortConvolution):
+                state_size += module.state_size
+        return state_size

fla/layers/hgrn2.py

@@ -0,0 +1,210 @@
+# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang
+
+# "HGRN2: Gated Linear RNNs with State Expansion"[https://arxiv.org/abs/2404.07904]
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange
+
+from fla.layers.utils import get_layer_cache, get_unpad_data, index_first_axis, pad_input, update_layer_cache
+from fla.modules import RMSNorm, ShortConvolution
+from fla.modules.activations import swish
+from fla.modules.layernorm import rms_norm_linear
+from fla.ops.gla import chunk_gla, fused_chunk_gla, fused_recurrent_gla
+
+if TYPE_CHECKING:
+    from transformers.processing_utils import Unpack
+
+    from fla.models.utils import Cache
+
+
+class HGRN2Attention(nn.Module):
+
+    def __init__(
+        self,
+        mode: str = 'chunk',
+        hidden_size: int = 1024,
+        num_heads: int | None = None,
+        expand_ratio: int | None = 128,
+        use_short_conv: bool = False,
+        conv_size: int = 4,
+        conv_bias: bool = False,
+        elementwise_affine: bool | None = True,
+        norm_eps: float = 1e-5,
+        layer_idx: int = None,
+    ) -> HGRN2Attention:
+        super().__init__()
+
+        self.mode = mode
+        self.hidden_size = hidden_size
+
+        if expand_ratio is not None:
+            num_heads = hidden_size // expand_ratio
+        elif expand_ratio is None and num_heads is not None:
+            expand_ratio = hidden_size // num_heads
+        elif expand_ratio is None and num_heads is None:
+            raise RuntimeError("One of `expand_ratio` or `num_heads` should be provided.")
+        self.num_heads = num_heads
+        self.expand_ratio = expand_ratio
+
+        self.use_short_conv = use_short_conv
+        self.conv_size = conv_size
+        self.conv_bias = conv_bias
+
+        self.forget_dim = int(self.num_heads * self.expand_ratio)
+        self.input_dim = hidden_size
+        self.layer_idx = layer_idx
+
+        assert mode in ['chunk', 'fused_recurrent', 'fused_chunk'], f"Not supported mode `{mode}`."
+        assert self.forget_dim % num_heads == 0, f"forget dim must be divisible by num_heads of {num_heads}"
+        assert self.input_dim % num_heads == 0, f"input dim must be divisible by num_heads of {num_heads}"
+
+        self.head_f_dim = self.expand_ratio
+        self.head_i_dim = self.hidden_size // num_heads
+
+        self.q_proj = nn.Linear(hidden_size, self.forget_dim, bias=False)
+        self.f_proj = nn.Linear(hidden_size, self.forget_dim, bias=False)
+        self.i_proj = nn.Linear(hidden_size, self.input_dim, bias=False)
+
+        if use_short_conv:
+            self.conv_size = conv_size
+            self.q_conv1d = ShortConvolution(
+                hidden_size=self.forget_dim,
+                kernel_size=conv_size,
+                bias=conv_bias,
+                activation=None,
+            )
+            self.f_conv1d = ShortConvolution(
+                hidden_size=self.forget_dim,
+                kernel_size=conv_size,
+                bias=conv_bias,
+                activation=None,
+            )
+            self.i_conv1d = ShortConvolution(
+                hidden_size=self.input_dim,
+                kernel_size=conv_size,
+                bias=conv_bias,
+                activation=None,
+            )
+
+        self.g_norm = RMSNorm(hidden_size=self.hidden_size, elementwise_affine=elementwise_affine,
+                              eps=norm_eps, dtype=torch.float32)
+        self.o_proj = nn.Linear(self.input_dim, hidden_size, bias=False)
+
+    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,
+        lower_bound: torch.Tensor | None = None,
+        **kwargs: Unpack[dict],
+    ) -> tuple[torch.Tensor, torch.Tensor | None, Cache | None]:
+        if attention_mask is not None:
+            assert len(attention_mask.shape) == 2, (
+                "Expected attention_mask as a 0-1 matrix with shape [batch_size, seq_len] "
+                "for padding purposes (0 indicating padding). "
+                "Arbitrary attention masks of shape [batch_size, seq_len, seq_len] are not allowed."
+            )
+
+        batch_size, q_len, _ = hidden_states.shape
+        mode = 'fused_recurrent' if hidden_states.shape[1] <= 64 else self.mode
+
+        last_state = get_layer_cache(self, past_key_values)
+
+        cu_seqlens = kwargs.get('cu_seqlens')
+        if attention_mask is not None:
+            indices, cu_seqlens, _ = get_unpad_data(attention_mask[:, -q_len:])
+            hidden_states = index_first_axis(rearrange(hidden_states, "b s ... -> (b s) ..."), indices).unsqueeze(0)
+
+        if self.use_short_conv:
+            conv_state_q, conv_state_f, conv_state_i = None, None, None
+            if last_state is not None:
+                conv_state_q, conv_state_f, conv_state_i = last_state['conv_state']
+            q, conv_state_q = self.q_conv1d(
+                x=self.q_proj(hidden_states),
+                cache=conv_state_q,
+                output_final_state=use_cache,
+                cu_seqlens=cu_seqlens,
+            )
+            f, conv_state_f = self.f_conv1d(
+                x=self.f_proj(hidden_states),
+                cache=conv_state_f,
+                output_final_state=use_cache,
+                cu_seqlens=cu_seqlens,
+            )
+            i, conv_state_i = self.i_conv1d(
+                x=self.i_proj(hidden_states),
+                cache=conv_state_i,
+                output_final_state=use_cache,
+                cu_seqlens=cu_seqlens,
+            )
+        else:
+            q = self.q_proj(hidden_states)
+            f = self.f_proj(hidden_states)
+            i = self.i_proj(hidden_states)
+
+        q = swish(q)
+
+        g = F.logsigmoid(f)
+        # the lower bound for the first layer is zero
+        if lower_bound is not None and self.layer_idx > 0:
+            g = torch.logaddexp(lower_bound.log(), torch.log1p(-lower_bound) + g)
+        k = 1 - g.exp()
+
+        q, k, g = map(lambda x: rearrange(x, '... (h d) -> ... h d', d=self.head_f_dim), (q, k.to(i), g))
+        i = rearrange(i, '... (h d) -> ... h d', d=self.head_i_dim)
+
+        recurrent_state = last_state['recurrent_state'] if last_state is not None else None
+        if mode == 'fused_recurrent':
+            o, recurrent_state = fused_recurrent_gla(
+                q=q,
+                k=k,
+                v=i,
+                gk=g,
+                initial_state=recurrent_state,
+                output_final_state=use_cache,
+                cu_seqlens=cu_seqlens,
+            )
+        elif mode == 'fused_chunk':
+            o, recurrent_state = fused_chunk_gla(
+                q=q,
+                k=k,
+                v=i,
+                g=g,
+                initial_state=recurrent_state,
+                output_final_state=use_cache,
+            )
+        elif mode == 'chunk':
+            o, recurrent_state = chunk_gla(
+                q=q,
+                k=k,
+                v=i,
+                g=g,
+                initial_state=recurrent_state,
+                output_final_state=use_cache,
+                cu_seqlens=cu_seqlens,
+            )
+        else:
+            raise NotImplementedError(f"Not supported mode `{mode}`.")
+
+        update_layer_cache(
+            self,
+            past_key_values,
+            recurrent_state=recurrent_state,
+            conv_state=(conv_state_q, conv_state_f, conv_state_i) if self.use_short_conv else None,
+            offset=q_len,
+        )
+
+        o = rearrange(o, '... h d -> ... (h d)')
+        o = rms_norm_linear(o, self.g_norm.weight, self.g_norm.bias, self.o_proj.weight, self.o_proj.bias)
+        if attention_mask is not None:
+            o = pad_input(o.squeeze(0), indices, batch_size, q_len)
+
+        return o, None, past_key_values

fla/layers/kda.py

@@ -0,0 +1,277 @@
+# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang
+
+from __future__ import annotations
+
+import math
+from typing import TYPE_CHECKING
+
+import torch
+import torch.nn as nn
+from einops import rearrange, repeat
+from torch.nn import functional as F
+
+from fla.layers.utils import get_layer_cache, get_unpad_data, index_first_axis, pad_input, update_layer_cache
+from fla.modules import FusedRMSNormGated, ShortConvolution
+from fla.ops.kda import chunk_kda, fused_recurrent_kda
+from fla.ops.kda.gate import fused_kda_gate
+
+if TYPE_CHECKING:
+    from transformers.processing_utils import Unpack
+
+    from fla.models.utils import Cache
+
+
+class KimiDeltaAttention(nn.Module):
+    """
+    Kimi Delta Attention (KDA) layer implementation.
+
+    Args:
+        hidden_size (int, Optional):
+            The hidden size of the input. Default: 2048.
+        expand_v (float, Optional):
+            The expansion ratio for the value dimension. Default: 1.0.
+        head_dim (int, Optional):
+            The dimension of each head. Default: 128.
+        num_heads (int, Optional):
+            The number of heads. Default: 16.
+        num_v_heads (int, Optional):
+            The number of heads for the value projection, equal to `num_heads` if `None`.
+            GVA (Grouped Value Attention) is applied if `num_v_heads` > `num_heads`. Default: `None`.
+        mode (str, Optional):
+            Which Kimi Delta Attention kernel to use.
+            Currently available: `chunk` and `fused_recurrent`.
+            Default: `chunk`.
+        use_short_conv (bool, Optional):
+            Whether to use short convolutions. Default: `True`.
+        allow_neg_eigval (bool, Optional):
+            Allow negative eigenvalues. Default: `False`. If set to `True`, the beta will be multiplied by 2.
+            See reference:
+            [Unlocking State-Tracking in Linear RNNs Through Negative Eigenvalues](https://arxiv.org/abs/2411.12537)
+        conv_size (int, Optional):
+            The kernel size of the short convolution, only used when `use_short_conv` is `True`. Default: 4.
+        conv_bias (bool, Optional):
+            Whether to use bias in the short convolution, only used when `use_short_conv` is `True`. Default: `False`.
+        layer_idx (int, Optional):
+            The index of the layer. Default: None.
+        norm_eps (float, Optional):
+            The epsilon value for the normalization layer. Default: 1e-5.
+    """
+
+    def __init__(
+        self,
+        hidden_size: int = 2048,
+        expand_v: float = 1,
+        head_dim: int = 128,
+        num_heads: int = 16,
+        num_v_heads: int = None,
+        mode: str = "chunk",
+        use_short_conv: bool = True,
+        allow_neg_eigval: bool = False,
+        conv_size: int = 4,
+        conv_bias: bool = False,
+        layer_idx: int = None,
+        norm_eps: float = 1e-5,
+        **kwargs,
+    ) -> KimiDeltaAttention:
+        super().__init__()
+
+        self.mode = mode
+        self.allow_neg_eigval = allow_neg_eigval
+        self.hidden_size = hidden_size
+        self.expand_v = expand_v
+
+        self.use_short_conv = use_short_conv
+        self.conv_size = conv_size
+        self.conv_bias = conv_bias
+
+        self.head_dim = head_dim
+        self.num_heads = num_heads
+        self.num_v_heads = num_v_heads if num_v_heads is not None else num_heads
+
+        self.head_k_dim = head_dim
+        self.head_v_dim = int(self.head_dim * self.expand_v)
+        self.key_dim = int(self.num_heads * self.head_k_dim)
+        self.value_dim = int(self.num_v_heads * self.head_v_dim)
+        self.layer_idx = layer_idx
+
+        # Consistency check: Ensure expand_v produces integer values
+        if not math.isclose(self.num_v_heads * self.head_dim * expand_v, self.value_dim, rel_tol=1e-5):
+            raise ValueError(
+                f"expand_v={expand_v} does not produce an integer value when multiplied by key_dim={self.key_dim}. "
+                f"Resulting value_dim would be {self.num_v_heads * self.head_dim * expand_v}, which is invalid for nn.Linear.",
+            )
+        if self.num_v_heads > self.num_heads and self.num_v_heads % self.num_heads != 0:
+            raise ValueError(
+                f"num_v_heads={self.num_v_heads} must be divisible by num_heads={self.num_heads}.",
+            )
+
+        if not math.isclose(head_dim * expand_v, self.head_v_dim, rel_tol=1e-5):
+            raise ValueError(
+                f"expand_v={expand_v} does not produce an integer value when multiplied by head_dim={head_dim}. "
+                f"Resulting head_v_dim would be {head_dim * expand_v}, which is invalid for FusedRMSNormGated.",
+            )
+        assert mode in ["chunk", "fused_recurrent"], f"Not supported mode `{mode}`."
+
+        self.q_proj = nn.Linear(hidden_size, self.key_dim, bias=False)
+        self.k_proj = nn.Linear(hidden_size, self.key_dim, bias=False)
+        self.v_proj = nn.Linear(hidden_size, self.value_dim, bias=False)
+
+        if use_short_conv:
+            self.q_conv1d = ShortConvolution(
+                hidden_size=self.key_dim,
+                kernel_size=conv_size,
+                bias=conv_bias,
+                activation="silu",
+            )
+            self.k_conv1d = ShortConvolution(
+                hidden_size=self.key_dim,
+                kernel_size=conv_size,
+                bias=conv_bias,
+                activation="silu",
+            )
+            self.v_conv1d = ShortConvolution(
+                hidden_size=self.value_dim,
+                kernel_size=conv_size,
+                bias=conv_bias,
+                activation="silu",
+            )
+
+        self.f_proj = nn.Sequential(
+            nn.Linear(hidden_size, self.head_v_dim, bias=False),
+            nn.Linear(self.head_v_dim, self.key_dim, bias=False),
+        )
+        self.b_proj = nn.Linear(hidden_size, self.num_heads, bias=False)
+
+        self.A_log = nn.Parameter(torch.log(torch.empty(self.num_heads, dtype=torch.float32).uniform_(1, 16)))
+        self.A_log._no_weight_decay = True
+        dt = torch.exp(
+            torch.rand(self.key_dim, dtype=torch.float32) * (math.log(0.1) - math.log(0.001)) + math.log(0.001)
+        ).clamp(min=1e-4)
+        inv_dt = dt + torch.log(-torch.expm1(-dt))
+        self.dt_bias = nn.Parameter(inv_dt)
+        self.dt_bias._no_weight_decay = True
+
+        self.g_proj = nn.Sequential(
+            nn.Linear(hidden_size, self.head_v_dim, bias=False),
+            nn.Linear(self.head_v_dim, self.value_dim, bias=True),
+        )
+        self.o_norm = FusedRMSNormGated(self.head_v_dim, activation="sigmoid", eps=norm_eps)
+        self.o_proj = nn.Linear(self.value_dim, hidden_size, bias=False)
+
+    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]:
+        if attention_mask is not None:
+            assert len(attention_mask.shape) == 2, (
+                "Expected attention_mask as a 0-1 matrix with shape [batch_size, seq_len] "
+                "for padding purposes (0 indicating padding). "
+                "Arbitrary attention masks of shape [batch_size, seq_len, seq_len] are not allowed."
+            )
+
+        batch_size, q_len, _ = hidden_states.shape
+        # change to inference mode.
+        mode = "fused_recurrent" if (q_len <= 64 and not self.training) else self.mode
+        if self.training:
+            assert mode == "chunk", "Only chunk mode is supported in training."
+
+        last_state = get_layer_cache(self, past_key_values)
+
+        cu_seqlens = kwargs.get("cu_seqlens")
+        if attention_mask is not None:
+            indices, cu_seqlens, _ = get_unpad_data(attention_mask[:, -q_len:])
+            hidden_states = index_first_axis(rearrange(hidden_states, "b s ... -> (b s) ..."), indices).unsqueeze(0)
+
+        if self.use_short_conv:
+            conv_state_q, conv_state_k, conv_state_v = None, None, None
+            if last_state is not None:
+                conv_state_q, conv_state_k, conv_state_v = last_state["conv_state"]
+            q, conv_state_q = self.q_conv1d(
+                x=self.q_proj(hidden_states),
+                cache=conv_state_q,
+                output_final_state=use_cache,
+                cu_seqlens=cu_seqlens,
+            )
+            k, conv_state_k = self.k_conv1d(
+                x=self.k_proj(hidden_states),
+                cache=conv_state_k,
+                output_final_state=use_cache,
+                cu_seqlens=cu_seqlens,
+            )
+            v, conv_state_v = self.v_conv1d(
+                x=self.v_proj(hidden_states),
+                cache=conv_state_v,
+                output_final_state=use_cache,
+                cu_seqlens=cu_seqlens,
+            )
+        else:
+            q = F.silu(self.q_proj(hidden_states))
+            k = F.silu(self.k_proj(hidden_states))
+            v = F.silu(self.v_proj(hidden_states))
+
+        g = self.f_proj(hidden_states)
+        beta = self.b_proj(hidden_states).sigmoid()
+
+        q, k, g = (rearrange(x, "... (h d) -> ... h d", d=self.head_k_dim) for x in (q, k, g))
+        v = rearrange(v, "... (h d) -> ... h d", d=self.head_v_dim)
+
+        # for multi-value attention, we repeat the inputs for simplicity.
+        if self.num_v_heads > self.num_heads:
+            q, k, g = (repeat(x, "... h d -> ... (h g) d", g=self.num_v_heads // self.num_heads) for x in (q, k, g))
+            beta = repeat(beta, "... h -> ... (h g)", g=self.num_v_heads // self.num_heads)
+
+        if self.allow_neg_eigval:
+            beta = beta * 2.0
+
+        recurrent_state = last_state["recurrent_state"] if last_state is not None else None
+        if mode == "chunk":
+            o, recurrent_state = chunk_kda(
+                q=q,
+                k=k,
+                v=v,
+                g=g,
+                beta=beta,
+                A_log=self.A_log,
+                dt_bias=self.dt_bias,
+                initial_state=recurrent_state,
+                output_final_state=use_cache,
+                use_qk_l2norm_in_kernel=True,
+                use_gate_in_kernel=True,
+                cu_seqlens=cu_seqlens,
+            )
+        elif mode == "fused_recurrent":
+            g = fused_kda_gate(g=g, A_log=self.A_log, dt_bias=self.dt_bias)
+            o, recurrent_state = fused_recurrent_kda(
+                q=q,
+                k=k,
+                v=v,
+                g=g,
+                beta=beta,
+                initial_state=recurrent_state,
+                output_final_state=use_cache,
+                use_qk_l2norm_in_kernel=True,
+                cu_seqlens=cu_seqlens,
+            )
+        else:
+            raise NotImplementedError(f"Not supported mode `{mode}`.")
+
+        update_layer_cache(
+            self,
+            past_key_values,
+            recurrent_state=recurrent_state,
+            conv_state=(conv_state_q, conv_state_k, conv_state_v) if self.use_short_conv else None,
+            offset=q_len,
+        )
+
+        o = self.o_norm(o, rearrange(self.g_proj(hidden_states), "... (h d) -> ... h d", d=self.head_v_dim))
+        o = rearrange(o, "b t h d -> b t (h d)")
+        o = self.o_proj(o)
+        if attention_mask is not None:
+            o = pad_input(o.squeeze(0), indices, batch_size, q_len)
+
+        return o, None, past_key_values

fla/layers/lightnet.py

@@ -0,0 +1,238 @@
+# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang
+
+# ["You Only Scan Once: Efficient Multi-dimension Sequential Modeling with LightNet"](https://arxiv.org/abs/2405.21022)
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange
+
+from fla.layers.utils import get_layer_cache, update_layer_cache
+from fla.modules import FusedRMSNormGated, ShortConvolution
+from fla.modules.fused_norm_gate import rms_norm_swish_gate_linear
+from fla.ops.gla import chunk_gla, fused_recurrent_gla
+
+if TYPE_CHECKING:
+    from transformers.processing_utils import Unpack
+
+    from fla.models.utils import Cache
+
+
+class LightNetAttention(nn.Module):
+
+    def __init__(
+        self,
+        mode: str = 'chunk',
+        hidden_size: int = 1024,
+        num_heads: int | None = None,
+        expand_ratio: int | None = 128,
+        use_short_conv: bool = False,
+        conv_size: int = 4,
+        conv_bias: bool = False,
+        gate_low_rank_dim: int = 128,
+        elementwise_affine: bool | None = True,
+        norm_eps: float = 1e-5,
+        layer_idx: int = None,
+    ) -> LightNetAttention:
+        super().__init__()
+
+        self.mode = mode
+        self.hidden_size = hidden_size
+
+        if expand_ratio is None and num_heads is not None:
+            expand_ratio = hidden_size // num_heads
+        elif expand_ratio is not None and num_heads is None:
+            num_heads = hidden_size // expand_ratio
+        elif expand_ratio is None and num_heads is None:
+            raise RuntimeError("One of `expand_ratio` or `num_heads` should be provided.")
+        self.num_heads = num_heads
+        self.expand_ratio = expand_ratio
+
+        self.use_short_conv = use_short_conv
+        self.conv_size = conv_size
+        self.conv_bias = conv_bias
+
+        self.key_dim = int(self.num_heads * self.expand_ratio)
+        self.value_dim = hidden_size
+        self.gate_low_rank_dim = gate_low_rank_dim
+        self.layer_idx = layer_idx
+
+        assert mode in ['chunk', 'fused_chunk'], f"Not supported mode `{mode}`."
+        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_f_dim = self.expand_ratio
+        self.head_i_dim = self.hidden_size // num_heads
+
+        self.q_proj = nn.Linear(hidden_size, self.key_dim, bias=False)
+        self.k_proj = nn.Linear(hidden_size, self.key_dim, bias=False)
+        self.v_proj = nn.Linear(hidden_size, self.value_dim, bias=False)
+
+        if use_short_conv:
+            self.conv_size = conv_size
+            self.q_conv1d = ShortConvolution(
+                hidden_size=self.key_dim,
+                kernel_size=conv_size,
+                bias=conv_bias,
+                activation=None,
+            )
+            self.k_conv1d = ShortConvolution(
+                hidden_size=self.key_dim,
+                kernel_size=conv_size,
+                bias=conv_bias,
+                activation=None,
+            )
+            self.v_conv1d = ShortConvolution(
+                hidden_size=self.value_dim,
+                kernel_size=conv_size,
+                bias=conv_bias,
+                activation=None,
+            )
+
+        self.g_proj = nn.Sequential(
+            nn.Linear(hidden_size, gate_low_rank_dim, bias=False),
+            nn.Linear(gate_low_rank_dim, hidden_size, bias=False),
+        )
+        self.g_norm = FusedRMSNormGated(
+            hidden_size=hidden_size,
+            elementwise_affine=elementwise_affine,
+            eps=norm_eps,
+        )
+        self.o_proj = nn.Linear(self.value_dim, hidden_size, bias=False)
+
+    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]:
+        if attention_mask is not None:
+            assert len(attention_mask.shape) == 2, (
+                "Expected attention_mask as a 0-1 matrix with shape [batch_size, seq_len] "
+                "for padding purposes (0 indicating padding). "
+                "Arbitrary attention masks of shape [batch_size, seq_len, seq_len] are not allowed."
+            )
+
+        # launching the triton kernel for just one token will actually be slower
+        mode = 'fused_recurrent' if hidden_states.shape[1] <= 64 else self.mode
+
+        last_state = get_layer_cache(self, past_key_values)
+
+        cu_seqlens = kwargs.get('cu_seqlens')
+        if self.use_short_conv:
+            conv_state_q, conv_state_k, conv_state_v = None, None, None
+            if last_state is not None:
+                conv_state_q, conv_state_k, conv_state_v = last_state['conv_state']
+            conv_mask = attention_mask[:, -hidden_states.shape[1]:] if attention_mask is not None else None
+            q, conv_state_q = self.q_conv1d(
+                x=self.q_proj(hidden_states),
+                mask=conv_mask,
+                cache=conv_state_q,
+                output_final_state=use_cache,
+                cu_seqlens=cu_seqlens,
+            )
+            k, conv_state_k = self.k_conv1d(
+                x=self.k_proj(hidden_states),
+                mask=conv_mask,
+                cache=conv_state_k,
+                output_final_state=use_cache,
+                cu_seqlens=cu_seqlens,
+            )
+            v, conv_state_v = self.v_conv1d(
+                x=self.v_proj(hidden_states),
+                mask=conv_mask,
+                cache=conv_state_v,
+                output_final_state=use_cache,
+                cu_seqlens=cu_seqlens,
+            )
+        else:
+            q = self.q_proj(hidden_states)
+            k = self.k_proj(hidden_states)
+            v = self.v_proj(hidden_states)
+
+        # dealing with left-padding
+        if attention_mask is not None:
+            v = v.mul(attention_mask[:, -v.shape[-2]:, None])
+
+        q = F.silu(q)
+        q, k = map(lambda x: rearrange(x, '... (h d) -> ... h d', d=self.head_f_dim), (q, k))
+        v = rearrange(v, '... (h d) -> ... h d', d=self.head_i_dim)
+        # TODO: this 2 steps took huge amount of time, which should be optimized
+        last_z = last_state['ffn_state'] if last_state is not None and last_state.get('ffn_state') is not None else None
+        if last_z is not None:
+            # Decode path: continue logcumsumexp from cached state
+            z = torch.logaddexp(last_z, k.float())
+            k, g = torch.exp(k - z).to(k.dtype), (last_z - z).to(k.dtype)
+        else:
+            # Prefill path: mask padding positions to -inf so they don't affect logcumsumexp
+            if cu_seqlens is not None:
+                raise NotImplementedError("LightNet does not support variable-length sequences for now.")
+            k_float = k.float()
+            if attention_mask is not None:
+                pad_mask = attention_mask[:, -k.shape[1]:, None, None]  # (B, T, 1, 1)
+                k_for_z = k_float.masked_fill(pad_mask == 0, float('-inf'))
+            else:
+                k_for_z = k_float
+            z = k_for_z.logcumsumexp(1)
+            k_new = torch.exp(k_float - z)
+            g_new = torch.cat((z[:, :1], z[:, :-1]), 1) - z
+            # NaN/inf arise at fully-masked positions (-inf - (-inf)), zero them out
+            k = torch.nan_to_num(k_new, nan=0.0, posinf=0.0).to(k.dtype)
+            g = torch.nan_to_num(g_new, nan=0.0, posinf=0.0, neginf=0.0).to(k.dtype)
+
+        recurrent_state = last_state['recurrent_state'] if last_state is not None else None
+        if mode == 'fused_recurrent':
+            o, recurrent_state = fused_recurrent_gla(
+                q=q,
+                k=k,
+                v=v,
+                gk=g,
+                initial_state=recurrent_state,
+                output_final_state=use_cache,
+                cu_seqlens=cu_seqlens,
+            )
+        elif mode == 'chunk':
+            o, recurrent_state = chunk_gla(
+                q=q,
+                k=k,
+                v=v,
+                g=g,
+                initial_state=recurrent_state,
+                output_final_state=use_cache,
+                cu_seqlens=cu_seqlens,
+            )
+        else:
+            raise NotImplementedError(f"Not supported mode `{mode}`.")
+
+        update_layer_cache(
+            self,
+            past_key_values,
+            recurrent_state=recurrent_state,
+            conv_state=(conv_state_q, conv_state_k, conv_state_v) if self.use_short_conv else None,
+            ffn_state=z[:, -1:],
+            offset=q.shape[1],
+        )
+
+        o = rms_norm_swish_gate_linear(
+            rearrange(o, 'b t h d -> b t (h d)'),
+            self.g_proj(hidden_states),
+            self.g_norm.weight,
+            self.g_norm.bias,
+            self.o_proj.weight,
+            self.o_proj.bias,
+        )
+        return o, None, past_key_values
+
+    def state_size(self, **kwargs) -> int:
+        state_size = self.key_dim * self.head_i_dim
+        for module in self.children():
+            if isinstance(module, ShortConvolution):
+                state_size += module.state_size
+        return state_size

fla/layers/linear_attn.py

@@ -0,0 +1,196 @@
+# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange, repeat
+
+from fla.layers.utils import get_layer_cache, update_layer_cache
+from fla.modules import RMSNorm
+from fla.modules.feature_map import DPFPFeatureMap, HadamardFeatureMap, HedgehogFeatureMap, T2RFeatureMap
+from fla.ops.linear_attn import chunk_linear_attn, fused_chunk_linear_attn, fused_recurrent_linear_attn
+
+if TYPE_CHECKING:
+    from fla.models.utils import Cache
+
+
+class LinearAttention(nn.Module):
+
+    def __init__(
+        self,
+        mode: str = 'chunk',
+        hidden_size: str = 1024,
+        expand_k: float = 1.0,
+        expand_v: float = 1.0,
+        num_heads: int = 8,
+        num_kv_heads: int | None = None,
+        feature_map: str = 'elementwise_product',
+        tie_feature_map_qk: bool = False,
+        output_norm: str = 'rmsnorm',
+        norm_q: bool = False,
+        norm_k: bool = False,
+        do_feature_map_norm: bool = False,
+        elementwise_affine: bool = True,
+        norm_eps: float = 1e-5,
+        layer_idx: int | None = None,
+        **kwargs,
+    ):
+        super().__init__()
+
+        self.hidden_size = hidden_size
+        self.mode = mode
+        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.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 mode in ['chunk', 'fused_chunk', 'fused_recurrent'], f"Not supported mode `{mode}`."
+        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
+        self.do_feature_map_norm = do_feature_map_norm
+        self.layer_idx = layer_idx
+
+        if feature_map == 'hedgehog':
+            if tie_feature_map_qk:
+                self.feature_map_q = self.feature_map_k = HedgehogFeatureMap(head_dim=self.head_k_dim)
+            else:
+                self.feature_map_q = HedgehogFeatureMap(head_dim=self.head_k_dim)
+                self.feature_map_k = HedgehogFeatureMap(head_dim=self.head_k_dim)
+
+        elif feature_map == 't2r':
+            if tie_feature_map_qk:
+                self.feature_map_q = self.feature_map_k = T2RFeatureMap(head_dim=self.head_k_dim)
+            else:
+                self.feature_map_q = T2RFeatureMap(head_dim=self.head_k_dim)
+                self.feature_map_k = T2RFeatureMap(head_dim=self.head_k_dim)
+
+        elif feature_map == 'elementwise_product':
+            if tie_feature_map_qk:
+                self.feature_map_q = self.feature_map_k = HadamardFeatureMap(head_dim=self.head_k_dim)
+            else:
+                self.feature_map_q = HadamardFeatureMap(head_dim=self.head_k_dim)
+                self.feature_map_k = HadamardFeatureMap(head_dim=self.head_k_dim)
+
+        elif feature_map == 'dpfp':
+            self.feature_map_q = DPFPFeatureMap(head_dim=self.head_k_dim)
+            self.feature_map_k = DPFPFeatureMap(head_dim=self.head_k_dim)
+
+        elif feature_map == 'elu':
+            def elu(x):
+                return F.elu(x) + 1
+            self.feature_map_q = elu
+            self.feature_map_k = elu
+
+        elif feature_map == 'relu':
+            self.feature_map_q = nn.ReLU()
+            self.feature_map_k = nn.ReLU()
+
+        elif feature_map == 'identity':
+            self.feature_map_q = nn.Identity()
+            self.feature_map_k = nn.Identity()
+        else:
+            raise NotImplementedError(f"Not supported feature map `{feature_map}`.")
+
+        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)
+
+        if output_norm == 'rmsnorm':
+            self.norm = RMSNorm(hidden_size=self.head_v_dim, elementwise_affine=elementwise_affine,
+                                eps=norm_eps, dtype=torch.float32)
+        elif output_norm == 'identity':
+            self.norm = nn.Identity()
+        else:
+            raise NotImplementedError(f"Not supported output norm `{output_norm}`.")
+
+        self.o_proj = nn.Linear(self.value_dim, hidden_size, bias=False)
+
+        self.norm_q = norm_q
+        self.norm_k = norm_k
+
+    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,
+    ) -> tuple[torch.Tensor, torch.Tensor | None, Cache | None]:
+        # Match other recurrent layers: use the recurrent kernel for decode/small chunks.
+        mode = 'fused_recurrent' if hidden_states.shape[1] <= 64 else self.mode
+        last_state = get_layer_cache(self, past_key_values)
+
+        q = self.q_proj(hidden_states)
+        k = self.k_proj(hidden_states)
+        v = self.v_proj(hidden_states)
+
+        if attention_mask is not None:
+            v = v.mul(attention_mask[:, -v.shape[-2]:, None])
+
+        q = rearrange(q, '... (h d) -> ... h d', d=self.head_k_dim)
+        if self.num_kv_groups > 1:
+            k = repeat(k, '... (h d) -> ... (h g) d', d=self.head_k_dim, g=self.num_kv_groups)
+            v = repeat(v, '... (h d) -> ... (h g) d', d=self.head_v_dim, g=self.num_kv_groups)
+        else:
+            k = rearrange(k, '... (h d) -> ... h d', d=self.head_k_dim)
+            v = rearrange(v, '... (h d) -> ... h d', d=self.head_v_dim)
+
+        q = self.feature_map_q(q)
+        k = self.feature_map_k(k)
+
+        if self.norm_q:
+            q = q / (q.sum(-1, True) + 1e-4)
+        if self.norm_k:
+            k = k / (k.sum(-1, True) + 1e-4)
+
+        recurrent_state = last_state['recurrent_state'] if last_state is not None else None
+        if mode == 'chunk':
+            o, final_state = chunk_linear_attn(
+                q=q,
+                k=k,
+                v=v,
+                initial_state=recurrent_state,
+                output_final_state=use_cache,
+                normalize=self.do_feature_map_norm,
+            )
+        elif mode == 'fused_chunk':
+            o, final_state = fused_chunk_linear_attn(
+                q=q,
+                k=k,
+                v=v,
+                initial_state=recurrent_state,
+                output_final_state=use_cache,
+                normalize=self.do_feature_map_norm,
+            )
+        elif mode == 'fused_recurrent':
+            o, final_state = fused_recurrent_linear_attn(
+                q=q,
+                k=k,
+                v=v,
+                initial_state=recurrent_state,
+                output_final_state=use_cache,
+                normalize=self.do_feature_map_norm,
+            )
+        else:
+            raise NotImplementedError
+        update_layer_cache(
+            self,
+            past_key_values,
+            recurrent_state=final_state,
+            offset=q.shape[1],
+        )
+        o = self.norm(o)
+        o = rearrange(o, '... h d -> ... (h d)')
+        o = self.o_proj(o)
+        return o, None, past_key_values

fla/layers/log_linear_mamba2.py

@@ -0,0 +1,684 @@
+from __future__ import annotations
+
+import math
+from typing import TYPE_CHECKING
+
+import torch
+import torch.nn as nn
+from einops import rearrange
+from transformers.activations import ACT2FN
+from transformers.utils import logging
+
+from fla.layers.mamba2 import apply_mask_to_padding_states, causal_conv1d_fn, causal_conv1d_update, is_fast_path_available
+from fla.layers.utils import get_layer_cache, update_layer_cache
+from fla.modules.layernorm_gated import RMSNormGated, rmsnorm_fn
+from fla.ops.log_linear_attn.chunk import LogLinearAttentionState, chunk_log_linear_attn
+
+if TYPE_CHECKING:
+    from fla.models.utils import Cache
+
+logger = logging.get_logger(__name__)
+
+
+def ceil_log(x: int, b: int) -> int:
+    return math.ceil(math.log(x, b))
+
+
+def get_num_levels(length: int, base: int) -> int:
+    return ceil_log(length, base) + 1
+
+
+MAX_SEQUENCE_LENGTH = 2048 * 8
+LAMBDA_LEVEL_BASE = 2
+MAX_NUM_LEVELS = get_num_levels(length=MAX_SEQUENCE_LENGTH, base=LAMBDA_LEVEL_BASE)
+
+
+def hmamba_chunk_scan_combined(
+    x: torch.Tensor,
+    dt: torch.Tensor,
+    A: torch.Tensor,
+    B: torch.Tensor,
+    C: torch.Tensor,
+    dl: torch.Tensor,
+    L: torch.Tensor,
+    chunk_size: int,
+    D: torch.Tensor | None = None,
+    z: torch.Tensor | None = None,
+    dt_bias: torch.Tensor | None = None,
+    initial_states: LogLinearAttentionState | None = None,
+    seq_idx: torch.Tensor | None = None,
+    cu_seqlens: torch.Tensor | None = None,
+    dt_softplus: bool = False,
+    dt_limit: tuple[float, float] = (0.0, float("inf")),
+    return_final_states: bool = False,
+):
+    if z is not None:
+        raise NotImplementedError
+    if seq_idx is not None:
+        raise NotImplementedError
+    if cu_seqlens is not None:
+        raise NotImplementedError
+    if dt_softplus is not True:
+        raise NotImplementedError
+    if tuple(dt_limit) != (0.0, float("inf")):
+        raise NotImplementedError
+    if chunk_size != 64:
+        raise NotImplementedError
+    if not B.shape == C.shape:
+        raise ValueError("B and C must have the same shape")
+
+    if D is not None:
+        if D.dim() != 1:
+            raise ValueError
+        D = rearrange(D, "h -> 1 1 h 1")
+        D_residual = x * D
+
+    if dt_bias is not None:
+        dt = dt + rearrange(dt_bias, "h -> 1 1 h")
+    if dt_softplus:
+        dt = torch.nn.functional.softplus(dt)
+    if dt_limit != (0.0, float("inf")):
+        dt = torch.clamp(dt, min=dt_limit[0], max=dt_limit[1])
+    x = (x * rearrange(dt, "b l h -> b l h 1")).to(x.dtype)
+    A = rearrange(A, "h -> 1 1 h") * dt
+
+    L = torch.nn.functional.softplus(rearrange(L, "h ell -> 1 1 h ell") * dl).to(L.dtype)
+
+    y, state = chunk_log_linear_attn(
+        q=C,
+        k=B,
+        v=x,
+        g=A,
+        level_scales=L,
+        initial_state=initial_states,
+        output_final_state=return_final_states,
+        cu_seqlens=cu_seqlens,
+    )
+
+    if D is not None:
+        y = y + D_residual
+
+    return y, state
+
+
+def hmamba_split_conv1d_scan_combined(
+    zxbcdtdl: torch.Tensor,
+    conv1d_weight: torch.Tensor,
+    conv1d_bias: torch.Tensor,
+    dt_bias: torch.Tensor,
+    A: torch.Tensor,
+    L: torch.Tensor,
+    D: torch.Tensor,
+    chunk_size: int,
+    initial_states: torch.Tensor | None = None,
+    seq_idx: torch.Tensor | None = None,
+    dt_limit: tuple[float, float] = (0.0, float("inf")),
+    return_final_states: bool = False,
+    activation: str = "silu",
+    rmsnorm_weight: torch.Tensor | None = None,
+    rmsnorm_eps: float = 1e-6,
+    outproj_weight: torch.Tensor | None = None,
+    outproj_bias: torch.Tensor | None = None,
+    headdim: int | None = None,
+    ngroups: int = 1,
+    norm_before_gate: bool = True,
+    conv1d_fn=None,
+    conv_backend: str = "cuda",
+) -> torch.Tensor:
+    """
+    Argument:
+        zxbcdtdl: (batch, seqlen, 2 * dim + 2 * ngroups * dstate + nheads) where dim == nheads * headdim
+        conv1d_weight: (dim + 2 * ngroups * dstate, width)
+        conv1d_bias: (dim + 2 * ngroups * dstate,)
+        dt_bias: (nheads,)
+        A: (nheads)
+        L: (nheads, nlevels)
+        D: (nheads, headdim) or (nheads,)
+        initial_states: (batch, nheads, headdim, dstate)
+        seq_idx: (batch, seqlen), int32
+        rmsnorm_weight: (dim,)
+        outproj_weight: (out_dim, dim)
+        outproj_bias: (out_dim,)
+        headdim: if D is 1D, headdim must be passed in
+        norm_before_gate: if True, we do RMSNorm(x) * F.silu(z). If False, we do RMSNorm(x * F.silu(z))
+    Return:
+        out: (batch, seqlen, dim)
+    """
+    if initial_states is not None:
+        raise NotImplementedError
+    if seq_idx is not None:
+        raise NotImplementedError
+    if dt_limit != (0.0, float("inf")):
+        raise NotImplementedError
+    if return_final_states is not False:
+        raise NotImplementedError
+    if norm_before_gate is not False:
+        raise NotImplementedError
+    if rmsnorm_weight is None:
+        raise NotImplementedError
+    if activation not in ["silu", "swish"]:
+        raise NotImplementedError
+
+    batch, seqlen, _ = zxbcdtdl.shape
+    dlambda = L.shape[-1]
+    (nheads,) = D.shape
+    dim = nheads * headdim
+    dstate = (zxbcdtdl.shape[-1] - 2 * dim - nheads - nheads * dlambda) // ngroups // 2
+
+    if D.dim() != 1:
+        raise ValueError
+    if headdim is None:
+        raise ValueError
+    if nheads % ngroups != 0:
+        raise ValueError
+    if zxbcdtdl.shape != (
+        batch,
+        seqlen,
+        2 * dim + 2 * ngroups * dstate + nheads + nheads * dlambda,
+    ):
+        raise ValueError
+    if dt_bias.shape != (nheads,):
+        raise ValueError
+    if A.shape != (nheads,):
+        raise ValueError
+    if L.shape != (nheads, dlambda):
+        raise ValueError
+    if D.shape != (nheads,):
+        raise ValueError
+    if rmsnorm_weight is None:
+        raise ValueError
+
+    zxBCdtl_splits = [dim, dim + 2 * ngroups * dstate, nheads, nheads * dlambda]
+    xBC_splits = [dim, ngroups * dstate, ngroups * dstate]
+    z, xBC, dt, dl = torch.split(zxbcdtdl, zxBCdtl_splits, dim=-1)
+    _conv_fn = conv1d_fn if conv1d_fn is not None else causal_conv1d_fn
+    _conv_out = _conv_fn(
+        rearrange(xBC, "b s d -> b d s"),
+        conv1d_weight,
+        bias=conv1d_bias,
+        activation=activation,
+        seq_idx=seq_idx,
+    )
+    if conv_backend == 'triton':
+        _conv_out = _conv_out[0]
+    xBC = rearrange(_conv_out, "b d s -> b s d")
+    x, B, C = torch.split(xBC, xBC_splits, dim=-1)
+    x = rearrange(x, "b l (h p) -> b l h p", h=nheads, p=headdim)
+    B = rearrange(B, "b l (g n) -> b l g n", g=ngroups, n=dstate)
+    C = rearrange(C, "b l (g n) -> b l g n", g=ngroups, n=dstate)
+    dl = rearrange(dl, "b l (h ell) -> b l h ell", h=nheads, ell=dlambda)
+    y, _ = hmamba_chunk_scan_combined(
+        x=x,
+        dt=dt,
+        A=A,
+        B=B,
+        C=C,
+        dl=dl,
+        L=L,
+        chunk_size=chunk_size,
+        D=D,
+        z=z if rmsnorm_weight is None else None,
+        dt_bias=dt_bias,
+        dt_softplus=True,
+        seq_idx=seq_idx,
+        cu_seqlens=None,
+        dt_limit=dt_limit,
+        return_final_states=return_final_states,
+    )
+
+    y = rearrange(y, "b l h p -> b l (h p)")
+    if rmsnorm_weight is not None:
+        y = rmsnorm_fn(
+            x=y,
+            weight=rmsnorm_weight,
+            bias=None,
+            z=z,
+            eps=rmsnorm_eps,
+            group_size=None,
+            norm_before_gate=False,
+        )
+    out = torch.nn.functional.linear(y, outproj_weight, outproj_bias)
+    return out
+
+
+class LogLinearMamba2(nn.Module):
+    """
+    Compute ∆, A, B, C, and D the state space parameters and compute the `contextualized_states`.
+    A, D are input independent (see Mamba paper [1] Section 3.5.2 "Interpretation of A" for why A isn't selective)
+    ∆, B, C are input-dependent (this is a key difference between Mamba and the linear time invariant S4,
+    and is why Mamba is called **selective** state spaces)
+    """
+
+    def __init__(
+        self,
+        num_heads: int,
+        head_dim: int = 64,
+        hidden_size: int = 2048,
+        state_size: int = 128,
+        expand: int = 2,
+        n_groups: int = 1,
+        conv_kernel: int = 4,
+        use_conv_bias: bool = False,
+        hidden_act: str = "silu",
+        rms_norm: bool = True,
+        chunk_size: int = 64,
+        time_step_rank: float = 256,
+        time_step_limit: tuple[float, float] = (0.0, float("inf")),
+        time_step_min: float = 0.001,
+        time_step_max: float = 0.1,
+        use_bias: bool = True,
+        norm_eps: float = 1e-5,
+        layer_idx: int = None,
+        backend: str = "cuda",
+    ):
+        super().__init__()
+        self.num_heads = num_heads
+        self.hidden_size = hidden_size
+        self.ssm_state_size = state_size
+        self.conv_kernel_size = conv_kernel
+        self.intermediate_size = int(expand * self.hidden_size)
+        self.time_step_rank = int(time_step_rank)
+        self.layer_idx = layer_idx
+        self.use_conv_bias = use_conv_bias
+        self.activation = hidden_act
+        self.act = ACT2FN[hidden_act]
+
+        self.layer_norm_epsilon = norm_eps
+        self.rms_norm = rms_norm
+
+        self.n_groups = n_groups
+        self.head_dim = head_dim
+        self.chunk_size = chunk_size
+
+        self.time_step_limit = time_step_limit
+        self.time_step_min = time_step_min
+        self.time_step_max = time_step_max
+
+        self.conv_dim = self.intermediate_size + 2 * self.n_groups * self.ssm_state_size
+        self.conv1d = nn.Conv1d(
+            in_channels=self.conv_dim,
+            out_channels=self.conv_dim,
+            bias=use_conv_bias,
+            kernel_size=conv_kernel,
+            groups=self.conv_dim,
+            padding=conv_kernel - 1,
+        )
+
+        self.num_lambda_dims = MAX_NUM_LEVELS
+        self.lambda_level_module = None
+
+        # projection of the input hidden states
+        projection_size = (
+            self.intermediate_size
+            + self.conv_dim
+            + self.num_heads * (self.num_lambda_dims + 1)
+        )
+        self.in_proj = nn.Linear(
+            self.hidden_size,
+            projection_size,
+            bias=use_bias,
+        )
+        # selective projection used to make dt, B and C input dependant
+
+        # time step projection (discretization)
+        # instantiate once and copy inv_dt in init_weights of PretrainedModel
+        self.dt_bias = nn.Parameter(torch.ones(self.num_heads))
+
+        # S4D real initialization. These are not discretized!
+        # The core is to load them, compute the discrete states, then write the updated state. Keeps the memory bounded
+        A = torch.arange(1, self.num_heads + 1)
+        self.A_log = nn.Parameter(torch.log(A))
+        self.A_log._no_weight_decay = True
+
+        self.lambda_mode = "positive"
+        L = torch.ones(self.num_heads, self.num_lambda_dims)
+        self.L = nn.Parameter(L)
+        self.L._no_weight_decay = True
+
+        self.norm = RMSNormGated(
+            self.intermediate_size, eps=self.layer_norm_epsilon, norm_before_gate=False,
+        )
+        self.D = nn.Parameter(torch.ones(self.num_heads))
+        self.D._no_weight_decay = True
+
+        self.out_proj = nn.Linear(
+            self.intermediate_size, self.hidden_size, bias=use_bias,
+        )
+        self.use_bias = use_bias
+
+        if not is_fast_path_available:
+            logger.warning_once(
+                "The fast path is not available because one of "
+                "`(selective_state_update, causal_conv1d_fn, causal_conv1d_update)` is None. "
+                "Falling back to the naive implementation. "
+                "To install follow https://github.com/state-spaces/mamba/#installation and"
+                "https://github.com/Dao-AILab/causal-conv1d",
+            )
+        import os
+        backend = os.environ.get('FLA_CONV_BACKEND', backend)
+        assert backend in ['cuda', 'triton'], f"Unsupported backend: {backend}"
+        if backend == 'cuda' and causal_conv1d_fn is None:
+            logger.warning_once(
+                "The CUDA backend is not available because `causal_conv1d` is None. "
+                "Falling back to the Triton backend. "
+                "To install follow https://github.com/Dao-AILab/causal-conv1d",
+            )
+            backend = 'triton'
+        if backend == 'triton':
+            from fla.modules.convolution import causal_conv1d as causal_conv1d_triton
+            from fla.modules.convolution import causal_conv1d_update as causal_conv1d_update_triton
+            self.causal_conv1d_fn = causal_conv1d_triton
+            self.causal_conv1d_update = causal_conv1d_update_triton
+            logger.warning(
+                "LogLinearMamba2 does not recommend using Triton's conv1d backend, "
+                "as it is untested and may contain bugs.",
+            )
+        else:
+            self.causal_conv1d_fn = causal_conv1d_fn
+            self.causal_conv1d_update = causal_conv1d_update
+        self.backend = backend
+
+    def cuda_kernels_forward(
+        self,
+        hidden_states: torch.Tensor,
+        last_state: dict | None = None,
+        use_cache: bool = False,
+        attention_mask: torch.Tensor | None = None,
+    ):
+        if self.activation not in ["silu", "swish"]:
+            raise ValueError
+
+        # 1. Gated MLP's linear projection
+        # Only apply padding mask during prefill (last_state is None).
+        # During decode, attention_mask has shape (B, accumulated_len) which
+        # mismatches hidden_states (B, 1, D).
+        hidden_states = apply_mask_to_padding_states(
+            hidden_states=hidden_states,
+            attention_mask=attention_mask if last_state is None else None,
+        )
+        projected_states = self.in_proj(hidden_states)
+
+        # Set up dimensions for reshapes later
+        batch_size, seq_len, _ = hidden_states.shape
+        groups_time_state_size = self.n_groups * self.ssm_state_size
+        d_mlp = (
+            projected_states.shape[-1]
+            - 2 * self.intermediate_size
+            - 2 * self.n_groups * self.ssm_state_size
+            - self.num_heads * (self.num_lambda_dims + 1)
+        ) // 2
+        if d_mlp != 0:
+            raise ValueError
+
+        # Single step calculations via cache
+        if last_state is not None:
+            if hidden_states.shape[1] != 1:
+                raise ValueError("LogLinearMamba2 cached decoding only supports a single new token per step.")
+
+            gate, xBC, dt, dl = torch.split(
+                projected_states.squeeze(1),
+                [
+                    self.intermediate_size,
+                    self.conv_dim,
+                    self.num_heads,
+                    self.num_heads * self.num_lambda_dims,
+                ],
+                dim=-1,
+            )
+
+            # 2. Convolution sequence transformation
+            conv_state = last_state['conv_state']
+            xBC = self.causal_conv1d_update(
+                xBC,
+                conv_state,
+                rearrange(self.conv1d.weight, "d 1 w -> d w"),
+                self.conv1d.bias,
+                self.activation,
+            )
+
+            x, B, C = torch.split(
+                xBC,
+                [
+                    self.intermediate_size,
+                    groups_time_state_size,
+                    groups_time_state_size,
+                ],
+                dim=-1,
+            )
+
+            # 3. SSM transformation
+            A = -torch.exp(self.A_log.float())  # (nheads,)
+            B = rearrange(
+                B,
+                "b (g n) -> b g n",
+                b=batch_size,
+                g=self.n_groups,
+                n=self.ssm_state_size,
+            )
+            C = rearrange(
+                C,
+                "b (g n) -> b g n",
+                b=batch_size,
+                g=self.n_groups,
+                n=self.ssm_state_size,
+            )
+            x_reshaped = rearrange(
+                x,
+                "b (h p) -> b h p",
+                b=batch_size,
+                h=self.num_heads,
+                p=self.head_dim,
+            )
+            dl_reshaped = rearrange(
+                dl,
+                "b (h ell) -> b h ell",
+                b=batch_size,
+                h=self.num_heads,
+                ell=self.num_lambda_dims,
+            )
+            y, hssm_state = hmamba_chunk_scan_combined(
+                x_reshaped,
+                dt=dt,
+                A=A,
+                B=B,
+                C=C,
+                dl=dl_reshaped,
+                L=self.L,
+                D=self.D,
+                z=None,
+                dt_bias=self.dt_bias,
+                dt_softplus=True,
+                initial_states=last_state['recurrent_state'],
+                return_final_states=True,
+            )
+            y = rearrange(
+                y,
+                "b h p -> b (h p)",
+                b=batch_size,
+                h=self.num_heads,
+                p=self.head_dim,
+            )
+            y = self.norm(y, gate)
+
+            # 4. Final linear projection
+            out = self.out_proj(y)[:, None, ...]
+            return out, conv_state, hssm_state
+
+        # Fused calculations or step by step if no initialized cache is found
+        else:
+            A = -torch.exp(
+                self.A_log.float(),
+            )  # (num_heads) or (intermediate_size, state_size)
+            dt_limit_kwargs = (
+                {}
+                if self.time_step_limit == (0.0, float("inf"))
+                else {"dt_limit": self.time_step_limit}
+            )
+
+            # 2-4. Fused kernel for conv1d, SSM, and the final projection
+            if self.training and not use_cache:
+                out = torch.utils.checkpoint.checkpoint(
+                    hmamba_split_conv1d_scan_combined,
+                    use_reentrant=False,
+                    # function arguments
+                    zxbcdtdl=projected_states,
+                    conv1d_weight=rearrange(self.conv1d.weight, "d 1 w -> d w"),
+                    conv1d_bias=self.conv1d.bias,
+                    dt_bias=self.dt_bias,
+                    A=A,
+                    L=self.L,
+                    D=self.D,
+                    chunk_size=self.chunk_size,
+                    conv1d_fn=self.causal_conv1d_fn,
+                    conv_backend=self.backend,
+                    seq_idx=None,  # was seq_idx
+                    activation=self.activation,
+                    rmsnorm_weight=self.norm.weight,
+                    rmsnorm_eps=self.norm.eps,
+                    outproj_weight=self.out_proj.weight,
+                    outproj_bias=self.out_proj.bias,
+                    headdim=self.head_dim,
+                    ngroups=self.n_groups,
+                    norm_before_gate=False,
+                    return_final_states=False,
+                    **dt_limit_kwargs,
+                )
+                return out, None, None
+
+            else:
+                gate, xBC, dt, dl = torch.split(
+                    projected_states,
+                    [
+                        self.intermediate_size,
+                        self.conv_dim,
+                        self.num_heads,
+                        self.num_heads * self.num_lambda_dims,
+                    ],
+                    dim=-1,
+                )
+
+                # 2. Convolution sequence transformation
+                # Init cache
+                masked_xBC = apply_mask_to_padding_states(xBC, attention_mask)
+                new_conv_state = None
+                if use_cache:
+                    xBC_t = rearrange(masked_xBC, "b l d -> b d l")
+                    new_conv_state = torch.nn.functional.pad(
+                        xBC_t,
+                        (self.conv_kernel_size - xBC_t.shape[-1], 0),
+                    )
+
+                _conv1d_output = self.causal_conv1d_fn(
+                    x=xBC.transpose(1, 2),
+                    weight=rearrange(self.conv1d.weight, "d 1 w -> d w"),
+                    bias=self.conv1d.bias,
+                    activation=self.activation,
+                )
+                if self.backend == 'cuda':
+                    xBC = _conv1d_output.transpose(1, 2)
+                elif self.backend == 'triton':
+                    xBC = _conv1d_output[0].transpose(1, 2).contiguous()
+                else:
+                    raise ValueError(f"Unsupported backend: {self.backend}")
+
+                xBC = apply_mask_to_padding_states(
+                    hidden_states=xBC,
+                    attention_mask=attention_mask,
+                )
+
+                x, B, C = torch.split(
+                    xBC,
+                    [
+                        self.intermediate_size,
+                        groups_time_state_size,
+                        groups_time_state_size,
+                    ],
+                    dim=-1,
+                )
+
+                # 3. SSM transformation
+                y, hssm_state = hmamba_chunk_scan_combined(
+                    rearrange(
+                        x,
+                        "b l (h p) -> b l h p",
+                        b=batch_size,
+                        l=seq_len,
+                        p=self.head_dim,
+                    ),
+                    dt=dt,
+                    A=A,
+                    B=rearrange(
+                        B,
+                        "b l (g n) -> b l g n",
+                        b=batch_size,
+                        l=seq_len,
+                        g=self.n_groups,
+                    ),
+                    C=rearrange(
+                        C,
+                        "b l (g n) -> b l g n",
+                        b=batch_size,
+                        l=seq_len,
+                        g=self.n_groups,
+                    ),
+                    dl=rearrange(
+                        dl,
+                        "b l (h ell) -> b l h ell",
+                        b=batch_size,
+                        h=self.num_heads,
+                        ell=self.num_lambda_dims,
+                    ),
+                    L=self.L,
+                    chunk_size=self.chunk_size,
+                    D=self.D,
+                    z=None,
+                    seq_idx=None,
+                    return_final_states=True,
+                    dt_bias=self.dt_bias,
+                    dt_softplus=True,
+                    **dt_limit_kwargs,
+                )
+
+                y = rearrange(
+                    y,
+                    "b l h p -> b l (h p)",
+                    b=batch_size,
+                    l=seq_len,
+                    h=self.num_heads,
+                    p=self.head_dim,
+                )
+                # Multiply "gate" branch and apply extra normalization layer
+                y = self.norm(y, gate)
+
+                # 4. Final linear projection
+                out = self.out_proj(y)
+
+                return out, new_conv_state, hssm_state
+
+    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,
+    ) -> tuple[torch.Tensor, torch.Tensor | None, Cache | None]:
+        last_state = get_layer_cache(self, past_key_values)
+
+        if "cuda" in self.in_proj.weight.device.type:
+            output, conv_state, hssm_state = self.cuda_kernels_forward(
+                hidden_states, last_state, use_cache, attention_mask
+            )
+        else:
+            raise NotImplementedError
+
+        update_layer_cache(
+            self,
+            past_key_values,
+            recurrent_state=hssm_state,
+            conv_state=conv_state,
+            offset=hidden_states.shape[1],
+        )
+
+        return output, None, past_key_values

fla/layers/mamba.py

@@ -0,0 +1,395 @@
+# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang
+
+from __future__ import annotations
+
+import warnings
+from typing import TYPE_CHECKING
+
+import torch
+import torch.nn as nn
+from transformers.utils import logging
+
+from fla.layers.utils import get_layer_cache, update_layer_cache
+from fla.modules.activations import ACT2FN
+
+with warnings.catch_warnings():
+    warnings.simplefilter('ignore')
+    try:
+        from mamba_ssm.ops.selective_scan_interface import mamba_inner_fn, selective_scan_fn
+        from mamba_ssm.ops.triton.selective_state_update import selective_state_update
+    except ImportError:
+        selective_state_update, selective_scan_fn, mamba_inner_fn = None, None, None
+
+    try:
+        from causal_conv1d import causal_conv1d_fn, causal_conv1d_update
+    except ImportError:
+        causal_conv1d_update, causal_conv1d_fn = None, None
+    is_fast_path_available = all((
+        selective_state_update,
+        selective_scan_fn,
+        mamba_inner_fn,
+    ))
+if TYPE_CHECKING:
+    from transformers.processing_utils import Unpack
+
+    from fla.models.utils import Cache
+
+logger = logging.get_logger(__name__)
+
+
+class Mamba(nn.Module):
+    """
+    Compute ∆, A, B, C, and D the state space parameters and compute the `contextualized_states`.
+    A, D are input independent (see Mamba paper [1] Section 3.5.2 "Interpretation of A" for why A isn't selective)
+    ∆, B, C are input-dependent (this is a key difference between Mamba and the linear time invariant S4,
+    and is why Mamba is called **selective** state spaces)
+    """
+
+    def __init__(
+        self,
+        hidden_size: int = 2048,
+        state_size: int = 16,
+        conv_kernel: int = 4,
+        use_conv_bias: bool = True,
+        intermediate_size: int = 2048,
+        time_step_rank: int = 256,
+        use_bias: bool = True,
+        hidden_act: str = "silu",
+        layer_idx: int = None,
+        backend: str = "cuda",
+    ):
+        super().__init__()
+
+        self.hidden_size = hidden_size
+        self.ssm_state_size = state_size
+        self.conv_kernel_size = conv_kernel
+        self.use_conv_bias = use_conv_bias
+        self.intermediate_size = intermediate_size
+        self.time_step_rank = time_step_rank
+        self.use_bias = use_bias
+
+        self.conv1d = nn.Conv1d(
+            in_channels=self.intermediate_size,
+            out_channels=self.intermediate_size,
+            bias=use_conv_bias,
+            kernel_size=conv_kernel,
+            groups=self.intermediate_size,
+            padding=conv_kernel - 1,
+        )
+
+        self.activation = hidden_act
+        self.act = ACT2FN[hidden_act]
+
+        self.layer_idx = layer_idx
+
+        # projection of the input hidden states
+        self.in_proj = nn.Linear(self.hidden_size, self.intermediate_size * 2, bias=use_bias)
+        # selective projection used to make dt, B and C input dependant
+        self.x_proj = nn.Linear(self.intermediate_size, self.time_step_rank + self.ssm_state_size * 2, bias=False)
+        # time step projection (discretization)
+        self.dt_proj = nn.Linear(self.time_step_rank, self.intermediate_size, bias=True)
+
+        # S4D real initialization. These are not discretized!
+        # The core is to load them, compute the discrete states, then write the updated state. Keeps the memory bounded
+        A = torch.arange(1, self.ssm_state_size + 1, dtype=torch.float32)[None, :]
+        A = A.expand(self.intermediate_size, -1).contiguous()
+
+        self.A_log = nn.Parameter(torch.log(A))
+        self.D = nn.Parameter(torch.ones(self.intermediate_size))
+        self.out_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=use_bias)
+
+        if not is_fast_path_available:
+            logger.warning_once(
+                "The fast path is not available because on of "
+                "`(selective_state_update, selective_scan_fn, causal_conv1d_fn, causal_conv1d_update, mamba_inner_fn)`"
+                " is None. Falling back to the naive implementation. "
+                "To install follow https://github.com/state-spaces/mamba/#installation and"
+                " https://github.com/Dao-AILab/causal-conv1d",
+            )
+        import os
+        backend = os.environ.get('FLA_CONV_BACKEND', backend)
+        assert backend in ['cuda', 'triton'], f"Unsupported backend: {backend}"
+        if backend == 'cuda' and causal_conv1d_fn is None:
+            logger.warning_once(
+                "The CUDA backend is not available because `causal_conv1d` is None. "
+                "Falling back to the Triton backend. "
+                "To install follow https://github.com/Dao-AILab/causal-conv1d",
+            )
+            backend = 'triton'
+        if backend == 'triton':
+            from fla.modules.convolution import causal_conv1d as causal_conv1d_triton
+            from fla.modules.convolution import causal_conv1d_update as causal_conv1d_update_triton
+            self.causal_conv1d_fn = causal_conv1d_triton
+            self.causal_conv1d_update = causal_conv1d_update_triton
+        else:
+            self.causal_conv1d_fn = causal_conv1d_fn
+            self.causal_conv1d_update = causal_conv1d_update
+        self.backend = backend
+
+    def _to_causal_conv_layout(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        return hidden_states.transpose(1, 2).contiguous()
+
+    def _from_causal_conv_layout(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        return hidden_states.transpose(1, 2).contiguous()
+
+    def _build_conv_state(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        seq_len = hidden_states.shape[-1]
+        if seq_len >= self.conv_kernel_size:
+            return hidden_states[..., -self.conv_kernel_size:].contiguous()
+        return nn.functional.pad(hidden_states, (self.conv_kernel_size - seq_len, 0)).contiguous()
+
+    def cuda_kernels_forward(
+        self,
+        hidden_states: torch.Tensor,
+        last_state: dict | None = None,
+        use_cache: bool | None = False,
+        attention_mask: torch.LongTensor | None = None,
+        **kwargs: Unpack[dict],
+    ):
+        if last_state is not None and hidden_states.shape[1] != 1:
+            raise ValueError("Mamba cached decoding only supports a single new token per step.")
+
+        # 1. Gated MLP's linear projection
+        projected_states = self.in_proj(hidden_states).transpose(1, 2)
+
+        if self.training and not use_cache:
+            contextualized_states = mamba_inner_fn(
+                projected_states,
+                self.conv1d.weight,
+                self.conv1d.bias if self.use_conv_bias else None,
+                self.x_proj.weight,
+                self.dt_proj.weight,
+                self.out_proj.weight,
+                self.out_proj.bias.float() if self.use_bias else None,
+                -torch.exp(self.A_log.float()),
+                None,  # input-dependent B
+                None,  # input-dependent C
+                self.D.float(),
+                delta_bias=self.dt_proj.bias.float(),
+                delta_softplus=True,
+            )
+            return contextualized_states, None, None
+
+        hidden_states, gate = projected_states.chunk(2, dim=1)
+
+        if attention_mask is not None and last_state is None:
+            # Mask before the depthwise conv so cached/prefill conv inputs do not keep pad tokens.
+            hidden_states = hidden_states * attention_mask.unsqueeze(1)
+
+        # 2. Convolution sequence transformation
+        conv_inputs = hidden_states
+        conv_weights = self.conv1d.weight.view(self.conv1d.weight.size(0), self.conv1d.weight.size(2))
+        if last_state is not None:
+            conv_state = last_state['conv_state']
+            ssm_state = last_state['recurrent_state']
+
+            if self.backend == 'triton':
+                hidden_states, conv_state = self.causal_conv1d_update(
+                    x=self._to_causal_conv_layout(conv_inputs),
+                    cache=conv_state,
+                    weight=conv_weights,
+                    bias=self.conv1d.bias,
+                    activation=self.activation,
+                )
+                hidden_states = self._from_causal_conv_layout(hidden_states)
+            else:
+                hidden_states = self.causal_conv1d_update(
+                    conv_inputs.squeeze(-1),
+                    conv_state,
+                    conv_weights,
+                    self.conv1d.bias,
+                    self.activation,
+                )
+                hidden_states = hidden_states.unsqueeze(-1)
+        else:
+            conv_state = None
+            ssm_state = None
+            if self.backend == 'triton':
+                hidden_states, conv_state = self.causal_conv1d_fn(
+                    x=self._to_causal_conv_layout(conv_inputs),
+                    weight=conv_weights,
+                    bias=self.conv1d.bias,
+                    activation=self.activation,
+                    output_final_state=bool(use_cache),
+                )
+                hidden_states = self._from_causal_conv_layout(hidden_states)
+            else:
+                if use_cache:
+                    conv_state = self._build_conv_state(conv_inputs)
+                hidden_states = self.causal_conv1d_fn(
+                    conv_inputs, conv_weights, self.conv1d.bias, activation=self.activation,
+                )
+
+        if attention_mask is not None and last_state is None:
+            # Re-mask after the conv: causal kernels can regenerate non-zero values at masked positions,
+            # and those values would otherwise leak into x_proj and the SSM recurrence.
+            hidden_states = hidden_states * attention_mask.unsqueeze(1)
+
+        # 3. State Space Model sequence transformation
+        # 3.a. input varying initialization of time_step, B and C
+        ssm_parameters = self.x_proj(hidden_states.transpose(1, 2))
+        time_step, B, C = torch.split(
+            ssm_parameters, [self.time_step_rank, self.ssm_state_size, self.ssm_state_size], dim=-1,
+        )
+        discrete_time_step = self.dt_proj.weight @ time_step.transpose(1, 2)
+
+        A = -torch.exp(self.A_log.float())
+        # 3.c perform the recurrence y ← SSM(A, B, C)(x)
+        time_proj_bias = self.dt_proj.bias.float() if hasattr(self.dt_proj, "bias") else None
+        if last_state is not None:
+            scan_outputs = selective_state_update(
+                ssm_state,
+                hidden_states[..., 0],
+                discrete_time_step[..., 0],
+                A,
+                B[:, 0],
+                C[:, 0],
+                self.D,
+                gate[..., 0],
+                time_proj_bias,
+                dt_softplus=True,
+            ).unsqueeze(-1)
+        else:
+            scan_outputs, ssm_state = selective_scan_fn(
+                hidden_states,
+                discrete_time_step,
+                A,
+                B.transpose(1, 2),
+                C.transpose(1, 2),
+                self.D.float(),
+                gate,
+                time_proj_bias,
+                delta_softplus=True,
+                return_last_state=True,
+            )
+
+        # 4. Final linear projection
+        contextualized_states = self.out_proj(scan_outputs.transpose(1, 2))
+        return contextualized_states, conv_state, ssm_state
+
+    def slow_forward(
+        self,
+        input_states,
+        last_state: dict | None = None,
+        use_cache: bool | None = False,
+        attention_mask: torch.LongTensor | None = None,
+        **kwargs: Unpack[dict],
+    ):
+        if last_state is not None and input_states.shape[1] != 1:
+            raise ValueError("Mamba cached decoding only supports a single new token per step.")
+
+        batch_size, seq_len, _ = input_states.shape
+        dtype = input_states.dtype
+        # 1. Gated MLP's linear projection
+        # [batch, 2 * intermediate_size, seq_len]
+        projected_states = self.in_proj(input_states).transpose(1, 2)
+        hidden_states, gate = projected_states.chunk(2, dim=1)
+
+        if attention_mask is not None and last_state is None:
+            # Mask before the depthwise conv so cached/prefill conv inputs do not keep pad tokens.
+            hidden_states = hidden_states * attention_mask.unsqueeze(1)
+
+        # 2. Convolution sequence transformation
+        if last_state is not None:
+            conv_state = last_state['conv_state']
+            ssm_state = last_state['recurrent_state'].clone().to(hidden_states.device)
+
+            # decode path: single token
+            conv_state = conv_state.roll(shifts=-1, dims=-1)
+            conv_state[:, :, -1] = hidden_states[:, :, 0].to(conv_state.device)
+            hidden_states = torch.sum(conv_state * self.conv1d.weight[:, 0, :], dim=-1)
+            if self.use_conv_bias:
+                hidden_states += self.conv1d.bias
+            # [batch, intermediate_size, 1] : decoding
+            hidden_states = self.act(hidden_states).to(dtype).unsqueeze(-1)
+        elif use_cache:
+            ssm_state = torch.zeros(
+                (batch_size, self.intermediate_size, self.ssm_state_size),
+                device=hidden_states.device, dtype=dtype,
+            )
+            conv_state = self._build_conv_state(hidden_states)
+            # [batch, intermediate_size, seq_len]
+            hidden_states = self.act(self.conv1d(hidden_states)[..., :seq_len])
+        else:
+            ssm_state = torch.zeros(
+                (batch_size, self.intermediate_size, self.ssm_state_size),
+                device=hidden_states.device, dtype=dtype,
+            )
+            conv_state = None
+            # [batch, intermediate_size, seq_len]
+            hidden_states = self.act(self.conv1d(hidden_states)[..., :seq_len])
+
+        if attention_mask is not None and last_state is None:
+            # Re-mask after the conv: causal kernels can regenerate non-zero values at masked positions,
+            # and those values would otherwise leak into x_proj and the SSM recurrence.
+            hidden_states = hidden_states * attention_mask.unsqueeze(1)
+
+        # 3. State Space Model sequence transformation
+        # 3.a. Selection:  [batch, seq_len, self.time_step_rank + self.ssm_state_size * 2]
+        ssm_parameters = self.x_proj(hidden_states.transpose(1, 2))
+        time_step, B, C = torch.split(
+            ssm_parameters, [self.time_step_rank, self.ssm_state_size, self.ssm_state_size], dim=-1,
+        )
+        # [batch, seq_len, intermediate_size]
+        discrete_time_step = self.dt_proj(time_step)
+        # [batch, intermediate_size, seq_len]
+        discrete_time_step = nn.functional.softplus(discrete_time_step).transpose(1, 2)
+
+        # 3.b. Discretization: B and C to [batch, seq_len, intermediate_size, ssm_state_size] (SRAM)
+        # [intermediate_size, ssm_state_size]
+        A = -torch.exp(self.A_log.float())
+        # [batch, intermediate_size, seq_len, ssm_state_size]
+        discrete_A = torch.exp(A[None, :, None, :] * discrete_time_step[:, :, :, None])
+        # [batch, intermediate_size, seq_len, ssm_state_size]
+        discrete_B = discrete_time_step[:, :, :, None] * B[:, None, :, :].float()
+        deltaB_u = discrete_B * hidden_states[:, :, :, None].float()
+
+        # 3.c perform the recurrence y ← SSM(A, B, C)(x)
+        scan_outputs = []
+        for i in range(hidden_states.shape[-1]):
+            # [batch, intermediade_size, ssm_state]
+            ssm_state = discrete_A[:, :, i, :] * ssm_state + deltaB_u[:, :, i, :]
+            # [batch, intermediade_size, 1]
+            scan_output = torch.matmul(ssm_state.to(dtype), C[:, i, :].unsqueeze(-1))
+            scan_outputs.append(scan_output[:, :, 0])
+        # [batch, seq_len, intermediade_size]
+        scan_output = torch.stack(scan_outputs, dim=-1)
+        scan_output = scan_output + (hidden_states * self.D[None, :, None])
+        scan_output = (scan_output * self.act(gate))
+
+        # 4. Final linear projection
+        # [batch, seq_len, hidden_size]
+        contextualized_states = self.out_proj(scan_output.transpose(1, 2))
+        return contextualized_states, conv_state, ssm_state
+    # fmt: on
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: torch.LongTensor | 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]:
+        last_state = get_layer_cache(self, past_key_values)
+
+        if is_fast_path_available and "cuda" in self.x_proj.weight.device.type:
+            output, conv_state, ssm_state = self.cuda_kernels_forward(
+                hidden_states, last_state, use_cache, attention_mask, **kwargs
+            )
+        else:
+            output, conv_state, ssm_state = self.slow_forward(
+                hidden_states, last_state, use_cache, attention_mask, **kwargs
+            )
+
+        if use_cache and past_key_values is not None:
+            update_layer_cache(
+                self,
+                past_key_values,
+                recurrent_state=ssm_state,
+                conv_state=conv_state,
+                offset=hidden_states.shape[1],
+            )
+
+        return output, None, past_key_values

fla/layers/mamba2.py

@@ -0,0 +1,649 @@
+# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang
+
+from __future__ import annotations
+
+import math
+import warnings
+from typing import TYPE_CHECKING
+
+import torch
+import torch.nn as nn
+from transformers.utils import logging
+
+from fla.layers.utils import get_layer_cache, update_layer_cache
+from fla.modules.activations import ACT2FN
+from fla.modules.layernorm_gated import RMSNormGated
+
+with warnings.catch_warnings():
+    warnings.simplefilter('ignore')
+    try:
+        from mamba_ssm.ops.triton.selective_state_update import selective_state_update
+        from mamba_ssm.ops.triton.ssd_combined import mamba_chunk_scan_combined, mamba_split_conv1d_scan_combined
+    except ImportError:
+        selective_state_update, mamba_chunk_scan_combined, mamba_split_conv1d_scan_combined = None, None, None
+    try:
+        from causal_conv1d import causal_conv1d_fn, causal_conv1d_update
+    except ImportError:
+        causal_conv1d_update, causal_conv1d_fn = None, None
+    is_fast_path_available = selective_state_update is not None
+
+if TYPE_CHECKING:
+    from fla.models.utils import Cache
+
+logger = logging.get_logger(__name__)
+
+
+def apply_mask_to_padding_states(hidden_states, attention_mask):
+    """
+    Tunes out the hidden states for padding tokens, see https://github.com/state-spaces/mamba/issues/66
+    """
+    if attention_mask is not None and attention_mask.shape[1] > 1 and attention_mask.shape[0] > 1:
+        dtype = hidden_states.dtype
+        hidden_states = (hidden_states * attention_mask[:, :, None]).to(dtype)
+
+    return hidden_states
+
+
+def pad_tensor_by_size(input_tensor: torch.Tensor, pad_size: int):
+    """
+    Padding x tensor with `pad_size` on the seq_len dim (dim=1)
+
+    Assumes that we only have tensors of either size 4 or 3
+    """
+    pad_shape = (0, 0, 0, 0, 0, pad_size, 0, 0) if len(input_tensor.shape) == 4 else (0, 0, 0, pad_size, 0, 0)
+
+    return torch.nn.functional.pad(input_tensor, pad_shape, mode="constant", value=0)
+
+
+def reshape_into_chunks(input_tensor, pad_size, chunk_size):
+    """
+    Padding input_tensor with `pad_size` on the seq_len dim (dim=1) and
+    simultaneously splitting it into chunk sequences.
+
+    Assumes that we only have tensors of either size 4 or 3
+    """
+    # [bsz, seq_len, ...] -> [bsz, seq_len multiple of chunk_size, ...]
+    input_tensor = pad_tensor_by_size(input_tensor, pad_size)
+
+    if len(input_tensor.shape) == 3:
+        # [bsz, seq_len multiple of chunk_size, num_heads] -> [bsz, -1, chunk_size, num_heads]
+        return input_tensor.reshape(input_tensor.shape[0], -1, chunk_size, input_tensor.shape[2])
+    else:
+        # [bsz, seq_len multiple of chunk_size, num_heads, head_dim or state_size] ->
+        # [bsz, -1, chunk_size, num_heads, head_dim or state_size]
+        return input_tensor.reshape(
+            input_tensor.shape[0], -1, chunk_size, input_tensor.shape[2], input_tensor.shape[3],
+        )
+
+
+def segment_sum(input_tensor):
+    """
+    More stable segment sum calculation. Uses cumulative sums and masking instead of direct subtractions.
+    """
+    chunk_size = input_tensor.size(-1)
+    # 1. expand input tensor to have an additional dimension and repeat along that dimension
+    # [..., chunk_size] -> [..., chunk_size, chunk_size]
+    input_tensor = input_tensor[..., None].expand(*input_tensor.size(), chunk_size)
+    # 2. create a lower triangular mask with the diagonal set to 0 to 0 out elements above diag
+    mask = torch.tril(torch.ones(chunk_size, chunk_size, device=input_tensor.device, dtype=torch.bool), diagonal=-1)
+    input_tensor = input_tensor.masked_fill(~mask, 0)
+    # 3. compute actual cumsum
+    tensor_segsum = torch.cumsum(input_tensor, dim=-2)
+
+    # 4. apply mask to keep only the lower triangular part of the cumulative sum result (incl diagonal this time)
+    mask = torch.tril(torch.ones(chunk_size, chunk_size, device=input_tensor.device, dtype=torch.bool), diagonal=0)
+    tensor_segsum = tensor_segsum.masked_fill(~mask, -torch.inf)
+    return tensor_segsum
+
+
+class Mamba2(nn.Module):
+    """
+    Compute ∆, A, B, C, and D the state space parameters and compute the `contextualized_states`.
+    A, D are input independent (see Mamba paper [1] Section 3.5.2 "Interpretation of A" for why A isn't selective)
+    ∆, B, C are input-dependent (this is a key difference between Mamba and the linear time invariant S4,
+    and is why Mamba is called **selective** state spaces)
+    """
+
+    def __init__(
+        self,
+        num_heads: int,
+        head_dim: int = 64,
+        hidden_size: int = 2048,
+        state_size: int = 128,
+        expand: int = 2,
+        n_groups: int = 1,
+        conv_kernel: int = 4,
+        use_conv_bias: bool = False,
+        hidden_act: str = "silu",
+        rms_norm: bool = True,
+        chunk_size: int = 256,
+        time_step_rank: float = 256,
+        time_step_limit: tuple[float, float] = (0.0, float("inf")),
+        time_step_min: float = 0.001,
+        time_step_max: float = 0.1,
+        use_bias: bool = True,
+        norm_eps: float = 1e-5,
+        layer_idx: int = None,
+        backend: str = "cuda",
+    ) -> Mamba2:
+        super().__init__()
+
+        self.num_heads = num_heads
+        self.head_dim = head_dim
+        self.hidden_size = hidden_size
+        self.ssm_state_size = state_size
+        self.expand = expand
+        self.intermediate_size = int(expand * hidden_size)
+        self.n_groups = n_groups
+
+        self.conv_kernel_size = conv_kernel
+        self.use_conv_bias = use_conv_bias
+        self.activation = hidden_act
+        self.act = ACT2FN[hidden_act]
+
+        self.rms_norm = rms_norm
+        self.norm_eps = norm_eps
+
+        self.chunk_size = chunk_size
+
+        self.time_step_rank = int(time_step_rank)
+        self.time_step_limit = time_step_limit
+        self.time_step_min = time_step_min
+        self.time_step_max = time_step_max
+
+        self.conv_dim = self.intermediate_size + 2 * self.n_groups * self.ssm_state_size
+        self.conv1d = nn.Conv1d(
+            in_channels=self.conv_dim,
+            out_channels=self.conv_dim,
+            bias=use_conv_bias,
+            kernel_size=conv_kernel,
+            groups=self.conv_dim,
+            padding=conv_kernel - 1,
+        )
+
+        # projection of the input hidden states
+        projection_size = self.intermediate_size + self.conv_dim + self.num_heads
+        self.in_proj = nn.Linear(
+            self.hidden_size,
+            projection_size,
+            bias=use_bias,
+        )
+        # selective projection used to make dt, B and C input dependant
+
+        # time step projection (discretization)
+        # instantiate once and copy inv_dt in init_weights of PretrainedModel
+        # hard coded for now
+        dt_init_floor = 1e-4
+        dt = torch.exp(
+            torch.rand(self.num_heads) * (
+                math.log(self.time_step_max) - math.log(self.time_step_min)
+            ) + math.log(self.time_step_min)
+        )
+        dt = torch.clamp(dt, min=dt_init_floor)
+        # Inverse of softplus: https://github.com/pytorch/pytorch/issues/72759
+        inv_dt = dt + torch.log(-torch.expm1(-dt))
+        self.dt_bias = nn.Parameter(inv_dt)
+
+        # S4D real initialization. These are not discretized!
+        # The core is to load them, compute the discrete states, then write the updated state. Keeps the memory bounded
+        A = torch.empty(self.num_heads, dtype=torch.float32).uniform_(0, 16)
+        self.A_log = nn.Parameter(torch.log(A))
+        self.A_log._no_weight_decay = True
+        self.norm = RMSNormGated(
+            self.intermediate_size, eps=self.norm_eps, norm_before_gate=False,
+        )
+        self.D = nn.Parameter(torch.ones(self.num_heads))
+        self.D._no_weight_decay = True
+
+        self.out_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=use_bias)
+        self.use_bias = use_bias
+
+        self.layer_idx = layer_idx
+
+        if not is_fast_path_available:
+            logger.warning_once(
+                "The fast path is not available because one of "
+                "`(selective_state_update)` is None. "
+                "Falling back to the naive implementation. "
+                "To install follow https://github.com/state-spaces/mamba/#installation",
+            )
+        import os
+        backend = os.environ.get('FLA_CONV_BACKEND', backend)
+        assert backend in ['cuda', 'triton'], f"Unsupported backend: {backend}"
+        if backend == 'cuda' and causal_conv1d_fn is None:
+            logger.warning_once(
+                "The CUDA backend is not available because `causal_conv1d` is None. "
+                "Falling back to the Triton backend. "
+                "To install follow https://github.com/Dao-AILab/causal-conv1d",
+            )
+            backend = 'triton'
+        if backend == 'triton':
+            from fla.modules.convolution import causal_conv1d as causal_conv1d_triton
+            from fla.modules.convolution import causal_conv1d_update as causal_conv1d_update_triton
+            self.causal_conv1d_fn = causal_conv1d_triton
+            self.causal_conv1d_update = causal_conv1d_update_triton
+            logger.warning(
+                "Mamba2 does not recommend using Triton's conv1d backend, "
+                "as it is untested and may contain bugs.",
+            )
+        else:
+            self.causal_conv1d_fn = causal_conv1d_fn
+            self.causal_conv1d_update = causal_conv1d_update
+        self.backend = backend
+
+    def cuda_kernels_forward(
+        self,
+        hidden_states: torch.Tensor,
+        last_state: dict | None = None,
+        use_cache: bool = False,
+        attention_mask: torch.Tensor | None = None,
+    ):
+        # 1. Gated MLP's linear projection
+        projected_states = self.in_proj(hidden_states)
+
+        # Set up dimensions for reshapes later
+        batch_size, seq_len, _ = hidden_states.shape
+        groups_time_state_size = self.n_groups * self.ssm_state_size
+        d_mlp = (
+            projected_states.shape[-1]
+            - 2 * self.intermediate_size
+            - 2 * self.n_groups * self.ssm_state_size
+            - self.num_heads
+        ) // 2
+
+        # Single step calculations via cache (decode)
+        if last_state is not None:
+            if hidden_states.shape[1] != 1:
+                raise ValueError("Mamba2 cached decoding only supports a single new token per step.")
+            conv_state = last_state['conv_state']
+            ssm_state = last_state['recurrent_state']
+
+            _, _, gate, hidden_states_B_C, dt = projected_states.squeeze(1).split(
+                [d_mlp, d_mlp, self.intermediate_size, self.conv_dim, self.num_heads], dim=-1,
+            )
+
+            # 2. Convolution sequence transformation
+            hidden_states_B_C = self.causal_conv1d_update(
+                hidden_states_B_C.contiguous(),
+                conv_state,
+                self.conv1d.weight.squeeze(1),
+                self.conv1d.bias,
+                self.activation,
+            )
+
+            hidden_states, B, C = torch.split(
+                hidden_states_B_C,
+                [
+                    self.intermediate_size,
+                    groups_time_state_size,
+                    groups_time_state_size,
+                ],
+                dim=-1,
+            )
+
+            # 3. SSM transformation
+            A = -torch.exp(self.A_log.float())  # (nheads,)
+            A = A[:, None, ...][:, :, None].expand(-1, self.head_dim, self.ssm_state_size).to(dtype=torch.float32)
+            dt = dt[:, :, None].expand(-1, -1, self.head_dim)
+            dt_bias = self.dt_bias[:, None, ...].expand(-1, self.head_dim)
+            D = self.D[:, None, ...].expand(-1, self.head_dim)
+            B = B.view(batch_size, self.n_groups, B.shape[1] // self.n_groups)
+            C = C.view(batch_size, self.n_groups, C.shape[1] // self.n_groups)
+            hidden_states_reshaped = hidden_states.view(batch_size, self.num_heads, self.head_dim)
+
+            hidden_states = selective_state_update(
+                ssm_state,
+                hidden_states_reshaped,
+                dt,
+                A,
+                B,
+                C,
+                D,
+                z=None,
+                dt_bias=dt_bias,
+                dt_softplus=True,
+            )
+            hidden_states = hidden_states.view(batch_size, self.num_heads * self.head_dim)
+            hidden_states = self.norm(hidden_states, gate)
+
+            # 4. Final linear projection
+            out = self.out_proj(hidden_states)[:, None, ...]
+
+            # conv_state is updated in-place by causal_conv1d_update
+            # ssm_state is updated in-place by selective_state_update
+            return out, conv_state, ssm_state
+
+        # Fused calculations or step by step if no initialized cache is found (prefill)
+        else:
+            A = -torch.exp(self.A_log.float())  # (num_heads) or (intermediate_size, state_size)
+            dt_limit_kwargs = {} if self.time_step_limit == (0.0, float("inf")) else {"dt_limit": self.time_step_limit}
+
+            # 2-4. Fused kernel for conv1d, SSM, and the final projection
+            if self.training and not use_cache:
+                out = mamba_split_conv1d_scan_combined(
+                    projected_states,
+                    self.conv1d.weight.squeeze(1),
+                    self.conv1d.bias,
+                    self.dt_bias,
+                    A,
+                    D=self.D,
+                    chunk_size=self.chunk_size,
+                    seq_idx=None,  # was seq_idx
+                    activation=self.activation,
+                    rmsnorm_weight=self.norm.weight,
+                    rmsnorm_eps=self.norm.eps,
+                    outproj_weight=self.out_proj.weight,
+                    outproj_bias=self.out_proj.bias,
+                    headdim=self.head_dim,
+                    ngroups=self.n_groups,
+                    norm_before_gate=False,
+                    return_final_states=False,
+                    **dt_limit_kwargs,
+                )
+                return out, None, None
+
+            else:
+                _, _, gate, hidden_states_B_C, dt = projected_states.split(
+                    [d_mlp, d_mlp, self.intermediate_size, self.conv_dim, self.num_heads], dim=-1,
+                )
+
+                # 2. Convolution sequence transformation
+                hidden_states_B_C = apply_mask_to_padding_states(hidden_states_B_C, attention_mask)
+                # Compute conv_state for cache
+                new_conv_state = None
+                if use_cache:
+                    hidden_states_B_C_transposed = hidden_states_B_C.transpose(1, 2)
+                    new_conv_state = nn.functional.pad(
+                        hidden_states_B_C_transposed,
+                        (self.conv_kernel_size - hidden_states_B_C_transposed.shape[-1], 0),
+                    )
+
+                if self.activation not in ["silu", "swish"]:
+                    hidden_states_B_C = self.act(
+                        self.conv1d(hidden_states_B_C.transpose(1, 2))[..., :seq_len].transpose(1, 2),
+                    )
+                else:
+                    _conv1d_output = self.causal_conv1d_fn(
+                        x=hidden_states_B_C.transpose(1, 2).contiguous(),
+                        weight=self.conv1d.weight.squeeze(1),
+                        bias=self.conv1d.bias,
+                        activation=self.activation,
+                    )
+                    if self.backend == 'cuda':
+                        hidden_states_B_C = _conv1d_output
+                        hidden_states_B_C = hidden_states_B_C.transpose(1, 2)
+                    elif self.backend == 'triton':
+                        hidden_states_B_C, _ = _conv1d_output
+                        hidden_states_B_C = hidden_states_B_C.transpose(1, 2).contiguous()
+                    else:
+                        raise ValueError(f"Unsupported backend: {self.backend}")
+
+                hidden_states_B_C = (hidden_states_B_C * attention_mask[:, :, None]).to(hidden_states_B_C.dtype) \
+                    if attention_mask is not None and attention_mask.shape[1] > 1 and attention_mask.shape[0] > 1 \
+                    else hidden_states_B_C
+                hidden_states, B, C = torch.split(
+                    hidden_states_B_C,
+                    [self.intermediate_size, groups_time_state_size, groups_time_state_size],
+                    dim=-1,
+                )
+
+                # 3. SSM transformation
+                scan_output, ssm_state = mamba_chunk_scan_combined(
+                    hidden_states.view(batch_size, seq_len, -1, self.head_dim),
+                    dt,
+                    A,
+                    B.view(batch_size, seq_len, self.n_groups, -1),
+                    C.view(batch_size, seq_len, self.n_groups, -1),
+                    chunk_size=self.chunk_size,
+                    D=self.D,
+                    z=None,
+                    seq_idx=None,
+                    return_final_states=True,
+                    dt_bias=self.dt_bias,
+                    dt_softplus=True,
+                    **dt_limit_kwargs,
+                )
+
+                scan_output = scan_output.view(batch_size, seq_len, -1)
+                # Multiply "gate" branch and apply extra normalization layer
+                scan_output = self.norm(scan_output, gate)
+
+                # 4. Final linear projection
+                out = self.out_proj(scan_output)
+
+                return out, new_conv_state, ssm_state
+
+    # fmt: off
+    def torch_forward(
+        self,
+        input_states,
+        last_state: dict | None = None,
+        use_cache: bool = False,
+        attention_mask: torch.Tensor | None = None,
+    ):
+        batch_size, seq_len, _ = input_states.shape
+        dtype = input_states.dtype
+
+        # 1. Gated MLP's linear projection
+        projected_states = self.in_proj(input_states)
+        d_mlp = (projected_states.shape[-1] - 2 * self.intermediate_size -
+                 2 * self.n_groups * self.ssm_state_size - self.num_heads) // 2
+        _, _, gate, hidden_states_B_C, dt = projected_states.split(
+            [d_mlp, d_mlp, self.intermediate_size,  self.conv_dim, self.num_heads], dim=-1,
+        )
+
+        # 2. Convolution sequence transformation
+        if last_state is not None:
+            if input_states.shape[1] != 1:
+                raise ValueError("Mamba2 cached decoding only supports a single new token per step.")
+            # Decode path: single-step update
+            conv_state = last_state['conv_state']
+            ssm_state = last_state['recurrent_state']
+
+            conv_state = conv_state.roll(shifts=-1, dims=-1)
+            conv_state[:, :, -1] = hidden_states_B_C[:, 0, :].to(conv_state.device)
+
+            # We need to guarantee that anything regarding the cache is on the same device
+            conv_states_for_compute = conv_state.to(device=self.conv1d.weight.device)
+
+            hidden_states_B_C = torch.sum(
+                conv_states_for_compute * self.conv1d.weight.squeeze(1), dim=-1,
+            )
+            if self.use_conv_bias:
+                hidden_states_B_C = hidden_states_B_C + self.conv1d.bias
+            hidden_states_B_C = self.act(hidden_states_B_C)
+        else:
+            # Prefill path
+            hidden_states_B_C = apply_mask_to_padding_states(hidden_states_B_C, attention_mask)
+            new_conv_state = None
+            if use_cache:
+                hidden_states_B_C_transposed = hidden_states_B_C.transpose(1, 2)
+                new_conv_state = nn.functional.pad(
+                    hidden_states_B_C_transposed, (self.conv_kernel_size - hidden_states_B_C_transposed.shape[-1], 0),
+                )
+
+            hidden_states_B_C = self.act(self.conv1d(hidden_states_B_C.transpose(1, 2))[..., :seq_len].transpose(1, 2))
+
+        if last_state is None:
+            hidden_states_B_C = (hidden_states_B_C * attention_mask[:, :, None]).to(hidden_states_B_C.dtype) \
+                if attention_mask is not None and attention_mask.shape[1] > 1 and attention_mask.shape[0] > 1 \
+                else hidden_states_B_C
+        hidden_states, B, C = torch.split(
+            hidden_states_B_C,
+            [self.intermediate_size, self.n_groups * self.ssm_state_size, self.n_groups * self.ssm_state_size],
+            dim=-1,
+        )
+
+        # 3. SSM transformation
+        A = -torch.exp(self.A_log.float())                            # [num_heads]
+        if last_state is not None:
+            # Decode path
+            cache_device = ssm_state.device
+
+            # Note: there is no need to pad parameter matrices here, as there is just one new token
+            # for batched generation
+            dt = dt[:, 0, :][:, None, ...]
+            dt = dt.transpose(1, 2).expand(batch_size, dt.shape[-1], self.head_dim)
+            # [num_heads] -> [num_heads, head_dim]
+            dt_bias = self.dt_bias[..., None].expand(self.dt_bias.shape[0], self.head_dim)
+
+            dt = torch.nn.functional.softplus(dt + dt_bias.to(dt.dtype))
+            dt = torch.clamp(dt, self.time_step_limit[0], self.time_step_limit[1])
+            A = A[..., None, None].expand(self.num_heads, self.head_dim, self.ssm_state_size).to(dtype=torch.float32)
+            # [bsz, num_heads, head_dim, state_size]
+            dA = (torch.exp(dt[..., None] * A)).to(device=cache_device)
+
+            # Discretize B
+            # [bsz, n_groups * state_size] -> [bsz, n_groups, 1, state_size] ->
+            # -> [bsz, n_groups, group to head repetition factor, state_size] -> [bsz, num_heads, state_size]
+            B = B.reshape(batch_size, self.n_groups, -1)[..., None, :]
+            B = B.expand(batch_size, self.n_groups, self.num_heads // self.n_groups, B.shape[-1]).contiguous()
+            B = B.reshape(batch_size, -1, B.shape[-1])
+            # [bsz, num_heads, head_dim, state_size]
+            dB = dt[..., None] * B[..., None, :]
+
+            # Discretize x into dB
+            # [bsz, intermediate_size] -> [bsz, num_heads, head_dim]
+            hidden_states = hidden_states.reshape(batch_size, -1, self.head_dim)
+            dBx = (dB * hidden_states[..., None]).to(device=cache_device)
+
+            # State calculation
+            ssm_state = ssm_state * dA + dBx
+
+            # Subsequent output
+            # [bsz, n_groups * state_size] -> [bsz, num_heads, state_size]
+            C = C.reshape(batch_size, self.n_groups, -1)[..., None, :]
+            C = C.expand(batch_size, self.n_groups, self.num_heads // self.n_groups, C.shape[-1]).contiguous()
+            C = C.reshape(batch_size, -1, C.shape[-1])
+            # [bsz, num_heads, head_dim]
+
+            ssm_states_for_compute = ssm_state.to(device=C.device, dtype=C.dtype)  # Shape: [b, h, d, n]
+            # Reshape ssm_states to merge the first two dimensions
+            # Shape: [b*h, d, n]
+            ssm_states_reshaped = ssm_states_for_compute.view(batch_size * self.num_heads, self.head_dim, self.ssm_state_size)
+            C_reshaped = C.view(batch_size * self.num_heads, self.ssm_state_size, 1)  # Shape: [b*h, n, 1]
+            y = torch.bmm(ssm_states_reshaped, C_reshaped)
+            y = y.view(batch_size, self.num_heads, self.head_dim)
+
+            # D skip connection
+            # [num_heads] -> [num_heads, head_dim]
+            D = self.D[..., None].expand(self.D.shape[0], self.head_dim)
+            y = (y + hidden_states * D).to(y.dtype)
+
+            # [bsz, num_heads, head_dim] -> [bsz, 1, intermediate_size]
+            y = y.reshape(batch_size, -1)[:, None, ...]
+
+            scan_output = self.norm(y, gate)
+            contextualized_states = self.out_proj(scan_output.to(dtype))
+            return contextualized_states, conv_state, ssm_state
+        else:
+            # Prefill path
+            # begin ssd naive implementation without einsums
+            dt = nn.functional.softplus(dt + self.dt_bias)
+            dt = torch.clamp(dt, self.time_step_limit[0], self.time_step_limit[1])
+            hidden_states = hidden_states.reshape(batch_size, seq_len, -1, self.head_dim).float()
+            B = B.reshape(batch_size, seq_len, -1, self.ssm_state_size).float()
+            C = C.reshape(batch_size, seq_len, -1, self.ssm_state_size).float()
+            B = B.repeat(1, 1, self.num_heads // self.n_groups, 1)
+            C = C.repeat(1, 1, self.num_heads // self.n_groups, 1)
+            pad_size = (self.chunk_size - seq_len % self.chunk_size) % self.chunk_size
+
+            D_residual = self.D[..., None] * pad_tensor_by_size(hidden_states, pad_size)
+
+            # Discretize x and A
+            hidden_states = hidden_states * dt[..., None]
+            A = A.to(hidden_states.dtype) * dt
+
+            # Rearrange into blocks/chunks
+            hidden_states, A, B, C = [reshape_into_chunks(t, pad_size, self.chunk_size) for t in (hidden_states, A, B, C)]
+
+            # [bsz, -1, chunk_size, num_heads] -> [bsz, num_heads, -1, chunk_size]
+            A = A.permute(0, 3, 1, 2)
+            A_cumsum = torch.cumsum(A, dim=-1)
+
+            # 1. Compute the output for each intra-chunk (diagonal blocks)
+            # This is the analog of a causal mask
+            L = torch.exp(segment_sum(A))
+
+            # Contraction of C and B to get G (attention-weights like)
+            # shape: (b, c, l, s, h, n)
+            G_intermediate = C[:, :, :, None, :, :] * B[:, :, None, :, :, :]
+            G = G_intermediate.sum(dim=-1)  # shape: (b, c, l, s, h)
+
+            # Compute M, equivalent to applying attention mask to weights
+            M_intermediate = G[..., None] * L.permute(0, 2, 3, 4, 1)[..., None]
+            M = M_intermediate.sum(dim=-1)
+
+            # Compute Y_diag (apply to values)
+            Y_diag = (M[..., None] * hidden_states[:, :, None]).sum(dim=3)
+
+            # 2. Compute the state for each intra-chunk
+            # (right term of low-rank factorization of off-diagonal blocks; B terms)
+            decay_states = torch.exp(A_cumsum[:, :, :, -1:] - A_cumsum)
+            B_decay = B * decay_states.permute(0, -2, -1, 1)[..., None]
+            states = (B_decay[..., None, :] * hidden_states[..., None]).sum(dim=2)
+
+            # 3. Compute the inter-chunk SSM recurrence; produces correct SSM states at chunk boundaries
+            # (middle term of factorization of off-diag blocks; A terms)
+            previous_states = torch.zeros_like(states[:, :1])
+            states = torch.cat([previous_states, states], dim=1)
+            decay_chunk = torch.exp(segment_sum(nn.functional.pad(A_cumsum[:, :, :, -1], (1, 0))))
+            decay_chunk = decay_chunk.transpose(1, 3)
+            new_states = (decay_chunk[..., None, None] * states[:, :, None, ...]).sum(dim=1)
+            states, ssm_state = new_states[:, :-1], new_states[:, -1]
+
+            # 4. Compute state -> output conversion per chunk
+            # (left term of low-rank factorization of off-diagonal blocks; C terms)
+            state_decay_out = torch.exp(A_cumsum)
+            C_times_states = (C[..., None, :] * states[:, :, None, ...])
+            state_decay_out_permuted = state_decay_out.permute(0, 2, 3, 1)
+            Y_off = (C_times_states.sum(-1) * state_decay_out_permuted[..., None])
+
+            # Add output of intra-chunk and inter-chunk terms (diagonal and off-diagonal blocks)
+            y = Y_diag + Y_off
+            # [bsz, -1, self.chunk_size, num_heads, head_dim] -> [bsz, (padded) seq_len, num_heads, head_dim]
+            y = y.reshape(batch_size, -1, self.num_heads, self.head_dim)
+
+            y = y + D_residual
+            # Cutting off padded chunks
+            if pad_size > 0:
+                y = y[:, :seq_len, :, :]
+            y = y.reshape(batch_size, seq_len, -1)
+
+            scan_output = self.norm(y, gate)
+
+            # end ssd naive
+
+            # 4. Final linear projection
+            contextualized_states = self.out_proj(scan_output.to(dtype))  # [batch, seq_len, hidden_size]
+            return contextualized_states, new_conv_state if use_cache else None, ssm_state
+    # fmt: on
+
+    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,
+    ) -> tuple[torch.Tensor, torch.Tensor | None, Cache | None]:
+        last_state = get_layer_cache(self, past_key_values)
+
+        if is_fast_path_available and "cuda" in self.in_proj.weight.device.type:
+            output, conv_state, ssm_state = self.cuda_kernels_forward(hidden_states, last_state, use_cache, attention_mask)
+        else:
+            dtype = hidden_states.dtype
+            if last_state is None and attention_mask is not None and attention_mask.shape[1] > 1 and attention_mask.shape[0] > 1:
+                hidden_states = (hidden_states * attention_mask[:, :, None]).to(dtype)
+            output, conv_state, ssm_state = self.torch_forward(hidden_states, last_state, use_cache, attention_mask)
+
+        update_layer_cache(
+            self,
+            past_key_values,
+            recurrent_state=ssm_state,
+            conv_state=conv_state,
+            offset=hidden_states.shape[1],
+        )
+
+        return output, None, past_key_values

fla/layers/mesa_net.py

@@ -0,0 +1,221 @@
+# 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 torch.nn import functional as F
+
+from fla.layers.utils import get_layer_cache, get_unpad_data, index_first_axis, pad_input, update_layer_cache
+from fla.modules import FusedRMSNormGated, RMSNorm, ShortConvolution
+from fla.modules.l2norm import l2_norm
+from fla.ops.mesa_net import chunk_mesa_net, mesa_net_decoding_one_step
+
+if TYPE_CHECKING:
+    from transformers.processing_utils import Unpack
+
+    from fla.models.utils import Cache
+
+
+class MesaNet(nn.Module):
+    """
+    The layer implementaion for [MesaNet: Sequence Modeling by Locally Optimal Test-Time Training].  # noqa
+
+    Args:
+        hidden_size (int, Optional):
+            The hidden size of the input. Default: 2048.
+        expand_v (float, Optional):
+            The expansion ratio for the value dim. Default: 1.
+        num_heads (int, Optional):
+            The number of heads. Default: 16.
+        mode (str, Optional):
+            Which MesaNet kernel to use.
+            Currently available: `chunk`.
+            Default: `chunk`.
+        use_output_gate (bool, Optional):
+            Whether to use output gate. Default: `False`.
+        conv_size (int):
+            The kernel size of the short convolution. Default: 4.
+        layer_idx (int, Optional):
+            The index of the layer. Default: None.
+        norm_eps (float, Optional):
+            The epsilon value for the normalization layer. Default: 1e-5.
+        lambda_lower_bound (float):
+            The lower bound for the lambda parameter. Default: 0.25.
+        max_cg_step_training (int):
+            The maximum number of CG steps for training. Default: 30.
+        max_cg_step_decoding (int):
+            The maximum number of CG steps for decoding. Default: 30.
+    """
+
+    def __init__(
+        self,
+        hidden_size: int = 2048,
+        num_heads: int = 16,
+        head_dim: int = 128,
+        mode: str = 'chunk',
+        use_output_gate: bool = False,
+        use_short_conv: bool = True,
+        conv_size: int = 4,
+        conv_bias: bool = False,
+        layer_idx: int = None,
+        norm_eps: float = 1e-5,
+        lambda_lower_bound: float = 0.25,
+        max_cg_step_training: int = 30,
+        max_cg_step_decoding: int = 30,
+        **kwargs,
+    ) -> MesaNet:
+        super().__init__()
+
+        self.mode = mode
+        self.hidden_size = hidden_size
+        self.use_output_gate = use_output_gate
+        self.use_short_conv = use_short_conv
+        self.conv_size = conv_size
+        self.conv_bias = conv_bias
+        self.num_heads = num_heads
+        self.head_dim = head_dim
+        self.key_dim = self.num_heads * self.head_dim
+        self.value_dim = self.key_dim
+        self.head_k_dim = self.head_dim
+        self.head_v_dim = self.head_dim
+        self.layer_idx = layer_idx
+        self.lambda_lower_bound = lambda_lower_bound
+        self.max_cg_step_training = max_cg_step_training
+        self.max_cg_step_decoding = max_cg_step_decoding
+
+        self.q_proj = nn.Linear(hidden_size, self.key_dim, bias=False)
+        self.k_proj = nn.Linear(hidden_size, self.key_dim, bias=False)
+        self.v_proj = nn.Linear(hidden_size, self.value_dim, bias=False)
+        self.a_proj = nn.Linear(hidden_size, self.num_heads, bias=True)
+        self.b_proj = nn.Linear(hidden_size, self.num_heads, bias=True)
+
+        lambda_initial_value = 1.0
+        init_lamb_value = torch.log(torch.exp(torch.tensor(lambda_initial_value - lambda_lower_bound)) - 1.0)
+        init_lamb_params = torch.empty(self.key_dim, dtype=torch.float32).fill_(init_lamb_value)
+
+        self.lambda_params = nn.Parameter(init_lamb_params)
+        self.lambda_params._no_weight_decay = True
+
+        self.conv_size = conv_size
+        self.q_conv1d = ShortConvolution(
+            hidden_size=self.key_dim,
+            kernel_size=conv_size,
+            bias=self.conv_bias,
+            activation='silu',
+        )
+        self.k_conv1d = ShortConvolution(
+            hidden_size=self.key_dim,
+            kernel_size=conv_size,
+            bias=self.conv_bias,
+            activation='silu',
+        )
+        if use_output_gate:
+            self.g_proj = nn.Linear(hidden_size, self.value_dim, bias=False)
+            self.o_norm = FusedRMSNormGated(self.head_v_dim, eps=norm_eps)
+        else:
+            self.o_norm = RMSNorm(self.head_v_dim, eps=norm_eps, dtype=torch.float32)
+        self.o_proj = nn.Linear(self.value_dim, hidden_size, bias=False)
+
+    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]:
+        if attention_mask is not None:
+            assert len(attention_mask.shape) == 2, (
+                "Expected attention_mask as a 0-1 matrix with shape [batch_size, seq_len] "
+                "for padding purposes (0 indicating padding). "
+                "Arbitrary attention masks of shape [batch_size, seq_len, seq_len] are not allowed."
+            )
+
+        batch_size, q_len, _ = hidden_states.shape
+        last_state = get_layer_cache(self, past_key_values)
+
+        cu_seqlens = kwargs.get('cu_seqlens')
+        if attention_mask is not None:
+            indices, cu_seqlens, _ = get_unpad_data(attention_mask[:, -q_len:])
+            hidden_states = index_first_axis(rearrange(hidden_states, "b s ... -> (b s) ..."), indices).unsqueeze(0)
+
+        conv_state_q, conv_state_k = None, None
+        if last_state is not None:
+            conv_state_q, conv_state_k = last_state['conv_state']
+        q, conv_state_q = self.q_conv1d(
+            x=self.q_proj(hidden_states),
+            cache=conv_state_q,
+            output_final_state=use_cache,
+            cu_seqlens=cu_seqlens,
+        )
+        k, conv_state_k = self.k_conv1d(
+            x=self.k_proj(hidden_states),
+            cache=conv_state_k,
+            output_final_state=use_cache,
+            cu_seqlens=cu_seqlens,
+        )
+        v = self.v_proj(hidden_states)
+
+        q, k = map(lambda x: rearrange(x, '... (h d) -> ... h d', d=self.head_k_dim), (q, k))
+        v = rearrange(v, '... (h d) -> ... h d', d=self.head_v_dim)
+        beta = self.b_proj(hidden_states).float().sigmoid()
+        g = F.logsigmoid(self.a_proj(hidden_states).float())
+        lamb = F.softplus(self.lambda_params.float()) + self.lambda_lower_bound
+        lamb = lamb.reshape(self.num_heads, -1)
+
+        last_h_kk, last_h_kv = last_state['recurrent_state'] if last_state is not None else (None, None)
+
+        # prefilling or training
+        # Note that QK will be normalized inside the kernel to avoid saving the activations, thereby reducing the memory usage.
+        if last_state is None:
+            o, h_kk, h_kv = chunk_mesa_net(
+                q=q,
+                k=k,
+                v=v,
+                g=g,
+                beta=beta,
+                lamb=lamb,
+                output_final_state=use_cache,
+                max_CG_iteration=self.max_cg_step_training,
+                use_qk_l2norm_in_kernel=True,
+                cu_seqlens=cu_seqlens,
+            )
+        # decoding
+        else:
+            q = l2_norm(q)
+            k = l2_norm(k)
+            o, h_kk, h_kv = mesa_net_decoding_one_step(
+                q=q.squeeze(0),
+                k=k.squeeze(0),
+                v=v.squeeze(0),
+                g=g.squeeze(0),
+                beta=beta.squeeze(0),
+                lamb=lamb,
+                prev_h_kk=last_h_kk,
+                prev_h_kv=last_h_kv,
+                max_CG_iteration=self.max_cg_step_decoding,
+            )
+            o = o.unsqueeze(0).to(q)
+
+        update_layer_cache(
+            self,
+            past_key_values,
+            recurrent_state=(h_kk, h_kv),
+            conv_state=(conv_state_q, conv_state_k),
+            offset=q_len,
+        )
+        if self.use_output_gate:
+            g = rearrange(self.g_proj(hidden_states), '... (h d) -> ... h d', d=self.head_v_dim)
+            o = self.o_norm(o, g)
+        else:
+            o = self.o_norm(o)
+        o = rearrange(o, 'b t h d -> b t (h d)')
+        o = self.o_proj(o)
+        if attention_mask is not None:
+            o = pad_input(o.squeeze(0), indices, batch_size, q_len)
+        return o, None, past_key_values

fla/layers/mla.py

@@ -0,0 +1,225 @@
+# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang
+
+""" Implementing the Deepseek Multi Latent Attention (MLA) module. Reference:
+
+https://github.com/huggingface/transformers/blob/main/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py#L328
+"""
+
+from __future__ import annotations
+
+import math
+import warnings
+from typing import TYPE_CHECKING
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange, repeat
+from transformers.utils import logging
+
+from fla.layers.utils import pad_input, unpad_input
+from fla.modules import RMSNorm, RotaryEmbedding
+from fla.ops.utils.index import prepare_lens_from_mask
+
+if TYPE_CHECKING:
+    from fla.models.utils import Cache
+
+try:
+    from flash_attn import flash_attn_func, flash_attn_varlen_func
+except ImportError:
+    warnings.warn(
+        "Flash Attention is not installed. Please install it via `pip install flash-attn --no-build-isolation`",
+        category=ImportWarning,
+    )
+    flash_attn_func = None
+
+logger = logging.get_logger(__name__)
+
+
+def yarn_get_mscale(scale=1, mscale=1):
+    if scale <= 1:
+        return 1.0
+    return 0.1 * mscale * math.log(scale) + 1.0
+
+
+class MultiheadLatentAttention(nn.Module):
+    r"""
+    Multi-headed attention from [Deepseek V2](https://arxiv.org/abs/2405.04434)
+    """
+
+    def __init__(
+        self,
+        hidden_size: int = 2048,
+        num_heads: int = 16,
+        q_lora_rank: int | None = 1536,  # q lora rank is optional, None indicates no q lora
+        qk_rope_head_dim: int = 64,
+        kv_lora_rank: int = 512,  # following the original Deepseek paper
+        v_head_dim: int = 128,
+        qk_nope_head_dim: int = 128,
+        qk_head_dim: int | None = 192,  # qk_nope_head_dim + qk_rope_head_dim
+        window_size: int | None = None,
+        rope_theta: float = 10000.,
+        max_position_embeddings: int | None = None,
+        rope_scaling: dict | None = None,
+        layer_idx: int = None,
+    ) -> MultiheadLatentAttention:
+        super().__init__()
+
+        # sanity check
+        if qk_head_dim is not None:
+            assert qk_head_dim == qk_nope_head_dim + qk_rope_head_dim, \
+                f"qk_head_dim {qk_head_dim} != qk_nope_head_dim {qk_nope_head_dim} + qk_rope_head_dim {qk_rope_head_dim}"
+        else:
+            qk_head_dim = qk_nope_head_dim + qk_rope_head_dim
+
+        # attention params info
+        self.hidden_size = hidden_size
+        self.num_heads = num_heads
+        self.q_lora_rank = q_lora_rank
+        self.qk_rope_head_dim = qk_rope_head_dim
+        self.kv_lora_rank = kv_lora_rank
+        self.v_head_dim = v_head_dim
+        self.qk_nope_head_dim = qk_nope_head_dim
+        self.qk_head_dim = qk_head_dim
+
+        self.window_size = window_size
+        self.rope_theta = rope_theta
+        self.max_position_embeddings = max_position_embeddings
+        self.layer_idx = layer_idx
+
+        if flash_attn_func is None:
+            raise ImportError("Please install Flash Attention via `pip install flash-attn --no-build-isolation` first")
+
+        if q_lora_rank is not None:
+            self.q_proj = nn.Sequential(
+                nn.Linear(hidden_size, q_lora_rank, bias=False),
+                RMSNorm(q_lora_rank, dtype=torch.float32),
+                nn.Linear(q_lora_rank, self.num_heads * self.qk_head_dim, bias=False),
+            )
+        else:
+            self.q_proj = nn.Linear(hidden_size, self.num_heads * self.qk_head_dim, bias=False)
+
+        self.k_rope = nn.Linear(hidden_size, self.qk_rope_head_dim, bias=False)
+        self.kv_proj = nn.Sequential(
+            nn.Linear(hidden_size, self.kv_lora_rank, bias=False),
+            RMSNorm(self.kv_lora_rank, dtype=torch.float32),
+            nn.Linear(self.kv_lora_rank, self.num_heads * (self.qk_nope_head_dim + self.v_head_dim), bias=False),
+        )
+
+        self.o_proj = nn.Linear(self.num_heads * self.v_head_dim, hidden_size, bias=False)
+
+        self.scaling = self.qk_head_dim ** (-0.5)
+        if rope_scaling is not None and rope_scaling.get("rope_type", "default") != "default":
+            mscale_all_dim = rope_scaling.get("mscale_all_dim", 0)
+            scaling_factor = rope_scaling["factor"]
+            if mscale_all_dim:
+                mscale = yarn_get_mscale(scaling_factor, mscale_all_dim)
+                self.scaling = self.scaling * mscale * mscale
+
+        self.rotary = RotaryEmbedding(dim=self.qk_rope_head_dim, base=self.rope_theta)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: torch.Tensor | None,
+        past_key_values: Cache | None = None,
+        output_attentions: bool = False,
+        use_cache: bool = False,
+        **kwargs,
+    ) -> tuple[torch.Tensor, torch.Tensor | None, tuple[torch.Tensor] | None]:
+        # if attention_mask is not None, this is doing inference
+        if attention_mask is not None:
+            assert len(attention_mask.shape) == 2, (
+                "Expected attention_mask as a 0-1 matrix with shape [batch_size, seq_len] "
+                "for padding purposes (0 indicating padding). "
+                "Arbitrary attention masks of shape [batch_size, seq_len, seq_len] are not allowed."
+            )
+
+        # prepare q, k, v
+        batch_size, q_len, _ = hidden_states.shape
+
+        q_states = self.q_proj(hidden_states)
+        q_states = rearrange(q_states, '... (h d) -> ... h d', d=self.qk_head_dim)
+        q_pass, q_rot = torch.split(q_states, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1)
+        k_pass, k_rot = self.kv_proj(hidden_states), self.k_rope(hidden_states)
+
+        k_rot = rearrange(k_rot, 'b t d -> b t 1 d')
+        k_pass = rearrange(k_pass, '... (h d) -> ... h d', d=self.qk_nope_head_dim + self.v_head_dim)
+        k_pass, v = torch.split(k_pass, [self.qk_nope_head_dim, self.v_head_dim], dim=-1)
+
+        # apply rotary position embedding
+        seqlen_offset, max_seqlen = 0, q_len
+        if past_key_values is not None:
+            seqlen_offset = past_key_values.get_seq_length(self.layer_idx)
+            max_seqlen = q_len + seqlen_offset
+
+            if attention_mask is not None:
+                seqlen_offset = seqlen_offset + prepare_lens_from_mask(attention_mask) - attention_mask.shape[-1]
+                max_seqlen = q_len + max(seqlen_offset)
+
+        if self.max_position_embeddings is not None:
+            max_seqlen = max(max_seqlen, self.max_position_embeddings)
+        cu_seqlens = kwargs.get("cu_seqlens")
+        q_rot, k_rot = self.rotary(
+            q_rot, k_rot, seqlen_offset=seqlen_offset, max_seqlen=max_seqlen, cu_seqlens=cu_seqlens,
+        )
+
+        k_rot = repeat(k_rot, 'b t 1 d -> b t h d', h=self.num_heads)
+        q = torch.cat((q_pass, q_rot), dim=-1)
+        k = torch.cat((k_pass, k_rot), dim=-1)
+
+        # TODO: instead of caching the full k, v, we can actually only cache the compressed_kv and k_rot
+        # and recover the full k, v from compressed_kv and k_rot
+        if past_key_values is not None:
+            cache_has_content = past_key_values.get_seq_length(self.layer_idx) > 0
+            k_cached, v_cached = past_key_values.update(
+                attn_state=(k, v),
+                layer_idx=self.layer_idx,
+                offset=q_len,
+            )['attn_state']
+            if cache_has_content:
+                k, v = k_cached, v_cached
+
+        # Head dim match to use flash-attn
+        if self.qk_head_dim != self.v_head_dim:
+            v = F.pad(v, [0, self.qk_head_dim - self.v_head_dim])
+
+        # Contains at least one padding token in the sequence
+        if attention_mask is not None:
+            if q.shape[1] == 1 and self.window_size is not None:
+                attention_mask = attention_mask[:, -self.window_size:]
+            q, (k, v), indices_q, cu_seqlens, max_seq_lens = unpad_input(q, (k, v), attention_mask, q_len)
+            cu_seqlens_q, cu_seqlens_k = cu_seqlens
+            max_seqlen_q, max_seqlen_k = max_seq_lens
+            o = flash_attn_varlen_func(
+                q, k, v,
+                cu_seqlens_q=cu_seqlens_q,
+                cu_seqlens_k=cu_seqlens_k,
+                max_seqlen_q=max_seqlen_q,
+                max_seqlen_k=max_seqlen_k,
+                causal=True,
+                window_size=(-1, -1) if self.window_size is None else (self.window_size-1, 0),
+            )
+            o = pad_input(o, indices_q, batch_size, q_len)
+        elif cu_seqlens is not None:
+            o = flash_attn_varlen_func(
+                q.squeeze(0), k.squeeze(0), v.squeeze(0),
+                cu_seqlens_q=cu_seqlens,
+                cu_seqlens_k=cu_seqlens,
+                max_seqlen_q=max_seqlen,
+                max_seqlen_k=max_seqlen,
+                causal=True,
+                window_size=(-1, -1) if self.window_size is None else (self.window_size-1, 0),
+            ).unsqueeze(0)
+        else:
+            o = flash_attn_func(
+                q, k, v,
+                causal=True,
+                window_size=(-1, -1) if self.window_size is None else (self.window_size-1, 0),
+            )
+
+        if self.qk_head_dim != self.v_head_dim:
+            o = o[:, :, :, :self.v_head_dim]
+        o = o.reshape(batch_size, q_len, -1)
+        o = self.o_proj(o)
+        return o, None, past_key_values

fla/layers/mom.py

@@ -0,0 +1,831 @@
+
+from __future__ import annotations
+
+import math
+from typing import TYPE_CHECKING
+
+import torch
+import torch.nn as nn
+from einops import rearrange
+from torch.nn import functional as F
+
+from fla.modules import FusedRMSNormGated, RMSNorm, ShortConvolution
+from fla.ops.gated_delta_rule import chunk_gated_delta_rule, fused_recurrent_gated_delta_rule
+
+if TYPE_CHECKING:
+    from transformers.processing_utils import Unpack
+
+    from fla.models.utils import Cache
+
+from fla.layers.utils import get_layer_cache, get_unpad_data, index_first_axis, pad_input, unpad_input, update_layer_cache
+
+
+def _upad_input(
+    query_layer: torch.Tensor,
+    key_layer: torch.Tensor,
+    value_layer: torch.Tensor,
+    gate_layer: torch.Tensor,
+    beta_layer: torch.Tensor,
+    attention_mask: torch.Tensor,
+):
+    """
+    Unpads query, key, and values tensors, using a single dimension for all tokens even though they belong to
+    different batches.
+
+    This function is used instead of `flash_attn.bert_padding.unpad_input` in order to avoid the recomputation
+    of the same intermediary
+    tensors for query, key, value tensors.
+
+    Arguments:
+        query_layer (`torch.Tensor`):
+            Query state with padding. Shape: (batch_size, query_length, num_heads, head_dim).
+        key_layer (`torch.Tensor`):
+            Key state with padding. Shape: (batch_size, kv_seq_len, num_key_value_heads, head_dim).
+        value_layer (`torch.Tensor`):
+            Value state with padding. Shape: (batch_size, kv_seq_len, num_key_value_heads, head_dim).
+        attention_mask (`torch.Tensor`):
+            Boolean or int tensor of shape (batch_size, sequence_length), 1 means valid and 0 means not valid.
+        query_length (`int`):
+            Target length.
+
+    Return:
+        query_layer (`torch.Tensor`):
+            Query state without padding. Shape: (total_target_length, num_heads, head_dim).
+        key_layer (`torch.Tensor`):
+            Key state with padding. Shape: (total_source_length, num_key_value_heads, head_dim).
+        value_layer (`torch.Tensor`):
+            Value state with padding. Shape: (total_source_length, num_key_value_heads, head_dim).
+        indices_q (`torch.Tensor`):
+            The indices of non-masked tokens from the flattened input target sequence.
+        (cu_seqlens_q, cu_seqlens_k) (`Tuple[int]`):
+            The cumulative sequence lengths for the target (query) and source (key, value), used to index
+            into ragged (unpadded) tensors. `cu_seqlens` shape is (batch_size + 1,).
+        (max_seqlen_in_batch_q, max_seqlen_in_batch_k) (`Tuple[int]`):
+            Maximum sequence length in batch (`max_seqlen_in_batch_q` for the target sequence i.e. query,
+            `max_seqlen_in_batch_k` for the source sequence i.e. key/value).
+    """
+    query_length = query_layer.shape[1]
+    indices_k, cu_seqlens_k, max_seqlen_in_batch_k = get_unpad_data(attention_mask)
+    batch_size, kv_seq_len, dim = key_layer.shape
+    v_dim = value_layer.shape[-1]
+
+    key_layer = index_first_axis(key_layer.reshape(batch_size * kv_seq_len, dim), indices_k)
+    value_layer = index_first_axis(
+        value_layer.reshape(batch_size * kv_seq_len, v_dim), indices_k,
+    )
+    gate_layer = index_first_axis(gate_layer.reshape(batch_size * kv_seq_len, -1), indices_k)
+    beta_layer = index_first_axis(beta_layer.reshape(batch_size * kv_seq_len, -1), indices_k)
+    if query_length == kv_seq_len:
+        query_layer = index_first_axis(query_layer.reshape(batch_size * kv_seq_len, dim), indices_k)
+        cu_seqlens_q = cu_seqlens_k
+        max_seqlen_in_batch_q = max_seqlen_in_batch_k
+        indices_q = indices_k
+    elif query_length == 1:
+        max_seqlen_in_batch_q = 1
+        cu_seqlens_q = torch.arange(
+            batch_size + 1, dtype=torch.int32, device=query_layer.device,
+        )  # There is a memcpy here, that is very bad.
+        indices_q = cu_seqlens_q[:-1]
+        query_layer = query_layer.squeeze(1)
+    else:
+        # The -q_len: slice assumes left padding.
+        attention_mask = attention_mask[:, -query_length:]
+        query_layer, indices_q, cu_seqlens_q, max_seqlen_in_batch_q = unpad_input(query_layer, attention_mask)
+
+    return (
+        query_layer,
+        key_layer,
+        value_layer,
+        gate_layer,
+        beta_layer,
+        indices_q,
+        (cu_seqlens_q, cu_seqlens_k),
+        (max_seqlen_in_batch_q, max_seqlen_in_batch_k),
+    )
+
+
+def transform(
+    x: torch.Tensor,
+    routing_mask: torch.Tensor,
+    num_memories: int,
+    selected_memories: torch.Tensor,
+    attention_mask: torch.Tensor,
+):
+    """
+    Reorganize token embeddings into memory-aligned chunks.
+
+    Steps:
+        - Expand for top-k routing if needed.
+        - Mask out padded tokens via `attention_mask`.
+        - Sort tokens by (batch, memory).
+        - Gather and pad tokens per memory slot.
+
+    Args:
+        x: (batch, seq, hidden) input embeddings.
+        routing_mask: (batch, seq, num_memories) binary routing mask.
+        num_memories: number of memory slots.
+        selected_memories: memory indices per token,
+            (batch, seq) if k=1 else (batch, seq, topk).
+        attention_mask: (batch, seq) valid-token mask.
+
+    Returns:
+        transformed_x: (num_memories, batch, max_len, hidden) reorganized tokens.
+        truncation_indices: (batch*num_memories, max_len) gather indices.
+        sorted_indices: (batch*seq*topk,) global sort order.
+        max_len: int, max tokens per memory.
+        mask: (batch*num_memories, max_len) validity mask.
+        mask_2: (num_memories, batch, max_len) validity mask reshaped.
+    """
+    if selected_memories.dim() == 3:
+        # (batch, seq, topk)
+        topk = selected_memories.shape[2]
+        # x (batch, seq, hidden)
+        x = x.repeat_interleave(topk, dim=1)
+        # x (batch, seq * topk, hidden)
+        # (batch, seq, topk)
+        selected_memories = selected_memories.reshape(selected_memories.shape[0], -1)
+        # (batch, seq * topk)
+
+    if attention_mask is not None:
+        attention_mask = attention_mask[:, -routing_mask.shape[1]:]
+        # mask out the masked tokens
+        routing_mask[attention_mask.bitwise_not().unsqueeze(-1).expand(-1, -1, num_memories)] = 0
+
+    b, s, d = x.shape
+    x_flat = x.reshape(b * s, d)  # [b*s, d]
+
+    with torch.no_grad():
+        batch_indices = torch.arange(b, device=x.device).unsqueeze(-1)
+        batch_indices = batch_indices.repeat(1, s).reshape(-1)
+        if attention_mask is not None:
+            # sort the masked tokens to the end
+            batch_indices[attention_mask.repeat_interleave(topk, dim=1).bitwise_not().flatten()] = b
+        # (b * s)
+        memories_flat = selected_memories.reshape(-1)  # [b*s]
+
+        combined = batch_indices * (memories_flat.max() + 1) + memories_flat
+        sorted_indices = combined.argsort()
+
+    x_sorted = x_flat[sorted_indices]  # [b*s, d]
+    # (b*s, hidden) -> (b, s, hidd)
+    with torch.no_grad():
+        # routing_mask (b, s, num_memories)
+        batch_memory_tokens = routing_mask.sum(dim=1)
+        # (b, num_memories)
+        flatten_offset = batch_memory_tokens.flatten().cumsum(dim=0)
+        max_len = batch_memory_tokens.max()
+        indices = (
+            torch.arange(max_len, device=flatten_offset.device).unsqueeze(0).expand(b * num_memories, -1)
+            + torch.cat([torch.tensor([0], device=flatten_offset.device), flatten_offset[:-1]], dim=0).unsqueeze(1)
+        )
+        mask = indices < flatten_offset.unsqueeze(-1)
+        truncation_indices = torch.where(mask, indices, torch.zeros_like(indices))
+
+    gathered_x = torch.gather(x_sorted, 0, truncation_indices.reshape(-1).unsqueeze(-1).expand(-1, d))
+    transformed_x = gathered_x.reshape(b * num_memories, -1, d).reshape((b, num_memories, max_len, d)).transpose(0, 1)
+    # transformed_x = transformed_x * mask.unsqueeze(-1).expand_as(transformed_x)
+    # pad_x = torch.zeros((b * num_memories, capacity_len-max_len, d), dtype=transformed_x.dtype, device=transformed_x.device)
+    # pad_mask = torch.zeros((b * num_memories, capacity_len-max_len), dtype=transformed_x.dtype, device=transformed_x.device)
+    # left pad
+    # transformed_x = torch.cat((pad_x, transformed_x), dim=1).reshape((b, num_memories, capacity_len, d)).transpose(0, 1)
+    mask_2 = mask.reshape((b, num_memories, max_len)).transpose(0, 1)
+    # truncation_indices += capacity_len-max_len
+    # if attention_mask is not None:
+    #     mask_2
+
+    return transformed_x, truncation_indices, sorted_indices, max_len, mask, mask_2
+
+
+def reconstruct(
+    transformed_x,
+    indices: torch.Tensor,
+    sorted_indices: torch.Tensor,
+    batch_size: int,
+    seq_len: int,
+    topk: int,
+    routing_weights: torch.Tensor,
+    mask: torch.Tensor,
+):
+    '''
+    Reconstruct and mix transformed outputs back into the original input sequence shape.
+
+    Key operations:
+    1. Reshapes and transposes `transformed_x` to prepare for scattering.
+    2. Applies the `mask` to zero out invalid positions.
+    3. Uses `torch.scatter_add_` to scatter and sum the transformed outputs back to their original positions
+        based on `indices`.
+    4. Rearranges the scattered outputs using `sorted_indices` to ensure correct ordering.
+    5. Applies the `routing_weights` to weight the outputs.
+    6. Sums over the `topk` dimension to produce the final reconstructed output.
+
+    Args:
+        transformed_x (torch.Tensor):
+            The transformed output tensor from memory units or experts.
+            Shape: (num_memories, batch_size, capacity_len, hidden_size)
+        indices (torch.Tensor):
+            Indices used for scattering the transformed outputs back to their corresponding positions.
+            Shape: (batch*num_memories, max_len)
+        sorted_indices (torch.Tensor):
+            Sorting indices used to rearrange the scattered outputs back into the original sequence order.
+            Shape: (batch_size*seq_len*topk)
+        batch_size (int):
+            The size of the batch.
+        seq_len (int):
+            The length of the input sequence.
+        topk (int):
+            The number of top elements selected (`topk`) per token during the selection process.
+        routing_weights (torch.Tensor):
+            Routing weights assigned to the top-k selected outputs when reconstructing the final output.
+            Shape: (batch_size, seq_len, topk)
+        mask (torch.Tensor):
+            Boolean mask indicating valid positions in the sequence.
+            Shape: (batch*num_memories, max_len)
+
+    Returns:
+        restored_x (torch.Tensor):
+            The reconstructed output tensor in the original input sequence shape.
+            Shape: (batch_size, seq_len, hidden_size)
+    '''
+    transformed_x = transformed_x.transpose(0, 1).reshape(
+        (-1, transformed_x.shape[2], transformed_x.shape[3]))
+    b, s, k, d = batch_size, seq_len, topk, transformed_x.shape[2]
+    gathered_x = transformed_x.reshape(
+        (transformed_x.shape[0] * transformed_x.shape[1], transformed_x.shape[2]))
+    mask_expanded = mask.reshape(-1).unsqueeze(-1).expand_as(gathered_x)
+    gathered_x = gathered_x * mask_expanded
+
+    assert (indices >= 0).all(), "Indices should be non-negative"
+
+    resortd_x = torch.zeros((b * s * k, d), device=gathered_x.device, dtype=gathered_x.dtype).scatter_add_(
+        0,
+        indices.reshape(-1).unsqueeze(-1).expand(-1, d),
+        gathered_x,
+    )
+    assert (indices < resortd_x.size(0)).all(), "Indices should be less than resortd_x size"
+
+    inverse_indices = sorted_indices.argsort()
+    rearranged_x_flat = resortd_x[inverse_indices]
+    restored_x = rearranged_x_flat.reshape((b, s * k, d))
+    restored_x = restored_x.reshape(b, s, k, d) * routing_weights.reshape(b, s, k).unsqueeze(-1)
+    restored_x = restored_x.sum(dim=2)
+    return restored_x
+
+
+class MomAttention(nn.Module):
+    """
+    The layer implementaion for [MoM: Linear Sequence Modeling with Mixture-of-Memories](https://arxiv.org/abs/2502.13685).
+    """
+
+    def __init__(
+        self,
+        hidden_size: int = 2048,
+        head_dim: int = 256,
+        num_heads: int = 4,
+        expand_v: float = 2,
+        mode: str = 'chunk',
+        use_output_gate: bool = True,
+        use_short_conv: bool = True,
+        conv_size: int = 4,
+        conv_bias: bool = False,
+        layer_idx: int = None,
+        norm_eps: float = 1e-5,
+        num_memories: int = 8,
+        topk: int = 2,
+        capacity: float = 1.0,
+        shared_mem: bool = False,
+        single_kv_proj: bool = False,
+        **kwargs,
+    ) -> MomAttention:
+        super().__init__()
+        self.num_memories = num_memories
+        self.topk = topk
+        self.capacity = capacity
+        self.shared_mem = shared_mem
+        self.single_kv_proj = single_kv_proj
+
+        self.mode = mode
+
+        self.hidden_size = hidden_size
+        self.expand_v = expand_v
+
+        self.use_output_gate = use_output_gate
+        self.use_short_conv = use_short_conv
+        self.conv_size = conv_size
+        self.conv_bias = conv_bias
+
+        self.head_dim = head_dim
+        self.num_heads = num_heads
+
+        self.key_dim = int(self.num_heads * self.head_dim)
+        self.value_dim = int(self.key_dim * self.expand_v)
+        self.head_qk_dim = head_dim
+        self.head_v_dim = int(head_dim * self.expand_v)
+        self.layer_idx = layer_idx
+        self.silu = nn.SiLU()
+
+        assert mode in ['chunk', 'fused_recurrent'], f"Not suppoerted mode `{mode}`."
+
+        self.q_proj = nn.Linear(hidden_size, self.key_dim, bias=False)
+        self.gate = nn.Linear(self.hidden_size, self.num_memories, bias=False)
+        if self.single_kv_proj:
+            self.shared_k = nn.Linear(hidden_size, self.key_dim, bias=False)
+            self.shared_v = nn.Linear(hidden_size, self.value_dim, bias=False)
+            self.shared_b = nn.Linear(hidden_size, self.num_heads, bias=False)
+            self.shared_a = nn.Linear(hidden_size, self.num_heads, bias=False)
+        else:
+            self.k_proj = nn.ModuleList([
+                nn.Linear(self.hidden_size, self.key_dim, bias=False)
+                for _ in range(self.num_memories)
+            ])
+            self.v_proj = nn.ModuleList([
+                nn.Linear(self.hidden_size, self.value_dim, bias=False)
+                for _ in range(self.num_memories)
+            ])
+            self.b_proj = nn.ModuleList([
+                nn.Linear(self.hidden_size, self.num_heads, bias=False)
+                for _ in range(self.num_memories)
+            ])
+            self.a_proj = nn.ModuleList([
+                nn.Linear(self.hidden_size, self.num_heads, bias=False)
+                for _ in range(self.num_memories)
+            ])
+            if self.shared_mem:
+                self.shared_k = nn.Linear(hidden_size, self.key_dim, bias=False)
+                self.shared_v = nn.Linear(hidden_size, self.value_dim, bias=False)
+                self.shared_b = nn.Linear(hidden_size, self.num_heads, bias=False)
+                self.shared_a = nn.Linear(hidden_size, self.num_heads, bias=False)
+
+        A = torch.empty(self.num_heads, dtype=torch.float32).uniform_(0, 16)
+        self.A_log = nn.Parameter(torch.log(A))
+        self.A_log._no_weight_decay = True
+        # hard coded for now
+        dt_min = 0.001
+        dt_max = 0.1
+        dt_init_floor = 1e-4
+        dt = torch.exp(
+            torch.rand(self.num_heads) * (math.log(dt_max) - math.log(dt_min))
+            + math.log(dt_min),
+        )
+        dt = torch.clamp(dt, min=dt_init_floor)
+        # Inverse of softplus: https://github.com/pytorch/pytorch/issues/72759
+        inv_dt = dt + torch.log(-torch.expm1(-dt))
+        self.dt_bias = nn.Parameter(inv_dt)
+        # Just to be explicit. Without this we already don't put wd on dt_bias because of the check
+        # name.endswith("bias") in param_grouping.py
+        self.dt_bias._no_weight_decay = True
+
+        if use_short_conv:
+            self.conv_size = conv_size
+            self.q_conv1d = ShortConvolution(
+                hidden_size=self.key_dim,
+                kernel_size=conv_size,
+                bias=conv_bias,
+                activation='silu',
+            )
+            self.k_conv1d = ShortConvolution(
+                hidden_size=self.key_dim,
+                kernel_size=conv_size,
+                bias=conv_bias,
+                activation='silu',
+            )
+            self.v_conv1d = ShortConvolution(
+                hidden_size=self.value_dim,
+                kernel_size=conv_size,
+                bias=conv_bias,
+                activation='silu',
+            )
+        else:
+            raise UserWarning(
+                "ShortConvolution is crucial to the performance. "
+                "Do not turn it off, i.e., setting `use_short_conv=False` unless you know what you are doing.",
+            )
+        if use_output_gate:
+            self.g_proj = nn.Linear(hidden_size, self.value_dim, bias=False)
+            self.o_norm = FusedRMSNormGated(self.head_v_dim, eps=norm_eps)
+        else:
+            self.o_norm = RMSNorm(self.head_v_dim, eps=norm_eps, dtype=torch.float32)
+        self.o_proj = nn.Linear(self.value_dim, hidden_size, bias=False)
+        self.apply(self._initialize_weights)
+
+    def _initialize_weights(self, module: nn.Module):
+        if getattr(module, "_is_hf_initialized", False):
+            return
+        if isinstance(module, nn.Linear):
+            nn.init.xavier_uniform_(module.weight, gain=2 ** -2.5)
+            if module.bias is not None:
+                nn.init.zeros_(module.bias)
+        module._is_hf_initialized = True
+
+    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]:
+        if attention_mask is not None:
+            attention_mask = (attention_mask == 1)
+            assert len(attention_mask.shape) == 2, (
+                "Expected attention_mask as a 0-1 matrix with shape [batch_size, seq_len] "
+                "for padding purposes (0 indicating padding). "
+                "Arbitrary attention masks of shape [batch_size, seq_len, seq_len] are not allowed."
+            )
+
+        origin_cu_seqlens = kwargs.get('cu_seqlens')
+        if origin_cu_seqlens is not None:
+            hidden_states, attention_mask = self.cu2pad(hidden_states, origin_cu_seqlens)
+
+        mode = 'fused_recurrent' if (hidden_states.shape[1] <= 64 and not self.training) else self.mode
+        if self.training:
+            assert mode == 'chunk', "Only chunk mode is supported in training."
+
+        last_state = get_layer_cache(self, past_key_values)
+        # _, q_len = hidden_states.shape[0], hidden_states.shape[1]
+
+        # 🔍 topk gating
+        router_logits = self.gate(hidden_states)  # (bsz, q_len, num_memories)
+        scores = F.softmax(router_logits, dim=2, dtype=torch.float)
+        routing_weights, selected_memories = torch.topk(scores, self.topk, dim=-1)  # (bsz, seq, topk)
+        routing_weights /= routing_weights.sum(dim=-1, keepdim=True)
+        routing_weights = routing_weights.to(hidden_states.dtype)  # we cast back to the input dtype
+        routing_weights_full = torch.zeros(
+            routing_weights.shape[0],
+            routing_weights.shape[1],
+            self.num_memories,
+            dtype=routing_weights.dtype,
+            device=routing_weights.device,
+        ).scatter(-1, selected_memories, routing_weights)
+        routing_mask = routing_weights_full.bool().int()
+
+        # if self.use_output_gate:
+        #     o_g = self.g_proj(hidden_states)
+
+        batch_size, seq_len = hidden_states.shape[0], hidden_states.shape[1]
+
+        shared_hidden_states = hidden_states
+        hidden_states, indices, sorted_indices, max_len, mask, mask_2 = transform(
+            hidden_states, routing_mask, self.num_memories, selected_memories, attention_mask)
+
+        q = self.q_proj(hidden_states)
+        if self.single_kv_proj:
+            k = self.shared_k(hidden_states)
+            v = self.shared_v(hidden_states)
+            beta = self.shared_b(hidden_states).sigmoid()
+            g = -self.A_log.float().exp() * F.softplus(self.shared_a(hidden_states).float() + self.dt_bias)
+        else:
+            k = torch.stack([k_expert(hidden_states[i]) for i, k_expert in enumerate(self.k_proj)], dim=0)
+            v = torch.stack([v_expert(hidden_states[i]) for i, v_expert in enumerate(self.v_proj)], dim=0)
+            beta = torch.stack([b_expert(hidden_states[i]).sigmoid() for i, b_expert in enumerate(self.b_proj)], dim=0)
+            g = torch.stack([-self.A_log.float().exp() * F.softplus(a_expert(hidden_states[i]).float() + self.dt_bias)
+                            for i, a_expert in enumerate(self.a_proj)], dim=0)
+
+        q, k, v, g, beta, mask_2 = (rearrange(x, 'e b l ... ->  (e b) l ...') for x in (q, k, v, g, beta, mask_2))
+        cu_q, cu_k, cu_v, cu_g, cu_beta, indices_q, cu_seqlen_all, max_seq_lens = _upad_input(q, k, v, g, beta, mask_2)
+        cu_seqlens, reverse_indices = cu_seqlen_all[0].to(torch.long).unique(return_inverse=True)
+        cu_q, cu_k, cu_v, cu_g, cu_beta = (x.unsqueeze(0).contiguous() for x in (cu_q, cu_k, cu_v, cu_g, cu_beta))
+
+        if self.use_short_conv:
+            conv_state_q, conv_state_k, conv_state_v = [None, None], [None, None], [None, None]
+            if last_state is not None:
+                conv_state_q, conv_state_k, conv_state_v = last_state['conv_state']
+
+            conv_cu_seqlens = cu_seqlens
+            padded = False
+            if self.training:
+                conv_cu_seqlens = None
+            elif seq_len != 1 and (cu_seqlens[1:] - cu_seqlens[:-1]).min().item() < self.conv_size:
+                padded = True
+                conv_cu_seqlens, cu_q, cu_k, cu_v, pad_lengths = self.pad_for_conv(cu_seqlens, cu_q, cu_k, cu_v)
+
+            conv_q = self.prepare_recurrent_state(
+                conv_state_q[0],
+                conv_cu_seqlens,
+                cu_seqlen_all[0],
+                reverse_indices,
+                batch_size,
+            )
+            cu_q, conv_q_new = self.q_conv1d(
+                x=cu_q,
+                cache=conv_q,
+                output_final_state=use_cache,
+                cu_seqlens=conv_cu_seqlens,
+            )
+            conv_state_q[0] = self.handle_recurrent_state(
+                conv_state_q[0],
+                conv_q_new,
+                conv_cu_seqlens,
+                cu_seqlen_all[0],
+                reverse_indices,
+            )
+            conv_k = self.prepare_recurrent_state(
+                conv_state_k[0],
+                conv_cu_seqlens,
+                cu_seqlen_all[0],
+                reverse_indices,
+                batch_size,
+            )
+            cu_k, conv_k_new = self.k_conv1d(
+                x=cu_k,
+                cache=conv_k,
+                output_final_state=use_cache,
+                cu_seqlens=conv_cu_seqlens,
+            )
+            conv_state_k[0] = self.handle_recurrent_state(
+                conv_state_k[0],
+                conv_k_new,
+                conv_cu_seqlens,
+                cu_seqlen_all[0],
+                reverse_indices,
+            )
+            conv_v = self.prepare_recurrent_state(
+                conv_state_v[0],
+                conv_cu_seqlens,
+                cu_seqlen_all[0],
+                reverse_indices,
+                batch_size,
+            )
+            cu_v, conv_v_new = self.v_conv1d(
+                x=cu_v,
+                cache=conv_v,
+                output_final_state=use_cache,
+                cu_seqlens=conv_cu_seqlens,
+            )
+            conv_state_v[0] = self.handle_recurrent_state(
+                conv_state_v[0],
+                conv_v_new, conv_cu_seqlens,
+                cu_seqlen_all[0],
+                reverse_indices,
+            )
+
+            if padded:
+                cu_q, cu_k, cu_v = self.unpad_after_conv(conv_cu_seqlens, cu_seqlens, cu_q, cu_k, cu_v, pad_lengths)
+
+        else:
+            q, k, v = self.silu(q), self.silu(k), self.silu(v)
+
+        cu_q, cu_k, cu_v = map(lambda x: rearrange(x, 'b t (h d) -> b t h d', h=self.num_heads), (cu_q, cu_k, cu_v))
+
+        recurrent_state = last_state['recurrent_state'] if last_state is not None else [
+            None for _ in range(1 + self.shared_mem)]
+        if mode == 'chunk':
+            o, recurrent_state_ = chunk_gated_delta_rule(
+                q=cu_q,
+                k=cu_k,
+                v=cu_v,
+                g=cu_g,
+                beta=cu_beta,
+                initial_state=recurrent_state[0],
+                output_final_state=use_cache,
+                use_qk_l2norm_in_kernel=True,
+                cu_seqlens=cu_seqlens,
+            )
+            recurrent_state[0] = self.handle_recurrent_state(
+                recurrent_state[0],
+                recurrent_state_,
+                cu_seqlens,
+                cu_seqlen_all[0],
+                reverse_indices,
+            )
+
+        elif mode == 'fused_recurrent':
+            memories = self.prepare_recurrent_state(
+                recurrent_state[0],
+                cu_seqlens, cu_seqlen_all[0],
+                reverse_indices, batch_size,
+            )
+            o, recurrent_state_ = fused_recurrent_gated_delta_rule(
+                q=cu_q,
+                k=cu_k,
+                v=cu_v,
+                g=cu_g,
+                beta=cu_beta,
+                initial_state=memories,
+                output_final_state=use_cache,
+                use_qk_l2norm_in_kernel=True,
+                cu_seqlens=cu_seqlens,
+            )
+            recurrent_state[0] = self.handle_recurrent_state(
+                recurrent_state[0],
+                recurrent_state_,
+                cu_seqlens,
+                cu_seqlen_all[0],
+                reverse_indices,
+            )
+
+        o = o.squeeze(0).contiguous()
+        o = pad_input(o, indices_q, batch_size*self.num_memories, max_len)
+        o = rearrange(o, '(e b) l h d -> e b l (h d)', b=batch_size)
+        o = reconstruct(o, indices=indices, sorted_indices=sorted_indices, batch_size=batch_size,
+                        seq_len=seq_len, topk=self.topk, routing_weights=routing_weights, mask=mask)
+        o = rearrange(o, 'b l (h d) -> b l h d', h=self.num_heads)
+
+        if self.shared_mem:
+            shared_o = self.shared_o(shared_hidden_states, attention_mask, recurrent_state,
+                                     use_cache, conv_state_q, conv_state_k, conv_state_v)
+            o += shared_o
+
+        update_layer_cache(
+            self,
+            past_key_values,
+            recurrent_state=recurrent_state,
+            conv_state=(conv_state_q, conv_state_k, conv_state_v) if self.use_short_conv else None,
+            offset=q.shape[2],
+        )
+
+        if self.use_output_gate:
+            g = rearrange(self.g_proj(shared_hidden_states), '... (h d) -> ... h d', d=self.head_v_dim)
+            o = self.o_norm(o, g)
+        else:
+            o = self.o_norm(o)
+        o = rearrange(o, 'b t h d -> b t (h d)')
+        o = self.o_proj(o)
+
+        if origin_cu_seqlens is not None:
+            indices, _, _ = get_unpad_data(attention_mask[:, -seq_len:])
+            o = index_first_axis(rearrange(o, "b s ... -> (b s) ..."), indices).unsqueeze(0)
+
+        return o, None, past_key_values, router_logits.view(-1, self.num_memories)
+
+    def shared_o(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+        recurrent_state=None,
+        use_cache: bool | None = False,
+        conv_state_q=[None, None],
+        conv_state_k=[None, None],
+        conv_state_v=[None, None],
+        **kwargs,
+    ) -> torch.Tensor:
+        if attention_mask is not None:
+            assert len(attention_mask.shape) == 2, (
+                "Expected attention_mask as a 0-1 matrix with shape [batch_size, seq_len] "
+                "for padding purposes (0 indicating padding). "
+                "Arbitrary attention masks of shape [batch_size, seq_len, seq_len] are not allowed."
+            )
+
+        mode = 'fused_recurrent' if hidden_states.shape[1] <= 64 else self.mode
+        if self.training:
+            assert mode == 'chunk', "Only chunk mode is supported in training."
+
+        cu_seqlens = None
+        if attention_mask is not None:
+            batch_size, q_len = hidden_states.shape[0], hidden_states.shape[1]
+            indices, cu_seqlens, _ = get_unpad_data(attention_mask[:, -q_len:])
+            hidden_states = index_first_axis(rearrange(hidden_states, "b s ... -> (b s) ..."), indices).unsqueeze(0)
+
+        if self.use_short_conv:
+            q, conv_state_q[1] = self.q_conv1d(
+                x=self.q_proj(hidden_states),
+                cache=conv_state_q[1],
+                output_final_state=use_cache,
+                cu_seqlens=cu_seqlens,
+            )
+            k, conv_state_k[1] = self.k_conv1d(
+                x=self.shared_k(hidden_states),
+                cache=conv_state_k[1],
+                output_final_state=use_cache,
+                cu_seqlens=cu_seqlens,
+            )
+            v, conv_state_v[1] = self.v_conv1d(
+                x=self.shared_v(hidden_states),
+                cache=conv_state_v[1],
+                output_final_state=use_cache,
+                cu_seqlens=cu_seqlens,
+            )
+        else:
+            q = self.silu(self.q_proj(hidden_states))
+            k = self.silu(self.shared_k(hidden_states))
+            v = self.silu(self.shared_v(hidden_states))
+
+        q, k, v = map(lambda x: rearrange(x, 'b t (h d) -> b t h d', h=self.num_heads), (q, k, v))
+        beta = self.shared_b(hidden_states).sigmoid()
+        g = -self.A_log.float().exp() * F.softplus(self.shared_a(hidden_states).float() + self.dt_bias)
+
+        if mode == 'chunk':
+            o, recurrent_state[-1] = chunk_gated_delta_rule(
+                q=q,
+                k=k,
+                v=v,
+                g=g,
+                beta=beta,
+                initial_state=recurrent_state[-1],
+                output_final_state=use_cache,
+                cu_seqlens=cu_seqlens,
+                use_qk_l2norm_in_kernel=True,
+            )
+        elif mode == 'fused_recurrent':
+            o, recurrent_state[-1] = fused_recurrent_gated_delta_rule(
+                q=q,
+                k=k,
+                v=v,
+                g=g,
+                beta=beta,
+                initial_state=recurrent_state[-1],
+                output_final_state=use_cache,
+                cu_seqlens=cu_seqlens,
+                use_qk_l2norm_in_kernel=True,
+            )
+        else:
+            raise NotImplementedError(f"Not supported mode `{mode}`.")
+
+        if attention_mask is not None:
+            o = pad_input(o.squeeze(0), indices, batch_size, q_len)
+        return o
+
+    def cu2pad(self, x, cu_seqlens):
+        batch_size = cu_seqlens.shape[0] - 1
+        max_len = (cu_seqlens[1:] - cu_seqlens[:-1]).max().item()
+        indices = torch.tensor([], dtype=torch.long, device=x.device)
+        attention_mask = torch.ones((batch_size, max_len), dtype=torch.bool, device=x.device)
+        for i in range(batch_size):
+            seq_len = cu_seqlens[i+1] - cu_seqlens[i]
+            pad_len = max_len - seq_len
+            batch_indices = torch.arange(pad_len, max_len, device=x.device)
+            batch_indices = batch_indices + i * max_len
+            indices = torch.cat([indices, batch_indices])
+            attention_mask[i, :pad_len] = False
+        x = pad_input(x.squeeze(0), indices, batch_size, max_len)
+        return x, attention_mask
+
+    def pad_for_conv(self, cu_seqlens, cu_q, cu_k, cu_v):
+        lengths = cu_seqlens[1:] - cu_seqlens[:-1]
+        pad_lengths = torch.clamp(self.conv_size - lengths, min=0)
+        new_lengths = lengths + pad_lengths
+        new_cu_seqlens = torch.cat([
+            torch.tensor([0], device=cu_seqlens.device, dtype=cu_seqlens.dtype),
+            torch.cumsum(new_lengths, dim=0),
+        ])
+        final_total_len = new_cu_seqlens[-1].item()
+        new_q = torch.zeros((1, final_total_len, cu_q.shape[-1]), dtype=cu_q.dtype, device=cu_q.device)
+        new_k = torch.zeros((1, final_total_len, cu_k.shape[-1]), dtype=cu_k.dtype, device=cu_k.device)
+        new_v = torch.zeros((1, final_total_len, cu_v.shape[-1]), dtype=cu_v.dtype, device=cu_v.device)
+        num_sequences = len(lengths)
+        for i in range(num_sequences):
+            src_start = cu_seqlens[i]
+            src_end = cu_seqlens[i+1]
+            dest_start = new_cu_seqlens[i] + pad_lengths[i]
+            dest_end = new_cu_seqlens[i+1]
+            new_q[:, dest_start:dest_end, ...] = cu_q[:, src_start:src_end, ...]
+            new_k[:, dest_start:dest_end, ...] = cu_k[:, src_start:src_end, ...]
+            new_v[:, dest_start:dest_end, ...] = cu_v[:, src_start:src_end, ...]
+
+        return new_cu_seqlens, new_q, new_k, new_v, pad_lengths
+
+    def unpad_after_conv(self, conv_cu_seqlens, cu_seqlens, cu_q, cu_k, cu_v, pad_lengths):
+        original_total_len = cu_seqlens[-1].item()
+        orig_q = torch.empty((1, original_total_len, cu_q.shape[-1]), dtype=cu_q.dtype, device=cu_q.device)
+        orig_k = torch.empty((1, original_total_len, cu_k.shape[-1]), dtype=cu_k.dtype, device=cu_k.device)
+        orig_v = torch.empty((1, original_total_len, cu_v.shape[-1]), dtype=cu_v.dtype, device=cu_v.device)
+
+        num_sequences = len(pad_lengths)
+        for i in range(num_sequences):
+            dest_start = cu_seqlens[i]
+            dest_end = cu_seqlens[i+1]
+            src_start = conv_cu_seqlens[i] + pad_lengths[i]
+            src_end = conv_cu_seqlens[i+1]
+
+            orig_q[:, dest_start:dest_end, ...] = cu_q[:, src_start:src_end, ...]
+            orig_k[:, dest_start:dest_end, ...] = cu_k[:, src_start:src_end, ...]
+            orig_v[:, dest_start:dest_end, ...] = cu_v[:, src_start:src_end, ...]
+        return orig_q, orig_k, orig_v
+
+    def prepare_recurrent_state(self, recurrent_state, cu_seqlens, cu_seqlen_all, reverse_indices, batch_size):
+        if recurrent_state is None:
+            return None
+
+        if cu_seqlens is None:
+            return recurrent_state
+
+        total_len = len(cu_seqlen_all)
+        if len(cu_seqlens) != total_len:
+            # select memories that are activated
+            memories = torch.zeros_like(recurrent_state[:self.topk*batch_size])
+            mem_id = 0
+            for i in range(total_len-1):
+                if cu_seqlen_all[i] != cu_seqlen_all[i+1]:
+                    memories[mem_id] = recurrent_state[i]
+                    mem_id += 1
+            assert mem_id == self.topk * batch_size, f"The number of memories {mem_id} is not correct."
+        else:
+            memories = recurrent_state
+
+        return memories
+
+    def handle_recurrent_state(self, recurrent_state, recurrent_state_new, cu_seqlens, cu_seqlen_all, reverse_indices):
+        if recurrent_state_new is None:
+            return None
+        if cu_seqlens is None:
+            return recurrent_state_new
+        if recurrent_state is None:
+            recurrent_state = torch.zeros_like(recurrent_state_new[reverse_indices[1:]-1])
+        total_len = len(cu_seqlen_all)
+        if len(cu_seqlens) != total_len:
+            for i in range(total_len-1):
+                if cu_seqlen_all[i] != cu_seqlen_all[i+1]:
+                    recurrent_state[i] = recurrent_state_new[reverse_indices[i+1]-1]
+        else:
+            recurrent_state = recurrent_state_new
+        return recurrent_state

fla/layers/multiscale_retention.py

@@ -0,0 +1,303 @@
+# 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, repeat
+from transformers.activations import ACT2FN
+
+from fla.layers.utils import get_layer_cache, get_unpad_data, index_first_axis, pad_input, update_layer_cache
+from fla.modules import FusedRMSNormGated, RMSNorm, ShortConvolution
+from fla.modules.rotary import RotaryEmbedding
+from fla.ops.retention import chunk_retention, fused_chunk_retention, fused_recurrent_retention, parallel_retention
+from fla.ops.utils.index import prepare_lens_from_mask
+
+if TYPE_CHECKING:
+    from transformers.processing_utils import Unpack
+
+    from fla.models.utils import Cache
+
+
+class MultiScaleRetention(nn.Module):
+    r"""
+    The layer implementaion for [Retentive Network: A Successor to Transformer for Large Language Models](https://arxiv.org/pdf/2307.08621.pdf).  # noqa
+
+    Args:
+        mode (str, Optional):
+            Which Retention kernel to use.
+            Currently available: `chunk`, `fused_recurrent`, `parallel`, and `fused_chunk`.
+            Default: `chunk`.
+        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: 2.0.
+        num_heads (int, Optional):
+            The number of heads. Default: 8.
+        num_kv_heads (int, Optional):
+            The number of key/value heads, used for MQA. Default: None.
+        feature_map (str, Optional):
+            Feature map function applied to queries/keys. Default: None.
+        use_short_conv (bool, Optional):
+            Whether to use short convolutions. Default: `False`.
+        conv_size (int, Optional):
+            The kernel size of the short convolution, only used when `use_short_conv` is `True`. Default: 4.
+        conv_bias (bool, Optional):
+            Whether to use bias in the short convolution, only used when `use_short_conv` is `True`. Default: `False`.
+        use_output_gate (bool, Optional):
+            Whether to use output gate. Default: `True`.
+        gate_fn (str, Optional):
+            The activation function for the output gate. Default: `swish`.
+        elementwise_affine (bool, Optional):
+            If `True`, applies elementwise affine to LayerNorm with learnable parameters. Default: `True`.
+        norm_eps (float, Optional):
+            The epsilon value for the layernorm/rmsnorm layer. Default: 1e-5.
+        fuse_norm (bool, Optional):
+            Whether to fuse the norm and the output gate for better memory footprint. Default: `True`.
+        layer_idx (int, Optional):
+            The index of the layer. Default: None.
+    """
+
+    def __init__(
+        self,
+        mode: str = 'chunk',
+        hidden_size: int = 1024,
+        expand_k: float = 1.0,
+        expand_v: float = 2.0,
+        num_heads: int = 8,
+        num_kv_heads: int | None = None,
+        feature_map: str | None = None,
+        use_short_conv: bool = False,
+        conv_size: int = 4,
+        conv_bias: bool = False,
+        use_output_gate: bool = True,
+        gate_fn: str = 'swish',
+        elementwise_affine: bool | None = True,
+        norm_eps: float = 1e-5,
+        fuse_norm: bool = True,
+        layer_idx: int = None,
+        **kwargs,
+    ) -> MultiScaleRetention:
+        super().__init__()
+
+        self.mode = mode
+        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.feature_map_fn = ACT2FN[feature_map] if feature_map is not None else None
+
+        self.use_short_conv = use_short_conv
+        self.conv_size = conv_size
+        self.conv_bias = conv_bias
+        self.use_output_gate = use_output_gate
+
+        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
+        self.layer_idx = layer_idx
+
+        assert mode in ['chunk', 'fused_chunk', 'parallel', 'fused_recurrent'], f"Not supported mode `{mode}`."
+        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
+
+        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)
+        if self.use_output_gate:
+            self.g_proj = nn.Linear(hidden_size, self.value_dim, bias=False)
+
+        if use_short_conv:
+            self.conv_size = conv_size
+            self.q_conv1d = ShortConvolution(
+                hidden_size=self.key_dim,
+                kernel_size=conv_size,
+                bias=conv_bias,
+                activation='silu',
+            )
+            self.k_conv1d = ShortConvolution(
+                hidden_size=self.key_dim_per_group,
+                kernel_size=conv_size,
+                bias=conv_bias,
+                activation='silu',
+            )
+            self.v_conv1d = ShortConvolution(
+                hidden_size=self.value_dim_per_group,
+                kernel_size=conv_size,
+                bias=conv_bias,
+                activation='silu',
+            )
+
+        self.o_proj = nn.Linear(self.value_dim, hidden_size, bias=False)
+
+        if gate_fn == 'swish' and fuse_norm and use_output_gate:
+            self.g_norm_swish_gate = FusedRMSNormGated(
+                hidden_size=self.head_v_dim,
+                elementwise_affine=elementwise_affine,
+                eps=norm_eps,
+            )
+            self.fuse_norm_and_gate = True
+        else:
+            self.fuse_norm_and_gate = False
+            self.g_norm = RMSNorm(
+                hidden_size=self.head_v_dim,
+                elementwise_affine=elementwise_affine,
+                eps=norm_eps,
+                dtype=torch.float32
+            )
+            self.gate_fn = ACT2FN[gate_fn]
+
+        # TODO: fix this issue
+        # https://github.com/Dao-AILab/flash-attention/blob/main/flash_attn/ops/triton/rotary.py#L180
+        # Ideally, we would want to support arbitrary d_head_qk
+        assert self.head_k_dim <= 256, "head_k_dim must be less than or equal to 256"
+        self.rotary = RotaryEmbedding(dim=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]:
+        if attention_mask is not None:
+            assert len(attention_mask.shape) == 2, (
+                "Expected attention_mask as a 0-1 matrix with shape [batch_size, seq_len] "
+                "for padding purposes (0 indicating padding). "
+                "Arbitrary attention masks of shape [batch_size, seq_len, seq_len] are not allowed."
+            )
+
+        batch_size, q_len, _ = hidden_states.shape
+        mode = 'fused_recurrent' if q_len <= 64 else self.mode
+
+        last_state = get_layer_cache(self, past_key_values)
+
+        cu_seqlens = kwargs.get('cu_seqlens')
+        if attention_mask is not None:
+            indices, cu_seqlens, _ = get_unpad_data(attention_mask[:, -q_len:])
+            hidden_states = index_first_axis(rearrange(hidden_states, "b s ... -> (b s) ..."), indices).unsqueeze(0)
+
+        if self.use_short_conv:
+            conv_state_q, conv_state_k, conv_state_v = None, None, None
+            if last_state is not None:
+                conv_state_q, conv_state_k, conv_state_v = last_state['conv_state']
+            q, conv_state_q = self.q_conv1d(
+                x=self.q_proj(hidden_states),
+                cache=conv_state_q,
+                output_final_state=use_cache,
+                cu_seqlens=cu_seqlens,
+            )
+            k, conv_state_k = self.k_conv1d(
+                x=self.k_proj(hidden_states),
+                cache=conv_state_k,
+                output_final_state=use_cache,
+                cu_seqlens=cu_seqlens,
+            )
+            v, conv_state_v = self.v_conv1d(
+                x=self.v_proj(hidden_states),
+                cache=conv_state_v,
+                output_final_state=use_cache,
+                cu_seqlens=cu_seqlens,
+            )
+        else:
+            q = self.q_proj(hidden_states)
+            k = self.k_proj(hidden_states)
+            v = self.v_proj(hidden_states)
+
+        q = rearrange(q, '... (h d) -> ... h d', d=self.head_k_dim)
+        k = rearrange(k, '... (h d) -> ... h d', d=self.head_k_dim)
+        if self.feature_map_fn is not None:
+            q, k = map(self.feature_map_fn, (q, k))
+
+        seqlen_offset, max_seqlen = 0, q.shape[1]
+        if past_key_values is not None:
+            seqlen_offset = past_key_values.get_seq_length(self.layer_idx)
+            max_seqlen = q.shape[1] + seqlen_offset
+
+            if attention_mask is not None and seqlen_offset > 0:
+                # to deliminate the offsets of padding tokens
+                seqlen_offset = prepare_lens_from_mask(attention_mask) - q_len
+                max_seqlen = q.shape[1] + seqlen_offset.max().item()
+
+        q, k = self.rotary(q, k, seqlen_offset=seqlen_offset, max_seqlen=max_seqlen, cu_seqlens=cu_seqlens)
+
+        if self.num_kv_groups > 1:
+            k = repeat(k, '... h d -> ... (h g) d', g=self.num_kv_groups)
+            v = repeat(v, '... (h d) -> ... (h g) d', d=self.head_v_dim, g=self.num_kv_groups)
+        else:
+            v = rearrange(v, '... (h d) -> ... h d', d=self.head_v_dim)
+
+        recurrent_state = last_state['recurrent_state'] if last_state is not None else None
+        if mode == 'chunk':
+            o, recurrent_state = chunk_retention(
+                q=q,
+                k=k,
+                v=v,
+                initial_state=recurrent_state,
+                output_final_state=use_cache,
+                cu_seqlens=cu_seqlens,
+            )
+        elif mode == 'fused_chunk':
+            o, recurrent_state = fused_chunk_retention(
+                q=q,
+                k=k,
+                v=v,
+                initial_state=recurrent_state,
+                output_final_state=use_cache,
+                cu_seqlens=cu_seqlens,
+            )
+        elif mode == 'parallel':
+            o, recurrent_state = parallel_retention(
+                q=q,
+                k=k,
+                v=v,
+                cu_seqlens=cu_seqlens,
+            )
+        elif mode == 'fused_recurrent':
+            o, recurrent_state = fused_recurrent_retention(
+                q=q,
+                k=k,
+                v=v,
+                initial_state=recurrent_state,
+                output_final_state=use_cache,
+                cu_seqlens=cu_seqlens,
+            )
+        else:
+            raise NotImplementedError(f"Not supported mode `{mode}`.")
+
+        update_layer_cache(
+            self,
+            past_key_values,
+            recurrent_state=recurrent_state,
+            conv_state=(conv_state_q, conv_state_k, conv_state_v) if self.use_short_conv else None,
+            offset=q_len,
+        )
+
+        if self.use_output_gate:
+            g = self.g_proj(hidden_states)
+            if self.fuse_norm_and_gate:
+                g = rearrange(g, '... (h d) -> ... h d', d=self.head_v_dim)
+                o = self.g_norm_swish_gate(o, g)
+                o = rearrange(o, '... h d -> ... (h d)')
+            else:
+                o = rearrange(self.g_norm(o), '... h d -> ... (h d)')
+                o = o * self.gate_fn(g)
+        else:
+            o = rearrange(self.g_norm(o), '... h d -> ... (h d)')
+        o = self.o_proj(o)
+        if attention_mask is not None:
+            o = pad_input(o.squeeze(0), indices, batch_size, q_len)
+
+        return o, None, past_key_values

fla/layers/nsa.py

@@ -0,0 +1,137 @@
+# 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 transformers.utils import logging
+
+from fla.modules import RotaryEmbedding
+from fla.ops.nsa.parallel import parallel_nsa
+from fla.ops.utils.index import prepare_lens_from_mask
+
+if TYPE_CHECKING:
+    from fla.models.utils import Cache
+
+logger = logging.get_logger(__name__)
+
+
+class NativeSparseAttention(nn.Module):
+
+    def __init__(
+        self,
+        hidden_size: int = 2048,
+        num_heads: int = 64,
+        num_kv_heads: int | None = 4,
+        head_dim: int = 64,
+        qkv_bias: bool = False,
+        block_size: int | None = 64,
+        block_counts: torch.LongTensor | int | None = 16,
+        window_size: int | None = 512,
+        rope_theta: float | None = 10000.,
+        max_position_embeddings: int | None = None,
+        layer_idx: int = None,
+    ):
+        super().__init__()
+
+        self.hidden_size = hidden_size
+        self.num_heads = num_heads
+        if num_kv_heads is None:
+            self.num_kv_heads = self.num_heads
+        else:
+            self.num_kv_heads = num_kv_heads
+        self.num_kv_groups = num_heads // self.num_kv_heads
+        self.head_dim = head_dim
+        self.kv_dim = self.num_kv_heads * self.head_dim
+        self.qkv_bias = qkv_bias
+
+        self.block_size = block_size
+        self.block_counts = block_counts
+        self.window_size = window_size
+        self.rope_theta = rope_theta
+        self.max_position_embeddings = max_position_embeddings
+        self.layer_idx = layer_idx
+
+        self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=self.qkv_bias)
+        self.k_proj = nn.Linear(self.hidden_size, self.kv_dim, bias=self.qkv_bias)
+        self.v_proj = nn.Linear(self.hidden_size, self.kv_dim, bias=self.qkv_bias)
+        self.g_proj = nn.Linear(self.hidden_size, self.num_heads * 3, bias=False)
+        self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=False)
+
+        self.rotary = RotaryEmbedding(dim=self.head_dim, base=self.rope_theta)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: torch.LongTensor | None = None,
+        past_key_values: Cache | None = None,
+        output_attentions: bool = False,
+        use_cache: bool = False,
+        **kwargs,
+    ) -> tuple[torch.Tensor, torch.Tensor | None, tuple[torch.Tensor] | None]:
+        if attention_mask is not None:
+            assert len(attention_mask.shape) == 2, (
+                "Expected attention_mask as a 0-1 matrix with shape [batch_size, seq_len] "
+                "for padding purposes (0 indicating padding). "
+                "Arbitrary attention masks of shape [batch_size, seq_len, seq_len] are not allowed."
+            )
+
+        batch_size, seq_len, _ = hidden_states.size()
+
+        q = rearrange(self.q_proj(hidden_states), '... (h d) -> ... h d', d=self.head_dim)
+        k = rearrange(self.k_proj(hidden_states), '... (h d) -> ... h d', d=self.head_dim)
+        v = rearrange(self.v_proj(hidden_states), '... (h d) -> ... h d', d=self.head_dim)
+        g = rearrange(self.g_proj(hidden_states), '... (h d) -> ... h d', d=3)
+        g_cmp, g_slc, g_swa = g.sigmoid().unbind(-1)
+
+        cu_seqlens = kwargs.get('cu_seqlens')
+
+        seqlen_offset, max_seqlen = 0, seq_len
+        if past_key_values is not None:
+            seqlen_offset = past_key_values.get_seq_length(self.layer_idx)
+            max_seqlen = q.shape[1] + seqlen_offset
+
+            if attention_mask is not None:
+                # to deliminate the offsets of padding tokens
+                seqlen_offset = seqlen_offset + prepare_lens_from_mask(attention_mask) - attention_mask.shape[-1]
+                max_seqlen = q.shape[1] + max(seqlen_offset)
+
+        if self.max_position_embeddings is not None:
+            max_seqlen = max(max_seqlen, self.max_position_embeddings)
+        q, k = self.rotary(q, k, seqlen_offset=seqlen_offset, max_seqlen=max_seqlen, cu_seqlens=cu_seqlens)
+
+        if past_key_values is not None:
+            cache_has_content = past_key_values.get_seq_length(self.layer_idx) > 0
+            k_cached, v_cached = past_key_values.update(
+                attn_state=(k.flatten(-2, -1), v.flatten(-2, -1)),
+                layer_idx=self.layer_idx,
+                offset=seq_len,
+                cache_kwargs=dict(window_size=self.window_size),
+            )['attn_state']
+            if cache_has_content:
+                k, v = k_cached, v_cached
+                k = rearrange(k, '... (h d) -> ... h d', d=self.head_dim)
+                v = rearrange(v, '... (h d) -> ... h d', d=self.head_dim)
+
+        o = parallel_nsa(
+            q=q,
+            k=k,
+            v=v,
+            g_cmp=g_cmp,
+            g_slc=g_slc,
+            g_swa=g_swa,
+            block_size=self.block_size,
+            block_counts=self.block_counts,
+            window_size=self.window_size,
+            cu_seqlens=cu_seqlens,
+        )
+        o = o.reshape(batch_size, seq_len, -1)
+        o = self.o_proj(o)
+
+        if not output_attentions:
+            attentions = None
+
+        return o, attentions, past_key_values

fla/layers/path_attn.py

@@ -0,0 +1,216 @@
+# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange
+from transformers.utils import logging
+
+from fla.layers.utils import pad_input, unpad_input
+from fla.modules import RMSNorm, ShortConvolution
+from fla.modules.l2norm import l2_norm
+from fla.ops.attn.decoding import attn_decoding_one_step
+from fla.ops.path_attn.parallel import parallel_path_attn
+
+if TYPE_CHECKING:
+    from fla.models.utils import Cache
+
+logger = logging.get_logger(__name__)
+
+
+class PaTHAttention(nn.Module):
+    def __init__(
+        self,
+        hidden_size: int = 2048,
+        num_heads: int = 32,
+        num_kv_heads: int | None = None,
+        use_forget_gate: bool = False,
+        use_qk_norm: bool = False,
+        layer_idx: int = None,
+        use_low_rank_w: bool = True,
+        use_w_shortconv: bool = True,
+        conv_size: int = 3,
+        conv_bias: bool = False,
+    ):
+        super().__init__()
+
+        self.hidden_size = hidden_size
+        self.num_heads = num_heads
+        if num_kv_heads is None:
+            self.num_kv_heads = self.num_heads
+        else:
+            self.num_kv_heads = num_kv_heads
+        self.head_dim = self.hidden_size // self.num_heads
+        self.kv_dim = self.num_kv_heads * self.head_dim
+
+        self.layer_idx = layer_idx
+
+        self.q_proj = nn.Linear(self.hidden_size, self.hidden_size, bias=False)
+        self.k_proj = nn.Linear(self.hidden_size, self.kv_dim, bias=False)
+        self.v_proj = nn.Linear(self.hidden_size, self.kv_dim, bias=False)
+
+        # We use low-rank parameterization for the w_proj to reduce parameters in MHA settings.
+        if use_low_rank_w:
+            self.w_proj = nn.Sequential(
+                nn.Linear(self.hidden_size, 32, bias=False),
+                nn.Linear(32, self.kv_dim, bias=False),
+            )
+        # In MQA/GQA settings, key/value heads are shared, so we use a standard linear projection
+        # which doesn't introduce too many parameters
+        else:
+            self.w_proj = nn.Linear(self.hidden_size, self.kv_dim, bias=False)
+
+        # per head norm
+        if use_qk_norm:
+            self.maybe_q_norm = RMSNorm(self.head_dim, dtype=torch.float32)
+            self.maybe_k_norm = RMSNorm(self.head_dim, dtype=torch.float32)
+        else:
+            self.maybe_q_norm = nn.Identity()
+            self.maybe_k_norm = nn.Identity()
+
+        if use_w_shortconv:
+            self.w_conv1d = ShortConvolution(hidden_size=self.kv_dim, kernel_size=conv_size, bias=conv_bias, activation='silu')
+        self.use_w_shortconv = use_w_shortconv
+        self.bt_proj = nn.Linear(self.hidden_size, self.num_kv_heads, bias=True)
+        self.use_forget_gate = use_forget_gate
+        if use_forget_gate:
+            self.g_proj = nn.Linear(self.hidden_size, self.num_heads, bias=True)
+        self.o_proj = nn.Linear(self.hidden_size, self.hidden_size, bias=False)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: torch.LongTensor | None = None,
+        past_key_values: Cache | None = None,
+        output_attentions: bool = False,
+        use_cache: bool = False,
+        **kwargs,
+    ) -> tuple[torch.Tensor, torch.Tensor | None, tuple[torch.Tensor] | None]:
+        if use_cache:
+            assert past_key_values is not None, "past_key_values must be provided when use_cache is True"
+        if attention_mask is not None:
+            assert len(attention_mask.shape) == 2, (
+                "Expected attention_mask as a 0-1 matrix with shape [batch_size, seq_len] "
+                "for padding purposes (0 indicating padding). "
+                "Arbitrary attention masks of shape [batch_size, seq_len, seq_len] are not allowed."
+            )
+        batch_size, q_len, _ = hidden_states.size()
+        q = self.q_proj(hidden_states)
+        k = self.k_proj(hidden_states)
+        v = self.v_proj(hidden_states)
+        w = self.w_proj(hidden_states)
+        beta = self.bt_proj(hidden_states).float().sigmoid() * 2  # allowing negative eigenvalues
+        g = F.logsigmoid(self.g_proj(hidden_states).float()) if self.use_forget_gate else None
+        cu_seqlens = kwargs.get('cu_seqlens')
+        assert not (cu_seqlens is not None and attention_mask is not None), (
+            "cu_seqlens should not be provided when attention_mask is not None"
+        )
+        # Training
+        if attention_mask is None:
+            assert use_cache is False, "use_cache should be False in training"
+            if self.use_w_shortconv:
+                w, _ = self.w_conv1d(w, cache=None, output_final_state=False, cu_seqlens=cu_seqlens)
+            q = rearrange(q, '... (h d) -> ... h d', d=self.head_dim)
+            k = rearrange(k, '... (h d) -> ... h d', d=self.head_dim)
+            v = rearrange(v, '... (h d) -> ... h d', d=self.head_dim)
+            w = rearrange(w, '... (h d) -> ... h d', d=self.head_dim)
+            q, k = self.maybe_q_norm(q), self.maybe_k_norm(k)
+            w = l2_norm(w, output_dtype=torch.float32)
+            o, _ = parallel_path_attn(q=q, k=k, v=v, w=w, beta=beta, g=g, cu_seqlens=cu_seqlens)
+
+        # Prefilling or decoding
+        else:
+            assert self.training is False, "attention mask is not supported in training. Please use variable length input."
+            try:
+                last_state = past_key_values[self.layer_idx]
+            except KeyError:
+                last_state = None
+            # Decoding
+            if last_state is not None:
+                if g is not None:
+                    past_k, past_v, past_g = last_state['attn_state']
+                else:
+                    past_k, past_v = last_state['attn_state']
+                    past_g = None
+                w_conv_state = last_state['conv_state']
+                past_k = rearrange(past_k, '... (h d) -> ... h d', d=self.head_dim)
+                if self.use_w_shortconv:
+                    w, w_conv_state = self.w_conv1d(w, cache=w_conv_state, output_final_state=use_cache, cu_seqlens=cu_seqlens)
+                w = rearrange(w, '... (h d) -> ... h d', d=self.head_dim)
+                w = l2_norm(w, output_dtype=torch.float32)
+
+                @torch.compile
+                def rank_one_update(k, w, beta):
+                    original_dtype = k.dtype
+                    k = k.float()
+                    w = w.float()
+                    beta = beta.float()
+                    k = k - beta[..., None].float() * (k * w).sum(-1, keepdim=True) * w
+                    return k.to(original_dtype)
+
+                past_k = rank_one_update(past_k, w, beta)
+                past_k = rearrange(past_k, '... h d -> ... (h d)')
+                k = torch.cat([past_k, k], dim=1)
+                v = torch.cat([past_v, v], dim=1)
+                g = torch.cat([past_g, g], dim=1) if g is not None else None
+                past_key_values[self.layer_idx]['attn_state'] = (k, v, g) if g is not None else (k, v)
+                past_key_values.update(
+                    conv_state=w_conv_state,
+                    layer_idx=self.layer_idx,
+                    offset=q_len,
+                )
+                if g is not None:
+                    q, (k, v, g), indices_q, cu_seqlens, max_seq_lens = unpad_input(
+                        q, (k, v, g), attention_mask, q_len, keepdim=True)
+                    max_seqlen_q, max_seqlen_k = max_seq_lens
+                else:
+                    q, (k, v), indices_q, cu_seqlens, max_seq_lens = unpad_input(
+                        q, (k, v), attention_mask, q_len, keepdim=True)
+                    max_seqlen_q, max_seqlen_k = max_seq_lens
+                _, cu_seqlens = cu_seqlens
+                q = rearrange(q, '... (h d) -> ... h d', d=self.head_dim)
+                k = rearrange(k, '... (h d) -> ... h d', d=self.head_dim)
+                v = rearrange(v, '... (h d) -> ... h d', d=self.head_dim)
+                assert max_seqlen_q == 1, "only support q_len == 1 for decoding"
+                o = attn_decoding_one_step(q, k, v, g, cu_seqlens=cu_seqlens, do_gate_scale=True)  # reduced to fox's decoding
+            # Prefilling
+            else:
+                v_cache = v.clone()
+                g_cache = g.clone() if g is not None else None
+                if g is None:
+                    q, (k, v, w, beta), indices_q, cu_seqlens, max_seq_lens = unpad_input(
+                        q, (k, v, w, beta), attention_mask, q_len, keepdim=True)
+                else:
+                    q, (k, v, w, beta, g), indices_q, cu_seqlens, max_seq_lens = unpad_input(
+                        q, (k, v, w, beta, g), attention_mask, q_len, keepdim=True)
+                max_seqlen_q, max_seqlen_k = max_seq_lens
+                assert max_seqlen_q == max_seqlen_k, "max_seqlen_q should be equal to max_seqlen_k in prefilling"
+                _, cu_seqlens = cu_seqlens
+                if self.use_w_shortconv:
+                    w, w_conv_state = self.w_conv1d(w, cache=None, output_final_state=use_cache, cu_seqlens=cu_seqlens)
+                else:
+                    w_conv_state = None
+                q = rearrange(q, '... (h d) -> ... h d', d=self.head_dim)
+                k = rearrange(k, '... (h d) -> ... h d', d=self.head_dim)
+                v = rearrange(v, '... (h d) -> ... h d', d=self.head_dim)
+                w = rearrange(w, '... (h d) -> ... h d', d=self.head_dim)
+                w = l2_norm(w, output_dtype=torch.float32)
+                o, k_cache = parallel_path_attn(q=q, k=k, v=v, w=w, beta=beta, g=g,
+                                                cu_seqlens=cu_seqlens, use_cache=use_cache)
+                if use_cache:
+                    k_cache = pad_input(k_cache.squeeze(0), indices_q, batch_size, q_len)
+                    k_cache = rearrange(k_cache, '... h d -> ... (h d)')
+                    past_key_values.update(
+                        attn_state=(k_cache, v_cache, g_cache) if g_cache is not None else (k_cache, v_cache),
+                        conv_state=w_conv_state,
+                        layer_idx=self.layer_idx,
+                        offset=q_len,
+                    )
+            o = pad_input(o.squeeze(0), indices_q, batch_size, q_len)
+        o = rearrange(o, '... h d -> ... (h d)')
+        o = self.o_proj(o)
+        return o, None, past_key_values

fla/layers/quasar.py

@@ -0,0 +1,439 @@
+# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang
+# Modified for QuasarAttention
+
+from __future__ import annotations
+
+import contextlib
+import math
+import os
+from typing import TYPE_CHECKING
+
+import torch
+import torch.nn as nn
+from einops import rearrange, repeat
+from torch.nn import functional as F
+
+from fla.layers.utils import get_unpad_data, index_first_axis, pad_input
+
+
+def _quasar_debug_tensor(name: str, tensor: torch.Tensor, layer_idx: int | None) -> None:
+    if os.environ.get("QUASAR_DEBUG_FINITE", "0") != "1":
+        return
+    if tensor is None or torch.isfinite(tensor).all():
+        return
+    with torch.no_grad():
+        t = torch.nan_to_num(tensor.detach().float(), nan=0.0, posinf=0.0, neginf=0.0)
+        nonfinite = int((~torch.isfinite(tensor)).sum().item())
+        print(
+            f"[QUASAR DEBUG] layer={layer_idx} stage={name} nonfinite={nonfinite} "
+            f"min={float(t.min())} max={float(t.max())} mean={float(t.mean())}",
+            flush=True,
+        )
+
+
+class _TorchRMSNormGated(nn.Module):
+    def __init__(self, hidden_size: int, activation: str = "sigmoid", eps: float = 1e-5):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+        self.activation = activation
+        self.eps = eps
+
+    def reset_parameters(self) -> None:
+        self.weight.data.fill_(1.0)
+
+    def forward(self, x: torch.Tensor, gate: torch.Tensor) -> torch.Tensor:
+        dtype = x.dtype
+        y = torch.nan_to_num(
+            x.float(),
+            nan=0.0,
+            posinf=1e4,
+            neginf=-1e4,
+        ).clamp_(min=-1e4, max=1e4)
+        y = y * torch.rsqrt(y.square().mean(dim=-1, keepdim=True) + self.eps)
+        weight = torch.nan_to_num(
+            self.weight.float(),
+            nan=1.0,
+            posinf=1.0,
+            neginf=1.0,
+        ).clamp_(min=0.0, max=4.0)
+        y = y.to(dtype) * weight.to(dtype=dtype, device=x.device)
+        gate = torch.nan_to_num(
+            gate.float(),
+            nan=0.0,
+            posinf=30.0,
+            neginf=-30.0,
+        ).clamp_(min=-30.0, max=30.0)
+        if self.activation in {"swish", "silu"}:
+            gate = gate * torch.sigmoid(gate)
+        elif self.activation == "sigmoid":
+            gate = torch.sigmoid(gate)
+        return y * gate.to(dtype=dtype, device=x.device)
+
+
+def rotate_half(x):
+    """Rotates half the hidden dims of the input."""
+    x1 = x[..., : x.shape[-1] // 2]
+    x2 = x[..., x.shape[-1] // 2 :]
+    return torch.cat((-x2, x1), dim=-1)
+
+def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None):
+    """Applies Rotary Position Embedding to the query and key tensors."""
+    # cos, sin: [1, 1, seq_len, rotary_dim]
+    # q, k: [batch_size, seq_len, n_heads, head_dim]
+    rotary_dim = cos.shape[-1]
+    q_rot, q_pass = q[..., :rotary_dim], q[..., rotary_dim:]
+    k_rot, k_pass = k[..., :rotary_dim], k[..., rotary_dim:]
+    cos = cos.transpose(1, 2) # [1, seq_len, 1, rotary_dim]
+    sin = sin.transpose(1, 2) # [1, seq_len, 1, rotary_dim]
+    q_embed = (q_rot * cos) + (rotate_half(q_rot) * sin)
+    k_embed = (k_rot * cos) + (rotate_half(k_rot) * sin)
+    return torch.cat([q_embed, q_pass], dim=-1), torch.cat([k_embed, k_pass], dim=-1)
+
+if TYPE_CHECKING:
+    from transformers.processing_utils import Unpack
+
+    from fla.models.utils import Cache
+
+
+class QuasarAttention(nn.Module):
+    """
+    QuasarAttention layer implementation.
+
+    Args:
+        hidden_size (int, Optional):
+            The hidden size of the input. Default: 2048.
+        head_dim (int, Optional):
+            The dimension of each head. Default: 128.
+        num_heads (int, Optional):
+            The number of heads. Default: 16.
+        mode (str, Optional):
+            Which QuasarAttention kernel to use.
+            Currently available: `chunk` and `fused_recurrent`.
+            Default: `chunk`.
+        use_short_conv (bool, Optional):
+            Whether to use short convolutions. Default: `True`.
+        conv_size (int, Optional):
+            The kernel size of the short convolution, only used when `use_short_conv` is `True`. Default: 4.
+        conv_bias (bool, Optional):
+            Whether to use bias in the short convolution, only used when `use_short_conv` is `True`. Default: `False`.
+        layer_idx (int, Optional):
+            The index of the layer. Default: None.
+        norm_eps (float, Optional):
+            The epsilon value for the normalization layer. Default: 1e-5.
+    """
+
+    def __init__(
+        self,
+        hidden_size: int = 2048,
+        head_dim: int = 128,
+        num_heads: int = 16,
+        mode: str = "chunk",
+        use_short_conv: bool = True,
+        conv_size: int = 4,
+        conv_bias: bool = False,
+        layer_idx: int = None,
+        norm_eps: float = 1e-5,
+        **kwargs,
+    ) -> QuasarAttention:
+        super().__init__()
+
+        self.mode = mode
+        self.hidden_size = hidden_size
+
+        self.use_short_conv = use_short_conv
+        self.conv_size = conv_size
+        self.conv_bias = conv_bias
+
+        self.head_dim = head_dim
+        self.num_heads = num_heads
+        self.key_dim = int(self.num_heads * self.head_dim)
+        self.value_dim = int(self.num_heads * self.head_dim)
+        self.layer_idx = layer_idx
+
+        assert mode in ["chunk", "fused_recurrent"], f"Not supported mode `{mode}`."
+
+        self.q_proj = nn.Linear(hidden_size, self.key_dim, bias=False)
+        self.k_proj = nn.Linear(hidden_size, self.key_dim, bias=False)
+        self.v_proj = nn.Linear(hidden_size, self.value_dim, bias=False)
+
+        # KDA matching: Use SiLU on q, k, v for better learning if not using short conv
+        # (Short conv already has its own activation)
+        self.q_act = nn.SiLU()
+        self.k_act = nn.SiLU()
+        self.v_act = nn.SiLU()
+
+        if use_short_conv:
+            from fla.modules.convolution import ShortConvolution
+
+            self.q_conv1d = ShortConvolution(
+                hidden_size=self.key_dim,
+                kernel_size=conv_size,
+                bias=conv_bias,
+                activation="silu",
+            )
+            self.k_conv1d = ShortConvolution(
+                hidden_size=self.key_dim,
+                kernel_size=conv_size,
+                bias=conv_bias,
+                activation="silu",
+            )
+            self.v_conv1d = ShortConvolution(
+                hidden_size=self.value_dim,
+                kernel_size=conv_size,
+                bias=conv_bias,
+                activation="silu",
+            )
+
+        # Data-dependent Beta (Adaptive Decay)
+        # Instead of a static per-head parameter, we use a linear projection 
+        # to allow the model to learn contextual importance (read/write sharpness).
+        self.b_proj = nn.Linear(hidden_size, self.num_heads, bias=False)
+
+        # Learnable state decay (like KDA/Mamba A matrix)
+        self.A_log = nn.Parameter(torch.log(torch.empty(self.num_heads, dtype=torch.float32).uniform_(1, 16)))
+        self.A_log._no_weight_decay = True
+        self.dt_bias = nn.Parameter(torch.zeros(self.key_dim, dtype=torch.float32))
+        self.dt_bias._no_weight_decay = True
+
+        # KIMI matches: separate f_proj for kernel and g_proj for final output gating
+        self.f_proj = nn.Linear(hidden_size, self.key_dim, bias=False)
+        self.g_proj = nn.Sequential(
+            nn.Linear(hidden_size, self.head_dim, bias=False),
+            nn.Linear(self.head_dim, self.value_dim, bias=True),
+        )
+
+        self.o_norm = _TorchRMSNormGated(self.head_dim, activation="sigmoid", eps=norm_eps)
+        self.o_proj = nn.Linear(self.value_dim, hidden_size, bias=False)
+
+    def reset_parameters(self) -> None:
+        for module in self.children():
+            reset = getattr(module, "reset_parameters", None)
+            if callable(reset):
+                reset()
+
+    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]:
+        if attention_mask is not None:
+            assert len(attention_mask.shape) == 2, (
+                "Expected attention_mask as a 0-1 matrix with shape [batch_size, seq_len] "
+                "for padding purposes (0 indicating padding). "
+                "Arbitrary attention masks of shape [batch_size, seq_len, seq_len] are not allowed."
+            )
+
+        batch_size, q_len, _ = hidden_states.shape
+        mode = self.mode
+        if self.training and mode == "fused_recurrent":
+            # The fused recurrent Quasar path is forward-only in this tree.
+            # Training must use the chunk kernel until its backward exists.
+            mode = "chunk"
+
+        # Bailing hidden states can be very large after MoE/FSDP checkpoint
+        # restore. Quasar's delta-rule triangular solve is much more sensitive
+        # to projection scale than GQA/GLA, so sanitize and RMS-normalize only
+        # the Quasar branch input. The residual model path remains untouched.
+        input_dtype = hidden_states.dtype
+        hidden_states = torch.nan_to_num(
+            hidden_states.float(),
+            nan=0.0,
+            posinf=60.0,
+            neginf=-60.0,
+        ).clamp_(min=-60.0, max=60.0)
+        hidden_states = hidden_states * torch.rsqrt(
+            hidden_states.square().mean(dim=-1, keepdim=True) + 1e-6
+        )
+        hidden_states = hidden_states.to(dtype=input_dtype)
+        _quasar_debug_tensor("input_normed", hidden_states, self.layer_idx)
+
+        last_state = None
+        recurrent_state = None
+        conv_state_q, conv_state_k, conv_state_v = None, None, None
+        
+        if past_key_values is not None and self.layer_idx is not None:
+            if hasattr(past_key_values, "recurrent_states") and self.layer_idx in past_key_values.recurrent_states:
+                recurrent_state = past_key_values.recurrent_states[self.layer_idx]
+            if hasattr(past_key_values, "conv_states") and self.layer_idx in past_key_values.conv_states:
+                conv_state_q, conv_state_k, conv_state_v = past_key_values.conv_states[self.layer_idx]
+            else:
+                try:
+                    # Standard list/tuple cache (FLA style fallback)
+                    if len(past_key_values) > self.layer_idx:
+                        last_state = past_key_values[self.layer_idx]
+                        if isinstance(last_state, dict):
+                            recurrent_state = last_state.get("recurrent_state", None)
+                            convs = last_state.get("conv_state", None)
+                            if convs is not None:
+                                conv_state_q, conv_state_k, conv_state_v = convs
+                except TypeError:
+                    pass
+
+        cu_seqlens = kwargs.get("cu_seqlens")
+        if attention_mask is not None:
+            # Optimization: Skip unpadding if all tokens are valid (common in packed distillation)
+            if attention_mask.all():
+                indices, cu_seqlens = None, None
+            else:
+                indices, cu_seqlens, _ = get_unpad_data(attention_mask[:, -q_len:])
+                hidden_states = index_first_axis(rearrange(hidden_states, "b s ... -> (b s) ..."), indices).unsqueeze(0)
+        else:
+            indices = None
+
+        if self.use_short_conv:
+            q, conv_state_q = self.q_conv1d(
+                x=self.q_proj(hidden_states),
+                cache=conv_state_q,
+                output_final_state=use_cache,
+                cu_seqlens=cu_seqlens,
+            )
+            k, conv_state_k = self.k_conv1d(
+                x=self.k_proj(hidden_states),
+                cache=conv_state_k,
+                output_final_state=use_cache,
+                cu_seqlens=cu_seqlens,
+            )
+            v, conv_state_v = self.v_conv1d(
+                x=self.v_proj(hidden_states),
+                cache=conv_state_v,
+                output_final_state=use_cache,
+                cu_seqlens=cu_seqlens,
+            )
+        else:
+            q = self.q_act(self.q_proj(hidden_states))
+            k = self.k_act(self.k_proj(hidden_states))
+            v = self.v_act(self.v_proj(hidden_states))
+        _quasar_debug_tensor("q_proj", q, self.layer_idx)
+        _quasar_debug_tensor("k_proj", k, self.layer_idx)
+        _quasar_debug_tensor("v_proj", v, self.layer_idx)
+
+        q = rearrange(q, "... (h d) -> ... h d", d=self.head_dim)
+        k = rearrange(k, "... (h d) -> ... h d", d=self.head_dim)
+        v = rearrange(v, "... (h d) -> ... h d", d=self.head_dim)
+
+        # Apply RoPE if provided
+        cos = kwargs.get("cos")
+        sin = kwargs.get("sin")
+        if cos is not None and sin is not None:
+            if attention_mask is not None:
+                # Unpad cos/sin using the same indices
+                # cos/sin shape is [1, 1, seq_len, head_dim] or [batch_size, seq_len, head_dim]
+                if cos.shape[0] == 1 and cos.shape[1] == 1:
+                    # Broadcastable/Shared RoPE [1, 1, seq_len, head_dim]
+                    # We need to expand to [batch_size, seq_len, head_dim] before unpadding
+                    cos_expanded = cos.squeeze(1).expand(batch_size, -1, -1)
+                    sin_expanded = sin.squeeze(1).expand(batch_size, -1, -1)
+                    cos = index_first_axis(rearrange(cos_expanded, "b s d -> (b s) d"), indices).unsqueeze(0).unsqueeze(1)
+                    sin = index_first_axis(rearrange(sin_expanded, "b s d -> (b s) d"), indices).unsqueeze(0).unsqueeze(1)
+                else:
+                    # Already [batch_size, 1, seq_len, head_dim] or [batch_size, seq_len, head_dim]
+                    if cos.dim() == 4:
+                        cos = cos.squeeze(1)
+                        sin = sin.squeeze(1)
+                    cos = index_first_axis(rearrange(cos, "b s d -> (b s) d"), indices).unsqueeze(0).unsqueeze(1)
+                    sin = index_first_axis(rearrange(sin, "b s d -> (b s) d"), indices).unsqueeze(0).unsqueeze(1)
+            
+            q, k = apply_rotary_pos_emb(q, k, cos, sin)
+
+        # QK Normalization AFTER RoPE — ensures kernel receives unit-norm vectors
+        # regardless of any precision drift introduced by the rotation
+        q = F.normalize(q, p=2, dim=-1)
+        k = F.normalize(k, p=2, dim=-1)
+        _quasar_debug_tensor("q_norm", q, self.layer_idx)
+        _quasar_debug_tensor("k_norm", k, self.layer_idx)
+
+        # Adaptive Beta: Sigmoid(b_proj(x)) is bounded to (0, 1) to prevent explosions.
+        beta = self.b_proj(hidden_states).sigmoid()
+        _quasar_debug_tensor("beta", beta, self.layer_idx)
+
+        if mode == "chunk":
+            from fla.ops.quasar.chunk import chunk_quasar
+
+            o, recurrent_state = chunk_quasar(
+                q=q,
+                k=k,
+                v=v,
+                beta=beta,
+                A_log=self.A_log,
+                dt_bias=self.dt_bias,
+                initial_state=recurrent_state,
+                output_final_state=use_cache,
+                cu_seqlens=cu_seqlens,
+                use_qk_l2norm_in_kernel=True,
+            )
+            _quasar_debug_tensor("chunk_kernel_o", o, self.layer_idx)
+        elif mode == "fused_recurrent":
+            from fla.ops.quasar.fused_recurrent import fused_recurrent_quasar
+
+            # Use f_proj for kernel gate in fused mode
+            f_gate = self.f_proj(hidden_states)
+            f_gate = rearrange(f_gate, "... (h d) -> ... h d", d=self.head_dim)
+            o, recurrent_state = fused_recurrent_quasar(
+                q=q,
+                k=k,
+                v=v,
+                g=f_gate,
+                beta=beta,
+                A_log=self.A_log,
+                dt_bias=self.dt_bias,
+                initial_state=recurrent_state,
+                output_final_state=use_cache,
+                use_qk_l2norm_in_kernel=True,
+            )
+            _quasar_debug_tensor("fused_kernel_o", o, self.layer_idx)
+        else:
+            raise NotImplementedError(f"Not supported mode `{mode}`.")
+
+        o = torch.nan_to_num(
+            o.float(),
+            nan=0.0,
+            posinf=1e4,
+            neginf=-1e4,
+        ).clamp_(min=-1e4, max=1e4).to(dtype=v.dtype)
+        _quasar_debug_tensor("kernel_o_clamped", o, self.layer_idx)
+
+        if past_key_values is not None:
+            if hasattr(past_key_values, "update_quasar_state"):
+                past_key_values.update_quasar_state(
+                    self.layer_idx, 
+                    recurrent_state, 
+                    (conv_state_q, conv_state_k, conv_state_v) if self.use_short_conv else None
+                )
+            else:
+                with contextlib.suppress(TypeError):
+                    past_key_values.update(
+                        recurrent_state=recurrent_state,
+                        conv_state=(conv_state_q, conv_state_k, conv_state_v) if self.use_short_conv else None,
+                        layer_idx=self.layer_idx,
+                        offset=q_len,
+                    )
+
+        # Final output gating using g_proj
+        # Handle flattened inputs (unpadded) from FSDP/Flash-Linear-Attention
+        if hidden_states.dim() == 2:
+            # (N, D) -> (N, H, D/H)
+            g = self.g_proj(hidden_states)
+            g = rearrange(g, "n (h d) -> n h d", d=self.head_dim)
+            _quasar_debug_tensor("output_gate", g, self.layer_idx)
+            o = self.o_norm(o, g)
+            o = rearrange(o, "n h d -> n (h d)")
+        else:
+            # (B, S, D) -> (B, S, H, D/H)
+            g = self.g_proj(hidden_states)
+            g = rearrange(g, "b s (h d) -> b s h d", d=self.head_dim)
+            _quasar_debug_tensor("output_gate", g, self.layer_idx)
+            o = self.o_norm(o, g)
+            o = rearrange(o, "b s h d -> b s (h d)")
+        _quasar_debug_tensor("post_norm_gate", o, self.layer_idx)
+        
+        o = self.o_proj(o)
+        _quasar_debug_tensor("o_proj", o, self.layer_idx)
+        if attention_mask is not None:
+            o = pad_input(o.squeeze(0), indices, batch_size, q_len)
+
+        # LFM2 expects 2 return values (hidden_states, _)
+        return o, None

fla/layers/rebased.py

@@ -0,0 +1,144 @@
+# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang
+
+"""
+https://github.com/corl-team/rebased/blob/main/flash_linear_attention/fla/layers/rebased_fast.py
+"""
+
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+from einops import rearrange
+
+from fla.modules.feature_map import RebasedFeatureMap
+from fla.ops.linear_attn import chunk_linear_attn, fused_chunk_linear_attn
+from fla.ops.rebased import parallel_rebased
+
+
+class ReBasedLinearAttention(nn.Module):
+
+    def __init__(
+        self,
+        hidden_size: int,
+        l_max: int = 2048,
+        feature_dim: int = 16,
+        num_key_value_heads: int = 16,
+        num_heads: int = 16,
+        use_gamma: bool | None = True,
+        use_beta: bool | None = True,
+        normalize: bool | None = True,
+        causal: bool = True,
+        eps: float = 1e-5,
+        mode: str = "parallel",
+        layer_idx: int | None = None,
+        **kwargs,
+    ) -> ReBasedLinearAttention:
+        super().__init__()
+        self.hidden_size = hidden_size
+        self.l_max = l_max
+        self.mode = mode
+        assert self.mode in ["fused_chunk", "parallel", 'chunk']
+
+        self.feature_dim = feature_dim
+        self.num_key_value_heads = num_key_value_heads
+        self.num_heads = num_heads
+        self.head_dim = self.hidden_size // self.num_key_value_heads
+        self.use_gamma = use_gamma
+        self.use_beta = use_beta
+        self.normalize = normalize
+        self.causal = causal
+        self.eps = eps
+        self.mode = mode
+        self.layer_idx = layer_idx
+
+        self.feature_map = RebasedFeatureMap(self.feature_dim, use_gamma, use_beta, normalize)
+        self.q_proj = nn.Linear(self.hidden_size, self.feature_dim * self.num_heads, bias=False)
+        self.k_proj = nn.Linear(self.hidden_size, self.feature_dim * self.num_heads, bias=False)
+        self.v_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=False)
+        self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=False)
+        self.dropout = nn.Identity()
+
+    def forward(self, hidden_states: torch.Tensor, **kwargs):
+        mode = self.mode
+        q = rearrange(
+            self.q_proj(hidden_states),
+            "... (h d) -> ... h d",
+            h=self.num_heads,
+            d=self.feature_dim,
+        )
+        k = rearrange(
+            self.k_proj(hidden_states),
+            "... (h d) -> ... h d",
+            h=self.num_heads,
+            d=self.feature_dim,
+        )
+        v = rearrange(
+            self.v_proj(hidden_states),
+            "... (h d) -> ... h d",
+            h=self.num_key_value_heads,
+            d=self.head_dim,
+        )
+        q, k = self.feature_map(q, flatten=(mode != 'parallel')), self.feature_map(k, flatten=(mode != 'parallel'))
+        if mode == "fused_chunk":
+            o = fused_chunk_linear_attn(
+                q=q,
+                k=k,
+                v=v,
+                normalize=True,
+                scale=1,
+            )
+        elif mode == 'chunk':
+            o = chunk_linear_attn(
+                q=q,
+                k=k,
+                v=v,
+                normalize=True,
+                scale=1,
+            )
+        elif mode == 'parallel':
+            assert q.shape[-1] <= 128
+            o = parallel_rebased(
+                q=q,
+                k=k,
+                v=v,
+                eps=self.eps,
+                use_scale=True,
+                use_normalize=True,
+            )
+        o = rearrange(o, "... h d -> ... (h d)")
+        o = self.o_proj(o)
+        o = self.dropout(o)
+        return o
+
+    # https://github.com/HazyResearch/zoology/blob/main/zoology/mixers/based.py#L119
+    def forward_reference(
+        self,
+        hidden_states: torch.Tensor,
+        filters: torch.Tensor = None,
+        *args,
+        **kwargs,
+    ):
+        """
+        x (torch.Tensor): tensor of shape (b, d, t)
+        y (torch.Tensor): tensor of shape (b, d, t)
+        """
+        b, t, _ = hidden_states.size()
+        q, k, v = self.q_proj(hidden_states), self.k_proj(hidden_states), self.v_proj(hidden_states)
+
+        q = q.view(b, t, -1, self.feature_dim).transpose(1, 2)
+        k = k.view(b, t, -1, self.feature_dim).transpose(1, 2)
+        v = v.view(b, t, -1, self.head_dim).transpose(1, 2)
+
+        # Linear attention
+        q, k = self.feature_map(q), self.feature_map(k)
+        q, k, v = q.unsqueeze(-2), k.unsqueeze(-2), v.unsqueeze(-1)
+
+        # Compute attention
+        if self.causal:
+            y = ((q * (k * v).cumsum(2)).sum(-1) / ((q * k.cumsum(2)).sum(-1) + self.eps))
+        else:
+            y = ((q * (k * v).sum(2, True)).sum(-1) / ((q * k.sum(2, True)).sum(-1) + self.eps))
+        y = rearrange(y, 'b h t d -> b t (h d)')
+        y = self.o_proj(y.to(hidden_states.dtype))
+        y = self.dropout(y)
+        return y.to(hidden_states.dtype)

fla/layers/rodimus.py

@@ -0,0 +1,397 @@
+# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang
+
+from __future__ import annotations
+
+import warnings
+from typing import TYPE_CHECKING
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.utils.checkpoint
+from einops import rearrange, repeat
+from transformers.utils import logging
+
+from fla.layers.utils import (
+    get_layer_cache,
+    get_unpad_data,
+    index_first_axis,
+    pad_input,
+    require_cache_layer_idx,
+    unpad_input,
+    update_layer_cache,
+)
+from fla.modules import RMSNorm, RotaryEmbedding, ShortConvolution
+from fla.modules.layernorm_gated import RMSNormGated
+from fla.ops.gla import chunk_gla, fused_chunk_gla, fused_recurrent_gla
+
+if TYPE_CHECKING:
+    from transformers.processing_utils import Unpack
+
+    from fla.models.utils import Cache
+
+try:
+    from flash_attn import flash_attn_func, flash_attn_varlen_func
+except ImportError:
+    warnings.warn(
+        "Flash Attention is not installed. Please install it via `pip install flash-attn --no-build-isolation`",
+        category=ImportWarning,
+    )
+    flash_attn_func = None
+
+logger = logging.get_logger(__name__)
+
+
+def align_multiple(value, multiple_size=8):
+    if value % multiple_size != 0:
+        value += multiple_size - (value % multiple_size)
+    return value
+
+
+def autocast_to_fp16(x):
+    if x.dtype not in {torch.float16, torch.bfloat16}:
+        return x.to(dtype=torch.bfloat16)
+    else:
+        return x
+
+
+class RodimusAttention(nn.Module):
+    def __init__(
+        self,
+        block_type: str = 'rodimus',
+        mode: str = 'chunk',
+        hidden_size: int = 1024,
+        input_gate_low_rank: float | str | None = 'auto',
+        expand_ratio: int = 64,
+        use_short_conv: bool = True,
+        conv_size: int = 4,
+        conv_bias: bool = True,
+        norm_eps: float = 1e-5,
+        k_norm_eps: float | None = None,
+        residual_in_fp32: bool = True,
+        layer_idx: int = None,
+    ):
+        super().__init__()
+
+        self.block_type = block_type
+        self.mode = mode
+        self.hidden_size = hidden_size
+        self.d_inner = align_multiple(int(self.hidden_size * 2), 8)
+
+        self.expand_ratio = expand_ratio
+        self.input_gate_low_rank = max(self.hidden_size // 64, 16) if input_gate_low_rank == "auto" else input_gate_low_rank
+
+        self.use_short_conv = use_short_conv
+        self.conv_size = conv_size
+        self.conv_bias = conv_bias
+
+        self.norm_eps = norm_eps
+        self.k_norm_eps = k_norm_eps if k_norm_eps is not None else 1e-12
+        self.mem_size = expand_ratio
+
+        self.residual_in_fp32 = residual_in_fp32
+        self.layer_idx = layer_idx
+
+        assert mode in ['chunk', 'fused_recurrent', 'fused_chunk'], f"Not supported mode `{mode}`."
+
+        self.gate_proj = nn.Linear(self.hidden_size, self.d_inner, bias=False)
+        self.up_proj = nn.Linear(self.hidden_size, self.d_inner, bias=False)
+        self.activation_norm = RMSNormGated(hidden_size=self.d_inner, eps=norm_eps, norm_before_gate=False)
+        self.down_proj = nn.Linear(self.d_inner, self.hidden_size, bias=False)
+
+        if use_short_conv:
+            self.short_conv = ShortConvolution(
+                hidden_size=self.d_inner,
+                kernel_size=conv_size,
+                bias=conv_bias,
+                activation='silu',
+            )
+
+        self.residual_weight = nn.Parameter(torch.ones(
+            (self.d_inner, ), dtype=torch.float32 if self.residual_in_fp32 else None), requires_grad=True)
+
+        self.k_proj = nn.Linear(self.d_inner, self.mem_size, bias=False)
+        self.q_proj = nn.Linear(self.d_inner, self.mem_size, bias=False)
+
+        self.g_gate_proj = nn.Linear(self.d_inner, self.mem_size, bias=True)
+        self.tau_gate_proj = nn.Linear(self.d_inner, self.mem_size, bias=True)
+        self.i_gate_proj = nn.Sequential(
+            nn.Linear(self.d_inner, self.input_gate_low_rank, bias=False),
+            nn.Linear(self.input_gate_low_rank, self.d_inner, bias=True),
+            nn.Sigmoid(),
+        )
+
+    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]:
+        if attention_mask is not None:
+            assert len(attention_mask.shape) == 2, (
+                "Expected attention_mask as a 0-1 matrix with shape [batch_size, seq_len] "
+                "for padding purposes (0 indicating padding). "
+                "Arbitrary attention masks of shape [batch_size, seq_len, seq_len] are not allowed."
+            )
+
+        batch_size, q_len, _ = hidden_states.shape
+        # mode = 'fused_recurrent' if hidden_states.shape[1] <= 64 else self.mode
+        mode = 'fused_recurrent' if hidden_states.shape[1] == 1 else self.mode
+
+        last_state = get_layer_cache(self, past_key_values)
+
+        cu_seqlens = kwargs.get('cu_seqlens')
+        if attention_mask is not None:
+            indices, cu_seqlens, _ = get_unpad_data(attention_mask[:, -q_len:])
+            hidden_states = index_first_axis(rearrange(hidden_states, "b s ... -> (b s) ..."), indices).unsqueeze(0)
+
+        hidden_states, final_gate = self.up_proj(hidden_states), self.gate_proj(hidden_states)
+
+        if self.use_short_conv:
+            conv_state = None
+            if last_state is not None:
+                conv_state = last_state['conv_state']
+            shift_hidden_states, conv_state = self.short_conv(
+                x=hidden_states,
+                cache=conv_state,
+                output_final_state=use_cache,
+                cu_seqlens=cu_seqlens,
+            )
+        else:
+            shift_hidden_states = hidden_states
+
+        q = self.q_proj(shift_hidden_states)
+        k = self.k_proj(shift_hidden_states)
+        v = self.i_gate_proj(hidden_states) * hidden_states
+
+        g_gate = F.linear(shift_hidden_states, self.g_gate_proj.weight) + self.g_gate_proj.bias.float()
+        tau_gate = F.linear(shift_hidden_states, self.tau_gate_proj.weight) + self.tau_gate_proj.bias.float()
+
+        g_gate = F.softplus(g_gate)
+        it_gate = g_gate
+        rt_gate_log = -g_gate
+
+        tau_gate = F.sigmoid(tau_gate)
+        it_gate = it_gate ** tau_gate
+        rt_gate_log = rt_gate_log * tau_gate
+
+        k = F.normalize(k.float(), dim=-1, eps=self.k_norm_eps) * it_gate
+        q, k, v, rt_gate_log = map(lambda x: x.unsqueeze(1).transpose(1, 2), (q, k, v, rt_gate_log))
+
+        recurrent_state = last_state['recurrent_state'] if last_state is not None else None
+        if mode == 'fused_recurrent':
+            o, recurrent_state = fused_recurrent_gla(
+                q=q,
+                k=k,
+                v=v,
+                gk=rt_gate_log,
+                initial_state=recurrent_state,
+                output_final_state=use_cache,
+                cu_seqlens=cu_seqlens,
+                head_first=False,
+            )
+        elif mode == 'fused_chunk':
+            o, recurrent_state = fused_chunk_gla(
+                q=q,
+                k=k,
+                v=v,
+                g=rt_gate_log,
+                initial_state=recurrent_state,
+                output_final_state=use_cache,
+                head_first=False,
+            )
+        elif mode == 'chunk':
+            q, k, rt_gate_log = map(lambda x: x.to(v.dtype), (q, k, rt_gate_log))
+            o, recurrent_state = chunk_gla(
+                q=q,
+                k=k,
+                v=v,
+                g=rt_gate_log,
+                initial_state=recurrent_state,
+                output_final_state=use_cache,
+                cu_seqlens=cu_seqlens,
+                head_first=False,
+            )
+        else:
+            raise NotImplementedError(f"Not supported mode `{mode}`.")
+
+        rodimus_caches = None
+        if past_key_values is not None:
+            if self.block_type == 'rodimus':
+                update_layer_cache(
+                    self,
+                    past_key_values,
+                    recurrent_state=recurrent_state,
+                    conv_state=conv_state if self.use_short_conv else None,
+                    offset=q_len,
+                )
+            else:
+                rodimus_caches = (recurrent_state, conv_state if self.use_short_conv else None)
+
+        o = (o.transpose(1, 2).squeeze(1) + (shift_hidden_states.float()
+             if self.residual_in_fp32 else shift_hidden_states) * self.residual_weight).to(o.dtype)
+
+        o = self.activation_norm(o, final_gate)
+        o = self.down_proj(o)
+
+        if attention_mask is not None:
+            o = pad_input(o.squeeze(0), indices, batch_size, q_len)
+
+        if self.block_type == 'rodimus':
+            return o, None, past_key_values
+        else:
+            return o, None, (past_key_values, rodimus_caches)
+
+
+class SlidingWindowSharedKeyAttention(nn.Module):
+    def __init__(
+        self,
+        hidden_size: int = 2048,
+        num_heads: int = 32,
+        qkv_bias: bool = False,
+        qk_norm: bool = False,
+        window_size: int = 2048,
+        rope_theta: float | None = 10000.,
+        max_position_embeddings: int | None = None,
+        layer_idx: int = None,
+    ):
+        super().__init__()
+
+        self.hidden_size = hidden_size
+        self.num_heads = num_heads
+        self.head_dim = self.hidden_size // self.num_heads
+        self.qkv_bias = qkv_bias
+        self.qk_norm = qk_norm
+
+        self.window_size = window_size
+        self.rope_theta = rope_theta
+        self.max_position_embeddings = max_position_embeddings
+        self.layer_idx = layer_idx
+
+        self.q_proj = nn.Linear(self.hidden_size, self.hidden_size, bias=self.qkv_bias)
+        self.k_proj = nn.Linear(self.hidden_size, self.head_dim, bias=self.qkv_bias)
+        self.v_proj = nn.Linear(self.hidden_size, self.hidden_size, bias=self.qkv_bias)
+        self.o_proj = nn.Linear(self.hidden_size, self.hidden_size, bias=False)
+
+        if qk_norm:
+            self.q_norm = RMSNorm(self.head_dim, dtype=torch.float32)
+            self.k_norm = RMSNorm(self.head_dim, dtype=torch.float32)
+
+        self.rotary = RotaryEmbedding(dim=self.head_dim, base=self.rope_theta)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: torch.LongTensor | None = None,
+        past_key_values: Cache | None = None,
+        output_attentions: bool = False,
+        use_cache: bool = False,
+        **kwargs,
+    ) -> tuple[torch.Tensor, torch.Tensor | None, tuple[torch.Tensor] | None]:
+        rodimus_caches = kwargs.get('rodimus_caches')
+
+        if attention_mask is not None:
+            assert len(attention_mask.shape) == 2, (
+                "Expected attention_mask as a 0-1 matrix with shape [batch_size, seq_len] "
+                "for padding purposes (0 indicating padding). "
+                "Arbitrary attention masks of shape [batch_size, seq_len, seq_len] are not allowed."
+            )
+
+        batch_size, q_len, _ = hidden_states.size()
+
+        q = rearrange(self.q_proj(hidden_states), '... (h d) -> ... h d', d=self.head_dim)
+        k = rearrange(self.k_proj(hidden_states), '... (h d) -> ... h d', d=self.head_dim)
+        v = rearrange(self.v_proj(hidden_states), '... (h d) -> ... h d', d=self.head_dim)
+
+        if self.qk_norm:
+            q, k = self.q_norm(q), self.k_norm(k)
+
+        # equivalent to cu_seqlens in `flash_attn`
+        cu_seqlens = kwargs.get('cu_seqlens')
+
+        layer_idx = require_cache_layer_idx(self, past_key_values)
+        seqlen_offset, max_seqlen = 0, q.shape[1]
+        if past_key_values is not None:
+            seqlen_offset = past_key_values.get_seq_length(layer_idx)
+            max_seqlen = q.shape[1] + seqlen_offset
+
+            if attention_mask is not None:
+                # to deliminate the offsets of padding tokens
+                seqlen_offset = seqlen_offset + attention_mask.sum(-1) - attention_mask.shape[-1]
+                max_seqlen = q.shape[1] + max(seqlen_offset)
+
+        if self.max_position_embeddings is not None:
+            max_seqlen = max(max_seqlen, self.max_position_embeddings)
+        q, k = self.rotary(q, k, seqlen_offset=seqlen_offset, max_seqlen=max_seqlen, cu_seqlens=cu_seqlens)
+
+        if past_key_values is not None:
+            if rodimus_caches is not None:
+                recurrent_state, conv_state = rodimus_caches
+            else:
+                recurrent_state, conv_state = None, None
+
+            cache_has_content = past_key_values.get_seq_length(layer_idx) > 0
+            k_cached, v_cached = past_key_values.update(
+                recurrent_state=recurrent_state,
+                conv_state=conv_state,
+                attn_state=[k.flatten(-2, -1), v.flatten(-2, -1)],
+                layer_idx=layer_idx,
+                offset=q_len,
+                cache_kwargs=dict(window_size=self.window_size),
+            )['attn_state']
+            if cache_has_content:
+                k, v = k_cached, v_cached
+                k = rearrange(k, '... (h d) -> ... h d', d=self.head_dim)
+                v = rearrange(v, '... (h d) -> ... h d', d=self.head_dim)
+
+        if flash_attn_func is None:
+            raise ImportError("Please install Flash Attention via `pip install flash-attn --no-build-isolation` first")
+
+        q, k, v = map(autocast_to_fp16, (q, k, v))
+        k = repeat(k, "... h d -> ... (n h) d", n=self.num_heads)
+        # Contains at least one padding token in the sequence
+        if attention_mask is not None:
+            q, (k, v), indices_q, cu_seqlens, max_seq_lens = unpad_input(
+                q=q,
+                states=(k, v),
+                attention_mask=attention_mask[:, -max(self.window_size, q_len):],
+                q_len=q_len,
+            )
+            cu_seqlens_q, cu_seqlens_k = cu_seqlens
+            max_seqlen_q, max_seqlen_k = max_seq_lens
+            o = flash_attn_varlen_func(
+                q, k, v,
+                cu_seqlens_q=cu_seqlens_q,
+                cu_seqlens_k=cu_seqlens_k,
+                max_seqlen_q=max_seqlen_q,
+                max_seqlen_k=max_seqlen_k,
+                causal=True,
+                window_size=(-1, -1) if self.window_size is None else (self.window_size-1, 0),
+            )
+            o = pad_input(o, indices_q, batch_size, q_len)
+        elif cu_seqlens is not None:
+            o = flash_attn_varlen_func(
+                q.squeeze(0), k.squeeze(0), v.squeeze(0),
+                cu_seqlens_q=cu_seqlens,
+                cu_seqlens_k=cu_seqlens,
+                max_seqlen_q=max_seqlen,
+                max_seqlen_k=max_seqlen,
+                causal=True,
+                window_size=(-1, -1) if self.window_size is None else (self.window_size-1, 0),
+            ).unsqueeze(0)
+        else:
+            o = flash_attn_func(
+                q, k, v,
+                causal=True,
+                window_size=(-1, -1) if self.window_size is None else (self.window_size-1, 0),
+            )
+        o = o.reshape(batch_size, q_len, -1)
+        o = self.o_proj(o.to(dtype=self.o_proj.weight.dtype))
+
+        if not output_attentions:
+            attentions = None
+
+        return o, attentions, past_key_values

fla/layers/rwkv6.py

@@ -0,0 +1,360 @@
+# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang
+
+# "Eagle and Finch: RWKV with Matrix-Valued States and Dynamic Recurrence"[https://arxiv.org/abs/2404.05892]
+
+from __future__ import annotations
+
+import math
+import warnings
+from typing import TYPE_CHECKING
+
+import torch
+import torch.nn as nn
+from einops import rearrange
+
+from fla.layers.utils import get_layer_cache, update_layer_cache
+from fla.modules import GroupNorm
+from fla.modules.activations import ACT2FN
+from fla.modules.token_shift import token_shift
+from fla.ops.rwkv6 import chunk_rwkv6, fused_recurrent_rwkv6
+
+if TYPE_CHECKING:
+    from fla.models.utils import Cache
+
+
+class RWKV6Attention(nn.Module):
+
+    def __init__(
+        self,
+        mode: str = 'chunk',
+        hidden_size: int = 1024,
+        expand_k: float = 0.5,
+        expand_v: float = 1.0,
+        num_heads: int = 4,
+        gate_fn: str = 'swish',
+        proj_low_rank_dim: int = 32,
+        gate_low_rank_dim: int = 64,
+        fuse_norm: bool = True,
+        elementwise_affine: bool | None = True,
+        norm_eps: float = 1e-5,
+        layer_idx: int = None,
+        **kwargs,
+    ) -> RWKV6Attention:
+        super().__init__()
+
+        self.mode = mode
+        self.hidden_size = hidden_size
+        self.expand_k = expand_k
+        self.expand_v = expand_v
+        self.num_heads = num_heads
+        self.proj_low_rank_dim = proj_low_rank_dim
+        self.gate_low_rank_dim = gate_low_rank_dim
+
+        self.key_dim = int(hidden_size * expand_k)
+        self.value_dim = int(hidden_size * expand_v)
+        self.layer_idx = layer_idx
+
+        assert mode in ['chunk', 'fused_recurrent'], f"Not supported mode `{mode}`."
+        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
+
+        self.time_shift = nn.ZeroPad2d((0, 0, 1, -1))
+        self.x_proj = nn.Sequential(
+            LerpLinear(hidden_size, proj_low_rank_dim * 5),
+            nn.Tanh(),
+            nn.Linear(proj_low_rank_dim * 5, hidden_size, bias=False),
+        )
+        self.x_bias = nn.Parameter(torch.zeros(5, hidden_size))
+
+        self.r_proj = DDLerpLinear(hidden_size, self.key_dim)
+        self.w_proj = DDLerpLinear(hidden_size, self.key_dim, low_rank_dim=gate_low_rank_dim)
+        self.k_proj = DDLerpLinear(hidden_size, self.key_dim)
+        self.v_proj = DDLerpLinear(hidden_size, self.value_dim)
+        self.g_proj = DDLerpLinear(hidden_size, self.value_dim)
+        self.bonus = nn.Parameter(torch.zeros(num_heads, self.head_k_dim))
+
+        # TODO: fuse GroupNorm and output gate
+        self.g_norm = GroupNorm(self.num_heads, self.value_dim, elementwise_affine=elementwise_affine, bias=True, eps=norm_eps)
+        self.o_proj = nn.Linear(self.value_dim, hidden_size, bias=False)
+        self.gate_fn = ACT2FN[gate_fn]
+
+        try:
+            from transformers.modeling_utils import _init_weights
+        except ImportError:
+            _init_weights = True
+        if _init_weights:
+            self.apply(self._initialize_weights)
+
+        warnings.warn(
+            "According to Bo, you are using a potentially buggy FLA implementation of RWKV. "
+            "If you plan to report any numbers based on this implementation, we strongly recommend "
+            "cross-checking with the official repo: https://github.com/BlinkDL/RWKV-LM. "
+            "Bo may disagree with results reported from this version.",
+        )
+
+    def _initialize_weights(self, module: nn.Module):
+        if getattr(module, "_is_hf_initialized", False):
+            return
+        if isinstance(module, nn.Linear):
+            nn.init.xavier_uniform_(module.weight, gain=2 ** -2.5)
+            if module.bias is not None:
+                nn.init.zeros_(module.bias)
+        if isinstance(module, nn.Parameter):
+            nn.init.xavier_uniform_(module, gain=2 ** -2.5)
+        module._is_hf_initialized = True
+
+    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,
+        cu_seqlens: torch.LongTensor | None = None,
+        **kwargs,
+    ) -> tuple[torch.Tensor, torch.Tensor | None, Cache | None]:
+        if attention_mask is not None:
+            assert len(attention_mask.shape) == 2, (
+                "Expected attention_mask as a 0-1 matrix with shape [batch_size, seq_len] "
+                "for padding purposes (0 indicating padding). "
+                "Arbitrary attention masks of shape [batch_size, seq_len, seq_len] are not allowed."
+            )
+
+        batch_size, seq_len, hidden_size = hidden_states.shape
+        # launching the triton kernel for just one token will actually be slower
+        mode = 'fused_recurrent' if hidden_states.shape[1] <= 64 else self.mode
+
+        last_state = get_layer_cache(self, past_key_values)
+
+        if attention_mask is not None:
+            hidden_states = hidden_states.mul(attention_mask[:, -hidden_states.shape[-2]:, None])
+
+        if hidden_states.shape[1] == 1 and last_state is not None:
+            shifted = last_state['conv_state'].unsqueeze(1)
+            delta = shifted - hidden_states
+        elif last_state is None:
+            delta = token_shift(hidden_states, cu_seqlens)
+        else:
+            shifted = self.time_shift(hidden_states)
+            shifted[:, 0] = last_state['conv_state']
+            delta = shifted - hidden_states
+
+        x = self.x_proj[0](hidden_states, delta, cu_seqlens).view(batch_size, seq_len, -1, self.proj_low_rank_dim)
+        x = torch.einsum('b t n r, h n r-> b t n h', self.x_proj[1](x), self.x_proj[2].weight.view(hidden_size, 5, -1))
+
+        r, w, k, v, g = x.add_(self.x_bias).unbind(-2)
+        r = self.r_proj(hidden_states, r, delta, cu_seqlens)
+        w = self.w_proj(hidden_states, w, delta, cu_seqlens)
+        k = self.k_proj(hidden_states, k, delta, cu_seqlens)
+        v = self.v_proj(hidden_states, v, delta, cu_seqlens)
+        g = self.g_proj(hidden_states, g, delta, cu_seqlens)
+
+        # dealing with left-padding
+        if attention_mask is not None:
+            v = v.mul(attention_mask[:, -v.shape[-2]:, None])
+        r, w, k = map(lambda x: rearrange(x, 'b t (h d) -> b t h d', d=self.head_k_dim), (r, w, k))
+        v = rearrange(v, 'b t (h d) -> b t h d', d=self.head_v_dim)
+        w = -torch.exp(w)
+        u = self.bonus
+
+        recurrent_state = last_state['recurrent_state'] if last_state is not None else None
+
+        if mode == 'fused_recurrent':
+            o, recurrent_state = fused_recurrent_rwkv6(
+                r=r,
+                k=k,
+                v=v,
+                w=w,
+                u=u,
+                scale=1.,
+                initial_state=recurrent_state,
+                output_final_state=use_cache,
+                cu_seqlens=cu_seqlens,
+            )
+        elif mode == 'chunk':
+            o, recurrent_state = chunk_rwkv6(
+                r=r,
+                k=k,
+                v=v,
+                w=w,
+                u=u,
+                scale=1.,
+                initial_state=recurrent_state,
+                output_final_state=use_cache,
+                cu_seqlens=cu_seqlens,
+            )
+        else:
+            raise NotImplementedError(f"Not supported mode `{mode}`.")
+
+        update_layer_cache(
+            self,
+            past_key_values,
+            recurrent_state=recurrent_state,
+            conv_state=hidden_states[:, -1],
+            offset=seq_len,
+        )
+
+        o = self.g_norm(rearrange(o, '... h d -> ... (h d)')) * self.gate_fn(g)
+        o = self.o_proj(o)
+
+        return o, None, past_key_values
+
+
+class LoRA(nn.Module):
+
+    def __init__(
+        self,
+        input_dim: int,
+        output_dim: int,
+        low_rank_dim: int,
+        bias: bool | None = True,
+        activation: str | None = 'tanh',
+    ):
+        super().__init__()
+
+        self.input_dim = input_dim
+        self.output_dim = output_dim
+        self.low_rank_dim = low_rank_dim
+        self.bias = bias
+
+        if activation is None:
+            self.activation = nn.Identity()
+        elif activation == 'sigmoid':
+            self.activation = nn.Sigmoid()
+        elif activation == 'tanh':
+            self.activation = nn.Tanh()
+        elif activation == 'relu':
+            self.activation = nn.ReLU()
+        else:
+            raise ValueError(f"Not supported activation `{activation}`.")
+
+        self.lora = nn.Sequential(
+            nn.Linear(input_dim, low_rank_dim, bias=False),
+            self.activation,
+            nn.Linear(low_rank_dim, output_dim, bias=bias),
+        )
+        try:
+            from transformers.modeling_utils import _init_weights
+        except ImportError:
+            _init_weights = True
+        if _init_weights:
+            self.apply(self._initialize_weights)
+
+    def __repr__(self) -> str:
+        s = f"{self.__class__.__name__}("
+        s += f"input_dim={self.input_dim}, low_rank_dim={self.low_rank_dim}, output_dim={self.output_dim}"
+        if not self.bias:
+            s += f", bias={self.bias}"
+        s += ")"
+        return s
+
+    def _initialize_weights(self, module: nn.Module):
+        if getattr(module, "_is_hf_initialized", False):
+            return
+
+        # Initialize weights to zero as in original code
+        nn.init.zeros_(self.lora[0].weight)
+        original_dtype = self.lora[2].weight.dtype
+        shape = self.lora[2].weight.shape
+        # Convert to float32 for numerical stability in orthogonal init
+        weight_fp32 = self.lora[2].weight.float()
+
+        # Calculate gain based on dimensions
+        gain = math.sqrt(shape[1] / shape[0]) if shape[1] > shape[0] else 1
+
+        # Apply orthogonal initialization with scaling factor 0.1
+        nn.init.orthogonal_(weight_fp32, gain=gain * 0.1)
+
+        # Convert back to original dtype
+        self.lora[2].weight.data.copy_(weight_fp32.to(original_dtype))
+        # Set Lora[2] bias to zero
+        if self.lora[2].bias is not None:
+            nn.init.zeros_(self.lora[2].bias)
+
+        module._is_hf_initialized = True
+
+    def set_bias_value(self, value):
+        """Set bias to a specific value (for v0, w0 etc.)"""
+        if self.bias and self.lora[2].bias is not None:
+            if isinstance(value, torch.Tensor):
+                # Handle tensor values
+                self.lora[2].bias.data.copy_(value.to(self.lora[2].bias.dtype))
+            else:
+                # Handle scalar values
+                nn.init.constant_(self.lora[2].bias, value)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.lora(x)
+
+
+class LerpLinear(nn.Module):
+
+    def __init__(
+        self,
+        input_dim: int,
+        output_dim: int,
+        low_rank_dim: int | None = None,
+    ):
+        super().__init__()
+
+        self.input_dim = input_dim
+        self.output_dim = output_dim
+        self.low_rank_dim = low_rank_dim
+
+        self.time_shift = nn.ZeroPad2d((0, 0, 1, -1))
+        if low_rank_dim is None:
+            self.linear = nn.Linear(input_dim, output_dim, bias=False)
+        else:
+            self.linear = LoRA(input_dim, output_dim, low_rank_dim)
+        self.mu = nn.Parameter(torch.zeros(input_dim))
+
+    def __repr__(self) -> str:
+        s = f"{self.__class__.__name__}({self.input_dim}, {self.output_dim}"
+        if self.low_rank_dim is not None:
+            s += f", low_rank_dim={self.low_rank_dim}"
+        s += ")"
+        return s
+
+    def forward(self, x: torch.Tensor, delta: torch.Tensor | None = None,
+                cu_seqlens: torch.LongTensor | None = None) -> torch.Tensor:
+        if delta is None:
+            delta = token_shift(x, cu_seqlens)
+        return self.linear(x + delta * self.mu)
+
+
+class DDLerpLinear(nn.Module):
+
+    def __init__(
+        self,
+        input_dim: int,
+        output_dim: int,
+        low_rank_dim: int | None = None,
+    ):
+        super().__init__()
+
+        self.input_dim = input_dim
+        self.output_dim = output_dim
+        self.low_rank_dim = low_rank_dim
+
+        self.time_shift = nn.ZeroPad2d((0, 0, 1, -1))
+        if low_rank_dim is None:
+            self.linear = nn.Linear(input_dim, output_dim, bias=False)
+        else:
+            self.linear = LoRA(input_dim, output_dim, low_rank_dim)
+
+    def __repr__(self) -> str:
+        s = f"{self.__class__.__name__}({self.input_dim}, {self.output_dim}"
+        if self.low_rank_dim is not None:
+            s += f", low_rank_dim={self.low_rank_dim}"
+        s += ")"
+        return s
+
+    def forward(self, x: torch.Tensor, mu: torch.Tensor,
+                delta: torch.Tensor | None = None,
+                cu_seqlens: torch.LongTensor | None = None) -> torch.Tensor:
+        if delta is None:
+            delta = token_shift(x, cu_seqlens)
+        return self.linear(x + delta * mu)

fla/layers/rwkv7.py

@@ -0,0 +1,347 @@
+# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang
+
+from __future__ import annotations
+
+import warnings
+from typing import TYPE_CHECKING
+
+import torch
+import torch.nn as nn
+from einops import rearrange
+from torch.nn import functional as F
+
+from fla.layers.rwkv6 import LoRA
+from fla.layers.utils import get_layer_cache, update_layer_cache
+from fla.modules import GroupNorm
+from fla.modules.l2norm import l2_norm
+from fla.modules.token_shift import token_shift
+from fla.ops.rwkv7 import chunk_rwkv7, fused_mul_recurrent_rwkv7
+from fla.ops.rwkv7.fused_addcmul import fused_addcmul_rwkv7
+from fla.ops.rwkv7.fused_k_update import fused_k_rwkv7
+from fla.ops.rwkv7.gate_output_correction import gate_output_correction
+
+if TYPE_CHECKING:
+    from fla.models.utils import Cache
+
+
+class RWKV7Attention(nn.Module):
+
+    def __init__(
+        self,
+        mode: str = 'chunk',
+        hidden_size: int = 1024,
+        head_dim: int | None = 64,
+        num_heads: int | None = None,
+        decay_low_rank_dim: int | None = None,
+        gate_low_rank_dim: int | None = None,
+        a_low_rank_dim: int | None = None,
+        v_low_rank_dim: int | None = None,
+        elementwise_affine: bool | None = True,
+        norm_eps: float = 1e-5,
+        layer_idx: int = None,
+        fuse_norm: bool = False,
+        value_dim: int = None,
+        num_hidden_layers: int = None,
+        **kwargs,
+    ) -> RWKV7Attention:
+        super().__init__()
+
+        self.mode = mode
+        assert mode in ['chunk', 'fused_recurrent'], f"Not supported mode `{mode}`."
+        self.hidden_size = hidden_size
+
+        self.key_dim = hidden_size
+        self.value_dim = value_dim if value_dim is not None else hidden_size
+        if head_dim is None and num_heads is None:
+            raise ValueError("Either `head_dim` or `num_heads` must be specified.")
+        elif head_dim is not None:
+            self.head_dim = head_dim
+            self.num_heads = int(hidden_size // head_dim)
+        elif num_heads is not None:
+            self.head_dim = int(hidden_size // num_heads)
+            self.num_heads = num_heads
+        self.head_v_dim = int(self.value_dim // self.num_heads)
+
+        # Increase lora dimension for headdim>64
+        factor = self.head_dim / 64
+        if decay_low_rank_dim is None:
+            decay_low_rank_dim = max(32, int(round((2.5 * (hidden_size**0.5)) * factor / 32) * 32))
+            self.decay_low_rank_dim = decay_low_rank_dim
+        else:
+            self.decay_low_rank_dim = decay_low_rank_dim
+
+        if gate_low_rank_dim is None:
+            gate_low_rank_dim = max(32, int(round((5 * (hidden_size**0.5)) / 32) * 32))
+            self.gate_low_rank_dim = gate_low_rank_dim
+        else:
+            self.gate_low_rank_dim = gate_low_rank_dim
+
+        if a_low_rank_dim is None:
+            a_low_rank_dim = max(32, int(round((2.5 * (hidden_size**0.5)) * factor / 32) * 32))
+            self.a_low_rank_dim = a_low_rank_dim
+        else:
+            self.a_low_rank_dim = a_low_rank_dim
+
+        if v_low_rank_dim is None:
+            v_low_rank_dim = max(32, int(round((1.7 * (hidden_size**0.5)) * factor / 32) * 32))
+            self.v_low_rank_dim = v_low_rank_dim
+        else:
+            self.v_low_rank_dim = v_low_rank_dim
+
+        self.layer_idx = layer_idx
+        self.num_hidden_layers = num_hidden_layers
+        self.fuse_norm = fuse_norm
+
+        self.time_shift = nn.ZeroPad2d((0, 0, 1, -1))
+        self.x_r = nn.Parameter(torch.zeros(1, 1, hidden_size))
+        self.x_w = nn.Parameter(torch.zeros(1, 1, hidden_size))
+        self.x_k = nn.Parameter(torch.zeros(1, 1, hidden_size))
+        self.x_v = nn.Parameter(torch.zeros(1, 1, hidden_size))
+        self.x_a = nn.Parameter(torch.zeros(1, 1, hidden_size))
+        self.x_g = nn.Parameter(torch.zeros(1, 1, hidden_size))
+
+        self.k_k = nn.Parameter(torch.zeros(self.key_dim))
+        self.k_a = nn.Parameter(torch.zeros(self.key_dim))
+        self.r_k = nn.Parameter(torch.zeros(self.num_heads, self.head_dim))
+
+        self.r_proj = nn.Linear(hidden_size, self.key_dim, bias=False)
+        self.k_proj = nn.Linear(hidden_size, self.key_dim, bias=False)
+        self.v_proj = nn.Linear(hidden_size, self.value_dim, bias=False)
+        self.o_proj = nn.Linear(self.value_dim, hidden_size, bias=False)
+
+        self.w_lora = LoRA(hidden_size, self.key_dim, low_rank_dim=decay_low_rank_dim, activation='tanh')
+        if self.layer_idx != 0:
+            self.v_lora = LoRA(hidden_size, self.value_dim, low_rank_dim=v_low_rank_dim, activation=None)
+        self.a_lora = LoRA(hidden_size, self.key_dim, low_rank_dim=a_low_rank_dim, activation=None)
+        self.g_lora = LoRA(hidden_size, self.value_dim, low_rank_dim=gate_low_rank_dim, activation='sigmoid', bias=False)
+
+        if self.fuse_norm:
+            self.g_norm = GroupNorm(
+                num_groups=self.num_heads,
+                hidden_size=self.value_dim,
+                elementwise_affine=elementwise_affine,
+                eps=self.head_dim*norm_eps,
+                bias=True,
+            )
+        else:
+            self.g_norm = nn.GroupNorm(
+                num_groups=self.num_heads,
+                num_channels=self.value_dim,
+                eps=self.head_dim*norm_eps,
+                affine=elementwise_affine,
+            )
+
+        try:
+            from transformers.modeling_utils import _init_weights
+        except ImportError:
+            _init_weights = True
+        if _init_weights:
+            self.apply(self._initialize_weights)
+        for name, module in self.named_modules():
+            module._in_rwkv_module = True
+
+        warnings.warn(
+            "According to Bo, you are using a potentially buggy FLA implementation of RWKV. "
+            "If you plan to report any numbers based on this implementation, we strongly recommend "
+            "cross-checking with the official repo: https://github.com/BlinkDL/RWKV-LM. "
+            "Bo may disagree with results reported from this version.",
+        )
+
+    @torch.no_grad()
+    @torch.compiler.disable
+    def _initialize_weights(self, module: nn.Module):
+        if getattr(module, "_is_hf_initialized", False):
+            return
+
+        # Initialize only when we're processing the RWKV7Attention module itself
+        if isinstance(module, RWKV7Attention) and self.layer_idx is not None:
+            ratio_0_to_1 = self.layer_idx / (self.num_hidden_layers - 1)  # 0 to 1
+            ratio_1_to_almost0 = 1.0 - (self.layer_idx / self.num_hidden_layers)  # 1 to ~0
+
+            # Create position-based initialization tensor
+            ddd = torch.ones(1, 1, self.hidden_size, device=self.x_r.device)
+            www = torch.zeros(self.hidden_size, device=self.x_r.device)
+            zigzag = torch.zeros(self.hidden_size, device=self.x_r.device)
+            linear = torch.zeros(self.hidden_size, device=self.x_r.device)
+            for n in range(self.hidden_size):
+                linear[n] = n / (self.hidden_size-1) - 0.5
+                zigzag[n] = ((n % self.head_dim) - ((self.head_dim-1) / 2)) / ((self.head_dim-1) / 2)
+                zigzag[n] = zigzag[n] * abs(zigzag[n])
+                www[n] = -6 + 6 * (n / (self.hidden_size - 1)) ** (1 + 1 * ratio_0_to_1 ** 0.3)
+                ddd[0, 0, n] = n / self.hidden_size
+
+            # Initialize x_* parameters directly
+            self.x_r.data = (1.0 - torch.pow(ddd, 0.2 * ratio_1_to_almost0)).to(self.x_r.dtype)
+            self.x_w.data = (1.0 - torch.pow(ddd, 0.9 * ratio_1_to_almost0)).to(self.x_w.dtype)
+            self.x_k.data = (1.0 - torch.pow(ddd, 0.7 * ratio_1_to_almost0)).to(self.x_k.dtype)
+            self.x_v.data = (1.0 - torch.pow(ddd, 0.7 * ratio_1_to_almost0)).to(self.x_v.dtype)
+            self.x_a.data = (1.0 - torch.pow(ddd, 0.9 * ratio_1_to_almost0)).to(self.x_a.dtype)
+            self.x_g.data = (1.0 - torch.pow(ddd, 0.2 * ratio_1_to_almost0)).to(self.x_g.dtype)
+
+            # Initialize k_k, k_a, r_k
+            nn.init.constant_(self.k_a, 1.02)
+            nn.init.constant_(self.r_k, -0.04)
+            self.k_k.data.copy_((torch.zeros(self.hidden_size, device=self.k_k.device) +
+                                0.71 - linear*0.1).to(self.k_k.dtype))
+            # Set specific bias values for LoRA modules
+            # 0.5 comes from F.softplus
+            self.w_lora.set_bias_value(www + 0.5 + zigzag*2.5)
+            self.a_lora.set_bias_value(-0.19 + zigzag*0.3 + linear*0.4)
+
+            # v0 initialization - ones (for non-first layers)
+            if self.layer_idx != 0:
+                self.v_lora._initialize_weights(self.v_lora)
+                self.v_lora.set_bias_value(0.73 - linear*0.4)
+
+            # Initialize GroupNorm
+            self.g_norm.weight.data[:] = ((self.layer_idx + 1) / self.num_hidden_layers) ** 0.7
+
+            # Initialize Linear projections
+            self._orthogonal_init(self.r_proj.weight)
+            self._orthogonal_init(self.k_proj.weight, gain=0.1)
+            self._orthogonal_init(self.v_proj.weight)
+            self.o_proj.weight.data.zero_()
+
+            # Clean up temporary tensors to free memory
+            del ddd, www, zigzag, linear
+
+        module._is_hf_initialized = True
+
+    @staticmethod
+    def _orthogonal_init(weight, gain=1.0):
+        oringinal_dtype = weight.dtype
+        weight = weight.float()
+        nn.init.orthogonal_(weight, gain=gain)
+        weight = weight.to(oringinal_dtype)
+
+    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,
+        v_first: torch.Tensor = None,
+        cu_seqlens: torch.LongTensor | None = None,
+        **kwargs,
+    ) -> tuple[torch.Tensor, torch.Tensor | None, Cache | None]:
+        batch_size, seq_len, _ = hidden_states.shape
+        if attention_mask is not None:
+            assert len(attention_mask.shape) == 2, (
+                "Expected attention_mask as a 0-1 matrix with shape [batch_size, seq_len] "
+                "for padding purposes (0 indicating padding). "
+                "Arbitrary attention masks of shape [batch_size, seq_len, seq_len] are not allowed."
+            )
+            am = attention_mask.narrow(1, attention_mask.size(1) - seq_len, seq_len).unsqueeze(-1)
+
+        last_state = get_layer_cache(self, past_key_values)
+
+        if attention_mask is not None:
+            hidden_states = hidden_states.mul(am)
+
+        # delta [batch_size, seq_len, hidden_size]
+        # conv_cache [N, D]
+        if last_state is None:
+            conv_cache = None
+            recurrent_state = None
+        else:
+            conv_cache = last_state['conv_state']
+            recurrent_state = last_state['recurrent_state']
+
+        delta, conv_state = token_shift(
+            hidden_states, cu_seqlens, output_cache=True, cache=conv_cache,
+        )
+        xr, xw, xk, xv, xa, xg = fused_addcmul_rwkv7(hidden_states, delta, self.x_r, self.x_w,
+                                                     self.x_k, self.x_v, self.x_a, self.x_g)
+
+        r = self.r_proj(xr)
+        # Using bf16 for LoRA computation is numerically safe here because:
+        # 1. After sigmoid activation:
+        #    - Max absolute error (vs float32): 0.003
+        #    - Mean absolute error: 0.0004
+        # 2. Subsequent scaling by -0.6065 will further reduce relative error
+        #    (error scales linearly with constant multiplication)
+        # 3. Final compounded error remains within acceptable bounds for bf16 precision
+        # Empirical observation confirms bf16 introduces no practical degradation
+        w = -0.6065306597126334 * self.w_lora(xw).sigmoid()
+
+        k = self.k_proj(xk)
+        v = self.v_proj(xv)
+
+        if self.layer_idx == 0:
+            v_first = v
+        else:
+            v = torch.lerp(v, v_first, self.v_lora(xv).sigmoid())
+        a = self.a_lora(xa).sigmoid()
+        g = self.g_lora(xg)
+
+        if self.fuse_norm:
+            kk = l2_norm(rearrange(k * self.k_k, 'b t (h d) -> b t h d', d=self.head_dim))
+        else:
+            kk = F.normalize(rearrange(k * self.k_k, 'b t (h d) -> b t h d', d=self.head_dim), dim=-1, p=2.0)
+
+        # Prefer addcmul over expanded form for numerical stability in bf16:
+        # 1. Fused Multiply-Add (FMA) in addcmul reduces intermediate rounding:
+        #    - Single op vs original 3 ops (mul, sub, mul)
+        #    - 1 less intermediate value storage (bf16 write->read overhead)
+        # 2. Mathematically equivalent to k*(1 + (a-1)*self.k_a)
+        #    but with better precision preservation
+        # 3. Particularly crucial for bf16 where intermediate values easily lose precision
+        # 4. Pytorch method: k = k.addcmul(k * (a - 1), self.k_a)
+        k = fused_k_rwkv7(k, a, self.k_a)
+
+        # dealing with left-padding
+        if attention_mask is not None:
+            v = v * am
+
+        r, w, k, a = map(lambda x: rearrange(x, 'b t (h d) -> b t h d', d=self.head_dim), (r, w, k, a))
+        v = rearrange(v, 'b t (h d) -> b t h d', d=self.head_v_dim)
+
+        if self.training or seq_len >= 64:
+            # if training, use chunk mode no matter how short the sequence is
+            # launching the triton kernel for just one token will actually be slower
+            o, recurrent_state = chunk_rwkv7(
+                r=r,
+                w=w,
+                k=k,
+                v=v,
+                a=-kk,
+                b=kk * a,
+                scale=1.,
+                initial_state=recurrent_state,
+                output_final_state=use_cache,
+                cu_seqlens=cu_seqlens,
+                safe_gate=True,
+                chunk_size=64,
+            )
+        else:
+            o, recurrent_state = fused_mul_recurrent_rwkv7(
+                r=r,
+                w=w,
+                k=k,
+                v=v,
+                kk=kk,
+                a=a,
+                scale=1.,
+                initial_state=recurrent_state,
+                output_final_state=use_cache,
+                cu_seqlens=cu_seqlens,
+            )
+
+        update_layer_cache(
+            self,
+            past_key_values,
+            recurrent_state=recurrent_state,
+            conv_state=conv_state,
+            offset=r.shape[1],
+        )
+
+        if self.fuse_norm:
+            o = self.g_norm(rearrange(o, '... h d -> ... (h d)'))
+        else:
+            o = self.g_norm(rearrange(o, 'b t h d -> (b t) (h d)')).view(batch_size, seq_len, -1)
+
+        o = gate_output_correction(o, r, k, self.r_k, v, g)
+        o = self.o_proj(o)
+
+        return o, None, past_key_values, v_first

fla/layers/simple_gla.py

@@ -0,0 +1,273 @@
+# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange, repeat
+
+from fla.layers.utils import get_layer_cache, update_layer_cache
+from fla.modules import FusedRMSNormGated, RMSNorm, ShortConvolution
+from fla.modules.activations import ACT2FN
+from fla.ops.simple_gla import chunk_simple_gla, fused_recurrent_simple_gla
+
+if TYPE_CHECKING:
+    from fla.models.utils import Cache
+
+
+class SimpleGatedLinearAttention(nn.Module):
+    r"""
+    The layer implementaion for [Gated Linear Attention Transformers with Hardware-Efficient Training](https://arxiv.org/abs/2312.06635).  # noqa
+    This layer calls the simplified GLA kernel in which the gating is head-wise instead of elementwise.
+
+    Args:
+        mode (str, Optional):
+            Which GLA kernel to use.
+            Currently available: `chunk`.
+            Default: `chunk`.
+        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.
+        feature_map (str, Optional):
+            Feature map function applied to queries/keys. Default: None.
+        use_short_conv (bool, Optional):
+            Whether to use short convolutions. Default: `False`.
+        conv_size (int, Optional):
+            The kernel size of the short convolution, only used when `use_short_conv` is `True`. Default: 4.
+        conv_bias (bool, Optional):
+            Whether to use bias in the short convolution, only used when `use_short_conv` is `True`. Default: `False`.
+        gate_fn (str, Optional):
+            The activation function for the output gate. Default: `swish`.
+        elementwise_affine (bool, Optional):
+            If `True`, applies elementwise affine to LayerNorm with learnable parameters. Default: `True`.
+        norm_eps (float, Optional):
+            The epsilon value for the layernorm/rmsnorm layer. Default: 1e-5.
+        gate_logit_normalizer (int, Optional):
+            The normalizer for the gate logits, appied after `logsigmoid`. Default: 16.
+        fuse_norm (bool, Optional):
+            Whether to fuse the norm and the output gate for better memory footprint. Default: `True`.
+        layer_idx (int, Optional):
+            The index of the layer. Default: None.
+    """
+
+    def __init__(
+        self,
+        mode: str = 'chunk',
+        hidden_size: int = 1024,
+        expand_k: float = 1.,
+        expand_v: float = 1.,
+        num_heads: int = 4,
+        num_kv_heads: int | None = None,
+        feature_map: str | None = None,
+        use_short_conv: bool = True,
+        conv_size: int = 4,
+        conv_bias: bool = False,
+        gate_fn: str = 'swish',
+        elementwise_affine: bool | None = True,
+        norm_eps: float = 1e-5,
+        gate_logit_normalizer: int = 16,
+        fuse_norm: bool = True,
+        layer_idx: int = None,
+    ) -> SimpleGatedLinearAttention:
+        super().__init__()
+
+        self.mode = mode
+        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.feature_map_fn = ACT2FN[feature_map] if feature_map is not None else None
+
+        self.use_short_conv = use_short_conv
+        self.conv_size = conv_size
+        self.conv_bias = conv_bias
+
+        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
+        self.layer_idx = layer_idx
+
+        assert mode in ['chunk', "fused_recurrent"], f"Not supported mode `{mode}`."
+        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
+
+        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.g_proj = nn.Linear(hidden_size, self.value_dim, bias=False)
+
+        if use_short_conv:
+            self.conv_size = conv_size
+            self.q_conv1d = ShortConvolution(
+                hidden_size=self.key_dim,
+                kernel_size=conv_size,
+                bias=conv_bias,
+                activation='silu',
+            )
+            self.k_conv1d = ShortConvolution(
+                hidden_size=self.key_dim_per_group,
+                kernel_size=conv_size,
+                bias=conv_bias,
+                activation='silu',
+            )
+            self.v_conv1d = ShortConvolution(
+                hidden_size=self.value_dim_per_group,
+                kernel_size=conv_size,
+                bias=conv_bias,
+                activation='silu',
+            )
+
+        self.gk_proj = nn.Linear(hidden_size, self.num_heads)
+
+        if gate_fn == 'swish' and fuse_norm:
+            self.g_norm_swish_gate = FusedRMSNormGated(
+                hidden_size=self.head_v_dim,
+                elementwise_affine=elementwise_affine,
+                eps=norm_eps,
+            )
+            self.fuse_norm_and_gate = True
+        else:
+            self.fuse_norm_and_gate = False
+            self.g_norm = RMSNorm(
+                hidden_size=self.head_v_dim,
+                elementwise_affine=elementwise_affine,
+                eps=norm_eps,
+                dtype=torch.float32
+            )
+            self.gate_fn = ACT2FN[gate_fn]
+        self.o_proj = nn.Linear(self.value_dim, hidden_size, bias=False)
+
+        self.gate_logit_normalizer = gate_logit_normalizer
+
+    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,
+    ) -> tuple[torch.Tensor, torch.Tensor | None, Cache | None]:
+        if attention_mask is not None:
+            assert len(attention_mask.shape) == 2, (
+                "Expected attention_mask as a 0-1 matrix with shape [batch_size, seq_len] "
+                "for padding purposes (0 indicating padding). "
+                "Arbitrary attention masks of shape [batch_size, seq_len, seq_len] are not allowed."
+            )
+
+        # launching the triton kernel for just one token will actually be slower
+        mode = 'fused_recurrent' if hidden_states.shape[1] <= 64 else self.mode
+
+        last_state = get_layer_cache(self, past_key_values)
+
+        cu_seqlens = kwargs.get('cu_seqlens')
+        if self.use_short_conv:
+            conv_state_q, conv_state_k, conv_state_v = None, None, None
+            if last_state is not None:
+                conv_state_q, conv_state_k, conv_state_v = last_state['conv_state']
+            conv_mask = attention_mask[:, -hidden_states.shape[1]:] if attention_mask is not None else None
+            q, conv_state_q = self.q_conv1d(
+                x=self.q_proj(hidden_states),
+                mask=conv_mask,
+                cache=conv_state_q,
+                output_final_state=use_cache,
+                cu_seqlens=cu_seqlens,
+            )
+            k, conv_state_k = self.k_conv1d(
+                x=self.k_proj(hidden_states),
+                mask=conv_mask,
+                cache=conv_state_k,
+                output_final_state=use_cache,
+                cu_seqlens=cu_seqlens,
+            )
+            v, conv_state_v = self.v_conv1d(
+                x=self.v_proj(hidden_states),
+                mask=conv_mask,
+                cache=conv_state_v,
+                output_final_state=use_cache,
+                cu_seqlens=cu_seqlens,
+            )
+        else:
+            q = self.q_proj(hidden_states)
+            k = self.k_proj(hidden_states)
+            v = self.v_proj(hidden_states)
+        gk = self.gk_proj(hidden_states)
+
+        if self.feature_map_fn is not None:
+            q, k = map(self.feature_map_fn, (q, k))
+        # dealing with left-padding
+        if attention_mask is not None:
+            v = v.mul_(attention_mask[:, -v.shape[-2]:, None])
+        q = rearrange(q, '... (h d) -> ... h d', h=self.num_heads)
+        if self.num_kv_groups > 1:
+            k, v = (repeat(x, '... (h d) -> ... (h g) d', h=self.num_kv_heads, g=self.num_kv_groups) for x in (k, v))
+        else:
+            k, v = (rearrange(x, '... (h d) -> ... h d', h=self.num_kv_heads) for x in (k, v))
+        gk = F.logsigmoid(gk) / self.gate_logit_normalizer
+
+        recurrent_state = last_state['recurrent_state'] if last_state is not None else None
+        if mode == 'chunk':
+            o, recurrent_state = chunk_simple_gla(
+                q=q,
+                k=k,
+                v=v,
+                g=gk,
+                initial_state=recurrent_state,
+                output_final_state=use_cache,
+                cu_seqlens=cu_seqlens,
+            )
+        elif mode == 'fused_recurrent':
+            o, recurrent_state = fused_recurrent_simple_gla(
+                q=q,
+                k=k,
+                v=v,
+                g=gk,
+                initial_state=recurrent_state,
+                output_final_state=use_cache,
+                cu_seqlens=cu_seqlens,
+            )
+        else:
+            raise NotImplementedError(f"Not supported mode `{mode}`.")
+
+        update_layer_cache(
+            self,
+            past_key_values,
+            recurrent_state=recurrent_state,
+            conv_state=(conv_state_q, conv_state_k, conv_state_v) if self.use_short_conv else None,
+            offset=q.shape[1],
+        )
+
+        g = self.g_proj(hidden_states)
+        if self.fuse_norm_and_gate:
+            g = rearrange(g, 'b t (h d) -> b t h d', h=self.num_heads)
+            o = self.g_norm_swish_gate(o, g)
+            o = rearrange(o, 'b t h d -> b t (h d)')
+        else:
+            o = rearrange(self.g_norm(o), 'b t h d -> b t (h d)')
+            o = o * self.gate_fn(g)
+        o = self.o_proj(o)
+
+        return o, None, past_key_values
+
+    def state_size(self, **kwargs) -> int:
+        state_size = self.key_dim * self.head_v_dim
+        for module in self.children():
+            if isinstance(module, ShortConvolution):
+                state_size += module.state_size
+        return state_size

fla/layers/utils.py

@@ -0,0 +1,218 @@
+# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang
+
+# Code is adapted from flash-attn.bert_padding.py
+
+
+import torch
+from einops import rearrange, repeat
+
+from fla.ops.utils.index import prepare_cu_seqlens_from_mask, prepare_lens_from_mask
+from fla.utils import tensor_cache
+
+_LAYER_IDX_REQUIRED_MSG = "{cls} requires `layer_idx` when `past_key_values` is provided."
+
+
+class IndexFirstAxis(torch.autograd.Function):
+
+    @staticmethod
+    def forward(ctx, x, indices):
+        ctx.save_for_backward(indices)
+        assert x.ndim >= 2
+        ctx.first_axis_dim, other_shape = x.shape[0], x.shape[1:]
+        second_dim = other_shape.numel()
+        # TD [2022-03-04] For some reason torch.gather is a bit faster than indexing.
+        # return x[indices]
+        return torch.gather(
+            rearrange(x, "b ... -> b (...)"), 0, repeat(indices, "z -> z d", d=second_dim),
+        ).reshape(-1, *other_shape)
+
+    @staticmethod
+    def backward(ctx, do):
+        (indices,) = ctx.saved_tensors
+        assert do.ndim >= 2
+        other_shape = do.shape[1:]
+        do = rearrange(do, "b ... -> b (...)")
+        dx = torch.zeros(
+            [ctx.first_axis_dim, do.shape[1]],
+            device=do.device,
+            dtype=do.dtype,
+        )
+        # TD [2022-03-04] For some reason torch.scatter is a bit faster than indexing.
+        # dx[indices] = do
+        dx.scatter_(0, repeat(indices, "z -> z d", d=do.shape[1]), do)
+        return dx.reshape(ctx.first_axis_dim, *other_shape), None
+
+
+index_first_axis = IndexFirstAxis.apply
+
+
+class IndexPutFirstAxis(torch.autograd.Function):
+
+    @staticmethod
+    def forward(ctx, x, indices, first_axis_dim):
+        ctx.save_for_backward(indices)
+        assert indices.ndim == 1
+        assert x.ndim >= 2
+        y = torch.zeros(first_axis_dim, *x.shape[1:], device=x.device, dtype=x.dtype)
+        # TODO [2022-03-04] For some reason torch.scatter is a bit faster than indexing.
+        y[indices] = x
+        # y.scatter_(0, repeat(indices, 'z -> z d', d=x.shape[1]), x)
+        return y
+
+    @staticmethod
+    def backward(ctx, do):
+        (indices,) = ctx.saved_tensors
+        # TODO [2022-03-04] For some reason torch.gather is a bit faster than indexing.
+        dx = do[indices]
+        # dx = torch.gather(do, 0, repeat(indices, 'z -> z d', d=do.shape[1]))
+        return dx, None, None
+
+
+index_put_first_axis = IndexPutFirstAxis.apply
+
+
+@tensor_cache
+def get_unpad_data(
+    attention_mask: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor, int]:
+    """
+    Retrieves indexing data required to repad unpadded (ragged) tensors.
+
+    Args:
+        attention_mask (`torch.Tensor`):
+            Boolean or int tensor of shape (batch_size, sequence_length), 1 means valid and 0 means not valid.
+
+    Return:
+        indices (`torch.Tensor`):
+            The indices of non-masked tokens from the flattened input sequence.
+        cu_seqlens (`torch.Tensor`):
+            The cumulative sequence lengths, used to index into ragged (unpadded) tensors.
+            `cu_seqlens` shape is [batch_size + 1].
+        max_seqlen_in_batch (`int`):
+            Maximum sequence length in batch.
+    """
+    lens = prepare_lens_from_mask(attention_mask)
+    indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten()
+    max_seqlen_in_batch = lens.max().item()
+    cu_seqlens = prepare_cu_seqlens_from_mask(attention_mask)
+    return indices, cu_seqlens, max_seqlen_in_batch
+
+
+def unpad_input(
+    q: torch.Tensor,
+    states: tuple[torch.Tensor],
+    attention_mask: torch.Tensor,
+    q_len: int,
+    keepdim: bool = False,
+):
+    """
+    Unpads query, key, and values tensors, using a single dimension for all tokens
+    even though they belong to different batches.
+
+
+    Arguments:
+        q (`torch.Tensor`):
+            Query state with padding. Shape: [batch_size, q_len, ...].
+        states (`Tuple[torch.Tensor]`):
+            Attention state with padding. Shape: [batch_size, seq_len, ...].
+        attention_mask (`torch.Tensor`):
+            Boolean or int tensor of shape [batch_size, sequence_length], 1 means valid and 0 means not valid.
+        q_len (`int`):
+            Target length.
+        keepdim (`bool`):
+            Whether to keep the batch dimension. Default: `False`.
+
+    Return:
+        q (`torch.Tensor`):
+            Query state without padding.
+            Shape: [1, total_target_length, ...] if `keepdim=True` else [total_target_length, ...].
+        states (`Tuple[torch.Tensor]`):
+            Attention state without padding.
+            Shape: [1, total_source_length, ...] if `keepdim=True` else [total_source_length, ...].
+        indices_q (`torch.Tensor`):
+            The indices of non-masked tokens from the flattened input target sequence.
+        (cu_seqlens_q, cu_seqlens_k) (`Tuple[int]`):
+            The cumulative sequence lengths for the target (query) and source (key, value),
+            used to index into ragged (unpadded) tensors.
+            `cu_seqlens` shape is [batch_size + 1].
+        (max_seqlen_in_batch_q, max_seqlen_in_batch_k) (`Tuple[int]`):
+            Maximum sequence length in batch (`max_seqlen_in_batch_q` for the target sequence
+            i.e. query, `max_seqlen_in_batch_k` for the source sequence i.e. key/value).
+    """
+    indices_k, cu_seqlens_k, max_seqlen_in_batch_k = get_unpad_data(attention_mask)
+    batch_size, seq_len, *_ = states[0].shape
+
+    state = tuple(
+        index_first_axis(rearrange(s, "b s ... -> (b s) ..."), indices_k)
+        for s in states
+    )
+
+    if q_len == seq_len:
+        q = index_first_axis(rearrange(q, "b s ... -> (b s) ..."), indices_k)
+        cu_seqlens_q = cu_seqlens_k
+        max_seqlen_in_batch_q = max_seqlen_in_batch_k
+        indices_q = indices_k
+    elif q_len == 1:
+        max_seqlen_in_batch_q = 1
+        cu_seqlens_q = torch.arange(batch_size + 1, dtype=torch.int32, device=q.device)
+        indices_q = cu_seqlens_q[:-1]
+        q = q.squeeze(1)
+    else:
+        raise NotImplementedError("We only support either q_len == k_len (prefilling) or q_len == 1 (decoding)")
+
+    if keepdim:
+        q = q.unsqueeze(0)
+        state = tuple(s.unsqueeze(0) for s in state)
+
+    return (
+        q,
+        state,
+        indices_q,
+        (cu_seqlens_q, cu_seqlens_k),
+        (max_seqlen_in_batch_q, max_seqlen_in_batch_k),
+    )
+
+
+def pad_input(
+    hidden_states: torch.Tensor,
+    indices: torch.LongTensor,
+    batch_size: int,
+    seq_len: int,
+) -> torch.Tensor:
+    """
+    Args:
+        hidden_states ([total_tokens, ...]):
+            where total_tokens denotes the number of tokens in selected in attention_mask.
+        indices ([total_tokens]):
+            the indices that represent the non-masked tokens of the original padded input sequence.
+        batch_size (int):
+            batch_size size for the padded sequence.
+        seq_len (int):
+            maximum sequence length for the padded sequence.
+
+    Return:
+        hidden_states of shape [batch_size, seq_len, ...]
+    """
+    output = index_put_first_axis(hidden_states, indices, batch_size * seq_len)
+    return rearrange(output, "(b s) ... -> b s ...", b=batch_size)
+
+
+def require_cache_layer_idx(module, past_key_values):
+    layer_idx = getattr(module, "layer_idx", None)
+    if past_key_values is not None and layer_idx is None:
+        raise ValueError(_LAYER_IDX_REQUIRED_MSG.format(cls=module.__class__.__name__))
+    return layer_idx
+
+
+def get_layer_cache(module, past_key_values):
+    layer_idx = require_cache_layer_idx(module, past_key_values)
+    if past_key_values is not None and len(past_key_values) > layer_idx:
+        return past_key_values[layer_idx]
+    return None
+
+
+def update_layer_cache(module, past_key_values, **kwargs):
+    layer_idx = require_cache_layer_idx(module, past_key_values)
+    if past_key_values is not None:
+        return past_key_values.update(layer_idx=layer_idx, **kwargs)
+    return None

fla/models/__init__.py

@@ -0,0 +1,3 @@
+# Keep model package imports lazy.
+__all__ = []
+

fla/models/abc/__init__.py

@@ -0,0 +1,12 @@
+
+from transformers import AutoConfig, AutoModel, AutoModelForCausalLM
+
+from fla.models.abc.configuration_abc import ABCConfig
+from fla.models.abc.modeling_abc import ABCForCausalLM, ABCModel
+
+AutoConfig.register(ABCConfig.model_type, ABCConfig, exist_ok=True)
+AutoModel.register(ABCConfig, ABCModel, exist_ok=True)
+AutoModelForCausalLM.register(ABCConfig, ABCForCausalLM, exist_ok=True)
+
+
+__all__ = ['ABCConfig', 'ABCForCausalLM', 'ABCModel']

fla/models/abc/configuration_abc.py

@@ -0,0 +1,105 @@
+
+import warnings
+
+from transformers.configuration_utils import PretrainedConfig
+
+
+class ABCConfig(PretrainedConfig):
+
+    model_type = 'abc'
+    keys_to_ignore_at_inference = ['past_key_values']
+
+    def __init__(
+        self,
+        hidden_size: int = 2048,
+        gate_low_rank_dim: int = 16,
+        clamp_min: float = -32,
+        clamp_max: float = 32,
+        hidden_ratio: int | None = 4,
+        intermediate_size: int | None = None,
+        num_hidden_layers: int = 24,
+        num_heads: int = 4,
+        num_slots: int | None = 64,
+        use_short_conv: bool = False,
+        conv_size: int = 4,
+        exapnd_k: float = 0.5,
+        exapnd_v: float = 1,
+        hidden_act: str = "swish",
+        max_position_embeddings: int = 2048,
+        elementwise_affine: bool | None = True,
+        norm_eps: float = 1e-6,
+        use_rope: bool = True,
+        attn: dict | None = None,
+        use_cache: bool = True,
+        pad_token_id: int | None = None,
+        bos_token_id: int = 1,
+        eos_token_id: int = 2,
+        tie_word_embeddings: bool = False,
+        initializer_range: float = 0.02,
+        fuse_norm: bool = True,
+        fuse_swiglu: bool = True,
+        fuse_cross_entropy: bool = True,
+        fuse_linear_cross_entropy: bool = False,
+        use_l2warp: bool = False,
+        vocab_size: int = 32000,
+        **kwargs,
+    ):
+        self.hidden_size = hidden_size
+        self.gate_low_rank_dim = gate_low_rank_dim
+        self.clamp_min = clamp_min
+        self.clamp_max = clamp_max
+        self.hidden_ratio = hidden_ratio
+        self.intermediate_size = intermediate_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_heads = num_heads
+        self.num_slots = num_slots
+        self.use_short_conv = use_short_conv
+        self.conv_size = conv_size
+        self.expand_k = exapnd_k
+        self.expand_v = exapnd_v
+        self.hidden_act = hidden_act
+        self.max_position_embeddings = max_position_embeddings
+        self.elementwise_affine = elementwise_affine
+        self.norm_eps = norm_eps
+        self.use_rope = use_rope
+        self.attn = attn
+        self.use_cache = use_cache
+        self.initializer_range = initializer_range
+
+        self.fuse_norm = fuse_norm
+        self.fuse_swiglu = fuse_swiglu
+        self.fuse_cross_entropy = fuse_cross_entropy
+        self.fuse_linear_cross_entropy = fuse_linear_cross_entropy
+        self.use_l2warp = use_l2warp
+        self.vocab_size = vocab_size
+
+        if fuse_cross_entropy and fuse_linear_cross_entropy:
+            raise ValueError(
+                "`fuse_cross_entropy` and `fuse_linear_cross_entropy` cannot be True at the same time.",
+            )
+        if fuse_linear_cross_entropy:
+            warnings.warn(
+                "`fuse_linear_cross_entropy` is enabled, which can improves memory efficiency "
+                "at the potential cost of reduced precision. "
+                "If you observe issues like loss divergence, consider disabling this setting.",
+            )
+
+        if attn is not None:
+            if not isinstance(attn, dict):
+                raise ValueError("attn must be a dictionary")
+            if 'layers' not in attn:
+                raise ValueError("Layer indices must be provided to initialize hybrid attention layers")
+            if 'num_heads' not in attn:
+                raise ValueError("Number of heads must be provided to initialize hybrid attention layers")
+            attn['num_kv_heads'] = attn.get('num_kv_heads', attn['num_heads'])
+            attn['qkv_bias'] = attn.get('qkv_bias', False)
+            attn['window_size'] = attn.get('window_size', None)
+            attn['rope_theta'] = attn.get('rope_theta', 10000.)
+
+        super().__init__(
+            pad_token_id=pad_token_id,
+            bos_token_id=bos_token_id,
+            eos_token_id=eos_token_id,
+            tie_word_embeddings=tie_word_embeddings,
+            **kwargs,
+        )

fla/models/abc/modeling_abc.py

@@ -0,0 +1,371 @@
+from __future__ import annotations
+
+import math
+import warnings
+from typing import TYPE_CHECKING, Optional
+
+import torch
+import torch.nn as nn
+from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
+from transformers.modeling_utils import PreTrainedModel
+from transformers.utils import logging
+from transformers.utils.deprecation import deprecate_kwarg
+
+from fla.layers.abc import ABCAttention
+from fla.layers.attn import Attention
+from fla.models.abc.configuration_abc import ABCConfig
+from fla.models.utils import Cache, FLAGenerationMixin
+from fla.modules import FusedCrossEntropyLoss, FusedLinearCrossEntropyLoss, RMSNorm
+from fla.modules import GatedMLP as ABCMLP
+from fla.modules.l2warp import l2_warp
+
+try:
+    from transformers.modeling_layers import GradientCheckpointingLayer
+except ImportError:
+    from fla.models.modeling_layers import GradientCheckpointingLayer
+
+logger = logging.get_logger(__name__)
+
+if TYPE_CHECKING:
+    from transformers.processing_utils import Unpack
+
+
+class ABCBlock(GradientCheckpointingLayer):
+
+    def __init__(self, config: ABCConfig, layer_idx: int):
+        super().__init__()
+
+        self.config = config
+        self.layer_idx = layer_idx
+
+        self.attn_norm = (RMSNorm if config.fuse_norm else nn.RMSNorm)(config.hidden_size, eps=config.norm_eps)
+        if config.attn is not None and layer_idx in config.attn['layers']:
+            self.attn = Attention(
+                hidden_size=config.hidden_size,
+                num_heads=config.attn['num_heads'],
+                num_kv_heads=config.attn['num_kv_heads'],
+                qkv_bias=config.attn['qkv_bias'],
+                window_size=config.attn['window_size'],
+                rope_theta=config.attn['rope_theta'],
+                max_position_embeddings=config.max_position_embeddings,
+                layer_idx=layer_idx,
+            )
+        else:
+            self.attn = ABCAttention(
+                hidden_size=config.hidden_size,
+                expand_k=config.expand_k,
+                expand_v=config.expand_v,
+                num_heads=config.num_heads,
+                num_slots=config.num_slots,
+                use_short_conv=config.use_short_conv,
+                conv_size=config.conv_size,
+                gate_fn=config.hidden_act,
+                elementwise_affine=config.elementwise_affine,
+                norm_eps=config.norm_eps,
+                use_rope=config.use_rope,
+                clamp_min=config.clamp_min,
+                clamp_max=config.clamp_max,
+                fuse_norm=config.fuse_norm,
+                layer_idx=layer_idx,
+            )
+        self.mlp_norm = (RMSNorm if config.fuse_norm else nn.RMSNorm)(config.hidden_size, eps=config.norm_eps)
+        self.mlp = ABCMLP(
+            hidden_size=config.hidden_size,
+            hidden_ratio=config.hidden_ratio,
+            intermediate_size=config.intermediate_size,
+            hidden_act=config.hidden_act,
+            fuse_swiglu=config.fuse_swiglu,
+        )
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+        past_key_values: Cache | list[torch.FloatTensor] | None = None,
+        use_cache: bool | None = False,
+        output_attentions: bool | None = False,
+        **kwargs: Unpack[dict],
+    ) -> tuple[torch.FloatTensor, tuple[torch.FloatTensor, torch.FloatTensor] | None]:
+
+        residual = hidden_states
+
+        hidden_states = self.attn_norm(hidden_states)
+        hidden_states, attentions, past_key_values = self.attn(
+            hidden_states=hidden_states,
+            attention_mask=attention_mask,
+            past_key_values=past_key_values,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            **kwargs,
+        )
+        if self.config.fuse_norm:
+            hidden_states, residual = self.mlp_norm(hidden_states, residual, True)
+        else:
+            hidden_states = residual + hidden_states
+            residual = hidden_states
+            hidden_states = self.mlp_norm(hidden_states)
+        hidden_states = self.mlp(hidden_states)
+        hidden_states = residual + hidden_states
+
+        outputs = (hidden_states, attentions, past_key_values)
+
+        return outputs
+
+
+class ABCPreTrainedModel(PreTrainedModel):
+
+    config_class = ABCConfig
+    base_model_prefix = 'model'
+    supports_gradient_checkpointing = True
+    _no_split_modules = ['ABCBlock']
+    _supports_cache_class = True
+
+    def __init__(self, *inputs, **kwargs):
+        super().__init__(*inputs, **kwargs)
+
+    def _init_weights(
+        self,
+        module: nn.Module,
+        prenorm_residual_strategy: str | None = None,
+        num_residuals_per_layer: int = 2,
+    ):
+        if isinstance(module, (nn.Linear, nn.Conv1d)):
+            # Slightly different from the TF version which uses truncated_normal for initialization
+            # cf https://github.com/pytorch/pytorch/pull/5617
+            nn.init.normal_(module.weight, mean=0.0, std=self.config.initializer_range)
+            if module.bias is not None:
+                nn.init.zeros_(module.bias)
+        elif isinstance(module, nn.Embedding):
+            nn.init.normal_(module.weight, mean=0.0, std=self.config.initializer_range)
+        elif hasattr(module, 'reset_parameters'):
+            module.reset_parameters()
+
+        if prenorm_residual_strategy is not None:
+            # Reinitialize selected weights subject to the OpenAI GPT-2 Paper Scheme:
+            #   > A modified initialization which accounts for the accumulation on the residual path with model depth. Scale
+            #   > the weights of residual layers at initialization by a factor of 1/√N where N is the # of residual layers.
+            #   >   -- GPT-2 :: https://openai.com/blog/better-language-models/
+            #
+            # Reference (Megatron-LM): https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/model/gpt_model.py
+            p = None
+            if hasattr(module, 'o_proj'):
+                p = module.o_proj.weight
+            elif hasattr(module, 'down_proj'):
+                p = module.down_proj.weight
+            if p is not None:
+                # Special Scaled Initialization --> There are 2 Layer Norms per Transformer Block
+                # Following Pytorch init, except scale by 1/sqrt(2 * n_layer)
+                # We need to reinit p since this code could be called multiple times
+                # Having just p *= scale would repeatedly scale it down
+                if prenorm_residual_strategy == 'rescale':
+                    nn.init.kaiming_uniform_(p, a=math.sqrt(5))
+                    with torch.no_grad():
+                        p /= math.sqrt(num_residuals_per_layer * self.config.num_hidden_layers)
+                elif prenorm_residual_strategy == 'zero':
+                    nn.init.zeros_(p)
+                else:
+                    raise ValueError(f"Invalid prenorm_residual_strategy: {prenorm_residual_strategy}")
+
+
+class ABCModel(ABCPreTrainedModel):
+
+    def __init__(self, config: ABCConfig):
+        super().__init__(config)
+        self.padding_idx = config.pad_token_id
+        self.vocab_size = config.vocab_size
+
+        self.embeddings = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
+        self.layers = nn.ModuleList([ABCBlock(config, layer_idx) for layer_idx in range(config.num_hidden_layers)])
+        self.norm = (RMSNorm if config.fuse_norm else nn.RMSNorm)(config.hidden_size, eps=config.norm_eps)
+
+        self.gradient_checkpointing = False
+
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.embeddings
+
+    def set_input_embeddings(self, value):
+        self.embeddings = value
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor | None = None,
+        attention_mask: Optional[torch.Tensor] = None,  # noqa
+        inputs_embeds: torch.FloatTensor | None = None,
+        past_key_values: Cache | list[torch.FloatTensor] | None = None,
+        use_cache: bool | None = None,
+        output_attentions: bool | None = None,
+        output_hidden_states: bool | None = None,
+        return_dict: bool | None = None,
+        **kwargs: Unpack[dict],
+    ) -> tuple | BaseModelOutputWithPast:
+        if output_attentions:
+            warnings.warn("`ABCModel` does not `output_attentions` now, setting it to `False`.")
+            output_attentions = False
+        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 if not self.training else False)
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        # retrieve input_ids and inputs_embeds
+        if input_ids is not None and inputs_embeds is not None:
+            raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
+        if input_ids is None and inputs_embeds is None:
+            raise ValueError("You have to specify either input_ids or inputs_embeds")
+
+        if inputs_embeds is None:
+            inputs_embeds = self.embeddings(input_ids)
+        hidden_states = inputs_embeds
+
+        if use_cache and not isinstance(past_key_values, Cache):
+            past_key_values = Cache.from_legacy_cache(past_key_values)
+
+        all_hidden_states = () if output_hidden_states else None
+        all_attns = () if output_attentions else None
+        for layer in self.layers:
+            if output_hidden_states:
+                all_hidden_states += (hidden_states,)
+
+            hidden_states, attentions, past_key_values = layer(
+                hidden_states,
+                attention_mask,
+                past_key_values=past_key_values,
+                use_cache=use_cache,
+                output_attentions=output_attentions,
+                **kwargs,
+            )
+
+            if output_attentions:
+                all_attns += (attentions,)
+
+        hidden_states = self.norm(hidden_states)
+
+        # add hidden states from the last decoder layer
+        if output_hidden_states:
+            all_hidden_states += (hidden_states,)
+
+        if not return_dict:
+            return tuple(i for i in [hidden_states, past_key_values, all_hidden_states, all_attns] if i is not None)
+        return BaseModelOutputWithPast(
+            last_hidden_state=hidden_states,
+            past_key_values=past_key_values,
+            hidden_states=all_hidden_states,
+            attentions=all_attns,
+        )
+
+
+class ABCForCausalLM(ABCPreTrainedModel, FLAGenerationMixin):
+
+    _tied_weights_keys = ["lm_head.weight"]
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.model = ABCModel(config)
+        self.vocab_size = config.vocab_size
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+        self.criterion = None
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.model.embeddings
+
+    def set_input_embeddings(self, value):
+        self.model.embeddings = value
+
+    def get_output_embeddings(self):
+        return self.lm_head
+
+    def set_output_embeddings(self, new_embeddings):
+        self.lm_head = new_embeddings
+
+    def set_decoder(self, decoder):
+        self.model = decoder
+
+    def get_decoder(self):
+        return self.model
+
+    def generate(self, *args, **kwargs):
+        try:
+            return super().generate(*args, **kwargs)
+        except AttributeError as exception:
+            if 'past_key_values' in str(exception):
+                raise AttributeError(
+                    f"You tried to call `generate` with a decoding strategy that manipulates `past_key_values`, "
+                    f"which is not supported for {self.__class__.__name__}. "
+                    f"Try another generation strategy instead. "
+                    f"For the available generation strategies, check this doc: "
+                    f"https://huggingface.co/docs/transformers/en/generation_strategies#decoding-strategies",
+                )
+            else:
+                raise exception
+
+    @deprecate_kwarg("num_logits_to_keep", version="4.50", new_name="logits_to_keep")
+    def forward(
+        self,
+        input_ids: torch.LongTensor = None,
+        attention_mask: torch.Tensor | None = None,
+        inputs_embeds: torch.Tensor | None = None,
+        past_key_values: Cache | list[torch.FloatTensor] | None = None,
+        labels: torch.LongTensor | None = None,
+        use_cache: bool | None = None,
+        output_attentions: bool | None = None,
+        output_hidden_states: bool | None = None,
+        return_dict: bool | None = None,
+        logits_to_keep: int | None = 0,
+        **kwargs: Unpack[dict],
+    ) -> tuple | CausalLMOutputWithPast:
+        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
+        )
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        outputs = self.model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            inputs_embeds=inputs_embeds,
+            past_key_values=past_key_values,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+            **kwargs,
+        )
+
+        hidden_states = outputs[0]
+
+        loss, logits = None, None
+        if not self.config.fuse_linear_cross_entropy or labels is None:
+            logits = self.lm_head(hidden_states if logits_to_keep is None else hidden_states[:, -logits_to_keep:])
+        if labels is not None:
+            if getattr(self, 'criterion', None) is None:
+                if self.config.fuse_linear_cross_entropy:
+                    criterion = FusedLinearCrossEntropyLoss(use_l2warp=self.config.use_l2warp)
+                elif self.config.fuse_cross_entropy:
+                    criterion = FusedCrossEntropyLoss(inplace_backward=True)
+                else:
+                    criterion = nn.CrossEntropyLoss()
+            else:
+                criterion = self.criterion
+            labels = labels.to(hidden_states.device)
+            labels = torch.cat((labels[..., 1:], torch.full_like(labels[:, :1], criterion.ignore_index)), 1)
+            if self.config.fuse_linear_cross_entropy:
+                loss = criterion(hidden_states, labels, self.lm_head.weight, self.lm_head.bias)
+            else:
+                loss = criterion(logits.view(labels.numel(), -1), labels.view(-1))
+                loss = l2_warp(loss, logits) if self.config.use_l2warp else loss
+
+        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,
+        )

fla/models/bitnet/__init__.py

@@ -0,0 +1,12 @@
+
+from transformers import AutoConfig, AutoModel, AutoModelForCausalLM
+
+from fla.models.bitnet.configuration_bitnet import BitNetConfig
+from fla.models.bitnet.modeling_bitnet import BitNetForCausalLM, BitNetModel
+
+AutoConfig.register(BitNetConfig.model_type, BitNetConfig, exist_ok=True)
+AutoModel.register(BitNetConfig, BitNetModel, exist_ok=True)
+AutoModelForCausalLM.register(BitNetConfig, BitNetForCausalLM, exist_ok=True)
+
+
+__all__ = ['BitNetConfig', 'BitNetForCausalLM', 'BitNetModel']