Add custom GatedDeltaProduct implementation with state weaving

src/models/blocks.py

@@ -1,7 +1,10 @@
 import torch
 import torch.nn as nn
-from fla.models.gated_deltaproduct import GatedDeltaProductConfig
-from fla.models.gated_deltaproduct.modeling_gated_deltaproduct import GatedDeltaProductBlock
+
+from src.models.gated_deltaproduct import GatedDeltaProductConfig
+from src.models.gated_deltaproduct.modeling_gated_deltaproduct import (
+    GatedDeltaProductBlock,
+)
 
 
 class GatedDeltaProductEncoder(nn.Module):

src/models/gated_deltaproduct/README.md

@@ -0,0 +1,344 @@
+# Custom GatedDeltaProduct Implementation
+
+This directory contains a custom implementation of the GatedDeltaProduct layer, based on the [Flash Linear Attention (FLA)](https://github.com/fla-org/flash-linear-attention) library, with modifications specifically designed for **time series forecasting** tasks.
+
+## Overview
+
+Our custom implementation adds **hidden state weaving** functionality that enables information to flow across encoder layers, maintaining temporal continuity - a crucial feature for time series forecasting that differs from the general-purpose language modeling focus of the official FLA implementation.
+
+## Reference
+
+This implementation is based on:
+- **Official FLA Repository**: [https://github.com/fla-org/flash-linear-attention](https://github.com/fla-org/flash-linear-attention)
+- **Original Paper**: [DeltaProduct: Improving State-Tracking in Linear RNNs via Householder Products](https://arxiv.org/html/2502.10297v3) (Siems et al., 2025)
+
+---
+
+## What is DeltaProduct?
+
+DeltaProduct is a linear RNN architecture that uses **diagonal plus rank-nₕ** state-transition matrices, formed as products of `nₕ` generalized Householder transformations. This provides a tunable mechanism to balance expressivity and efficiency compared to diagonal-only architectures like Mamba or GLA.
+
+### Key Concepts
+
+- **Householder transformations**: Enable simultaneous token-channel mixing, overcoming the expressivity limitations of purely diagonal state-transition matrices
+- **Rank-nₕ structure**: Allows better expressivity than rank-1 (DeltaNet) while maintaining training efficiency. The parameter `nₕ` (number of Householder transformations) provides a tunable trade-off between expressivity and computational cost
+- **Gated variant**: Adds gating mechanisms for improved performance, allowing the model to control information flow through forget gates and output gates
+
+### Architecture Overview
+
+DeltaProduct improves upon earlier linear RNN architectures:
+
+- **Diagonal architectures** (Mamba, GLA, mLSTM): Use diagonal state-transition matrices for fast runtime but suffer from limited expressivity
+- **Rank-1 architectures** (DeltaNet, RWKV-7): Use diagonal plus rank-1 structure, enabling simultaneous token-channel mixing with only a slight decrease in training efficiency
+- **DeltaProduct**: Extends this to diagonal plus rank-nₕ structure, where multiple Householder transformations (nₕ ≥ 1) provide greater expressivity while maintaining computational efficiency
+
+The architecture interprets DeltaNet's recurrence as performing one step of online gradient descent per token on an associative recall loss. DeltaProduct instead takes multiple (`nₕ`) steps per token, naturally leading to the rank-nₕ structure.
+
+---
+
+## State Weaving Mechanism
+
+Unlike DeltaProduct's original design for autoregressive language modeling, time series forecasting across a full horizon does not require causal masking. To exploit this property, we introduce **state weaving**, a mechanism that enables bidirectional information flow across the entire sequence length without additional parameters or computational overhead.
+
+<div align="center">
+  <img src="https://iili.io/Ks86Z0X.png" alt="State Weaving Architecture" width="450"/>
+</div>
+
+*Figure: The TempoPFN architecture using stacked GatedDeltaProduct blocks with learnable initial states H₀ⁱ and state-weaving. The final hidden state of each layer Hₜⁱ is added to the learnable initial state of the next layer H₀ⁱ⁺¹, enabling bidirectional information flow.*
+
+### How State Weaving Works
+
+In our implementation, state weaving operates as follows:
+
+1. **Learnable Initial States**: Each encoder layer `i` has a learnable initial hidden state `H₀ⁱ` that is optimized during training.
+
+2. **State Propagation**: The final hidden state from layer `i`, denoted `Hₜⁱ`, is propagated forward and combined with the learnable initial state of the next layer:
+   ```
+   H₀ⁱ⁺¹ = H₀ⁱ⁺¹ + Hₜⁱ
+   ```
+
+3. **Bidirectional Information Flow**: This mechanism effectively lifts the causal constraint while maintaining computational efficiency. Information from later tokens can influence earlier layers through the accumulated hidden states, enabling the model to process the entire sequence (history + future horizon) coherently.
+
+4. **No Extra Overhead**: Unlike explicit bidirectional architectures, state weaving requires no additional parameters or computational overhead beyond the existing forward pass.
+
+This design is particularly powerful for time series forecasting, where:
+- The full prediction horizon is known at inference time
+- Coherent predictions across all future time steps are desired
+- Historical context should inform all future predictions simultaneously
+
+---
+
+## Key Differences from Official FLA
+
+### 1. **`initial_state` Parameter in Forward Method**
+
+#### Official FLA (`fla/layers/gated_deltaproduct.py`)
+```python
+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]:
+```
+**No `initial_state` parameter** - The official implementation only uses `recurrent_state` from `past_key_values`.
+
+#### Our Custom Implementation (`gated_deltaproduct.py`)
+```python
+def forward(
+    self,
+    hidden_states: torch.Tensor,
+    attention_mask: Optional[torch.Tensor] = None,
+    past_key_values: Optional[Cache] = None,
+    initial_state: Optional[torch.Tensor] = None,  # ← ADDED
+    use_cache: Optional[bool] = False,
+    output_attentions: Optional[bool] = False,
+    **kwargs: Unpack[Dict],
+) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Cache]]:
+```
+**Added `initial_state` parameter** - Allows external control of the initial recurrent state, enabling layer-to-layer state propagation.
+
+---
+
+### 2. **Usage of `initial_state` in Chunk Mode**
+
+#### Official FLA
+```python
+if mode == 'chunk':
+    o, recurrent_state = chunk_gated_delta_product(
+        q=q, k=k, v=v, g=g, beta=beta,
+        initial_state=recurrent_state,  # ← Only from past_key_values
+        output_final_state=use_cache,
+        cu_seqlens=cu_seqlens,
+        num_householder=self.num_householder,
+        use_qk_l2norm_in_kernel=True,
+    )
+```
+
+#### Our Custom Implementation
+```python
+if mode == "chunk":
+    o, recurrent_state = chunk_gated_delta_product(
+        q=q, k=k, v=v, g=g, beta=beta,
+        initial_state=initial_state,  # ← Uses external initial_state if provided
+        output_final_state=output_attentions,
+        cu_seqlens=cu_seqlens,
+        num_householder=self.num_householder,
+        use_qk_l2norm_in_kernel=True,
+    )
+```
+
+**Key Difference**: Our implementation prioritizes the externally provided `initial_state` over `recurrent_state` from `past_key_values`, enabling layer-to-layer state propagation.
+
+---
+
+### 3. **Return Value: Hidden State Output**
+
+#### Official FLA (`fla/models/gated_deltaproduct/modeling_gated_deltaproduct.py`)
+```python
+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]:
+    # ...
+    return outputs  # Returns (hidden_states, attentions, past_key_values)
+```
+
+**No `initial_state` parameter** - The block doesn't accept or return hidden states explicitly.
+
+#### Our Custom Implementation (`modeling_gated_deltaproduct.py`)
+```python
+def forward(
+    self,
+    hidden_states: torch.Tensor,
+    attention_mask: Optional[torch.Tensor] = None,
+    past_key_values: Optional[Union[Cache, List[torch.FloatTensor]]] = None,
+    use_cache: Optional[bool] = False,
+    output_attentions: Optional[bool] = False,
+    initial_state: Optional[torch.FloatTensor] = None,  # ← ADDED
+    **kwargs: Unpack[Dict],
+) -> Tuple[
+    torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]
+]:
+    # ...
+    hidden_states, attentions, past_key_values = self.attn(
+        # ...
+        initial_state=initial_state,  # ← Passed through
+        **kwargs,
+    )
+    # ...
+    return outputs  # Returns (hidden_states, attentions, past_key_values)
+```
+
+**Added `initial_state` parameter** - The block accepts and forwards `initial_state` to the attention layer.
+
+---
+
+### 4. **Hidden State Weaving Implementation**
+
+Our implementation supports two modes of hidden state weaving (controlled by the `weaving` parameter in encoder config):
+
+#### **Mode 1: Weaving Enabled (`weaving=True`)** - Default
+```python
+if self.encoder_config.get("weaving", True):
+    # initial hidden state is learnable
+    hidden_state = torch.zeros_like(
+        self.initial_hidden_state[0].repeat(batch_size * num_channels, 1, 1, 1)
+    )
+    for layer_idx, encoder_layer in enumerate(self.encoder_layers):
+        x, hidden_state = encoder_layer(
+            x,
+            hidden_state + self.initial_hidden_state[layer_idx].repeat(
+                batch_size * num_channels, 1, 1, 1
+            ),
+        )
+```
+
+**Key Features**:
+- Hidden state accumulates across layers
+- Each layer receives: `previous_hidden_state + learnable_initial_state[layer_idx]`
+- State persists between layers, allowing information to flow through the network
+
+#### **Mode 2: No Weaving (`weaving=False`)**
+```python
+else:
+    # initial hidden state is separately learnable for each layer
+    for layer_idx, encoder_layer in enumerate(self.encoder_layers):
+        initial_hidden_state = self.initial_hidden_state[layer_idx].repeat(
+            batch_size * num_channels, 1, 1, 1
+        )
+        x, _ = encoder_layer(x, initial_hidden_state)
+```
+
+**Key Features**:
+- Each layer uses its own independent learnable initial state
+- No accumulation between layers
+- Hidden state is discarded after each layer
+
+---
+
+### 5. **Learnable Initial Hidden States**
+
+Our implementation includes learnable initial states managed at the model level:
+
+```python
+num_initial_hidden_states = self.num_encoder_layers
+self.initial_hidden_state = nn.ParameterList(
+    [
+        nn.Parameter(
+            torch.randn(
+                1, self.encoder_config["num_heads"], head_k_dim, head_v_dim
+            )
+            / head_k_dim,
+            requires_grad=True,
+        )
+        for _ in range(num_initial_hidden_states)
+    ]
+)
+```
+
+**Key Features**:
+- One learnable parameter per encoder layer
+- Shape: `[1, num_heads, head_k_dim, head_v_dim]`
+- Initialized with small random values scaled by `head_k_dim`
+- These are trainable parameters that can be optimized during training
+
+---
+
+### 6. **Parameter Name Differences**
+
+- **Official FLA**: Uses `use_output_gate` parameter
+- **Our Implementation**: Uses `use_gate` parameter (renamed for clarity)
+
+---
+
+### 7. **Return Value Differences**
+
+#### Official FLA (`fla/layers/gated_deltaproduct.py`)
+```python
+return o, None, past_key_values  # Returns (output, None, past_key_values)
+```
+
+#### Our Custom Implementation (`gated_deltaproduct.py`)
+```python
+return o, recurrent_state, past_key_values  # Returns (output, recurrent_state, past_key_values)
+```
+
+**Key Difference**: Our implementation returns `recurrent_state` (the final hidden state) instead of `None`, enabling state propagation.
+
+---
+
+### 8. **Encoder Wrapper Return Values**
+
+Our `GatedDeltaProductEncoder` (in `src/models/blocks.py`) returns both the output and hidden state:
+
+```python
+x, last_hidden_state, _ = self.encoder_layer(
+    x, output_attentions=True, initial_state=initial_state
+)
+return x, last_hidden_state  # ← Returns hidden state for weaving
+```
+
+This allows state propagation between layers in the `TimeSeriesModel`.
+
+---
+
+## Summary Table
+
+| Feature | Official FLA | Our Custom Implementation |
+|---------|-------------|---------------------------|
+| `initial_state` in `forward()` | ❌ No | ✅ Yes |
+| `initial_state` in `GatedDeltaProductBlock.forward()` | ❌ No | ✅ Yes |
+| Hidden state weaving | ❌ No | ✅ Yes (configurable) |
+| Learnable initial states | ❌ No | ✅ Yes (`nn.ParameterList`) |
+| Returns `recurrent_state` | ❌ No (returns `None`) | ✅ Yes |
+| Layer-to-layer state propagation | ❌ No | ✅ Yes (when `weaving=True`) |
+| Parameter name | `use_output_gate` | `use_gate` |
+
+---
+
+## Why These Differences Matter for Time Series Forecasting
+
+1. **Temporal Continuity**: Hidden state weaving allows information to flow across layers, maintaining temporal patterns across the encoder stack. This is crucial for time series where historical context matters.
+
+2. **Learnable Initialization**: Learnable initial states allow the model to learn optimal starting points for the recurrent computation, which can be crucial for capturing time series patterns.
+
+3. **Flexible State Management**: The `weaving` parameter allows switching between:
+   - **Weaving mode**: Better for capturing long-term dependencies across layers
+   - **Independent mode**: Each layer processes independently, potentially more stable
+
+4. **State Propagation**: Returning and propagating hidden states enables the model to maintain context across multiple encoder layers, which is beneficial for time series forecasting where historical context matters.
+
+These modifications make our implementation better suited for time series forecasting tasks compared to the general-purpose language modeling focus of the official FLA implementation.
+
+---
+
+## Files in This Directory
+
+- **`gated_deltaproduct.py`**: Core GatedDeltaProduct layer implementation with `initial_state` support
+- **`modeling_gated_deltaproduct.py`**: GatedDeltaProductBlock wrapper that integrates the layer
+- **`configuration_gated_deltaproduct.py`**: Configuration class for the model
+- **`__init__.py`**: Module exports
+
+---
+
+## Usage
+
+See `src/models/model.py` and `src/models/blocks.py` for examples of how to use this custom implementation with hidden state weaving.
+
+To enable/disable weaving, set the `weaving` parameter in your encoder configuration:
+```python
+encoder_config = {
+    "weaving": True,  # Enable state propagation across layers
+    # ... other config parameters
+}
+```
+

src/models/gated_deltaproduct/__init__.py

@@ -0,0 +1,11 @@
+from src.models.gated_deltaproduct.configuration_gated_deltaproduct import (
+    GatedDeltaProductConfig,
+)
+from src.models.gated_deltaproduct.modeling_gated_deltaproduct import (
+    GatedDeltaProductBlock,
+)
+
+__all__ = [
+    "GatedDeltaProductConfig",
+    "GatedDeltaProductBlock",
+]

src/models/gated_deltaproduct/configuration_gated_deltaproduct.py

@@ -0,0 +1,108 @@
+import warnings
+
+from transformers.configuration_utils import PretrainedConfig
+
+
+class GatedDeltaProductConfig(PretrainedConfig):
+    model_type = "gated_deltaproduct"
+    keys_to_ignore_at_inference = ["past_key_values"]
+
+    def __init__(
+        self,
+        attn_mode: str = "chunk",
+        conv_size: int = 4,
+        head_dim: int = 256,
+        num_heads: int = 6,
+        hidden_size: int = 2048,
+        expand_v: float = 2.0,
+        use_gate: bool = True,  # Changed from use_output_gate to use_gate for custom implementation
+        use_short_conv: bool = True,
+        max_position_embeddings: int = 2048,
+        hidden_ratio: int | None = 4,
+        intermediate_size: int | None = None,
+        hidden_act: str = "swish",
+        num_hidden_layers: int = 21,
+        norm_eps: float = 1e-6,
+        attn: dict | None = None,
+        use_cache: bool = True,
+        pad_token_id: int = 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,
+        use_forget_gate: bool = False,
+        allow_neg_eigval: bool = False,
+        num_householder: int = 1,
+        **kwargs,
+    ):
+        self.attn_mode = attn_mode
+        self.conv_size = conv_size
+        self.head_dim = head_dim
+        self.num_heads = num_heads
+        self.hidden_size = hidden_size
+        self.expand_v = expand_v
+        self.use_gate = use_gate  # Changed from use_output_gate to use_gate
+        self.use_short_conv = use_short_conv
+        self.max_position_embeddings = max_position_embeddings
+
+        self.hidden_ratio = hidden_ratio
+        self.intermediate_size = intermediate_size
+        self.hidden_act = hidden_act
+        self.num_hidden_layers = num_hidden_layers
+        self.norm_eps = norm_eps
+        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.",
+            )
+
+        # DeltaProduct specific
+        self.allow_neg_eigval = allow_neg_eigval
+        self.num_householder = num_householder
+        self.use_forget_gate = use_forget_gate
+
+        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.0)
+
+        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,
+        )

src/models/gated_deltaproduct/gated_deltaproduct.py

@@ -0,0 +1,351 @@
+# -*- coding: utf-8 -*-
+# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang
+
+from __future__ import annotations
+
+import math
+import warnings
+from typing import TYPE_CHECKING, Dict, Optional, Tuple
+
+import torch
+import torch.nn as nn
+from einops import rearrange, repeat
+from fla.layers.utils import get_unpad_data, index_first_axis, pad_input
+from fla.modules import FusedRMSNormGated, RMSNorm, ShortConvolution
+from fla.ops.delta_rule import fused_recurrent_delta_rule
+from fla.ops.gated_delta_product import chunk_gated_delta_product
+from fla.ops.gated_delta_rule import fused_recurrent_gated_delta_rule
+from torch.nn import functional as F
+
+if TYPE_CHECKING:
+    from fla.models.utils import Cache
+    from transformers.processing_utils import Unpack
+
+
+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_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_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
+        self.init_hidden_state = nn.Parameter(
+            torch.randn(self.num_heads, self.head_dim, self.head_dim)
+        )
+
+        # 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 suppoerted 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_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)
+        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: Optional[torch.Tensor] = None,
+        past_key_values: Optional[Cache] = None,
+        initial_state: Optional[torch.Tensor] = None,
+        use_cache: Optional[bool] = False,
+        output_attentions: Optional[bool] = False,
+        **kwargs: Unpack[Dict],
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Cache]]:
+        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 = self.mode
+
+        if self.training:
+            assert mode == "chunk", "Only chunk mode is supported in training."
+
+        last_state = None
+        if past_key_values is not None and len(past_key_values) > self.layer_idx:
+            last_state = past_key_values[self.layer_idx]
+
+        cu_seqlens = kwargs.get("cu_seqlens", None)
+        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,
+            "... l (n h d) -> ... (l n) h d",
+            n=self.num_householder,
+            d=self.head_k_dim,
+        )
+        v = rearrange(
+            v,
+            "... l (n h d) -> ... (l 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.0
+
+        beta = rearrange(beta, "... l (n h) -> ... (l 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=initial_state,
+                output_final_state=output_attentions,
+                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 = torch.zeros(
+                    g.shape[0],
+                    g.shape[1],
+                    self.num_householder,
+                    g.shape[2],
+                    device=g.device,
+                    dtype=torch.float32,
+                )
+                g_new[:, :, 0] = g
+                g = rearrange(g_new, "... l n h -> ... (l 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, "... l n h d-> ... (l n) h d")
+            if self.use_forget_gate:
+                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,
+                )
+            else:
+                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 * self.num_householder
+                    if cu_seqlens is not None
+                    else None,
+                    use_qk_l2norm_in_kernel=True,
+                )
+            o = rearrange(o, "... (l n) h d -> ... l n h d", n=self.num_householder)[
+                ..., -1, :, :
+            ].contiguous()
+
+        if past_key_values is not None:
+            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,
+            )
+
+        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, recurrent_state, past_key_values

src/models/gated_deltaproduct/modeling_gated_deltaproduct.py

@@ -0,0 +1,105 @@
+# -*- coding: utf-8 -*-
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Dict, List, Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+from fla.layers.attn import Attention
+from fla.models.utils import Cache
+from fla.modules import GatedMLP as GatedDeltaProductMLP
+from fla.modules import RMSNorm
+
+from src.models.gated_deltaproduct.configuration_gated_deltaproduct import (
+    GatedDeltaProductConfig,
+)
+from src.models.gated_deltaproduct.gated_deltaproduct import GatedDeltaProduct
+
+if TYPE_CHECKING:
+    from transformers.processing_utils import Unpack
+
+
+class GatedDeltaProductBlock(nn.Module):
+    def __init__(self, config: GatedDeltaProductConfig, 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 = GatedDeltaProduct(
+                mode=config.attn_mode,
+                hidden_size=config.hidden_size,
+                expand_v=config.expand_v,
+                head_dim=config.head_dim,
+                num_heads=config.num_heads,
+                use_gate=config.use_gate,
+                use_forget_gate=config.use_forget_gate,
+                use_short_conv=config.use_short_conv,
+                conv_size=config.conv_size,
+                norm_eps=config.norm_eps,
+                allow_neg_eigval=config.allow_neg_eigval,
+                num_householder=config.num_householder,
+                layer_idx=layer_idx,
+            )
+        self.mlp_norm = (RMSNorm if config.fuse_norm else nn.RMSNorm)(
+            config.hidden_size, eps=config.norm_eps
+        )
+        self.mlp = GatedDeltaProductMLP(
+            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: Optional[torch.Tensor] = None,
+        past_key_values: Optional[Union[Cache, List[torch.FloatTensor]]] = None,
+        use_cache: Optional[bool] = False,
+        output_attentions: Optional[bool] = False,
+        initial_state: Optional[torch.FloatTensor] = None,
+        **kwargs: Unpack[Dict],
+    ) -> Tuple[
+        torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]
+    ]:
+        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,
+            initial_state=initial_state,
+            **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, **kwargs)
+        hidden_states = residual + hidden_states
+
+        outputs = (hidden_states, attentions, past_key_values)
+
+        return outputs