From f438a2e842ef5205e705f913d1558c35cb394af1 Mon Sep 17 00:00:00 2001
From: Cyril Vallez <cyril.vallez@gmail.com>
Date: Mon, 13 Jul 2026 18:29:08 +0900
Subject: [PATCH 01/43] kernel native

---
 .../models/qwen3_5/modeling_qwen3_5.py        | 58 +++++++++-----
 .../qwen3_5_moe/modeling_qwen3_5_moe.py       | 58 +++++++++-----
 .../models/qwen3_next/modeling_qwen3_next.py  | 79 +++++++++++--------
 .../models/qwen3_next/modular_qwen3_next.py   | 77 ++++++++++--------
 4 files changed, 167 insertions(+), 105 deletions(-)

diff --git a/src/transformers/models/qwen3_5/modeling_qwen3_5.py b/src/transformers/models/qwen3_5/modeling_qwen3_5.py
index 1f695935030c..f56c33ae37fd 100644
--- a/src/transformers/models/qwen3_5/modeling_qwen3_5.py
+++ b/src/transformers/models/qwen3_5/modeling_qwen3_5.py
@@ -32,7 +32,7 @@
 from ...activations import ACT2FN
 from ...cache_utils import Cache, DynamicCache
 from ...generation import GenerationMixin
-from ...integrations import use_kernel_forward_from_hub
+from ...integrations import use_kernel_forward_from_hub, use_kernel_func_from_hub
 from ...integrations.accelerate import force_accelerate_hooks
 from ...masking_utils import create_causal_mask, create_recurrent_attention_mask
 from ...modeling_flash_attention_utils import FlashAttentionKwargs
@@ -58,18 +58,13 @@
     maybe_autocast,
     merge_with_config_defaults,
 )
-from ...utils.import_utils import is_causal_conv1d_available, is_flash_linear_attention_available
+from ...utils.import_utils import is_flash_linear_attention_available
 from ...utils.output_capturing import capture_outputs
 from ...vision_utils import get_vision_bilinear_indices_and_weights, get_vision_cu_seqlens, get_vision_position_ids
 from ..auto.modeling_auto import AutoModel
 from .configuration_qwen3_5 import Qwen3_5Config, Qwen3_5TextConfig, Qwen3_5VisionConfig
 
 
-if is_causal_conv1d_available():
-    from causal_conv1d import causal_conv1d_fn, causal_conv1d_update
-else:
-    causal_conv1d_update, causal_conv1d_fn = None, None
-
 if is_flash_linear_attention_available():
     from fla.modules import FusedRMSNormGated
     from fla.ops.gated_delta_rule import chunk_gated_delta_rule, fused_recurrent_gated_delta_rule
@@ -216,17 +211,16 @@ def apply_mask_to_padding_states(hidden_states, attention_mask):
     return hidden_states
 
 
-is_fast_path_available = all(
-    (causal_conv1d_fn, causal_conv1d_update, chunk_gated_delta_rule, fused_recurrent_gated_delta_rule)
-)
+is_fast_path_available = all((chunk_gated_delta_rule, fused_recurrent_gated_delta_rule))
 
 
-def torch_causal_conv1d_update(
-    hidden_states,
-    conv_state,
-    weight,
-    bias=None,
-    activation=None,
+@use_kernel_func_from_hub("causal_conv1d_update")
+def causal_conv1d_update(
+    hidden_states: torch.Tensor,
+    conv_state: torch.Tensor,
+    weight: nn.Parameter,
+    bias: nn.Parameter | None = None,
+    activation: str | None = None,
 ):
     _, hidden_size, seq_len = hidden_states.shape
     state_len = conv_state.shape[-1]
@@ -234,9 +228,33 @@ def torch_causal_conv1d_update(
     hidden_states_new = torch.cat([conv_state, hidden_states], dim=-1).to(weight.dtype)
     conv_state.copy_(hidden_states_new[:, :, -state_len:])
     out = F.conv1d(hidden_states_new, weight.unsqueeze(1), bias, padding=0, groups=hidden_size)
-    out = F.silu(out[:, :, -seq_len:])
-    out = out.to(hidden_states.dtype)
-    return out
+    out = out[:, :, -seq_len:]
+    if activation is not None:
+        out = ACT2FN[activation](out)
+    return out.to(hidden_states.dtype)
+
+
+@use_kernel_func_from_hub("causal_conv1d_fn")
+def causal_conv1d_fn(
+    hidden_states: torch.Tensor,
+    weight: nn.Parameter,
+    bias: nn.Parameter | None = None,
+    activation: str | None = None,
+    **kwargs,
+):
+    _, hidden_size, seq_len = hidden_states.shape
+    padding = weight.shape[-1] - 1
+
+    out = F.conv1d(
+        hidden_states.to(weight.dtype),
+        weight=weight.unsqueeze(1),
+        bias=bias,
+        padding=padding,
+        groups=hidden_size,
+    )[:, :, :seq_len]
+    if activation is not None:
+        out = ACT2FN[activation](out)
+    return out.to(hidden_states.dtype)
 
 
 def l2norm(x: torch.FloatTensor, dim: int = -1, eps: float = 1e-6):
@@ -418,8 +436,6 @@ def __init__(self, config: Qwen3_5Config, layer_idx: int):
 
         self.out_proj = nn.Linear(self.value_dim, self.hidden_size, bias=False)
 
-        self.causal_conv1d_fn = causal_conv1d_fn
-        self.causal_conv1d_update = causal_conv1d_update or torch_causal_conv1d_update
         self.chunk_gated_delta_rule = chunk_gated_delta_rule or torch_chunk_gated_delta_rule
         self.recurrent_gated_delta_rule = fused_recurrent_gated_delta_rule or torch_recurrent_gated_delta_rule
 
diff --git a/src/transformers/models/qwen3_5_moe/modeling_qwen3_5_moe.py b/src/transformers/models/qwen3_5_moe/modeling_qwen3_5_moe.py
index 643af3afd48b..d601fb7c6ac7 100644
--- a/src/transformers/models/qwen3_5_moe/modeling_qwen3_5_moe.py
+++ b/src/transformers/models/qwen3_5_moe/modeling_qwen3_5_moe.py
@@ -32,7 +32,7 @@
 from ...activations import ACT2FN
 from ...cache_utils import Cache, DynamicCache
 from ...generation import GenerationMixin
-from ...integrations import use_experts_implementation, use_kernel_forward_from_hub
+from ...integrations import use_experts_implementation, use_kernel_forward_from_hub, use_kernel_func_from_hub
 from ...integrations.accelerate import force_accelerate_hooks
 from ...masking_utils import create_causal_mask, create_recurrent_attention_mask
 from ...modeling_flash_attention_utils import FlashAttentionKwargs
@@ -55,18 +55,13 @@
     maybe_autocast,
     merge_with_config_defaults,
 )
-from ...utils.import_utils import is_causal_conv1d_available, is_flash_linear_attention_available
+from ...utils.import_utils import is_flash_linear_attention_available
 from ...utils.output_capturing import OutputRecorder, capture_outputs
 from ...vision_utils import get_vision_bilinear_indices_and_weights, get_vision_cu_seqlens, get_vision_position_ids
 from ..auto.modeling_auto import AutoModel
 from .configuration_qwen3_5_moe import Qwen3_5MoeConfig, Qwen3_5MoeTextConfig, Qwen3_5MoeVisionConfig
 
 
-if is_causal_conv1d_available():
-    from causal_conv1d import causal_conv1d_fn, causal_conv1d_update
-else:
-    causal_conv1d_update, causal_conv1d_fn = None, None
-
 if is_flash_linear_attention_available():
     from fla.modules import FusedRMSNormGated
     from fla.ops.gated_delta_rule import chunk_gated_delta_rule, fused_recurrent_gated_delta_rule
@@ -213,17 +208,16 @@ def apply_mask_to_padding_states(hidden_states, attention_mask):
     return hidden_states
 
 
-is_fast_path_available = all(
-    (causal_conv1d_fn, causal_conv1d_update, chunk_gated_delta_rule, fused_recurrent_gated_delta_rule)
-)
+is_fast_path_available = all((chunk_gated_delta_rule, fused_recurrent_gated_delta_rule))
 
 
-def torch_causal_conv1d_update(
-    hidden_states,
-    conv_state,
-    weight,
-    bias=None,
-    activation=None,
+@use_kernel_func_from_hub("causal_conv1d_update")
+def causal_conv1d_update(
+    hidden_states: torch.Tensor,
+    conv_state: torch.Tensor,
+    weight: nn.Parameter,
+    bias: nn.Parameter | None = None,
+    activation: str | None = None,
 ):
     _, hidden_size, seq_len = hidden_states.shape
     state_len = conv_state.shape[-1]
@@ -231,9 +225,33 @@ def torch_causal_conv1d_update(
     hidden_states_new = torch.cat([conv_state, hidden_states], dim=-1).to(weight.dtype)
     conv_state.copy_(hidden_states_new[:, :, -state_len:])
     out = F.conv1d(hidden_states_new, weight.unsqueeze(1), bias, padding=0, groups=hidden_size)
-    out = F.silu(out[:, :, -seq_len:])
-    out = out.to(hidden_states.dtype)
-    return out
+    out = out[:, :, -seq_len:]
+    if activation is not None:
+        out = ACT2FN[activation](out)
+    return out.to(hidden_states.dtype)
+
+
+@use_kernel_func_from_hub("causal_conv1d_fn")
+def causal_conv1d_fn(
+    hidden_states: torch.Tensor,
+    weight: nn.Parameter,
+    bias: nn.Parameter | None = None,
+    activation: str | None = None,
+    **kwargs,
+):
+    _, hidden_size, seq_len = hidden_states.shape
+    padding = weight.shape[-1] - 1
+
+    out = F.conv1d(
+        hidden_states.to(weight.dtype),
+        weight=weight.unsqueeze(1),
+        bias=bias,
+        padding=padding,
+        groups=hidden_size,
+    )[:, :, :seq_len]
+    if activation is not None:
+        out = ACT2FN[activation](out)
+    return out.to(hidden_states.dtype)
 
 
 def l2norm(x: torch.FloatTensor, dim: int = -1, eps: float = 1e-6):
@@ -415,8 +433,6 @@ def __init__(self, config: Qwen3_5MoeConfig, layer_idx: int):
 
         self.out_proj = nn.Linear(self.value_dim, self.hidden_size, bias=False)
 
-        self.causal_conv1d_fn = causal_conv1d_fn
-        self.causal_conv1d_update = causal_conv1d_update or torch_causal_conv1d_update
         self.chunk_gated_delta_rule = chunk_gated_delta_rule or torch_chunk_gated_delta_rule
         self.recurrent_gated_delta_rule = fused_recurrent_gated_delta_rule or torch_recurrent_gated_delta_rule
 
diff --git a/src/transformers/models/qwen3_next/modeling_qwen3_next.py b/src/transformers/models/qwen3_next/modeling_qwen3_next.py
index 95c506aaf8e1..0881208a4a23 100644
--- a/src/transformers/models/qwen3_next/modeling_qwen3_next.py
+++ b/src/transformers/models/qwen3_next/modeling_qwen3_next.py
@@ -29,7 +29,7 @@
 from ...activations import ACT2FN
 from ...cache_utils import Cache, DynamicCache
 from ...generation import GenerationMixin
-from ...integrations import use_experts_implementation
+from ...integrations import use_experts_implementation, use_kernel_func_from_hub
 from ...integrations.accelerate import force_accelerate_hooks
 from ...masking_utils import create_causal_mask, create_recurrent_attention_mask
 from ...modeling_flash_attention_utils import FlashAttentionKwargs
@@ -45,16 +45,11 @@
 from ...processing_utils import Unpack
 from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, logging
 from ...utils.generic import maybe_autocast, merge_with_config_defaults
-from ...utils.import_utils import is_causal_conv1d_available, is_flash_linear_attention_available
+from ...utils.import_utils import is_flash_linear_attention_available
 from ...utils.output_capturing import OutputRecorder, capture_outputs
 from .configuration_qwen3_next import Qwen3NextConfig
 
 
-if is_causal_conv1d_available():
-    from causal_conv1d import causal_conv1d_fn, causal_conv1d_update
-else:
-    causal_conv1d_update, causal_conv1d_fn = None, None
-
 if is_flash_linear_attention_available():
     from fla.modules import FusedRMSNormGated
     from fla.ops.gated_delta_rule import chunk_gated_delta_rule, fused_recurrent_gated_delta_rule
@@ -342,17 +337,16 @@ def apply_mask_to_padding_states(hidden_states, attention_mask):
     return hidden_states
 
 
-is_fast_path_available = all(
-    (causal_conv1d_fn, causal_conv1d_update, chunk_gated_delta_rule, fused_recurrent_gated_delta_rule)
-)
+is_fast_path_available = all((chunk_gated_delta_rule, fused_recurrent_gated_delta_rule))
 
 
-def torch_causal_conv1d_update(
-    hidden_states,
-    conv_state,
-    weight,
-    bias=None,
-    activation=None,
+@use_kernel_func_from_hub("causal_conv1d_update")
+def causal_conv1d_update(
+    hidden_states: torch.Tensor,
+    conv_state: torch.Tensor,
+    weight: nn.Parameter,
+    bias: nn.Parameter | None = None,
+    activation: str | None = None,
 ):
     _, hidden_size, seq_len = hidden_states.shape
     state_len = conv_state.shape[-1]
@@ -360,9 +354,33 @@ def torch_causal_conv1d_update(
     hidden_states_new = torch.cat([conv_state, hidden_states], dim=-1).to(weight.dtype)
     conv_state.copy_(hidden_states_new[:, :, -state_len:])
     out = F.conv1d(hidden_states_new, weight.unsqueeze(1), bias, padding=0, groups=hidden_size)
-    out = F.silu(out[:, :, -seq_len:])
-    out = out.to(hidden_states.dtype)
-    return out
+    out = out[:, :, -seq_len:]
+    if activation is not None:
+        out = ACT2FN[activation](out)
+    return out.to(hidden_states.dtype)
+
+
+@use_kernel_func_from_hub("causal_conv1d_fn")
+def causal_conv1d_fn(
+    hidden_states: torch.Tensor,
+    weight: nn.Parameter,
+    bias: nn.Parameter | None = None,
+    activation: str | None = None,
+    **kwargs,
+):
+    _, hidden_size, seq_len = hidden_states.shape
+    padding = weight.shape[-1] - 1
+
+    out = F.conv1d(
+        hidden_states.to(weight.dtype),
+        weight=weight.unsqueeze(1),
+        bias=bias,
+        padding=padding,
+        groups=hidden_size,
+    )[:, :, :seq_len]
+    if activation is not None:
+        out = ACT2FN[activation](out)
+    return out.to(hidden_states.dtype)
 
 
 def l2norm(x: torch.FloatTensor, dim: int = -1, eps: float = 1e-6):
@@ -549,8 +567,6 @@ def __init__(self, config: Qwen3NextConfig, layer_idx: int):
 
         self.out_proj = nn.Linear(self.value_dim, self.hidden_size, bias=False)
 
-        self.causal_conv1d_fn = causal_conv1d_fn
-        self.causal_conv1d_update = causal_conv1d_update or torch_causal_conv1d_update
         self.chunk_gated_delta_rule = chunk_gated_delta_rule or torch_chunk_gated_delta_rule
         self.recurrent_gated_delta_rule = fused_recurrent_gated_delta_rule or torch_recurrent_gated_delta_rule
 
@@ -626,7 +642,7 @@ def forward(
 
         if use_precomputed_states and seq_len == 1:
             # Single-token cached decode: the fused per-step kernel updates the conv state in-place.
-            mixed_qkv = self.causal_conv1d_update(
+            mixed_qkv = causal_conv1d_update(
                 mixed_qkv,
                 conv_state,
                 self.conv1d.weight.squeeze(1),
@@ -643,16 +659,15 @@ def forward(
             if cache_params is not None:
                 new_conv_state = F.pad(mixed_qkv, (self.conv_kernel_size - mixed_qkv.shape[-1], 0))
                 cache_params.update_conv_state(new_conv_state, self.layer_idx)
-            if self.causal_conv1d_fn is not None:
-                mixed_qkv = self.causal_conv1d_fn(
-                    x=mixed_qkv,
-                    weight=self.conv1d.weight.squeeze(1),
-                    bias=self.conv1d.bias,
-                    activation=self.activation,
-                    seq_idx=kwargs.get("seq_idx"),
-                )
-            else:
-                mixed_qkv = F.silu(self.conv1d(mixed_qkv)[:, :, : mixed_qkv.shape[-1]])
+
+            mixed_qkv = causal_conv1d_fn(
+                mixed_qkv,
+                self.conv1d.weight.squeeze(1),
+                self.conv1d.bias,
+                activation=self.activation,
+                seq_idx=kwargs.get("seq_idx"),
+            )
+
             if use_precomputed_states:
                 mixed_qkv = mixed_qkv[:, :, -seq_len:]
 
diff --git a/src/transformers/models/qwen3_next/modular_qwen3_next.py b/src/transformers/models/qwen3_next/modular_qwen3_next.py
index 86f31ad1431d..6266587a3ff9 100644
--- a/src/transformers/models/qwen3_next/modular_qwen3_next.py
+++ b/src/transformers/models/qwen3_next/modular_qwen3_next.py
@@ -23,6 +23,7 @@
 from ... import initialization as init
 from ...activations import ACT2FN
 from ...cache_utils import Cache, DynamicCache
+from ...integrations import use_kernel_func_from_hub
 from ...integrations.accelerate import force_accelerate_hooks
 from ...masking_utils import create_causal_mask, create_recurrent_attention_mask
 from ...modeling_flash_attention_utils import FlashAttentionKwargs
@@ -32,7 +33,6 @@
 from ...utils import TransformersKwargs, auto_docstring, logging
 from ...utils.generic import merge_with_config_defaults, no_inherit_decorator
 from ...utils.import_utils import (
-    is_causal_conv1d_available,
     is_flash_linear_attention_available,
 )
 from ...utils.output_capturing import OutputRecorder, capture_outputs
@@ -55,11 +55,6 @@
 from .configuration_qwen3_next import Qwen3NextConfig
 
 
-if is_causal_conv1d_available():
-    from causal_conv1d import causal_conv1d_fn, causal_conv1d_update
-else:
-    causal_conv1d_update, causal_conv1d_fn = None, None
-
 if is_flash_linear_attention_available():
     from fla.modules import FusedRMSNormGated
     from fla.ops.gated_delta_rule import chunk_gated_delta_rule, fused_recurrent_gated_delta_rule
@@ -68,9 +63,7 @@
     FusedRMSNormGated = None
 
 
-is_fast_path_available = all(
-    (causal_conv1d_fn, causal_conv1d_update, chunk_gated_delta_rule, fused_recurrent_gated_delta_rule)
-)
+is_fast_path_available = all((chunk_gated_delta_rule, fused_recurrent_gated_delta_rule))
 
 logger = logging.get_logger(__name__)
 
@@ -188,12 +181,13 @@ def forward(
         return attn_output, attn_weights
 
 
-def torch_causal_conv1d_update(
-    hidden_states,
-    conv_state,
-    weight,
-    bias=None,
-    activation=None,
+@use_kernel_func_from_hub("causal_conv1d_update")
+def causal_conv1d_update(
+    hidden_states: torch.Tensor,
+    conv_state: torch.Tensor,
+    weight: nn.Parameter,
+    bias: nn.Parameter | None = None,
+    activation: str | None = None,
 ):
     _, hidden_size, seq_len = hidden_states.shape
     state_len = conv_state.shape[-1]
@@ -201,9 +195,33 @@ def torch_causal_conv1d_update(
     hidden_states_new = torch.cat([conv_state, hidden_states], dim=-1).to(weight.dtype)
     conv_state.copy_(hidden_states_new[:, :, -state_len:])
     out = F.conv1d(hidden_states_new, weight.unsqueeze(1), bias, padding=0, groups=hidden_size)
-    out = F.silu(out[:, :, -seq_len:])
-    out = out.to(hidden_states.dtype)
-    return out
+    out = out[:, :, -seq_len:]
+    if activation is not None:
+        out = ACT2FN[activation](out)
+    return out.to(hidden_states.dtype)
+
+
+@use_kernel_func_from_hub("causal_conv1d_fn")
+def causal_conv1d_fn(
+    hidden_states: torch.Tensor,
+    weight: nn.Parameter,
+    bias: nn.Parameter | None = None,
+    activation: str | None = None,
+    **kwargs,
+):
+    _, hidden_size, seq_len = hidden_states.shape
+    padding = weight.shape[-1] - 1
+
+    out = F.conv1d(
+        hidden_states.to(weight.dtype),
+        weight=weight.unsqueeze(1),
+        bias=bias,
+        padding=padding,
+        groups=hidden_size,
+    )[:, :, :seq_len]
+    if activation is not None:
+        out = ACT2FN[activation](out)
+    return out.to(hidden_states.dtype)
 
 
 def l2norm(x: torch.FloatTensor, dim: int = -1, eps: float = 1e-6):
@@ -390,8 +408,6 @@ def __init__(self, config: Qwen3NextConfig, layer_idx: int):
 
         self.out_proj = nn.Linear(self.value_dim, self.hidden_size, bias=False)
 
-        self.causal_conv1d_fn = causal_conv1d_fn
-        self.causal_conv1d_update = causal_conv1d_update or torch_causal_conv1d_update
         self.chunk_gated_delta_rule = chunk_gated_delta_rule or torch_chunk_gated_delta_rule
         self.recurrent_gated_delta_rule = fused_recurrent_gated_delta_rule or torch_recurrent_gated_delta_rule
 
@@ -467,7 +483,7 @@ def forward(
 
         if use_precomputed_states and seq_len == 1:
             # Single-token cached decode: the fused per-step kernel updates the conv state in-place.
-            mixed_qkv = self.causal_conv1d_update(
+            mixed_qkv = causal_conv1d_update(
                 mixed_qkv,
                 conv_state,
                 self.conv1d.weight.squeeze(1),
@@ -484,16 +500,15 @@ def forward(
             if cache_params is not None:
                 new_conv_state = F.pad(mixed_qkv, (self.conv_kernel_size - mixed_qkv.shape[-1], 0))
                 cache_params.update_conv_state(new_conv_state, self.layer_idx)
-            if self.causal_conv1d_fn is not None:
-                mixed_qkv = self.causal_conv1d_fn(
-                    x=mixed_qkv,
-                    weight=self.conv1d.weight.squeeze(1),
-                    bias=self.conv1d.bias,
-                    activation=self.activation,
-                    seq_idx=kwargs.get("seq_idx"),
-                )
-            else:
-                mixed_qkv = F.silu(self.conv1d(mixed_qkv)[:, :, : mixed_qkv.shape[-1]])
+
+            mixed_qkv = causal_conv1d_fn(
+                mixed_qkv,
+                self.conv1d.weight.squeeze(1),
+                self.conv1d.bias,
+                activation=self.activation,
+                seq_idx=kwargs.get("seq_idx"),
+            )
+
             if use_precomputed_states:
                 mixed_qkv = mixed_qkv[:, :, -seq_len:]
 

From 7c1335b8e8da5085fa460dd7d3be4f841717038a Mon Sep 17 00:00:00 2001
From: Cyril Vallez <cyril.vallez@gmail.com>
Date: Mon, 13 Jul 2026 18:31:08 +0900
Subject: [PATCH 02/43] fix warning

---
 src/transformers/models/qwen3_5/modeling_qwen3_5.py         | 5 ++---
 src/transformers/models/qwen3_5_moe/modeling_qwen3_5_moe.py | 5 ++---
 src/transformers/models/qwen3_next/modeling_qwen3_next.py   | 5 ++---
 src/transformers/models/qwen3_next/modular_qwen3_next.py    | 5 ++---
 4 files changed, 8 insertions(+), 12 deletions(-)

diff --git a/src/transformers/models/qwen3_5/modeling_qwen3_5.py b/src/transformers/models/qwen3_5/modeling_qwen3_5.py
index f56c33ae37fd..dba996a00d99 100644
--- a/src/transformers/models/qwen3_5/modeling_qwen3_5.py
+++ b/src/transformers/models/qwen3_5/modeling_qwen3_5.py
@@ -441,9 +441,8 @@ def __init__(self, config: Qwen3_5Config, layer_idx: int):
 
         if not is_fast_path_available:
             logger.warning_once(
-                "The fast path is not available because one of the required library is not installed. Falling back to "
-                "torch implementation. To install follow https://github.com/fla-org/flash-linear-attention#installation and"
-                " https://github.com/Dao-AILab/causal-conv1d"
+                "The fast path is not available because the required library is not installed. Falling back to "
+                "torch implementation. To install follow https://github.com/fla-org/flash-linear-attention#installation"
             )
 
         self.layer_type = config.layer_types[layer_idx]
diff --git a/src/transformers/models/qwen3_5_moe/modeling_qwen3_5_moe.py b/src/transformers/models/qwen3_5_moe/modeling_qwen3_5_moe.py
index d601fb7c6ac7..b3b2b631ea91 100644
--- a/src/transformers/models/qwen3_5_moe/modeling_qwen3_5_moe.py
+++ b/src/transformers/models/qwen3_5_moe/modeling_qwen3_5_moe.py
@@ -438,9 +438,8 @@ def __init__(self, config: Qwen3_5MoeConfig, layer_idx: int):
 
         if not is_fast_path_available:
             logger.warning_once(
-                "The fast path is not available because one of the required library is not installed. Falling back to "
-                "torch implementation. To install follow https://github.com/fla-org/flash-linear-attention#installation and"
-                " https://github.com/Dao-AILab/causal-conv1d"
+                "The fast path is not available because the required library is not installed. Falling back to "
+                "torch implementation. To install follow https://github.com/fla-org/flash-linear-attention#installation"
             )
 
         self.layer_type = config.layer_types[layer_idx]
diff --git a/src/transformers/models/qwen3_next/modeling_qwen3_next.py b/src/transformers/models/qwen3_next/modeling_qwen3_next.py
index 0881208a4a23..e94b08a2c120 100644
--- a/src/transformers/models/qwen3_next/modeling_qwen3_next.py
+++ b/src/transformers/models/qwen3_next/modeling_qwen3_next.py
@@ -572,9 +572,8 @@ def __init__(self, config: Qwen3NextConfig, layer_idx: int):
 
         if not is_fast_path_available:
             logger.warning_once(
-                "The fast path is not available because one of the required library is not installed. Falling back to "
-                "torch implementation. To install follow https://github.com/fla-org/flash-linear-attention#installation and"
-                " https://github.com/Dao-AILab/causal-conv1d"
+                "The fast path is not available because the required library is not installed. Falling back to "
+                "torch implementation. To install follow https://github.com/fla-org/flash-linear-attention#installation"
             )
 
         self.layer_type = config.layer_types[layer_idx]
diff --git a/src/transformers/models/qwen3_next/modular_qwen3_next.py b/src/transformers/models/qwen3_next/modular_qwen3_next.py
index 6266587a3ff9..32a7c69964f6 100644
--- a/src/transformers/models/qwen3_next/modular_qwen3_next.py
+++ b/src/transformers/models/qwen3_next/modular_qwen3_next.py
@@ -413,9 +413,8 @@ def __init__(self, config: Qwen3NextConfig, layer_idx: int):
 
         if not is_fast_path_available:
             logger.warning_once(
-                "The fast path is not available because one of the required library is not installed. Falling back to "
-                "torch implementation. To install follow https://github.com/fla-org/flash-linear-attention#installation and"
-                " https://github.com/Dao-AILab/causal-conv1d"
+                "The fast path is not available because the required library is not installed. Falling back to "
+                "torch implementation. To install follow https://github.com/fla-org/flash-linear-attention#installation"
             )
 
         self.layer_type = config.layer_types[layer_idx]

From f8ff0db041cf17c35fa6eb0c1ed659ff6de3d301 Mon Sep 17 00:00:00 2001
From: Cyril Vallez <cyril.vallez@gmail.com>
Date: Mon, 13 Jul 2026 18:35:27 +0900
Subject: [PATCH 03/43] have to fix other modular to be coherent

---
 .../models/qwen3_5/modeling_qwen3_5.py        | 21 ++++++++---------
 .../models/qwen3_5/modular_qwen3_5.py         | 23 ++++++++++---------
 .../qwen3_5_moe/modeling_qwen3_5_moe.py       | 21 ++++++++---------
 3 files changed, 32 insertions(+), 33 deletions(-)

diff --git a/src/transformers/models/qwen3_5/modeling_qwen3_5.py b/src/transformers/models/qwen3_5/modeling_qwen3_5.py
index dba996a00d99..653d6dd2747c 100644
--- a/src/transformers/models/qwen3_5/modeling_qwen3_5.py
+++ b/src/transformers/models/qwen3_5/modeling_qwen3_5.py
@@ -487,7 +487,7 @@ def forward(
 
         if use_precomputed_states and seq_len == 1:
             # Single-token cached decode: the fused per-step kernel updates the conv state in-place.
-            mixed_qkv = self.causal_conv1d_update(
+            mixed_qkv = causal_conv1d_update(
                 mixed_qkv,
                 conv_state,
                 self.conv1d.weight.squeeze(1),
@@ -504,16 +504,15 @@ def forward(
             if cache_params is not None:
                 new_conv_state = F.pad(mixed_qkv, (self.conv_kernel_size - mixed_qkv.shape[-1], 0))
                 cache_params.update_conv_state(new_conv_state, self.layer_idx)
-            if self.causal_conv1d_fn is not None:
-                mixed_qkv = self.causal_conv1d_fn(
-                    x=mixed_qkv,
-                    weight=self.conv1d.weight.squeeze(1),
-                    bias=self.conv1d.bias,
-                    activation=self.activation,
-                    seq_idx=kwargs.get("seq_idx"),
-                )
-            else:
-                mixed_qkv = F.silu(self.conv1d(mixed_qkv)[:, :, : mixed_qkv.shape[-1]])
+
+            mixed_qkv = causal_conv1d_fn(
+                mixed_qkv,
+                self.conv1d.weight.squeeze(1),
+                self.conv1d.bias,
+                activation=self.activation,
+                seq_idx=kwargs.get("seq_idx"),
+            )
+
             if use_precomputed_states:
                 mixed_qkv = mixed_qkv[:, :, -seq_len:]
 
diff --git a/src/transformers/models/qwen3_5/modular_qwen3_5.py b/src/transformers/models/qwen3_5/modular_qwen3_5.py
index 2b88938d533d..d81edf5015a3 100644
--- a/src/transformers/models/qwen3_5/modular_qwen3_5.py
+++ b/src/transformers/models/qwen3_5/modular_qwen3_5.py
@@ -46,6 +46,8 @@
     Qwen3NextPreTrainedModel,
     Qwen3NextRMSNorm,
     apply_mask_to_padding_states,
+    causal_conv1d_fn,
+    causal_conv1d_update,
 )
 from ..qwen3_vl.configuration_qwen3_vl import Qwen3VLConfig, Qwen3VLVisionConfig
 from ..qwen3_vl.modeling_qwen3_vl import (
@@ -251,7 +253,7 @@ def forward(
 
         if use_precomputed_states and seq_len == 1:
             # Single-token cached decode: the fused per-step kernel updates the conv state in-place.
-            mixed_qkv = self.causal_conv1d_update(
+            mixed_qkv = causal_conv1d_update(
                 mixed_qkv,
                 conv_state,
                 self.conv1d.weight.squeeze(1),
@@ -268,16 +270,15 @@ def forward(
             if cache_params is not None:
                 new_conv_state = F.pad(mixed_qkv, (self.conv_kernel_size - mixed_qkv.shape[-1], 0))
                 cache_params.update_conv_state(new_conv_state, self.layer_idx)
-            if self.causal_conv1d_fn is not None:
-                mixed_qkv = self.causal_conv1d_fn(
-                    x=mixed_qkv,
-                    weight=self.conv1d.weight.squeeze(1),
-                    bias=self.conv1d.bias,
-                    activation=self.activation,
-                    seq_idx=kwargs.get("seq_idx"),
-                )
-            else:
-                mixed_qkv = F.silu(self.conv1d(mixed_qkv)[:, :, : mixed_qkv.shape[-1]])
+
+            mixed_qkv = causal_conv1d_fn(
+                mixed_qkv,
+                self.conv1d.weight.squeeze(1),
+                self.conv1d.bias,
+                activation=self.activation,
+                seq_idx=kwargs.get("seq_idx"),
+            )
+
             if use_precomputed_states:
                 mixed_qkv = mixed_qkv[:, :, -seq_len:]
 
diff --git a/src/transformers/models/qwen3_5_moe/modeling_qwen3_5_moe.py b/src/transformers/models/qwen3_5_moe/modeling_qwen3_5_moe.py
index b3b2b631ea91..e12ee94ea268 100644
--- a/src/transformers/models/qwen3_5_moe/modeling_qwen3_5_moe.py
+++ b/src/transformers/models/qwen3_5_moe/modeling_qwen3_5_moe.py
@@ -484,7 +484,7 @@ def forward(
 
         if use_precomputed_states and seq_len == 1:
             # Single-token cached decode: the fused per-step kernel updates the conv state in-place.
-            mixed_qkv = self.causal_conv1d_update(
+            mixed_qkv = causal_conv1d_update(
                 mixed_qkv,
                 conv_state,
                 self.conv1d.weight.squeeze(1),
@@ -501,16 +501,15 @@ def forward(
             if cache_params is not None:
                 new_conv_state = F.pad(mixed_qkv, (self.conv_kernel_size - mixed_qkv.shape[-1], 0))
                 cache_params.update_conv_state(new_conv_state, self.layer_idx)
-            if self.causal_conv1d_fn is not None:
-                mixed_qkv = self.causal_conv1d_fn(
-                    x=mixed_qkv,
-                    weight=self.conv1d.weight.squeeze(1),
-                    bias=self.conv1d.bias,
-                    activation=self.activation,
-                    seq_idx=kwargs.get("seq_idx"),
-                )
-            else:
-                mixed_qkv = F.silu(self.conv1d(mixed_qkv)[:, :, : mixed_qkv.shape[-1]])
+
+            mixed_qkv = causal_conv1d_fn(
+                mixed_qkv,
+                self.conv1d.weight.squeeze(1),
+                self.conv1d.bias,
+                activation=self.activation,
+                seq_idx=kwargs.get("seq_idx"),
+            )
+
             if use_precomputed_states:
                 mixed_qkv = mixed_qkv[:, :, -seq_len:]
 

From 9a8127e887e92062a74e51425aaf07f5f7d2d40e Mon Sep 17 00:00:00 2001
From: Cyril Vallez <cyril.vallez@gmail.com>
Date: Mon, 13 Jul 2026 18:41:13 +0900
Subject: [PATCH 04/43] remove useless

---
 src/transformers/models/qwen3_5/modeling_qwen3_5.py         | 1 -
 src/transformers/models/qwen3_5_moe/modeling_qwen3_5_moe.py | 1 -
 src/transformers/models/qwen3_next/modeling_qwen3_next.py   | 1 -
 src/transformers/models/qwen3_next/modular_qwen3_next.py    | 1 -
 4 files changed, 4 deletions(-)

diff --git a/src/transformers/models/qwen3_5/modeling_qwen3_5.py b/src/transformers/models/qwen3_5/modeling_qwen3_5.py
index 653d6dd2747c..a648a30bdd23 100644
--- a/src/transformers/models/qwen3_5/modeling_qwen3_5.py
+++ b/src/transformers/models/qwen3_5/modeling_qwen3_5.py
@@ -403,7 +403,6 @@ def __init__(self, config: Qwen3_5Config, layer_idx: int):
         self.conv_kernel_size = config.linear_conv_kernel_dim
         self.layer_idx = layer_idx
         self.activation = config.hidden_act
-        self.act = ACT2FN[config.hidden_act]
         self.layer_norm_epsilon = config.rms_norm_eps
 
         # QKV
diff --git a/src/transformers/models/qwen3_5_moe/modeling_qwen3_5_moe.py b/src/transformers/models/qwen3_5_moe/modeling_qwen3_5_moe.py
index e12ee94ea268..029e7c6ca572 100644
--- a/src/transformers/models/qwen3_5_moe/modeling_qwen3_5_moe.py
+++ b/src/transformers/models/qwen3_5_moe/modeling_qwen3_5_moe.py
@@ -400,7 +400,6 @@ def __init__(self, config: Qwen3_5MoeConfig, layer_idx: int):
         self.conv_kernel_size = config.linear_conv_kernel_dim
         self.layer_idx = layer_idx
         self.activation = config.hidden_act
-        self.act = ACT2FN[config.hidden_act]
         self.layer_norm_epsilon = config.rms_norm_eps
 
         # QKV
diff --git a/src/transformers/models/qwen3_next/modeling_qwen3_next.py b/src/transformers/models/qwen3_next/modeling_qwen3_next.py
index e94b08a2c120..5b3d4d7201a2 100644
--- a/src/transformers/models/qwen3_next/modeling_qwen3_next.py
+++ b/src/transformers/models/qwen3_next/modeling_qwen3_next.py
@@ -528,7 +528,6 @@ def __init__(self, config: Qwen3NextConfig, layer_idx: int):
         self.conv_kernel_size = config.linear_conv_kernel_dim
         self.layer_idx = layer_idx
         self.activation = config.hidden_act
-        self.act = ACT2FN[config.hidden_act]
         self.layer_norm_epsilon = config.rms_norm_eps
 
         # QKV
diff --git a/src/transformers/models/qwen3_next/modular_qwen3_next.py b/src/transformers/models/qwen3_next/modular_qwen3_next.py
index 32a7c69964f6..02a370b2d34c 100644
--- a/src/transformers/models/qwen3_next/modular_qwen3_next.py
+++ b/src/transformers/models/qwen3_next/modular_qwen3_next.py
@@ -369,7 +369,6 @@ def __init__(self, config: Qwen3NextConfig, layer_idx: int):
         self.conv_kernel_size = config.linear_conv_kernel_dim
         self.layer_idx = layer_idx
         self.activation = config.hidden_act
-        self.act = ACT2FN[config.hidden_act]
         self.layer_norm_epsilon = config.rms_norm_eps
 
         # QKV

From 72b45caf021e8f37d7a1facd1885471d60113ae4 Mon Sep 17 00:00:00 2001
From: Cyril Vallez <cyril.vallez@gmail.com>
Date: Mon, 13 Jul 2026 23:17:20 +0900
Subject: [PATCH 05/43] use native lib as well

---
 src/transformers/integrations/hub_kernels.py  |  3 +--
 .../models/qwen3_next/modular_qwen3_next.py   | 11 +++++-----
 src/transformers/utils/generic.py             | 20 +++++++++++++++++++
 3 files changed, 27 insertions(+), 7 deletions(-)

diff --git a/src/transformers/integrations/hub_kernels.py b/src/transformers/integrations/hub_kernels.py
index 25885dbb2b05..87c30ec543e7 100644
--- a/src/transformers/integrations/hub_kernels.py
+++ b/src/transformers/integrations/hub_kernels.py
@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import functools
+import importlib
 import os
 import re
 import sys
@@ -531,8 +532,6 @@ def lazy_load_kernel(kernel_name: str, mapping: dict[str, ModuleType | None] = _
 
     else:
         # Try to import is_{kernel_name}_available from ..utils
-        import importlib
-
         new_kernel_name = kernel_name.replace("-", "_")
         func_name = f"is_{new_kernel_name}_available"
 
diff --git a/src/transformers/models/qwen3_next/modular_qwen3_next.py b/src/transformers/models/qwen3_next/modular_qwen3_next.py
index 02a370b2d34c..463761254842 100644
--- a/src/transformers/models/qwen3_next/modular_qwen3_next.py
+++ b/src/transformers/models/qwen3_next/modular_qwen3_next.py
@@ -23,7 +23,7 @@
 from ... import initialization as init
 from ...activations import ACT2FN
 from ...cache_utils import Cache, DynamicCache
-from ...integrations import use_kernel_func_from_hub
+from ...integrations import use_kernel_func_from_hub, use_kernelized_func
 from ...integrations.accelerate import force_accelerate_hooks
 from ...masking_utils import create_causal_mask, create_recurrent_attention_mask
 from ...modeling_flash_attention_utils import FlashAttentionKwargs
@@ -31,10 +31,8 @@
 from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
 from ...processing_utils import Unpack
 from ...utils import TransformersKwargs, auto_docstring, logging
-from ...utils.generic import merge_with_config_defaults, no_inherit_decorator
-from ...utils.import_utils import (
-    is_flash_linear_attention_available,
-)
+from ...utils.generic import merge_with_config_defaults, no_inherit_decorator, replace_with_function_from_package
+from ...utils.import_utils import is_flash_linear_attention_available
 from ...utils.output_capturing import OutputRecorder, capture_outputs
 from ..bamba.modeling_bamba import apply_mask_to_padding_states, apply_rotary_pos_emb
 from ..gemma2.modeling_gemma2 import Gemma2RotaryEmbedding
@@ -182,6 +180,7 @@ def forward(
 
 
 @use_kernel_func_from_hub("causal_conv1d_update")
+@replace_with_function_from_package("causal_conv1d_update", "causal_conv1d")
 def causal_conv1d_update(
     hidden_states: torch.Tensor,
     conv_state: torch.Tensor,
@@ -202,6 +201,7 @@ def causal_conv1d_update(
 
 
 @use_kernel_func_from_hub("causal_conv1d_fn")
+@replace_with_function_from_package("causal_conv1d_update", "causal_conv1d")
 def causal_conv1d_fn(
     hidden_states: torch.Tensor,
     weight: nn.Parameter,
@@ -355,6 +355,7 @@ def torch_recurrent_gated_delta_rule(
     return core_attn_out, last_recurrent_state
 
 
+@use_kernelized_func([causal_conv1d_update, causal_conv1d_fn])
 class Qwen3NextGatedDeltaNet(nn.Module):
     def __init__(self, config: Qwen3NextConfig, layer_idx: int):
         super().__init__()
diff --git a/src/transformers/utils/generic.py b/src/transformers/utils/generic.py
index 9a5a97a023fd..7cb8e0a46949 100644
--- a/src/transformers/utils/generic.py
+++ b/src/transformers/utils/generic.py
@@ -17,6 +17,7 @@
 
 from __future__ import annotations
 
+import importlib
 import inspect
 import json
 import os
@@ -1156,3 +1157,22 @@ def wrapper(*args, **kwargs):
         return wrapper
 
     return decorator
+
+
+def replace_with_function_from_package(function_name: str, package: str):
+    """
+    Decorator that tries to replace the decorated function with `function_name` imported from `package`, if it's available. If not,
+    simply returns the decorated function.
+    Useful to define explicit torch fallback functions, while still using an optimized kernel imported from somewhere else if available.
+    """
+
+    def decorator(func: Callable) -> Callable:
+        try:
+            module = importlib.import_module(package)
+            function = getattr(module, function_name)
+        except Exception:
+            function = func
+
+        return function
+
+    return decorator

From 54d80dfb9d3de813a0432f3decc705b854003f34 Mon Sep 17 00:00:00 2001
From: Cyril Vallez <cyril.vallez@gmail.com>
Date: Mon, 13 Jul 2026 23:18:59 +0900
Subject: [PATCH 06/43] modular

---
 src/transformers/models/qwen3_5/modeling_qwen3_5.py        | 3 +++
 .../models/qwen3_5_moe/modeling_qwen3_5_moe.py             | 3 +++
 src/transformers/models/qwen3_next/modeling_qwen3_next.py  | 7 +++++--
 3 files changed, 11 insertions(+), 2 deletions(-)

diff --git a/src/transformers/models/qwen3_5/modeling_qwen3_5.py b/src/transformers/models/qwen3_5/modeling_qwen3_5.py
index a648a30bdd23..7324b9c23818 100644
--- a/src/transformers/models/qwen3_5/modeling_qwen3_5.py
+++ b/src/transformers/models/qwen3_5/modeling_qwen3_5.py
@@ -57,6 +57,7 @@
     is_flash_attention_requested,
     maybe_autocast,
     merge_with_config_defaults,
+    replace_with_function_from_package,
 )
 from ...utils.import_utils import is_flash_linear_attention_available
 from ...utils.output_capturing import capture_outputs
@@ -215,6 +216,7 @@ def apply_mask_to_padding_states(hidden_states, attention_mask):
 
 
 @use_kernel_func_from_hub("causal_conv1d_update")
+@replace_with_function_from_package("causal_conv1d_update", "causal_conv1d")
 def causal_conv1d_update(
     hidden_states: torch.Tensor,
     conv_state: torch.Tensor,
@@ -235,6 +237,7 @@ def causal_conv1d_update(
 
 
 @use_kernel_func_from_hub("causal_conv1d_fn")
+@replace_with_function_from_package("causal_conv1d_update", "causal_conv1d")
 def causal_conv1d_fn(
     hidden_states: torch.Tensor,
     weight: nn.Parameter,
diff --git a/src/transformers/models/qwen3_5_moe/modeling_qwen3_5_moe.py b/src/transformers/models/qwen3_5_moe/modeling_qwen3_5_moe.py
index 029e7c6ca572..4742680964f2 100644
--- a/src/transformers/models/qwen3_5_moe/modeling_qwen3_5_moe.py
+++ b/src/transformers/models/qwen3_5_moe/modeling_qwen3_5_moe.py
@@ -54,6 +54,7 @@
     is_flash_attention_requested,
     maybe_autocast,
     merge_with_config_defaults,
+    replace_with_function_from_package,
 )
 from ...utils.import_utils import is_flash_linear_attention_available
 from ...utils.output_capturing import OutputRecorder, capture_outputs
@@ -212,6 +213,7 @@ def apply_mask_to_padding_states(hidden_states, attention_mask):
 
 
 @use_kernel_func_from_hub("causal_conv1d_update")
+@replace_with_function_from_package("causal_conv1d_update", "causal_conv1d")
 def causal_conv1d_update(
     hidden_states: torch.Tensor,
     conv_state: torch.Tensor,
@@ -232,6 +234,7 @@ def causal_conv1d_update(
 
 
 @use_kernel_func_from_hub("causal_conv1d_fn")
+@replace_with_function_from_package("causal_conv1d_update", "causal_conv1d")
 def causal_conv1d_fn(
     hidden_states: torch.Tensor,
     weight: nn.Parameter,
diff --git a/src/transformers/models/qwen3_next/modeling_qwen3_next.py b/src/transformers/models/qwen3_next/modeling_qwen3_next.py
index 5b3d4d7201a2..007318584ab9 100644
--- a/src/transformers/models/qwen3_next/modeling_qwen3_next.py
+++ b/src/transformers/models/qwen3_next/modeling_qwen3_next.py
@@ -29,7 +29,7 @@
 from ...activations import ACT2FN
 from ...cache_utils import Cache, DynamicCache
 from ...generation import GenerationMixin
-from ...integrations import use_experts_implementation, use_kernel_func_from_hub
+from ...integrations import use_experts_implementation, use_kernel_func_from_hub, use_kernelized_func
 from ...integrations.accelerate import force_accelerate_hooks
 from ...masking_utils import create_causal_mask, create_recurrent_attention_mask
 from ...modeling_flash_attention_utils import FlashAttentionKwargs
@@ -44,7 +44,7 @@
 from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
 from ...processing_utils import Unpack
 from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, logging
-from ...utils.generic import maybe_autocast, merge_with_config_defaults
+from ...utils.generic import maybe_autocast, merge_with_config_defaults, replace_with_function_from_package
 from ...utils.import_utils import is_flash_linear_attention_available
 from ...utils.output_capturing import OutputRecorder, capture_outputs
 from .configuration_qwen3_next import Qwen3NextConfig
@@ -341,6 +341,7 @@ def apply_mask_to_padding_states(hidden_states, attention_mask):
 
 
 @use_kernel_func_from_hub("causal_conv1d_update")
+@replace_with_function_from_package("causal_conv1d_update", "causal_conv1d")
 def causal_conv1d_update(
     hidden_states: torch.Tensor,
     conv_state: torch.Tensor,
@@ -361,6 +362,7 @@ def causal_conv1d_update(
 
 
 @use_kernel_func_from_hub("causal_conv1d_fn")
+@replace_with_function_from_package("causal_conv1d_update", "causal_conv1d")
 def causal_conv1d_fn(
     hidden_states: torch.Tensor,
     weight: nn.Parameter,
@@ -514,6 +516,7 @@ def torch_recurrent_gated_delta_rule(
     return core_attn_out, last_recurrent_state
 
 
+@use_kernelized_func([causal_conv1d_update, causal_conv1d_fn])
 class Qwen3NextGatedDeltaNet(nn.Module):
     def __init__(self, config: Qwen3NextConfig, layer_idx: int):
         super().__init__()

From d2beaf6423d3617cd4b133e0423671fa08903742 Mon Sep 17 00:00:00 2001
From: Cyril Vallez <cyril.vallez@gmail.com>
Date: Mon, 13 Jul 2026 23:30:44 +0900
Subject: [PATCH 07/43] combine them

---
 src/transformers/integrations/hub_kernels.py  | 21 +++++++++++++++++++
 .../models/qwen3_5/modeling_qwen3_5.py        | 10 ++++-----
 .../qwen3_5_moe/modeling_qwen3_5_moe.py       | 10 ++++-----
 .../models/qwen3_next/modeling_qwen3_next.py  | 11 +++++-----
 .../models/qwen3_next/modular_qwen3_next.py   | 11 +++++-----
 src/transformers/utils/generic.py             | 20 ------------------
 6 files changed, 39 insertions(+), 44 deletions(-)

diff --git a/src/transformers/integrations/hub_kernels.py b/src/transformers/integrations/hub_kernels.py
index 87c30ec543e7..732b56f30181 100644
--- a/src/transformers/integrations/hub_kernels.py
+++ b/src/transformers/integrations/hub_kernels.py
@@ -853,6 +853,27 @@ def _noop_forward(self, *args, **kwargs):
     kernel_config.kernel_mapping = new_mapping
 
 
+def use_kernel_func_from_hub_with_fallback(func_name: str, package: str):
+    """
+    Decorator that tries to replace the decorated function with `func_name` imported from `package`, if it's available. If not,
+    simply returns the decorated function.
+    Useful to define explicit torch fallback functions, while still using an optimized kernel imported from somewhere else if available.
+    """
+
+    kernel_wrapper_decorator = use_kernel_func_from_hub(func_name)
+
+    def decorator(func: Callable) -> Callable:
+        try:
+            module = importlib.import_module(package)
+            function = getattr(module, func_name)
+        except Exception:
+            function = func
+
+        return kernel_wrapper_decorator(function)
+
+    return decorator
+
+
 __all__ = [
     "LayerRepository",
     "get_kernel",
diff --git a/src/transformers/models/qwen3_5/modeling_qwen3_5.py b/src/transformers/models/qwen3_5/modeling_qwen3_5.py
index 7324b9c23818..0441d04002fc 100644
--- a/src/transformers/models/qwen3_5/modeling_qwen3_5.py
+++ b/src/transformers/models/qwen3_5/modeling_qwen3_5.py
@@ -32,8 +32,9 @@
 from ...activations import ACT2FN
 from ...cache_utils import Cache, DynamicCache
 from ...generation import GenerationMixin
-from ...integrations import use_kernel_forward_from_hub, use_kernel_func_from_hub
+from ...integrations import use_kernel_forward_from_hub
 from ...integrations.accelerate import force_accelerate_hooks
+from ...integrations.hub_kernels import use_kernel_func_from_hub_with_fallback
 from ...masking_utils import create_causal_mask, create_recurrent_attention_mask
 from ...modeling_flash_attention_utils import FlashAttentionKwargs
 from ...modeling_layers import (
@@ -57,7 +58,6 @@
     is_flash_attention_requested,
     maybe_autocast,
     merge_with_config_defaults,
-    replace_with_function_from_package,
 )
 from ...utils.import_utils import is_flash_linear_attention_available
 from ...utils.output_capturing import capture_outputs
@@ -215,8 +215,7 @@ def apply_mask_to_padding_states(hidden_states, attention_mask):
 is_fast_path_available = all((chunk_gated_delta_rule, fused_recurrent_gated_delta_rule))
 
 
-@use_kernel_func_from_hub("causal_conv1d_update")
-@replace_with_function_from_package("causal_conv1d_update", "causal_conv1d")
+@use_kernel_func_from_hub_with_fallback("causal_conv1d_update", "causal_conv1d")
 def causal_conv1d_update(
     hidden_states: torch.Tensor,
     conv_state: torch.Tensor,
@@ -236,8 +235,7 @@ def causal_conv1d_update(
     return out.to(hidden_states.dtype)
 
 
-@use_kernel_func_from_hub("causal_conv1d_fn")
-@replace_with_function_from_package("causal_conv1d_update", "causal_conv1d")
+@use_kernel_func_from_hub_with_fallback("causal_conv1d_fn", "causal_conv1d")
 def causal_conv1d_fn(
     hidden_states: torch.Tensor,
     weight: nn.Parameter,
diff --git a/src/transformers/models/qwen3_5_moe/modeling_qwen3_5_moe.py b/src/transformers/models/qwen3_5_moe/modeling_qwen3_5_moe.py
index 4742680964f2..ff6bb720eaf5 100644
--- a/src/transformers/models/qwen3_5_moe/modeling_qwen3_5_moe.py
+++ b/src/transformers/models/qwen3_5_moe/modeling_qwen3_5_moe.py
@@ -32,8 +32,9 @@
 from ...activations import ACT2FN
 from ...cache_utils import Cache, DynamicCache
 from ...generation import GenerationMixin
-from ...integrations import use_experts_implementation, use_kernel_forward_from_hub, use_kernel_func_from_hub
+from ...integrations import use_experts_implementation, use_kernel_forward_from_hub
 from ...integrations.accelerate import force_accelerate_hooks
+from ...integrations.hub_kernels import use_kernel_func_from_hub_with_fallback
 from ...masking_utils import create_causal_mask, create_recurrent_attention_mask
 from ...modeling_flash_attention_utils import FlashAttentionKwargs
 from ...modeling_layers import GradientCheckpointingLayer
@@ -54,7 +55,6 @@
     is_flash_attention_requested,
     maybe_autocast,
     merge_with_config_defaults,
-    replace_with_function_from_package,
 )
 from ...utils.import_utils import is_flash_linear_attention_available
 from ...utils.output_capturing import OutputRecorder, capture_outputs
@@ -212,8 +212,7 @@ def apply_mask_to_padding_states(hidden_states, attention_mask):
 is_fast_path_available = all((chunk_gated_delta_rule, fused_recurrent_gated_delta_rule))
 
 
-@use_kernel_func_from_hub("causal_conv1d_update")
-@replace_with_function_from_package("causal_conv1d_update", "causal_conv1d")
+@use_kernel_func_from_hub_with_fallback("causal_conv1d_update", "causal_conv1d")
 def causal_conv1d_update(
     hidden_states: torch.Tensor,
     conv_state: torch.Tensor,
@@ -233,8 +232,7 @@ def causal_conv1d_update(
     return out.to(hidden_states.dtype)
 
 
-@use_kernel_func_from_hub("causal_conv1d_fn")
-@replace_with_function_from_package("causal_conv1d_update", "causal_conv1d")
+@use_kernel_func_from_hub_with_fallback("causal_conv1d_fn", "causal_conv1d")
 def causal_conv1d_fn(
     hidden_states: torch.Tensor,
     weight: nn.Parameter,
diff --git a/src/transformers/models/qwen3_next/modeling_qwen3_next.py b/src/transformers/models/qwen3_next/modeling_qwen3_next.py
index 007318584ab9..8bc3ae758c76 100644
--- a/src/transformers/models/qwen3_next/modeling_qwen3_next.py
+++ b/src/transformers/models/qwen3_next/modeling_qwen3_next.py
@@ -29,8 +29,9 @@
 from ...activations import ACT2FN
 from ...cache_utils import Cache, DynamicCache
 from ...generation import GenerationMixin
-from ...integrations import use_experts_implementation, use_kernel_func_from_hub, use_kernelized_func
+from ...integrations import use_experts_implementation, use_kernelized_func
 from ...integrations.accelerate import force_accelerate_hooks
+from ...integrations.hub_kernels import use_kernel_func_from_hub_with_fallback
 from ...masking_utils import create_causal_mask, create_recurrent_attention_mask
 from ...modeling_flash_attention_utils import FlashAttentionKwargs
 from ...modeling_layers import (
@@ -44,7 +45,7 @@
 from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
 from ...processing_utils import Unpack
 from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, logging
-from ...utils.generic import maybe_autocast, merge_with_config_defaults, replace_with_function_from_package
+from ...utils.generic import maybe_autocast, merge_with_config_defaults
 from ...utils.import_utils import is_flash_linear_attention_available
 from ...utils.output_capturing import OutputRecorder, capture_outputs
 from .configuration_qwen3_next import Qwen3NextConfig
@@ -340,8 +341,7 @@ def apply_mask_to_padding_states(hidden_states, attention_mask):
 is_fast_path_available = all((chunk_gated_delta_rule, fused_recurrent_gated_delta_rule))
 
 
-@use_kernel_func_from_hub("causal_conv1d_update")
-@replace_with_function_from_package("causal_conv1d_update", "causal_conv1d")
+@use_kernel_func_from_hub_with_fallback("causal_conv1d_update", "causal_conv1d")
 def causal_conv1d_update(
     hidden_states: torch.Tensor,
     conv_state: torch.Tensor,
@@ -361,8 +361,7 @@ def causal_conv1d_update(
     return out.to(hidden_states.dtype)
 
 
-@use_kernel_func_from_hub("causal_conv1d_fn")
-@replace_with_function_from_package("causal_conv1d_update", "causal_conv1d")
+@use_kernel_func_from_hub_with_fallback("causal_conv1d_fn", "causal_conv1d")
 def causal_conv1d_fn(
     hidden_states: torch.Tensor,
     weight: nn.Parameter,
diff --git a/src/transformers/models/qwen3_next/modular_qwen3_next.py b/src/transformers/models/qwen3_next/modular_qwen3_next.py
index 463761254842..9658c252d0c2 100644
--- a/src/transformers/models/qwen3_next/modular_qwen3_next.py
+++ b/src/transformers/models/qwen3_next/modular_qwen3_next.py
@@ -23,15 +23,16 @@
 from ... import initialization as init
 from ...activations import ACT2FN
 from ...cache_utils import Cache, DynamicCache
-from ...integrations import use_kernel_func_from_hub, use_kernelized_func
+from ...integrations import use_kernelized_func
 from ...integrations.accelerate import force_accelerate_hooks
+from ...integrations.hub_kernels import use_kernel_func_from_hub_with_fallback
 from ...masking_utils import create_causal_mask, create_recurrent_attention_mask
 from ...modeling_flash_attention_utils import FlashAttentionKwargs
 from ...modeling_outputs import MoeCausalLMOutputWithPast, MoeModelOutputWithPast
 from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
 from ...processing_utils import Unpack
 from ...utils import TransformersKwargs, auto_docstring, logging
-from ...utils.generic import merge_with_config_defaults, no_inherit_decorator, replace_with_function_from_package
+from ...utils.generic import merge_with_config_defaults, no_inherit_decorator
 from ...utils.import_utils import is_flash_linear_attention_available
 from ...utils.output_capturing import OutputRecorder, capture_outputs
 from ..bamba.modeling_bamba import apply_mask_to_padding_states, apply_rotary_pos_emb
@@ -179,8 +180,7 @@ def forward(
         return attn_output, attn_weights
 
 
-@use_kernel_func_from_hub("causal_conv1d_update")
-@replace_with_function_from_package("causal_conv1d_update", "causal_conv1d")
+@use_kernel_func_from_hub_with_fallback("causal_conv1d_update", "causal_conv1d")
 def causal_conv1d_update(
     hidden_states: torch.Tensor,
     conv_state: torch.Tensor,
@@ -200,8 +200,7 @@ def causal_conv1d_update(
     return out.to(hidden_states.dtype)
 
 
-@use_kernel_func_from_hub("causal_conv1d_fn")
-@replace_with_function_from_package("causal_conv1d_update", "causal_conv1d")
+@use_kernel_func_from_hub_with_fallback("causal_conv1d_fn", "causal_conv1d")
 def causal_conv1d_fn(
     hidden_states: torch.Tensor,
     weight: nn.Parameter,
diff --git a/src/transformers/utils/generic.py b/src/transformers/utils/generic.py
index 7cb8e0a46949..9a5a97a023fd 100644
--- a/src/transformers/utils/generic.py
+++ b/src/transformers/utils/generic.py
@@ -17,7 +17,6 @@
 
 from __future__ import annotations
 
-import importlib
 import inspect
 import json
 import os
@@ -1157,22 +1156,3 @@ def wrapper(*args, **kwargs):
         return wrapper
 
     return decorator
-
-
-def replace_with_function_from_package(function_name: str, package: str):
-    """
-    Decorator that tries to replace the decorated function with `function_name` imported from `package`, if it's available. If not,
-    simply returns the decorated function.
-    Useful to define explicit torch fallback functions, while still using an optimized kernel imported from somewhere else if available.
-    """
-
-    def decorator(func: Callable) -> Callable:
-        try:
-            module = importlib.import_module(package)
-            function = getattr(module, function_name)
-        except Exception:
-            function = func
-
-        return function
-
-    return decorator

From c256c47b5d7aeb8670f11735983d23d07a7adb90 Mon Sep 17 00:00:00 2001
From: Cyril Vallez <cyril.vallez@gmail.com>
Date: Mon, 13 Jul 2026 23:36:51 +0900
Subject: [PATCH 08/43] doc

---
 src/transformers/integrations/hub_kernels.py | 11 ++++++++---
 1 file changed, 8 insertions(+), 3 deletions(-)

diff --git a/src/transformers/integrations/hub_kernels.py b/src/transformers/integrations/hub_kernels.py
index 732b56f30181..d774575ad3cd 100644
--- a/src/transformers/integrations/hub_kernels.py
+++ b/src/transformers/integrations/hub_kernels.py
@@ -855,9 +855,14 @@ def _noop_forward(self, *args, **kwargs):
 
 def use_kernel_func_from_hub_with_fallback(func_name: str, package: str):
     """
-    Decorator that tries to replace the decorated function with `func_name` imported from `package`, if it's available. If not,
-    simply returns the decorated function.
-    Useful to define explicit torch fallback functions, while still using an optimized kernel imported from somewhere else if available.
+    Similar to `use_kernel_func_from_hub`, but first tries to replace the decorated function with `func_name` imported from `package`,
+    if it's available. If not, simply applies `use_kernel_func_from_hub` on the decorated function.
+    Useful to define explicit torch fallback functions, while still using an optimized implementations from either `kernels` or
+    an auxiliary package (e.g. `causal_conv1d`) if available.
+    This means that the precedence order will be the following:
+    - 1st `kernels`, if it's available and `use_kernels=True`
+    - if the above is not True, then `func_name` imported from `package` if `package` is available
+    - if None of the above are True, the base decorated function
     """
 
     kernel_wrapper_decorator = use_kernel_func_from_hub(func_name)

From 30e0ff74e97487c9a64694e919fb406a99f2cf45 Mon Sep 17 00:00:00 2001
From: vasqu <antonprogamer@gmail.com>
Date: Mon, 13 Jul 2026 20:42:50 +0200
Subject: [PATCH 09/43] ignore kwargs, allow og

---
 src/transformers/integrations/__init__.py     |  2 +
 src/transformers/integrations/hub_kernels.py  | 54 ++++++++++---------
 .../models/qwen3_5/modeling_qwen3_5.py        |  3 +-
 .../qwen3_5_moe/modeling_qwen3_5_moe.py       |  7 ++-
 .../models/qwen3_next/modeling_qwen3_next.py  |  3 +-
 .../models/qwen3_next/modular_qwen3_next.py   |  3 +-
 6 files changed, 38 insertions(+), 34 deletions(-)

diff --git a/src/transformers/integrations/__init__.py b/src/transformers/integrations/__init__.py
index 2ecc31ae54cf..5d6fc98f1f25 100755
--- a/src/transformers/integrations/__init__.py
+++ b/src/transformers/integrations/__init__.py
@@ -80,6 +80,7 @@
         "replace_kernel_forward_from_hub",
         "use_kernel_forward_from_hub",
         "use_kernel_func_from_hub",
+        "use_kernel_func_from_hub_with_fallback",
         "use_kernelized_func",
     ],
     "integration_utils": [
@@ -241,6 +242,7 @@
         replace_kernel_forward_from_hub,
         use_kernel_forward_from_hub,
         use_kernel_func_from_hub,
+        use_kernel_func_from_hub_with_fallback,
         use_kernelized_func,
     )
     from .integration_utils import (
diff --git a/src/transformers/integrations/hub_kernels.py b/src/transformers/integrations/hub_kernels.py
index d774575ad3cd..c69528c001c1 100644
--- a/src/transformers/integrations/hub_kernels.py
+++ b/src/transformers/integrations/hub_kernels.py
@@ -13,6 +13,7 @@
 # limitations under the License.
 import functools
 import importlib
+import inspect
 import os
 import re
 import sys
@@ -32,6 +33,7 @@
     is_kernels_available,
     is_rocm_platform,
     is_torch_available,
+    resolve_internal_import,
 )
 from .flash_attention import flash_attention_forward
 
@@ -654,6 +656,32 @@ def new_init(self, *args, **kwargs):
     return decorator
 
 
+def use_kernel_func_from_hub_with_fallback(package: str, func_name: str, internal_path: str | None = None):
+    # TODO: change when we sync to 0.16.x+
+    kernel_wrapper_decorator = use_kernel_func_from_hub(func_name)
+
+    def decorator(torch_function: Callable) -> Callable:
+        implementation = None
+        try:
+            module = importlib.import_module(package)
+            implementation = resolve_internal_import(module, internal_path or func_name)
+        except Exception:
+            implementation = torch_function
+        finally:
+            implementation = torch_function if implementation is None else implementation
+
+        applicable_params = inspect.signature(implementation).parameters
+
+        @functools.wraps(torch_function)
+        def wrapped(*args, **kwargs):
+            kwargs = {k: v for k, v in kwargs.items() if k in applicable_params}
+            return implementation(*args, **kwargs)
+
+        return kernel_wrapper_decorator(wrapped)
+
+    return decorator
+
+
 # Whether to allow hub kernels coming from untrusted repos, i.e. repos outside `kernels-community`
 ALLOW_ALL_KERNELS = False
 
@@ -853,32 +881,6 @@ def _noop_forward(self, *args, **kwargs):
     kernel_config.kernel_mapping = new_mapping
 
 
-def use_kernel_func_from_hub_with_fallback(func_name: str, package: str):
-    """
-    Similar to `use_kernel_func_from_hub`, but first tries to replace the decorated function with `func_name` imported from `package`,
-    if it's available. If not, simply applies `use_kernel_func_from_hub` on the decorated function.
-    Useful to define explicit torch fallback functions, while still using an optimized implementations from either `kernels` or
-    an auxiliary package (e.g. `causal_conv1d`) if available.
-    This means that the precedence order will be the following:
-    - 1st `kernels`, if it's available and `use_kernels=True`
-    - if the above is not True, then `func_name` imported from `package` if `package` is available
-    - if None of the above are True, the base decorated function
-    """
-
-    kernel_wrapper_decorator = use_kernel_func_from_hub(func_name)
-
-    def decorator(func: Callable) -> Callable:
-        try:
-            module = importlib.import_module(package)
-            function = getattr(module, func_name)
-        except Exception:
-            function = func
-
-        return kernel_wrapper_decorator(function)
-
-    return decorator
-
-
 __all__ = [
     "LayerRepository",
     "get_kernel",
diff --git a/src/transformers/models/qwen3_5/modeling_qwen3_5.py b/src/transformers/models/qwen3_5/modeling_qwen3_5.py
index 0441d04002fc..dad105e8b954 100644
--- a/src/transformers/models/qwen3_5/modeling_qwen3_5.py
+++ b/src/transformers/models/qwen3_5/modeling_qwen3_5.py
@@ -32,9 +32,8 @@
 from ...activations import ACT2FN
 from ...cache_utils import Cache, DynamicCache
 from ...generation import GenerationMixin
-from ...integrations import use_kernel_forward_from_hub
+from ...integrations import use_kernel_forward_from_hub, use_kernel_func_from_hub_with_fallback
 from ...integrations.accelerate import force_accelerate_hooks
-from ...integrations.hub_kernels import use_kernel_func_from_hub_with_fallback
 from ...masking_utils import create_causal_mask, create_recurrent_attention_mask
 from ...modeling_flash_attention_utils import FlashAttentionKwargs
 from ...modeling_layers import (
diff --git a/src/transformers/models/qwen3_5_moe/modeling_qwen3_5_moe.py b/src/transformers/models/qwen3_5_moe/modeling_qwen3_5_moe.py
index ff6bb720eaf5..2ad29909ec05 100644
--- a/src/transformers/models/qwen3_5_moe/modeling_qwen3_5_moe.py
+++ b/src/transformers/models/qwen3_5_moe/modeling_qwen3_5_moe.py
@@ -32,9 +32,12 @@
 from ...activations import ACT2FN
 from ...cache_utils import Cache, DynamicCache
 from ...generation import GenerationMixin
-from ...integrations import use_experts_implementation, use_kernel_forward_from_hub
+from ...integrations import (
+    use_experts_implementation,
+    use_kernel_forward_from_hub,
+    use_kernel_func_from_hub_with_fallback,
+)
 from ...integrations.accelerate import force_accelerate_hooks
-from ...integrations.hub_kernels import use_kernel_func_from_hub_with_fallback
 from ...masking_utils import create_causal_mask, create_recurrent_attention_mask
 from ...modeling_flash_attention_utils import FlashAttentionKwargs
 from ...modeling_layers import GradientCheckpointingLayer
diff --git a/src/transformers/models/qwen3_next/modeling_qwen3_next.py b/src/transformers/models/qwen3_next/modeling_qwen3_next.py
index 8bc3ae758c76..e618ad1462e8 100644
--- a/src/transformers/models/qwen3_next/modeling_qwen3_next.py
+++ b/src/transformers/models/qwen3_next/modeling_qwen3_next.py
@@ -29,9 +29,8 @@
 from ...activations import ACT2FN
 from ...cache_utils import Cache, DynamicCache
 from ...generation import GenerationMixin
-from ...integrations import use_experts_implementation, use_kernelized_func
+from ...integrations import use_experts_implementation, use_kernel_func_from_hub_with_fallback, use_kernelized_func
 from ...integrations.accelerate import force_accelerate_hooks
-from ...integrations.hub_kernels import use_kernel_func_from_hub_with_fallback
 from ...masking_utils import create_causal_mask, create_recurrent_attention_mask
 from ...modeling_flash_attention_utils import FlashAttentionKwargs
 from ...modeling_layers import (
diff --git a/src/transformers/models/qwen3_next/modular_qwen3_next.py b/src/transformers/models/qwen3_next/modular_qwen3_next.py
index 9658c252d0c2..6e9e7bffba37 100644
--- a/src/transformers/models/qwen3_next/modular_qwen3_next.py
+++ b/src/transformers/models/qwen3_next/modular_qwen3_next.py
@@ -23,9 +23,8 @@
 from ... import initialization as init
 from ...activations import ACT2FN
 from ...cache_utils import Cache, DynamicCache
-from ...integrations import use_kernelized_func
+from ...integrations import use_kernel_func_from_hub_with_fallback, use_kernelized_func
 from ...integrations.accelerate import force_accelerate_hooks
-from ...integrations.hub_kernels import use_kernel_func_from_hub_with_fallback
 from ...masking_utils import create_causal_mask, create_recurrent_attention_mask
 from ...modeling_flash_attention_utils import FlashAttentionKwargs
 from ...modeling_outputs import MoeCausalLMOutputWithPast, MoeModelOutputWithPast

From 38c2ece9b1a404cc972870d42ed2378ac651466e Mon Sep 17 00:00:00 2001
From: vasqu <antonprogamer@gmail.com>
Date: Wed, 29 Jul 2026 17:23:33 +0000
Subject: [PATCH 10/43] fix

---
 src/transformers/integrations/hub_kernels.py | 26 ++++++++++++++++++++
 1 file changed, 26 insertions(+)

diff --git a/src/transformers/integrations/hub_kernels.py b/src/transformers/integrations/hub_kernels.py
index 62373488f7e8..7bd9252ce783 100644
--- a/src/transformers/integrations/hub_kernels.py
+++ b/src/transformers/integrations/hub_kernels.py
@@ -640,6 +640,32 @@ def get_kernel(
     )
 
 
+def use_kernel_func_from_hub_with_fallback(package: str, func_name: str, internal_path: str | None = None):
+    # TODO: change when we sync to 0.16.x+
+    kernel_wrapper_decorator = use_kernel_func_from_hub(func_name)
+
+    def decorator(torch_function: Callable) -> Callable:
+        implementation = None
+        try:
+            module = importlib.import_module(package)
+            implementation = resolve_internal_import(module, internal_path or func_name)
+        except Exception:
+            implementation = torch_function
+        finally:
+            implementation = torch_function if implementation is None else implementation
+
+        applicable_params = inspect.signature(implementation).parameters
+
+        @functools.wraps(torch_function)
+        def wrapped(*args, **kwargs):
+            kwargs = {k: v for k, v in kwargs.items() if k in applicable_params}
+            return implementation(*args, **kwargs)
+
+        return kernel_wrapper_decorator(wrapped)
+
+    return decorator
+
+
 # Whether to allow hub kernels coming from untrusted repos, i.e. repos outside `kernels-community`
 ALLOW_ALL_KERNELS = False
 

From 7788cac43ab86b14571be5939667de9afdfec5b8 Mon Sep 17 00:00:00 2001
From: vasqu <antonprogamer@gmail.com>
Date: Wed, 29 Jul 2026 21:01:34 +0000
Subject: [PATCH 11/43] gdn like paths (missing conv of olmo hybrid)

---
 src/transformers/integrations/hub_kernels.py  | 52 ++++++++++++--
 .../olmo_hybrid/modeling_olmo_hybrid.py       | 59 ++++++---------
 .../models/olmo_hybrid/modular_olmo_hybrid.py | 35 ++-------
 .../models/qwen3_5/modeling_qwen3_5.py        | 65 +++++++----------
 .../models/qwen3_5/modular_qwen3_5.py         | 17 +++--
 .../qwen3_5_moe/modeling_qwen3_5_moe.py       | 60 ++++++----------
 .../models/qwen3_next/modeling_qwen3_next.py  | 71 ++++++++-----------
 .../models/qwen3_next/modular_qwen3_next.py   | 61 +++++++---------
 8 files changed, 189 insertions(+), 231 deletions(-)

diff --git a/src/transformers/integrations/hub_kernels.py b/src/transformers/integrations/hub_kernels.py
index 7bd9252ce783..9ddeb8094bda 100644
--- a/src/transformers/integrations/hub_kernels.py
+++ b/src/transformers/integrations/hub_kernels.py
@@ -179,6 +179,34 @@ def _build_kernel_mapping() -> dict:
                     ),
                 },
             },
+            "chunk_gated_delta_rule": {
+                "cuda": {
+                    Mode.TRAINING: LayerRepository(
+                        repo_id="kernels-community/fla",
+                        layer_name="chunk_gated_delta_rule",
+                        version=1,
+                    ),
+                    Mode.INFERENCE: LayerRepository(
+                        repo_id="kernels-community/fla",
+                        layer_name="chunk_gated_delta_rule",
+                        version=1,
+                    ),
+                },
+            },
+            "recurrent_gated_delta_rule": {
+                "cuda": {
+                    Mode.TRAINING: LayerRepository(
+                        repo_id="kernels-community/fla",
+                        layer_name="recurrent_gated_delta_rule",
+                        version=1,
+                    ),
+                    Mode.INFERENCE: LayerRepository(
+                        repo_id="kernels-community/fla",
+                        layer_name="recurrent_gated_delta_rule",
+                        version=1,
+                    ),
+                },
+            },
             "SwiGLUMLP": {
                 "cuda": {
                     Mode.INFERENCE | Mode.TORCH_COMPILE: LayerRepository(
@@ -269,6 +297,20 @@ def _build_kernel_mapping() -> dict:
                     ),
                 },
             },
+            "RMSNormGated": {
+                "cuda": {
+                    Mode.TRAINING: LayerRepository(
+                        repo_id="kernels-community/fla",
+                        layer_name="FusedRMSNormGated",
+                        version=1,
+                    ),
+                    Mode.INFERENCE: LayerRepository(
+                        repo_id="kernels-community/fla",
+                        layer_name="FusedRMSNormGated",
+                        version=1,
+                    ),
+                },
+            },
             "MegaBlocksMoeMLP": {
                 "cuda": {
                     Mode.TRAINING: LayerRepository(
@@ -640,15 +682,17 @@ def get_kernel(
     )
 
 
-def use_kernel_func_from_hub_with_fallback(package: str, func_name: str, internal_path: str | None = None):
-    # TODO: change when we sync to 0.16.x+
-    kernel_wrapper_decorator = use_kernel_func_from_hub(func_name)
+def use_kernel_func_from_hub_with_fallback(func_name: str, package: str, internal_path: str | None = None):
+    kernel_wrapper_decorator = use_kernel_forward_from_hub(func_name)
+
+    # Allow internal path prefix if given to resolve non __init__ imports
+    full_path = func_name if internal_path is None else f"{internal_path}.{func_name}"
 
     def decorator(torch_function: Callable) -> Callable:
         implementation = None
         try:
             module = importlib.import_module(package)
-            implementation = resolve_internal_import(module, internal_path or func_name)
+            implementation = resolve_internal_import(module, full_path)
         except Exception:
             implementation = torch_function
         finally:
diff --git a/src/transformers/models/olmo_hybrid/modeling_olmo_hybrid.py b/src/transformers/models/olmo_hybrid/modeling_olmo_hybrid.py
index 0d86b0f6c209..0e7d5e64e963 100644
--- a/src/transformers/models/olmo_hybrid/modeling_olmo_hybrid.py
+++ b/src/transformers/models/olmo_hybrid/modeling_olmo_hybrid.py
@@ -31,14 +31,14 @@
 from ...activations import ACT2FN
 from ...cache_utils import Cache
 from ...generation import GenerationMixin
-from ...integrations import use_kernel_forward_from_hub, use_kernelized_func
+from ...integrations import use_kernel_forward_from_hub, use_kernel_func_from_hub_with_fallback, use_kernelized_func
 from ...masking_utils import create_causal_mask, create_recurrent_attention_mask
 from ...modeling_layers import GradientCheckpointingLayer
 from ...modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
 from ...modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update
 from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
 from ...processing_utils import Unpack
-from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, logging
+from ...utils import TransformersKwargs, auto_docstring, can_return_tuple
 from ...utils.generic import maybe_autocast, merge_with_config_defaults
 from ...utils.import_utils import is_flash_linear_attention_available
 from ...utils.output_capturing import capture_outputs
@@ -46,17 +46,11 @@
 
 
 if is_flash_linear_attention_available():
-    from fla.modules import FusedRMSNormGated, ShortConvolution
-    from fla.ops.gated_delta_rule import chunk_gated_delta_rule, fused_recurrent_gated_delta_rule
+    from fla.modules import ShortConvolution
 else:
-    chunk_gated_delta_rule, fused_recurrent_gated_delta_rule = None, None
-    FusedRMSNormGated = None
     ShortConvolution = None
 
 
-logger = logging.get_logger(__name__)
-
-
 class OlmoHybridDynamicCache:
     """
     Cache for hybrid model supporting both attention KV cache and linear attention state.
@@ -162,11 +156,13 @@ def get_query_offset(self, layer_idx: int = 0) -> int:
         return self.get_seq_length(layer_idx=layer_idx)
 
 
+@use_kernel_forward_from_hub("RMSNormGated")
 class OlmoHybridRMSNormGated(nn.Module):
     def __init__(self, hidden_size, eps=1e-6, **kwargs):
         super().__init__()
         self.weight = nn.Parameter(torch.ones(hidden_size))
         self.variance_epsilon = eps
+        self.activation = "silu"
 
     def forward(self, hidden_states, gate=None):
         input_dtype = hidden_states.dtype
@@ -175,7 +171,7 @@ def forward(self, hidden_states, gate=None):
         # Norm before gate
         hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
         hidden_states = self.weight * hidden_states.to(input_dtype)
-        hidden_states = hidden_states * F.silu(gate.to(torch.float32))
+        hidden_states = hidden_states * ACT2FN[self.activation](gate.to(torch.float32))
 
         return hidden_states.to(input_dtype)
 
@@ -503,6 +499,7 @@ def l2norm(x: torch.FloatTensor, dim: int = -1, eps: float = 1e-6):
     return x * inv_norm
 
 
+@use_kernel_func_from_hub_with_fallback("chunk_gated_delta_rule", "fla", internal_path="ops.gated_delta_rule")
 def torch_chunk_gated_delta_rule(
     query,
     key,
@@ -584,8 +581,17 @@ def torch_chunk_gated_delta_rule(
     return core_attn_out, last_recurrent_state
 
 
+@use_kernel_func_from_hub_with_fallback("recurrent_gated_delta_rule", "fla", internal_path="ops.gated_delta_rule")
 def torch_recurrent_gated_delta_rule(
-    query, key, value, g, beta, initial_state, output_final_state, use_qk_l2norm_in_kernel=False
+    query,
+    key,
+    value,
+    g,
+    beta,
+    initial_state,
+    output_final_state,
+    use_qk_l2norm_in_kernel=False,
+    **kwargs,
 ):
     initial_dtype = query.dtype
     if use_qk_l2norm_in_kernel:
@@ -628,11 +634,6 @@ def torch_recurrent_gated_delta_rule(
     return core_attn_out, last_recurrent_state
 
 
-is_fast_path_available = all(
-    (ShortConvolution, chunk_gated_delta_rule, fused_recurrent_gated_delta_rule, FusedRMSNormGated)
-)
-
-
 class OlmoHybridGatedDeltaNet(nn.Module):
     """
     GatedDeltaNet linear attention for OLMo Hybrid.
@@ -668,6 +669,7 @@ def __init__(self, config: OlmoHybridConfig, layer_idx: int):
 
         self.o_proj = nn.Linear(self.value_dim, self.hidden_size, bias=False)
 
+        # TODO: can be moved into kernels as well
         Conv1dClass = ShortConvolution if ShortConvolution is not None else OlmoHybridShortConvolution
 
         self.q_conv1d = Conv1dClass(
@@ -703,24 +705,7 @@ def __init__(self, config: OlmoHybridConfig, layer_idx: int):
         self.dt_bias = nn.Parameter(inv_dt)
 
         # Output norm - NOTE: FLA's FusedRMSNormGated uses eps=1e-5 by default
-        self.o_norm = (
-            OlmoHybridRMSNormGated(self.head_v_dim, eps=1e-5)
-            if FusedRMSNormGated is None
-            else FusedRMSNormGated(
-                self.head_v_dim,
-                eps=1e-5,
-            )
-        )
-
-        self.chunk_gated_delta_rule = chunk_gated_delta_rule or torch_chunk_gated_delta_rule
-        self.recurrent_gated_delta_rule = fused_recurrent_gated_delta_rule or torch_recurrent_gated_delta_rule
-
-        if not is_fast_path_available:
-            logger.warning_once(
-                "The fast path is not available because one of the required libraries is not installed. "
-                "Falling back to torch implementation. To install, follow: "
-                "https://github.com/fla-org/flash-linear-attention#installation"
-            )
+        self.o_norm = OlmoHybridRMSNormGated(self.head_v_dim, eps=1e-5)
 
         self.layer_type = config.layer_types[layer_idx]
 
@@ -782,7 +767,7 @@ def forward(
         g = -self.A_log.float().exp() * F.softplus(self.a_proj(hidden_states).float() + self.dt_bias)
 
         if use_precomputed and seq_len == 1:
-            output, new_recurrent_state = self.recurrent_gated_delta_rule(
+            output, new_recurrent_state = torch_recurrent_gated_delta_rule(
                 q,
                 k,
                 v,
@@ -791,9 +776,10 @@ def forward(
                 initial_state=recurrent_state,
                 output_final_state=use_cache,
                 use_qk_l2norm_in_kernel=True,
+                **kwargs,
             )
         else:
-            output, new_recurrent_state = self.chunk_gated_delta_rule(
+            output, new_recurrent_state = torch_chunk_gated_delta_rule(
                 q,
                 k,
                 v,
@@ -802,6 +788,7 @@ def forward(
                 initial_state=recurrent_state if use_precomputed else None,
                 output_final_state=use_cache,
                 use_qk_l2norm_in_kernel=True,
+                **kwargs,
             )
 
         if cache_params is not None:
diff --git a/src/transformers/models/olmo_hybrid/modular_olmo_hybrid.py b/src/transformers/models/olmo_hybrid/modular_olmo_hybrid.py
index 505eb16a06a7..656b4251f01a 100644
--- a/src/transformers/models/olmo_hybrid/modular_olmo_hybrid.py
+++ b/src/transformers/models/olmo_hybrid/modular_olmo_hybrid.py
@@ -58,17 +58,10 @@
 
 
 if is_flash_linear_attention_available():
-    from fla.modules import FusedRMSNormGated, ShortConvolution
-    from fla.ops.gated_delta_rule import chunk_gated_delta_rule, fused_recurrent_gated_delta_rule
+    from fla.modules import ShortConvolution
 else:
-    chunk_gated_delta_rule, fused_recurrent_gated_delta_rule = None, None
-    FusedRMSNormGated = None
     ShortConvolution = None
 
-is_fast_path_available = all(
-    (ShortConvolution, chunk_gated_delta_rule, fused_recurrent_gated_delta_rule, FusedRMSNormGated)
-)
-
 
 logger = logging.get_logger(__name__)
 
@@ -481,6 +474,7 @@ def __init__(self, config: OlmoHybridConfig, layer_idx: int):
 
         self.o_proj = nn.Linear(self.value_dim, self.hidden_size, bias=False)
 
+        # TODO: can be moved into kernels as well
         Conv1dClass = ShortConvolution if ShortConvolution is not None else OlmoHybridShortConvolution
 
         self.q_conv1d = Conv1dClass(
@@ -516,24 +510,7 @@ def __init__(self, config: OlmoHybridConfig, layer_idx: int):
         self.dt_bias = nn.Parameter(inv_dt)
 
         # Output norm - NOTE: FLA's FusedRMSNormGated uses eps=1e-5 by default
-        self.o_norm = (
-            OlmoHybridRMSNormGated(self.head_v_dim, eps=1e-5)
-            if FusedRMSNormGated is None
-            else FusedRMSNormGated(
-                self.head_v_dim,
-                eps=1e-5,
-            )
-        )
-
-        self.chunk_gated_delta_rule = chunk_gated_delta_rule or torch_chunk_gated_delta_rule
-        self.recurrent_gated_delta_rule = fused_recurrent_gated_delta_rule or torch_recurrent_gated_delta_rule
-
-        if not is_fast_path_available:
-            logger.warning_once(
-                "The fast path is not available because one of the required libraries is not installed. "
-                "Falling back to torch implementation. To install, follow: "
-                "https://github.com/fla-org/flash-linear-attention#installation"
-            )
+        self.o_norm = OlmoHybridRMSNormGated(self.head_v_dim, eps=1e-5)
 
         self.layer_type = config.layer_types[layer_idx]
 
@@ -595,7 +572,7 @@ def forward(
         g = -self.A_log.float().exp() * F.softplus(self.a_proj(hidden_states).float() + self.dt_bias)
 
         if use_precomputed and seq_len == 1:
-            output, new_recurrent_state = self.recurrent_gated_delta_rule(
+            output, new_recurrent_state = torch_recurrent_gated_delta_rule(
                 q,
                 k,
                 v,
@@ -604,9 +581,10 @@ def forward(
                 initial_state=recurrent_state,
                 output_final_state=use_cache,
                 use_qk_l2norm_in_kernel=True,
+                **kwargs,
             )
         else:
-            output, new_recurrent_state = self.chunk_gated_delta_rule(
+            output, new_recurrent_state = torch_chunk_gated_delta_rule(
                 q,
                 k,
                 v,
@@ -615,6 +593,7 @@ def forward(
                 initial_state=recurrent_state if use_precomputed else None,
                 output_final_state=use_cache,
                 use_qk_l2norm_in_kernel=True,
+                **kwargs,
             )
 
         if cache_params is not None:
diff --git a/src/transformers/models/qwen3_5/modeling_qwen3_5.py b/src/transformers/models/qwen3_5/modeling_qwen3_5.py
index 4d0f6359cf73..e225715310dd 100644
--- a/src/transformers/models/qwen3_5/modeling_qwen3_5.py
+++ b/src/transformers/models/qwen3_5/modeling_qwen3_5.py
@@ -32,7 +32,7 @@
 from ...activations import ACT2FN
 from ...cache_utils import Cache, DynamicCache
 from ...generation import GenerationMixin
-from ...integrations import use_kernel_forward_from_hub, use_kernel_func_from_hub_with_fallback
+from ...integrations import use_kernel_forward_from_hub, use_kernel_func_from_hub_with_fallback, use_kernelized_func
 from ...integrations.accelerate import force_accelerate_hooks
 from ...masking_utils import create_causal_mask, create_recurrent_attention_mask
 from ...modeling_flash_attention_utils import FlashAttentionKwargs
@@ -50,7 +50,7 @@
 from ...modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update
 from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
 from ...processing_utils import Unpack
-from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, logging, torch_compilable_check
+from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, torch_compilable_check
 from ...utils.generic import (
     accepts_precomputed_kwargs,
     get_max_seqlen,
@@ -58,7 +58,6 @@
     maybe_autocast,
     merge_with_config_defaults,
 )
-from ...utils.import_utils import is_causal_conv1d_available, is_flash_linear_attention_available
 from ...utils.output_capturing import capture_outputs
 from ...vision_utils import (
     get_vision_attention_seqlens,
@@ -69,17 +68,6 @@
 from .configuration_qwen3_5 import Qwen3_5Config, Qwen3_5TextConfig, Qwen3_5VisionConfig
 
 
-if is_flash_linear_attention_available():
-    from fla.modules import FusedRMSNormGated
-    from fla.ops.gated_delta_rule import chunk_gated_delta_rule, fused_recurrent_gated_delta_rule
-else:
-    chunk_gated_delta_rule, fused_recurrent_gated_delta_rule = None, None
-    FusedRMSNormGated = None
-
-
-logger = logging.get_logger(__name__)
-
-
 class Qwen3_5VisionRotaryEmbedding(nn.Module):
     inv_freq: torch.Tensor  # fix linting for `register_buffer`
 
@@ -186,11 +174,13 @@ def apply_interleaved_mrope(self, freqs, mrope_section):
         return freqs_t
 
 
+@use_kernel_forward_from_hub("RMSNormGated")
 class Qwen3_5RMSNormGated(nn.Module):
     def __init__(self, hidden_size, eps=1e-6, **kwargs):
         super().__init__()
         self.weight = nn.Parameter(torch.ones(hidden_size))
         self.variance_epsilon = eps
+        self.activation = "silu"
 
     def forward(self, hidden_states, gate=None):
         input_dtype = hidden_states.dtype
@@ -199,7 +189,7 @@ def forward(self, hidden_states, gate=None):
         # Norm before gate
         hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
         hidden_states = self.weight * hidden_states.to(input_dtype)
-        hidden_states = hidden_states * F.silu(gate.to(torch.float32))
+        hidden_states = hidden_states * ACT2FN[self.activation](gate.to(torch.float32))
 
         return hidden_states.to(input_dtype)
 
@@ -265,6 +255,7 @@ def l2norm(x: torch.FloatTensor, dim: int = -1, eps: float = 1e-6):
     return x * inv_norm
 
 
+@use_kernel_func_from_hub_with_fallback("chunk_gated_delta_rule", "fla", internal_path="ops.gated_delta_rule")
 def torch_chunk_gated_delta_rule(
     query,
     key,
@@ -346,8 +337,17 @@ def torch_chunk_gated_delta_rule(
     return core_attn_out, last_recurrent_state
 
 
+@use_kernel_func_from_hub_with_fallback("recurrent_gated_delta_rule", "fla", internal_path="ops.gated_delta_rule")
 def torch_recurrent_gated_delta_rule(
-    query, key, value, g, beta, initial_state, output_final_state, use_qk_l2norm_in_kernel=False
+    query,
+    key,
+    value,
+    g,
+    beta,
+    initial_state,
+    output_final_state,
+    use_qk_l2norm_in_kernel=False,
+    **kwargs,
 ):
     initial_dtype = query.dtype
     if use_qk_l2norm_in_kernel:
@@ -391,6 +391,9 @@ def torch_recurrent_gated_delta_rule(
 
 
 @use_kernel_forward_from_hub("Qwen3_5GatedDeltaNet")
+@use_kernelized_func(
+    [torch_recurrent_gated_delta_rule, torch_chunk_gated_delta_rule, causal_conv1d_fn, causal_conv1d_update]
+)
 class Qwen3_5GatedDeltaNet(nn.Module):
     def __init__(self, config: Qwen3_5Config, layer_idx: int):
         super().__init__()
@@ -425,27 +428,9 @@ def __init__(self, config: Qwen3_5Config, layer_idx: int):
         A = torch.empty(self.num_v_heads).uniform_(0, 16)
         self.A_log = nn.Parameter(torch.log(A))
 
-        self.norm = (
-            Qwen3_5RMSNormGated(self.head_v_dim, eps=self.layer_norm_epsilon)
-            if FusedRMSNormGated is None
-            else FusedRMSNormGated(
-                self.head_v_dim,
-                eps=self.layer_norm_epsilon,
-                activation=self.activation,
-            )
-        )
-
+        self.norm = Qwen3_5RMSNormGated(self.head_v_dim, eps=self.layer_norm_epsilon)
         self.out_proj = nn.Linear(self.value_dim, self.hidden_size, bias=False)
 
-        self.chunk_gated_delta_rule = chunk_gated_delta_rule or torch_chunk_gated_delta_rule
-        self.recurrent_gated_delta_rule = fused_recurrent_gated_delta_rule or torch_recurrent_gated_delta_rule
-
-        if not is_flash_linear_attention_available() or not is_causal_conv1d_available():
-            logger.warning_once(
-                "The fast path is not available because the required library is not installed. Falling back to "
-                "torch implementation. To install follow https://github.com/fla-org/flash-linear-attention#installation"
-            )
-
         self.layer_type = config.layer_types[layer_idx]
 
         self.in_proj_qkv = nn.Linear(self.hidden_size, self.key_dim * 2 + self.value_dim, bias=False)
@@ -497,7 +482,7 @@ def forward(
                 self.conv1d.weight.squeeze(1),
                 self.conv1d.bias,
                 activation=self.activation,
-                seq_idx=kwargs.get("seq_idx"),
+                **kwargs,
             )
 
             # Drop the additional previous states
@@ -528,7 +513,7 @@ def forward(
 
         recurrent_state = cache_params.layers[self.layer_idx].recurrent_states[0] if use_precomputed_states else None
         if use_precomputed_states and seq_len == 1:
-            core_attn_out, last_recurrent_state = self.recurrent_gated_delta_rule(
+            core_attn_out, last_recurrent_state = torch_recurrent_gated_delta_rule(
                 query,
                 key,
                 value,
@@ -537,9 +522,10 @@ def forward(
                 initial_state=recurrent_state,
                 output_final_state=cache_params is not None,
                 use_qk_l2norm_in_kernel=True,
+                **kwargs,
             )
         else:
-            core_attn_out, last_recurrent_state = self.chunk_gated_delta_rule(
+            core_attn_out, last_recurrent_state = torch_chunk_gated_delta_rule(
                 query,
                 key,
                 value,
@@ -548,8 +534,7 @@ def forward(
                 initial_state=recurrent_state,
                 output_final_state=cache_params is not None,
                 use_qk_l2norm_in_kernel=True,
-                # The chunked FLA kernel takes a single `cu_seqlens` arg; for packed self-attention this matches q-side lengths.
-                cu_seqlens=kwargs.get("cu_seq_lens_q"),
+                **kwargs,
             )
 
         # Update cache
diff --git a/src/transformers/models/qwen3_5/modular_qwen3_5.py b/src/transformers/models/qwen3_5/modular_qwen3_5.py
index b1d854af7a5e..de3eaa20a234 100644
--- a/src/transformers/models/qwen3_5/modular_qwen3_5.py
+++ b/src/transformers/models/qwen3_5/modular_qwen3_5.py
@@ -22,7 +22,7 @@
 
 from ... import initialization as init
 from ...cache_utils import Cache, DynamicCache
-from ...integrations import use_kernel_forward_from_hub
+from ...integrations import use_kernel_forward_from_hub, use_kernelized_func
 from ...masking_utils import create_causal_mask, create_recurrent_attention_mask
 from ...modeling_layers import (
     GenericForSequenceClassification,
@@ -52,6 +52,8 @@
     apply_mask_to_padding_states,
     causal_conv1d_fn,
     causal_conv1d_update,
+    torch_chunk_gated_delta_rule,
+    torch_recurrent_gated_delta_rule,
 )
 from ..qwen3_vl.configuration_qwen3_vl import Qwen3VLConfig, Qwen3VLVisionConfig
 from ..qwen3_vl.modeling_qwen3_vl import (
@@ -206,6 +208,9 @@ def compute_default_rope_parameters(
 
 
 @use_kernel_forward_from_hub("Qwen3_5GatedDeltaNet")
+@use_kernelized_func(
+    [torch_recurrent_gated_delta_rule, torch_chunk_gated_delta_rule, causal_conv1d_fn, causal_conv1d_update]
+)
 class Qwen3_5GatedDeltaNet(Qwen3NextGatedDeltaNet):
     def __init__(self, config: Qwen3_5Config, layer_idx: int):
         super().__init__(config, layer_idx)
@@ -266,7 +271,7 @@ def forward(
                 self.conv1d.weight.squeeze(1),
                 self.conv1d.bias,
                 activation=self.activation,
-                seq_idx=kwargs.get("seq_idx"),
+                **kwargs,
             )
 
             # Drop the additional previous states
@@ -297,7 +302,7 @@ def forward(
 
         recurrent_state = cache_params.layers[self.layer_idx].recurrent_states[0] if use_precomputed_states else None
         if use_precomputed_states and seq_len == 1:
-            core_attn_out, last_recurrent_state = self.recurrent_gated_delta_rule(
+            core_attn_out, last_recurrent_state = torch_recurrent_gated_delta_rule(
                 query,
                 key,
                 value,
@@ -306,9 +311,10 @@ def forward(
                 initial_state=recurrent_state,
                 output_final_state=cache_params is not None,
                 use_qk_l2norm_in_kernel=True,
+                **kwargs,
             )
         else:
-            core_attn_out, last_recurrent_state = self.chunk_gated_delta_rule(
+            core_attn_out, last_recurrent_state = torch_chunk_gated_delta_rule(
                 query,
                 key,
                 value,
@@ -317,8 +323,7 @@ def forward(
                 initial_state=recurrent_state,
                 output_final_state=cache_params is not None,
                 use_qk_l2norm_in_kernel=True,
-                # The chunked FLA kernel takes a single `cu_seqlens` arg; for packed self-attention this matches q-side lengths.
-                cu_seqlens=kwargs.get("cu_seq_lens_q"),
+                **kwargs,
             )
 
         # Update cache
diff --git a/src/transformers/models/qwen3_5_moe/modeling_qwen3_5_moe.py b/src/transformers/models/qwen3_5_moe/modeling_qwen3_5_moe.py
index b2316fdeff52..2899393e7f66 100644
--- a/src/transformers/models/qwen3_5_moe/modeling_qwen3_5_moe.py
+++ b/src/transformers/models/qwen3_5_moe/modeling_qwen3_5_moe.py
@@ -51,7 +51,7 @@
 from ...modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update
 from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
 from ...processing_utils import Unpack
-from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, logging, torch_compilable_check
+from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, torch_compilable_check
 from ...utils.generic import (
     accepts_precomputed_kwargs,
     get_max_seqlen,
@@ -59,7 +59,6 @@
     maybe_autocast,
     merge_with_config_defaults,
 )
-from ...utils.import_utils import is_causal_conv1d_available, is_flash_linear_attention_available
 from ...utils.output_capturing import OutputRecorder, capture_outputs
 from ...vision_utils import (
     get_vision_attention_seqlens,
@@ -70,17 +69,6 @@
 from .configuration_qwen3_5_moe import Qwen3_5MoeConfig, Qwen3_5MoeTextConfig, Qwen3_5MoeVisionConfig
 
 
-if is_flash_linear_attention_available():
-    from fla.modules import FusedRMSNormGated
-    from fla.ops.gated_delta_rule import chunk_gated_delta_rule, fused_recurrent_gated_delta_rule
-else:
-    chunk_gated_delta_rule, fused_recurrent_gated_delta_rule = None, None
-    FusedRMSNormGated = None
-
-
-logger = logging.get_logger(__name__)
-
-
 class Qwen3_5MoeVisionRotaryEmbedding(nn.Module):
     inv_freq: torch.Tensor  # fix linting for `register_buffer`
 
@@ -187,11 +175,13 @@ def apply_interleaved_mrope(self, freqs, mrope_section):
         return freqs_t
 
 
+@use_kernel_forward_from_hub("RMSNormGated")
 class Qwen3_5MoeRMSNormGated(nn.Module):
     def __init__(self, hidden_size, eps=1e-6, **kwargs):
         super().__init__()
         self.weight = nn.Parameter(torch.ones(hidden_size))
         self.variance_epsilon = eps
+        self.activation = "silu"
 
     def forward(self, hidden_states, gate=None):
         input_dtype = hidden_states.dtype
@@ -200,7 +190,7 @@ def forward(self, hidden_states, gate=None):
         # Norm before gate
         hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
         hidden_states = self.weight * hidden_states.to(input_dtype)
-        hidden_states = hidden_states * F.silu(gate.to(torch.float32))
+        hidden_states = hidden_states * ACT2FN[self.activation](gate.to(torch.float32))
 
         return hidden_states.to(input_dtype)
 
@@ -266,6 +256,7 @@ def l2norm(x: torch.FloatTensor, dim: int = -1, eps: float = 1e-6):
     return x * inv_norm
 
 
+@use_kernel_func_from_hub_with_fallback("chunk_gated_delta_rule", "fla", internal_path="ops.gated_delta_rule")
 def torch_chunk_gated_delta_rule(
     query,
     key,
@@ -347,8 +338,17 @@ def torch_chunk_gated_delta_rule(
     return core_attn_out, last_recurrent_state
 
 
+@use_kernel_func_from_hub_with_fallback("recurrent_gated_delta_rule", "fla", internal_path="ops.gated_delta_rule")
 def torch_recurrent_gated_delta_rule(
-    query, key, value, g, beta, initial_state, output_final_state, use_qk_l2norm_in_kernel=False
+    query,
+    key,
+    value,
+    g,
+    beta,
+    initial_state,
+    output_final_state,
+    use_qk_l2norm_in_kernel=False,
+    **kwargs,
 ):
     initial_dtype = query.dtype
     if use_qk_l2norm_in_kernel:
@@ -426,27 +426,9 @@ def __init__(self, config: Qwen3_5MoeConfig, layer_idx: int):
         A = torch.empty(self.num_v_heads).uniform_(0, 16)
         self.A_log = nn.Parameter(torch.log(A))
 
-        self.norm = (
-            Qwen3_5MoeRMSNormGated(self.head_v_dim, eps=self.layer_norm_epsilon)
-            if FusedRMSNormGated is None
-            else FusedRMSNormGated(
-                self.head_v_dim,
-                eps=self.layer_norm_epsilon,
-                activation=self.activation,
-            )
-        )
-
+        self.norm = Qwen3_5MoeRMSNormGated(self.head_v_dim, eps=self.layer_norm_epsilon)
         self.out_proj = nn.Linear(self.value_dim, self.hidden_size, bias=False)
 
-        self.chunk_gated_delta_rule = chunk_gated_delta_rule or torch_chunk_gated_delta_rule
-        self.recurrent_gated_delta_rule = fused_recurrent_gated_delta_rule or torch_recurrent_gated_delta_rule
-
-        if not is_flash_linear_attention_available() or not is_causal_conv1d_available():
-            logger.warning_once(
-                "The fast path is not available because the required library is not installed. Falling back to "
-                "torch implementation. To install follow https://github.com/fla-org/flash-linear-attention#installation"
-            )
-
         self.layer_type = config.layer_types[layer_idx]
 
         self.in_proj_qkv = nn.Linear(self.hidden_size, self.key_dim * 2 + self.value_dim, bias=False)
@@ -498,7 +480,7 @@ def forward(
                 self.conv1d.weight.squeeze(1),
                 self.conv1d.bias,
                 activation=self.activation,
-                seq_idx=kwargs.get("seq_idx"),
+                **kwargs,
             )
 
             # Drop the additional previous states
@@ -529,7 +511,7 @@ def forward(
 
         recurrent_state = cache_params.layers[self.layer_idx].recurrent_states[0] if use_precomputed_states else None
         if use_precomputed_states and seq_len == 1:
-            core_attn_out, last_recurrent_state = self.recurrent_gated_delta_rule(
+            core_attn_out, last_recurrent_state = torch_recurrent_gated_delta_rule(
                 query,
                 key,
                 value,
@@ -538,9 +520,10 @@ def forward(
                 initial_state=recurrent_state,
                 output_final_state=cache_params is not None,
                 use_qk_l2norm_in_kernel=True,
+                **kwargs,
             )
         else:
-            core_attn_out, last_recurrent_state = self.chunk_gated_delta_rule(
+            core_attn_out, last_recurrent_state = torch_chunk_gated_delta_rule(
                 query,
                 key,
                 value,
@@ -549,8 +532,7 @@ def forward(
                 initial_state=recurrent_state,
                 output_final_state=cache_params is not None,
                 use_qk_l2norm_in_kernel=True,
-                # The chunked FLA kernel takes a single `cu_seqlens` arg; for packed self-attention this matches q-side lengths.
-                cu_seqlens=kwargs.get("cu_seq_lens_q"),
+                **kwargs,
             )
 
         # Update cache
diff --git a/src/transformers/models/qwen3_next/modeling_qwen3_next.py b/src/transformers/models/qwen3_next/modeling_qwen3_next.py
index 15eec2537efd..198b5fc36a9c 100644
--- a/src/transformers/models/qwen3_next/modeling_qwen3_next.py
+++ b/src/transformers/models/qwen3_next/modeling_qwen3_next.py
@@ -29,7 +29,12 @@
 from ...activations import ACT2FN
 from ...cache_utils import Cache, DynamicCache
 from ...generation import GenerationMixin
-from ...integrations import use_experts_implementation, use_kernel_func_from_hub_with_fallback, use_kernelized_func
+from ...integrations import (
+    use_experts_implementation,
+    use_kernel_forward_from_hub,
+    use_kernel_func_from_hub_with_fallback,
+    use_kernelized_func,
+)
 from ...integrations.accelerate import force_accelerate_hooks
 from ...masking_utils import create_causal_mask, create_recurrent_attention_mask
 from ...modeling_flash_attention_utils import FlashAttentionKwargs
@@ -43,29 +48,19 @@
 from ...modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update
 from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
 from ...processing_utils import Unpack
-from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, logging
+from ...utils import TransformersKwargs, auto_docstring, can_return_tuple
 from ...utils.generic import maybe_autocast, merge_with_config_defaults
-from ...utils.import_utils import is_causal_conv1d_available, is_flash_linear_attention_available
 from ...utils.output_capturing import OutputRecorder, capture_outputs
 from .configuration_qwen3_next import Qwen3NextConfig
 
 
-if is_flash_linear_attention_available():
-    from fla.modules import FusedRMSNormGated
-    from fla.ops.gated_delta_rule import chunk_gated_delta_rule, fused_recurrent_gated_delta_rule
-else:
-    chunk_gated_delta_rule, fused_recurrent_gated_delta_rule = None, None
-    FusedRMSNormGated = None
-
-
-logger = logging.get_logger(__name__)
-
-
+@use_kernel_forward_from_hub("RMSNormGated")
 class Qwen3NextRMSNormGated(nn.Module):
     def __init__(self, hidden_size, eps=1e-6, **kwargs):
         super().__init__()
         self.weight = nn.Parameter(torch.ones(hidden_size))
         self.variance_epsilon = eps
+        self.activation = "silu"
 
     def forward(self, hidden_states, gate=None):
         input_dtype = hidden_states.dtype
@@ -74,7 +69,7 @@ def forward(self, hidden_states, gate=None):
         # Norm before gate
         hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
         hidden_states = self.weight * hidden_states.to(input_dtype)
-        hidden_states = hidden_states * F.silu(gate.to(torch.float32))
+        hidden_states = hidden_states * ACT2FN[self.activation](gate.to(torch.float32))
 
         return hidden_states.to(input_dtype)
 
@@ -387,6 +382,7 @@ def l2norm(x: torch.FloatTensor, dim: int = -1, eps: float = 1e-6):
     return x * inv_norm
 
 
+@use_kernel_func_from_hub_with_fallback("chunk_gated_delta_rule", "fla", internal_path="ops.gated_delta_rule")
 def torch_chunk_gated_delta_rule(
     query,
     key,
@@ -468,8 +464,17 @@ def torch_chunk_gated_delta_rule(
     return core_attn_out, last_recurrent_state
 
 
+@use_kernel_func_from_hub_with_fallback("recurrent_gated_delta_rule", "fla", internal_path="ops.gated_delta_rule")
 def torch_recurrent_gated_delta_rule(
-    query, key, value, g, beta, initial_state, output_final_state, use_qk_l2norm_in_kernel=False
+    query,
+    key,
+    value,
+    g,
+    beta,
+    initial_state,
+    output_final_state,
+    use_qk_l2norm_in_kernel=False,
+    **kwargs,
 ):
     initial_dtype = query.dtype
     if use_qk_l2norm_in_kernel:
@@ -512,7 +517,9 @@ def torch_recurrent_gated_delta_rule(
     return core_attn_out, last_recurrent_state
 
 
-@use_kernelized_func([causal_conv1d_update, causal_conv1d_fn])
+@use_kernelized_func(
+    [torch_recurrent_gated_delta_rule, torch_chunk_gated_delta_rule, causal_conv1d_fn, causal_conv1d_update]
+)
 class Qwen3NextGatedDeltaNet(nn.Module):
     def __init__(self, config: Qwen3NextConfig, layer_idx: int):
         super().__init__()
@@ -553,27 +560,9 @@ def __init__(self, config: Qwen3NextConfig, layer_idx: int):
         A = torch.empty(self.num_v_heads).uniform_(0, 16)
         self.A_log = nn.Parameter(torch.log(A))
 
-        self.norm = (
-            Qwen3NextRMSNormGated(self.head_v_dim, eps=self.layer_norm_epsilon)
-            if FusedRMSNormGated is None
-            else FusedRMSNormGated(
-                self.head_v_dim,
-                eps=self.layer_norm_epsilon,
-                activation=self.activation,
-            )
-        )
-
+        self.norm = Qwen3NextRMSNormGated(self.head_v_dim, eps=self.layer_norm_epsilon)
         self.out_proj = nn.Linear(self.value_dim, self.hidden_size, bias=False)
 
-        self.chunk_gated_delta_rule = chunk_gated_delta_rule or torch_chunk_gated_delta_rule
-        self.recurrent_gated_delta_rule = fused_recurrent_gated_delta_rule or torch_recurrent_gated_delta_rule
-
-        if not is_flash_linear_attention_available() or not is_causal_conv1d_available():
-            logger.warning_once(
-                "The fast path is not available because the required library is not installed. Falling back to "
-                "torch implementation. To install follow https://github.com/fla-org/flash-linear-attention#installation"
-            )
-
         self.layer_type = config.layer_types[layer_idx]
 
     def fix_query_key_value_ordering(self, mixed_qkvz, mixed_ba):
@@ -648,7 +637,7 @@ def forward(
                 self.conv1d.weight.squeeze(1),
                 self.conv1d.bias,
                 activation=self.activation,
-                seq_idx=kwargs.get("seq_idx"),
+                **kwargs,
             )
 
             # Drop the additional previous states
@@ -678,7 +667,7 @@ def forward(
 
         recurrent_state = cache_params.layers[self.layer_idx].recurrent_states[0] if use_precomputed_states else None
         if use_precomputed_states and seq_len == 1:
-            core_attn_out, last_recurrent_state = self.recurrent_gated_delta_rule(
+            core_attn_out, last_recurrent_state = torch_recurrent_gated_delta_rule(
                 query,
                 key,
                 value,
@@ -687,9 +676,10 @@ def forward(
                 initial_state=recurrent_state,
                 output_final_state=cache_params is not None,
                 use_qk_l2norm_in_kernel=True,
+                **kwargs,
             )
         else:
-            core_attn_out, last_recurrent_state = self.chunk_gated_delta_rule(
+            core_attn_out, last_recurrent_state = torch_chunk_gated_delta_rule(
                 query,
                 key,
                 value,
@@ -698,8 +688,7 @@ def forward(
                 initial_state=recurrent_state,
                 output_final_state=cache_params is not None,
                 use_qk_l2norm_in_kernel=True,
-                # The chunked FLA kernel takes a single `cu_seqlens` arg; for packed self-attention this matches q-side lengths.
-                cu_seqlens=kwargs.get("cu_seq_lens_q"),
+                **kwargs,
             )
 
         # Update cache
diff --git a/src/transformers/models/qwen3_next/modular_qwen3_next.py b/src/transformers/models/qwen3_next/modular_qwen3_next.py
index 9b3a203c2220..cde375b6a4c3 100644
--- a/src/transformers/models/qwen3_next/modular_qwen3_next.py
+++ b/src/transformers/models/qwen3_next/modular_qwen3_next.py
@@ -23,7 +23,7 @@
 from ... import initialization as init
 from ...activations import ACT2FN
 from ...cache_utils import Cache, DynamicCache
-from ...integrations import use_kernel_func_from_hub_with_fallback, use_kernelized_func
+from ...integrations import use_kernel_forward_from_hub, use_kernel_func_from_hub_with_fallback, use_kernelized_func
 from ...integrations.accelerate import force_accelerate_hooks
 from ...masking_utils import create_causal_mask, create_recurrent_attention_mask
 from ...modeling_flash_attention_utils import FlashAttentionKwargs
@@ -32,7 +32,6 @@
 from ...processing_utils import Unpack
 from ...utils import TransformersKwargs, auto_docstring, logging
 from ...utils.generic import merge_with_config_defaults, no_inherit_decorator
-from ...utils.import_utils import is_causal_conv1d_available, is_flash_linear_attention_available
 from ...utils.output_capturing import OutputRecorder, capture_outputs
 from ..bamba.modeling_bamba import apply_mask_to_padding_states, apply_rotary_pos_emb
 from ..gemma2.modeling_gemma2 import Gemma2RotaryEmbedding
@@ -53,22 +52,16 @@
 from .configuration_qwen3_next import Qwen3NextConfig
 
 
-if is_flash_linear_attention_available():
-    from fla.modules import FusedRMSNormGated
-    from fla.ops.gated_delta_rule import chunk_gated_delta_rule, fused_recurrent_gated_delta_rule
-else:
-    chunk_gated_delta_rule, fused_recurrent_gated_delta_rule = None, None
-    FusedRMSNormGated = None
-
-
 logger = logging.get_logger(__name__)
 
 
+@use_kernel_forward_from_hub("RMSNormGated")
 class Qwen3NextRMSNormGated(nn.Module):
     def __init__(self, hidden_size, eps=1e-6, **kwargs):
         super().__init__()
         self.weight = nn.Parameter(torch.ones(hidden_size))
         self.variance_epsilon = eps
+        self.activation = "silu"
 
     def forward(self, hidden_states, gate=None):
         input_dtype = hidden_states.dtype
@@ -77,7 +70,7 @@ def forward(self, hidden_states, gate=None):
         # Norm before gate
         hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
         hidden_states = self.weight * hidden_states.to(input_dtype)
-        hidden_states = hidden_states * F.silu(gate.to(torch.float32))
+        hidden_states = hidden_states * ACT2FN[self.activation](gate.to(torch.float32))
 
         return hidden_states.to(input_dtype)
 
@@ -226,6 +219,7 @@ def l2norm(x: torch.FloatTensor, dim: int = -1, eps: float = 1e-6):
     return x * inv_norm
 
 
+@use_kernel_func_from_hub_with_fallback("chunk_gated_delta_rule", "fla", internal_path="ops.gated_delta_rule")
 def torch_chunk_gated_delta_rule(
     query,
     key,
@@ -307,8 +301,17 @@ def torch_chunk_gated_delta_rule(
     return core_attn_out, last_recurrent_state
 
 
+@use_kernel_func_from_hub_with_fallback("recurrent_gated_delta_rule", "fla", internal_path="ops.gated_delta_rule")
 def torch_recurrent_gated_delta_rule(
-    query, key, value, g, beta, initial_state, output_final_state, use_qk_l2norm_in_kernel=False
+    query,
+    key,
+    value,
+    g,
+    beta,
+    initial_state,
+    output_final_state,
+    use_qk_l2norm_in_kernel=False,
+    **kwargs,
 ):
     initial_dtype = query.dtype
     if use_qk_l2norm_in_kernel:
@@ -351,7 +354,9 @@ def torch_recurrent_gated_delta_rule(
     return core_attn_out, last_recurrent_state
 
 
-@use_kernelized_func([causal_conv1d_update, causal_conv1d_fn])
+@use_kernelized_func(
+    [torch_recurrent_gated_delta_rule, torch_chunk_gated_delta_rule, causal_conv1d_fn, causal_conv1d_update]
+)
 class Qwen3NextGatedDeltaNet(nn.Module):
     def __init__(self, config: Qwen3NextConfig, layer_idx: int):
         super().__init__()
@@ -392,27 +397,9 @@ def __init__(self, config: Qwen3NextConfig, layer_idx: int):
         A = torch.empty(self.num_v_heads).uniform_(0, 16)
         self.A_log = nn.Parameter(torch.log(A))
 
-        self.norm = (
-            Qwen3NextRMSNormGated(self.head_v_dim, eps=self.layer_norm_epsilon)
-            if FusedRMSNormGated is None
-            else FusedRMSNormGated(
-                self.head_v_dim,
-                eps=self.layer_norm_epsilon,
-                activation=self.activation,
-            )
-        )
-
+        self.norm = Qwen3NextRMSNormGated(self.head_v_dim, eps=self.layer_norm_epsilon)
         self.out_proj = nn.Linear(self.value_dim, self.hidden_size, bias=False)
 
-        self.chunk_gated_delta_rule = chunk_gated_delta_rule or torch_chunk_gated_delta_rule
-        self.recurrent_gated_delta_rule = fused_recurrent_gated_delta_rule or torch_recurrent_gated_delta_rule
-
-        if not is_flash_linear_attention_available() or not is_causal_conv1d_available():
-            logger.warning_once(
-                "The fast path is not available because the required library is not installed. Falling back to "
-                "torch implementation. To install follow https://github.com/fla-org/flash-linear-attention#installation"
-            )
-
         self.layer_type = config.layer_types[layer_idx]
 
     def fix_query_key_value_ordering(self, mixed_qkvz, mixed_ba):
@@ -487,7 +474,7 @@ def forward(
                 self.conv1d.weight.squeeze(1),
                 self.conv1d.bias,
                 activation=self.activation,
-                seq_idx=kwargs.get("seq_idx"),
+                **kwargs,
             )
 
             # Drop the additional previous states
@@ -517,7 +504,7 @@ def forward(
 
         recurrent_state = cache_params.layers[self.layer_idx].recurrent_states[0] if use_precomputed_states else None
         if use_precomputed_states and seq_len == 1:
-            core_attn_out, last_recurrent_state = self.recurrent_gated_delta_rule(
+            core_attn_out, last_recurrent_state = torch_recurrent_gated_delta_rule(
                 query,
                 key,
                 value,
@@ -526,9 +513,10 @@ def forward(
                 initial_state=recurrent_state,
                 output_final_state=cache_params is not None,
                 use_qk_l2norm_in_kernel=True,
+                **kwargs,
             )
         else:
-            core_attn_out, last_recurrent_state = self.chunk_gated_delta_rule(
+            core_attn_out, last_recurrent_state = torch_chunk_gated_delta_rule(
                 query,
                 key,
                 value,
@@ -537,8 +525,7 @@ def forward(
                 initial_state=recurrent_state,
                 output_final_state=cache_params is not None,
                 use_qk_l2norm_in_kernel=True,
-                # The chunked FLA kernel takes a single `cu_seqlens` arg; for packed self-attention this matches q-side lengths.
-                cu_seqlens=kwargs.get("cu_seq_lens_q"),
+                **kwargs,
             )
 
         # Update cache

From cf49de645f21909d1dd630afc828f1ee7cba3c42 Mon Sep 17 00:00:00 2001
From: vasqu <antonprogamer@gmail.com>
Date: Wed, 29 Jul 2026 21:45:13 +0000
Subject: [PATCH 12/43] olmo hybrid fused conv style

---
 src/transformers/conversion_mapping.py        |   9 +
 .../olmo_hybrid/modeling_olmo_hybrid.py       | 322 ++++++------------
 .../models/olmo_hybrid/modular_olmo_hybrid.py | 282 ++++-----------
 3 files changed, 176 insertions(+), 437 deletions(-)

diff --git a/src/transformers/conversion_mapping.py b/src/transformers/conversion_mapping.py
index 3ea0571245fb..dd2120050fb1 100755
--- a/src/transformers/conversion_mapping.py
+++ b/src/transformers/conversion_mapping.py
@@ -1017,6 +1017,15 @@ def _build_checkpoint_conversion_mapping():
         "olmo_hybrid": [
             WeightRenaming("attention_layer_norm", "input_layernorm"),
             WeightRenaming("feedforward_layer_norm", "post_attention_layernorm"),
+            WeightConverter(
+                source_patterns=[
+                    "linear_attn.q_conv1d.weight",
+                    "linear_attn.k_conv1d.weight",
+                    "linear_attn.v_conv1d.weight",
+                ],
+                target_patterns="linear_attn.conv1d.weight",
+                operations=[Concatenate(dim=0)],
+            ),
         ],
         "qwen3_5_text": [PrefixChange(prefix_to_remove="language_model", model_prefix="model")],
         "sam3_tracker": [
diff --git a/src/transformers/models/olmo_hybrid/modeling_olmo_hybrid.py b/src/transformers/models/olmo_hybrid/modeling_olmo_hybrid.py
index 0e7d5e64e963..13b96a1a39f3 100644
--- a/src/transformers/models/olmo_hybrid/modeling_olmo_hybrid.py
+++ b/src/transformers/models/olmo_hybrid/modeling_olmo_hybrid.py
@@ -21,7 +21,7 @@
 
 import math
 from collections.abc import Callable
-from typing import Any, Optional
+from typing import Optional
 
 import torch
 import torch.nn as nn
@@ -29,7 +29,7 @@
 
 from ... import initialization as init
 from ...activations import ACT2FN
-from ...cache_utils import Cache
+from ...cache_utils import Cache, DynamicCache
 from ...generation import GenerationMixin
 from ...integrations import use_kernel_forward_from_hub, use_kernel_func_from_hub_with_fallback, use_kernelized_func
 from ...masking_utils import create_causal_mask, create_recurrent_attention_mask
@@ -40,122 +40,10 @@
 from ...processing_utils import Unpack
 from ...utils import TransformersKwargs, auto_docstring, can_return_tuple
 from ...utils.generic import maybe_autocast, merge_with_config_defaults
-from ...utils.import_utils import is_flash_linear_attention_available
 from ...utils.output_capturing import capture_outputs
 from .configuration_olmo_hybrid import OlmoHybridConfig
 
 
-if is_flash_linear_attention_available():
-    from fla.modules import ShortConvolution
-else:
-    ShortConvolution = None
-
-
-class OlmoHybridDynamicCache:
-    """
-    Cache for hybrid model supporting both attention KV cache and linear attention state.
-
-    The main difference is that this cache stores separate conv states for q, k, v (instead of a single conv_states).
-    """
-
-    is_compileable = False
-
-    def __init__(self, config: OlmoHybridConfig):
-        super().__init__()
-        self.layer_types = config.layer_types
-        self.transformer_layers = [
-            i for i in range(config.num_hidden_layers) if self.layer_types[i] == "full_attention"
-        ]
-        self.last_linear_layer = len(self.layer_types) - 1 - self.layer_types[::-1].index("linear_attention")
-        self.recurrent_states = [None for _ in range(config.num_hidden_layers)]
-        self.key_cache = [None for _ in range(config.num_hidden_layers)]
-        self.value_cache = [None for _ in range(config.num_hidden_layers)]
-        # Replace single conv_states with separate q, k, v conv states
-        self.conv_states_q = [None for _ in range(config.num_hidden_layers)]
-        self.conv_states_k = [None for _ in range(config.num_hidden_layers)]
-        self.conv_states_v = [None for _ in range(config.num_hidden_layers)]
-
-    def __len__(self):
-        return len(self.layer_types)
-
-    def update(
-        self,
-        key_states: torch.Tensor,
-        value_states: torch.Tensor,
-        layer_idx: int,
-        cache_kwargs: dict[str, Any] | None = None,
-    ) -> tuple[torch.Tensor, torch.Tensor]:
-        if self.key_cache[layer_idx] is None:
-            self.key_cache[layer_idx] = key_states
-            self.value_cache[layer_idx] = value_states
-        else:
-            self.key_cache[layer_idx] = torch.cat([self.key_cache[layer_idx], key_states], dim=2)
-            self.value_cache[layer_idx] = torch.cat([self.value_cache[layer_idx], value_states], dim=2)
-
-        return self.key_cache[layer_idx], self.value_cache[layer_idx]
-
-    def reorder_cache(self, beam_idx: torch.LongTensor):
-        """Reorders the cache for beam search, given the selected beam indices."""
-        batch_size = beam_idx.shape[0]
-        for layer_idx in range(len(self.key_cache)):
-            if self.key_cache[layer_idx] is not None:
-                if self.key_cache[layer_idx].shape[0] < batch_size:
-                    expand_ratio = batch_size // self.key_cache[layer_idx].shape[0]
-                    self.key_cache[layer_idx] = self.key_cache[layer_idx].repeat_interleave(expand_ratio, dim=0)
-                    self.value_cache[layer_idx] = self.value_cache[layer_idx].repeat_interleave(expand_ratio, dim=0)
-                device = self.key_cache[layer_idx].device
-                self.key_cache[layer_idx] = self.key_cache[layer_idx].index_select(0, beam_idx.to(device))
-                self.value_cache[layer_idx] = self.value_cache[layer_idx].index_select(0, beam_idx.to(device))
-            if self.conv_states_q[layer_idx] is not None:
-                if self.conv_states_q[layer_idx].shape[0] < batch_size:
-                    expand_ratio = batch_size // self.conv_states_q[layer_idx].shape[0]
-                    self.conv_states_q[layer_idx] = self.conv_states_q[layer_idx].repeat_interleave(
-                        expand_ratio, dim=0
-                    )
-                    self.conv_states_k[layer_idx] = self.conv_states_k[layer_idx].repeat_interleave(
-                        expand_ratio, dim=0
-                    )
-                    self.conv_states_v[layer_idx] = self.conv_states_v[layer_idx].repeat_interleave(
-                        expand_ratio, dim=0
-                    )
-                    self.recurrent_states[layer_idx] = self.recurrent_states[layer_idx].repeat_interleave(
-                        expand_ratio, dim=0
-                    )
-                device = self.conv_states_q[layer_idx].device
-                self.conv_states_q[layer_idx] = self.conv_states_q[layer_idx].index_select(0, beam_idx.to(device))
-                self.conv_states_k[layer_idx] = self.conv_states_k[layer_idx].index_select(0, beam_idx.to(device))
-                self.conv_states_v[layer_idx] = self.conv_states_v[layer_idx].index_select(0, beam_idx.to(device))
-                self.recurrent_states[layer_idx] = self.recurrent_states[layer_idx].index_select(
-                    0, beam_idx.to(device)
-                )
-
-    def get_seq_length(self, layer_idx: int | None = 0) -> int:
-        """Returns the sequence length of the cached states. A layer index can be optionally passed."""
-        # take any layer that contains cache and not empty tensor
-        layer_idx = self.transformer_layers[0] if layer_idx not in self.transformer_layers else layer_idx
-        if len(self.key_cache) <= layer_idx or self.key_cache[layer_idx] is None:
-            return 0
-        return self.key_cache[layer_idx].shape[-2]
-
-    def get_mask_sizes(self, query_length: int, layer_idx: int) -> tuple[int, int]:
-        """
-        Return a tuple (kv_length, kv_offset) corresponding to the length and offset that will be returned for
-        the given layer at `layer_idx`.
-        The masks are then prepared according to the given lengths (kv_length, kv_offset) and patterns for each layer.
-        """
-        kv_offset = 0
-        past_seen_tokens = self.get_seq_length(layer_idx)
-        kv_length = query_length + past_seen_tokens
-        return kv_length, kv_offset
-
-    def has_previous_state(self):
-        """We have a previous state if the last linear (conv) layer was already updated."""
-        return self.conv_states_q[self.last_linear_layer] is not None
-
-    def get_query_offset(self, layer_idx: int = 0) -> int:
-        return self.get_seq_length(layer_idx=layer_idx)
-
-
 @use_kernel_forward_from_hub("RMSNormGated")
 class OlmoHybridRMSNormGated(nn.Module):
     def __init__(self, hidden_size, eps=1e-6, **kwargs):
@@ -197,66 +85,6 @@ def extra_repr(self):
         return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}"
 
 
-class OlmoHybridShortConvolution(nn.Conv1d):
-    def __init__(
-        self,
-        hidden_size: int,
-        kernel_size: int,
-        bias: bool = False,
-        activation: str | None = "silu",
-    ):
-        super().__init__(
-            in_channels=hidden_size,
-            out_channels=hidden_size,
-            kernel_size=kernel_size,
-            groups=hidden_size,
-            padding=kernel_size - 1,
-            bias=bias,
-        )
-        self.hidden_size = hidden_size
-        self.conv_kernel_size = kernel_size
-        self.act_fn = ACT2FN[activation]
-
-    def forward(
-        self,
-        hidden_states: torch.Tensor,
-        cache: torch.Tensor | None = None,
-        use_precomputed: bool = False,
-        **kwargs,
-    ) -> tuple[torch.Tensor, torch.Tensor]:
-        seq_len, dim = hidden_states.shape[-2:]
-
-        hidden_states = hidden_states.transpose(1, 2)
-
-        if use_precomputed and seq_len == 1:
-            # Single-token decode: rolling-window update against the cached context.
-            x_with_state = torch.cat([cache, hidden_states], dim=-1)
-            out = F.conv1d(
-                x_with_state,
-                self.weight,
-                self.bias,
-                padding=0,
-                groups=dim,
-            )
-            conv_state = x_with_state[:, :, 1:]
-        else:
-            # Multi-token forward (prefill, or chunked-tokens decode when the cache has prior state).
-            if use_precomputed:
-                # Cached chunked-tokens decode: prepend the cached conv context so the causal conv
-                # sees the correct left-context rather than zero-padding. Dropped from the output
-                # at the end of this branch.
-                hidden_states = torch.cat([cache, hidden_states], dim=-1)
-            out = F.conv1d(hidden_states, self.weight, self.bias, padding=self.conv_kernel_size - 1, groups=dim)
-            out = out[:, :, : hidden_states.shape[-1]]
-            conv_state = F.pad(hidden_states, (self.conv_kernel_size - 1 - hidden_states.shape[-1], 0))
-            if use_precomputed:
-                out = out[:, :, -seq_len:]
-
-        out = self.act_fn(out)
-
-        return out.transpose(1, 2), conv_state
-
-
 def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
     """
     This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,
@@ -493,6 +321,49 @@ def apply_mask_to_padding_states(hidden_states, attention_mask):
     return hidden_states
 
 
+@use_kernel_func_from_hub_with_fallback("causal_conv1d_update", "causal_conv1d")
+def causal_conv1d_update(
+    hidden_states: torch.Tensor,
+    conv_state: torch.Tensor,
+    weight: nn.Parameter,
+    bias: nn.Parameter | None = None,
+    activation: str | None = None,
+):
+    _, hidden_size, seq_len = hidden_states.shape
+    state_len = conv_state.shape[-1]
+
+    hidden_states_new = torch.cat([conv_state, hidden_states], dim=-1).to(weight.dtype)
+    conv_state.copy_(hidden_states_new[:, :, -state_len:])
+    out = F.conv1d(hidden_states_new, weight.unsqueeze(1), bias, padding=0, groups=hidden_size)
+    out = out[:, :, -seq_len:]
+    if activation is not None:
+        out = ACT2FN[activation](out)
+    return out.to(hidden_states.dtype)
+
+
+@use_kernel_func_from_hub_with_fallback("causal_conv1d_fn", "causal_conv1d")
+def causal_conv1d_fn(
+    hidden_states: torch.Tensor,
+    weight: nn.Parameter,
+    bias: nn.Parameter | None = None,
+    activation: str | None = None,
+    **kwargs,
+):
+    _, hidden_size, seq_len = hidden_states.shape
+    padding = weight.shape[-1] - 1
+
+    out = F.conv1d(
+        hidden_states.to(weight.dtype),
+        weight=weight.unsqueeze(1),
+        bias=bias,
+        padding=padding,
+        groups=hidden_size,
+    )[:, :, :seq_len]
+    if activation is not None:
+        out = ACT2FN[activation](out)
+    return out.to(hidden_states.dtype)
+
+
 def l2norm(x: torch.FloatTensor, dim: int = -1, eps: float = 1e-6):
     """This function is intended to align with the l2norm implementation in the FLA library."""
     inv_norm = torch.rsqrt((x * x).sum(dim=dim, keepdim=True) + eps)
@@ -634,13 +505,15 @@ def torch_recurrent_gated_delta_rule(
     return core_attn_out, last_recurrent_state
 
 
+@use_kernelized_func(
+    [torch_recurrent_gated_delta_rule, torch_chunk_gated_delta_rule, causal_conv1d_fn, causal_conv1d_update]
+)
 class OlmoHybridGatedDeltaNet(nn.Module):
     """
     GatedDeltaNet linear attention for OLMo Hybrid.
 
     Key differences from Qwen3NextGatedDeltaNet:
     - Fully separate q/k/v/a/b projections (vs. fused qkvz + partially split ba)
-    - Per-projection conv1d for q, k, v (vs. single conv1d over concatenated qkv)
     - Dedicated g_proj gate (vs. z derived from the fused qkvz projection)
     - Supports allow_neg_eigval: scales beta by 2.0 to allow range [0, 2]
     """
@@ -658,6 +531,7 @@ def __init__(self, config: OlmoHybridConfig, layer_idx: int):
         self.conv_kernel_size = config.linear_conv_kernel_dim
         self.allow_neg_eigval = config.linear_allow_neg_eigval
         self.eps = config.rms_norm_eps
+        self.activation = config.hidden_act
 
         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)
@@ -669,26 +543,14 @@ def __init__(self, config: OlmoHybridConfig, layer_idx: int):
 
         self.o_proj = nn.Linear(self.value_dim, self.hidden_size, bias=False)
 
-        # TODO: can be moved into kernels as well
-        Conv1dClass = ShortConvolution if ShortConvolution is not None else OlmoHybridShortConvolution
-
-        self.q_conv1d = Conv1dClass(
-            hidden_size=self.key_dim,
-            kernel_size=self.conv_kernel_size,
-            bias=False,
-            activation="silu",
-        )
-        self.k_conv1d = Conv1dClass(
-            hidden_size=self.key_dim,
-            kernel_size=self.conv_kernel_size,
+        self.conv_dim = self.key_dim * 2 + self.value_dim
+        self.conv1d = nn.Conv1d(
+            in_channels=self.conv_dim,
+            out_channels=self.conv_dim,
             bias=False,
-            activation="silu",
-        )
-        self.v_conv1d = Conv1dClass(
-            hidden_size=self.value_dim,
             kernel_size=self.conv_kernel_size,
-            bias=False,
-            activation="silu",
+            groups=self.conv_dim,
+            padding=self.conv_kernel_size - 1,
         )
 
         A = torch.empty(self.num_v_heads, dtype=torch.float32).uniform_(
@@ -712,7 +574,7 @@ def __init__(self, config: OlmoHybridConfig, layer_idx: int):
     def forward(
         self,
         hidden_states: torch.Tensor,
-        cache_params: OlmoHybridDynamicCache | None = None,
+        cache_params: Cache | None = None,
         attention_mask: torch.Tensor | None = None,
         **kwargs: Unpack[TransformersKwargs],
     ) -> torch.Tensor:
@@ -725,32 +587,55 @@ def forward(
         # Reads "we have cached conv/recurrent state to continue from". Single-token vs multi-token
         # branching lives inside `ShortConvolution` and in the recurrent-vs-chunk kernel dispatch
         # below, each of which gates on `seq_len == 1` locally.
-        use_precomputed = use_cache and cache_params.has_previous_state()
+        use_precomputed_states = use_cache and cache_params.has_previous_state()
+
+        mixed_qkv = torch.cat(
+            [
+                self.q_proj(hidden_states),
+                self.k_proj(hidden_states),
+                self.v_proj(hidden_states),
+            ],
+            dim=-1,
+        ).transpose(1, 2)
+
+        if use_precomputed_states:
+            conv_state = cache_params.layers[self.layer_idx].conv_states[0]
+            recurrent_state = cache_params.layers[self.layer_idx].recurrent_states[0]
+
+        # Single token decode path
+        if use_precomputed_states and seq_len == 1:
+            mixed_qkv = causal_conv1d_update(
+                mixed_qkv,
+                conv_state,
+                weight=self.conv1d.weight.squeeze(1),
+                bias=self.conv1d.bias,
+                activation=self.activation,
+            )
+        # Multi token prefill or simple "full" prefill
+        else:
+            # Concatenated state for prefill
+            if cache_params is not None:
+                mixed_qkv = cache_params.update_conv_state(
+                    mixed_qkv, self.layer_idx, conv_kernel_size=self.conv_kernel_size
+                )
 
-        conv_state_q = cache_params.conv_states_q[self.layer_idx] if cache_params else None
-        conv_state_k = cache_params.conv_states_k[self.layer_idx] if cache_params else None
-        conv_state_v = cache_params.conv_states_v[self.layer_idx] if cache_params else None
-        recurrent_state = cache_params.recurrent_states[self.layer_idx] if cache_params else None
+            mixed_qkv = causal_conv1d_fn(
+                mixed_qkv,
+                weight=self.conv1d.weight.squeeze(1),
+                bias=self.conv1d.bias,
+                activation=self.activation,
+                **kwargs,
+            )
 
-        q = self.q_proj(hidden_states)
-        k = self.k_proj(hidden_states)
-        v = self.v_proj(hidden_states)
+            # Cut out any tail
+            mixed_qkv = mixed_qkv[:, :, -seq_len:]
 
-        q, new_conv_state_q = self.q_conv1d(
-            q, cache=conv_state_q, use_precomputed=use_precomputed, output_final_state=use_cache
-        )
-        k, new_conv_state_k = self.k_conv1d(
-            k, cache=conv_state_k, use_precomputed=use_precomputed, output_final_state=use_cache
-        )
-        v, new_conv_state_v = self.v_conv1d(
-            v, cache=conv_state_v, use_precomputed=use_precomputed, output_final_state=use_cache
+        q, k, v = torch.split(
+            mixed_qkv.transpose(1, 2),
+            [self.key_dim, self.key_dim, self.value_dim],
+            dim=-1,
         )
 
-        if cache_params is not None:
-            cache_params.conv_states_q[self.layer_idx] = new_conv_state_q
-            cache_params.conv_states_k[self.layer_idx] = new_conv_state_k
-            cache_params.conv_states_v[self.layer_idx] = new_conv_state_v
-
         q = q.view(batch_size, seq_len, -1, self.head_k_dim)
         k = k.view(batch_size, seq_len, -1, self.head_k_dim)
         v = v.view(batch_size, seq_len, -1, self.head_v_dim)
@@ -766,7 +651,7 @@ def forward(
 
         g = -self.A_log.float().exp() * F.softplus(self.a_proj(hidden_states).float() + self.dt_bias)
 
-        if use_precomputed and seq_len == 1:
+        if use_precomputed_states and seq_len == 1:
             output, new_recurrent_state = torch_recurrent_gated_delta_rule(
                 q,
                 k,
@@ -785,7 +670,7 @@ def forward(
                 v,
                 g=g,
                 beta=beta,
-                initial_state=recurrent_state if use_precomputed else None,
+                initial_state=recurrent_state if use_precomputed_states else None,
                 output_final_state=use_cache,
                 use_qk_l2norm_in_kernel=True,
                 **kwargs,
@@ -889,6 +774,7 @@ def forward(
             hidden_states=hidden_states,
             cache_params=past_key_values,
             attention_mask=attention_mask,
+            **kwargs,
         )
         hidden_states = residual + hidden_states
 
@@ -978,7 +864,7 @@ def forward(
             inputs_embeds = self.embed_tokens(input_ids)
 
         if use_cache and past_key_values is None:
-            past_key_values = OlmoHybridDynamicCache(config=self.config)
+            past_key_values = DynamicCache(config=self.config)
 
         if position_ids is None:
             past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
diff --git a/src/transformers/models/olmo_hybrid/modular_olmo_hybrid.py b/src/transformers/models/olmo_hybrid/modular_olmo_hybrid.py
index 656b4251f01a..809a3de51517 100644
--- a/src/transformers/models/olmo_hybrid/modular_olmo_hybrid.py
+++ b/src/transformers/models/olmo_hybrid/modular_olmo_hybrid.py
@@ -16,7 +16,6 @@
 
 import math
 from collections.abc import Callable
-from typing import Any
 
 import torch
 import torch.nn as nn
@@ -24,16 +23,15 @@
 from huggingface_hub.dataclasses import strict
 
 from ... import initialization as init
-from ...activations import ACT2FN
-from ...cache_utils import Cache
+from ...cache_utils import Cache, DynamicCache
 from ...configuration_utils import PreTrainedConfig, remap_legacy_layer_types
+from ...integrations import use_kernelized_func
 from ...masking_utils import create_causal_mask, create_recurrent_attention_mask
 from ...modeling_outputs import BaseModelOutputWithPast
 from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
 from ...processing_utils import Unpack
 from ...utils import TransformersKwargs, auto_docstring, logging
 from ...utils.generic import maybe_autocast, merge_with_config_defaults
-from ...utils.import_utils import is_flash_linear_attention_available
 from ...utils.output_capturing import capture_outputs
 from ..llama.configuration_llama import LlamaConfig
 from ..llama.modeling_llama import LlamaDecoderLayer
@@ -52,17 +50,13 @@
     Qwen3NextPreTrainedModel,
     Qwen3NextRMSNormGated,
     apply_mask_to_padding_states,
+    causal_conv1d_fn,
+    causal_conv1d_update,
     torch_chunk_gated_delta_rule,
     torch_recurrent_gated_delta_rule,
 )
 
 
-if is_flash_linear_attention_available():
-    from fla.modules import ShortConvolution
-else:
-    ShortConvolution = None
-
-
 logger = logging.get_logger(__name__)
 
 
@@ -183,111 +177,6 @@ def validate_architecture(self):
             raise ValueError("OLMoHybrid expects at least one attention layer.")
 
 
-class OlmoHybridDynamicCache:
-    """
-    Cache for hybrid model supporting both attention KV cache and linear attention state.
-
-    The main difference is that this cache stores separate conv states for q, k, v (instead of a single conv_states).
-    """
-
-    is_compileable = False
-
-    def __init__(self, config: OlmoHybridConfig):
-        super().__init__()
-        self.layer_types = config.layer_types
-        self.transformer_layers = [
-            i for i in range(config.num_hidden_layers) if self.layer_types[i] == "full_attention"
-        ]
-        self.last_linear_layer = len(self.layer_types) - 1 - self.layer_types[::-1].index("linear_attention")
-        self.recurrent_states = [None for _ in range(config.num_hidden_layers)]
-        self.key_cache = [None for _ in range(config.num_hidden_layers)]
-        self.value_cache = [None for _ in range(config.num_hidden_layers)]
-        # Replace single conv_states with separate q, k, v conv states
-        self.conv_states_q = [None for _ in range(config.num_hidden_layers)]
-        self.conv_states_k = [None for _ in range(config.num_hidden_layers)]
-        self.conv_states_v = [None for _ in range(config.num_hidden_layers)]
-
-    def __len__(self):
-        return len(self.layer_types)
-
-    def update(
-        self,
-        key_states: torch.Tensor,
-        value_states: torch.Tensor,
-        layer_idx: int,
-        cache_kwargs: dict[str, Any] | None = None,
-    ) -> tuple[torch.Tensor, torch.Tensor]:
-        if self.key_cache[layer_idx] is None:
-            self.key_cache[layer_idx] = key_states
-            self.value_cache[layer_idx] = value_states
-        else:
-            self.key_cache[layer_idx] = torch.cat([self.key_cache[layer_idx], key_states], dim=2)
-            self.value_cache[layer_idx] = torch.cat([self.value_cache[layer_idx], value_states], dim=2)
-
-        return self.key_cache[layer_idx], self.value_cache[layer_idx]
-
-    def reorder_cache(self, beam_idx: torch.LongTensor):
-        """Reorders the cache for beam search, given the selected beam indices."""
-        batch_size = beam_idx.shape[0]
-        for layer_idx in range(len(self.key_cache)):
-            if self.key_cache[layer_idx] is not None:
-                if self.key_cache[layer_idx].shape[0] < batch_size:
-                    expand_ratio = batch_size // self.key_cache[layer_idx].shape[0]
-                    self.key_cache[layer_idx] = self.key_cache[layer_idx].repeat_interleave(expand_ratio, dim=0)
-                    self.value_cache[layer_idx] = self.value_cache[layer_idx].repeat_interleave(expand_ratio, dim=0)
-                device = self.key_cache[layer_idx].device
-                self.key_cache[layer_idx] = self.key_cache[layer_idx].index_select(0, beam_idx.to(device))
-                self.value_cache[layer_idx] = self.value_cache[layer_idx].index_select(0, beam_idx.to(device))
-            if self.conv_states_q[layer_idx] is not None:
-                if self.conv_states_q[layer_idx].shape[0] < batch_size:
-                    expand_ratio = batch_size // self.conv_states_q[layer_idx].shape[0]
-                    self.conv_states_q[layer_idx] = self.conv_states_q[layer_idx].repeat_interleave(
-                        expand_ratio, dim=0
-                    )
-                    self.conv_states_k[layer_idx] = self.conv_states_k[layer_idx].repeat_interleave(
-                        expand_ratio, dim=0
-                    )
-                    self.conv_states_v[layer_idx] = self.conv_states_v[layer_idx].repeat_interleave(
-                        expand_ratio, dim=0
-                    )
-                    self.recurrent_states[layer_idx] = self.recurrent_states[layer_idx].repeat_interleave(
-                        expand_ratio, dim=0
-                    )
-                device = self.conv_states_q[layer_idx].device
-                self.conv_states_q[layer_idx] = self.conv_states_q[layer_idx].index_select(0, beam_idx.to(device))
-                self.conv_states_k[layer_idx] = self.conv_states_k[layer_idx].index_select(0, beam_idx.to(device))
-                self.conv_states_v[layer_idx] = self.conv_states_v[layer_idx].index_select(0, beam_idx.to(device))
-                self.recurrent_states[layer_idx] = self.recurrent_states[layer_idx].index_select(
-                    0, beam_idx.to(device)
-                )
-
-    def get_seq_length(self, layer_idx: int | None = 0) -> int:
-        """Returns the sequence length of the cached states. A layer index can be optionally passed."""
-        # take any layer that contains cache and not empty tensor
-        layer_idx = self.transformer_layers[0] if layer_idx not in self.transformer_layers else layer_idx
-        if len(self.key_cache) <= layer_idx or self.key_cache[layer_idx] is None:
-            return 0
-        return self.key_cache[layer_idx].shape[-2]
-
-    def get_mask_sizes(self, query_length: int, layer_idx: int) -> tuple[int, int]:
-        """
-        Return a tuple (kv_length, kv_offset) corresponding to the length and offset that will be returned for
-        the given layer at `layer_idx`.
-        The masks are then prepared according to the given lengths (kv_length, kv_offset) and patterns for each layer.
-        """
-        kv_offset = 0
-        past_seen_tokens = self.get_seq_length(layer_idx)
-        kv_length = query_length + past_seen_tokens
-        return kv_length, kv_offset
-
-    def has_previous_state(self):
-        """We have a previous state if the last linear (conv) layer was already updated."""
-        return self.conv_states_q[self.last_linear_layer] is not None
-
-    def get_query_offset(self, layer_idx: int = 0) -> int:
-        return self.get_seq_length(layer_idx=layer_idx)
-
-
 class OlmoHybridRMSNormGated(Qwen3NextRMSNormGated):
     pass
 
@@ -296,66 +185,6 @@ class OlmoHybridRMSNorm(Olmo3RMSNorm):
     pass
 
 
-class OlmoHybridShortConvolution(nn.Conv1d):
-    def __init__(
-        self,
-        hidden_size: int,
-        kernel_size: int,
-        bias: bool = False,
-        activation: str | None = "silu",
-    ):
-        super().__init__(
-            in_channels=hidden_size,
-            out_channels=hidden_size,
-            kernel_size=kernel_size,
-            groups=hidden_size,
-            padding=kernel_size - 1,
-            bias=bias,
-        )
-        self.hidden_size = hidden_size
-        self.conv_kernel_size = kernel_size
-        self.act_fn = ACT2FN[activation]
-
-    def forward(
-        self,
-        hidden_states: torch.Tensor,
-        cache: torch.Tensor | None = None,
-        use_precomputed: bool = False,
-        **kwargs,
-    ) -> tuple[torch.Tensor, torch.Tensor]:
-        seq_len, dim = hidden_states.shape[-2:]
-
-        hidden_states = hidden_states.transpose(1, 2)
-
-        if use_precomputed and seq_len == 1:
-            # Single-token decode: rolling-window update against the cached context.
-            x_with_state = torch.cat([cache, hidden_states], dim=-1)
-            out = F.conv1d(
-                x_with_state,
-                self.weight,
-                self.bias,
-                padding=0,
-                groups=dim,
-            )
-            conv_state = x_with_state[:, :, 1:]
-        else:
-            # Multi-token forward (prefill, or chunked-tokens decode when the cache has prior state).
-            if use_precomputed:
-                # Cached chunked-tokens decode: prepend the cached conv context so the causal conv
-                # sees the correct left-context rather than zero-padding. Dropped from the output
-                # at the end of this branch.
-                hidden_states = torch.cat([cache, hidden_states], dim=-1)
-            out = F.conv1d(hidden_states, self.weight, self.bias, padding=self.conv_kernel_size - 1, groups=dim)
-            out = out[:, :, : hidden_states.shape[-1]]
-            conv_state = F.pad(hidden_states, (self.conv_kernel_size - 1 - hidden_states.shape[-1], 0))
-            if use_precomputed:
-                out = out[:, :, -seq_len:]
-
-        out = self.act_fn(out)
-
-        return out.transpose(1, 2), conv_state
-
-
 class OlmoHybridAttention(Olmo3Attention):
     """
     Multi-headed attention for OLMo Hybrid that supports optional RoPE (NoPE mode).
@@ -439,13 +268,15 @@ def forward(self, x, position_ids):
         return cos, sin
 
 
+@use_kernelized_func(
+    [torch_recurrent_gated_delta_rule, torch_chunk_gated_delta_rule, causal_conv1d_fn, causal_conv1d_update]
+)
 class OlmoHybridGatedDeltaNet(nn.Module):
     """
     GatedDeltaNet linear attention for OLMo Hybrid.
 
     Key differences from Qwen3NextGatedDeltaNet:
     - Fully separate q/k/v/a/b projections (vs. fused qkvz + partially split ba)
-    - Per-projection conv1d for q, k, v (vs. single conv1d over concatenated qkv)
     - Dedicated g_proj gate (vs. z derived from the fused qkvz projection)
     - Supports allow_neg_eigval: scales beta by 2.0 to allow range [0, 2]
     """
@@ -463,6 +294,7 @@ def __init__(self, config: OlmoHybridConfig, layer_idx: int):
         self.conv_kernel_size = config.linear_conv_kernel_dim
         self.allow_neg_eigval = config.linear_allow_neg_eigval
         self.eps = config.rms_norm_eps
+        self.activation = config.hidden_act
 
         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)
@@ -474,26 +306,14 @@ def __init__(self, config: OlmoHybridConfig, layer_idx: int):
 
         self.o_proj = nn.Linear(self.value_dim, self.hidden_size, bias=False)
 
-        # TODO: can be moved into kernels as well
-        Conv1dClass = ShortConvolution if ShortConvolution is not None else OlmoHybridShortConvolution
-
-        self.q_conv1d = Conv1dClass(
-            hidden_size=self.key_dim,
-            kernel_size=self.conv_kernel_size,
+        self.conv_dim = self.key_dim * 2 + self.value_dim
+        self.conv1d = nn.Conv1d(
+            in_channels=self.conv_dim,
+            out_channels=self.conv_dim,
             bias=False,
-            activation="silu",
-        )
-        self.k_conv1d = Conv1dClass(
-            hidden_size=self.key_dim,
-            kernel_size=self.conv_kernel_size,
-            bias=False,
-            activation="silu",
-        )
-        self.v_conv1d = Conv1dClass(
-            hidden_size=self.value_dim,
             kernel_size=self.conv_kernel_size,
-            bias=False,
-            activation="silu",
+            groups=self.conv_dim,
+            padding=self.conv_kernel_size - 1,
         )
 
         A = torch.empty(self.num_v_heads, dtype=torch.float32).uniform_(
@@ -517,7 +337,7 @@ def __init__(self, config: OlmoHybridConfig, layer_idx: int):
     def forward(
         self,
         hidden_states: torch.Tensor,
-        cache_params: OlmoHybridDynamicCache | None = None,
+        cache_params: Cache | None = None,
         attention_mask: torch.Tensor | None = None,
         **kwargs: Unpack[TransformersKwargs],
     ) -> torch.Tensor:
@@ -530,31 +350,54 @@ def forward(
         # Reads "we have cached conv/recurrent state to continue from". Single-token vs multi-token
         # branching lives inside `ShortConvolution` and in the recurrent-vs-chunk kernel dispatch
         # below, each of which gates on `seq_len == 1` locally.
-        use_precomputed = use_cache and cache_params.has_previous_state()
+        use_precomputed_states = use_cache and cache_params.has_previous_state()
 
-        conv_state_q = cache_params.conv_states_q[self.layer_idx] if cache_params else None
-        conv_state_k = cache_params.conv_states_k[self.layer_idx] if cache_params else None
-        conv_state_v = cache_params.conv_states_v[self.layer_idx] if cache_params else None
-        recurrent_state = cache_params.recurrent_states[self.layer_idx] if cache_params else None
+        mixed_qkv = torch.cat(
+            [
+                self.q_proj(hidden_states),
+                self.k_proj(hidden_states),
+                self.v_proj(hidden_states),
+            ],
+            dim=-1,
+        ).transpose(1, 2)
+
+        if use_precomputed_states:
+            conv_state = cache_params.layers[self.layer_idx].conv_states[0]
+            recurrent_state = cache_params.layers[self.layer_idx].recurrent_states[0]
+
+        # Single token decode path
+        if use_precomputed_states and seq_len == 1:
+            mixed_qkv = causal_conv1d_update(
+                mixed_qkv,
+                conv_state,
+                weight=self.conv1d.weight.squeeze(1),
+                bias=self.conv1d.bias,
+                activation=self.activation,
+            )
+        # Multi token prefill or simple "full" prefill
+        else:
+            # Concatenated state for prefill
+            if cache_params is not None:
+                mixed_qkv = cache_params.update_conv_state(
+                    mixed_qkv, self.layer_idx, conv_kernel_size=self.conv_kernel_size
+                )
 
-        q = self.q_proj(hidden_states)
-        k = self.k_proj(hidden_states)
-        v = self.v_proj(hidden_states)
+            mixed_qkv = causal_conv1d_fn(
+                mixed_qkv,
+                weight=self.conv1d.weight.squeeze(1),
+                bias=self.conv1d.bias,
+                activation=self.activation,
+                **kwargs,
+            )
 
-        q, new_conv_state_q = self.q_conv1d(
-            q, cache=conv_state_q, use_precomputed=use_precomputed, output_final_state=use_cache
-        )
-        k, new_conv_state_k = self.k_conv1d(
-            k, cache=conv_state_k, use_precomputed=use_precomputed, output_final_state=use_cache
-        )
-        v, new_conv_state_v = self.v_conv1d(
-            v, cache=conv_state_v, use_precomputed=use_precomputed, output_final_state=use_cache
-        )
+            # Cut out any tail
+            mixed_qkv = mixed_qkv[:, :, -seq_len:]
 
-        if cache_params is not None:
-            cache_params.conv_states_q[self.layer_idx] = new_conv_state_q
-            cache_params.conv_states_k[self.layer_idx] = new_conv_state_k
-            cache_params.conv_states_v[self.layer_idx] = new_conv_state_v
+        q, k, v = torch.split(
+            mixed_qkv.transpose(1, 2),
+            [self.key_dim, self.key_dim, self.value_dim],
+            dim=-1,
+        )
 
         q = q.view(batch_size, seq_len, -1, self.head_k_dim)
         k = k.view(batch_size, seq_len, -1, self.head_k_dim)
@@ -571,7 +414,7 @@ def forward(
 
         g = -self.A_log.float().exp() * F.softplus(self.a_proj(hidden_states).float() + self.dt_bias)
 
-        if use_precomputed and seq_len == 1:
+        if use_precomputed_states and seq_len == 1:
             output, new_recurrent_state = torch_recurrent_gated_delta_rule(
                 q,
                 k,
@@ -590,7 +433,7 @@ def forward(
                 v,
                 g=g,
                 beta=beta,
-                initial_state=recurrent_state if use_precomputed else None,
+                initial_state=recurrent_state if use_precomputed_states else None,
                 output_final_state=use_cache,
                 use_qk_l2norm_in_kernel=True,
                 **kwargs,
@@ -647,6 +490,7 @@ def forward(
             hidden_states=hidden_states,
             cache_params=past_key_values,
             attention_mask=attention_mask,
+            **kwargs,
         )
         hidden_states = residual + hidden_states
 
@@ -725,7 +569,7 @@ def forward(
             inputs_embeds = self.embed_tokens(input_ids)
 
         if use_cache and past_key_values is None:
-            past_key_values = OlmoHybridDynamicCache(config=self.config)
+            past_key_values = DynamicCache(config=self.config)
 
         if position_ids is None:
             past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0

From 73db89ba8223c70eac8d5762e4a745b12ee7d778 Mon Sep 17 00:00:00 2001
From: vasqu <antonprogamer@gmail.com>
Date: Wed, 29 Jul 2026 21:53:16 +0000
Subject: [PATCH 13/43] fix

---
 tests/models/olmo_hybrid/test_modeling_olmo_hybrid.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/tests/models/olmo_hybrid/test_modeling_olmo_hybrid.py b/tests/models/olmo_hybrid/test_modeling_olmo_hybrid.py
index 56622d65407b..78e495cd36fc 100644
--- a/tests/models/olmo_hybrid/test_modeling_olmo_hybrid.py
+++ b/tests/models/olmo_hybrid/test_modeling_olmo_hybrid.py
@@ -34,11 +34,11 @@
 
     from transformers import (
         Cache,
+        DynamicCache,
         OlmoHybridForCausalLM,
         OlmoHybridModel,
     )
     from transformers.models.olmo_hybrid.modeling_olmo_hybrid import (
-        OlmoHybridDynamicCache,
         OlmoHybridRotaryEmbedding,
     )
 
@@ -87,7 +87,7 @@ def test_linear_attention_multi_token_cached_forward_matches_single_token(self):
         prompt = ids_tensor((1, prefill_len), config.vocab_size).to(torch_device)
         next_token = ids_tensor((1, 1), config.vocab_size).to(torch_device)
 
-        cache_single = OlmoHybridDynamicCache(config=config)
+        cache_single = DynamicCache(config=config)
         with torch.no_grad():
             model(input_ids=prompt, past_key_values=cache_single, use_cache=True)
             single_out = model(input_ids=next_token, past_key_values=cache_single, use_cache=True)
@@ -95,7 +95,7 @@ def test_linear_attention_multi_token_cached_forward_matches_single_token(self):
 
         distractors = ids_tensor((1, 7), config.vocab_size).to(torch_device)
         multi_input = torch.cat([next_token, distractors], dim=1)
-        cache_multi = OlmoHybridDynamicCache(config=config)
+        cache_multi = DynamicCache(config=config)
         with torch.no_grad():
             model(input_ids=prompt, past_key_values=cache_multi, use_cache=True)
             multi_out = model(input_ids=multi_input, past_key_values=cache_multi, use_cache=True)
@@ -106,7 +106,7 @@ def test_linear_attention_multi_token_cached_forward_matches_single_token(self):
     # === Cache helper methods (same pattern as Qwen3Next) ===
     def _check_past_key_values_for_generate(self, batch_size, past_key_values, seq_length, config):
         """OlmoHybrid has a special Cache as it alternates with gated deltanet layers"""
-        self.assertIsInstance(past_key_values, OlmoHybridDynamicCache)
+        self.assertIsInstance(past_key_values, DynamicCache)
 
         num_heads = getattr(config, "num_key_value_heads", config.num_attention_heads)
         head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)

From 8896696789213a46c29549b8c1cbdfb1f796e170 Mon Sep 17 00:00:00 2001
From: vasqu <antonprogamer@gmail.com>
Date: Wed, 29 Jul 2026 22:14:28 +0000
Subject: [PATCH 14/43] fix

---
 .../models/olmo_hybrid/modeling_olmo_hybrid.py             | 7 ++++---
 src/transformers/models/olmo_hybrid/modular_olmo_hybrid.py | 7 ++++---
 tests/models/olmo_hybrid/test_modeling_olmo_hybrid.py      | 1 +
 3 files changed, 9 insertions(+), 6 deletions(-)

diff --git a/src/transformers/models/olmo_hybrid/modeling_olmo_hybrid.py b/src/transformers/models/olmo_hybrid/modeling_olmo_hybrid.py
index 13b96a1a39f3..090f6cb3727c 100644
--- a/src/transformers/models/olmo_hybrid/modeling_olmo_hybrid.py
+++ b/src/transformers/models/olmo_hybrid/modeling_olmo_hybrid.py
@@ -652,7 +652,7 @@ def forward(
         g = -self.A_log.float().exp() * F.softplus(self.a_proj(hidden_states).float() + self.dt_bias)
 
         if use_precomputed_states and seq_len == 1:
-            output, new_recurrent_state = torch_recurrent_gated_delta_rule(
+            output, last_recurrent_state = torch_recurrent_gated_delta_rule(
                 q,
                 k,
                 v,
@@ -664,7 +664,7 @@ def forward(
                 **kwargs,
             )
         else:
-            output, new_recurrent_state = torch_chunk_gated_delta_rule(
+            output, last_recurrent_state = torch_chunk_gated_delta_rule(
                 q,
                 k,
                 v,
@@ -676,8 +676,9 @@ def forward(
                 **kwargs,
             )
 
+        # Update cache
         if cache_params is not None:
-            cache_params.recurrent_states[self.layer_idx] = new_recurrent_state
+            cache_params.update_recurrent_state(last_recurrent_state, self.layer_idx)
 
         gate = self.g_proj(hidden_states)
         output = output.reshape(-1, self.head_v_dim)
diff --git a/src/transformers/models/olmo_hybrid/modular_olmo_hybrid.py b/src/transformers/models/olmo_hybrid/modular_olmo_hybrid.py
index 809a3de51517..d0fab6a6109c 100644
--- a/src/transformers/models/olmo_hybrid/modular_olmo_hybrid.py
+++ b/src/transformers/models/olmo_hybrid/modular_olmo_hybrid.py
@@ -415,7 +415,7 @@ def forward(
         g = -self.A_log.float().exp() * F.softplus(self.a_proj(hidden_states).float() + self.dt_bias)
 
         if use_precomputed_states and seq_len == 1:
-            output, new_recurrent_state = torch_recurrent_gated_delta_rule(
+            output, last_recurrent_state = torch_recurrent_gated_delta_rule(
                 q,
                 k,
                 v,
@@ -427,7 +427,7 @@ def forward(
                 **kwargs,
             )
         else:
-            output, new_recurrent_state = torch_chunk_gated_delta_rule(
+            output, last_recurrent_state = torch_chunk_gated_delta_rule(
                 q,
                 k,
                 v,
@@ -439,8 +439,9 @@ def forward(
                 **kwargs,
             )
 
+        # Update cache
         if cache_params is not None:
-            cache_params.recurrent_states[self.layer_idx] = new_recurrent_state
+            cache_params.update_recurrent_state(last_recurrent_state, self.layer_idx)
 
         gate = self.g_proj(hidden_states)
         output = output.reshape(-1, self.head_v_dim)
diff --git a/tests/models/olmo_hybrid/test_modeling_olmo_hybrid.py b/tests/models/olmo_hybrid/test_modeling_olmo_hybrid.py
index 78e495cd36fc..80c123c5e1b5 100644
--- a/tests/models/olmo_hybrid/test_modeling_olmo_hybrid.py
+++ b/tests/models/olmo_hybrid/test_modeling_olmo_hybrid.py
@@ -58,6 +58,7 @@ def __init__(self, parent):
         self.linear_value_head_dim = 8
         self.linear_conv_kernel_dim = 4
         self.linear_allow_neg_eigval = False
+        self.hidden_act = "silu"
 
 
 @require_torch

From d3b27a664556236a2f4b4cee3cda11424cf39a38 Mon Sep 17 00:00:00 2001
From: vasqu <antonprogamer@gmail.com>
Date: Wed, 29 Jul 2026 23:27:29 +0000
Subject: [PATCH 15/43] poc mamba2, kernel must compile but torch seems to
 match

---
 .../models/mamba2/modeling_mamba2.py          | 436 +++++++++++++-----
 1 file changed, 309 insertions(+), 127 deletions(-)

diff --git a/src/transformers/models/mamba2/modeling_mamba2.py b/src/transformers/models/mamba2/modeling_mamba2.py
index 9ecae3a54415..444319bba8b6 100644
--- a/src/transformers/models/mamba2/modeling_mamba2.py
+++ b/src/transformers/models/mamba2/modeling_mamba2.py
@@ -24,7 +24,7 @@
 from ...activations import ACT2FN
 from ...cache_utils import Cache, DynamicCache
 from ...generation import GenerationMixin
-from ...integrations import lazy_load_kernel
+from ...integrations import lazy_load_kernel, use_kernel_func_from_hub_with_fallback
 from ...integrations.accelerate import force_accelerate_hooks
 from ...masking_utils import create_recurrent_attention_mask
 from ...modeling_layers import GradientCheckpointingLayer
@@ -162,6 +162,204 @@ def causal_conv1d_fn(
     return out.to(hidden_states.dtype)
 
 
+@use_kernel_func_from_hub_with_fallback(
+    "mamba_split_conv1d_scan_combined",
+    "mamba_ssm",
+    internal_path="ops.triton.ssd_combined",
+)
+def mamba2_split_conv1d_scan_combined(
+    zxbcdt: torch.Tensor,
+    conv1d_weight: torch.Tensor,
+    conv1d_bias: torch.Tensor | None,
+    dt_bias: torch.Tensor,
+    A: torch.Tensor,
+    D: torch.Tensor,
+    chunk_size: int,
+    initial_states: 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,
+    **kwargs,
+):
+    return None
+
+
+@use_kernel_func_from_hub_with_fallback(
+    "selective_state_update",
+    "mamba_ssm",
+    internal_path="ops.triton.selective_state_update",
+)
+def mamba2_selective_state_update(
+    state: torch.Tensor,
+    hidden_states: torch.Tensor,
+    dt: torch.Tensor,
+    A: torch.Tensor,
+    B: torch.Tensor,
+    C: torch.Tensor,
+    D: torch.Tensor | None = None,
+    dt_bias: torch.Tensor | None = None,
+    dt_softplus: bool = False,
+    **kwargs,
+):
+    batch_size, num_heads, head_dim = hidden_states.shape[:2]
+    num_groups = B.shape[1]
+    state_size = B.shape[-1]
+
+    if dt_bias is not None:
+        dt = dt + dt_bias.to(dt.dtype)
+    if dt_softplus:
+        dt = F.softplus(dt)
+    dt = dt[..., None]
+
+    # Discretize A
+    dA = torch.exp(dt.float() * A.float()).to(device=state.device)
+
+    # Discretize B
+    B = B.reshape(batch_size, num_groups, 1, state_size)
+    B = B.expand(batch_size, num_groups, num_heads // num_groups, state_size).contiguous()
+    B = B.reshape(batch_size, num_heads, 1, state_size)
+    dB = dt * B
+
+    # Discretize x into dB
+    dBx = (dB * hidden_states[..., None]).to(device=state.device)
+
+    # State calculation
+    ssm_states = state * dA + dBx
+    state.copy_(ssm_states.to(state.dtype))
+
+    # Subsequent output
+    C = C.reshape(batch_size, num_groups, 1, state_size)
+    C = C.expand(batch_size, num_groups, num_heads // num_groups, state_size).contiguous()
+    C = C.reshape(batch_size, num_heads, state_size)
+
+    # Reshape ssm_states to merge the first two dimensions
+    ssm_states = ssm_states.to(device=C.device, dtype=C.dtype)
+    ssm_states_reshaped = ssm_states.view(batch_size * num_heads, head_dim, state_size)
+    C_reshaped = C.view(batch_size * num_heads, state_size, 1)
+    out = torch.bmm(ssm_states_reshaped, C_reshaped)
+    out = out.view(batch_size, num_heads, head_dim)
+
+    # D skip connection
+    if D is not None:
+        out = (out + hidden_states * D).to(out.dtype)
+
+    return out.to(hidden_states.dtype)
+
+
+@use_kernel_func_from_hub_with_fallback(
+    "mamba_chunk_scan_combined",
+    "mamba_ssm",
+    internal_path="ops.triton.ssd_combined",
+)
+def mamba2_chunk_scan(
+    hidden_states: torch.Tensor,
+    dt: torch.Tensor,
+    A: torch.Tensor,
+    B: torch.Tensor,
+    C: torch.Tensor,
+    chunk_size: int,
+    D: torch.Tensor | None = None,
+    dt_bias: torch.Tensor | None = None,
+    initial_states: torch.Tensor | None = None,
+    dt_softplus: bool = False,
+    dt_limit: tuple[float, float] = (0.0, float("inf")),
+    return_final_states: bool = False,
+    **kwargs,
+):
+    input_dtype = hidden_states.dtype
+    batch_size, sequence_length, num_heads, head_dim = hidden_states.shape
+    num_groups = B.shape[2]
+
+    if dt_bias is not None:
+        dt = dt + dt_bias.to(dt.dtype)
+    if dt_softplus:
+        dt = F.softplus(dt)
+    dt = torch.clamp(dt, min=dt_limit[0], max=dt_limit[1])
+
+    hidden_states = hidden_states.float()
+    B = B.float().repeat_interleave(num_heads // num_groups, dim=2, output_size=num_heads)
+    C = C.float().repeat_interleave(num_heads // num_groups, dim=2, output_size=num_heads)
+
+    pad_size = (chunk_size - sequence_length % chunk_size) % chunk_size
+    D_residual = None
+    if D is not None:
+        D_residual = D[..., None] * pad_tensor_by_size(hidden_states, pad_size)
+
+    # Discretize x and A
+    hidden_states = hidden_states * dt[..., None].float()
+    A = A.to(hidden_states.dtype) * dt.float()
+
+     # Rearrange into blocks/chunks
+    hidden_states, A, B, C = [
+        reshape_into_chunks(tensor, pad_size, chunk_size) for tensor in (hidden_states, A, B, C)
+    ]
+
+    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)
+    G = (C[:, :, :, None, :, :] * B[:, :, None, :, :, :]).sum(dim=-1)
+
+    # Compute M, equivalent to applying attention mask to weights
+    M = (G[..., None] * L.permute(0, 2, 3, 4, 1)[..., None]).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 = (
+        initial_states[:, None].to(dtype=states.dtype, device=states.device)
+        if initial_states is not None
+        else torch.zeros_like(states[:, :1])
+    )
+    states = torch.cat([previous_states, states], dim=1)
+    decay_chunk = torch.exp(segment_sum(F.pad(A_cumsum[:, :, :, -1], (1, 0)))).transpose(1, 3)
+    new_states = (decay_chunk[..., None, None] * states[:, :, None, ...]).sum(dim=1)
+    states, final_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, ...]
+    Y_off = C_times_states.sum(-1) * state_decay_out.permute(0, 2, 3, 1)[..., None]
+
+    # Add output of intra-chunk and inter-chunk terms (diagonal and off-diagonal blocks)
+    output = Y_diag + Y_off
+    output = output.reshape(batch_size, -1, num_heads, head_dim)
+
+    if D_residual is not None:
+        output = output + D_residual
+
+    # Cutting off padded chunks
+    if pad_size > 0:
+        output = output[:, :sequence_length]
+
+    output = output.to(input_dtype)
+
+    if return_final_states:
+        return output, final_state
+
+    return output
+
+
 class Mamba2Mixer(nn.Module):
     """
     Compute ∆, A, B, C, and D the state space parameters and compute the `contextualized_states`.
@@ -433,6 +631,7 @@ def cuda_kernels_forward(
             # 4. Final linear projection
             out = self.out_proj(scan_output)
 
+        print("here")
         return out
 
     def torch_forward(
@@ -450,148 +649,131 @@ def torch_forward(
         hidden_states = apply_mask_to_padding_states(hidden_states, attention_mask)
         projected_states = self.in_proj(hidden_states)
 
+        A = -torch.exp(self.A_log.float())
+        dt_limit_kwargs = {} if self.time_step_limit == (0.0, float("inf")) else {"dt_limit": self.time_step_limit}
+        if self.training and cache_params is None:
+            fused_output = mamba2_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=kwargs.get("seq_idx"),  # TODO: kwargs
+                activation=self.activation,
+                rmsnorm_weight=self.norm.weight,
+                rmsnorm_eps=self.norm.variance_epsilon,
+                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,
+            )
+
+            # Only kernels can use this shortcircuit, fallback to normal torch otherwise
+            if fused_output is not None:
+                return fused_output
+
         gate, hidden_states_B_C, dt = projected_states.split(
             [self.intermediate_size, self.conv_dim, self.num_heads], dim=-1
         )
 
+        if use_precomputed_states:
+            conv_state = cache_params.layers[self.layer_idx].conv_states[0]
+            recurrent_state = cache_params.layers[self.layer_idx].recurrent_states[0]
+
         # 2. Convolution sequence transformation
-        hidden_states_B_C = self._convolution(hidden_states_B_C, cache_params, attention_mask, **kwargs)
-        hidden_states_B_C = apply_mask_to_padding_states(hidden_states_B_C, attention_mask)
+        if use_precomputed_states and seq_len == 1:
+            hidden_states_B_C = causal_conv1d_update(
+                hidden_states_B_C,
+                conv_state,
+                self.conv1d.weight.squeeze(1),
+                self.conv1d.bias,
+                activation=self.activation,
+            )
+        else:
+            if cache_params is not None:
+                hidden_states_B_C = cache_params.update_conv_state(
+                    hidden_states_B_C,
+                    self.layer_idx,
+                    conv_kernel_size=self.conv_kernel_size,
+                )
+
+            hidden_states_B_C = causal_conv1d_fn(
+                hidden_states_B_C,
+                self.conv1d.weight.squeeze(1),
+                self.conv1d.bias,
+                activation=self.activation,
+                **kwargs,
+            )
+
+            if cache_params is not None:
+                hidden_states_B_C = hidden_states_B_C[:, :, -seq_len:]
+
+        # 3. SSM transformation
+        hidden_states_B_C = apply_mask_to_padding_states(hidden_states_B_C.transpose(1, 2), attention_mask)
         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())
+        # Recurrent form
         if use_precomputed_states and seq_len == 1:
-            # We need to guarantee that anything regarding the cache is on the same device
-            cache_device = cache_params.layers[self.layer_idx].device
-
-            # Note: there is no need to pad parameter matrices here, as there is just one new token for batched generation
-            dt = dt.transpose(1, 2).expand(batch_size, dt.shape[-1], self.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))[..., None]
-            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)
-            dA = (torch.exp(dt * A)).to(device=cache_device)
-
-            # Discretize B
-            B = B.reshape(batch_size, self.n_groups, 1, -1)
-            B = B.expand(batch_size, self.n_groups, self.num_heads // self.n_groups, B.shape[-1]).contiguous()
-            B = B.reshape(batch_size, -1, 1, B.shape[-1])
-            dB = dt * B
-
-            # Discretize x into dB
-            hidden_states = hidden_states.reshape(batch_size, -1, self.head_dim)
-            dBx = (dB * hidden_states[..., None]).to(device=cache_device)
-
-            # State calculation
-            ssm_states = cache_params.layers[self.layer_idx].recurrent_states[0] * dA + dBx
-            ssm_states = cache_params.update_recurrent_state(ssm_states, layer_idx=self.layer_idx)
-
-            # Subsequent output
-            C = C.reshape(batch_size, self.n_groups, 1, -1)
-            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])
-
-            # Reshape ssm_states to merge the first two dimensions
-            ssm_states = ssm_states.to(device=C.device, dtype=C.dtype)
-            ssm_states_reshaped = ssm_states.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)
-            y = torch.bmm(ssm_states_reshaped, C_reshaped)
-            y = y.view(batch_size, self.num_heads, self.head_dim)
-
-            # D skip connection
-            D = self.D[..., None].expand(self.D.shape[0], self.head_dim)
-            y = (y + hidden_states * D).to(y.dtype)
-
-            y = y.reshape(batch_size, 1, -1)
+            hidden_states = hidden_states.view(batch_size, self.num_heads, self.head_dim)
+            dt = dt.transpose(1, 2).expand(-1, -1, self.head_dim)
+            A = A[:, None, None].expand(-1, self.head_dim, self.ssm_state_size)
+            B = B.view(batch_size, self.n_groups, self.ssm_state_size)
+            C = C.view(batch_size, self.n_groups, self.ssm_state_size)
+            D = self.D[:, None].expand(-1, self.head_dim)
+            dt_bias = self.dt_bias[:, None].expand(-1, self.head_dim)
+
+            scan_output = mamba2_selective_state_update(
+                recurrent_state,
+                hidden_states,
+                dt,
+                A,
+                B,
+                C,
+                D,
+                z=None,
+                dt_bias=dt_bias,
+                dt_softplus=True,
+                **kwargs,
+            )
+            scan_output = scan_output.view(batch_size, 1, -1)
+        # Chunk form
         else:
-            # 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_interleave(self.num_heads // self.n_groups, dim=2, output_size=self.num_heads)
-            C = C.repeat_interleave(self.num_heads // self.n_groups, dim=2, output_size=self.num_heads)
-            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)
-            ]
-
-            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)
-            G_intermediate = C[:, :, :, None, :, :] * B[:, :, None, :, :, :]
-            G = G_intermediate.sum(dim=-1)
-
-            # 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 = (
-                cache_params.layers[self.layer_idx]
-                .recurrent_states[0][:, None]
-                .to(dtype=states.dtype, device=states.device)
-                if use_precomputed_states
-                else torch.zeros_like(states[:, :1])
+            output_final_state = cache_params is not None
+            scan_result = mamba2_chunk_scan(
+                hidden_states.view(batch_size, seq_len, self.num_heads, self.head_dim),
+                dt,
+                A,
+                B.view(batch_size, seq_len, self.n_groups, self.ssm_state_size),
+                C.view(batch_size, seq_len, self.n_groups, self.ssm_state_size),
+                chunk_size=self.chunk_size,
+                D=self.D,
+                z=None,
+                return_final_states=output_final_state,
+                dt_bias=self.dt_bias,
+                dt_softplus=True,
+                initial_states=recurrent_state,
+                dt_limit=self.time_step_limit,
+                **kwargs,
             )
-            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
-            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)
 
-            # Init cache
-            if ssm_state is not None and cache_params is not None:
-                cache_params.update_recurrent_state(ssm_state, layer_idx=self.layer_idx)
+            if output_final_state:
+                scan_output, final_state = scan_result
+                cache_params.update_recurrent_state(final_state, layer_idx=self.layer_idx)
+            else:
+                scan_output = scan_result
+
+            scan_output = scan_output.reshape(batch_size, seq_len, -1)
 
-        scan_output = self.norm(y, gate)
+        scan_output = self.norm(scan_output, gate)
 
         # 4. Final linear projection
         contextualized_states = self.out_proj(scan_output.to(dtype))

From f0aa5812b196905fde667a4a7ee008f3db5f3681 Mon Sep 17 00:00:00 2001
From: vasqu <antonprogamer@gmail.com>
Date: Wed, 29 Jul 2026 23:56:22 +0000
Subject: [PATCH 16/43] kernels match --> hf kernels will be needed

---
 .../models/mamba2/modeling_mamba2.py          | 223 +-----------------
 1 file changed, 10 insertions(+), 213 deletions(-)

diff --git a/src/transformers/models/mamba2/modeling_mamba2.py b/src/transformers/models/mamba2/modeling_mamba2.py
index 444319bba8b6..4a8e56dd1241 100644
--- a/src/transformers/models/mamba2/modeling_mamba2.py
+++ b/src/transformers/models/mamba2/modeling_mamba2.py
@@ -24,13 +24,12 @@
 from ...activations import ACT2FN
 from ...cache_utils import Cache, DynamicCache
 from ...generation import GenerationMixin
-from ...integrations import lazy_load_kernel, use_kernel_func_from_hub_with_fallback
+from ...integrations import use_kernel_func_from_hub_with_fallback
 from ...integrations.accelerate import force_accelerate_hooks
 from ...masking_utils import create_recurrent_attention_mask
 from ...modeling_layers import GradientCheckpointingLayer
 from ...modeling_utils import PreTrainedModel
-from ...utils import ModelOutput, auto_docstring, is_torchdynamo_compiling, logging
-from ...utils.import_utils import resolve_internal_import
+from ...utils import ModelOutput, auto_docstring, logging
 from .configuration_mamba2 import Mamba2Config
 
 
@@ -121,6 +120,7 @@ def forward(self, hidden_states, gate=None):
         return self.weight * hidden_states.to(input_dtype)
 
 
+@use_kernel_func_from_hub_with_fallback("causal_conv1d_update", "causal_conv1d")
 def causal_conv1d_update(
     hidden_states: torch.Tensor,
     conv_state: torch.Tensor,
@@ -140,6 +140,7 @@ def causal_conv1d_update(
     return out.to(hidden_states.dtype)
 
 
+@use_kernel_func_from_hub_with_fallback("causal_conv1d_fn", "causal_conv1d")
 def causal_conv1d_fn(
     hidden_states: torch.Tensor,
     weight: nn.Parameter,
@@ -162,6 +163,7 @@ def causal_conv1d_fn(
     return out.to(hidden_states.dtype)
 
 
+# TODO: layer references for mamba2
 @use_kernel_func_from_hub_with_fallback(
     "mamba_split_conv1d_scan_combined",
     "mamba_ssm",
@@ -208,7 +210,7 @@ def mamba2_selective_state_update(
     dt_softplus: bool = False,
     **kwargs,
 ):
-    batch_size, num_heads, head_dim = hidden_states.shape[:2]
+    batch_size, num_heads, head_dim = hidden_states.shape
     num_groups = B.shape[1]
     state_size = B.shape[-1]
 
@@ -425,37 +427,6 @@ def __init__(self, config: Mamba2Config, layer_idx: int, initialize_mixer_weight
         self.out_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=config.use_bias)
         self.use_bias = config.use_bias
 
-        global causal_conv1d, causal_conv1d_update, causal_conv1d_fn
-        causal_conv1d = lazy_load_kernel("causal-conv1d")
-        causal_conv1d_update = getattr(causal_conv1d, "causal_conv1d_update", causal_conv1d_update)
-        causal_conv1d_fn = getattr(causal_conv1d, "causal_conv1d_fn", causal_conv1d_fn)
-
-        global mamba_ssm, selective_state_update, mamba_chunk_scan_combined, mamba_split_conv1d_scan_combined
-        mamba_ssm = lazy_load_kernel("mamba-ssm")
-        selective_state_update = resolve_internal_import(
-            mamba_ssm, chained_path="ops.triton.selective_state_update.selective_state_update"
-        )
-        mamba_chunk_scan_combined = resolve_internal_import(
-            mamba_ssm, chained_path="ops.triton.ssd_combined.mamba_chunk_scan_combined"
-        )
-        mamba_split_conv1d_scan_combined = resolve_internal_import(
-            mamba_ssm, chained_path="ops.triton.ssd_combined.mamba_split_conv1d_scan_combined"
-        )
-
-        global is_fast_path_available
-        is_fast_path_available = (
-            all((selective_state_update, mamba_chunk_scan_combined, mamba_split_conv1d_scan_combined))
-            and hasattr(causal_conv1d, "causal_conv1d_update")
-            and hasattr(causal_conv1d, "causal_conv1d_fn")
-        )
-
-        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"
-            )
-
         self.layer_type = config.layer_types[layer_idx]
 
     @torch.no_grad()
@@ -474,167 +445,8 @@ def init_mamba2_weights(self):
         inv_dt = dt + torch.log(-torch.expm1(-dt))
         init.copy_(self.dt_bias, inv_dt)
 
-    def _convolution(
-        self,
-        hidden_states: torch.Tensor,
-        cache_params: Cache | None = None,
-        attention_mask: torch.LongTensor | None = None,
-        **kwargs,
-    ):
-        seq_len = hidden_states.shape[1]
-        hidden_states = apply_mask_to_padding_states(hidden_states, attention_mask)
-        hidden_states = hidden_states.transpose(1, 2)
-
-        use_precomputed_states = cache_params is not None and cache_params.has_previous_state(self.layer_idx)
-
-        if use_precomputed_states and seq_len == 1 and not cache_params.layers[self.layer_idx].record_past:
-            conv_state = cache_params.layers[self.layer_idx].conv_states[0]
-            hidden_states = causal_conv1d_update(
-                hidden_states,
-                conv_state,
-                self.conv1d.weight.squeeze(1),
-                self.conv1d.bias,
-                self.activation,
-            )
-        else:
-            if cache_params is not None:
-                hidden_states = cache_params.update_conv_state(
-                    hidden_states, self.layer_idx, conv_kernel_size=self.conv_kernel_size
-                )
-
-            hidden_states = causal_conv1d_fn(
-                hidden_states,
-                self.conv1d.weight.squeeze(1),
-                self.conv1d.bias,
-                activation=self.activation,
-                seq_idx=kwargs.get("seq_idx"),
-            )
-
-            # Drop the additional previous states
-            if cache_params is not None:
-                hidden_states = hidden_states[:, :, -seq_len:]
-
-        hidden_states = hidden_states.transpose(1, 2)
-        return hidden_states
-
-    def cuda_kernels_forward(
-        self,
-        hidden_states: torch.Tensor,
-        cache_params: Cache | None = None,
-        attention_mask: torch.Tensor | None = None,
-        **kwargs,
-    ):
-        # 1. Gated MLP's linear projection
-        hidden_states = apply_mask_to_padding_states(hidden_states, attention_mask)
-        projected_states = self.in_proj(hidden_states)
-
-        A = -torch.exp(self.A_log.float())
-        dt_limit_kwargs = {} if self.time_step_limit == (0.0, float("inf")) else {"dt_limit": self.time_step_limit}
-        # Fused kernel for conv1d, SSM, and the final projection
-        if self.training and cache_params is None:
-            return 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=kwargs.get("seq_idx"),
-                activation=self.activation,
-                rmsnorm_weight=self.norm.weight,
-                rmsnorm_eps=self.norm.variance_epsilon,
-                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,
-            )
-
-        # Set up dimensions for reshapes later
-        batch_size, seq_len, _ = hidden_states.shape
-        groups_time_state_size = self.n_groups * self.ssm_state_size
-        use_precomputed_states = cache_params is not None and cache_params.has_previous_state(self.layer_idx)
-
-        gate, hidden_states_B_C, dt = projected_states.split(
-            [self.intermediate_size, self.conv_dim, self.num_heads], dim=-1
-        )
-
-        # Apply the conv
-        hidden_states_B_C = self._convolution(hidden_states_B_C, cache_params, attention_mask, **kwargs)
-        hidden_states_B_C = apply_mask_to_padding_states(hidden_states_B_C, attention_mask)
-        hidden_states, B, C = torch.split(
-            hidden_states_B_C,
-            [self.intermediate_size, groups_time_state_size, groups_time_state_size],
-            dim=-1,
-        )
-
-        recurrent_state = cache_params.layers[self.layer_idx].recurrent_states[0] if use_precomputed_states else None
-        # Single step calculations via cache
-        if use_precomputed_states and seq_len == 1:
-            # 3. SSM transformation
-            A = A[:, None, ...][:, :, None].expand(-1, self.head_dim, self.ssm_state_size).to(dtype=torch.float32)
-            dt = dt.transpose(1, 2).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[2] // self.n_groups)
-            C = C.view(batch_size, self.n_groups, C.shape[2] // self.n_groups)
-            hidden_states_reshaped = hidden_states.view(batch_size, self.num_heads, self.head_dim)
-            hidden_states = selective_state_update(
-                recurrent_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, 1, self.num_heads * self.head_dim)
-            hidden_states = self.norm(hidden_states, gate)
-
-            # 4. Final linear projection
-            out = self.out_proj(hidden_states)
-
-        # Fused calculations or step by step if no initialized cache is found
-        else:
-            # 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=kwargs.get("seq_idx"),
-                return_final_states=True,
-                dt_bias=self.dt_bias,
-                dt_softplus=True,
-                initial_states=recurrent_state,
-                **dt_limit_kwargs,
-            )
-
-            # Init cache
-            if ssm_state is not None and cache_params is not None:
-                cache_params.update_recurrent_state(ssm_state, layer_idx=self.layer_idx)
-
-            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)
-
-        print("here")
-        return out
-
-    def torch_forward(
+    @force_accelerate_hooks("conv1d")
+    def forward(
         self,
         hidden_states: torch.Tensor,
         cache_params: Cache | None = None,
@@ -686,6 +498,7 @@ def torch_forward(
             recurrent_state = cache_params.layers[self.layer_idx].recurrent_states[0]
 
         # 2. Convolution sequence transformation
+        hidden_states_B_C = hidden_states_B_C.transpose(1, 2)
         if use_precomputed_states and seq_len == 1:
             hidden_states_B_C = causal_conv1d_update(
                 hidden_states_B_C,
@@ -760,7 +573,7 @@ def torch_forward(
                 return_final_states=output_final_state,
                 dt_bias=self.dt_bias,
                 dt_softplus=True,
-                initial_states=recurrent_state,
+                initial_states=recurrent_state if cache_params is not None else None,
                 dt_limit=self.time_step_limit,
                 **kwargs,
             )
@@ -779,22 +592,6 @@ def torch_forward(
         contextualized_states = self.out_proj(scan_output.to(dtype))
         return contextualized_states
 
-    @force_accelerate_hooks("conv1d")
-    def forward(
-        self,
-        hidden_states,
-        cache_params: Cache | None = None,
-        attention_mask: torch.Tensor | None = None,
-        **kwargs,
-    ):
-        if is_fast_path_available and "cuda" in self.in_proj.weight.device.type and not is_torchdynamo_compiling():
-            return self.cuda_kernels_forward(hidden_states, cache_params, attention_mask, **kwargs)
-        if kwargs.get("seq_idx") is not None:
-            raise NotImplementedError(
-                "`seq_idx` support requires fast path support. Please install `mamba_ssm` and `causal_conv1d`"
-            )
-        return self.torch_forward(hidden_states, cache_params, attention_mask, **kwargs)
-
 
 class Mamba2RMSNorm(nn.Module):
     def __init__(self, hidden_size, eps=1e-6):

From 263f3077f423f5635cc185bd7be02bacfb741f4d Mon Sep 17 00:00:00 2001
From: vasqu <antonprogamer@gmail.com>
Date: Wed, 29 Jul 2026 23:57:00 +0000
Subject: [PATCH 17/43] style

---
 src/transformers/models/mamba2/modeling_mamba2.py | 6 ++----
 1 file changed, 2 insertions(+), 4 deletions(-)

diff --git a/src/transformers/models/mamba2/modeling_mamba2.py b/src/transformers/models/mamba2/modeling_mamba2.py
index 4a8e56dd1241..59442da76e20 100644
--- a/src/transformers/models/mamba2/modeling_mamba2.py
+++ b/src/transformers/models/mamba2/modeling_mamba2.py
@@ -298,10 +298,8 @@ def mamba2_chunk_scan(
     hidden_states = hidden_states * dt[..., None].float()
     A = A.to(hidden_states.dtype) * dt.float()
 
-     # Rearrange into blocks/chunks
-    hidden_states, A, B, C = [
-        reshape_into_chunks(tensor, pad_size, chunk_size) for tensor in (hidden_states, A, B, C)
-    ]
+    # Rearrange into blocks/chunks
+    hidden_states, A, B, C = [reshape_into_chunks(tensor, pad_size, chunk_size) for tensor in (hidden_states, A, B, C)]
 
     A = A.permute(0, 3, 1, 2)
     A_cumsum = torch.cumsum(A, dim=-1)

From a189da9e9a728827c3b317d118442278b0eda5a7 Mon Sep 17 00:00:00 2001
From: vasqu <antonprogamer@gmail.com>
Date: Thu, 30 Jul 2026 12:11:24 +0000
Subject: [PATCH 18/43] quick fixes

---
 .../models/mamba2/modeling_mamba2.py          |  2 +-
 .../olmo_hybrid/test_modeling_olmo_hybrid.py  | 45 ++++++-------------
 2 files changed, 15 insertions(+), 32 deletions(-)

diff --git a/src/transformers/models/mamba2/modeling_mamba2.py b/src/transformers/models/mamba2/modeling_mamba2.py
index 59442da76e20..74eccf0afa78 100644
--- a/src/transformers/models/mamba2/modeling_mamba2.py
+++ b/src/transformers/models/mamba2/modeling_mamba2.py
@@ -571,7 +571,7 @@ def forward(
                 return_final_states=output_final_state,
                 dt_bias=self.dt_bias,
                 dt_softplus=True,
-                initial_states=recurrent_state if cache_params is not None else None,
+                initial_states=recurrent_state if use_precomputed_states else None,
                 dt_limit=self.time_step_limit,
                 **kwargs,
             )
diff --git a/tests/models/olmo_hybrid/test_modeling_olmo_hybrid.py b/tests/models/olmo_hybrid/test_modeling_olmo_hybrid.py
index 80c123c5e1b5..1e96ae7b2f09 100644
--- a/tests/models/olmo_hybrid/test_modeling_olmo_hybrid.py
+++ b/tests/models/olmo_hybrid/test_modeling_olmo_hybrid.py
@@ -33,7 +33,6 @@
     import torch
 
     from transformers import (
-        Cache,
         DynamicCache,
         OlmoHybridForCausalLM,
         OlmoHybridModel,
@@ -66,6 +65,20 @@ class OlmoHybridModelTest(CausalLMModelTest, unittest.TestCase):
     model_tester_class = OlmoHybridModelTester
     rotary_embedding_layer = OlmoHybridRotaryEmbedding if is_torch_available() else None
 
+    def _get_conv_state_shape(self, batch_size: int, config):
+        conv_kernel = config.linear_conv_kernel_dim
+        key_dim = config.linear_key_head_dim * config.linear_num_key_heads
+        value_dim = config.linear_value_head_dim * config.linear_num_value_heads
+        # We have 3 conv states per layer, with different shapes
+        return [
+            (batch_size, key_dim, conv_kernel),
+            (batch_size, key_dim, conv_kernel),
+            (batch_size, value_dim, conv_kernel),
+        ]
+
+    def _get_recurrent_state_shape(self, batch_size: int, config):
+        return (batch_size, config.linear_num_value_heads, config.linear_key_head_dim, config.linear_value_head_dim)
+
     @unittest.skip("Float8 quantization + TP numerical noise exceeds match threshold")
     def test_tp_generation_quantized(self):
         pass
@@ -104,36 +117,6 @@ def test_linear_attention_multi_token_cached_forward_matches_single_token(self):
 
         torch.testing.assert_close(under_test_first, ref_first, rtol=1e-4, atol=1e-4)
 
-    # === Cache helper methods (same pattern as Qwen3Next) ===
-    def _check_past_key_values_for_generate(self, batch_size, past_key_values, seq_length, config):
-        """OlmoHybrid has a special Cache as it alternates with gated deltanet layers"""
-        self.assertIsInstance(past_key_values, DynamicCache)
-
-        num_heads = getattr(config, "num_key_value_heads", config.num_attention_heads)
-        head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)
-        expected_shape = (batch_size, num_heads, seq_length, head_dim)
-
-        attention_layer_indices = past_key_values.transformer_layers
-        self.assertListEqual(
-            [past_key_values.key_cache[idx].shape for idx in attention_layer_indices],
-            [expected_shape] * len(attention_layer_indices),
-        )
-        self.assertListEqual(
-            [past_key_values.value_cache[idx].shape for idx in attention_layer_indices],
-            [expected_shape] * len(attention_layer_indices),
-        )
-
-    def _check_caches_are_equal(self, cache1: Cache, cache2: Cache):
-        """OlmoHybrid has a special Cache as it alternates with gated deltanet layers"""
-        if not len(cache1) == len(cache2):
-            raise ValueError("Both caches do not have the same number of layers.")
-
-        num_layers = len(cache1)
-        for idx in range(num_layers):
-            if cache1.key_cache[idx] is not None:
-                torch.testing.assert_close(cache1.key_cache[idx], cache2.key_cache[idx])
-                torch.testing.assert_close(cache1.value_cache[idx], cache2.value_cache[idx])
-
     # === Override test_attention_outputs (same pattern as Qwen3Next) ===
     def test_attention_outputs(self):
         """Needs to be overwritten as OlmoHybrid alternates between attention layers and gated deltanet layers."""

From 4ebeb1113d45c563e85ea3b7df8c2459b380acdb Mon Sep 17 00:00:00 2001
From: vasqu <antonprogamer@gmail.com>
Date: Thu, 30 Jul 2026 13:20:01 +0000
Subject: [PATCH 19/43] mamba2 works

---
 src/transformers/integrations/hub_kernels.py  | 42 +++++++++++++++++++
 .../models/mamba2/modeling_mamba2.py          | 19 ++++++---
 2 files changed, 56 insertions(+), 5 deletions(-)

diff --git a/src/transformers/integrations/hub_kernels.py b/src/transformers/integrations/hub_kernels.py
index 9ddeb8094bda..d8f3da0e095b 100644
--- a/src/transformers/integrations/hub_kernels.py
+++ b/src/transformers/integrations/hub_kernels.py
@@ -207,6 +207,48 @@ def _build_kernel_mapping() -> dict:
                     ),
                 },
             },
+            "mamba_chunk_scan_combined": {
+                "cuda": {
+                    Mode.TRAINING: LayerRepository(
+                        repo_id="kernels-community/mamba-ssm",
+                        layer_name="mamba_chunk_scan_combined",
+                        version=1,
+                    ),
+                    Mode.INFERENCE: LayerRepository(
+                        repo_id="kernels-community/mamba-ssm",
+                        layer_name="mamba_chunk_scan_combined",
+                        version=1,
+                    ),
+                },
+            },
+            "mamba_split_conv1d_scan_combined": {
+                "cuda": {
+                    Mode.TRAINING: LayerRepository(
+                        repo_id="kernels-community/mamba-ssm",
+                        layer_name="mamba_split_conv1d_scan_combined",
+                        version=1,
+                    ),
+                    Mode.INFERENCE: LayerRepository(
+                        repo_id="kernels-community/mamba-ssm",
+                        layer_name="mamba_split_conv1d_scan_combined",
+                        version=1,
+                    ),
+                },
+            },
+            "selective_state_update": {
+                "cuda": {
+                    Mode.TRAINING: LayerRepository(
+                        repo_id="kernels-community/mamba-ssm",
+                        layer_name="selective_state_update",
+                        version=1,
+                    ),
+                    Mode.INFERENCE: LayerRepository(
+                        repo_id="kernels-community/mamba-ssm",
+                        layer_name="selective_state_update",
+                        version=1,
+                    ),
+                },
+            },
             "SwiGLUMLP": {
                 "cuda": {
                     Mode.INFERENCE | Mode.TORCH_COMPILE: LayerRepository(
diff --git a/src/transformers/models/mamba2/modeling_mamba2.py b/src/transformers/models/mamba2/modeling_mamba2.py
index 74eccf0afa78..185542b062f4 100644
--- a/src/transformers/models/mamba2/modeling_mamba2.py
+++ b/src/transformers/models/mamba2/modeling_mamba2.py
@@ -24,7 +24,7 @@
 from ...activations import ACT2FN
 from ...cache_utils import Cache, DynamicCache
 from ...generation import GenerationMixin
-from ...integrations import use_kernel_func_from_hub_with_fallback
+from ...integrations import use_kernel_func_from_hub_with_fallback, use_kernelized_func
 from ...integrations.accelerate import force_accelerate_hooks
 from ...masking_utils import create_recurrent_attention_mask
 from ...modeling_layers import GradientCheckpointingLayer
@@ -163,7 +163,6 @@ def causal_conv1d_fn(
     return out.to(hidden_states.dtype)
 
 
-# TODO: layer references for mamba2
 @use_kernel_func_from_hub_with_fallback(
     "mamba_split_conv1d_scan_combined",
     "mamba_ssm",
@@ -360,6 +359,15 @@ def mamba2_chunk_scan(
     return output
 
 
+@use_kernelized_func(
+    [
+        causal_conv1d_fn,
+        causal_conv1d_update,
+        mamba2_split_conv1d_scan_combined,
+        mamba2_selective_state_update,
+        mamba2_chunk_scan,
+    ]
+)
 class Mamba2Mixer(nn.Module):
     """
     Compute ∆, A, B, C, and D the state space parameters and compute the `contextualized_states`.
@@ -460,7 +468,9 @@ def forward(
         projected_states = self.in_proj(hidden_states)
 
         A = -torch.exp(self.A_log.float())
-        dt_limit_kwargs = {} if self.time_step_limit == (0.0, float("inf")) else {"dt_limit": self.time_step_limit}
+        fused_kwargs = (
+            kwargs | {} if self.time_step_limit == (0.0, float("inf")) else kwargs | {"dt_limit": self.time_step_limit}
+        )
         if self.training and cache_params is None:
             fused_output = mamba2_split_conv1d_scan_combined(
                 projected_states,
@@ -470,7 +480,6 @@ def forward(
                 A,
                 D=self.D,
                 chunk_size=self.chunk_size,
-                seq_idx=kwargs.get("seq_idx"),  # TODO: kwargs
                 activation=self.activation,
                 rmsnorm_weight=self.norm.weight,
                 rmsnorm_eps=self.norm.variance_epsilon,
@@ -480,7 +489,7 @@ def forward(
                 ngroups=self.n_groups,
                 norm_before_gate=False,
                 return_final_states=False,
-                **dt_limit_kwargs,
+                **fused_kwargs,
             )
 
             # Only kernels can use this shortcircuit, fallback to normal torch otherwise

From 1786c75b8b944e8951989a407b20a185d386fe4d Mon Sep 17 00:00:00 2001
From: vasqu <antonprogamer@gmail.com>
Date: Thu, 30 Jul 2026 14:46:05 +0000
Subject: [PATCH 20/43] mamba2 suite of models

---
 .../models/bamba/modeling_bamba.py            | 625 ++++++++---------
 .../models/bamba/modular_bamba.py             |  32 +-
 .../models/falcon_h1/modeling_falcon_h1.py    | 634 +++++++++---------
 .../models/falcon_h1/modular_falcon_h1.py     | 333 +++------
 .../modeling_granitemoehybrid.py              | 594 ++++++++--------
 .../models/mamba2/modeling_mamba2.py          |   5 +
 .../models/nemotron_h/modeling_nemotron_h.py  | 594 ++++++++--------
 .../models/nemotron_h/modular_nemotron_h.py   |  20 +-
 .../models/zamba2/modeling_zamba2.py          | 591 ++++++++--------
 9 files changed, 1579 insertions(+), 1849 deletions(-)

diff --git a/src/transformers/models/bamba/modeling_bamba.py b/src/transformers/models/bamba/modeling_bamba.py
index f2602342a539..c8da7690884e 100644
--- a/src/transformers/models/bamba/modeling_bamba.py
+++ b/src/transformers/models/bamba/modeling_bamba.py
@@ -24,7 +24,6 @@
 # limitations under the License.
 
 from collections.abc import Callable
-from typing import TypedDict
 
 import torch
 import torch.nn.functional as F
@@ -34,7 +33,7 @@
 from ...activations import ACT2FN
 from ...cache_utils import Cache, DynamicCache
 from ...generation import GenerationMixin
-from ...integrations import lazy_load_kernel, use_kernel_forward_from_hub
+from ...integrations import use_kernel_forward_from_hub, use_kernel_func_from_hub_with_fallback, use_kernelized_func
 from ...integrations.accelerate import force_accelerate_hooks
 from ...masking_utils import create_causal_mask, create_recurrent_attention_mask
 from ...modeling_layers import GradientCheckpointingLayer
@@ -42,41 +41,13 @@
 from ...modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update
 from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
 from ...processing_utils import Unpack
-from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, is_torchdynamo_compiling, logging
+from ...utils import auto_docstring, can_return_tuple
 from ...utils.deprecation import deprecate_kwarg
-from ...utils.generic import maybe_autocast, merge_with_config_defaults
-from ...utils.import_utils import resolve_internal_import
+from ...utils.generic import TransformersKwargs, maybe_autocast, merge_with_config_defaults
 from ...utils.output_capturing import capture_outputs
 from .configuration_bamba import BambaConfig
 
 
-logger = logging.get_logger(__name__)
-
-
-class BambaFlashAttentionKwargs(TypedDict, total=False):
-    """
-    Keyword arguments for advanced Flash Attention, causal-conv1d, and mamba_ssm kernel usage.
-    Use cases include padding-free training and fewer `torch.compile` graph breaks.
-
-    cu_seq_lens_q (`torch.LongTensor`):
-        Gets cumulative sequence length for query state.
-    cu_seq_lens_k (`torch.LongTensor`):
-        Gets cumulative sequence length for key state.
-    max_length_q (`int`):
-        Maximum sequence length for query state.
-    max_length_k (`int`):
-        Maximum sequence length for key state.
-    seq_idx (`torch.IntTensor`):
-        Index of each packed sequence.
-    """
-
-    cu_seq_lens_q: torch.LongTensor
-    cu_seq_lens_k: torch.LongTensor
-    max_length_q: int
-    max_length_k: int
-    seq_idx: torch.IntTensor
-
-
 class BambaRotaryEmbedding(nn.Module):
     inv_freq: torch.Tensor  # fix linting for `register_buffer`
 
@@ -372,6 +343,7 @@ def apply_mask_to_padding_states(hidden_states, attention_mask):
     return hidden_states
 
 
+@use_kernel_func_from_hub_with_fallback("causal_conv1d_update", "causal_conv1d")
 def causal_conv1d_update(
     hidden_states: torch.Tensor,
     conv_state: torch.Tensor,
@@ -391,6 +363,7 @@ def causal_conv1d_update(
     return out.to(hidden_states.dtype)
 
 
+@use_kernel_func_from_hub_with_fallback("causal_conv1d_fn", "causal_conv1d")
 def causal_conv1d_fn(
     hidden_states: torch.Tensor,
     weight: nn.Parameter,
@@ -413,6 +386,215 @@ def causal_conv1d_fn(
     return out.to(hidden_states.dtype)
 
 
+@use_kernel_func_from_hub_with_fallback(
+    "mamba_split_conv1d_scan_combined",
+    "mamba_ssm",
+    internal_path="ops.triton.ssd_combined",
+)
+def bamba_split_conv1d_scan_combined(
+    zxbcdt: torch.Tensor,
+    conv1d_weight: torch.Tensor,
+    conv1d_bias: torch.Tensor | None,
+    dt_bias: torch.Tensor,
+    A: torch.Tensor,
+    D: torch.Tensor,
+    chunk_size: int,
+    initial_states: 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,
+    **kwargs,
+):
+    return None
+
+
+@use_kernel_func_from_hub_with_fallback(
+    "selective_state_update",
+    "mamba_ssm",
+    internal_path="ops.triton.selective_state_update",
+)
+def bamba_selective_state_update(
+    state: torch.Tensor,
+    hidden_states: torch.Tensor,
+    dt: torch.Tensor,
+    A: torch.Tensor,
+    B: torch.Tensor,
+    C: torch.Tensor,
+    D: torch.Tensor | None = None,
+    dt_bias: torch.Tensor | None = None,
+    dt_softplus: bool = False,
+    z: torch.Tensor | None = None,
+    **kwargs,
+):
+    batch_size, num_heads, head_dim = hidden_states.shape
+    num_groups = B.shape[1]
+    state_size = B.shape[-1]
+
+    if dt_bias is not None:
+        dt = dt + dt_bias.to(dt.dtype)
+    if dt_softplus:
+        dt = F.softplus(dt)
+    dt = dt[..., None]
+
+    # Discretize A
+    dA = torch.exp(dt.float() * A.float()).to(device=state.device)
+
+    # Discretize B
+    B = B.reshape(batch_size, num_groups, 1, state_size)
+    B = B.expand(batch_size, num_groups, num_heads // num_groups, state_size).contiguous()
+    B = B.reshape(batch_size, num_heads, 1, state_size)
+    dB = dt * B
+
+    # Discretize x into dB
+    dBx = (dB * hidden_states[..., None]).to(device=state.device)
+
+    # State calculation
+    ssm_states = state * dA + dBx
+    state.copy_(ssm_states.to(state.dtype))
+
+    # Subsequent output
+    C = C.reshape(batch_size, num_groups, 1, state_size)
+    C = C.expand(batch_size, num_groups, num_heads // num_groups, state_size).contiguous()
+    C = C.reshape(batch_size, num_heads, state_size)
+
+    # Reshape ssm_states to merge the first two dimensions
+    ssm_states = ssm_states.to(device=C.device, dtype=C.dtype)
+    ssm_states_reshaped = ssm_states.view(batch_size * num_heads, head_dim, state_size)
+    C_reshaped = C.view(batch_size * num_heads, state_size, 1)
+    out = torch.bmm(ssm_states_reshaped, C_reshaped)
+    out = out.view(batch_size, num_heads, head_dim)
+
+    # D skip connection
+    if D is not None:
+        out = (out + hidden_states * D).to(out.dtype)
+
+    if z is not None:
+        out = out * F.silu(z)
+
+    return out.to(hidden_states.dtype)
+
+
+@use_kernel_func_from_hub_with_fallback(
+    "mamba_chunk_scan_combined",
+    "mamba_ssm",
+    internal_path="ops.triton.ssd_combined",
+)
+def bamba_chunk_scan(
+    hidden_states: torch.Tensor,
+    dt: torch.Tensor,
+    A: torch.Tensor,
+    B: torch.Tensor,
+    C: torch.Tensor,
+    chunk_size: int,
+    D: torch.Tensor | None = None,
+    dt_bias: torch.Tensor | None = None,
+    initial_states: torch.Tensor | None = None,
+    dt_softplus: bool = False,
+    dt_limit: tuple[float, float] = (0.0, float("inf")),
+    return_final_states: bool = False,
+    **kwargs,
+):
+    input_dtype = hidden_states.dtype
+    batch_size, sequence_length, num_heads, head_dim = hidden_states.shape
+    num_groups = B.shape[2]
+
+    if dt_bias is not None:
+        dt = dt + dt_bias.to(dt.dtype)
+    if dt_softplus:
+        dt = F.softplus(dt)
+    dt = torch.clamp(dt, min=dt_limit[0], max=dt_limit[1])
+
+    hidden_states = hidden_states.float()
+    B = B.float().repeat_interleave(num_heads // num_groups, dim=2, output_size=num_heads)
+    C = C.float().repeat_interleave(num_heads // num_groups, dim=2, output_size=num_heads)
+
+    pad_size = (chunk_size - sequence_length % chunk_size) % chunk_size
+    D_residual = None
+    if D is not None:
+        D_residual = D[..., None] * pad_tensor_by_size(hidden_states, pad_size)
+
+    # Discretize x and A
+    hidden_states = hidden_states * dt[..., None].float()
+    A = A.to(hidden_states.dtype) * dt.float()
+
+    # Rearrange into blocks/chunks
+    hidden_states, A, B, C = [reshape_into_chunks(tensor, pad_size, chunk_size) for tensor in (hidden_states, A, B, C)]
+
+    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)
+    G = (C[:, :, :, None, :, :] * B[:, :, None, :, :, :]).sum(dim=-1)
+
+    # Compute M, equivalent to applying attention mask to weights
+    M = (G[..., None] * L.permute(0, 2, 3, 4, 1)[..., None]).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 = (
+        initial_states[:, None].to(dtype=states.dtype, device=states.device)
+        if initial_states is not None
+        else torch.zeros_like(states[:, :1])
+    )
+    states = torch.cat([previous_states, states], dim=1)
+    decay_chunk = torch.exp(segment_sum(F.pad(A_cumsum[:, :, :, -1], (1, 0)))).transpose(1, 3)
+    new_states = (decay_chunk[..., None, None] * states[:, :, None, ...]).sum(dim=1)
+    states, final_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, ...]
+    Y_off = C_times_states.sum(-1) * state_decay_out.permute(0, 2, 3, 1)[..., None]
+
+    # Add output of intra-chunk and inter-chunk terms (diagonal and off-diagonal blocks)
+    output = Y_diag + Y_off
+    output = output.reshape(batch_size, -1, num_heads, head_dim)
+
+    if D_residual is not None:
+        output = output + D_residual
+
+    # Cutting off padded chunks
+    if pad_size > 0:
+        output = output[:, :sequence_length]
+
+    output = output.to(input_dtype)
+
+    if return_final_states:
+        return output, final_state
+
+    return output
+
+
+@use_kernelized_func(
+    [
+        causal_conv1d_fn,
+        causal_conv1d_update,
+        bamba_split_conv1d_scan_combined,
+        bamba_selective_state_update,
+        bamba_chunk_scan,
+    ]
+)
 class BambaMixer(nn.Module):
     """
     Compute ∆, A, B, C, and D the state space parameters and compute the `contextualized_states`.
@@ -473,37 +655,6 @@ def __init__(self, config: BambaConfig, layer_idx: int, initialize_mixer_weights
             self.init_bamba_weights()
         self.out_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=self.use_bias)
 
-        global causal_conv1d, causal_conv1d_update, causal_conv1d_fn
-        causal_conv1d = lazy_load_kernel("causal-conv1d")
-        causal_conv1d_update = getattr(causal_conv1d, "causal_conv1d_update", causal_conv1d_update)
-        causal_conv1d_fn = getattr(causal_conv1d, "causal_conv1d_fn", causal_conv1d_fn)
-
-        global mamba_ssm, selective_state_update, mamba_chunk_scan_combined, mamba_split_conv1d_scan_combined
-        mamba_ssm = lazy_load_kernel("mamba-ssm")
-        selective_state_update = resolve_internal_import(
-            mamba_ssm, chained_path="ops.triton.selective_state_update.selective_state_update"
-        )
-        mamba_chunk_scan_combined = resolve_internal_import(
-            mamba_ssm, chained_path="ops.triton.ssd_combined.mamba_chunk_scan_combined"
-        )
-        mamba_split_conv1d_scan_combined = resolve_internal_import(
-            mamba_ssm, chained_path="ops.triton.ssd_combined.mamba_split_conv1d_scan_combined"
-        )
-
-        global is_fast_path_available
-        is_fast_path_available = (
-            all((selective_state_update, mamba_chunk_scan_combined, mamba_split_conv1d_scan_combined))
-            and hasattr(causal_conv1d, "causal_conv1d_update")
-            and hasattr(causal_conv1d, "causal_conv1d_fn")
-        )
-
-        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"
-            )
-
         self.layer_type = config.layer_types[layer_idx]
 
     @torch.no_grad()
@@ -513,65 +664,28 @@ def init_bamba_weights(self):
         init.ones_(self.D)
         init.ones_(self.dt_bias)
 
-    def _convolution(
+    @force_accelerate_hooks("conv1d")
+    def forward(
         self,
         hidden_states: torch.Tensor,
         cache_params: Cache | None = None,
-        attention_mask: torch.LongTensor | None = None,
+        attention_mask: torch.Tensor | None = None,
         **kwargs,
     ):
-        seq_len = hidden_states.shape[1]
-        hidden_states = apply_mask_to_padding_states(hidden_states, attention_mask)
-        hidden_states = hidden_states.transpose(1, 2)
-
+        batch_size, seq_len, _ = hidden_states.shape
+        dtype = hidden_states.dtype
         use_precomputed_states = cache_params is not None and cache_params.has_previous_state(self.layer_idx)
 
-        if use_precomputed_states and seq_len == 1 and not cache_params.layers[self.layer_idx].record_past:
-            conv_state = cache_params.layers[self.layer_idx].conv_states[0]
-            hidden_states = causal_conv1d_update(
-                hidden_states,
-                conv_state,
-                self.conv1d.weight.squeeze(1),
-                self.conv1d.bias,
-                self.activation,
-            )
-        else:
-            if cache_params is not None:
-                hidden_states = cache_params.update_conv_state(
-                    hidden_states, self.layer_idx, conv_kernel_size=self.conv_kernel_size
-                )
-
-            hidden_states = causal_conv1d_fn(
-                hidden_states,
-                self.conv1d.weight.squeeze(1),
-                self.conv1d.bias,
-                activation=self.activation,
-                seq_idx=kwargs.get("seq_idx"),
-            )
-
-            # Drop the additional previous states
-            if cache_params is not None:
-                hidden_states = hidden_states[:, :, -seq_len:]
-
-        hidden_states = hidden_states.transpose(1, 2)
-        return hidden_states
-
-    def cuda_kernels_forward(
-        self,
-        hidden_states: torch.Tensor,
-        cache_params: Cache | None = None,
-        attention_mask: torch.Tensor | None = None,
-        **kwargs,
-    ):
         # 1. Gated MLP's linear projection
         hidden_states = apply_mask_to_padding_states(hidden_states, attention_mask)
         projected_states = self.in_proj(hidden_states)
 
         A = -torch.exp(self.A_log.float())
-        dt_limit_kwargs = {} if self.time_step_limit == (0.0, float("inf")) else {"dt_limit": self.time_step_limit}
-        # Fused kernel for conv1d, SSM, and the final projection
+        fused_kwargs = (
+            kwargs | {} if self.time_step_limit == (0.0, float("inf")) else kwargs | {"dt_limit": self.time_step_limit}
+        )
         if self.training and cache_params is None:
-            return mamba_split_conv1d_scan_combined(
+            fused_output = bamba_split_conv1d_scan_combined(
                 projected_states,
                 self.conv1d.weight.squeeze(1),
                 self.conv1d.bias,
@@ -579,7 +693,6 @@ def cuda_kernels_forward(
                 A,
                 D=self.D,
                 chunk_size=self.chunk_size,
-                seq_idx=kwargs.get("seq_idx"),
                 activation=self.activation,
                 rmsnorm_weight=self.norm.weight,
                 rmsnorm_eps=self.norm.variance_epsilon,
@@ -589,41 +702,71 @@ def cuda_kernels_forward(
                 ngroups=self.n_groups,
                 norm_before_gate=False,
                 return_final_states=False,
-                **dt_limit_kwargs,
+                **fused_kwargs,
             )
 
-        # Set up dimensions for reshapes later
-        batch_size, seq_len, _ = hidden_states.shape
-        groups_time_state_size = self.n_groups * self.ssm_state_size
-        use_precomputed_states = cache_params is not None and cache_params.has_previous_state(self.layer_idx)
+            # Only kernels can use this shortcircuit, fallback to normal torch otherwise
+            if fused_output is not None:
+                return fused_output
 
         gate, hidden_states_B_C, dt = projected_states.split(
             [self.intermediate_size, self.conv_dim, self.num_heads], dim=-1
         )
 
-        # Apply the conv
-        hidden_states_B_C = self._convolution(hidden_states_B_C, cache_params, attention_mask, **kwargs)
-        hidden_states_B_C = apply_mask_to_padding_states(hidden_states_B_C, attention_mask)
+        if use_precomputed_states:
+            conv_state = cache_params.layers[self.layer_idx].conv_states[0]
+            recurrent_state = cache_params.layers[self.layer_idx].recurrent_states[0]
+
+        # 2. Convolution sequence transformation
+        hidden_states_B_C = hidden_states_B_C.transpose(1, 2)
+        if use_precomputed_states and seq_len == 1:
+            hidden_states_B_C = causal_conv1d_update(
+                hidden_states_B_C,
+                conv_state,
+                self.conv1d.weight.squeeze(1),
+                self.conv1d.bias,
+                activation=self.activation,
+            )
+        else:
+            if cache_params is not None:
+                hidden_states_B_C = cache_params.update_conv_state(
+                    hidden_states_B_C,
+                    self.layer_idx,
+                    conv_kernel_size=self.conv_kernel_size,
+                )
+
+            hidden_states_B_C = causal_conv1d_fn(
+                hidden_states_B_C,
+                self.conv1d.weight.squeeze(1),
+                self.conv1d.bias,
+                activation=self.activation,
+                **kwargs,
+            )
+
+            if cache_params is not None:
+                hidden_states_B_C = hidden_states_B_C[:, :, -seq_len:]
+
+        # 3. SSM transformation
+        hidden_states_B_C = apply_mask_to_padding_states(hidden_states_B_C.transpose(1, 2), attention_mask)
         hidden_states, B, C = torch.split(
             hidden_states_B_C,
-            [self.intermediate_size, groups_time_state_size, groups_time_state_size],
+            [self.intermediate_size, self.n_groups * self.ssm_state_size, self.n_groups * self.ssm_state_size],
             dim=-1,
         )
 
-        recurrent_state = cache_params.layers[self.layer_idx].recurrent_states[0] if use_precomputed_states else None
-        # Single step calculations via cache
+        # Recurrent form
         if use_precomputed_states and seq_len == 1:
-            # 3. SSM transformation
-            A = A[:, None, ...][:, :, None].expand(-1, self.head_dim, self.ssm_state_size).to(dtype=torch.float32)
+            hidden_states = hidden_states.view(batch_size, self.num_heads, self.head_dim)
             dt = dt.transpose(1, 2).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[2] // self.n_groups)
-            C = C.view(batch_size, self.n_groups, C.shape[2] // self.n_groups)
-            hidden_states_reshaped = hidden_states.view(batch_size, self.num_heads, self.head_dim)
-            hidden_states = selective_state_update(
+            A = A[:, None, None].expand(-1, self.head_dim, self.ssm_state_size)
+            B = B.view(batch_size, self.n_groups, self.ssm_state_size)
+            C = C.view(batch_size, self.n_groups, self.ssm_state_size)
+            D = self.D[:, None].expand(-1, self.head_dim)
+            dt_bias = self.dt_bias[:, None].expand(-1, self.head_dim)
+
+            scan_output = bamba_selective_state_update(
                 recurrent_state,
-                hidden_states_reshaped,
+                hidden_states,
                 dt,
                 A,
                 B,
@@ -632,224 +775,44 @@ def cuda_kernels_forward(
                 z=None,
                 dt_bias=dt_bias,
                 dt_softplus=True,
+                **kwargs,
             )
-            hidden_states = hidden_states.view(batch_size, 1, self.num_heads * self.head_dim)
-            hidden_states = self.norm(hidden_states, gate)
-
-            # 4. Final linear projection
-            out = self.out_proj(hidden_states)
+            scan_output = scan_output.view(batch_size, 1, -1)
 
-        # Fused calculations or step by step if no initialized cache is found
+        # Chunk form
         else:
-            # 3. SSM transformation
-            scan_output, ssm_state = mamba_chunk_scan_combined(
-                hidden_states.view(batch_size, seq_len, -1, self.head_dim),
+            output_final_state = cache_params is not None
+            scan_result = bamba_chunk_scan(
+                hidden_states.view(batch_size, seq_len, self.num_heads, 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),
+                B.view(batch_size, seq_len, self.n_groups, self.ssm_state_size),
+                C.view(batch_size, seq_len, self.n_groups, self.ssm_state_size),
                 chunk_size=self.chunk_size,
                 D=self.D,
                 z=None,
-                seq_idx=kwargs.get("seq_idx"),
-                return_final_states=True,
+                return_final_states=output_final_state,
                 dt_bias=self.dt_bias,
                 dt_softplus=True,
-                initial_states=recurrent_state,
-                **dt_limit_kwargs,
+                initial_states=recurrent_state if use_precomputed_states else None,
+                dt_limit=self.time_step_limit,
+                **kwargs,
             )
 
-            # Init cache
-            if ssm_state is not None and cache_params is not None:
-                cache_params.update_recurrent_state(ssm_state, layer_idx=self.layer_idx)
-
-            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
-
-    def torch_forward(
-        self,
-        hidden_states: torch.Tensor,
-        cache_params: Cache | None = None,
-        attention_mask: torch.Tensor | None = None,
-        **kwargs,
-    ):
-        batch_size, seq_len, _ = hidden_states.shape
-        dtype = hidden_states.dtype
-        use_precomputed_states = cache_params is not None and cache_params.has_previous_state(self.layer_idx)
-
-        # 1. Gated MLP's linear projection
-        hidden_states = apply_mask_to_padding_states(hidden_states, attention_mask)
-        projected_states = self.in_proj(hidden_states)
+            if output_final_state:
+                scan_output, final_state = scan_result
+                cache_params.update_recurrent_state(final_state, layer_idx=self.layer_idx)
+            else:
+                scan_output = scan_result
 
-        gate, hidden_states_B_C, dt = projected_states.split(
-            [self.intermediate_size, self.conv_dim, self.num_heads], dim=-1
-        )
-
-        # 2. Convolution sequence transformation
-        hidden_states_B_C = self._convolution(hidden_states_B_C, cache_params, attention_mask, **kwargs)
-        hidden_states_B_C = apply_mask_to_padding_states(hidden_states_B_C, attention_mask)
-        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,
-        )
+            scan_output = scan_output.reshape(batch_size, seq_len, -1)
 
-        # 3. SSM transformation
-        A = -torch.exp(self.A_log.float())
-        if use_precomputed_states and seq_len == 1:
-            # We need to guarantee that anything regarding the cache is on the same device
-            cache_device = cache_params.layers[self.layer_idx].device
-
-            # Note: there is no need to pad parameter matrices here, as there is just one new token for batched generation
-            dt = dt.transpose(1, 2).expand(batch_size, dt.shape[-1], self.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))[..., None]
-            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)
-            dA = (torch.exp(dt * A)).to(device=cache_device)
-
-            # Discretize B
-            B = B.reshape(batch_size, self.n_groups, 1, -1)
-            B = B.expand(batch_size, self.n_groups, self.num_heads // self.n_groups, B.shape[-1]).contiguous()
-            B = B.reshape(batch_size, -1, 1, B.shape[-1])
-            dB = dt * B
-
-            # Discretize x into dB
-            hidden_states = hidden_states.reshape(batch_size, -1, self.head_dim)
-            dBx = (dB * hidden_states[..., None]).to(device=cache_device)
-
-            # State calculation
-            ssm_states = cache_params.layers[self.layer_idx].recurrent_states[0] * dA + dBx
-            ssm_states = cache_params.update_recurrent_state(ssm_states, layer_idx=self.layer_idx)
-
-            # Subsequent output
-            C = C.reshape(batch_size, self.n_groups, 1, -1)
-            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])
-
-            # Reshape ssm_states to merge the first two dimensions
-            ssm_states = ssm_states.to(device=C.device, dtype=C.dtype)
-            ssm_states_reshaped = ssm_states.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)
-            y = torch.bmm(ssm_states_reshaped, C_reshaped)
-            y = y.view(batch_size, self.num_heads, self.head_dim)
-
-            # D skip connection
-            D = self.D[..., None].expand(self.D.shape[0], self.head_dim)
-            y = (y + hidden_states * D).to(y.dtype)
-
-            y = y.reshape(batch_size, 1, -1)
-        else:
-            # 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_interleave(self.num_heads // self.n_groups, dim=2, output_size=self.num_heads)
-            C = C.repeat_interleave(self.num_heads // self.n_groups, dim=2, output_size=self.num_heads)
-            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)
-            ]
-
-            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)
-            G_intermediate = C[:, :, :, None, :, :] * B[:, :, None, :, :, :]
-            G = G_intermediate.sum(dim=-1)
-
-            # 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 = (
-                cache_params.layers[self.layer_idx]
-                .recurrent_states[0][:, None]
-                .to(dtype=states.dtype, device=states.device)
-                if use_precomputed_states
-                else 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
-            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)
-
-            # Init cache
-            if ssm_state is not None and cache_params is not None:
-                cache_params.update_recurrent_state(ssm_state, layer_idx=self.layer_idx)
-
-        scan_output = self.norm(y, gate)
+        scan_output = self.norm(scan_output, gate)
 
         # 4. Final linear projection
         contextualized_states = self.out_proj(scan_output.to(dtype))
         return contextualized_states
 
-    @force_accelerate_hooks("conv1d")
-    def forward(
-        self,
-        hidden_states,
-        cache_params: Cache | None = None,
-        attention_mask: torch.Tensor | None = None,
-        **kwargs,
-    ):
-        if is_fast_path_available and "cuda" in self.in_proj.weight.device.type and not is_torchdynamo_compiling():
-            return self.cuda_kernels_forward(hidden_states, cache_params, attention_mask, **kwargs)
-        if kwargs.get("seq_idx") is not None:
-            raise NotImplementedError(
-                "`seq_idx` support requires fast path support. Please install `mamba_ssm` and `causal_conv1d`"
-            )
-        return self.torch_forward(hidden_states, cache_params, attention_mask, **kwargs)
-
 
 class BambaMLP(nn.Module):
     def __init__(self, config):
@@ -914,7 +877,7 @@ def forward(
         past_key_values: Cache | None = None,
         use_cache: bool | None = False,
         position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
-        **kwargs: Unpack[BambaFlashAttentionKwargs],
+        **kwargs: Unpack[TransformersKwargs],
     ) -> tuple[torch.FloatTensor, tuple[torch.FloatTensor, torch.FloatTensor] | None]:
         residual = hidden_states
 
@@ -1006,7 +969,7 @@ def forward(
         past_key_values: Cache | None = None,
         inputs_embeds: torch.FloatTensor | None = None,
         use_cache: bool | None = None,
-        **kwargs: Unpack[BambaFlashAttentionKwargs],
+        **kwargs: Unpack[TransformersKwargs],
     ) -> BaseModelOutputWithPast:
         if (input_ids is None) ^ (inputs_embeds is not None):
             raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
diff --git a/src/transformers/models/bamba/modular_bamba.py b/src/transformers/models/bamba/modular_bamba.py
index 980a27bda429..c2fd89e285ed 100644
--- a/src/transformers/models/bamba/modular_bamba.py
+++ b/src/transformers/models/bamba/modular_bamba.py
@@ -18,8 +18,6 @@
 # limitations under the License.
 """PyTorch Bamba model."""
 
-from typing import TypedDict
-
 import torch
 from torch import nn
 
@@ -30,7 +28,7 @@
 from ...modeling_utils import PreTrainedModel
 from ...processing_utils import Unpack
 from ...utils import auto_docstring, can_return_tuple, logging
-from ...utils.generic import merge_with_config_defaults, no_inherit_decorator
+from ...utils.generic import TransformersKwargs, merge_with_config_defaults, no_inherit_decorator
 from ...utils.output_capturing import capture_outputs
 from ..jamba.modeling_jamba import JambaAttentionDecoderLayer
 from ..llama.modeling_llama import (
@@ -51,30 +49,6 @@
 logger = logging.get_logger(__name__)
 
 
-class BambaFlashAttentionKwargs(TypedDict, total=False):
-    """
-    Keyword arguments for advanced Flash Attention, causal-conv1d, and mamba_ssm kernel usage.
-    Use cases include padding-free training and fewer `torch.compile` graph breaks.
-
-    cu_seq_lens_q (`torch.LongTensor`):
-        Gets cumulative sequence length for query state.
-    cu_seq_lens_k (`torch.LongTensor`):
-        Gets cumulative sequence length for key state.
-    max_length_q (`int`):
-        Maximum sequence length for query state.
-    max_length_k (`int`):
-        Maximum sequence length for key state.
-    seq_idx (`torch.IntTensor`):
-        Index of each packed sequence.
-    """
-
-    cu_seq_lens_q: torch.LongTensor
-    cu_seq_lens_k: torch.LongTensor
-    max_length_q: int
-    max_length_k: int
-    seq_idx: torch.IntTensor
-
-
 class BambaRotaryEmbedding(LlamaRotaryEmbedding):
     def compute_default_rope_parameters(config: BambaConfig, device=None, **kwargs) -> tuple[torch.Tensor, float]:
         """
@@ -233,7 +207,7 @@ def forward(
         past_key_values: Cache | None = None,
         use_cache: bool | None = False,
         position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
-        **kwargs: Unpack[BambaFlashAttentionKwargs],
+        **kwargs: Unpack[TransformersKwargs],
     ) -> tuple[torch.FloatTensor, tuple[torch.FloatTensor, torch.FloatTensor] | None]:
         residual = hidden_states
 
@@ -325,7 +299,7 @@ def forward(
         past_key_values: Cache | None = None,
         inputs_embeds: torch.FloatTensor | None = None,
         use_cache: bool | None = None,
-        **kwargs: Unpack[BambaFlashAttentionKwargs],
+        **kwargs: Unpack[TransformersKwargs],
     ) -> BaseModelOutputWithPast:
         if (input_ids is None) ^ (inputs_embeds is not None):
             raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
diff --git a/src/transformers/models/falcon_h1/modeling_falcon_h1.py b/src/transformers/models/falcon_h1/modeling_falcon_h1.py
index 5c0b35faafbf..98dfe7b0352f 100644
--- a/src/transformers/models/falcon_h1/modeling_falcon_h1.py
+++ b/src/transformers/models/falcon_h1/modeling_falcon_h1.py
@@ -33,7 +33,7 @@
 from ...activations import ACT2FN
 from ...cache_utils import Cache, DynamicCache
 from ...generation import GenerationMixin
-from ...integrations import lazy_load_kernel, use_kernel_forward_from_hub, use_kernelized_func
+from ...integrations import use_kernel_forward_from_hub, use_kernel_func_from_hub_with_fallback, use_kernelized_func
 from ...integrations.accelerate import force_accelerate_hooks
 from ...masking_utils import create_causal_mask, create_recurrent_attention_mask
 from ...modeling_flash_attention_utils import FlashAttentionKwargs
@@ -42,17 +42,13 @@
 from ...modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update
 from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
 from ...processing_utils import Unpack
-from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, is_torchdynamo_compiling, logging
+from ...utils import TransformersKwargs, auto_docstring, can_return_tuple
 from ...utils.deprecation import deprecate_kwarg
 from ...utils.generic import maybe_autocast, merge_with_config_defaults
-from ...utils.import_utils import resolve_internal_import
 from ...utils.output_capturing import capture_outputs
 from .configuration_falcon_h1 import FalconH1Config
 
 
-logger = logging.get_logger(__name__)
-
-
 class FalconH1RotaryEmbedding(nn.Module):
     inv_freq: torch.Tensor  # fix linting for `register_buffer`
 
@@ -354,6 +350,7 @@ def apply_mask_to_padding_states(hidden_states, attention_mask):
     return hidden_states
 
 
+@use_kernel_func_from_hub_with_fallback("causal_conv1d_update", "causal_conv1d")
 def causal_conv1d_update(
     hidden_states: torch.Tensor,
     conv_state: torch.Tensor,
@@ -373,6 +370,7 @@ def causal_conv1d_update(
     return out.to(hidden_states.dtype)
 
 
+@use_kernel_func_from_hub_with_fallback("causal_conv1d_fn", "causal_conv1d")
 def causal_conv1d_fn(
     hidden_states: torch.Tensor,
     weight: nn.Parameter,
@@ -395,6 +393,215 @@ def causal_conv1d_fn(
     return out.to(hidden_states.dtype)
 
 
+@use_kernel_func_from_hub_with_fallback(
+    "mamba_split_conv1d_scan_combined",
+    "mamba_ssm",
+    internal_path="ops.triton.ssd_combined",
+)
+def falcon_h1_split_conv1d_scan_combined(
+    zxbcdt: torch.Tensor,
+    conv1d_weight: torch.Tensor,
+    conv1d_bias: torch.Tensor | None,
+    dt_bias: torch.Tensor,
+    A: torch.Tensor,
+    D: torch.Tensor,
+    chunk_size: int,
+    initial_states: 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,
+    **kwargs,
+):
+    return None
+
+
+@use_kernel_func_from_hub_with_fallback(
+    "selective_state_update",
+    "mamba_ssm",
+    internal_path="ops.triton.selective_state_update",
+)
+def falcon_h1_selective_state_update(
+    state: torch.Tensor,
+    hidden_states: torch.Tensor,
+    dt: torch.Tensor,
+    A: torch.Tensor,
+    B: torch.Tensor,
+    C: torch.Tensor,
+    D: torch.Tensor | None = None,
+    dt_bias: torch.Tensor | None = None,
+    dt_softplus: bool = False,
+    z: torch.Tensor | None = None,
+    **kwargs,
+):
+    batch_size, num_heads, head_dim = hidden_states.shape
+    num_groups = B.shape[1]
+    state_size = B.shape[-1]
+
+    if dt_bias is not None:
+        dt = dt + dt_bias.to(dt.dtype)
+    if dt_softplus:
+        dt = F.softplus(dt)
+    dt = dt[..., None]
+
+    # Discretize A
+    dA = torch.exp(dt.float() * A.float()).to(device=state.device)
+
+    # Discretize B
+    B = B.reshape(batch_size, num_groups, 1, state_size)
+    B = B.expand(batch_size, num_groups, num_heads // num_groups, state_size).contiguous()
+    B = B.reshape(batch_size, num_heads, 1, state_size)
+    dB = dt * B
+
+    # Discretize x into dB
+    dBx = (dB * hidden_states[..., None]).to(device=state.device)
+
+    # State calculation
+    ssm_states = state * dA + dBx
+    state.copy_(ssm_states.to(state.dtype))
+
+    # Subsequent output
+    C = C.reshape(batch_size, num_groups, 1, state_size)
+    C = C.expand(batch_size, num_groups, num_heads // num_groups, state_size).contiguous()
+    C = C.reshape(batch_size, num_heads, state_size)
+
+    # Reshape ssm_states to merge the first two dimensions
+    ssm_states = ssm_states.to(device=C.device, dtype=C.dtype)
+    ssm_states_reshaped = ssm_states.view(batch_size * num_heads, head_dim, state_size)
+    C_reshaped = C.view(batch_size * num_heads, state_size, 1)
+    out = torch.bmm(ssm_states_reshaped, C_reshaped)
+    out = out.view(batch_size, num_heads, head_dim)
+
+    # D skip connection
+    if D is not None:
+        out = (out + hidden_states * D).to(out.dtype)
+
+    if z is not None:
+        out = out * F.silu(z)
+
+    return out.to(hidden_states.dtype)
+
+
+@use_kernel_func_from_hub_with_fallback(
+    "mamba_chunk_scan_combined",
+    "mamba_ssm",
+    internal_path="ops.triton.ssd_combined",
+)
+def falcon_h1_chunk_scan(
+    hidden_states: torch.Tensor,
+    dt: torch.Tensor,
+    A: torch.Tensor,
+    B: torch.Tensor,
+    C: torch.Tensor,
+    chunk_size: int,
+    D: torch.Tensor | None = None,
+    dt_bias: torch.Tensor | None = None,
+    initial_states: torch.Tensor | None = None,
+    dt_softplus: bool = False,
+    dt_limit: tuple[float, float] = (0.0, float("inf")),
+    return_final_states: bool = False,
+    **kwargs,
+):
+    input_dtype = hidden_states.dtype
+    batch_size, sequence_length, num_heads, head_dim = hidden_states.shape
+    num_groups = B.shape[2]
+
+    if dt_bias is not None:
+        dt = dt + dt_bias.to(dt.dtype)
+    if dt_softplus:
+        dt = F.softplus(dt)
+    dt = torch.clamp(dt, min=dt_limit[0], max=dt_limit[1])
+
+    hidden_states = hidden_states.float()
+    B = B.float().repeat_interleave(num_heads // num_groups, dim=2, output_size=num_heads)
+    C = C.float().repeat_interleave(num_heads // num_groups, dim=2, output_size=num_heads)
+
+    pad_size = (chunk_size - sequence_length % chunk_size) % chunk_size
+    D_residual = None
+    if D is not None:
+        D_residual = D[..., None] * pad_tensor_by_size(hidden_states, pad_size)
+
+    # Discretize x and A
+    hidden_states = hidden_states * dt[..., None].float()
+    A = A.to(hidden_states.dtype) * dt.float()
+
+    # Rearrange into blocks/chunks
+    hidden_states, A, B, C = [reshape_into_chunks(tensor, pad_size, chunk_size) for tensor in (hidden_states, A, B, C)]
+
+    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)
+    G = (C[:, :, :, None, :, :] * B[:, :, None, :, :, :]).sum(dim=-1)
+
+    # Compute M, equivalent to applying attention mask to weights
+    M = (G[..., None] * L.permute(0, 2, 3, 4, 1)[..., None]).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 = (
+        initial_states[:, None].to(dtype=states.dtype, device=states.device)
+        if initial_states is not None
+        else torch.zeros_like(states[:, :1])
+    )
+    states = torch.cat([previous_states, states], dim=1)
+    decay_chunk = torch.exp(segment_sum(F.pad(A_cumsum[:, :, :, -1], (1, 0)))).transpose(1, 3)
+    new_states = (decay_chunk[..., None, None] * states[:, :, None, ...]).sum(dim=1)
+    states, final_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, ...]
+    Y_off = C_times_states.sum(-1) * state_decay_out.permute(0, 2, 3, 1)[..., None]
+
+    # Add output of intra-chunk and inter-chunk terms (diagonal and off-diagonal blocks)
+    output = Y_diag + Y_off
+    output = output.reshape(batch_size, -1, num_heads, head_dim)
+
+    if D_residual is not None:
+        output = output + D_residual
+
+    # Cutting off padded chunks
+    if pad_size > 0:
+        output = output[:, :sequence_length]
+
+    output = output.to(input_dtype)
+
+    if return_final_states:
+        return output, final_state
+
+    return output
+
+
+@use_kernelized_func(
+    [
+        causal_conv1d_fn,
+        causal_conv1d_update,
+        falcon_h1_split_conv1d_scan_combined,
+        falcon_h1_selective_state_update,
+        falcon_h1_chunk_scan,
+    ]
+)
 class FalconH1Mixer(nn.Module):
     """
     FalconH1Mixer is identical to classic Mamba2 mixer classes but differs on two different things
@@ -461,37 +668,6 @@ def __init__(self, config: FalconH1Config, layer_idx: int, initialize_mixer_weig
             self.init_falcon_h1_weights()
         self.out_proj = nn.Linear(self.intermediate_size, config.hidden_size, bias=config.projectors_bias)
 
-        global causal_conv1d, causal_conv1d_update, causal_conv1d_fn
-        causal_conv1d = lazy_load_kernel("causal-conv1d")
-        causal_conv1d_update = getattr(causal_conv1d, "causal_conv1d_update", causal_conv1d_update)
-        causal_conv1d_fn = getattr(causal_conv1d, "causal_conv1d_fn", causal_conv1d_fn)
-
-        global mamba_ssm, selective_state_update, mamba_chunk_scan_combined, mamba_split_conv1d_scan_combined
-        mamba_ssm = lazy_load_kernel("mamba-ssm")
-        selective_state_update = resolve_internal_import(
-            mamba_ssm, chained_path="ops.triton.selective_state_update.selective_state_update"
-        )
-        mamba_chunk_scan_combined = resolve_internal_import(
-            mamba_ssm, chained_path="ops.triton.ssd_combined.mamba_chunk_scan_combined"
-        )
-        mamba_split_conv1d_scan_combined = resolve_internal_import(
-            mamba_ssm, chained_path="ops.triton.ssd_combined.mamba_split_conv1d_scan_combined"
-        )
-
-        global is_fast_path_available
-        is_fast_path_available = (
-            all((selective_state_update, mamba_chunk_scan_combined, mamba_split_conv1d_scan_combined))
-            and hasattr(causal_conv1d, "causal_conv1d_update")
-            and hasattr(causal_conv1d, "causal_conv1d_fn")
-        )
-
-        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"
-            )
-
         self.layer_type = config.layer_types[layer_idx]
         self.groups_time_state_size = config.mamba_n_groups * self.ssm_state_size
         self.mamba_rms_norm = config.mamba_rms_norm
@@ -504,68 +680,31 @@ def init_falcon_h1_weights(self):
         init.ones_(self.D)
         init.ones_(self.dt_bias)
 
-    def _convolution(
+    @force_accelerate_hooks("conv1d")
+    def forward(
         self,
         hidden_states: torch.Tensor,
         cache_params: Cache | None = None,
-        attention_mask: torch.LongTensor | None = None,
+        attention_mask: torch.Tensor | None = None,
         **kwargs,
     ):
-        seq_len = hidden_states.shape[1]
-        hidden_states = apply_mask_to_padding_states(hidden_states, attention_mask)
-        hidden_states = hidden_states.transpose(1, 2)
-
+        batch_size, seq_len, _ = hidden_states.shape
+        dtype = hidden_states.dtype
         use_precomputed_states = cache_params is not None and cache_params.has_previous_state(self.layer_idx)
 
-        if use_precomputed_states and seq_len == 1 and not cache_params.layers[self.layer_idx].record_past:
-            conv_state = cache_params.layers[self.layer_idx].conv_states[0]
-            hidden_states = causal_conv1d_update(
-                hidden_states,
-                conv_state,
-                self.conv1d.weight.squeeze(1),
-                self.conv1d.bias,
-                self.activation,
-            )
-        else:
-            if cache_params is not None:
-                hidden_states = cache_params.update_conv_state(
-                    hidden_states, self.layer_idx, conv_kernel_size=self.conv_kernel_size
-                )
-
-            hidden_states = causal_conv1d_fn(
-                hidden_states,
-                self.conv1d.weight.squeeze(1),
-                self.conv1d.bias,
-                activation=self.activation,
-                seq_idx=kwargs.get("seq_idx"),
-            )
-
-            # Drop the additional previous states
-            if cache_params is not None:
-                hidden_states = hidden_states[:, :, -seq_len:]
-
-        hidden_states = hidden_states.transpose(1, 2)
-        return hidden_states
-
-    def cuda_kernels_forward(
-        self,
-        hidden_states: torch.Tensor,
-        cache_params: Cache | None = None,
-        attention_mask: torch.Tensor | None = None,
-        seq_idx: torch.IntTensor | None = None,
-        **kwargs,
-    ):
         # 1. Gated MLP's linear projection
         hidden_states = apply_mask_to_padding_states(hidden_states, attention_mask)
-        # Add Multipliers
+        # Key difference 1: Additional Multipliers
         hidden_states = hidden_states * self.ssm_in_multiplier
         projected_states = self.in_proj(hidden_states)
-        projected_states = projected_states * self.mup_vector  # ADD Mup Multipliers
+        projected_states = projected_states * self.mup_vector
 
-        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}
+        A = -torch.exp(self.A_log.float())
+        fused_kwargs = (
+            kwargs | {} if self.time_step_limit == (0.0, float("inf")) else kwargs | {"dt_limit": self.time_step_limit}
+        )
         if self.training and cache_params is None:
-            out = mamba_split_conv1d_scan_combined(  # noqa
+            fused_output = falcon_h1_split_conv1d_scan_combined(  # noqa F821
                 projected_states,
                 self.conv1d.weight.squeeze(1),
                 self.conv1d.bias,
@@ -573,7 +712,6 @@ def cuda_kernels_forward(
                 A,
                 D=self.D,
                 chunk_size=self.chunk_size,
-                seq_idx=seq_idx,
                 activation=self.activation,
                 rmsnorm_weight=self.norm.weight if self.mamba_rms_norm else None,
                 rmsnorm_eps=self.norm.variance_epsilon if self.mamba_rms_norm else None,
@@ -583,283 +721,125 @@ def cuda_kernels_forward(
                 ngroups=self.n_groups,
                 norm_before_gate=False,
                 return_final_states=False,
-                **dt_limit_kwargs,
+                **fused_kwargs,
             )
 
-        # Set up dimensions for reshapes later
-        batch_size, seq_len, _ = hidden_states.shape
-        groups_time_state_size = self.n_groups * self.ssm_state_size
-        use_precomputed_states = cache_params is not None and cache_params.has_previous_state(self.layer_idx)
+            # Only kernels can use this shortcircuit, fallback to normal torch otherwise
+            if fused_output is not None:
+                return fused_output
 
         gate, hidden_states_B_C, dt = projected_states.split(
             [self.intermediate_size, self.conv_dim, self.num_heads], dim=-1
         )
 
-        # Apply the conv
-        hidden_states_B_C = self._convolution(hidden_states_B_C, cache_params, attention_mask, **kwargs)
-        hidden_states_B_C = apply_mask_to_padding_states(hidden_states_B_C, attention_mask)
+        if use_precomputed_states:
+            conv_state = cache_params.layers[self.layer_idx].conv_states[0]
+            recurrent_state = cache_params.layers[self.layer_idx].recurrent_states[0]
+
+        # 2. Convolution sequence transformation
+        hidden_states_B_C = hidden_states_B_C.transpose(1, 2)
+        if use_precomputed_states and seq_len == 1:
+            hidden_states_B_C = causal_conv1d_update(
+                hidden_states_B_C,
+                conv_state,
+                self.conv1d.weight.squeeze(1),
+                self.conv1d.bias,
+                activation=self.activation,
+            )
+        else:
+            if cache_params is not None:
+                hidden_states_B_C = cache_params.update_conv_state(
+                    hidden_states_B_C,
+                    self.layer_idx,
+                    conv_kernel_size=self.conv_kernel_size,
+                )
+
+            hidden_states_B_C = causal_conv1d_fn(
+                hidden_states_B_C,
+                self.conv1d.weight.squeeze(1),
+                self.conv1d.bias,
+                activation=self.activation,
+                **kwargs,
+            )
+
+            if cache_params is not None:
+                hidden_states_B_C = hidden_states_B_C[:, :, -seq_len:]
+
+        # 3. SSM transformation
+        hidden_states_B_C = apply_mask_to_padding_states(hidden_states_B_C.transpose(1, 2), attention_mask)
         hidden_states, B, C = torch.split(
             hidden_states_B_C,
-            [self.intermediate_size, groups_time_state_size, groups_time_state_size],
+            [self.intermediate_size, self.n_groups * self.ssm_state_size, self.n_groups * self.ssm_state_size],
             dim=-1,
         )
 
-        recurrent_state = cache_params.layers[self.layer_idx].recurrent_states[0] if use_precomputed_states else None
-        # getting projected states from cache if it exists
+        # Recurrent form
         if use_precomputed_states and seq_len == 1:
-            # 3. SSM transformation
-            A = A[:, None, ...][:, :, None].expand(-1, self.head_dim, self.ssm_state_size).to(dtype=torch.float32)
+            hidden_states = hidden_states.view(batch_size, self.num_heads, self.head_dim)
             dt = dt.transpose(1, 2).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)
+            A = A[:, None, ...][:, :, None].expand(-1, self.head_dim, self.ssm_state_size).to(dtype=torch.float32)
             B = B.view(batch_size, self.n_groups, B.shape[2] // self.n_groups)
             C = C.view(batch_size, self.n_groups, C.shape[2] // self.n_groups)
-            hidden_states_reshaped = hidden_states.view(batch_size, self.num_heads, self.head_dim)
-            hidden_states = selective_state_update(  # noqa
+            D = self.D[:, None, ...].expand(-1, self.head_dim)
+            dt_bias = self.dt_bias[:, None, ...].expand(-1, self.head_dim)
+
+            scan_output = falcon_h1_selective_state_update(  # noqa F821
                 recurrent_state,
-                hidden_states_reshaped,
+                hidden_states,
                 dt,
                 A,
                 B,
                 C,
                 D,
+                # Key difference 2: Potential z gate into kernel
                 z=gate.view(batch_size, self.num_heads, self.head_dim) if not self.mamba_rms_norm else None,
                 dt_bias=dt_bias,
                 dt_softplus=True,
             )
-            hidden_states = hidden_states.view(batch_size, 1, self.num_heads * self.head_dim)
+            scan_output = scan_output.view(batch_size, 1, self.num_heads * self.head_dim)
 
+            # Key difference 3: Norm based handling as optional between z and norm
             if self.mamba_rms_norm:
-                hidden_states = self.norm(hidden_states, gate)
+                scan_output = self.norm(scan_output, gate)
 
-            # 4. Final linear projection
-            out = self.out_proj(hidden_states)
-        # Fused calculations or step by step if no initialized cache is found
+        # Chunk form
         else:
-            time_step = nn.functional.softplus(dt + self.dt_bias)
-            # This is a hack to make sure multi-GPU inference works with HF accelerate
-            # see: https://github.com/Dao-AILab/flash-attention/issues/523 for more details
-            with torch.cuda.device(hidden_states.device):
-                scan_output, ssm_state = mamba_chunk_scan_combined(  # noqa
-                    hidden_states.view(batch_size, seq_len, -1, self.head_dim),
-                    time_step,
-                    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,
-                    initial_states=recurrent_state,
-                    **dt_limit_kwargs,
-                )
-            if ssm_state is not None and cache_params is not None:
-                ssm_state = cache_params.update_recurrent_state(ssm_state, self.layer_idx)
-            scan_output = scan_output.view(batch_size, seq_len, -1)
-            # Multiply "gate" branch and apply extra normalization layer
-            if self.mamba_rms_norm:
-                out = self.norm(scan_output, gate)
-            else:
-                out = scan_output * torch.nn.functional.silu(gate)
-            out = self.out_proj(out)
-        return out
-
-    def torch_forward(
-        self,
-        input_states,
-        cache_params: Cache | None = None,
-        attention_mask: torch.Tensor | None = None,
-        **kwargs,
-    ):
-        batch_size, seq_len, _ = input_states.shape
-        dtype = input_states.dtype
-        use_precomputed_states = cache_params is not None and cache_params.has_previous_state(self.layer_idx)
-
-        # 1. Gated MLP's linear projection
-        input_states = apply_mask_to_padding_states(input_states, attention_mask)
-        # Add Multipliers
-        input_states = input_states * self.ssm_in_multiplier
-        projected_states = self.in_proj(input_states)
-        projected_states = projected_states * self.mup_vector  # ADD Mup Multipliers
-        gate, hidden_states_B_C, dt = projected_states.split(
-            [self.intermediate_size, self.conv_dim, self.num_heads], dim=-1
-        )
+            output_final_state = cache_params is not None
+            scan_result = falcon_h1_chunk_scan(  # noqa F821
+                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,
+                return_final_states=output_final_state,
+                dt_bias=self.dt_bias,
+                dt_softplus=True,
+                initial_states=recurrent_state if use_precomputed_states else None,
+                dt_limit=self.time_step_limit,
+                **kwargs,
+            )
 
-        # 2. Convolution sequence transformation
-        hidden_states_B_C = self._convolution(hidden_states_B_C, cache_params, attention_mask, **kwargs)
-        hidden_states_B_C = apply_mask_to_padding_states(hidden_states_B_C, attention_mask)
-        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,
-        )
+            if output_final_state:
+                scan_output, ssm_state = scan_result
+                cache_params.update_recurrent_state(ssm_state, self.layer_idx)
+            else:
+                scan_output = scan_result
 
-        # 3. SSM transformation
-        A = -torch.exp(self.A_log.float())  # [num_heads]
-        if use_precomputed_states and seq_len == 1:
-            # We need to guarantee that anything regarding the cache is on the same device
-            cache_device = cache_params.layers[self.layer_idx].recurrent_states[0].device
-
-            # Note: there is no need to pad parameter matrices here, as there is just one new token for batched generation
-            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))[..., None]
-            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 * A)).to(device=cache_device)
-
-            # Discretize B
-            B = B.reshape(batch_size, self.n_groups, 1, -1)
-            B = B.expand(batch_size, self.n_groups, self.num_heads // self.n_groups, B.shape[-1]).contiguous()
-            B = B.reshape(batch_size, -1, 1, B.shape[-1])
-            dB = dt * B
-
-            # Discretize x into dB
-            hidden_states = hidden_states.reshape(batch_size, -1, self.head_dim)
-            dBx = (dB * hidden_states[..., None]).to(device=cache_device)
-
-            # State calculation
-            ssm_states = cache_params.layers[self.layer_idx].recurrent_states[0] * dA + dBx
-            ssm_states = cache_params.update_recurrent_state(ssm_states, self.layer_idx)
-
-            # Subsequent output
-            # [bsz, n_groups * state_size] -> [bsz, num_heads, state_size]
-            C = C.reshape(batch_size, self.n_groups, 1, -1)
-            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 = ssm_states.to(device=C.device, dtype=C.dtype)  # Shape: [b, h, d, n]
-            # Reshape ssm_states to merge the first two dimensions
-            ssm_states_reshaped = ssm_states.view(
-                batch_size * self.num_heads, self.head_dim, self.ssm_state_size
-            )  # Shape: [b*h, d, n]
-            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, -1)
-        else:
-            # 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_interleave(self.num_heads // self.n_groups, dim=2, output_size=self.num_heads)
-            C = C.repeat_interleave(self.num_heads // self.n_groups, dim=2, output_size=self.num_heads)
-            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)
-            G_intermediate = C[:, :, :, None, :, :] * B[:, :, None, :, :, :]  # shape: (b, c, l, s, h, n)
-            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 = (
-                cache_params.layers[self.layer_idx]
-                .recurrent_states[0][:, None]
-                .to(dtype=states.dtype, device=states.device)
-                if use_precomputed_states
-                else 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)
-
-            # Init cache
-            if ssm_state is not None and cache_params is not None:
-                ssm_state = cache_params.update_recurrent_state(ssm_state, self.layer_idx)
-
-        if self.mamba_rms_norm:
-            scan_output = self.norm(y, gate)
-        else:
-            scan_output = y * torch.nn.functional.silu(gate)
+            scan_output = scan_output.view(batch_size, seq_len, -1)
 
-        # end ssd naive
+            # Key difference 3: Norm based handling as optional between z and norm
+            if self.mamba_rms_norm:
+                scan_output = self.norm(scan_output, gate)
+            else:
+                scan_output = scan_output * torch.nn.functional.silu(gate)
 
         # 4. Final linear projection
-        contextualized_states = self.out_proj(scan_output.to(dtype))  # [batch, seq_len, hidden_size]
+        contextualized_states = self.out_proj(scan_output.to(dtype))
         return contextualized_states
 
-    @force_accelerate_hooks("conv1d")
-    def forward(
-        self,
-        hidden_states,
-        cache_params: Cache | None = None,
-        attention_mask: torch.Tensor | None = None,
-        **kwargs,
-    ):
-        if is_fast_path_available and "cuda" in self.in_proj.weight.device.type and not is_torchdynamo_compiling():
-            return self.cuda_kernels_forward(hidden_states, cache_params, attention_mask, **kwargs)
-        if kwargs.get("seq_idx") is not None:
-            raise NotImplementedError(
-                "`seq_idx` support requires fast path support. Please install `mamba_ssm` and `causal_conv1d`"
-            )
-        return self.torch_forward(hidden_states, cache_params, attention_mask, **kwargs)
-
 
 class FalconH1MLP(nn.Module):
     def __init__(self, config: FalconH1Config):
diff --git a/src/transformers/models/falcon_h1/modular_falcon_h1.py b/src/transformers/models/falcon_h1/modular_falcon_h1.py
index 216e1dfb2e14..5214e68780aa 100644
--- a/src/transformers/models/falcon_h1/modular_falcon_h1.py
+++ b/src/transformers/models/falcon_h1/modular_falcon_h1.py
@@ -26,6 +26,7 @@
 
 from ... import initialization as init
 from ...cache_utils import Cache, DynamicCache
+from ...integrations.accelerate import force_accelerate_hooks
 from ...masking_utils import create_causal_mask, create_recurrent_attention_mask
 from ...modeling_flash_attention_utils import FlashAttentionKwargs
 from ...modeling_layers import GradientCheckpointingLayer
@@ -48,9 +49,8 @@
 from ..mamba2.modeling_mamba2 import (
     MambaRMSNormGated,
     apply_mask_to_padding_states,
-    pad_tensor_by_size,
-    reshape_into_chunks,
-    segment_sum,
+    causal_conv1d_fn,
+    causal_conv1d_update,
 )
 from .configuration_falcon_h1 import FalconH1Config
 
@@ -172,25 +172,31 @@ def __init__(self, config: FalconH1Config, layer_idx: int, initialize_mixer_weig
         self.out_proj = nn.Linear(self.intermediate_size, config.hidden_size, bias=config.projectors_bias)
         self.ssm_in_multiplier = config.ssm_in_multiplier
 
-    def cuda_kernels_forward(
+    @force_accelerate_hooks("conv1d")
+    def forward(
         self,
         hidden_states: torch.Tensor,
         cache_params: Cache | None = None,
         attention_mask: torch.Tensor | None = None,
-        seq_idx: torch.IntTensor | None = None,
         **kwargs,
     ):
+        batch_size, seq_len, _ = hidden_states.shape
+        dtype = hidden_states.dtype
+        use_precomputed_states = cache_params is not None and cache_params.has_previous_state(self.layer_idx)
+
         # 1. Gated MLP's linear projection
         hidden_states = apply_mask_to_padding_states(hidden_states, attention_mask)
-        # Add Multipliers
+        # Key difference 1: Additional Multipliers
         hidden_states = hidden_states * self.ssm_in_multiplier
         projected_states = self.in_proj(hidden_states)
-        projected_states = projected_states * self.mup_vector  # ADD Mup Multipliers
+        projected_states = projected_states * self.mup_vector
 
-        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}
+        A = -torch.exp(self.A_log.float())
+        fused_kwargs = (
+            kwargs | {} if self.time_step_limit == (0.0, float("inf")) else kwargs | {"dt_limit": self.time_step_limit}
+        )
         if self.training and cache_params is None:
-            out = mamba_split_conv1d_scan_combined(  # noqa
+            fused_output = falcon_h1_split_conv1d_scan_combined(  # noqa F821
                 projected_states,
                 self.conv1d.weight.squeeze(1),
                 self.conv1d.bias,
@@ -198,7 +204,6 @@ def cuda_kernels_forward(
                 A,
                 D=self.D,
                 chunk_size=self.chunk_size,
-                seq_idx=seq_idx,
                 activation=self.activation,
                 rmsnorm_weight=self.norm.weight if self.mamba_rms_norm else None,
                 rmsnorm_eps=self.norm.variance_epsilon if self.mamba_rms_norm else None,
@@ -208,265 +213,123 @@ def cuda_kernels_forward(
                 ngroups=self.n_groups,
                 norm_before_gate=False,
                 return_final_states=False,
-                **dt_limit_kwargs,
+                **fused_kwargs,
             )
 
-        # Set up dimensions for reshapes later
-        batch_size, seq_len, _ = hidden_states.shape
-        groups_time_state_size = self.n_groups * self.ssm_state_size
-        use_precomputed_states = cache_params is not None and cache_params.has_previous_state(self.layer_idx)
+            # Only kernels can use this shortcircuit, fallback to normal torch otherwise
+            if fused_output is not None:
+                return fused_output
 
         gate, hidden_states_B_C, dt = projected_states.split(
             [self.intermediate_size, self.conv_dim, self.num_heads], dim=-1
         )
 
-        # Apply the conv
-        hidden_states_B_C = self._convolution(hidden_states_B_C, cache_params, attention_mask, **kwargs)
-        hidden_states_B_C = apply_mask_to_padding_states(hidden_states_B_C, attention_mask)
+        if use_precomputed_states:
+            conv_state = cache_params.layers[self.layer_idx].conv_states[0]
+            recurrent_state = cache_params.layers[self.layer_idx].recurrent_states[0]
+
+        # 2. Convolution sequence transformation
+        hidden_states_B_C = hidden_states_B_C.transpose(1, 2)
+        if use_precomputed_states and seq_len == 1:
+            hidden_states_B_C = causal_conv1d_update(
+                hidden_states_B_C,
+                conv_state,
+                self.conv1d.weight.squeeze(1),
+                self.conv1d.bias,
+                activation=self.activation,
+            )
+        else:
+            if cache_params is not None:
+                hidden_states_B_C = cache_params.update_conv_state(
+                    hidden_states_B_C,
+                    self.layer_idx,
+                    conv_kernel_size=self.conv_kernel_size,
+                )
+
+            hidden_states_B_C = causal_conv1d_fn(
+                hidden_states_B_C,
+                self.conv1d.weight.squeeze(1),
+                self.conv1d.bias,
+                activation=self.activation,
+                **kwargs,
+            )
+
+            if cache_params is not None:
+                hidden_states_B_C = hidden_states_B_C[:, :, -seq_len:]
+
+        # 3. SSM transformation
+        hidden_states_B_C = apply_mask_to_padding_states(hidden_states_B_C.transpose(1, 2), attention_mask)
         hidden_states, B, C = torch.split(
             hidden_states_B_C,
-            [self.intermediate_size, groups_time_state_size, groups_time_state_size],
+            [self.intermediate_size, self.n_groups * self.ssm_state_size, self.n_groups * self.ssm_state_size],
             dim=-1,
         )
 
-        recurrent_state = cache_params.layers[self.layer_idx].recurrent_states[0] if use_precomputed_states else None
-        # getting projected states from cache if it exists
+        # Recurrent form
         if use_precomputed_states and seq_len == 1:
-            # 3. SSM transformation
-            A = A[:, None, ...][:, :, None].expand(-1, self.head_dim, self.ssm_state_size).to(dtype=torch.float32)
+            hidden_states = hidden_states.view(batch_size, self.num_heads, self.head_dim)
             dt = dt.transpose(1, 2).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)
+            A = A[:, None, ...][:, :, None].expand(-1, self.head_dim, self.ssm_state_size).to(dtype=torch.float32)
             B = B.view(batch_size, self.n_groups, B.shape[2] // self.n_groups)
             C = C.view(batch_size, self.n_groups, C.shape[2] // self.n_groups)
-            hidden_states_reshaped = hidden_states.view(batch_size, self.num_heads, self.head_dim)
-            hidden_states = selective_state_update(  # noqa
+            D = self.D[:, None, ...].expand(-1, self.head_dim)
+            dt_bias = self.dt_bias[:, None, ...].expand(-1, self.head_dim)
+
+            scan_output = falcon_h1_selective_state_update(  # noqa F821
                 recurrent_state,
-                hidden_states_reshaped,
+                hidden_states,
                 dt,
                 A,
                 B,
                 C,
                 D,
+                # Key difference 2: Potential z gate into kernel
                 z=gate.view(batch_size, self.num_heads, self.head_dim) if not self.mamba_rms_norm else None,
                 dt_bias=dt_bias,
                 dt_softplus=True,
             )
-            hidden_states = hidden_states.view(batch_size, 1, self.num_heads * self.head_dim)
+            scan_output = scan_output.view(batch_size, 1, self.num_heads * self.head_dim)
 
+            # Key difference 3: Norm based handling as optional between z and norm
             if self.mamba_rms_norm:
-                hidden_states = self.norm(hidden_states, gate)
+                scan_output = self.norm(scan_output, gate)
 
-            # 4. Final linear projection
-            out = self.out_proj(hidden_states)
-        # Fused calculations or step by step if no initialized cache is found
+        # Chunk form
         else:
-            time_step = nn.functional.softplus(dt + self.dt_bias)
-            # This is a hack to make sure multi-GPU inference works with HF accelerate
-            # see: https://github.com/Dao-AILab/flash-attention/issues/523 for more details
-            with torch.cuda.device(hidden_states.device):
-                scan_output, ssm_state = mamba_chunk_scan_combined(  # noqa
-                    hidden_states.view(batch_size, seq_len, -1, self.head_dim),
-                    time_step,
-                    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,
-                    initial_states=recurrent_state,
-                    **dt_limit_kwargs,
-                )
-            if ssm_state is not None and cache_params is not None:
-                ssm_state = cache_params.update_recurrent_state(ssm_state, self.layer_idx)
-            scan_output = scan_output.view(batch_size, seq_len, -1)
-            # Multiply "gate" branch and apply extra normalization layer
-            if self.mamba_rms_norm:
-                out = self.norm(scan_output, gate)
-            else:
-                out = scan_output * torch.nn.functional.silu(gate)
-            out = self.out_proj(out)
-        return out
-
-    def torch_forward(
-        self,
-        input_states,
-        cache_params: Cache | None = None,
-        attention_mask: torch.Tensor | None = None,
-        **kwargs,
-    ):
-        batch_size, seq_len, _ = input_states.shape
-        dtype = input_states.dtype
-        use_precomputed_states = cache_params is not None and cache_params.has_previous_state(self.layer_idx)
-
-        # 1. Gated MLP's linear projection
-        input_states = apply_mask_to_padding_states(input_states, attention_mask)
-        # Add Multipliers
-        input_states = input_states * self.ssm_in_multiplier
-        projected_states = self.in_proj(input_states)
-        projected_states = projected_states * self.mup_vector  # ADD Mup Multipliers
-        gate, hidden_states_B_C, dt = projected_states.split(
-            [self.intermediate_size, self.conv_dim, self.num_heads], dim=-1
-        )
+            output_final_state = cache_params is not None
+            scan_result = falcon_h1_chunk_scan(  # noqa F821
+                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,
+                return_final_states=output_final_state,
+                dt_bias=self.dt_bias,
+                dt_softplus=True,
+                initial_states=recurrent_state if use_precomputed_states else None,
+                dt_limit=self.time_step_limit,
+                **kwargs,
+            )
 
-        # 2. Convolution sequence transformation
-        hidden_states_B_C = self._convolution(hidden_states_B_C, cache_params, attention_mask, **kwargs)
-        hidden_states_B_C = apply_mask_to_padding_states(hidden_states_B_C, attention_mask)
-        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,
-        )
+            if output_final_state:
+                scan_output, ssm_state = scan_result
+                cache_params.update_recurrent_state(ssm_state, self.layer_idx)
+            else:
+                scan_output = scan_result
 
-        # 3. SSM transformation
-        A = -torch.exp(self.A_log.float())  # [num_heads]
-        if use_precomputed_states and seq_len == 1:
-            # We need to guarantee that anything regarding the cache is on the same device
-            cache_device = cache_params.layers[self.layer_idx].recurrent_states[0].device
-
-            # Note: there is no need to pad parameter matrices here, as there is just one new token for batched generation
-            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))[..., None]
-            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 * A)).to(device=cache_device)
-
-            # Discretize B
-            B = B.reshape(batch_size, self.n_groups, 1, -1)
-            B = B.expand(batch_size, self.n_groups, self.num_heads // self.n_groups, B.shape[-1]).contiguous()
-            B = B.reshape(batch_size, -1, 1, B.shape[-1])
-            dB = dt * B
-
-            # Discretize x into dB
-            hidden_states = hidden_states.reshape(batch_size, -1, self.head_dim)
-            dBx = (dB * hidden_states[..., None]).to(device=cache_device)
-
-            # State calculation
-            ssm_states = cache_params.layers[self.layer_idx].recurrent_states[0] * dA + dBx
-            ssm_states = cache_params.update_recurrent_state(ssm_states, self.layer_idx)
-
-            # Subsequent output
-            # [bsz, n_groups * state_size] -> [bsz, num_heads, state_size]
-            C = C.reshape(batch_size, self.n_groups, 1, -1)
-            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 = ssm_states.to(device=C.device, dtype=C.dtype)  # Shape: [b, h, d, n]
-            # Reshape ssm_states to merge the first two dimensions
-            ssm_states_reshaped = ssm_states.view(
-                batch_size * self.num_heads, self.head_dim, self.ssm_state_size
-            )  # Shape: [b*h, d, n]
-            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, -1)
-        else:
-            # 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_interleave(self.num_heads // self.n_groups, dim=2, output_size=self.num_heads)
-            C = C.repeat_interleave(self.num_heads // self.n_groups, dim=2, output_size=self.num_heads)
-            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)
-            G_intermediate = C[:, :, :, None, :, :] * B[:, :, None, :, :, :]  # shape: (b, c, l, s, h, n)
-            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 = (
-                cache_params.layers[self.layer_idx]
-                .recurrent_states[0][:, None]
-                .to(dtype=states.dtype, device=states.device)
-                if use_precomputed_states
-                else 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)
-
-            # Init cache
-            if ssm_state is not None and cache_params is not None:
-                ssm_state = cache_params.update_recurrent_state(ssm_state, self.layer_idx)
-
-        if self.mamba_rms_norm:
-            scan_output = self.norm(y, gate)
-        else:
-            scan_output = y * torch.nn.functional.silu(gate)
+            scan_output = scan_output.view(batch_size, seq_len, -1)
 
-        # end ssd naive
+            # Key difference 3: Norm based handling as optional between z and norm
+            if self.mamba_rms_norm:
+                scan_output = self.norm(scan_output, gate)
+            else:
+                scan_output = scan_output * torch.nn.functional.silu(gate)
 
         # 4. Final linear projection
-        contextualized_states = self.out_proj(scan_output.to(dtype))  # [batch, seq_len, hidden_size]
+        contextualized_states = self.out_proj(scan_output.to(dtype))
         return contextualized_states
 
 
diff --git a/src/transformers/models/granitemoehybrid/modeling_granitemoehybrid.py b/src/transformers/models/granitemoehybrid/modeling_granitemoehybrid.py
index 723bcd0b112a..7fb2590cd51f 100644
--- a/src/transformers/models/granitemoehybrid/modeling_granitemoehybrid.py
+++ b/src/transformers/models/granitemoehybrid/modeling_granitemoehybrid.py
@@ -30,9 +30,9 @@
 from ...cache_utils import Cache, DynamicCache
 from ...generation import GenerationMixin
 from ...integrations import (
-    lazy_load_kernel,
     use_experts_implementation,
     use_kernel_forward_from_hub,
+    use_kernel_func_from_hub_with_fallback,
     use_kernelized_func,
 )
 from ...integrations.accelerate import force_accelerate_hooks
@@ -42,17 +42,13 @@
 from ...modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update
 from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
 from ...processing_utils import Unpack
-from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, is_torchdynamo_compiling, logging
+from ...utils import TransformersKwargs, auto_docstring, can_return_tuple
 from ...utils.deprecation import deprecate_kwarg
 from ...utils.generic import maybe_autocast, merge_with_config_defaults
-from ...utils.import_utils import resolve_internal_import
 from ...utils.output_capturing import capture_outputs
 from .configuration_granitemoehybrid import GraniteMoeHybridConfig
 
 
-logger = logging.get_logger(__name__)
-
-
 def rotate_half(x):
     """Rotates half the hidden dims of the input."""
     x1 = x[..., : x.shape[-1] // 2]
@@ -258,6 +254,7 @@ def apply_mask_to_padding_states(hidden_states, attention_mask):
     return hidden_states
 
 
+@use_kernel_func_from_hub_with_fallback("causal_conv1d_update", "causal_conv1d")
 def causal_conv1d_update(
     hidden_states: torch.Tensor,
     conv_state: torch.Tensor,
@@ -277,6 +274,7 @@ def causal_conv1d_update(
     return out.to(hidden_states.dtype)
 
 
+@use_kernel_func_from_hub_with_fallback("causal_conv1d_fn", "causal_conv1d")
 def causal_conv1d_fn(
     hidden_states: torch.Tensor,
     weight: nn.Parameter,
@@ -299,6 +297,215 @@ def causal_conv1d_fn(
     return out.to(hidden_states.dtype)
 
 
+@use_kernel_func_from_hub_with_fallback(
+    "mamba_split_conv1d_scan_combined",
+    "mamba_ssm",
+    internal_path="ops.triton.ssd_combined",
+)
+def granitemoehybrid_split_conv1d_scan_combined(
+    zxbcdt: torch.Tensor,
+    conv1d_weight: torch.Tensor,
+    conv1d_bias: torch.Tensor | None,
+    dt_bias: torch.Tensor,
+    A: torch.Tensor,
+    D: torch.Tensor,
+    chunk_size: int,
+    initial_states: 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,
+    **kwargs,
+):
+    return None
+
+
+@use_kernel_func_from_hub_with_fallback(
+    "selective_state_update",
+    "mamba_ssm",
+    internal_path="ops.triton.selective_state_update",
+)
+def granitemoehybrid_selective_state_update(
+    state: torch.Tensor,
+    hidden_states: torch.Tensor,
+    dt: torch.Tensor,
+    A: torch.Tensor,
+    B: torch.Tensor,
+    C: torch.Tensor,
+    D: torch.Tensor | None = None,
+    dt_bias: torch.Tensor | None = None,
+    dt_softplus: bool = False,
+    z: torch.Tensor | None = None,
+    **kwargs,
+):
+    batch_size, num_heads, head_dim = hidden_states.shape
+    num_groups = B.shape[1]
+    state_size = B.shape[-1]
+
+    if dt_bias is not None:
+        dt = dt + dt_bias.to(dt.dtype)
+    if dt_softplus:
+        dt = F.softplus(dt)
+    dt = dt[..., None]
+
+    # Discretize A
+    dA = torch.exp(dt.float() * A.float()).to(device=state.device)
+
+    # Discretize B
+    B = B.reshape(batch_size, num_groups, 1, state_size)
+    B = B.expand(batch_size, num_groups, num_heads // num_groups, state_size).contiguous()
+    B = B.reshape(batch_size, num_heads, 1, state_size)
+    dB = dt * B
+
+    # Discretize x into dB
+    dBx = (dB * hidden_states[..., None]).to(device=state.device)
+
+    # State calculation
+    ssm_states = state * dA + dBx
+    state.copy_(ssm_states.to(state.dtype))
+
+    # Subsequent output
+    C = C.reshape(batch_size, num_groups, 1, state_size)
+    C = C.expand(batch_size, num_groups, num_heads // num_groups, state_size).contiguous()
+    C = C.reshape(batch_size, num_heads, state_size)
+
+    # Reshape ssm_states to merge the first two dimensions
+    ssm_states = ssm_states.to(device=C.device, dtype=C.dtype)
+    ssm_states_reshaped = ssm_states.view(batch_size * num_heads, head_dim, state_size)
+    C_reshaped = C.view(batch_size * num_heads, state_size, 1)
+    out = torch.bmm(ssm_states_reshaped, C_reshaped)
+    out = out.view(batch_size, num_heads, head_dim)
+
+    # D skip connection
+    if D is not None:
+        out = (out + hidden_states * D).to(out.dtype)
+
+    if z is not None:
+        out = out * F.silu(z)
+
+    return out.to(hidden_states.dtype)
+
+
+@use_kernel_func_from_hub_with_fallback(
+    "mamba_chunk_scan_combined",
+    "mamba_ssm",
+    internal_path="ops.triton.ssd_combined",
+)
+def granitemoehybrid_chunk_scan(
+    hidden_states: torch.Tensor,
+    dt: torch.Tensor,
+    A: torch.Tensor,
+    B: torch.Tensor,
+    C: torch.Tensor,
+    chunk_size: int,
+    D: torch.Tensor | None = None,
+    dt_bias: torch.Tensor | None = None,
+    initial_states: torch.Tensor | None = None,
+    dt_softplus: bool = False,
+    dt_limit: tuple[float, float] = (0.0, float("inf")),
+    return_final_states: bool = False,
+    **kwargs,
+):
+    input_dtype = hidden_states.dtype
+    batch_size, sequence_length, num_heads, head_dim = hidden_states.shape
+    num_groups = B.shape[2]
+
+    if dt_bias is not None:
+        dt = dt + dt_bias.to(dt.dtype)
+    if dt_softplus:
+        dt = F.softplus(dt)
+    dt = torch.clamp(dt, min=dt_limit[0], max=dt_limit[1])
+
+    hidden_states = hidden_states.float()
+    B = B.float().repeat_interleave(num_heads // num_groups, dim=2, output_size=num_heads)
+    C = C.float().repeat_interleave(num_heads // num_groups, dim=2, output_size=num_heads)
+
+    pad_size = (chunk_size - sequence_length % chunk_size) % chunk_size
+    D_residual = None
+    if D is not None:
+        D_residual = D[..., None] * pad_tensor_by_size(hidden_states, pad_size)
+
+    # Discretize x and A
+    hidden_states = hidden_states * dt[..., None].float()
+    A = A.to(hidden_states.dtype) * dt.float()
+
+    # Rearrange into blocks/chunks
+    hidden_states, A, B, C = [reshape_into_chunks(tensor, pad_size, chunk_size) for tensor in (hidden_states, A, B, C)]
+
+    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)
+    G = (C[:, :, :, None, :, :] * B[:, :, None, :, :, :]).sum(dim=-1)
+
+    # Compute M, equivalent to applying attention mask to weights
+    M = (G[..., None] * L.permute(0, 2, 3, 4, 1)[..., None]).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 = (
+        initial_states[:, None].to(dtype=states.dtype, device=states.device)
+        if initial_states is not None
+        else torch.zeros_like(states[:, :1])
+    )
+    states = torch.cat([previous_states, states], dim=1)
+    decay_chunk = torch.exp(segment_sum(F.pad(A_cumsum[:, :, :, -1], (1, 0)))).transpose(1, 3)
+    new_states = (decay_chunk[..., None, None] * states[:, :, None, ...]).sum(dim=1)
+    states, final_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, ...]
+    Y_off = C_times_states.sum(-1) * state_decay_out.permute(0, 2, 3, 1)[..., None]
+
+    # Add output of intra-chunk and inter-chunk terms (diagonal and off-diagonal blocks)
+    output = Y_diag + Y_off
+    output = output.reshape(batch_size, -1, num_heads, head_dim)
+
+    if D_residual is not None:
+        output = output + D_residual
+
+    # Cutting off padded chunks
+    if pad_size > 0:
+        output = output[:, :sequence_length]
+
+    output = output.to(input_dtype)
+
+    if return_final_states:
+        return output, final_state
+
+    return output
+
+
+@use_kernelized_func(
+    [
+        causal_conv1d_fn,
+        causal_conv1d_update,
+        granitemoehybrid_split_conv1d_scan_combined,
+        granitemoehybrid_selective_state_update,
+        granitemoehybrid_chunk_scan,
+    ]
+)
 class GraniteMoeHybridMambaLayer(nn.Module):
     """
     Compute ∆, A, B, C, and D the state space parameters and compute the `contextualized_states`.
@@ -359,37 +566,6 @@ def __init__(self, config: GraniteMoeHybridConfig, layer_idx: int, initialize_mi
             self.init_granitemoehybrid_weights()
         self.out_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=self.use_bias)
 
-        global causal_conv1d, causal_conv1d_update, causal_conv1d_fn
-        causal_conv1d = lazy_load_kernel("causal-conv1d")
-        causal_conv1d_update = getattr(causal_conv1d, "causal_conv1d_update", causal_conv1d_update)
-        causal_conv1d_fn = getattr(causal_conv1d, "causal_conv1d_fn", causal_conv1d_fn)
-
-        global mamba_ssm, selective_state_update, mamba_chunk_scan_combined, mamba_split_conv1d_scan_combined
-        mamba_ssm = lazy_load_kernel("mamba-ssm")
-        selective_state_update = resolve_internal_import(
-            mamba_ssm, chained_path="ops.triton.selective_state_update.selective_state_update"
-        )
-        mamba_chunk_scan_combined = resolve_internal_import(
-            mamba_ssm, chained_path="ops.triton.ssd_combined.mamba_chunk_scan_combined"
-        )
-        mamba_split_conv1d_scan_combined = resolve_internal_import(
-            mamba_ssm, chained_path="ops.triton.ssd_combined.mamba_split_conv1d_scan_combined"
-        )
-
-        global is_fast_path_available
-        is_fast_path_available = (
-            all((selective_state_update, mamba_chunk_scan_combined, mamba_split_conv1d_scan_combined))
-            and hasattr(causal_conv1d, "causal_conv1d_update")
-            and hasattr(causal_conv1d, "causal_conv1d_fn")
-        )
-
-        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"
-            )
-
         self.layer_type = config.layer_types[layer_idx]
 
     @torch.no_grad()
@@ -399,65 +575,28 @@ def init_granitemoehybrid_weights(self):
         init.ones_(self.D)
         init.ones_(self.dt_bias)
 
-    def _convolution(
+    @force_accelerate_hooks("conv1d")
+    def forward(
         self,
         hidden_states: torch.Tensor,
         cache_params: Cache | None = None,
-        attention_mask: torch.LongTensor | None = None,
+        attention_mask: torch.Tensor | None = None,
         **kwargs,
     ):
-        seq_len = hidden_states.shape[1]
-        hidden_states = apply_mask_to_padding_states(hidden_states, attention_mask)
-        hidden_states = hidden_states.transpose(1, 2)
-
+        batch_size, seq_len, _ = hidden_states.shape
+        dtype = hidden_states.dtype
         use_precomputed_states = cache_params is not None and cache_params.has_previous_state(self.layer_idx)
 
-        if use_precomputed_states and seq_len == 1 and not cache_params.layers[self.layer_idx].record_past:
-            conv_state = cache_params.layers[self.layer_idx].conv_states[0]
-            hidden_states = causal_conv1d_update(
-                hidden_states,
-                conv_state,
-                self.conv1d.weight.squeeze(1),
-                self.conv1d.bias,
-                self.activation,
-            )
-        else:
-            if cache_params is not None:
-                hidden_states = cache_params.update_conv_state(
-                    hidden_states, self.layer_idx, conv_kernel_size=self.conv_kernel_size
-                )
-
-            hidden_states = causal_conv1d_fn(
-                hidden_states,
-                self.conv1d.weight.squeeze(1),
-                self.conv1d.bias,
-                activation=self.activation,
-                seq_idx=kwargs.get("seq_idx"),
-            )
-
-            # Drop the additional previous states
-            if cache_params is not None:
-                hidden_states = hidden_states[:, :, -seq_len:]
-
-        hidden_states = hidden_states.transpose(1, 2)
-        return hidden_states
-
-    def cuda_kernels_forward(
-        self,
-        hidden_states: torch.Tensor,
-        cache_params: Cache | None = None,
-        attention_mask: torch.Tensor | None = None,
-        **kwargs,
-    ):
         # 1. Gated MLP's linear projection
         hidden_states = apply_mask_to_padding_states(hidden_states, attention_mask)
         projected_states = self.in_proj(hidden_states)
 
         A = -torch.exp(self.A_log.float())
-        dt_limit_kwargs = {} if self.time_step_limit == (0.0, float("inf")) else {"dt_limit": self.time_step_limit}
-        # Fused kernel for conv1d, SSM, and the final projection
+        fused_kwargs = (
+            kwargs | {} if self.time_step_limit == (0.0, float("inf")) else kwargs | {"dt_limit": self.time_step_limit}
+        )
         if self.training and cache_params is None:
-            return mamba_split_conv1d_scan_combined(
+            fused_output = granitemoehybrid_split_conv1d_scan_combined(
                 projected_states,
                 self.conv1d.weight.squeeze(1),
                 self.conv1d.bias,
@@ -465,7 +604,6 @@ def cuda_kernels_forward(
                 A,
                 D=self.D,
                 chunk_size=self.chunk_size,
-                seq_idx=kwargs.get("seq_idx"),
                 activation=self.activation,
                 rmsnorm_weight=self.norm.weight,
                 rmsnorm_eps=self.norm.variance_epsilon,
@@ -475,41 +613,71 @@ def cuda_kernels_forward(
                 ngroups=self.n_groups,
                 norm_before_gate=False,
                 return_final_states=False,
-                **dt_limit_kwargs,
+                **fused_kwargs,
             )
 
-        # Set up dimensions for reshapes later
-        batch_size, seq_len, _ = hidden_states.shape
-        groups_time_state_size = self.n_groups * self.ssm_state_size
-        use_precomputed_states = cache_params is not None and cache_params.has_previous_state(self.layer_idx)
+            # Only kernels can use this shortcircuit, fallback to normal torch otherwise
+            if fused_output is not None:
+                return fused_output
 
         gate, hidden_states_B_C, dt = projected_states.split(
             [self.intermediate_size, self.conv_dim, self.num_heads], dim=-1
         )
 
-        # Apply the conv
-        hidden_states_B_C = self._convolution(hidden_states_B_C, cache_params, attention_mask, **kwargs)
-        hidden_states_B_C = apply_mask_to_padding_states(hidden_states_B_C, attention_mask)
+        if use_precomputed_states:
+            conv_state = cache_params.layers[self.layer_idx].conv_states[0]
+            recurrent_state = cache_params.layers[self.layer_idx].recurrent_states[0]
+
+        # 2. Convolution sequence transformation
+        hidden_states_B_C = hidden_states_B_C.transpose(1, 2)
+        if use_precomputed_states and seq_len == 1:
+            hidden_states_B_C = causal_conv1d_update(
+                hidden_states_B_C,
+                conv_state,
+                self.conv1d.weight.squeeze(1),
+                self.conv1d.bias,
+                activation=self.activation,
+            )
+        else:
+            if cache_params is not None:
+                hidden_states_B_C = cache_params.update_conv_state(
+                    hidden_states_B_C,
+                    self.layer_idx,
+                    conv_kernel_size=self.conv_kernel_size,
+                )
+
+            hidden_states_B_C = causal_conv1d_fn(
+                hidden_states_B_C,
+                self.conv1d.weight.squeeze(1),
+                self.conv1d.bias,
+                activation=self.activation,
+                **kwargs,
+            )
+
+            if cache_params is not None:
+                hidden_states_B_C = hidden_states_B_C[:, :, -seq_len:]
+
+        # 3. SSM transformation
+        hidden_states_B_C = apply_mask_to_padding_states(hidden_states_B_C.transpose(1, 2), attention_mask)
         hidden_states, B, C = torch.split(
             hidden_states_B_C,
-            [self.intermediate_size, groups_time_state_size, groups_time_state_size],
+            [self.intermediate_size, self.n_groups * self.ssm_state_size, self.n_groups * self.ssm_state_size],
             dim=-1,
         )
 
-        recurrent_state = cache_params.layers[self.layer_idx].recurrent_states[0] if use_precomputed_states else None
-        # Single step calculations via cache
+        # Recurrent form
         if use_precomputed_states and seq_len == 1:
-            # 3. SSM transformation
-            A = A[:, None, ...][:, :, None].expand(-1, self.head_dim, self.ssm_state_size).to(dtype=torch.float32)
+            hidden_states = hidden_states.view(batch_size, self.num_heads, self.head_dim)
             dt = dt.transpose(1, 2).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[2] // self.n_groups)
-            C = C.view(batch_size, self.n_groups, C.shape[2] // self.n_groups)
-            hidden_states_reshaped = hidden_states.view(batch_size, self.num_heads, self.head_dim)
-            hidden_states = selective_state_update(
+            A = A[:, None, None].expand(-1, self.head_dim, self.ssm_state_size)
+            B = B.view(batch_size, self.n_groups, self.ssm_state_size)
+            C = C.view(batch_size, self.n_groups, self.ssm_state_size)
+            D = self.D[:, None].expand(-1, self.head_dim)
+            dt_bias = self.dt_bias[:, None].expand(-1, self.head_dim)
+
+            scan_output = granitemoehybrid_selective_state_update(
                 recurrent_state,
-                hidden_states_reshaped,
+                hidden_states,
                 dt,
                 A,
                 B,
@@ -518,224 +686,44 @@ def cuda_kernels_forward(
                 z=None,
                 dt_bias=dt_bias,
                 dt_softplus=True,
+                **kwargs,
             )
-            hidden_states = hidden_states.view(batch_size, 1, self.num_heads * self.head_dim)
-            hidden_states = self.norm(hidden_states, gate)
+            scan_output = scan_output.view(batch_size, 1, -1)
 
-            # 4. Final linear projection
-            out = self.out_proj(hidden_states)
-
-        # Fused calculations or step by step if no initialized cache is found
+        # Chunk form
         else:
-            # 3. SSM transformation
-            scan_output, ssm_state = mamba_chunk_scan_combined(
-                hidden_states.view(batch_size, seq_len, -1, self.head_dim),
+            output_final_state = cache_params is not None
+            scan_result = granitemoehybrid_chunk_scan(
+                hidden_states.view(batch_size, seq_len, self.num_heads, 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),
+                B.view(batch_size, seq_len, self.n_groups, self.ssm_state_size),
+                C.view(batch_size, seq_len, self.n_groups, self.ssm_state_size),
                 chunk_size=self.chunk_size,
                 D=self.D,
                 z=None,
-                seq_idx=kwargs.get("seq_idx"),
-                return_final_states=True,
+                return_final_states=output_final_state,
                 dt_bias=self.dt_bias,
                 dt_softplus=True,
-                initial_states=recurrent_state,
-                **dt_limit_kwargs,
+                initial_states=recurrent_state if use_precomputed_states else None,
+                dt_limit=self.time_step_limit,
+                **kwargs,
             )
 
-            # Init cache
-            if ssm_state is not None and cache_params is not None:
-                cache_params.update_recurrent_state(ssm_state, layer_idx=self.layer_idx)
+            if output_final_state:
+                scan_output, final_state = scan_result
+                cache_params.update_recurrent_state(final_state, layer_idx=self.layer_idx)
+            else:
+                scan_output = scan_result
 
-            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)
+            scan_output = scan_output.reshape(batch_size, seq_len, -1)
 
-            # 4. Final linear projection
-            out = self.out_proj(scan_output)
-
-        return out
-
-    def torch_forward(
-        self,
-        hidden_states: torch.Tensor,
-        cache_params: Cache | None = None,
-        attention_mask: torch.Tensor | None = None,
-        **kwargs,
-    ):
-        batch_size, seq_len, _ = hidden_states.shape
-        dtype = hidden_states.dtype
-        use_precomputed_states = cache_params is not None and cache_params.has_previous_state(self.layer_idx)
-
-        # 1. Gated MLP's linear projection
-        hidden_states = apply_mask_to_padding_states(hidden_states, attention_mask)
-        projected_states = self.in_proj(hidden_states)
-
-        gate, hidden_states_B_C, dt = projected_states.split(
-            [self.intermediate_size, self.conv_dim, self.num_heads], dim=-1
-        )
-
-        # 2. Convolution sequence transformation
-        hidden_states_B_C = self._convolution(hidden_states_B_C, cache_params, attention_mask, **kwargs)
-        hidden_states_B_C = apply_mask_to_padding_states(hidden_states_B_C, attention_mask)
-        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())
-        if use_precomputed_states and seq_len == 1:
-            # We need to guarantee that anything regarding the cache is on the same device
-            cache_device = cache_params.layers[self.layer_idx].device
-
-            # Note: there is no need to pad parameter matrices here, as there is just one new token for batched generation
-            dt = dt.transpose(1, 2).expand(batch_size, dt.shape[-1], self.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))[..., None]
-            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)
-            dA = (torch.exp(dt * A)).to(device=cache_device)
-
-            # Discretize B
-            B = B.reshape(batch_size, self.n_groups, 1, -1)
-            B = B.expand(batch_size, self.n_groups, self.num_heads // self.n_groups, B.shape[-1]).contiguous()
-            B = B.reshape(batch_size, -1, 1, B.shape[-1])
-            dB = dt * B
-
-            # Discretize x into dB
-            hidden_states = hidden_states.reshape(batch_size, -1, self.head_dim)
-            dBx = (dB * hidden_states[..., None]).to(device=cache_device)
-
-            # State calculation
-            ssm_states = cache_params.layers[self.layer_idx].recurrent_states[0] * dA + dBx
-            ssm_states = cache_params.update_recurrent_state(ssm_states, layer_idx=self.layer_idx)
-
-            # Subsequent output
-            C = C.reshape(batch_size, self.n_groups, 1, -1)
-            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])
-
-            # Reshape ssm_states to merge the first two dimensions
-            ssm_states = ssm_states.to(device=C.device, dtype=C.dtype)
-            ssm_states_reshaped = ssm_states.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)
-            y = torch.bmm(ssm_states_reshaped, C_reshaped)
-            y = y.view(batch_size, self.num_heads, self.head_dim)
-
-            # D skip connection
-            D = self.D[..., None].expand(self.D.shape[0], self.head_dim)
-            y = (y + hidden_states * D).to(y.dtype)
-
-            y = y.reshape(batch_size, 1, -1)
-        else:
-            # 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_interleave(self.num_heads // self.n_groups, dim=2, output_size=self.num_heads)
-            C = C.repeat_interleave(self.num_heads // self.n_groups, dim=2, output_size=self.num_heads)
-            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)
-            ]
-
-            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)
-            G_intermediate = C[:, :, :, None, :, :] * B[:, :, None, :, :, :]
-            G = G_intermediate.sum(dim=-1)
-
-            # 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 = (
-                cache_params.layers[self.layer_idx]
-                .recurrent_states[0][:, None]
-                .to(dtype=states.dtype, device=states.device)
-                if use_precomputed_states
-                else 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
-            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)
-
-            # Init cache
-            if ssm_state is not None and cache_params is not None:
-                cache_params.update_recurrent_state(ssm_state, layer_idx=self.layer_idx)
-
-        scan_output = self.norm(y, gate)
+        scan_output = self.norm(scan_output, gate)
 
         # 4. Final linear projection
         contextualized_states = self.out_proj(scan_output.to(dtype))
         return contextualized_states
 
-    @force_accelerate_hooks("conv1d")
-    def forward(
-        self,
-        hidden_states,
-        cache_params: Cache | None = None,
-        attention_mask: torch.Tensor | None = None,
-        **kwargs,
-    ):
-        if is_fast_path_available and "cuda" in self.in_proj.weight.device.type and not is_torchdynamo_compiling():
-            return self.cuda_kernels_forward(hidden_states, cache_params, attention_mask, **kwargs)
-        if kwargs.get("seq_idx") is not None:
-            raise NotImplementedError(
-                "`seq_idx` support requires fast path support. Please install `mamba_ssm` and `causal_conv1d`"
-            )
-        return self.torch_forward(hidden_states, cache_params, attention_mask, **kwargs)
-
 
 class GraniteMoeHybridRMSNormGated(torch.nn.Module):
     def __init__(self, hidden_size, eps=1e-6):
diff --git a/src/transformers/models/mamba2/modeling_mamba2.py b/src/transformers/models/mamba2/modeling_mamba2.py
index 185542b062f4..311d84e4baa1 100644
--- a/src/transformers/models/mamba2/modeling_mamba2.py
+++ b/src/transformers/models/mamba2/modeling_mamba2.py
@@ -207,6 +207,7 @@ def mamba2_selective_state_update(
     D: torch.Tensor | None = None,
     dt_bias: torch.Tensor | None = None,
     dt_softplus: bool = False,
+    z: torch.Tensor | None = None,
     **kwargs,
 ):
     batch_size, num_heads, head_dim = hidden_states.shape
@@ -251,6 +252,9 @@ def mamba2_selective_state_update(
     if D is not None:
         out = (out + hidden_states * D).to(out.dtype)
 
+    if z is not None:
+        out = out * F.silu(z)
+
     return out.to(hidden_states.dtype)
 
 
@@ -565,6 +569,7 @@ def forward(
                 **kwargs,
             )
             scan_output = scan_output.view(batch_size, 1, -1)
+
         # Chunk form
         else:
             output_final_state = cache_params is not None
diff --git a/src/transformers/models/nemotron_h/modeling_nemotron_h.py b/src/transformers/models/nemotron_h/modeling_nemotron_h.py
index 47b6be933324..a7a4afc842fc 100644
--- a/src/transformers/models/nemotron_h/modeling_nemotron_h.py
+++ b/src/transformers/models/nemotron_h/modeling_nemotron_h.py
@@ -32,9 +32,9 @@
 from ...cache_utils import Cache, DynamicCache
 from ...generation import GenerationMixin
 from ...integrations import (
-    lazy_load_kernel,
     use_experts_implementation,
     use_kernel_forward_from_hub,
+    use_kernel_func_from_hub_with_fallback,
     use_kernelized_func,
 )
 from ...integrations.accelerate import force_accelerate_hooks
@@ -44,16 +44,12 @@
 from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
 from ...models.zamba2.modeling_zamba2 import Zamba2RMSNormGated
 from ...processing_utils import Unpack
-from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, is_torchdynamo_compiling, logging
+from ...utils import TransformersKwargs, auto_docstring, can_return_tuple
 from ...utils.generic import merge_with_config_defaults
-from ...utils.import_utils import resolve_internal_import
 from ...utils.output_capturing import capture_outputs
 from .configuration_nemotron_h import NemotronHConfig
 
 
-logger = logging.get_logger(__name__)
-
-
 # Helper methods for segment sum computation
 
 
@@ -120,6 +116,7 @@ def apply_mask_to_padding_states(hidden_states, attention_mask):
     return hidden_states
 
 
+@use_kernel_func_from_hub_with_fallback("causal_conv1d_update", "causal_conv1d")
 def causal_conv1d_update(
     hidden_states: torch.Tensor,
     conv_state: torch.Tensor,
@@ -139,6 +136,7 @@ def causal_conv1d_update(
     return out.to(hidden_states.dtype)
 
 
+@use_kernel_func_from_hub_with_fallback("causal_conv1d_fn", "causal_conv1d")
 def causal_conv1d_fn(
     hidden_states: torch.Tensor,
     weight: nn.Parameter,
@@ -161,9 +159,215 @@ def causal_conv1d_fn(
     return out.to(hidden_states.dtype)
 
 
-is_fast_path_available = False
+@use_kernel_func_from_hub_with_fallback(
+    "mamba_split_conv1d_scan_combined",
+    "mamba_ssm",
+    internal_path="ops.triton.ssd_combined",
+)
+def nemotron_h_mamba2_split_conv1d_scan_combined(
+    zxbcdt: torch.Tensor,
+    conv1d_weight: torch.Tensor,
+    conv1d_bias: torch.Tensor | None,
+    dt_bias: torch.Tensor,
+    A: torch.Tensor,
+    D: torch.Tensor,
+    chunk_size: int,
+    initial_states: 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,
+    **kwargs,
+):
+    return None
 
 
+@use_kernel_func_from_hub_with_fallback(
+    "selective_state_update",
+    "mamba_ssm",
+    internal_path="ops.triton.selective_state_update",
+)
+def nemotron_h_mamba2_selective_state_update(
+    state: torch.Tensor,
+    hidden_states: torch.Tensor,
+    dt: torch.Tensor,
+    A: torch.Tensor,
+    B: torch.Tensor,
+    C: torch.Tensor,
+    D: torch.Tensor | None = None,
+    dt_bias: torch.Tensor | None = None,
+    dt_softplus: bool = False,
+    z: torch.Tensor | None = None,
+    **kwargs,
+):
+    batch_size, num_heads, head_dim = hidden_states.shape
+    num_groups = B.shape[1]
+    state_size = B.shape[-1]
+
+    if dt_bias is not None:
+        dt = dt + dt_bias.to(dt.dtype)
+    if dt_softplus:
+        dt = F.softplus(dt)
+    dt = dt[..., None]
+
+    # Discretize A
+    dA = torch.exp(dt.float() * A.float()).to(device=state.device)
+
+    # Discretize B
+    B = B.reshape(batch_size, num_groups, 1, state_size)
+    B = B.expand(batch_size, num_groups, num_heads // num_groups, state_size).contiguous()
+    B = B.reshape(batch_size, num_heads, 1, state_size)
+    dB = dt * B
+
+    # Discretize x into dB
+    dBx = (dB * hidden_states[..., None]).to(device=state.device)
+
+    # State calculation
+    ssm_states = state * dA + dBx
+    state.copy_(ssm_states.to(state.dtype))
+
+    # Subsequent output
+    C = C.reshape(batch_size, num_groups, 1, state_size)
+    C = C.expand(batch_size, num_groups, num_heads // num_groups, state_size).contiguous()
+    C = C.reshape(batch_size, num_heads, state_size)
+
+    # Reshape ssm_states to merge the first two dimensions
+    ssm_states = ssm_states.to(device=C.device, dtype=C.dtype)
+    ssm_states_reshaped = ssm_states.view(batch_size * num_heads, head_dim, state_size)
+    C_reshaped = C.view(batch_size * num_heads, state_size, 1)
+    out = torch.bmm(ssm_states_reshaped, C_reshaped)
+    out = out.view(batch_size, num_heads, head_dim)
+
+    # D skip connection
+    if D is not None:
+        out = (out + hidden_states * D).to(out.dtype)
+
+    if z is not None:
+        out = out * F.silu(z)
+
+    return out.to(hidden_states.dtype)
+
+
+@use_kernel_func_from_hub_with_fallback(
+    "mamba_chunk_scan_combined",
+    "mamba_ssm",
+    internal_path="ops.triton.ssd_combined",
+)
+def nemotron_h_mamba2_chunk_scan(
+    hidden_states: torch.Tensor,
+    dt: torch.Tensor,
+    A: torch.Tensor,
+    B: torch.Tensor,
+    C: torch.Tensor,
+    chunk_size: int,
+    D: torch.Tensor | None = None,
+    dt_bias: torch.Tensor | None = None,
+    initial_states: torch.Tensor | None = None,
+    dt_softplus: bool = False,
+    dt_limit: tuple[float, float] = (0.0, float("inf")),
+    return_final_states: bool = False,
+    **kwargs,
+):
+    input_dtype = hidden_states.dtype
+    batch_size, sequence_length, num_heads, head_dim = hidden_states.shape
+    num_groups = B.shape[2]
+
+    if dt_bias is not None:
+        dt = dt + dt_bias.to(dt.dtype)
+    if dt_softplus:
+        dt = F.softplus(dt)
+    dt = torch.clamp(dt, min=dt_limit[0], max=dt_limit[1])
+
+    hidden_states = hidden_states.float()
+    B = B.float().repeat_interleave(num_heads // num_groups, dim=2, output_size=num_heads)
+    C = C.float().repeat_interleave(num_heads // num_groups, dim=2, output_size=num_heads)
+
+    pad_size = (chunk_size - sequence_length % chunk_size) % chunk_size
+    D_residual = None
+    if D is not None:
+        D_residual = D[..., None] * pad_tensor_by_size(hidden_states, pad_size)
+
+    # Discretize x and A
+    hidden_states = hidden_states * dt[..., None].float()
+    A = A.to(hidden_states.dtype) * dt.float()
+
+    # Rearrange into blocks/chunks
+    hidden_states, A, B, C = [reshape_into_chunks(tensor, pad_size, chunk_size) for tensor in (hidden_states, A, B, C)]
+
+    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)
+    G = (C[:, :, :, None, :, :] * B[:, :, None, :, :, :]).sum(dim=-1)
+
+    # Compute M, equivalent to applying attention mask to weights
+    M = (G[..., None] * L.permute(0, 2, 3, 4, 1)[..., None]).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 = (
+        initial_states[:, None].to(dtype=states.dtype, device=states.device)
+        if initial_states is not None
+        else torch.zeros_like(states[:, :1])
+    )
+    states = torch.cat([previous_states, states], dim=1)
+    decay_chunk = torch.exp(segment_sum(F.pad(A_cumsum[:, :, :, -1], (1, 0)))).transpose(1, 3)
+    new_states = (decay_chunk[..., None, None] * states[:, :, None, ...]).sum(dim=1)
+    states, final_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, ...]
+    Y_off = C_times_states.sum(-1) * state_decay_out.permute(0, 2, 3, 1)[..., None]
+
+    # Add output of intra-chunk and inter-chunk terms (diagonal and off-diagonal blocks)
+    output = Y_diag + Y_off
+    output = output.reshape(batch_size, -1, num_heads, head_dim)
+
+    if D_residual is not None:
+        output = output + D_residual
+
+    # Cutting off padded chunks
+    if pad_size > 0:
+        output = output[:, :sequence_length]
+
+    output = output.to(input_dtype)
+
+    if return_final_states:
+        return output, final_state
+
+    return output
+
+
+@use_kernelized_func(
+    [
+        causal_conv1d_fn,
+        causal_conv1d_update,
+        nemotron_h_mamba2_split_conv1d_scan_combined,
+        nemotron_h_mamba2_selective_state_update,
+        nemotron_h_mamba2_chunk_scan,
+    ]
+)
 class NemotronHMamba2Mixer(nn.Module):
     """
     Compute ∆, A, B, C, and D the state space parameters and compute the `contextualized_states`.
@@ -224,37 +428,6 @@ def __init__(self, config: NemotronHConfig, layer_idx: int | None = None, initia
             self.init_nemotron_h_mamba2_weights()
         self.out_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=config.use_bias)
 
-        global causal_conv1d, causal_conv1d_update, causal_conv1d_fn
-        causal_conv1d = lazy_load_kernel("causal-conv1d")
-        causal_conv1d_update = getattr(causal_conv1d, "causal_conv1d_update", causal_conv1d_update)
-        causal_conv1d_fn = getattr(causal_conv1d, "causal_conv1d_fn", causal_conv1d_fn)
-
-        global mamba_ssm, selective_state_update, mamba_chunk_scan_combined, mamba_split_conv1d_scan_combined
-        mamba_ssm = lazy_load_kernel("mamba-ssm")
-        selective_state_update = resolve_internal_import(
-            mamba_ssm, chained_path="ops.triton.selective_state_update.selective_state_update"
-        )
-        mamba_chunk_scan_combined = resolve_internal_import(
-            mamba_ssm, chained_path="ops.triton.ssd_combined.mamba_chunk_scan_combined"
-        )
-        mamba_split_conv1d_scan_combined = resolve_internal_import(
-            mamba_ssm, chained_path="ops.triton.ssd_combined.mamba_split_conv1d_scan_combined"
-        )
-
-        global is_fast_path_available
-        is_fast_path_available = (
-            all((selective_state_update, mamba_chunk_scan_combined, mamba_split_conv1d_scan_combined))
-            and hasattr(causal_conv1d, "causal_conv1d_update")
-            and hasattr(causal_conv1d, "causal_conv1d_fn")
-        )
-
-        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"
-            )
-
         self.layer_type = config.layer_types[layer_idx]
         self.use_mem_eff_path = True
 
@@ -265,65 +438,28 @@ def init_nemotron_h_mamba2_weights(self):
         init.ones_(self.D)
         init.ones_(self.dt_bias)
 
-    def _convolution(
+    @force_accelerate_hooks("conv1d")
+    def forward(
         self,
         hidden_states: torch.Tensor,
         cache_params: Cache | None = None,
-        attention_mask: torch.LongTensor | None = None,
+        attention_mask: torch.Tensor | None = None,
         **kwargs,
     ):
-        seq_len = hidden_states.shape[1]
-        hidden_states = apply_mask_to_padding_states(hidden_states, attention_mask)
-        hidden_states = hidden_states.transpose(1, 2)
-
+        batch_size, seq_len, _ = hidden_states.shape
+        dtype = hidden_states.dtype
         use_precomputed_states = cache_params is not None and cache_params.has_previous_state(self.layer_idx)
 
-        if use_precomputed_states and seq_len == 1 and not cache_params.layers[self.layer_idx].record_past:
-            conv_state = cache_params.layers[self.layer_idx].conv_states[0]
-            hidden_states = causal_conv1d_update(
-                hidden_states,
-                conv_state,
-                self.conv1d.weight.squeeze(1),
-                self.conv1d.bias,
-                self.activation,
-            )
-        else:
-            if cache_params is not None:
-                hidden_states = cache_params.update_conv_state(
-                    hidden_states, self.layer_idx, conv_kernel_size=self.conv_kernel_size
-                )
-
-            hidden_states = causal_conv1d_fn(
-                hidden_states,
-                self.conv1d.weight.squeeze(1),
-                self.conv1d.bias,
-                activation=self.activation,
-                seq_idx=kwargs.get("seq_idx"),
-            )
-
-            # Drop the additional previous states
-            if cache_params is not None:
-                hidden_states = hidden_states[:, :, -seq_len:]
-
-        hidden_states = hidden_states.transpose(1, 2)
-        return hidden_states
-
-    def cuda_kernels_forward(
-        self,
-        hidden_states: torch.Tensor,
-        cache_params: Cache | None = None,
-        attention_mask: torch.Tensor | None = None,
-        **kwargs,
-    ):
         # 1. Gated MLP's linear projection
         hidden_states = apply_mask_to_padding_states(hidden_states, attention_mask)
         projected_states = self.in_proj(hidden_states)
 
         A = -torch.exp(self.A_log.float())
-        dt_limit_kwargs = {} if self.time_step_limit == (0.0, float("inf")) else {"dt_limit": self.time_step_limit}
-        # Fused kernel for conv1d, SSM, and the final projection
+        fused_kwargs = (
+            kwargs | {} if self.time_step_limit == (0.0, float("inf")) else kwargs | {"dt_limit": self.time_step_limit}
+        )
         if self.training and cache_params is None:
-            return mamba_split_conv1d_scan_combined(
+            fused_output = nemotron_h_mamba2_split_conv1d_scan_combined(
                 projected_states,
                 self.conv1d.weight.squeeze(1),
                 self.conv1d.bias,
@@ -331,7 +467,6 @@ def cuda_kernels_forward(
                 A,
                 D=self.D,
                 chunk_size=self.chunk_size,
-                seq_idx=kwargs.get("seq_idx"),
                 activation=self.activation,
                 rmsnorm_weight=self.norm.weight,
                 rmsnorm_eps=self.norm.variance_epsilon,
@@ -341,41 +476,71 @@ def cuda_kernels_forward(
                 ngroups=self.n_groups,
                 norm_before_gate=False,
                 return_final_states=False,
-                **dt_limit_kwargs,
+                **fused_kwargs,
             )
 
-        # Set up dimensions for reshapes later
-        batch_size, seq_len, _ = hidden_states.shape
-        groups_time_state_size = self.n_groups * self.ssm_state_size
-        use_precomputed_states = cache_params is not None and cache_params.has_previous_state(self.layer_idx)
+            # Only kernels can use this shortcircuit, fallback to normal torch otherwise
+            if fused_output is not None:
+                return fused_output
 
         gate, hidden_states_B_C, dt = projected_states.split(
             [self.intermediate_size, self.conv_dim, self.num_heads], dim=-1
         )
 
-        # Apply the conv
-        hidden_states_B_C = self._convolution(hidden_states_B_C, cache_params, attention_mask, **kwargs)
-        hidden_states_B_C = apply_mask_to_padding_states(hidden_states_B_C, attention_mask)
+        if use_precomputed_states:
+            conv_state = cache_params.layers[self.layer_idx].conv_states[0]
+            recurrent_state = cache_params.layers[self.layer_idx].recurrent_states[0]
+
+        # 2. Convolution sequence transformation
+        hidden_states_B_C = hidden_states_B_C.transpose(1, 2)
+        if use_precomputed_states and seq_len == 1:
+            hidden_states_B_C = causal_conv1d_update(
+                hidden_states_B_C,
+                conv_state,
+                self.conv1d.weight.squeeze(1),
+                self.conv1d.bias,
+                activation=self.activation,
+            )
+        else:
+            if cache_params is not None:
+                hidden_states_B_C = cache_params.update_conv_state(
+                    hidden_states_B_C,
+                    self.layer_idx,
+                    conv_kernel_size=self.conv_kernel_size,
+                )
+
+            hidden_states_B_C = causal_conv1d_fn(
+                hidden_states_B_C,
+                self.conv1d.weight.squeeze(1),
+                self.conv1d.bias,
+                activation=self.activation,
+                **kwargs,
+            )
+
+            if cache_params is not None:
+                hidden_states_B_C = hidden_states_B_C[:, :, -seq_len:]
+
+        # 3. SSM transformation
+        hidden_states_B_C = apply_mask_to_padding_states(hidden_states_B_C.transpose(1, 2), attention_mask)
         hidden_states, B, C = torch.split(
             hidden_states_B_C,
-            [self.intermediate_size, groups_time_state_size, groups_time_state_size],
+            [self.intermediate_size, self.n_groups * self.ssm_state_size, self.n_groups * self.ssm_state_size],
             dim=-1,
         )
 
-        recurrent_state = cache_params.layers[self.layer_idx].recurrent_states[0] if use_precomputed_states else None
-        # Single step calculations via cache
+        # Recurrent form
         if use_precomputed_states and seq_len == 1:
-            # 3. SSM transformation
-            A = A[:, None, ...][:, :, None].expand(-1, self.head_dim, self.ssm_state_size).to(dtype=torch.float32)
+            hidden_states = hidden_states.view(batch_size, self.num_heads, self.head_dim)
             dt = dt.transpose(1, 2).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[2] // self.n_groups)
-            C = C.view(batch_size, self.n_groups, C.shape[2] // self.n_groups)
-            hidden_states_reshaped = hidden_states.view(batch_size, self.num_heads, self.head_dim)
-            hidden_states = selective_state_update(
+            A = A[:, None, None].expand(-1, self.head_dim, self.ssm_state_size)
+            B = B.view(batch_size, self.n_groups, self.ssm_state_size)
+            C = C.view(batch_size, self.n_groups, self.ssm_state_size)
+            D = self.D[:, None].expand(-1, self.head_dim)
+            dt_bias = self.dt_bias[:, None].expand(-1, self.head_dim)
+
+            scan_output = nemotron_h_mamba2_selective_state_update(
                 recurrent_state,
-                hidden_states_reshaped,
+                hidden_states,
                 dt,
                 A,
                 B,
@@ -384,225 +549,44 @@ def cuda_kernels_forward(
                 z=None,
                 dt_bias=dt_bias,
                 dt_softplus=True,
+                **kwargs,
             )
-            hidden_states = hidden_states.view(batch_size, 1, self.num_heads * self.head_dim)
-            hidden_states = self.norm(hidden_states, gate)
+            scan_output = scan_output.view(batch_size, 1, -1)
 
-            # 4. Final linear projection
-            out = self.out_proj(hidden_states)
-
-        # Fused calculations or step by step if no initialized cache is found
+        # Chunk form
         else:
-            # 3. SSM transformation
-            scan_output, ssm_state = mamba_chunk_scan_combined(
-                hidden_states.view(batch_size, seq_len, -1, self.head_dim),
+            output_final_state = cache_params is not None
+            scan_result = nemotron_h_mamba2_chunk_scan(
+                hidden_states.view(batch_size, seq_len, self.num_heads, 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),
+                B.view(batch_size, seq_len, self.n_groups, self.ssm_state_size),
+                C.view(batch_size, seq_len, self.n_groups, self.ssm_state_size),
                 chunk_size=self.chunk_size,
                 D=self.D,
                 z=None,
-                seq_idx=kwargs.get("seq_idx"),
-                return_final_states=True,
+                return_final_states=output_final_state,
                 dt_bias=self.dt_bias,
                 dt_softplus=True,
-                initial_states=recurrent_state,
-                **dt_limit_kwargs,
+                initial_states=recurrent_state if use_precomputed_states else None,
+                dt_limit=self.time_step_limit,
+                **kwargs,
             )
 
-            # Init cache
-            if ssm_state is not None and cache_params is not None:
-                cache_params.update_recurrent_state(ssm_state, layer_idx=self.layer_idx)
-
-            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
-
-    def torch_forward(
-        self,
-        hidden_states: torch.Tensor,
-        cache_params: Cache | None = None,
-        attention_mask: torch.Tensor | None = None,
-        **kwargs,
-    ):
-        batch_size, seq_len, _ = hidden_states.shape
-        dtype = hidden_states.dtype
-        use_precomputed_states = cache_params is not None and cache_params.has_previous_state(self.layer_idx)
-
-        # 1. Gated MLP's linear projection
-        hidden_states = apply_mask_to_padding_states(hidden_states, attention_mask)
-        projected_states = self.in_proj(hidden_states)
+            if output_final_state:
+                scan_output, final_state = scan_result
+                cache_params.update_recurrent_state(final_state, layer_idx=self.layer_idx)
+            else:
+                scan_output = scan_result
 
-        gate, hidden_states_B_C, dt = projected_states.split(
-            [self.intermediate_size, self.conv_dim, self.num_heads], dim=-1
-        )
+            scan_output = scan_output.reshape(batch_size, seq_len, -1)
 
-        # 2. Convolution sequence transformation
-        hidden_states_B_C = self._convolution(hidden_states_B_C, cache_params, attention_mask, **kwargs)
-        hidden_states_B_C = apply_mask_to_padding_states(hidden_states_B_C, attention_mask)
-        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())
-        if use_precomputed_states and seq_len == 1:
-            # We need to guarantee that anything regarding the cache is on the same device
-            cache_device = cache_params.layers[self.layer_idx].device
-
-            # Note: there is no need to pad parameter matrices here, as there is just one new token for batched generation
-            dt = dt.transpose(1, 2).expand(batch_size, dt.shape[-1], self.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))[..., None]
-            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)
-            dA = (torch.exp(dt * A)).to(device=cache_device)
-
-            # Discretize B
-            B = B.reshape(batch_size, self.n_groups, 1, -1)
-            B = B.expand(batch_size, self.n_groups, self.num_heads // self.n_groups, B.shape[-1]).contiguous()
-            B = B.reshape(batch_size, -1, 1, B.shape[-1])
-            dB = dt * B
-
-            # Discretize x into dB
-            hidden_states = hidden_states.reshape(batch_size, -1, self.head_dim)
-            dBx = (dB * hidden_states[..., None]).to(device=cache_device)
-
-            # State calculation
-            ssm_states = cache_params.layers[self.layer_idx].recurrent_states[0] * dA + dBx
-            ssm_states = cache_params.update_recurrent_state(ssm_states, layer_idx=self.layer_idx)
-
-            # Subsequent output
-            C = C.reshape(batch_size, self.n_groups, 1, -1)
-            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])
-
-            # Reshape ssm_states to merge the first two dimensions
-            ssm_states = ssm_states.to(device=C.device, dtype=C.dtype)
-            ssm_states_reshaped = ssm_states.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)
-            y = torch.bmm(ssm_states_reshaped, C_reshaped)
-            y = y.view(batch_size, self.num_heads, self.head_dim)
-
-            # D skip connection
-            D = self.D[..., None].expand(self.D.shape[0], self.head_dim)
-            y = (y + hidden_states * D).to(y.dtype)
-
-            y = y.reshape(batch_size, 1, -1)
-        else:
-            # 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_interleave(self.num_heads // self.n_groups, dim=2, output_size=self.num_heads)
-            C = C.repeat_interleave(self.num_heads // self.n_groups, dim=2, output_size=self.num_heads)
-            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)
-            ]
-
-            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)
-            G_intermediate = C[:, :, :, None, :, :] * B[:, :, None, :, :, :]
-            G = G_intermediate.sum(dim=-1)
-
-            # 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 = (
-                cache_params.layers[self.layer_idx]
-                .recurrent_states[0][:, None]
-                .to(dtype=states.dtype, device=states.device)
-                if use_precomputed_states
-                else 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
-            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)
-
-            # Init cache
-            if ssm_state is not None and cache_params is not None:
-                cache_params.update_recurrent_state(ssm_state, layer_idx=self.layer_idx)
-
-        scan_output = self.norm(y, gate)
+        scan_output = self.norm(scan_output, gate)
 
         # 4. Final linear projection
         contextualized_states = self.out_proj(scan_output.to(dtype))
         return contextualized_states
 
-    @force_accelerate_hooks("conv1d")
-    def forward(
-        self,
-        hidden_states,
-        cache_params: Cache | None = None,
-        attention_mask: torch.Tensor | None = None,
-        **kwargs,
-    ):
-        if is_fast_path_available and "cuda" in self.in_proj.weight.device.type and not is_torchdynamo_compiling():
-            # Use cuda stream to avoid NaN when using multiple GPUs, which is caused by multi-GPU synchronization issue.
-            # Mamba might launch on the default cuda stream that not strictly respect the current Pytorch cuda stream.
-            # This leads to kernel reading uninitialized memory before the data transfer is complete.
-            with torch.cuda.stream(torch.cuda.default_stream(hidden_states.device)):
-                return self.cuda_kernels_forward(hidden_states, cache_params, attention_mask)
-
-        return self.torch_forward(hidden_states, cache_params, attention_mask)
-
 
 @use_kernel_forward_from_hub("RMSNorm")
 class NemotronHRMSNorm(nn.Module):
diff --git a/src/transformers/models/nemotron_h/modular_nemotron_h.py b/src/transformers/models/nemotron_h/modular_nemotron_h.py
index f84cb2d29994..4531b4eff6d0 100644
--- a/src/transformers/models/nemotron_h/modular_nemotron_h.py
+++ b/src/transformers/models/nemotron_h/modular_nemotron_h.py
@@ -35,7 +35,7 @@
 from ...models.zamba.modeling_zamba import ZambaForCausalLM
 from ...models.zamba2.modeling_zamba2 import Zamba2MambaMixer, Zamba2RMSNormGated
 from ...processing_utils import Unpack
-from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, is_torchdynamo_compiling, logging
+from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, logging
 from ...utils.generic import merge_with_config_defaults
 from ...utils.output_capturing import capture_outputs
 from .configuration_nemotron_h import NemotronHConfig
@@ -43,8 +43,6 @@
 
 logger = logging.get_logger(__name__)
 
-is_fast_path_available = False
-
 
 class NemotronHMamba2Mixer(Zamba2MambaMixer):
     def __init__(self, config: NemotronHConfig, layer_idx: int | None = None, initialize_mixer_weights: bool = True):
@@ -81,22 +79,6 @@ def __init__(self, config: NemotronHConfig, layer_idx: int | None = None, initia
         )
         self.out_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=config.use_bias)
 
-    def forward(
-        self,
-        hidden_states,
-        cache_params: Cache | None = None,
-        attention_mask: torch.Tensor | None = None,
-        **kwargs,
-    ):
-        if is_fast_path_available and "cuda" in self.in_proj.weight.device.type and not is_torchdynamo_compiling():
-            # Use cuda stream to avoid NaN when using multiple GPUs, which is caused by multi-GPU synchronization issue.
-            # Mamba might launch on the default cuda stream that not strictly respect the current Pytorch cuda stream.
-            # This leads to kernel reading uninitialized memory before the data transfer is complete.
-            with torch.cuda.stream(torch.cuda.default_stream(hidden_states.device)):
-                return self.cuda_kernels_forward(hidden_states, cache_params, attention_mask)
-
-        return self.torch_forward(hidden_states, cache_params, attention_mask)
-
 
 class NemotronHRMSNorm(LlamaRMSNorm):
     pass
diff --git a/src/transformers/models/zamba2/modeling_zamba2.py b/src/transformers/models/zamba2/modeling_zamba2.py
index e82f147d86a4..d2d953d52221 100644
--- a/src/transformers/models/zamba2/modeling_zamba2.py
+++ b/src/transformers/models/zamba2/modeling_zamba2.py
@@ -30,7 +30,7 @@
 from ...activations import ACT2FN
 from ...cache_utils import Cache, DynamicCache
 from ...generation import GenerationMixin
-from ...integrations import lazy_load_kernel, use_kernel_forward_from_hub
+from ...integrations import use_kernel_forward_from_hub, use_kernel_func_from_hub_with_fallback, use_kernelized_func
 from ...integrations.accelerate import force_accelerate_hooks
 from ...masking_utils import create_causal_mask, create_recurrent_attention_mask
 from ...modeling_layers import GradientCheckpointingLayer
@@ -38,10 +38,9 @@
 from ...modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update
 from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
 from ...processing_utils import Unpack
-from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, is_torchdynamo_compiling, logging
+from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, logging
 from ...utils.deprecation import deprecate_kwarg
 from ...utils.generic import maybe_autocast, merge_with_config_defaults
-from ...utils.import_utils import resolve_internal_import
 from ...utils.output_capturing import capture_outputs
 from .configuration_zamba2 import Zamba2Config
 
@@ -409,6 +408,7 @@ def apply_mask_to_padding_states(hidden_states, attention_mask):
     return hidden_states
 
 
+@use_kernel_func_from_hub_with_fallback("causal_conv1d_update", "causal_conv1d")
 def causal_conv1d_update(
     hidden_states: torch.Tensor,
     conv_state: torch.Tensor,
@@ -428,6 +428,7 @@ def causal_conv1d_update(
     return out.to(hidden_states.dtype)
 
 
+@use_kernel_func_from_hub_with_fallback("causal_conv1d_fn", "causal_conv1d")
 def causal_conv1d_fn(
     hidden_states: torch.Tensor,
     weight: nn.Parameter,
@@ -450,6 +451,215 @@ def causal_conv1d_fn(
     return out.to(hidden_states.dtype)
 
 
+@use_kernel_func_from_hub_with_fallback(
+    "mamba_split_conv1d_scan_combined",
+    "mamba_ssm",
+    internal_path="ops.triton.ssd_combined",
+)
+def zamba2_split_conv1d_scan_combined(
+    zxbcdt: torch.Tensor,
+    conv1d_weight: torch.Tensor,
+    conv1d_bias: torch.Tensor | None,
+    dt_bias: torch.Tensor,
+    A: torch.Tensor,
+    D: torch.Tensor,
+    chunk_size: int,
+    initial_states: 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,
+    **kwargs,
+):
+    return None
+
+
+@use_kernel_func_from_hub_with_fallback(
+    "selective_state_update",
+    "mamba_ssm",
+    internal_path="ops.triton.selective_state_update",
+)
+def zamba2_selective_state_update(
+    state: torch.Tensor,
+    hidden_states: torch.Tensor,
+    dt: torch.Tensor,
+    A: torch.Tensor,
+    B: torch.Tensor,
+    C: torch.Tensor,
+    D: torch.Tensor | None = None,
+    dt_bias: torch.Tensor | None = None,
+    dt_softplus: bool = False,
+    z: torch.Tensor | None = None,
+    **kwargs,
+):
+    batch_size, num_heads, head_dim = hidden_states.shape
+    num_groups = B.shape[1]
+    state_size = B.shape[-1]
+
+    if dt_bias is not None:
+        dt = dt + dt_bias.to(dt.dtype)
+    if dt_softplus:
+        dt = F.softplus(dt)
+    dt = dt[..., None]
+
+    # Discretize A
+    dA = torch.exp(dt.float() * A.float()).to(device=state.device)
+
+    # Discretize B
+    B = B.reshape(batch_size, num_groups, 1, state_size)
+    B = B.expand(batch_size, num_groups, num_heads // num_groups, state_size).contiguous()
+    B = B.reshape(batch_size, num_heads, 1, state_size)
+    dB = dt * B
+
+    # Discretize x into dB
+    dBx = (dB * hidden_states[..., None]).to(device=state.device)
+
+    # State calculation
+    ssm_states = state * dA + dBx
+    state.copy_(ssm_states.to(state.dtype))
+
+    # Subsequent output
+    C = C.reshape(batch_size, num_groups, 1, state_size)
+    C = C.expand(batch_size, num_groups, num_heads // num_groups, state_size).contiguous()
+    C = C.reshape(batch_size, num_heads, state_size)
+
+    # Reshape ssm_states to merge the first two dimensions
+    ssm_states = ssm_states.to(device=C.device, dtype=C.dtype)
+    ssm_states_reshaped = ssm_states.view(batch_size * num_heads, head_dim, state_size)
+    C_reshaped = C.view(batch_size * num_heads, state_size, 1)
+    out = torch.bmm(ssm_states_reshaped, C_reshaped)
+    out = out.view(batch_size, num_heads, head_dim)
+
+    # D skip connection
+    if D is not None:
+        out = (out + hidden_states * D).to(out.dtype)
+
+    if z is not None:
+        out = out * F.silu(z)
+
+    return out.to(hidden_states.dtype)
+
+
+@use_kernel_func_from_hub_with_fallback(
+    "mamba_chunk_scan_combined",
+    "mamba_ssm",
+    internal_path="ops.triton.ssd_combined",
+)
+def zamba2_chunk_scan(
+    hidden_states: torch.Tensor,
+    dt: torch.Tensor,
+    A: torch.Tensor,
+    B: torch.Tensor,
+    C: torch.Tensor,
+    chunk_size: int,
+    D: torch.Tensor | None = None,
+    dt_bias: torch.Tensor | None = None,
+    initial_states: torch.Tensor | None = None,
+    dt_softplus: bool = False,
+    dt_limit: tuple[float, float] = (0.0, float("inf")),
+    return_final_states: bool = False,
+    **kwargs,
+):
+    input_dtype = hidden_states.dtype
+    batch_size, sequence_length, num_heads, head_dim = hidden_states.shape
+    num_groups = B.shape[2]
+
+    if dt_bias is not None:
+        dt = dt + dt_bias.to(dt.dtype)
+    if dt_softplus:
+        dt = F.softplus(dt)
+    dt = torch.clamp(dt, min=dt_limit[0], max=dt_limit[1])
+
+    hidden_states = hidden_states.float()
+    B = B.float().repeat_interleave(num_heads // num_groups, dim=2, output_size=num_heads)
+    C = C.float().repeat_interleave(num_heads // num_groups, dim=2, output_size=num_heads)
+
+    pad_size = (chunk_size - sequence_length % chunk_size) % chunk_size
+    D_residual = None
+    if D is not None:
+        D_residual = D[..., None] * pad_tensor_by_size(hidden_states, pad_size)
+
+    # Discretize x and A
+    hidden_states = hidden_states * dt[..., None].float()
+    A = A.to(hidden_states.dtype) * dt.float()
+
+    # Rearrange into blocks/chunks
+    hidden_states, A, B, C = [reshape_into_chunks(tensor, pad_size, chunk_size) for tensor in (hidden_states, A, B, C)]
+
+    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)
+    G = (C[:, :, :, None, :, :] * B[:, :, None, :, :, :]).sum(dim=-1)
+
+    # Compute M, equivalent to applying attention mask to weights
+    M = (G[..., None] * L.permute(0, 2, 3, 4, 1)[..., None]).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 = (
+        initial_states[:, None].to(dtype=states.dtype, device=states.device)
+        if initial_states is not None
+        else torch.zeros_like(states[:, :1])
+    )
+    states = torch.cat([previous_states, states], dim=1)
+    decay_chunk = torch.exp(segment_sum(F.pad(A_cumsum[:, :, :, -1], (1, 0)))).transpose(1, 3)
+    new_states = (decay_chunk[..., None, None] * states[:, :, None, ...]).sum(dim=1)
+    states, final_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, ...]
+    Y_off = C_times_states.sum(-1) * state_decay_out.permute(0, 2, 3, 1)[..., None]
+
+    # Add output of intra-chunk and inter-chunk terms (diagonal and off-diagonal blocks)
+    output = Y_diag + Y_off
+    output = output.reshape(batch_size, -1, num_heads, head_dim)
+
+    if D_residual is not None:
+        output = output + D_residual
+
+    # Cutting off padded chunks
+    if pad_size > 0:
+        output = output[:, :sequence_length]
+
+    output = output.to(input_dtype)
+
+    if return_final_states:
+        return output, final_state
+
+    return output
+
+
+@use_kernelized_func(
+    [
+        causal_conv1d_fn,
+        causal_conv1d_update,
+        zamba2_split_conv1d_scan_combined,
+        zamba2_selective_state_update,
+        zamba2_chunk_scan,
+    ]
+)
 class Zamba2MambaMixer(nn.Module):
     """
     Compute ∆, A, B, C, and D the state space parameters and compute the `contextualized_states`.
@@ -510,37 +720,6 @@ def __init__(self, config: Zamba2Config, layer_idx: int | None = None, initializ
             self.init_zamba2_weights()
         self.out_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=config.add_bias_linear)
 
-        global causal_conv1d, causal_conv1d_update, causal_conv1d_fn
-        causal_conv1d = lazy_load_kernel("causal-conv1d")
-        causal_conv1d_update = getattr(causal_conv1d, "causal_conv1d_update", causal_conv1d_update)
-        causal_conv1d_fn = getattr(causal_conv1d, "causal_conv1d_fn", causal_conv1d_fn)
-
-        global mamba_ssm, selective_state_update, mamba_chunk_scan_combined, mamba_split_conv1d_scan_combined
-        mamba_ssm = lazy_load_kernel("mamba-ssm")
-        selective_state_update = resolve_internal_import(
-            mamba_ssm, chained_path="ops.triton.selective_state_update.selective_state_update"
-        )
-        mamba_chunk_scan_combined = resolve_internal_import(
-            mamba_ssm, chained_path="ops.triton.ssd_combined.mamba_chunk_scan_combined"
-        )
-        mamba_split_conv1d_scan_combined = resolve_internal_import(
-            mamba_ssm, chained_path="ops.triton.ssd_combined.mamba_split_conv1d_scan_combined"
-        )
-
-        global is_fast_path_available
-        is_fast_path_available = (
-            all((selective_state_update, mamba_chunk_scan_combined, mamba_split_conv1d_scan_combined))
-            and hasattr(causal_conv1d, "causal_conv1d_update")
-            and hasattr(causal_conv1d, "causal_conv1d_fn")
-        )
-
-        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"
-            )
-
         self.layer_type = config.layer_types[layer_idx]
 
     @torch.no_grad()
@@ -550,65 +729,28 @@ def init_zamba2_weights(self):
         init.ones_(self.D)
         init.ones_(self.dt_bias)
 
-    def _convolution(
+    @force_accelerate_hooks("conv1d")
+    def forward(
         self,
         hidden_states: torch.Tensor,
         cache_params: Cache | None = None,
-        attention_mask: torch.LongTensor | None = None,
+        attention_mask: torch.Tensor | None = None,
         **kwargs,
     ):
-        seq_len = hidden_states.shape[1]
-        hidden_states = apply_mask_to_padding_states(hidden_states, attention_mask)
-        hidden_states = hidden_states.transpose(1, 2)
-
+        batch_size, seq_len, _ = hidden_states.shape
+        dtype = hidden_states.dtype
         use_precomputed_states = cache_params is not None and cache_params.has_previous_state(self.layer_idx)
 
-        if use_precomputed_states and seq_len == 1 and not cache_params.layers[self.layer_idx].record_past:
-            conv_state = cache_params.layers[self.layer_idx].conv_states[0]
-            hidden_states = causal_conv1d_update(
-                hidden_states,
-                conv_state,
-                self.conv1d.weight.squeeze(1),
-                self.conv1d.bias,
-                self.activation,
-            )
-        else:
-            if cache_params is not None:
-                hidden_states = cache_params.update_conv_state(
-                    hidden_states, self.layer_idx, conv_kernel_size=self.conv_kernel_size
-                )
-
-            hidden_states = causal_conv1d_fn(
-                hidden_states,
-                self.conv1d.weight.squeeze(1),
-                self.conv1d.bias,
-                activation=self.activation,
-                seq_idx=kwargs.get("seq_idx"),
-            )
-
-            # Drop the additional previous states
-            if cache_params is not None:
-                hidden_states = hidden_states[:, :, -seq_len:]
-
-        hidden_states = hidden_states.transpose(1, 2)
-        return hidden_states
-
-    def cuda_kernels_forward(
-        self,
-        hidden_states: torch.Tensor,
-        cache_params: Cache | None = None,
-        attention_mask: torch.Tensor | None = None,
-        **kwargs,
-    ):
         # 1. Gated MLP's linear projection
         hidden_states = apply_mask_to_padding_states(hidden_states, attention_mask)
         projected_states = self.in_proj(hidden_states)
 
         A = -torch.exp(self.A_log.float())
-        dt_limit_kwargs = {} if self.time_step_limit == (0.0, float("inf")) else {"dt_limit": self.time_step_limit}
-        # Fused kernel for conv1d, SSM, and the final projection
+        fused_kwargs = (
+            kwargs | {} if self.time_step_limit == (0.0, float("inf")) else kwargs | {"dt_limit": self.time_step_limit}
+        )
         if self.training and cache_params is None:
-            return mamba_split_conv1d_scan_combined(
+            fused_output = zamba2_split_conv1d_scan_combined(
                 projected_states,
                 self.conv1d.weight.squeeze(1),
                 self.conv1d.bias,
@@ -616,7 +758,6 @@ def cuda_kernels_forward(
                 A,
                 D=self.D,
                 chunk_size=self.chunk_size,
-                seq_idx=kwargs.get("seq_idx"),
                 activation=self.activation,
                 rmsnorm_weight=self.norm.weight,
                 rmsnorm_eps=self.norm.variance_epsilon,
@@ -626,41 +767,71 @@ def cuda_kernels_forward(
                 ngroups=self.n_groups,
                 norm_before_gate=False,
                 return_final_states=False,
-                **dt_limit_kwargs,
+                **fused_kwargs,
             )
 
-        # Set up dimensions for reshapes later
-        batch_size, seq_len, _ = hidden_states.shape
-        groups_time_state_size = self.n_groups * self.ssm_state_size
-        use_precomputed_states = cache_params is not None and cache_params.has_previous_state(self.layer_idx)
+            # Only kernels can use this shortcircuit, fallback to normal torch otherwise
+            if fused_output is not None:
+                return fused_output
 
         gate, hidden_states_B_C, dt = projected_states.split(
             [self.intermediate_size, self.conv_dim, self.num_heads], dim=-1
         )
 
-        # Apply the conv
-        hidden_states_B_C = self._convolution(hidden_states_B_C, cache_params, attention_mask, **kwargs)
-        hidden_states_B_C = apply_mask_to_padding_states(hidden_states_B_C, attention_mask)
+        if use_precomputed_states:
+            conv_state = cache_params.layers[self.layer_idx].conv_states[0]
+            recurrent_state = cache_params.layers[self.layer_idx].recurrent_states[0]
+
+        # 2. Convolution sequence transformation
+        hidden_states_B_C = hidden_states_B_C.transpose(1, 2)
+        if use_precomputed_states and seq_len == 1:
+            hidden_states_B_C = causal_conv1d_update(
+                hidden_states_B_C,
+                conv_state,
+                self.conv1d.weight.squeeze(1),
+                self.conv1d.bias,
+                activation=self.activation,
+            )
+        else:
+            if cache_params is not None:
+                hidden_states_B_C = cache_params.update_conv_state(
+                    hidden_states_B_C,
+                    self.layer_idx,
+                    conv_kernel_size=self.conv_kernel_size,
+                )
+
+            hidden_states_B_C = causal_conv1d_fn(
+                hidden_states_B_C,
+                self.conv1d.weight.squeeze(1),
+                self.conv1d.bias,
+                activation=self.activation,
+                **kwargs,
+            )
+
+            if cache_params is not None:
+                hidden_states_B_C = hidden_states_B_C[:, :, -seq_len:]
+
+        # 3. SSM transformation
+        hidden_states_B_C = apply_mask_to_padding_states(hidden_states_B_C.transpose(1, 2), attention_mask)
         hidden_states, B, C = torch.split(
             hidden_states_B_C,
-            [self.intermediate_size, groups_time_state_size, groups_time_state_size],
+            [self.intermediate_size, self.n_groups * self.ssm_state_size, self.n_groups * self.ssm_state_size],
             dim=-1,
         )
 
-        recurrent_state = cache_params.layers[self.layer_idx].recurrent_states[0] if use_precomputed_states else None
-        # Single step calculations via cache
+        # Recurrent form
         if use_precomputed_states and seq_len == 1:
-            # 3. SSM transformation
-            A = A[:, None, ...][:, :, None].expand(-1, self.head_dim, self.ssm_state_size).to(dtype=torch.float32)
+            hidden_states = hidden_states.view(batch_size, self.num_heads, self.head_dim)
             dt = dt.transpose(1, 2).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[2] // self.n_groups)
-            C = C.view(batch_size, self.n_groups, C.shape[2] // self.n_groups)
-            hidden_states_reshaped = hidden_states.view(batch_size, self.num_heads, self.head_dim)
-            hidden_states = selective_state_update(
+            A = A[:, None, None].expand(-1, self.head_dim, self.ssm_state_size)
+            B = B.view(batch_size, self.n_groups, self.ssm_state_size)
+            C = C.view(batch_size, self.n_groups, self.ssm_state_size)
+            D = self.D[:, None].expand(-1, self.head_dim)
+            dt_bias = self.dt_bias[:, None].expand(-1, self.head_dim)
+
+            scan_output = zamba2_selective_state_update(
                 recurrent_state,
-                hidden_states_reshaped,
+                hidden_states,
                 dt,
                 A,
                 B,
@@ -669,224 +840,44 @@ def cuda_kernels_forward(
                 z=None,
                 dt_bias=dt_bias,
                 dt_softplus=True,
+                **kwargs,
             )
-            hidden_states = hidden_states.view(batch_size, 1, self.num_heads * self.head_dim)
-            hidden_states = self.norm(hidden_states, gate)
-
-            # 4. Final linear projection
-            out = self.out_proj(hidden_states)
+            scan_output = scan_output.view(batch_size, 1, -1)
 
-        # Fused calculations or step by step if no initialized cache is found
+        # Chunk form
         else:
-            # 3. SSM transformation
-            scan_output, ssm_state = mamba_chunk_scan_combined(
-                hidden_states.view(batch_size, seq_len, -1, self.head_dim),
+            output_final_state = cache_params is not None
+            scan_result = zamba2_chunk_scan(
+                hidden_states.view(batch_size, seq_len, self.num_heads, 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),
+                B.view(batch_size, seq_len, self.n_groups, self.ssm_state_size),
+                C.view(batch_size, seq_len, self.n_groups, self.ssm_state_size),
                 chunk_size=self.chunk_size,
                 D=self.D,
                 z=None,
-                seq_idx=kwargs.get("seq_idx"),
-                return_final_states=True,
+                return_final_states=output_final_state,
                 dt_bias=self.dt_bias,
                 dt_softplus=True,
-                initial_states=recurrent_state,
-                **dt_limit_kwargs,
+                initial_states=recurrent_state if use_precomputed_states else None,
+                dt_limit=self.time_step_limit,
+                **kwargs,
             )
 
-            # Init cache
-            if ssm_state is not None and cache_params is not None:
-                cache_params.update_recurrent_state(ssm_state, layer_idx=self.layer_idx)
-
-            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
-
-    def torch_forward(
-        self,
-        hidden_states: torch.Tensor,
-        cache_params: Cache | None = None,
-        attention_mask: torch.Tensor | None = None,
-        **kwargs,
-    ):
-        batch_size, seq_len, _ = hidden_states.shape
-        dtype = hidden_states.dtype
-        use_precomputed_states = cache_params is not None and cache_params.has_previous_state(self.layer_idx)
+            if output_final_state:
+                scan_output, final_state = scan_result
+                cache_params.update_recurrent_state(final_state, layer_idx=self.layer_idx)
+            else:
+                scan_output = scan_result
 
-        # 1. Gated MLP's linear projection
-        hidden_states = apply_mask_to_padding_states(hidden_states, attention_mask)
-        projected_states = self.in_proj(hidden_states)
+            scan_output = scan_output.reshape(batch_size, seq_len, -1)
 
-        gate, hidden_states_B_C, dt = projected_states.split(
-            [self.intermediate_size, self.conv_dim, self.num_heads], dim=-1
-        )
-
-        # 2. Convolution sequence transformation
-        hidden_states_B_C = self._convolution(hidden_states_B_C, cache_params, attention_mask, **kwargs)
-        hidden_states_B_C = apply_mask_to_padding_states(hidden_states_B_C, attention_mask)
-        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())
-        if use_precomputed_states and seq_len == 1:
-            # We need to guarantee that anything regarding the cache is on the same device
-            cache_device = cache_params.layers[self.layer_idx].device
-
-            # Note: there is no need to pad parameter matrices here, as there is just one new token for batched generation
-            dt = dt.transpose(1, 2).expand(batch_size, dt.shape[-1], self.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))[..., None]
-            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)
-            dA = (torch.exp(dt * A)).to(device=cache_device)
-
-            # Discretize B
-            B = B.reshape(batch_size, self.n_groups, 1, -1)
-            B = B.expand(batch_size, self.n_groups, self.num_heads // self.n_groups, B.shape[-1]).contiguous()
-            B = B.reshape(batch_size, -1, 1, B.shape[-1])
-            dB = dt * B
-
-            # Discretize x into dB
-            hidden_states = hidden_states.reshape(batch_size, -1, self.head_dim)
-            dBx = (dB * hidden_states[..., None]).to(device=cache_device)
-
-            # State calculation
-            ssm_states = cache_params.layers[self.layer_idx].recurrent_states[0] * dA + dBx
-            ssm_states = cache_params.update_recurrent_state(ssm_states, layer_idx=self.layer_idx)
-
-            # Subsequent output
-            C = C.reshape(batch_size, self.n_groups, 1, -1)
-            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])
-
-            # Reshape ssm_states to merge the first two dimensions
-            ssm_states = ssm_states.to(device=C.device, dtype=C.dtype)
-            ssm_states_reshaped = ssm_states.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)
-            y = torch.bmm(ssm_states_reshaped, C_reshaped)
-            y = y.view(batch_size, self.num_heads, self.head_dim)
-
-            # D skip connection
-            D = self.D[..., None].expand(self.D.shape[0], self.head_dim)
-            y = (y + hidden_states * D).to(y.dtype)
-
-            y = y.reshape(batch_size, 1, -1)
-        else:
-            # 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_interleave(self.num_heads // self.n_groups, dim=2, output_size=self.num_heads)
-            C = C.repeat_interleave(self.num_heads // self.n_groups, dim=2, output_size=self.num_heads)
-            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)
-            ]
-
-            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)
-            G_intermediate = C[:, :, :, None, :, :] * B[:, :, None, :, :, :]
-            G = G_intermediate.sum(dim=-1)
-
-            # 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 = (
-                cache_params.layers[self.layer_idx]
-                .recurrent_states[0][:, None]
-                .to(dtype=states.dtype, device=states.device)
-                if use_precomputed_states
-                else 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
-            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)
-
-            # Init cache
-            if ssm_state is not None and cache_params is not None:
-                cache_params.update_recurrent_state(ssm_state, layer_idx=self.layer_idx)
-
-        scan_output = self.norm(y, gate)
+        scan_output = self.norm(scan_output, gate)
 
         # 4. Final linear projection
         contextualized_states = self.out_proj(scan_output.to(dtype))
         return contextualized_states
 
-    @force_accelerate_hooks("conv1d")
-    def forward(
-        self,
-        hidden_states,
-        cache_params: Cache | None = None,
-        attention_mask: torch.Tensor | None = None,
-        **kwargs,
-    ):
-        if is_fast_path_available and "cuda" in self.in_proj.weight.device.type and not is_torchdynamo_compiling():
-            return self.cuda_kernels_forward(hidden_states, cache_params, attention_mask, **kwargs)
-        if kwargs.get("seq_idx") is not None:
-            raise NotImplementedError(
-                "`seq_idx` support requires fast path support. Please install `mamba_ssm` and `causal_conv1d`"
-            )
-        return self.torch_forward(hidden_states, cache_params, attention_mask, **kwargs)
-
 
 class Zamba2MLP(nn.Module):
     def __init__(self, config: Zamba2Config, num_fwd_mem_blocks=None, block_id: int | None = None):

From 0eca0d072738e548b67c9af71b64d445af8b0a0c Mon Sep 17 00:00:00 2001
From: vasqu <antonprogamer@gmail.com>
Date: Thu, 30 Jul 2026 14:49:10 +0000
Subject: [PATCH 21/43] oops

---
 tests/models/olmo_hybrid/test_modeling_olmo_hybrid.py | 7 +------
 1 file changed, 1 insertion(+), 6 deletions(-)

diff --git a/tests/models/olmo_hybrid/test_modeling_olmo_hybrid.py b/tests/models/olmo_hybrid/test_modeling_olmo_hybrid.py
index 1e96ae7b2f09..b56d572c29bf 100644
--- a/tests/models/olmo_hybrid/test_modeling_olmo_hybrid.py
+++ b/tests/models/olmo_hybrid/test_modeling_olmo_hybrid.py
@@ -69,12 +69,7 @@ def _get_conv_state_shape(self, batch_size: int, config):
         conv_kernel = config.linear_conv_kernel_dim
         key_dim = config.linear_key_head_dim * config.linear_num_key_heads
         value_dim = config.linear_value_head_dim * config.linear_num_value_heads
-        # We have 3 conv states per layer, with different shapes
-        return [
-            (batch_size, key_dim, conv_kernel),
-            (batch_size, key_dim, conv_kernel),
-            (batch_size, value_dim, conv_kernel),
-        ]
+        return (batch_size, key_dim * 2 + value_dim, conv_kernel,)
 
     def _get_recurrent_state_shape(self, batch_size: int, config):
         return (batch_size, config.linear_num_value_heads, config.linear_key_head_dim, config.linear_value_head_dim)

From 82bfadecd0b136c8a75f73efe75ab949d2802165 Mon Sep 17 00:00:00 2001
From: vasqu <antonprogamer@gmail.com>
Date: Thu, 30 Jul 2026 15:04:46 +0000
Subject: [PATCH 22/43] conv1ds across other models

---
 .../models/inkling/modeling_inkling.py        | 11 +++--
 .../models/inkling/modular_inkling.py         | 47 +-----------------
 src/transformers/models/lfm2/modeling_lfm2.py |  9 ++--
 src/transformers/models/lfm2/modular_lfm2.py  | 49 ++-----------------
 .../models/lfm2_moe/modeling_lfm2_moe.py      | 14 ++++--
 .../models/lfm2_moe/modular_lfm2_moe.py       | 11 -----
 .../olmo_hybrid/test_modeling_olmo_hybrid.py  |  6 ++-
 7 files changed, 33 insertions(+), 114 deletions(-)

diff --git a/src/transformers/models/inkling/modeling_inkling.py b/src/transformers/models/inkling/modeling_inkling.py
index 4b42c90ff177..6df67c93cbb9 100644
--- a/src/transformers/models/inkling/modeling_inkling.py
+++ b/src/transformers/models/inkling/modeling_inkling.py
@@ -31,7 +31,12 @@
 from ...activations import ACT2FN
 from ...cache_utils import Cache, DynamicCache
 from ...generation import GenerationMixin
-from ...integrations import use_experts_implementation, use_kernel_forward_from_hub, use_kernelized_func
+from ...integrations import (
+    use_experts_implementation,
+    use_kernel_forward_from_hub,
+    use_kernel_func_from_hub_with_fallback,
+    use_kernelized_func,
+)
 from ...integrations.accelerate import force_accelerate_hooks
 from ...masking_utils import create_causal_mask, create_recurrent_attention_mask, create_sliding_window_causal_mask
 from ...modeling_layers import GradientCheckpointingLayer
@@ -432,7 +437,7 @@ def apply_mask_to_padding_states(hidden_states, attention_mask):
     return hidden_states
 
 
-@use_kernel_forward_from_hub("causal_conv1d_update")
+@use_kernel_func_from_hub_with_fallback("causal_conv1d_update", "causal_conv1d")
 def causal_conv1d_update(
     hidden_states: torch.Tensor,
     conv_state: torch.Tensor,
@@ -452,7 +457,7 @@ def causal_conv1d_update(
     return out.to(hidden_states.dtype)
 
 
-@use_kernel_forward_from_hub("causal_conv1d_fn")
+@use_kernel_func_from_hub_with_fallback("causal_conv1d_fn", "causal_conv1d")
 def causal_conv1d_fn(
     hidden_states: torch.Tensor,
     weight: nn.Parameter,
diff --git a/src/transformers/models/inkling/modular_inkling.py b/src/transformers/models/inkling/modular_inkling.py
index 42fb95a03197..f0f94e17cd38 100644
--- a/src/transformers/models/inkling/modular_inkling.py
+++ b/src/transformers/models/inkling/modular_inkling.py
@@ -26,7 +26,7 @@
 from ...cache_utils import Cache, DynamicCache
 from ...configuration_utils import PreTrainedConfig
 from ...generation import GenerationMixin
-from ...integrations import use_kernel_forward_from_hub, use_kernelized_func
+from ...integrations import use_kernelized_func
 from ...integrations.accelerate import force_accelerate_hooks
 from ...masking_utils import create_causal_mask, create_recurrent_attention_mask, create_sliding_window_causal_mask
 from ...modeling_layers import GradientCheckpointingLayer
@@ -51,7 +51,7 @@
 from ..higgs_audio_v2.modeling_higgs_audio_v2 import HiggsAudioV2Embeddings
 from ..llama.modeling_llama import LlamaRMSNorm, repeat_kv
 from ..mixtral.modeling_mixtral import MixtralExperts
-from ..qwen3_next.modeling_qwen3_next import apply_mask_to_padding_states
+from ..qwen3_next.modeling_qwen3_next import apply_mask_to_padding_states, causal_conv1d_fn, causal_conv1d_update
 
 
 logger = logging.get_logger(__name__)
@@ -533,49 +533,6 @@ def forward(self, hidden_states) -> torch.Tensor:
         return hidden_states
 
 
-@use_kernel_forward_from_hub("causal_conv1d_update")
-def causal_conv1d_update(
-    hidden_states: torch.Tensor,
-    conv_state: torch.Tensor,
-    weight: nn.Parameter,
-    bias: nn.Parameter | None = None,
-    activation: str | None = None,
-):
-    _, hidden_size, seq_len = hidden_states.shape
-    state_len = conv_state.shape[-1]
-
-    hidden_states_new = torch.cat([conv_state, hidden_states], dim=-1).to(weight.dtype)
-    conv_state.copy_(hidden_states_new[:, :, -state_len:])
-    out = F.conv1d(hidden_states_new, weight.unsqueeze(1), bias, padding=0, groups=hidden_size)
-    out = out[:, :, -seq_len:]
-    if activation is not None:
-        out = ACT2FN[activation](out)
-    return out.to(hidden_states.dtype)
-
-
-@use_kernel_forward_from_hub("causal_conv1d_fn")
-def causal_conv1d_fn(
-    hidden_states: torch.Tensor,
-    weight: nn.Parameter,
-    bias: nn.Parameter | None = None,
-    activation: str | None = None,
-    **kwargs,
-):
-    _, hidden_size, seq_len = hidden_states.shape
-    padding = weight.shape[-1] - 1
-
-    out = F.conv1d(
-        hidden_states.to(weight.dtype),
-        weight=weight.unsqueeze(1),
-        bias=bias,
-        padding=padding,
-        groups=hidden_size,
-    )[:, :, :seq_len]
-    if activation is not None:
-        out = ACT2FN[activation](out)
-    return out.to(hidden_states.dtype)
-
-
 @use_kernelized_func([causal_conv1d_update, causal_conv1d_fn])
 class InklingShortConvolution(nn.Module):
     def __init__(self, hidden_size: int, conv_kernel_size: int, layer_idx: int, conv_idx: int):
diff --git a/src/transformers/models/lfm2/modeling_lfm2.py b/src/transformers/models/lfm2/modeling_lfm2.py
index 644c9985650b..a0fef79b7f73 100644
--- a/src/transformers/models/lfm2/modeling_lfm2.py
+++ b/src/transformers/models/lfm2/modeling_lfm2.py
@@ -26,7 +26,7 @@
 from ...activations import ACT2FN
 from ...cache_utils import Cache, DynamicCache
 from ...generation import GenerationMixin
-from ...integrations import use_kernel_forward_from_hub, use_kernelized_func
+from ...integrations import use_kernel_forward_from_hub, use_kernel_func_from_hub_with_fallback, use_kernelized_func
 from ...integrations.accelerate import force_accelerate_hooks
 from ...masking_utils import create_causal_mask, create_recurrent_attention_mask
 from ...modeling_layers import GradientCheckpointingLayer
@@ -36,7 +36,7 @@
 from ...processing_utils import Unpack
 from ...utils import TransformersKwargs, auto_docstring, can_return_tuple
 from ...utils.deprecation import deprecate_kwarg
-from ...utils.generic import maybe_autocast, maybe_replace_from_package, merge_with_config_defaults
+from ...utils.generic import maybe_autocast, merge_with_config_defaults
 from ...utils.output_capturing import capture_outputs
 from .configuration_lfm2 import Lfm2Config
 
@@ -279,7 +279,7 @@ def apply_mask_to_padding_states(hidden_states, attention_mask):
     return hidden_states
 
 
-@maybe_replace_from_package("causal_conv1d", "causal_conv1d_update")
+@use_kernel_func_from_hub_with_fallback("causal_conv1d_update", "causal_conv1d")
 def causal_conv1d_update(
     hidden_states: torch.Tensor,
     conv_state: torch.Tensor,
@@ -299,7 +299,7 @@ def causal_conv1d_update(
     return out.to(hidden_states.dtype)
 
 
-@maybe_replace_from_package("causal_conv1d", "causal_conv1d_fn")
+@use_kernel_func_from_hub_with_fallback("causal_conv1d_fn", "causal_conv1d")
 def causal_conv1d_fn(
     hidden_states: torch.Tensor,
     weight: nn.Parameter,
@@ -322,6 +322,7 @@ def causal_conv1d_fn(
     return out.to(hidden_states.dtype)
 
 
+@use_kernelized_func([causal_conv1d_update, causal_conv1d_fn])
 class Lfm2ShortConv(nn.Module):
     def __init__(
         self,
diff --git a/src/transformers/models/lfm2/modular_lfm2.py b/src/transformers/models/lfm2/modular_lfm2.py
index 9d87d263f8cd..cf9eae0f119f 100644
--- a/src/transformers/models/lfm2/modular_lfm2.py
+++ b/src/transformers/models/lfm2/modular_lfm2.py
@@ -17,8 +17,8 @@
 import torch.nn.functional as F
 from torch import nn
 
-from ...activations import ACT2FN
 from ...cache_utils import Cache, DynamicCache
+from ...integrations import use_kernelized_func
 from ...integrations.accelerate import force_accelerate_hooks
 from ...masking_utils import create_causal_mask, create_recurrent_attention_mask
 from ...modeling_layers import GradientCheckpointingLayer
@@ -26,8 +26,6 @@
 from ...modeling_utils import ALL_ATTENTION_FUNCTIONS
 from ...processing_utils import Unpack
 from ...utils import TransformersKwargs, logging
-from ...utils.generic import maybe_replace_from_package
-from ..bamba.modeling_bamba import apply_mask_to_padding_states
 from ..gemma2.modeling_gemma2 import Gemma2RotaryEmbedding
 from ..llama.modeling_llama import (
     LlamaAttention,
@@ -38,6 +36,7 @@
     apply_rotary_pos_emb,
     eager_attention_forward,
 )
+from ..qwen3_next.modeling_qwen3_next import apply_mask_to_padding_states, causal_conv1d_fn, causal_conv1d_update
 from .configuration_lfm2 import Lfm2Config
 
 
@@ -124,49 +123,7 @@ def forward(
         return output, attn_weights
 
 
-@maybe_replace_from_package("causal_conv1d", "causal_conv1d_update")
-def causal_conv1d_update(
-    hidden_states: torch.Tensor,
-    conv_state: torch.Tensor,
-    weight: nn.Parameter,
-    bias: nn.Parameter | None = None,
-    activation: str | None = None,
-):
-    _, hidden_size, seq_len = hidden_states.shape
-    state_len = conv_state.shape[-1]
-
-    hidden_states_new = torch.cat([conv_state, hidden_states], dim=-1).to(weight.dtype)
-    conv_state.copy_(hidden_states_new[:, :, -state_len:])
-    out = F.conv1d(hidden_states_new, weight.unsqueeze(1), bias, padding=0, groups=hidden_size)
-    out = out[:, :, -seq_len:]
-    if activation is not None:
-        out = ACT2FN[activation](out)
-    return out.to(hidden_states.dtype)
-
-
-@maybe_replace_from_package("causal_conv1d", "causal_conv1d_fn")
-def causal_conv1d_fn(
-    hidden_states: torch.Tensor,
-    weight: nn.Parameter,
-    bias: nn.Parameter | None = None,
-    activation: str | None = None,
-    **kwargs,
-):
-    _, hidden_size, seq_len = hidden_states.shape
-    padding = weight.shape[-1] - 1
-
-    out = F.conv1d(
-        hidden_states.to(weight.dtype),
-        weight=weight.unsqueeze(1),
-        bias=bias,
-        padding=padding,
-        groups=hidden_size,
-    )[:, :, :seq_len]
-    if activation is not None:
-        out = ACT2FN[activation](out)
-    return out.to(hidden_states.dtype)
-
-
+@use_kernelized_func([causal_conv1d_update, causal_conv1d_fn])
 class Lfm2ShortConv(nn.Module):
     def __init__(
         self,
diff --git a/src/transformers/models/lfm2_moe/modeling_lfm2_moe.py b/src/transformers/models/lfm2_moe/modeling_lfm2_moe.py
index 9cb0e09ce9a3..327699f253d8 100644
--- a/src/transformers/models/lfm2_moe/modeling_lfm2_moe.py
+++ b/src/transformers/models/lfm2_moe/modeling_lfm2_moe.py
@@ -28,7 +28,12 @@
 from ...activations import ACT2FN
 from ...cache_utils import Cache, DynamicCache
 from ...generation import GenerationMixin
-from ...integrations import use_experts_implementation, use_kernel_forward_from_hub, use_kernelized_func
+from ...integrations import (
+    use_experts_implementation,
+    use_kernel_forward_from_hub,
+    use_kernel_func_from_hub_with_fallback,
+    use_kernelized_func,
+)
 from ...integrations.accelerate import force_accelerate_hooks
 from ...masking_utils import create_causal_mask, create_recurrent_attention_mask
 from ...modeling_layers import GradientCheckpointingLayer
@@ -38,7 +43,7 @@
 from ...processing_utils import Unpack
 from ...utils import TransformersKwargs, auto_docstring, can_return_tuple
 from ...utils.deprecation import deprecate_kwarg
-from ...utils.generic import maybe_autocast, maybe_replace_from_package, merge_with_config_defaults
+from ...utils.generic import maybe_autocast, merge_with_config_defaults
 from ...utils.output_capturing import capture_outputs
 from .configuration_lfm2_moe import Lfm2MoeConfig
 
@@ -359,7 +364,7 @@ def apply_mask_to_padding_states(hidden_states, attention_mask):
     return hidden_states
 
 
-@maybe_replace_from_package("causal_conv1d", "causal_conv1d_update")
+@use_kernel_func_from_hub_with_fallback("causal_conv1d_update", "causal_conv1d")
 def causal_conv1d_update(
     hidden_states: torch.Tensor,
     conv_state: torch.Tensor,
@@ -379,7 +384,7 @@ def causal_conv1d_update(
     return out.to(hidden_states.dtype)
 
 
-@maybe_replace_from_package("causal_conv1d", "causal_conv1d_fn")
+@use_kernel_func_from_hub_with_fallback("causal_conv1d_fn", "causal_conv1d")
 def causal_conv1d_fn(
     hidden_states: torch.Tensor,
     weight: nn.Parameter,
@@ -402,6 +407,7 @@ def causal_conv1d_fn(
     return out.to(hidden_states.dtype)
 
 
+@use_kernelized_func([causal_conv1d_update, causal_conv1d_fn])
 class Lfm2MoeShortConv(nn.Module):
     def __init__(
         self,
diff --git a/src/transformers/models/lfm2_moe/modular_lfm2_moe.py b/src/transformers/models/lfm2_moe/modular_lfm2_moe.py
index 7649a495abe9..04281469b953 100644
--- a/src/transformers/models/lfm2_moe/modular_lfm2_moe.py
+++ b/src/transformers/models/lfm2_moe/modular_lfm2_moe.py
@@ -23,7 +23,6 @@
 from ...modeling_utils import PreTrainedModel
 from ...processing_utils import Unpack
 from ...utils import TransformersKwargs, logging
-from ...utils.import_utils import is_causal_conv1d_available
 from ..lfm2.modeling_lfm2 import (
     Lfm2Attention,
     Lfm2DecoderLayer,
@@ -38,16 +37,6 @@
 from .configuration_lfm2_moe import Lfm2MoeConfig
 
 
-if is_causal_conv1d_available():
-    from causal_conv1d import causal_conv1d_fn, causal_conv1d_update
-else:
-    causal_conv1d_fn, causal_conv1d_update = None, None
-
-
-kernel_modules = (causal_conv1d_fn, causal_conv1d_update)
-is_fast_path_available = all(kernel_modules)
-
-
 logger = logging.get_logger(__name__)
 
 
diff --git a/tests/models/olmo_hybrid/test_modeling_olmo_hybrid.py b/tests/models/olmo_hybrid/test_modeling_olmo_hybrid.py
index b56d572c29bf..1efa84e4547d 100644
--- a/tests/models/olmo_hybrid/test_modeling_olmo_hybrid.py
+++ b/tests/models/olmo_hybrid/test_modeling_olmo_hybrid.py
@@ -69,7 +69,11 @@ def _get_conv_state_shape(self, batch_size: int, config):
         conv_kernel = config.linear_conv_kernel_dim
         key_dim = config.linear_key_head_dim * config.linear_num_key_heads
         value_dim = config.linear_value_head_dim * config.linear_num_value_heads
-        return (batch_size, key_dim * 2 + value_dim, conv_kernel,)
+        return (
+            batch_size,
+            key_dim * 2 + value_dim,
+            conv_kernel,
+        )
 
     def _get_recurrent_state_shape(self, batch_size: int, config):
         return (batch_size, config.linear_num_value_heads, config.linear_key_head_dim, config.linear_value_head_dim)

From 520e340b00ae3d3619b21fc535f1d7298202f9fb Mon Sep 17 00:00:00 2001
From: vasqu <antonprogamer@gmail.com>
Date: Thu, 30 Jul 2026 16:22:14 +0000
Subject: [PATCH 23/43] let's try this

---
 src/transformers/integrations/hub_kernels.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/transformers/integrations/hub_kernels.py b/src/transformers/integrations/hub_kernels.py
index d8f3da0e095b..9ddf3a54fc1c 100644
--- a/src/transformers/integrations/hub_kernels.py
+++ b/src/transformers/integrations/hub_kernels.py
@@ -740,7 +740,7 @@ def decorator(torch_function: Callable) -> Callable:
         finally:
             implementation = torch_function if implementation is None else implementation
 
-        applicable_params = inspect.signature(implementation).parameters
+        applicable_params = tuple(inspect.signature(implementation).parameters)
 
         @functools.wraps(torch_function)
         def wrapped(*args, **kwargs):

From f69acfedac898a3a181cdc17921c9a76e9ab12cb Mon Sep 17 00:00:00 2001
From: vasqu <antonprogamer@gmail.com>
Date: Thu, 30 Jul 2026 18:18:18 +0000
Subject: [PATCH 24/43] mamba base implementation

---
 src/transformers/integrations/hub_kernels.py  |  29 +
 .../models/mamba/modeling_mamba.py            | 501 ++++++++++--------
 2 files changed, 299 insertions(+), 231 deletions(-)

diff --git a/src/transformers/integrations/hub_kernels.py b/src/transformers/integrations/hub_kernels.py
index 9ddf3a54fc1c..835e7fc5e9b6 100644
--- a/src/transformers/integrations/hub_kernels.py
+++ b/src/transformers/integrations/hub_kernels.py
@@ -235,6 +235,34 @@ def _build_kernel_mapping() -> dict:
                     ),
                 },
             },
+            "mamba_inner_fn": {
+                "cuda": {
+                    Mode.TRAINING: LayerRepository(
+                        repo_id="kernels-community/mamba-ssm",
+                        layer_name="mamba_inner_fn",
+                        version=1,
+                    ),
+                    Mode.INFERENCE: LayerRepository(
+                        repo_id="kernels-community/mamba-ssm",
+                        layer_name="mamba_inner_fn",
+                        version=1,
+                    ),
+                },
+            },
+            "selective_scan_fn": {
+                "cuda": {
+                    Mode.TRAINING: LayerRepository(
+                        repo_id="kernels-community/mamba-ssm",
+                        layer_name="selective_scan_fn",
+                        version=1,
+                    ),
+                    Mode.INFERENCE: LayerRepository(
+                        repo_id="kernels-community/mamba-ssm",
+                        layer_name="selective_scan_fn",
+                        version=1,
+                    ),
+                },
+            },
             "selective_state_update": {
                 "cuda": {
                     Mode.TRAINING: LayerRepository(
@@ -740,6 +768,7 @@ def decorator(torch_function: Callable) -> Callable:
         finally:
             implementation = torch_function if implementation is None else implementation
 
+        # Make it "frozen" like to let dynamo not try to look into any ordering
         applicable_params = tuple(inspect.signature(implementation).parameters)
 
         @functools.wraps(torch_function)
diff --git a/src/transformers/models/mamba/modeling_mamba.py b/src/transformers/models/mamba/modeling_mamba.py
index 45c36e8f9ac6..538d6e6ee129 100644
--- a/src/transformers/models/mamba/modeling_mamba.py
+++ b/src/transformers/models/mamba/modeling_mamba.py
@@ -25,7 +25,7 @@
 from ...activations import ACT2FN
 from ...cache_utils import Cache, DynamicCache
 from ...generation import GenerationMixin
-from ...integrations import lazy_load_kernel
+from ...integrations import use_kernel_func_from_hub_with_fallback, use_kernelized_func
 from ...integrations.accelerate import force_accelerate_hooks
 from ...modeling_layers import GradientCheckpointingLayer
 from ...modeling_utils import PreTrainedModel
@@ -38,25 +38,26 @@
     is_mambapy_available,
     is_torch_greater_or_equal,
     is_tracing,
-    resolve_internal_import,
 )
 from .configuration_mamba import MambaConfig
 
 
 logger = logging.get_logger(__name__)
 
-if is_torch_greater_or_equal("2.9.0"):
-    from torch._higher_order_ops.associative_scan import associative_scan
-else:
-    associative_scan = None
 
+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
+    """
+    # NOTE: attention mask is a 2D boolean tensor
+    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)
 
-if is_mambapy_available():
-    from mambapy.pscan import pscan
-else:
-    pscan = None
+    return hidden_states
 
 
+@use_kernel_func_from_hub_with_fallback("causal_conv1d_update", "causal_conv1d")
 def causal_conv1d_update(
     hidden_states: torch.Tensor,
     conv_state: torch.Tensor,
@@ -76,6 +77,7 @@ def causal_conv1d_update(
     return out.to(hidden_states.dtype)
 
 
+@use_kernel_func_from_hub_with_fallback("causal_conv1d_fn", "causal_conv1d")
 def causal_conv1d_fn(
     hidden_states: torch.Tensor,
     weight: nn.Parameter,
@@ -98,6 +100,187 @@ def causal_conv1d_fn(
     return out.to(hidden_states.dtype)
 
 
+@use_kernel_func_from_hub_with_fallback(
+    "mamba_inner_fn",
+    "mamba_ssm",
+    internal_path="ops.selective_scan_interface",
+)
+def mamba_inner_fn(
+    xz: torch.Tensor,
+    conv1d_weight: torch.Tensor,
+    conv1d_bias: torch.Tensor | None,
+    x_proj_weight: torch.Tensor,
+    delta_proj_weight: torch.Tensor,
+    out_proj_weight: torch.Tensor,
+    out_proj_bias: torch.Tensor | None,
+    A: torch.Tensor,
+    B: torch.Tensor | None = None,
+    C: torch.Tensor | None = None,
+    D: torch.Tensor | None = None,
+    delta_bias: torch.Tensor | None = None,
+    delta_softplus: bool = True,
+    **kwargs,
+):
+    return None
+
+
+@use_kernel_func_from_hub_with_fallback(
+    "selective_state_update",
+    "mamba_ssm",
+    internal_path="ops.triton.selective_state_update",
+)
+def mamba_selective_state_update(
+    state: torch.Tensor,
+    hidden_states: torch.Tensor,
+    dt: torch.Tensor,
+    A: torch.Tensor,
+    B: torch.Tensor,
+    C: torch.Tensor,
+    D: torch.Tensor | None = None,
+    dt_bias: torch.Tensor | None = None,
+    dt_softplus: bool = False,
+    z: torch.Tensor | None = None,
+    **kwargs,
+):
+    input_dtype = hidden_states.dtype
+
+    if dt_bias is not None:
+        dt = dt + dt_bias.to(dt.dtype)
+    if dt_softplus:
+        dt = F.softplus(dt)
+
+    # Discretize A
+    dA = torch.exp(dt.float()[..., None] * A.float()).to(device=state.device)
+
+    # Discretize B
+    dB = dt.float()[..., None] * B.float()[:, None, :]
+    # Discretize x into dB
+    dBx = dB * hidden_states.float()[..., None]
+
+    # State calculation
+    ssm_state = state.float() * dA + dBx
+    state.copy_(ssm_state.to(state.dtype))
+
+    # Subsequent output
+    out = torch.matmul(ssm_state.to(C.dtype), C.unsqueeze(-1)).squeeze(-1)
+
+    # D skip connection
+    if D is not None:
+        out = out + hidden_states * D
+
+    if z is not None:
+        out = out * F.silu(z)
+
+    return out.to(input_dtype)
+
+
+@use_kernel_func_from_hub_with_fallback(
+    "selective_scan_fn",
+    "mamba_ssm",
+    internal_path="ops.selective_scan_interface",
+)
+def mamba_selective_scan(
+    hidden_states: torch.Tensor,
+    dt: torch.Tensor,
+    A: torch.Tensor,
+    B: torch.Tensor,
+    C: torch.Tensor,
+    D: torch.Tensor | None = None,
+    z: torch.Tensor | None = None,
+    delta_bias: torch.Tensor | None = None,
+    delta_softplus: bool = False,
+    return_last_state: bool = False,
+    use_mambapy: bool = False,
+    use_associative_scan: bool = False,
+    **kwargs,
+):
+    # Torch only alternatives to the recurrent path
+    if is_torch_greater_or_equal("2.9.0"):
+        from torch._higher_order_ops.associative_scan import associative_scan
+    else:
+        associative_scan = None
+
+    if is_mambapy_available():
+        from mambapy.pscan import pscan
+    else:
+        pscan = None
+
+    batch_size, intermediate_size, seq_len = hidden_states.shape
+    input_dtype = hidden_states.dtype
+
+    if delta_bias is not None:
+        dt = dt + delta_bias.to(dt.dtype)[..., None]
+    if delta_softplus:
+        dt = F.softplus(dt)
+
+    # Discretize A and B for the entire sequence
+    discrete_A = torch.exp(A[None, :, None, :] * dt[:, :, :, None])
+    discrete_B = dt[:, :, :, None] * B[:, None, :, :].float()
+    deltaB_u = discrete_B * hidden_states[:, :, :, None].float()
+
+    # TODO: check these out
+    if use_mambapy and pscan is not None:
+        all_states = pscan(discrete_A.transpose(1, 2), deltaB_u.transpose(1, 2))
+
+        scan_output = (all_states @ C.unsqueeze(-1)).squeeze(3).transpose(1, 2)
+        ssm_state = all_states[:, -1]
+
+    elif use_associative_scan and associative_scan is not None and is_tracing(hidden_states):
+
+        def combine_fn(left, right):
+            a_left, b_left = left
+            a_right, b_right = right
+            return a_left * a_right, a_right * b_left + b_right
+
+        combine_mode = "pointwise" if discrete_A.device.type in ("cuda", "xpu") else "generic"
+        _, all_states = associative_scan(
+            combine_fn,
+            (discrete_A, deltaB_u),
+            dim=2,
+            combine_mode=combine_mode,
+        )
+
+        scan_output = torch.matmul(all_states.transpose(1, 2).to(input_dtype), C.transpose(1, 2).unsqueeze(-1))
+        scan_output = scan_output.squeeze(-1).transpose(1, 2)
+        ssm_state = all_states[:, :, -1]
+
+    # Recurrent iteration
+    else:
+        # "Initial hidden state" is not supported by the kernel path, so use
+        # the same zero initialization as the kernel
+        ssm_state = torch.zeros(
+            batch_size,
+            intermediate_size,
+            A.shape[-1],
+            dtype=input_dtype,
+            device=hidden_states.device,
+        )
+
+        scan_outputs = []
+        for index in range(seq_len):
+            # State calculation
+            ssm_state = discrete_A[:, :, index] * ssm_state + deltaB_u[:, :, index]
+
+            # Subsequent output
+            scan_output = torch.matmul(ssm_state.to(input_dtype), C[:, :, index].unsqueeze(-1))
+            scan_outputs.append(scan_output[:, :, 0])
+        scan_output = torch.stack(scan_outputs, dim=-1)
+
+    if D is not None:
+        scan_output = scan_output + hidden_states * D[None, :, None]
+
+    if z is not None:
+        scan_output = scan_output * F.silu(z)
+
+    if return_last_state:
+        return scan_output, ssm_state
+
+    return scan_output
+
+
+@use_kernelized_func(
+    [mamba_inner_fn, mamba_selective_scan, mamba_selective_state_update, causal_conv1d_fn, causal_conv1d_update]
+)
 class MambaMixer(nn.Module):
     """
     Compute ∆, A, B, C, and D the state space parameters and compute the `contextualized_states`.
@@ -147,28 +330,6 @@ def __init__(self, config: MambaConfig, layer_idx: int, initialize_mixer_weights
         self.out_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=config.use_bias)
         self.use_bias = config.use_bias
 
-        global causal_conv1d, causal_conv1d_update, causal_conv1d_fn
-        causal_conv1d = lazy_load_kernel("causal-conv1d")
-        causal_conv1d_update = getattr(causal_conv1d, "causal_conv1d_update", causal_conv1d_update)
-        causal_conv1d_fn = getattr(causal_conv1d, "causal_conv1d_fn", causal_conv1d_fn)
-
-        global mamba_ssm, selective_state_update, selective_scan_fn, mamba_inner_fn
-        mamba_ssm = lazy_load_kernel("mamba-ssm")
-        selective_state_update = resolve_internal_import(
-            mamba_ssm, chained_path="ops.triton.selective_state_update.selective_state_update"
-        )
-        selective_scan_fn = getattr(mamba_ssm, "selective_scan_fn", None)
-        mamba_inner_fn = getattr(mamba_ssm, "mamba_inner_fn", None)
-
-        global is_fast_path_available
-        is_fast_path_available = (
-            all((selective_state_update, selective_scan_fn, mamba_inner_fn))
-            and hasattr(causal_conv1d, "causal_conv1d_update")
-            and hasattr(causal_conv1d, "causal_conv1d_fn")
-        )
-
-        self.warn_slow_implementation()
-
         self.layer_type = config.layer_types[layer_idx]
 
     @torch.no_grad()
@@ -193,80 +354,25 @@ def init_mamba_weights(self):
         inv_dt = dt + torch.log(-torch.expm1(-dt))
         init.copy_(self.dt_proj.bias, inv_dt)
 
-    def warn_slow_implementation(self):
-        if not is_fast_path_available:
-            if self.use_mambapy:
-                if is_mambapy_available():
-                    logger.warning_once(
-                        "The fast path is not available because one of `(selective_state_update, selective_scan_fn, causal_conv1d_fn, causal_conv1d_update, mamba_inner_fn)`"
-                        " is None. Falling back to the mamba.py backend. To install follow https://github.com/state-spaces/mamba/#installation for mamba-ssm and"
-                        " install the kernels library using `pip install kernels` or https://github.com/Dao-AILab/causal-conv1d for causal-conv1d"
-                    )
-                else:
-                    raise ImportError(
-                        "use_mambapy is set to True but the mambapy package is not installed. To install it follow https://github.com/alxndrTL/mamba.py."
-                    )
-            else:
-                logger.warning_once(
-                    "The fast path is not available because one of `(selective_state_update, selective_scan_fn, causal_conv1d_fn, causal_conv1d_update, mamba_inner_fn)`"
-                    " is None. Falling back to the sequential implementation of Mamba, as use_mambapy is set to False. To install follow https://github.com/state-spaces/mamba/#installation for mamba-ssm and"
-                    " install the kernels library using `pip install kernels` or https://github.com/Dao-AILab/causal-conv1d for causal-conv1d. For the mamba.py backend, follow https://github.com/alxndrTL/mamba.py."
-                )
-
-    def _convolution(
+    @force_accelerate_hooks("conv1d")
+    def forward(
         self,
         hidden_states: torch.Tensor,
         cache_params: Cache | None = None,
         attention_mask: torch.LongTensor | None = None,
         **kwargs,
     ):
-        seq_len = hidden_states.shape[-1]
-        if attention_mask is not None:
-            hidden_states = hidden_states * attention_mask.unsqueeze(1)
-
+        seq_len = hidden_states.shape[1]
+        dtype = hidden_states.dtype
         use_precomputed_states = cache_params is not None and cache_params.has_previous_state(self.layer_idx)
 
-        if use_precomputed_states and seq_len == 1 and not cache_params.layers[self.layer_idx].record_past:
-            conv_state = cache_params.layers[self.layer_idx].conv_states[0]
-            hidden_states = causal_conv1d_update(
-                hidden_states,
-                conv_state,
-                self.conv1d.weight.squeeze(1),
-                self.conv1d.bias,
-                self.activation,
-            )
-        else:
-            if cache_params is not None:
-                hidden_states = cache_params.update_conv_state(
-                    hidden_states, self.layer_idx, conv_kernel_size=self.conv_kernel_size
-                )
-
-            hidden_states = causal_conv1d_fn(
-                hidden_states,
-                self.conv1d.weight.squeeze(1),
-                self.conv1d.bias,
-                activation=self.activation,
-                seq_idx=kwargs.get("seq_idx"),
-            )
-
-            # Drop the additional previous states
-            if cache_params is not None:
-                hidden_states = hidden_states[:, :, -seq_len:]
-
-        return hidden_states
-
-    def cuda_kernels_forward(
-        self,
-        hidden_states: torch.Tensor,
-        cache_params: Cache | None = None,
-        attention_mask: torch.LongTensor | None = None,
-        **kwargs,
-    ):
         # 1. Gated MLP's linear projection
+        hidden_states = apply_mask_to_padding_states(hidden_states, attention_mask)
         projected_states = self.in_proj(hidden_states).transpose(1, 2)
 
-        if self.training and cache_params is None:  # Doesn't support outputting the states -> used for training
-            return mamba_inner_fn(
+        A = -torch.exp(self.A_log.float())
+        if self.training and cache_params is None:
+            fused_output = mamba_inner_fn(
                 projected_states,
                 self.conv1d.weight,
                 self.conv1d.bias if self.use_conv_bias else None,
@@ -274,7 +380,7 @@ def cuda_kernels_forward(
                 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()),
+                A,
                 None,  # input-dependent B
                 None,  # input-dependent C
                 self.D.float(),
@@ -282,165 +388,98 @@ def cuda_kernels_forward(
                 delta_softplus=True,
             )
 
-        hidden_states, gate = projected_states.chunk(2, dim=1)
-        use_precomputed_states = cache_params is not None and cache_params.has_previous_state(self.layer_idx)
+            # Only kernels can use this shortcircuit, fallback to normal torch otherwise
+            if fused_output is not None:
+                return fused_output
 
-        # Apply the conv
-        hidden_states = self._convolution(hidden_states, cache_params, attention_mask, **kwargs)
+        hidden_states_B_C, gate = projected_states.chunk(2, dim=1)
 
-        if attention_mask is not None:
-            hidden_states = hidden_states * attention_mask.unsqueeze(1)
+        if use_precomputed_states:
+            conv_state = cache_params.layers[self.layer_idx].conv_states[0]
+            recurrent_state = cache_params.layers[self.layer_idx].recurrent_states[0]
 
-        # 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))
+        # 2. Convolution sequence transformation
+        if use_precomputed_states and seq_len == 1:
+            hidden_states_B_C = causal_conv1d_update(
+                hidden_states_B_C,
+                conv_state,
+                self.conv1d.weight.squeeze(1),
+                self.conv1d.bias,
+                activation=self.activation,
+            )
+        else:
+            if cache_params is not None:
+                hidden_states_B_C = cache_params.update_conv_state(
+                    hidden_states_B_C,
+                    self.layer_idx,
+                    conv_kernel_size=self.conv_kernel_size,
+                )
+
+            hidden_states_B_C = causal_conv1d_fn(
+                hidden_states_B_C,
+                self.conv1d.weight.squeeze(1),
+                self.conv1d.bias,
+                activation=self.activation,
+                **kwargs,
+            )
+
+            if cache_params is not None:
+                hidden_states_B_C = hidden_states_B_C[:, :, -seq_len:]
+
+        # 3. SSM transformation
+        hidden_states_B_C = apply_mask_to_padding_states(hidden_states_B_C.transpose(1, 2), attention_mask)
         time_step, B, C = torch.split(
-            ssm_parameters, [self.time_step_rank, self.ssm_state_size, self.ssm_state_size], dim=-1
+            self.x_proj(hidden_states_B_C),
+            [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 use_precomputed_states:
-            scan_outputs = selective_state_update(
-                cache_params.layers[self.layer_idx].recurrent_states[0],
-                hidden_states[..., 0],
-                discrete_time_step[..., 0],
+        time_step = self.dt_proj.weight @ time_step.transpose(1, 2)
+        time_proj_bias = self.dt_proj.bias.float() if self.dt_proj.bias is not None else None
+
+        # Recurrent form
+        if use_precomputed_states and seq_len == 1:
+            scan_output = mamba_selective_state_update(
+                recurrent_state,
+                hidden_states_B_C.transpose(1, 2)[..., 0],
+                time_step[..., 0],
                 A,
                 B[:, 0],
                 C[:, 0],
                 self.D,
-                gate[..., 0],
-                time_proj_bias,
+                z=gate[..., 0],
+                dt_bias=time_proj_bias,
                 dt_softplus=True,
             ).unsqueeze(-1)
+
+        # Full sequence form
         else:
-            scan_outputs, ssm_state = selective_scan_fn(
-                hidden_states,
-                discrete_time_step,
+            output_final_state = cache_params is not None
+            scan_result = mamba_selective_scan(
+                hidden_states_B_C.transpose(1, 2),
+                time_step,
                 A,
                 B.transpose(1, 2),
                 C.transpose(1, 2),
-                self.D.float(),
-                gate,
-                time_proj_bias,
+                D=self.D.float(),
+                z=gate,
+                delta_bias=time_proj_bias,
                 delta_softplus=True,
-                return_last_state=True,
+                return_last_state=output_final_state,
+                use_mambapy=self.use_mambapy,
+                use_associative_scan=self.use_associative_scan,
             )
-            if ssm_state is not None and cache_params is not None:
-                cache_params.update_recurrent_state(ssm_state, self.layer_idx)
 
-        # 4. Final linear projection
-        contextualized_states = self.out_proj(scan_outputs.transpose(1, 2))
-        return contextualized_states
-
-    def slow_forward(
-        self,
-        hidden_states: torch.Tensor,
-        cache_params: Cache | None = None,
-        attention_mask: torch.LongTensor | None = None,
-        **kwargs,
-    ):
-        batch_size, seq_len, _ = hidden_states.shape
-        dtype = hidden_states.dtype
-        use_precomputed_states = cache_params is not None and cache_params.has_previous_state(self.layer_idx)
-        # 1. Gated MLP's linear projection
-        projected_states = self.in_proj(hidden_states).transpose(1, 2)
-        hidden_states, gate = projected_states.chunk(2, dim=1)
-
-        # Apply the convolution
-        hidden_states = self._convolution(hidden_states, cache_params, attention_mask, **kwargs)
-
-        if attention_mask is not None:
-            hidden_states = hidden_states * attention_mask.unsqueeze(1)
-
-        if use_precomputed_states:
-            ssm_state = cache_params.layers[self.layer_idx].recurrent_states[0].clone()
-        else:
-            ssm_state = torch.zeros(
-                (batch_size, self.intermediate_size, self.ssm_state_size), device=hidden_states.device, dtype=dtype
-            )
-
-        # 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
-        )
-        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)
-        A = -torch.exp(self.A_log.float())  # [intermediate_size, ssm_state_size]
-        # Both discrete_A/B are shape [batch, intermediate_size, seq_len, ssm_state_size]
-        discrete_A = torch.exp(A[None, :, None, :] * discrete_time_step[:, :, :, None])
-        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)
-        if self.use_mambapy and self.training and cache_params is None:
-            # [batch, seq_len, intermediate_size, ssm_state_size]
-            hs = pscan(discrete_A.transpose(1, 2), deltaB_u.transpose(1, 2))
-
-            scan_output = (hs @ C.unsqueeze(-1)).squeeze(3).transpose(1, 2)  # [batch, intermediate_size, seq_len]
-            scan_output = scan_output + hidden_states * self.D[None, :, None]
-            scan_output = scan_output * self.act(gate)
-        else:
-            # Use associative_scan for parallel computation when available
-            if (
-                self.use_associative_scan
-                and associative_scan is not None
-                and is_tracing(hidden_states)
-                and cache_params is None
-            ):
-
-                def combine_fn(left, right):
-                    a_left, b_left = left
-                    a_right, b_right = right
-                    return (a_left * a_right, a_right * b_left + b_right)
-
-                combine_mode = "pointwise" if discrete_A.device.type in ("cuda", "xpu") else "generic"
-                _, all_h = associative_scan(combine_fn, (discrete_A, deltaB_u), dim=2, combine_mode=combine_mode)
-                # all_h: [B, D, S, N] -> output: [B, D, S]
-                scan_output = (
-                    torch.matmul(all_h.permute(0, 2, 1, 3).to(dtype), C.unsqueeze(-1)).squeeze(-1).permute(0, 2, 1)
-                )
-                ssm_state = all_h[:, :, -1, :]
+            if output_final_state:
+                scan_output, final_state = scan_result
+                cache_params.update_recurrent_state(final_state, self.layer_idx)
             else:
-                # Sequential loop for decoding or when associative_scan unavailable
-                scan_outputs = []
-                for i in range(seq_len):
-                    # [batch, intermediate_size, ssm_state]
-                    ssm_state = discrete_A[:, :, i, :] * ssm_state + deltaB_u[:, :, i, :]
-                    # [batch, intermediate_size, 1]
-                    scan_output = torch.matmul(ssm_state.to(dtype), C[:, i, :].unsqueeze(-1))
-                    scan_outputs.append(scan_output[:, :, 0])
-                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)
-
-            if cache_params is not None:
-                cache_params.update_recurrent_state(ssm_state, self.layer_idx)
+                scan_output = scan_result
 
         # 4. Final linear projection
-        contextualized_states = self.out_proj(scan_output.transpose(1, 2))
+        contextualized_states = self.out_proj(scan_output.transpose(1, 2).to(dtype))
         return contextualized_states
 
-    @force_accelerate_hooks("conv1d")
-    def forward(
-        self,
-        hidden_states,
-        cache_params: Cache | None = None,
-        attention_mask: torch.LongTensor | None = None,
-        **kwargs,
-    ):
-        if is_fast_path_available and "cuda" in self.x_proj.weight.device.type and not is_tracing(hidden_states):
-            return self.cuda_kernels_forward(hidden_states, cache_params, attention_mask, **kwargs)
-        return self.slow_forward(hidden_states, cache_params, attention_mask, **kwargs)
-
 
 class MambaRMSNorm(nn.Module):
     def __init__(self, hidden_size, eps=1e-6):

From bfbf9f1512a412c133daea3a390aba38741cf123 Mon Sep 17 00:00:00 2001
From: vasqu <antonprogamer@gmail.com>
Date: Mon, 3 Aug 2026 17:54:22 +0000
Subject: [PATCH 25/43] fixups as per review comments

---
 src/transformers/integrations/hub_kernels.py  | 13 +++++++++++
 .../models/bamba/modeling_bamba.py            | 23 ++++++++-----------
 .../models/falcon_h1/modeling_falcon_h1.py    | 21 ++++++++---------
 .../models/falcon_h1/modular_falcon_h1.py     |  9 +++++---
 .../modeling_granitemoehybrid.py              | 23 ++++++++-----------
 .../models/mamba/modeling_mamba.py            |  5 +---
 .../models/mamba2/modeling_mamba2.py          |  5 +---
 .../models/nemotron_h/modeling_nemotron_h.py  | 23 ++++++++-----------
 .../olmo_hybrid/modeling_olmo_hybrid.py       |  4 ++--
 .../models/qwen3_5/modeling_qwen3_5.py        |  4 ++--
 .../qwen3_5_moe/modeling_qwen3_5_moe.py       |  4 ++--
 .../models/qwen3_next/modeling_qwen3_next.py  |  4 ++--
 .../models/qwen3_next/modular_qwen3_next.py   |  4 ++--
 .../models/zamba2/modeling_zamba2.py          | 23 ++++++++-----------
 utils/modular_model_converter.py              | 11 ++++++++-
 15 files changed, 90 insertions(+), 86 deletions(-)

diff --git a/src/transformers/integrations/hub_kernels.py b/src/transformers/integrations/hub_kernels.py
index 835e7fc5e9b6..dd5e35f1c414 100644
--- a/src/transformers/integrations/hub_kernels.py
+++ b/src/transformers/integrations/hub_kernels.py
@@ -61,6 +61,18 @@
 _kernels_enabled = _TRANSFORMERS_USE_HUB_KERNELS in ENV_VARS_TRUE_VALUES
 
 
+# Maps from func name to the internal module path
+_KERNELS_INTERNAL_PATH_MAPPINGS = {
+    "chunk_gated_delta_rule": "ops.gated_delta_rule",
+    "recurrent_gated_delta_rule": "ops.gated_delta_rule",
+    "mamba_split_conv1d_scan_combined": "ops.triton.ssd_combined",
+    "selective_state_update": "ops.triton.selective_state_update",
+    "mamba_chunk_scan_combined": "ops.triton.ssd_combined",
+    "mamba_inner_fn": "ops.selective_scan_interface",
+    "selective_scan_fn": "ops.selective_scan_interface",
+}
+
+
 if is_kernels_available():
     from kernels import (
         CUDAProperties,
@@ -756,6 +768,7 @@ def use_kernel_func_from_hub_with_fallback(func_name: str, package: str, interna
     kernel_wrapper_decorator = use_kernel_forward_from_hub(func_name)
 
     # Allow internal path prefix if given to resolve non __init__ imports
+    internal_path = _KERNELS_INTERNAL_PATH_MAPPINGS.get(func_name, internal_path)  # defaults
     full_path = func_name if internal_path is None else f"{internal_path}.{func_name}"
 
     def decorator(torch_function: Callable) -> Callable:
diff --git a/src/transformers/models/bamba/modeling_bamba.py b/src/transformers/models/bamba/modeling_bamba.py
index c8da7690884e..8e334b2e22aa 100644
--- a/src/transformers/models/bamba/modeling_bamba.py
+++ b/src/transformers/models/bamba/modeling_bamba.py
@@ -389,9 +389,8 @@ def causal_conv1d_fn(
 @use_kernel_func_from_hub_with_fallback(
     "mamba_split_conv1d_scan_combined",
     "mamba_ssm",
-    internal_path="ops.triton.ssd_combined",
 )
-def bamba_split_conv1d_scan_combined(
+def mamba2_split_conv1d_scan_combined(
     zxbcdt: torch.Tensor,
     conv1d_weight: torch.Tensor,
     conv1d_bias: torch.Tensor | None,
@@ -418,9 +417,8 @@ def bamba_split_conv1d_scan_combined(
 @use_kernel_func_from_hub_with_fallback(
     "selective_state_update",
     "mamba_ssm",
-    internal_path="ops.triton.selective_state_update",
 )
-def bamba_selective_state_update(
+def mamba2_selective_state_update(
     state: torch.Tensor,
     hidden_states: torch.Tensor,
     dt: torch.Tensor,
@@ -484,9 +482,8 @@ def bamba_selective_state_update(
 @use_kernel_func_from_hub_with_fallback(
     "mamba_chunk_scan_combined",
     "mamba_ssm",
-    internal_path="ops.triton.ssd_combined",
 )
-def bamba_chunk_scan(
+def mamba2_chunk_scan(
     hidden_states: torch.Tensor,
     dt: torch.Tensor,
     A: torch.Tensor,
@@ -590,9 +587,9 @@ def bamba_chunk_scan(
     [
         causal_conv1d_fn,
         causal_conv1d_update,
-        bamba_split_conv1d_scan_combined,
-        bamba_selective_state_update,
-        bamba_chunk_scan,
+        mamba2_split_conv1d_scan_combined,
+        mamba2_selective_state_update,
+        mamba2_chunk_scan,
     ]
 )
 class BambaMixer(nn.Module):
@@ -685,7 +682,7 @@ def forward(
             kwargs | {} if self.time_step_limit == (0.0, float("inf")) else kwargs | {"dt_limit": self.time_step_limit}
         )
         if self.training and cache_params is None:
-            fused_output = bamba_split_conv1d_scan_combined(
+            fused_output = mamba2_split_conv1d_scan_combined(
                 projected_states,
                 self.conv1d.weight.squeeze(1),
                 self.conv1d.bias,
@@ -719,7 +716,7 @@ def forward(
 
         # 2. Convolution sequence transformation
         hidden_states_B_C = hidden_states_B_C.transpose(1, 2)
-        if use_precomputed_states and seq_len == 1:
+        if use_precomputed_states and seq_len == 1 and not cache_params.layers[self.layer_idx].record_past:
             hidden_states_B_C = causal_conv1d_update(
                 hidden_states_B_C,
                 conv_state,
@@ -764,7 +761,7 @@ def forward(
             D = self.D[:, None].expand(-1, self.head_dim)
             dt_bias = self.dt_bias[:, None].expand(-1, self.head_dim)
 
-            scan_output = bamba_selective_state_update(
+            scan_output = mamba2_selective_state_update(
                 recurrent_state,
                 hidden_states,
                 dt,
@@ -782,7 +779,7 @@ def forward(
         # Chunk form
         else:
             output_final_state = cache_params is not None
-            scan_result = bamba_chunk_scan(
+            scan_result = mamba2_chunk_scan(
                 hidden_states.view(batch_size, seq_len, self.num_heads, self.head_dim),
                 dt,
                 A,
diff --git a/src/transformers/models/falcon_h1/modeling_falcon_h1.py b/src/transformers/models/falcon_h1/modeling_falcon_h1.py
index 98dfe7b0352f..aa15dd947ef6 100644
--- a/src/transformers/models/falcon_h1/modeling_falcon_h1.py
+++ b/src/transformers/models/falcon_h1/modeling_falcon_h1.py
@@ -396,9 +396,8 @@ def causal_conv1d_fn(
 @use_kernel_func_from_hub_with_fallback(
     "mamba_split_conv1d_scan_combined",
     "mamba_ssm",
-    internal_path="ops.triton.ssd_combined",
 )
-def falcon_h1_split_conv1d_scan_combined(
+def mamba2_split_conv1d_scan_combined(
     zxbcdt: torch.Tensor,
     conv1d_weight: torch.Tensor,
     conv1d_bias: torch.Tensor | None,
@@ -425,9 +424,8 @@ def falcon_h1_split_conv1d_scan_combined(
 @use_kernel_func_from_hub_with_fallback(
     "selective_state_update",
     "mamba_ssm",
-    internal_path="ops.triton.selective_state_update",
 )
-def falcon_h1_selective_state_update(
+def mamba2_selective_state_update(
     state: torch.Tensor,
     hidden_states: torch.Tensor,
     dt: torch.Tensor,
@@ -491,9 +489,8 @@ def falcon_h1_selective_state_update(
 @use_kernel_func_from_hub_with_fallback(
     "mamba_chunk_scan_combined",
     "mamba_ssm",
-    internal_path="ops.triton.ssd_combined",
 )
-def falcon_h1_chunk_scan(
+def mamba2_chunk_scan(
     hidden_states: torch.Tensor,
     dt: torch.Tensor,
     A: torch.Tensor,
@@ -597,9 +594,9 @@ def falcon_h1_chunk_scan(
     [
         causal_conv1d_fn,
         causal_conv1d_update,
-        falcon_h1_split_conv1d_scan_combined,
-        falcon_h1_selective_state_update,
-        falcon_h1_chunk_scan,
+        mamba2_split_conv1d_scan_combined,
+        mamba2_selective_state_update,
+        mamba2_chunk_scan,
     ]
 )
 class FalconH1Mixer(nn.Module):
@@ -704,7 +701,7 @@ def forward(
             kwargs | {} if self.time_step_limit == (0.0, float("inf")) else kwargs | {"dt_limit": self.time_step_limit}
         )
         if self.training and cache_params is None:
-            fused_output = falcon_h1_split_conv1d_scan_combined(  # noqa F821
+            fused_output = mamba2_split_conv1d_scan_combined(
                 projected_states,
                 self.conv1d.weight.squeeze(1),
                 self.conv1d.bias,
@@ -783,7 +780,7 @@ def forward(
             D = self.D[:, None, ...].expand(-1, self.head_dim)
             dt_bias = self.dt_bias[:, None, ...].expand(-1, self.head_dim)
 
-            scan_output = falcon_h1_selective_state_update(  # noqa F821
+            scan_output = mamba2_selective_state_update(
                 recurrent_state,
                 hidden_states,
                 dt,
@@ -805,7 +802,7 @@ def forward(
         # Chunk form
         else:
             output_final_state = cache_params is not None
-            scan_result = falcon_h1_chunk_scan(  # noqa F821
+            scan_result = mamba2_chunk_scan(
                 hidden_states.view(batch_size, seq_len, -1, self.head_dim),
                 dt,
                 A,
diff --git a/src/transformers/models/falcon_h1/modular_falcon_h1.py b/src/transformers/models/falcon_h1/modular_falcon_h1.py
index 5214e68780aa..3a00b3694991 100644
--- a/src/transformers/models/falcon_h1/modular_falcon_h1.py
+++ b/src/transformers/models/falcon_h1/modular_falcon_h1.py
@@ -51,6 +51,9 @@
     apply_mask_to_padding_states,
     causal_conv1d_fn,
     causal_conv1d_update,
+    mamba2_chunk_scan,
+    mamba2_selective_state_update,
+    mamba2_split_conv1d_scan_combined,
 )
 from .configuration_falcon_h1 import FalconH1Config
 
@@ -196,7 +199,7 @@ def forward(
             kwargs | {} if self.time_step_limit == (0.0, float("inf")) else kwargs | {"dt_limit": self.time_step_limit}
         )
         if self.training and cache_params is None:
-            fused_output = falcon_h1_split_conv1d_scan_combined(  # noqa F821
+            fused_output = mamba2_split_conv1d_scan_combined(
                 projected_states,
                 self.conv1d.weight.squeeze(1),
                 self.conv1d.bias,
@@ -275,7 +278,7 @@ def forward(
             D = self.D[:, None, ...].expand(-1, self.head_dim)
             dt_bias = self.dt_bias[:, None, ...].expand(-1, self.head_dim)
 
-            scan_output = falcon_h1_selective_state_update(  # noqa F821
+            scan_output = mamba2_selective_state_update(
                 recurrent_state,
                 hidden_states,
                 dt,
@@ -297,7 +300,7 @@ def forward(
         # Chunk form
         else:
             output_final_state = cache_params is not None
-            scan_result = falcon_h1_chunk_scan(  # noqa F821
+            scan_result = mamba2_chunk_scan(
                 hidden_states.view(batch_size, seq_len, -1, self.head_dim),
                 dt,
                 A,
diff --git a/src/transformers/models/granitemoehybrid/modeling_granitemoehybrid.py b/src/transformers/models/granitemoehybrid/modeling_granitemoehybrid.py
index 7fb2590cd51f..b8c12a1baf67 100644
--- a/src/transformers/models/granitemoehybrid/modeling_granitemoehybrid.py
+++ b/src/transformers/models/granitemoehybrid/modeling_granitemoehybrid.py
@@ -300,9 +300,8 @@ def causal_conv1d_fn(
 @use_kernel_func_from_hub_with_fallback(
     "mamba_split_conv1d_scan_combined",
     "mamba_ssm",
-    internal_path="ops.triton.ssd_combined",
 )
-def granitemoehybrid_split_conv1d_scan_combined(
+def mamba2_split_conv1d_scan_combined(
     zxbcdt: torch.Tensor,
     conv1d_weight: torch.Tensor,
     conv1d_bias: torch.Tensor | None,
@@ -329,9 +328,8 @@ def granitemoehybrid_split_conv1d_scan_combined(
 @use_kernel_func_from_hub_with_fallback(
     "selective_state_update",
     "mamba_ssm",
-    internal_path="ops.triton.selective_state_update",
 )
-def granitemoehybrid_selective_state_update(
+def mamba2_selective_state_update(
     state: torch.Tensor,
     hidden_states: torch.Tensor,
     dt: torch.Tensor,
@@ -395,9 +393,8 @@ def granitemoehybrid_selective_state_update(
 @use_kernel_func_from_hub_with_fallback(
     "mamba_chunk_scan_combined",
     "mamba_ssm",
-    internal_path="ops.triton.ssd_combined",
 )
-def granitemoehybrid_chunk_scan(
+def mamba2_chunk_scan(
     hidden_states: torch.Tensor,
     dt: torch.Tensor,
     A: torch.Tensor,
@@ -501,9 +498,9 @@ def granitemoehybrid_chunk_scan(
     [
         causal_conv1d_fn,
         causal_conv1d_update,
-        granitemoehybrid_split_conv1d_scan_combined,
-        granitemoehybrid_selective_state_update,
-        granitemoehybrid_chunk_scan,
+        mamba2_split_conv1d_scan_combined,
+        mamba2_selective_state_update,
+        mamba2_chunk_scan,
     ]
 )
 class GraniteMoeHybridMambaLayer(nn.Module):
@@ -596,7 +593,7 @@ def forward(
             kwargs | {} if self.time_step_limit == (0.0, float("inf")) else kwargs | {"dt_limit": self.time_step_limit}
         )
         if self.training and cache_params is None:
-            fused_output = granitemoehybrid_split_conv1d_scan_combined(
+            fused_output = mamba2_split_conv1d_scan_combined(
                 projected_states,
                 self.conv1d.weight.squeeze(1),
                 self.conv1d.bias,
@@ -630,7 +627,7 @@ def forward(
 
         # 2. Convolution sequence transformation
         hidden_states_B_C = hidden_states_B_C.transpose(1, 2)
-        if use_precomputed_states and seq_len == 1:
+        if use_precomputed_states and seq_len == 1 and not cache_params.layers[self.layer_idx].record_past:
             hidden_states_B_C = causal_conv1d_update(
                 hidden_states_B_C,
                 conv_state,
@@ -675,7 +672,7 @@ def forward(
             D = self.D[:, None].expand(-1, self.head_dim)
             dt_bias = self.dt_bias[:, None].expand(-1, self.head_dim)
 
-            scan_output = granitemoehybrid_selective_state_update(
+            scan_output = mamba2_selective_state_update(
                 recurrent_state,
                 hidden_states,
                 dt,
@@ -693,7 +690,7 @@ def forward(
         # Chunk form
         else:
             output_final_state = cache_params is not None
-            scan_result = granitemoehybrid_chunk_scan(
+            scan_result = mamba2_chunk_scan(
                 hidden_states.view(batch_size, seq_len, self.num_heads, self.head_dim),
                 dt,
                 A,
diff --git a/src/transformers/models/mamba/modeling_mamba.py b/src/transformers/models/mamba/modeling_mamba.py
index 538d6e6ee129..1144aa310a68 100644
--- a/src/transformers/models/mamba/modeling_mamba.py
+++ b/src/transformers/models/mamba/modeling_mamba.py
@@ -103,7 +103,6 @@ def causal_conv1d_fn(
 @use_kernel_func_from_hub_with_fallback(
     "mamba_inner_fn",
     "mamba_ssm",
-    internal_path="ops.selective_scan_interface",
 )
 def mamba_inner_fn(
     xz: torch.Tensor,
@@ -127,7 +126,6 @@ def mamba_inner_fn(
 @use_kernel_func_from_hub_with_fallback(
     "selective_state_update",
     "mamba_ssm",
-    internal_path="ops.triton.selective_state_update",
 )
 def mamba_selective_state_update(
     state: torch.Tensor,
@@ -177,7 +175,6 @@ def mamba_selective_state_update(
 @use_kernel_func_from_hub_with_fallback(
     "selective_scan_fn",
     "mamba_ssm",
-    internal_path="ops.selective_scan_interface",
 )
 def mamba_selective_scan(
     hidden_states: torch.Tensor,
@@ -399,7 +396,7 @@ def forward(
             recurrent_state = cache_params.layers[self.layer_idx].recurrent_states[0]
 
         # 2. Convolution sequence transformation
-        if use_precomputed_states and seq_len == 1:
+        if use_precomputed_states and seq_len == 1 and not cache_params.layers[self.layer_idx].record_past:
             hidden_states_B_C = causal_conv1d_update(
                 hidden_states_B_C,
                 conv_state,
diff --git a/src/transformers/models/mamba2/modeling_mamba2.py b/src/transformers/models/mamba2/modeling_mamba2.py
index 311d84e4baa1..526ed4bcbebd 100644
--- a/src/transformers/models/mamba2/modeling_mamba2.py
+++ b/src/transformers/models/mamba2/modeling_mamba2.py
@@ -166,7 +166,6 @@ def causal_conv1d_fn(
 @use_kernel_func_from_hub_with_fallback(
     "mamba_split_conv1d_scan_combined",
     "mamba_ssm",
-    internal_path="ops.triton.ssd_combined",
 )
 def mamba2_split_conv1d_scan_combined(
     zxbcdt: torch.Tensor,
@@ -195,7 +194,6 @@ def mamba2_split_conv1d_scan_combined(
 @use_kernel_func_from_hub_with_fallback(
     "selective_state_update",
     "mamba_ssm",
-    internal_path="ops.triton.selective_state_update",
 )
 def mamba2_selective_state_update(
     state: torch.Tensor,
@@ -261,7 +259,6 @@ def mamba2_selective_state_update(
 @use_kernel_func_from_hub_with_fallback(
     "mamba_chunk_scan_combined",
     "mamba_ssm",
-    internal_path="ops.triton.ssd_combined",
 )
 def mamba2_chunk_scan(
     hidden_states: torch.Tensor,
@@ -510,7 +507,7 @@ def forward(
 
         # 2. Convolution sequence transformation
         hidden_states_B_C = hidden_states_B_C.transpose(1, 2)
-        if use_precomputed_states and seq_len == 1:
+        if use_precomputed_states and seq_len == 1 and not cache_params.layers[self.layer_idx].record_past:
             hidden_states_B_C = causal_conv1d_update(
                 hidden_states_B_C,
                 conv_state,
diff --git a/src/transformers/models/nemotron_h/modeling_nemotron_h.py b/src/transformers/models/nemotron_h/modeling_nemotron_h.py
index a7a4afc842fc..c13eb29ab06e 100644
--- a/src/transformers/models/nemotron_h/modeling_nemotron_h.py
+++ b/src/transformers/models/nemotron_h/modeling_nemotron_h.py
@@ -162,9 +162,8 @@ def causal_conv1d_fn(
 @use_kernel_func_from_hub_with_fallback(
     "mamba_split_conv1d_scan_combined",
     "mamba_ssm",
-    internal_path="ops.triton.ssd_combined",
 )
-def nemotron_h_mamba2_split_conv1d_scan_combined(
+def mamba2_split_conv1d_scan_combined(
     zxbcdt: torch.Tensor,
     conv1d_weight: torch.Tensor,
     conv1d_bias: torch.Tensor | None,
@@ -191,9 +190,8 @@ def nemotron_h_mamba2_split_conv1d_scan_combined(
 @use_kernel_func_from_hub_with_fallback(
     "selective_state_update",
     "mamba_ssm",
-    internal_path="ops.triton.selective_state_update",
 )
-def nemotron_h_mamba2_selective_state_update(
+def mamba2_selective_state_update(
     state: torch.Tensor,
     hidden_states: torch.Tensor,
     dt: torch.Tensor,
@@ -257,9 +255,8 @@ def nemotron_h_mamba2_selective_state_update(
 @use_kernel_func_from_hub_with_fallback(
     "mamba_chunk_scan_combined",
     "mamba_ssm",
-    internal_path="ops.triton.ssd_combined",
 )
-def nemotron_h_mamba2_chunk_scan(
+def mamba2_chunk_scan(
     hidden_states: torch.Tensor,
     dt: torch.Tensor,
     A: torch.Tensor,
@@ -363,9 +360,9 @@ def nemotron_h_mamba2_chunk_scan(
     [
         causal_conv1d_fn,
         causal_conv1d_update,
-        nemotron_h_mamba2_split_conv1d_scan_combined,
-        nemotron_h_mamba2_selective_state_update,
-        nemotron_h_mamba2_chunk_scan,
+        mamba2_split_conv1d_scan_combined,
+        mamba2_selective_state_update,
+        mamba2_chunk_scan,
     ]
 )
 class NemotronHMamba2Mixer(nn.Module):
@@ -459,7 +456,7 @@ def forward(
             kwargs | {} if self.time_step_limit == (0.0, float("inf")) else kwargs | {"dt_limit": self.time_step_limit}
         )
         if self.training and cache_params is None:
-            fused_output = nemotron_h_mamba2_split_conv1d_scan_combined(
+            fused_output = mamba2_split_conv1d_scan_combined(
                 projected_states,
                 self.conv1d.weight.squeeze(1),
                 self.conv1d.bias,
@@ -493,7 +490,7 @@ def forward(
 
         # 2. Convolution sequence transformation
         hidden_states_B_C = hidden_states_B_C.transpose(1, 2)
-        if use_precomputed_states and seq_len == 1:
+        if use_precomputed_states and seq_len == 1 and not cache_params.layers[self.layer_idx].record_past:
             hidden_states_B_C = causal_conv1d_update(
                 hidden_states_B_C,
                 conv_state,
@@ -538,7 +535,7 @@ def forward(
             D = self.D[:, None].expand(-1, self.head_dim)
             dt_bias = self.dt_bias[:, None].expand(-1, self.head_dim)
 
-            scan_output = nemotron_h_mamba2_selective_state_update(
+            scan_output = mamba2_selective_state_update(
                 recurrent_state,
                 hidden_states,
                 dt,
@@ -556,7 +553,7 @@ def forward(
         # Chunk form
         else:
             output_final_state = cache_params is not None
-            scan_result = nemotron_h_mamba2_chunk_scan(
+            scan_result = mamba2_chunk_scan(
                 hidden_states.view(batch_size, seq_len, self.num_heads, self.head_dim),
                 dt,
                 A,
diff --git a/src/transformers/models/olmo_hybrid/modeling_olmo_hybrid.py b/src/transformers/models/olmo_hybrid/modeling_olmo_hybrid.py
index 6363acf4081d..a2fe2acd85c2 100644
--- a/src/transformers/models/olmo_hybrid/modeling_olmo_hybrid.py
+++ b/src/transformers/models/olmo_hybrid/modeling_olmo_hybrid.py
@@ -361,7 +361,7 @@ def l2norm(x: torch.FloatTensor, dim: int = -1, eps: float = 1e-6):
     return x * inv_norm
 
 
-@use_kernel_func_from_hub_with_fallback("chunk_gated_delta_rule", "fla", internal_path="ops.gated_delta_rule")
+@use_kernel_func_from_hub_with_fallback("chunk_gated_delta_rule", "fla")
 def torch_chunk_gated_delta_rule(
     query,
     key,
@@ -443,7 +443,7 @@ def torch_chunk_gated_delta_rule(
     return core_attn_out, last_recurrent_state
 
 
-@use_kernel_func_from_hub_with_fallback("recurrent_gated_delta_rule", "fla", internal_path="ops.gated_delta_rule")
+@use_kernel_func_from_hub_with_fallback("recurrent_gated_delta_rule", "fla")
 def torch_recurrent_gated_delta_rule(
     query,
     key,
diff --git a/src/transformers/models/qwen3_5/modeling_qwen3_5.py b/src/transformers/models/qwen3_5/modeling_qwen3_5.py
index 793d395ae6c7..9deeab69b152 100644
--- a/src/transformers/models/qwen3_5/modeling_qwen3_5.py
+++ b/src/transformers/models/qwen3_5/modeling_qwen3_5.py
@@ -249,7 +249,7 @@ def l2norm(x: torch.FloatTensor, dim: int = -1, eps: float = 1e-6):
     return x * inv_norm
 
 
-@use_kernel_func_from_hub_with_fallback("chunk_gated_delta_rule", "fla", internal_path="ops.gated_delta_rule")
+@use_kernel_func_from_hub_with_fallback("chunk_gated_delta_rule", "fla")
 def torch_chunk_gated_delta_rule(
     query,
     key,
@@ -331,7 +331,7 @@ def torch_chunk_gated_delta_rule(
     return core_attn_out, last_recurrent_state
 
 
-@use_kernel_func_from_hub_with_fallback("recurrent_gated_delta_rule", "fla", internal_path="ops.gated_delta_rule")
+@use_kernel_func_from_hub_with_fallback("recurrent_gated_delta_rule", "fla")
 def torch_recurrent_gated_delta_rule(
     query,
     key,
diff --git a/src/transformers/models/qwen3_5_moe/modeling_qwen3_5_moe.py b/src/transformers/models/qwen3_5_moe/modeling_qwen3_5_moe.py
index 511066852d3d..17b16a0cc2c7 100644
--- a/src/transformers/models/qwen3_5_moe/modeling_qwen3_5_moe.py
+++ b/src/transformers/models/qwen3_5_moe/modeling_qwen3_5_moe.py
@@ -250,7 +250,7 @@ def l2norm(x: torch.FloatTensor, dim: int = -1, eps: float = 1e-6):
     return x * inv_norm
 
 
-@use_kernel_func_from_hub_with_fallback("chunk_gated_delta_rule", "fla", internal_path="ops.gated_delta_rule")
+@use_kernel_func_from_hub_with_fallback("chunk_gated_delta_rule", "fla")
 def torch_chunk_gated_delta_rule(
     query,
     key,
@@ -332,7 +332,7 @@ def torch_chunk_gated_delta_rule(
     return core_attn_out, last_recurrent_state
 
 
-@use_kernel_func_from_hub_with_fallback("recurrent_gated_delta_rule", "fla", internal_path="ops.gated_delta_rule")
+@use_kernel_func_from_hub_with_fallback("recurrent_gated_delta_rule", "fla")
 def torch_recurrent_gated_delta_rule(
     query,
     key,
diff --git a/src/transformers/models/qwen3_next/modeling_qwen3_next.py b/src/transformers/models/qwen3_next/modeling_qwen3_next.py
index 2cdca2d4056d..71e76d97e02f 100644
--- a/src/transformers/models/qwen3_next/modeling_qwen3_next.py
+++ b/src/transformers/models/qwen3_next/modeling_qwen3_next.py
@@ -373,7 +373,7 @@ def l2norm(x: torch.FloatTensor, dim: int = -1, eps: float = 1e-6):
     return x * inv_norm
 
 
-@use_kernel_func_from_hub_with_fallback("chunk_gated_delta_rule", "fla", internal_path="ops.gated_delta_rule")
+@use_kernel_func_from_hub_with_fallback("chunk_gated_delta_rule", "fla")
 def torch_chunk_gated_delta_rule(
     query,
     key,
@@ -455,7 +455,7 @@ def torch_chunk_gated_delta_rule(
     return core_attn_out, last_recurrent_state
 
 
-@use_kernel_func_from_hub_with_fallback("recurrent_gated_delta_rule", "fla", internal_path="ops.gated_delta_rule")
+@use_kernel_func_from_hub_with_fallback("recurrent_gated_delta_rule", "fla")
 def torch_recurrent_gated_delta_rule(
     query,
     key,
diff --git a/src/transformers/models/qwen3_next/modular_qwen3_next.py b/src/transformers/models/qwen3_next/modular_qwen3_next.py
index 7049e7e7f0c7..3b10b4a54a67 100644
--- a/src/transformers/models/qwen3_next/modular_qwen3_next.py
+++ b/src/transformers/models/qwen3_next/modular_qwen3_next.py
@@ -206,7 +206,7 @@ def l2norm(x: torch.FloatTensor, dim: int = -1, eps: float = 1e-6):
     return x * inv_norm
 
 
-@use_kernel_func_from_hub_with_fallback("chunk_gated_delta_rule", "fla", internal_path="ops.gated_delta_rule")
+@use_kernel_func_from_hub_with_fallback("chunk_gated_delta_rule", "fla")
 def torch_chunk_gated_delta_rule(
     query,
     key,
@@ -288,7 +288,7 @@ def torch_chunk_gated_delta_rule(
     return core_attn_out, last_recurrent_state
 
 
-@use_kernel_func_from_hub_with_fallback("recurrent_gated_delta_rule", "fla", internal_path="ops.gated_delta_rule")
+@use_kernel_func_from_hub_with_fallback("recurrent_gated_delta_rule", "fla")
 def torch_recurrent_gated_delta_rule(
     query,
     key,
diff --git a/src/transformers/models/zamba2/modeling_zamba2.py b/src/transformers/models/zamba2/modeling_zamba2.py
index d2d953d52221..ec95cd71ca2b 100644
--- a/src/transformers/models/zamba2/modeling_zamba2.py
+++ b/src/transformers/models/zamba2/modeling_zamba2.py
@@ -454,9 +454,8 @@ def causal_conv1d_fn(
 @use_kernel_func_from_hub_with_fallback(
     "mamba_split_conv1d_scan_combined",
     "mamba_ssm",
-    internal_path="ops.triton.ssd_combined",
 )
-def zamba2_split_conv1d_scan_combined(
+def mamba2_split_conv1d_scan_combined(
     zxbcdt: torch.Tensor,
     conv1d_weight: torch.Tensor,
     conv1d_bias: torch.Tensor | None,
@@ -483,9 +482,8 @@ def zamba2_split_conv1d_scan_combined(
 @use_kernel_func_from_hub_with_fallback(
     "selective_state_update",
     "mamba_ssm",
-    internal_path="ops.triton.selective_state_update",
 )
-def zamba2_selective_state_update(
+def mamba2_selective_state_update(
     state: torch.Tensor,
     hidden_states: torch.Tensor,
     dt: torch.Tensor,
@@ -549,9 +547,8 @@ def zamba2_selective_state_update(
 @use_kernel_func_from_hub_with_fallback(
     "mamba_chunk_scan_combined",
     "mamba_ssm",
-    internal_path="ops.triton.ssd_combined",
 )
-def zamba2_chunk_scan(
+def mamba2_chunk_scan(
     hidden_states: torch.Tensor,
     dt: torch.Tensor,
     A: torch.Tensor,
@@ -655,9 +652,9 @@ def zamba2_chunk_scan(
     [
         causal_conv1d_fn,
         causal_conv1d_update,
-        zamba2_split_conv1d_scan_combined,
-        zamba2_selective_state_update,
-        zamba2_chunk_scan,
+        mamba2_split_conv1d_scan_combined,
+        mamba2_selective_state_update,
+        mamba2_chunk_scan,
     ]
 )
 class Zamba2MambaMixer(nn.Module):
@@ -750,7 +747,7 @@ def forward(
             kwargs | {} if self.time_step_limit == (0.0, float("inf")) else kwargs | {"dt_limit": self.time_step_limit}
         )
         if self.training and cache_params is None:
-            fused_output = zamba2_split_conv1d_scan_combined(
+            fused_output = mamba2_split_conv1d_scan_combined(
                 projected_states,
                 self.conv1d.weight.squeeze(1),
                 self.conv1d.bias,
@@ -784,7 +781,7 @@ def forward(
 
         # 2. Convolution sequence transformation
         hidden_states_B_C = hidden_states_B_C.transpose(1, 2)
-        if use_precomputed_states and seq_len == 1:
+        if use_precomputed_states and seq_len == 1 and not cache_params.layers[self.layer_idx].record_past:
             hidden_states_B_C = causal_conv1d_update(
                 hidden_states_B_C,
                 conv_state,
@@ -829,7 +826,7 @@ def forward(
             D = self.D[:, None].expand(-1, self.head_dim)
             dt_bias = self.dt_bias[:, None].expand(-1, self.head_dim)
 
-            scan_output = zamba2_selective_state_update(
+            scan_output = mamba2_selective_state_update(
                 recurrent_state,
                 hidden_states,
                 dt,
@@ -847,7 +844,7 @@ def forward(
         # Chunk form
         else:
             output_final_state = cache_params is not None
-            scan_result = zamba2_chunk_scan(
+            scan_result = mamba2_chunk_scan(
                 hidden_states.view(batch_size, seq_len, self.num_heads, self.head_dim),
                 dt,
                 A,
diff --git a/utils/modular_model_converter.py b/utils/modular_model_converter.py
index 405e636d9999..a30cec1e37f4 100644
--- a/utils/modular_model_converter.py
+++ b/utils/modular_model_converter.py
@@ -77,7 +77,16 @@ def get_module_source_from_name(module_name: str) -> str:
 
 # Some exceptions to never replace, usually some package names that may contain model names (they may be used outside
 # `from xxx import y`)
-NAMES_TO_NEVER_REPLACE = ("mamba_ssm", "mamba-ssm", "mamba_inner_fn")
+NAMES_TO_NEVER_REPLACE = (
+    "mamba_ssm",
+    "mamba-ssm",
+    "mamba_inner_fn",
+    "mamba_selective_state_update",
+    "mamba_selective_scan",
+    "mamba2_split_conv1d_scan_combined",
+    "mamba2_selective_state_update",
+    "mamba2_chunk_scan",
+)
 
 
 def preserve_case_replace(text, patterns: dict, default_name: str):

From 511d4a61a1cd658fe2c12b70afdade863bd2e00a Mon Sep 17 00:00:00 2001
From: vasqu <antonprogamer@gmail.com>
Date: Mon, 3 Aug 2026 19:06:56 +0000
Subject: [PATCH 26/43] fixup mamba tests and other issues

---
 src/transformers/integrations/accelerate.py   | 37 +++++++++++--------
 .../models/mamba/modeling_mamba.py            | 12 ++++--
 tests/models/mamba/test_modeling_mamba.py     | 19 +++++-----
 3 files changed, 39 insertions(+), 29 deletions(-)

diff --git a/src/transformers/integrations/accelerate.py b/src/transformers/integrations/accelerate.py
index 476ddf95ba34..9741e7772a45 100644
--- a/src/transformers/integrations/accelerate.py
+++ b/src/transformers/integrations/accelerate.py
@@ -920,31 +920,38 @@ def check_tied_parameters_on_same_device(tied_params, device_map):
             )
 
 
-def force_accelerate_hooks(child_module_name: str) -> Callable:
+def force_accelerate_hooks(child_module_names: str | list[str]) -> Callable:
     """
-    Decorator to forcefully fire the accelerate hooks of `child_module_name`, before entering the forward of the parent itself.
-    Indeed, the hooks of the child are only fired through the `forward` child's method, so if the child weights are used directly,
+    Decorator to forcefully fire the accelerate hooks of `child_module_names`, before entering the forward of the parent itself.
+    Indeed, the hooks of a child are only fired through the `forward` child's method, so if the child weights are used directly,
     as is the case inside `causal_conv1d_fn` and `causal_conv1d_update` for example, they will not be fired. This may cause device
     issues, especially in the case of offloading, that this decorator will correct.
     """
 
+    if isinstance(child_module_names, str):
+        child_module_names = [child_module_names]
+
     def decorator(forward_func: Callable) -> Callable:
         def wrapped(self, *args, **kwargs):
-            hooked_module = getattr(self, child_module_name)
-            hook = getattr(hooked_module, "_hf_hook", None)
-            if hook is not None:
-                # Note that here we only call the hook with the module, not `*args` not `**kwargs`, as we assume the `forward`
-                # on which this decorator is applied is responsible to move the args and kwargs with its own hook if any. This makes
-                # sense as the module decorated with this should have all internal modules on the same device
-                hook.pre_forward(hooked_module)
+            hooked_modules = []
+            for child_module_name in child_module_names:
+                hooked_module = getattr(self, child_module_name)
+                hook = getattr(hooked_module, "_hf_hook", None)
+                hooked_modules.append((hooked_module, hook))
+                if hook is not None:
+                    # Note that here we only call the hook with the module, not `*args` not `**kwargs`, as we assume the `forward`
+                    # on which this decorator is applied is responsible to move the args and kwargs with its own hook if any. This makes
+                    # sense as the module decorated with this should have all internal modules on the same device
+                    hook.pre_forward(hooked_module)
 
             output = forward_func(self, *args, **kwargs)
 
-            if hook is not None:
-                # Note that here we only call the hook with the module, not `output`, as we assume the `forward` on which
-                # this decorator is applied is responsible to move the output with its own hook if any. This makes sense
-                # as the module decorated with this should have all internal modules on the same device
-                hook.post_forward(hooked_module, ())
+            for hooked_module, hook in reversed(hooked_modules):
+                if hook is not None:
+                    # Note that here we only call the hook with the module, not `output`, as we assume the `forward` on which
+                    # this decorator is applied is responsible to move the output with its own hook if any. This makes sense
+                    # as the module decorated with this should have all internal modules on the same device
+                    hook.post_forward(hooked_module, ())
 
             return output
 
diff --git a/src/transformers/models/mamba/modeling_mamba.py b/src/transformers/models/mamba/modeling_mamba.py
index 1144aa310a68..55a37103e545 100644
--- a/src/transformers/models/mamba/modeling_mamba.py
+++ b/src/transformers/models/mamba/modeling_mamba.py
@@ -50,7 +50,7 @@ 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
     """
     # NOTE: attention mask is a 2D boolean tensor
-    if attention_mask is not None and attention_mask.shape[1] > 1 and attention_mask.shape[0] > 1:
+    if attention_mask is not None:
         dtype = hidden_states.dtype
         hidden_states = (hidden_states * attention_mask[:, :, None]).to(dtype)
 
@@ -210,6 +210,10 @@ def mamba_selective_scan(
     if delta_softplus:
         dt = F.softplus(dt)
 
+    # We need to transpose on the basis of the original kernel layout
+    B = B.transpose(1, 2)
+    C = C.transpose(1, 2)
+
     # Discretize A and B for the entire sequence
     discrete_A = torch.exp(A[None, :, None, :] * dt[:, :, :, None])
     discrete_B = dt[:, :, :, None] * B[:, None, :, :].float()
@@ -237,7 +241,7 @@ def combine_fn(left, right):
             combine_mode=combine_mode,
         )
 
-        scan_output = torch.matmul(all_states.transpose(1, 2).to(input_dtype), C.transpose(1, 2).unsqueeze(-1))
+        scan_output = torch.matmul(all_states.transpose(1, 2).to(input_dtype), C.unsqueeze(-1))
         scan_output = scan_output.squeeze(-1).transpose(1, 2)
         ssm_state = all_states[:, :, -1]
 
@@ -259,7 +263,7 @@ def combine_fn(left, right):
             ssm_state = discrete_A[:, :, index] * ssm_state + deltaB_u[:, :, index]
 
             # Subsequent output
-            scan_output = torch.matmul(ssm_state.to(input_dtype), C[:, :, index].unsqueeze(-1))
+            scan_output = torch.matmul(ssm_state.to(input_dtype), C[:, index, :].unsqueeze(-1))
             scan_outputs.append(scan_output[:, :, 0])
         scan_output = torch.stack(scan_outputs, dim=-1)
 
@@ -351,7 +355,7 @@ def init_mamba_weights(self):
         inv_dt = dt + torch.log(-torch.expm1(-dt))
         init.copy_(self.dt_proj.bias, inv_dt)
 
-    @force_accelerate_hooks("conv1d")
+    @force_accelerate_hooks(["conv1d", "dt_proj"])
     def forward(
         self,
         hidden_states: torch.Tensor,
diff --git a/tests/models/mamba/test_modeling_mamba.py b/tests/models/mamba/test_modeling_mamba.py
index cf17bc753769..fa10d3662d4e 100644
--- a/tests/models/mamba/test_modeling_mamba.py
+++ b/tests/models/mamba/test_modeling_mamba.py
@@ -20,7 +20,7 @@
 from parameterized import parameterized
 
 from transformers import AutoTokenizer, MambaConfig, is_torch_available
-from transformers.testing_utils import require_torch, slow, torch_device
+from transformers.testing_utils import require_torch, require_torch_greater_or_equal, slow, torch_device
 
 from ...generation.test_utils import GenerationTesterMixin
 from ...test_configuration_common import ConfigTester
@@ -180,7 +180,11 @@ def create_and_check_mamba_cached_slow_forward_and_backwards(
         self, config, input_ids, *args, gradient_checkpointing=False
     ):
         model = MambaModel(config)
-        model.to(torch_device)
+
+        # force torch path in any case
+        model.to("cpu")
+        input_ids = input_ids.to("cpu")
+
         if gradient_checkpointing:
             model.gradient_checkpointing_enable()
 
@@ -190,7 +194,7 @@ def create_and_check_mamba_cached_slow_forward_and_backwards(
 
         # use cache
         token_emb = model.embeddings(input_ids)
-        outputs = model.layers[0].mixer.slow_forward(token_emb, cache)
+        outputs = model.layers[0].mixer(token_emb, cache)
 
         loss = torch.log1p(torch.abs(outputs.sum()))
         self.parent.assertEqual(loss.shape, ())
@@ -473,12 +477,9 @@ def test_compile_mamba_cache(self):
         output_sentence = self.tokenizer.decode(output[0].tolist())
         self.assertEqual(output_sentence, expected_output)
 
+    @require_torch_greater_or_equal("2.9.0")
     @pytest.mark.torch_compile_test
     def test_compile_associative_scan_no_cache(self):
-        from transformers.models.mamba.modeling_mamba import associative_scan
-
-        if associative_scan is None:
-            self.skipTest("associative_scan is not available in this PyTorch version.")
         if torch_device == "cpu":
             self.skipTest("Associative scan compile test requires a torch accelerator.")
 
@@ -498,13 +499,11 @@ def test_compile_associative_scan_no_cache(self):
         output_sentence = self.tokenizer.decode(output[0].tolist())
         self.assertEqual(output_sentence, expected_output)
 
+    @require_torch_greater_or_equal("2.9.0")
     @pytest.mark.torch_compile_test
     def test_associative_scan_matches_sequential(self):
         """Compiled generate with use_associative_scan=False vs =True produces the same text."""
-        from transformers.models.mamba.modeling_mamba import associative_scan
 
-        if associative_scan is None:
-            self.skipTest("associative_scan is not available (requires torch >= 2.9.0).")
         if torch_device == "cpu":
             self.skipTest("Associative scan test requires a torch accelerator.")
 

From 3ee52c354204a30bae5ca5ef7cd838b75c446b91 Mon Sep 17 00:00:00 2001
From: vasqu <antonprogamer@gmail.com>
Date: Mon, 3 Aug 2026 19:13:46 +0000
Subject: [PATCH 27/43] make no shape check by default (single padded sample
 should also work)

---
 src/transformers/models/bamba/modeling_bamba.py                 | 2 +-
 src/transformers/models/falcon_h1/modeling_falcon_h1.py         | 2 +-
 .../models/granitemoehybrid/modeling_granitemoehybrid.py        | 2 +-
 src/transformers/models/inkling/modeling_inkling.py             | 2 +-
 src/transformers/models/lfm2/modeling_lfm2.py                   | 2 +-
 src/transformers/models/lfm2_moe/modeling_lfm2_moe.py           | 2 +-
 src/transformers/models/mamba2/modeling_mamba2.py               | 2 +-
 src/transformers/models/minimax/modeling_minimax.py             | 2 +-
 src/transformers/models/nemotron_h/modeling_nemotron_h.py       | 2 +-
 src/transformers/models/olmo_hybrid/modeling_olmo_hybrid.py     | 2 +-
 src/transformers/models/qwen3_5/modeling_qwen3_5.py             | 2 +-
 src/transformers/models/qwen3_5_moe/modeling_qwen3_5_moe.py     | 2 +-
 src/transformers/models/qwen3_next/modeling_qwen3_next.py       | 2 +-
 src/transformers/models/zamba2/modeling_zamba2.py               | 2 +-
 14 files changed, 14 insertions(+), 14 deletions(-)

diff --git a/src/transformers/models/bamba/modeling_bamba.py b/src/transformers/models/bamba/modeling_bamba.py
index 8e334b2e22aa..decf58d2bcda 100644
--- a/src/transformers/models/bamba/modeling_bamba.py
+++ b/src/transformers/models/bamba/modeling_bamba.py
@@ -336,7 +336,7 @@ 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
     """
     # NOTE: attention mask is a 2D boolean tensor
-    if attention_mask is not None and attention_mask.shape[1] > 1 and attention_mask.shape[0] > 1:
+    if attention_mask is not None:
         dtype = hidden_states.dtype
         hidden_states = (hidden_states * attention_mask[:, :, None]).to(dtype)
 
diff --git a/src/transformers/models/falcon_h1/modeling_falcon_h1.py b/src/transformers/models/falcon_h1/modeling_falcon_h1.py
index aa15dd947ef6..03f7216b034e 100644
--- a/src/transformers/models/falcon_h1/modeling_falcon_h1.py
+++ b/src/transformers/models/falcon_h1/modeling_falcon_h1.py
@@ -343,7 +343,7 @@ 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
     """
     # NOTE: attention mask is a 2D boolean tensor
-    if attention_mask is not None and attention_mask.shape[1] > 1 and attention_mask.shape[0] > 1:
+    if attention_mask is not None:
         dtype = hidden_states.dtype
         hidden_states = (hidden_states * attention_mask[:, :, None]).to(dtype)
 
diff --git a/src/transformers/models/granitemoehybrid/modeling_granitemoehybrid.py b/src/transformers/models/granitemoehybrid/modeling_granitemoehybrid.py
index b8c12a1baf67..f3685dc728b4 100644
--- a/src/transformers/models/granitemoehybrid/modeling_granitemoehybrid.py
+++ b/src/transformers/models/granitemoehybrid/modeling_granitemoehybrid.py
@@ -247,7 +247,7 @@ 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
     """
     # NOTE: attention mask is a 2D boolean tensor
-    if attention_mask is not None and attention_mask.shape[1] > 1 and attention_mask.shape[0] > 1:
+    if attention_mask is not None:
         dtype = hidden_states.dtype
         hidden_states = (hidden_states * attention_mask[:, :, None]).to(dtype)
 
diff --git a/src/transformers/models/inkling/modeling_inkling.py b/src/transformers/models/inkling/modeling_inkling.py
index 6df67c93cbb9..c68483bf155f 100644
--- a/src/transformers/models/inkling/modeling_inkling.py
+++ b/src/transformers/models/inkling/modeling_inkling.py
@@ -430,7 +430,7 @@ 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
     """
     # NOTE: attention mask is a 2D boolean tensor
-    if attention_mask is not None and attention_mask.shape[1] > 1 and attention_mask.shape[0] > 1:
+    if attention_mask is not None:
         dtype = hidden_states.dtype
         hidden_states = (hidden_states * attention_mask[:, :, None]).to(dtype)
 
diff --git a/src/transformers/models/lfm2/modeling_lfm2.py b/src/transformers/models/lfm2/modeling_lfm2.py
index a0fef79b7f73..c605c86acaaa 100644
--- a/src/transformers/models/lfm2/modeling_lfm2.py
+++ b/src/transformers/models/lfm2/modeling_lfm2.py
@@ -272,7 +272,7 @@ 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
     """
     # NOTE: attention mask is a 2D boolean tensor
-    if attention_mask is not None and attention_mask.shape[1] > 1 and attention_mask.shape[0] > 1:
+    if attention_mask is not None:
         dtype = hidden_states.dtype
         hidden_states = (hidden_states * attention_mask[:, :, None]).to(dtype)
 
diff --git a/src/transformers/models/lfm2_moe/modeling_lfm2_moe.py b/src/transformers/models/lfm2_moe/modeling_lfm2_moe.py
index 327699f253d8..07ef4510abb0 100644
--- a/src/transformers/models/lfm2_moe/modeling_lfm2_moe.py
+++ b/src/transformers/models/lfm2_moe/modeling_lfm2_moe.py
@@ -357,7 +357,7 @@ 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
     """
     # NOTE: attention mask is a 2D boolean tensor
-    if attention_mask is not None and attention_mask.shape[1] > 1 and attention_mask.shape[0] > 1:
+    if attention_mask is not None:
         dtype = hidden_states.dtype
         hidden_states = (hidden_states * attention_mask[:, :, None]).to(dtype)
 
diff --git a/src/transformers/models/mamba2/modeling_mamba2.py b/src/transformers/models/mamba2/modeling_mamba2.py
index 526ed4bcbebd..5ae3991df372 100644
--- a/src/transformers/models/mamba2/modeling_mamba2.py
+++ b/src/transformers/models/mamba2/modeling_mamba2.py
@@ -95,7 +95,7 @@ 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
     """
     # NOTE: attention mask is a 2D boolean tensor
-    if attention_mask is not None and attention_mask.shape[1] > 1 and attention_mask.shape[0] > 1:
+    if attention_mask is not None:
         dtype = hidden_states.dtype
         hidden_states = (hidden_states * attention_mask[:, :, None]).to(dtype)
 
diff --git a/src/transformers/models/minimax/modeling_minimax.py b/src/transformers/models/minimax/modeling_minimax.py
index b7a773043e28..446c07efb5e4 100644
--- a/src/transformers/models/minimax/modeling_minimax.py
+++ b/src/transformers/models/minimax/modeling_minimax.py
@@ -112,7 +112,7 @@ 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
     """
     # NOTE: attention mask is a 2D boolean tensor
-    if attention_mask is not None and attention_mask.shape[1] > 1 and attention_mask.shape[0] > 1:
+    if attention_mask is not None:
         dtype = hidden_states.dtype
         hidden_states = (hidden_states * attention_mask[:, :, None]).to(dtype)
 
diff --git a/src/transformers/models/nemotron_h/modeling_nemotron_h.py b/src/transformers/models/nemotron_h/modeling_nemotron_h.py
index c13eb29ab06e..6cba8aa1ff6a 100644
--- a/src/transformers/models/nemotron_h/modeling_nemotron_h.py
+++ b/src/transformers/models/nemotron_h/modeling_nemotron_h.py
@@ -109,7 +109,7 @@ 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
     """
     # NOTE: attention mask is a 2D boolean tensor
-    if attention_mask is not None and attention_mask.shape[1] > 1 and attention_mask.shape[0] > 1:
+    if attention_mask is not None:
         dtype = hidden_states.dtype
         hidden_states = (hidden_states * attention_mask[:, :, None]).to(dtype)
 
diff --git a/src/transformers/models/olmo_hybrid/modeling_olmo_hybrid.py b/src/transformers/models/olmo_hybrid/modeling_olmo_hybrid.py
index a2fe2acd85c2..4abd5d6232c3 100644
--- a/src/transformers/models/olmo_hybrid/modeling_olmo_hybrid.py
+++ b/src/transformers/models/olmo_hybrid/modeling_olmo_hybrid.py
@@ -305,7 +305,7 @@ 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
     """
     # NOTE: attention mask is a 2D boolean tensor
-    if attention_mask is not None and attention_mask.shape[1] > 1 and attention_mask.shape[0] > 1:
+    if attention_mask is not None:
         dtype = hidden_states.dtype
         hidden_states = (hidden_states * attention_mask[:, :, None]).to(dtype)
 
diff --git a/src/transformers/models/qwen3_5/modeling_qwen3_5.py b/src/transformers/models/qwen3_5/modeling_qwen3_5.py
index 9deeab69b152..831fa13c4b89 100644
--- a/src/transformers/models/qwen3_5/modeling_qwen3_5.py
+++ b/src/transformers/models/qwen3_5/modeling_qwen3_5.py
@@ -193,7 +193,7 @@ 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
     """
     # NOTE: attention mask is a 2D boolean tensor
-    if attention_mask is not None and attention_mask.shape[1] > 1 and attention_mask.shape[0] > 1:
+    if attention_mask is not None:
         dtype = hidden_states.dtype
         hidden_states = (hidden_states * attention_mask[:, :, None]).to(dtype)
 
diff --git a/src/transformers/models/qwen3_5_moe/modeling_qwen3_5_moe.py b/src/transformers/models/qwen3_5_moe/modeling_qwen3_5_moe.py
index 17b16a0cc2c7..e12ba4482672 100644
--- a/src/transformers/models/qwen3_5_moe/modeling_qwen3_5_moe.py
+++ b/src/transformers/models/qwen3_5_moe/modeling_qwen3_5_moe.py
@@ -194,7 +194,7 @@ 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
     """
     # NOTE: attention mask is a 2D boolean tensor
-    if attention_mask is not None and attention_mask.shape[1] > 1 and attention_mask.shape[0] > 1:
+    if attention_mask is not None:
         dtype = hidden_states.dtype
         hidden_states = (hidden_states * attention_mask[:, :, None]).to(dtype)
 
diff --git a/src/transformers/models/qwen3_next/modeling_qwen3_next.py b/src/transformers/models/qwen3_next/modeling_qwen3_next.py
index 71e76d97e02f..00e88fd9cd9c 100644
--- a/src/transformers/models/qwen3_next/modeling_qwen3_next.py
+++ b/src/transformers/models/qwen3_next/modeling_qwen3_next.py
@@ -317,7 +317,7 @@ 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
     """
     # NOTE: attention mask is a 2D boolean tensor
-    if attention_mask is not None and attention_mask.shape[1] > 1 and attention_mask.shape[0] > 1:
+    if attention_mask is not None:
         dtype = hidden_states.dtype
         hidden_states = (hidden_states * attention_mask[:, :, None]).to(dtype)
 
diff --git a/src/transformers/models/zamba2/modeling_zamba2.py b/src/transformers/models/zamba2/modeling_zamba2.py
index ec95cd71ca2b..8c8634e4b9e2 100644
--- a/src/transformers/models/zamba2/modeling_zamba2.py
+++ b/src/transformers/models/zamba2/modeling_zamba2.py
@@ -401,7 +401,7 @@ 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
     """
     # NOTE: attention mask is a 2D boolean tensor
-    if attention_mask is not None and attention_mask.shape[1] > 1 and attention_mask.shape[0] > 1:
+    if attention_mask is not None:
         dtype = hidden_states.dtype
         hidden_states = (hidden_states * attention_mask[:, :, None]).to(dtype)
 

From 34051c4d11733c4e19b63fda7ddcb21fcf99425d Mon Sep 17 00:00:00 2001
From: vasqu <antonprogamer@gmail.com>
Date: Mon, 3 Aug 2026 19:16:41 +0000
Subject: [PATCH 28/43] remove todo

---
 src/transformers/models/mamba/modeling_mamba.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/src/transformers/models/mamba/modeling_mamba.py b/src/transformers/models/mamba/modeling_mamba.py
index 55a37103e545..595545ab2d9a 100644
--- a/src/transformers/models/mamba/modeling_mamba.py
+++ b/src/transformers/models/mamba/modeling_mamba.py
@@ -219,7 +219,6 @@ def mamba_selective_scan(
     discrete_B = dt[:, :, :, None] * B[:, None, :, :].float()
     deltaB_u = discrete_B * hidden_states[:, :, :, None].float()
 
-    # TODO: check these out
     if use_mambapy and pscan is not None:
         all_states = pscan(discrete_A.transpose(1, 2), deltaB_u.transpose(1, 2))
 

From c3a6c9824c13a1e917f0fcfa70decedb03a5fcec Mon Sep 17 00:00:00 2001
From: vasqu <antonprogamer@gmail.com>
Date: Mon, 3 Aug 2026 21:19:18 +0000
Subject: [PATCH 29/43] propogate mamba1

---
 .../falcon_mamba/modeling_falcon_mamba.py     | 606 +++++++++---------
 .../falcon_mamba/modular_falcon_mamba.py      | 333 ++++------
 .../models/jamba/modeling_jamba.py            | 466 ++++++++------
 .../models/jamba/modular_jamba.py             | 240 +++----
 .../models/mamba/modeling_mamba.py            |   4 +
 .../test_modeling_falcon_mamba.py             |  16 +-
 tests/models/jamba/test_modeling_jamba.py     |   1 +
 utils/modular_model_converter.py              |   1 +
 8 files changed, 811 insertions(+), 856 deletions(-)

diff --git a/src/transformers/models/falcon_mamba/modeling_falcon_mamba.py b/src/transformers/models/falcon_mamba/modeling_falcon_mamba.py
index 87773fab7100..ac833f31db9c 100644
--- a/src/transformers/models/falcon_mamba/modeling_falcon_mamba.py
+++ b/src/transformers/models/falcon_mamba/modeling_falcon_mamba.py
@@ -30,34 +30,45 @@
 from ...activations import ACT2FN
 from ...cache_utils import Cache, DynamicCache
 from ...generation import GenerationMixin
-from ...integrations import lazy_load_kernel
+from ...integrations import use_kernel_forward_from_hub, use_kernel_func_from_hub_with_fallback, use_kernelized_func
 from ...integrations.accelerate import force_accelerate_hooks
 from ...modeling_layers import GradientCheckpointingLayer
 from ...modeling_utils import PreTrainedModel
-from ...utils import ModelOutput, auto_docstring, logging
-from ...utils.import_utils import (
-    is_mambapy_available,
-    is_torch_greater_or_equal,
-    is_tracing,
-    resolve_internal_import,
-)
+from ...utils import ModelOutput, auto_docstring
+from ...utils.import_utils import is_mambapy_available, is_torch_greater_or_equal, is_tracing
 from .configuration_falcon_mamba import FalconMambaConfig
 
 
-if is_torch_greater_or_equal("2.9.0"):
-    from torch._higher_order_ops.associative_scan import associative_scan
-else:
-    associative_scan = None
+class FalconMambaWeightlessRMSNorm(torch.nn.Module):
+    def __init__(self, hidden_size, eps: float = 1e-6):
+        super().__init__()
+        self.eps = eps
+        # Dummy weights that are not used (only for imitating on kernels path)
+        self.register_buffer("weight", torch.ones(hidden_size, requires_grad=False), persistent=False)
+
+    def _norm(self, x):
+        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
+
+    def forward(self, x):
+        return self._norm(x.float()).type_as(x)
+
+    def extra_repr(self):
+        return f"eps={self.eps}"
 
-if is_mambapy_available():
-    from mambapy.pscan import pscan
-else:
-    pscan = None
 
+def apply_mask_to_padding_states(hidden_states, attention_mask):
+    """
+    Tunes out the hidden states for padding tokens, see https://github.com/state-spaces/falcon_mamba/issues/66
+    """
+    # NOTE: attention mask is a 2D boolean tensor
+    if attention_mask is not None:
+        dtype = hidden_states.dtype
+        hidden_states = (hidden_states * attention_mask[:, :, None]).to(dtype)
 
-logger = logging.get_logger(__name__)
+    return hidden_states
 
 
+@use_kernel_func_from_hub_with_fallback("causal_conv1d_update", "causal_conv1d")
 def causal_conv1d_update(
     hidden_states: torch.Tensor,
     conv_state: torch.Tensor,
@@ -77,6 +88,7 @@ def causal_conv1d_update(
     return out.to(hidden_states.dtype)
 
 
+@use_kernel_func_from_hub_with_fallback("causal_conv1d_fn", "causal_conv1d")
 def causal_conv1d_fn(
     hidden_states: torch.Tensor,
     weight: nn.Parameter,
@@ -99,25 +111,191 @@ def causal_conv1d_fn(
     return out.to(hidden_states.dtype)
 
 
-def rms_forward(hidden_states, variance_epsilon=1e-6):
-    """
-    Calculates simple RMSNorm with no learnable weights. `MambaRMSNorm` will
-    leverage this in order to multiply the final result with the RMSNorm weight
-
-    Args:
-        hidden_states (`torch.Tensor`):
-            Hidden states to normalize
-        variance_epsilon (`float`):
-            The eps value to add in the square root scaling factor
-    """
+@use_kernel_func_from_hub_with_fallback(
+    "mamba_inner_fn",
+    "mamba_ssm",
+)
+def mamba_inner_fn(
+    xz: torch.Tensor,
+    conv1d_weight: torch.Tensor,
+    conv1d_bias: torch.Tensor | None,
+    x_proj_weight: torch.Tensor,
+    delta_proj_weight: torch.Tensor,
+    out_proj_weight: torch.Tensor,
+    out_proj_bias: torch.Tensor | None,
+    A: torch.Tensor,
+    B: torch.Tensor | None = None,
+    C: torch.Tensor | None = None,
+    D: torch.Tensor | None = None,
+    delta_bias: torch.Tensor | None = None,
+    delta_softplus: bool = True,
+    b_rms_weight: torch.Tensor | None = None,
+    c_rms_weight: torch.Tensor | None = None,
+    dt_rms_weight: torch.Tensor | None = None,
+    b_c_dt_rms_eps: float = 1e-6,
+    **kwargs,
+):
+    return None
+
+
+@use_kernel_func_from_hub_with_fallback(
+    "selective_state_update",
+    "mamba_ssm",
+)
+def mamba_selective_state_update(
+    state: torch.Tensor,
+    hidden_states: torch.Tensor,
+    dt: torch.Tensor,
+    A: torch.Tensor,
+    B: torch.Tensor,
+    C: torch.Tensor,
+    D: torch.Tensor | None = None,
+    dt_bias: torch.Tensor | None = None,
+    dt_softplus: bool = False,
+    z: torch.Tensor | None = None,
+    **kwargs,
+):
     input_dtype = hidden_states.dtype
-    hidden_states = hidden_states.to(torch.float32)
 
-    variance = hidden_states.pow(2).mean(-1, keepdim=True)
-    hidden_states = hidden_states * torch.rsqrt(variance + variance_epsilon)
-    return hidden_states.to(input_dtype)
+    if dt_bias is not None:
+        dt = dt + dt_bias.to(dt.dtype)
+    if dt_softplus:
+        dt = F.softplus(dt)
+
+    # Discretize A
+    dA = torch.exp(dt.float()[..., None] * A.float()).to(device=state.device)
+
+    # Discretize B
+    dB = dt.float()[..., None] * B.float()[:, None, :]
+    # Discretize x into dB
+    dBx = dB * hidden_states.float()[..., None]
+
+    # State calculation
+    ssm_state = state.float() * dA + dBx
+    state.copy_(ssm_state.to(state.dtype))
+
+    # Subsequent output
+    out = torch.matmul(ssm_state.to(C.dtype), C.unsqueeze(-1)).squeeze(-1)
+
+    # D skip connection
+    if D is not None:
+        out = out + hidden_states * D
+
+    if z is not None:
+        out = out * F.silu(z)
+
+    return out.to(input_dtype)
+
+
+@use_kernel_func_from_hub_with_fallback(
+    "selective_scan_fn",
+    "mamba_ssm",
+)
+def mamba_selective_scan(
+    hidden_states: torch.Tensor,
+    dt: torch.Tensor,
+    A: torch.Tensor,
+    B: torch.Tensor,
+    C: torch.Tensor,
+    D: torch.Tensor | None = None,
+    z: torch.Tensor | None = None,
+    delta_bias: torch.Tensor | None = None,
+    delta_softplus: bool = False,
+    return_last_state: bool = False,
+    use_falcon_mambapy: bool = False,
+    use_associative_scan: bool = False,
+    **kwargs,
+):
+    # Torch only alternatives to the recurrent path
+    if is_torch_greater_or_equal("2.9.0"):
+        from torch._higher_order_ops.associative_scan import associative_scan
+    else:
+        associative_scan = None
+
+    if is_mambapy_available():
+        from mambapy.pscan import pscan
+    else:
+        pscan = None
+
+    batch_size, intermediate_size, seq_len = hidden_states.shape
+    input_dtype = hidden_states.dtype
+
+    if delta_bias is not None:
+        dt = dt + delta_bias.to(dt.dtype)[..., None]
+    if delta_softplus:
+        dt = F.softplus(dt)
+
+    # We need to transpose on the basis of the original kernel layout
+    B = B.transpose(1, 2)
+    C = C.transpose(1, 2)
+
+    # Discretize A and B for the entire sequence
+    discrete_A = torch.exp(A[None, :, None, :] * dt[:, :, :, None])
+    discrete_B = dt[:, :, :, None] * B[:, None, :, :].float()
+    deltaB_u = discrete_B * hidden_states[:, :, :, None].float()
+
+    if use_falcon_mambapy and pscan is not None:
+        all_states = pscan(discrete_A.transpose(1, 2), deltaB_u.transpose(1, 2))
+
+        scan_output = (all_states @ C.unsqueeze(-1)).squeeze(3).transpose(1, 2)
+        ssm_state = all_states[:, -1]
 
+    elif use_associative_scan and associative_scan is not None and is_tracing(hidden_states):
 
+        def combine_fn(left, right):
+            a_left, b_left = left
+            a_right, b_right = right
+            return a_left * a_right, a_right * b_left + b_right
+
+        combine_mode = "pointwise" if discrete_A.device.type in ("cuda", "xpu") else "generic"
+        _, all_states = associative_scan(
+            combine_fn,
+            (discrete_A, deltaB_u),
+            dim=2,
+            combine_mode=combine_mode,
+        )
+
+        scan_output = torch.matmul(all_states.transpose(1, 2).to(input_dtype), C.unsqueeze(-1))
+        scan_output = scan_output.squeeze(-1).transpose(1, 2)
+        ssm_state = all_states[:, :, -1]
+
+    # Recurrent iteration
+    else:
+        # "Initial hidden state" is not supported by the kernel path, so use
+        # the same zero initialization as the kernel
+        ssm_state = torch.zeros(
+            batch_size,
+            intermediate_size,
+            A.shape[-1],
+            dtype=input_dtype,
+            device=hidden_states.device,
+        )
+
+        scan_outputs = []
+        for index in range(seq_len):
+            # State calculation
+            ssm_state = discrete_A[:, :, index] * ssm_state + deltaB_u[:, :, index]
+
+            # Subsequent output
+            scan_output = torch.matmul(ssm_state.to(input_dtype), C[:, index, :].unsqueeze(-1))
+            scan_outputs.append(scan_output[:, :, 0])
+        scan_output = torch.stack(scan_outputs, dim=-1)
+
+    if D is not None:
+        scan_output = scan_output + hidden_states * D[None, :, None]
+
+    if z is not None:
+        scan_output = scan_output * F.silu(z)
+
+    if return_last_state:
+        return scan_output, ssm_state
+
+    return scan_output
+
+
+@use_kernelized_func(
+    [mamba_inner_fn, mamba_selective_scan, mamba_selective_state_update, causal_conv1d_fn, causal_conv1d_update]
+)
 class FalconMambaMixer(nn.Module):
     """
     Compute ∆, A, B, C, and D the state space parameters and compute the `contextualized_states`.
@@ -167,32 +345,10 @@ def __init__(self, config: FalconMambaConfig, layer_idx: int, initialize_mixer_w
         self.out_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=config.use_bias)
         self.use_bias = config.use_bias
 
-        global causal_conv1d, causal_conv1d_update, causal_conv1d_fn
-        causal_conv1d = lazy_load_kernel("causal-conv1d")
-        causal_conv1d_update = getattr(causal_conv1d, "causal_conv1d_update", causal_conv1d_update)
-        causal_conv1d_fn = getattr(causal_conv1d, "causal_conv1d_fn", causal_conv1d_fn)
-
-        global mamba_ssm, selective_state_update, selective_scan_fn, mamba_inner_fn
-        mamba_ssm = lazy_load_kernel("mamba-ssm")
-        selective_state_update = resolve_internal_import(
-            mamba_ssm, chained_path="ops.triton.selective_state_update.selective_state_update"
-        )
-        selective_scan_fn = getattr(mamba_ssm, "selective_scan_fn", None)
-        mamba_inner_fn = getattr(mamba_ssm, "mamba_inner_fn", None)
-
-        global is_fast_path_available
-        is_fast_path_available = (
-            all((selective_state_update, selective_scan_fn, mamba_inner_fn))
-            and hasattr(causal_conv1d, "causal_conv1d_update")
-            and hasattr(causal_conv1d, "causal_conv1d_fn")
-        )
-
-        self.warn_slow_implementation()
-
         self.layer_type = config.layer_types[layer_idx]
-        # Triton expects to pass RMS weights even if they are non learnable, thus we need to create these weights here
-        self.register_buffer("b_c_rms", torch.ones(self.ssm_state_size, requires_grad=False), persistent=False)
-        self.register_buffer("dt_rms", torch.ones(self.intermediate_size, requires_grad=False), persistent=False)
+        self.dt_layernorm = FalconMambaWeightlessRMSNorm(self.intermediate_size, eps=config.mixer_rms_eps)
+        self.b_layernorm = FalconMambaWeightlessRMSNorm(self.ssm_state_size, eps=config.mixer_rms_eps)
+        self.c_layernorm = FalconMambaWeightlessRMSNorm(self.ssm_state_size, eps=config.mixer_rms_eps)
         self.rms_eps = config.mixer_rms_eps
 
     @torch.no_grad()
@@ -216,86 +372,29 @@ def init_falcon_mamba_weights(self):
         # Inverse of softplus: https://github.com/pytorch/pytorch/issues/72759
         inv_dt = dt + torch.log(-torch.expm1(-dt))
         init.copy_(self.dt_proj.bias, inv_dt)
-        init.ones_(self.b_c_rms)
-        init.ones_(self.dt_rms)
-
-    def warn_slow_implementation(self):
-        if not is_fast_path_available:  # noqa
-            if self.use_falcon_mambapy:
-                if is_mambapy_available():
-                    logger.warning_once(
-                        "The fast path is not available because one of `(selective_state_update, selective_scan_fn, causal_conv1d_fn, causal_conv1d_update, mamba_inner_fn)`"
-                        " is None. Falling back to the mamba.py backend. The recommended way to enable the fast path is `pip install kernels`, which provides the"
-                        " FalconMamba kernels (loaded on demand). Alternatively, install mamba-ssm (https://github.com/state-spaces/mamba/#installation) and"
-                        " causal-conv1d (https://github.com/Dao-AILab/causal-conv1d)."
-                    )
-                else:
-                    raise ImportError(
-                        "use_mambapy is set to True but the mambapy package is not installed. To install it follow https://github.com/alxndrTL/mamba.py."
-                    )
-            else:
-                logger.warning_once(
-                    "The fast path is not available because one of `(selective_state_update, selective_scan_fn, causal_conv1d_fn, causal_conv1d_update, mamba_inner_fn)`"
-                    " is None. Falling back to the sequential implementation of Mamba, as use_mambapy is set to False. The recommended way to enable the fast path is"
-                    " `pip install kernels`, which provides the FalconMamba kernels (loaded on demand). Alternatively, install mamba-ssm"
-                    " (https://github.com/state-spaces/mamba/#installation) and causal-conv1d (https://github.com/Dao-AILab/causal-conv1d)."
-                    " For the mamba.py backend, follow https://github.com/alxndrTL/mamba.py."
-                )
+        init.ones_(self.dt_layernorm.weight)
+        init.ones_(self.b_layernorm.weight)
+        init.ones_(self.c_layernorm.weight)
 
-    def _convolution(
+    @force_accelerate_hooks(["conv1d", "dt_proj"])
+    def forward(
         self,
         hidden_states: torch.Tensor,
         cache_params: Cache | None = None,
-        attention_mask: torch.LongTensor | None = None,
+        attention_mask: torch.Tensor | None = None,
         **kwargs,
     ):
-        seq_len = hidden_states.shape[-1]
-        if attention_mask is not None:
-            hidden_states = hidden_states * attention_mask.unsqueeze(1)
-
+        seq_len = hidden_states.shape[1]
+        dtype = hidden_states.dtype
         use_precomputed_states = cache_params is not None and cache_params.has_previous_state(self.layer_idx)
 
-        if use_precomputed_states and seq_len == 1 and not cache_params.layers[self.layer_idx].record_past:
-            conv_state = cache_params.layers[self.layer_idx].conv_states[0]
-            hidden_states = causal_conv1d_update(
-                hidden_states,
-                conv_state,
-                self.conv1d.weight.squeeze(1),
-                self.conv1d.bias,
-                self.activation,
-            )
-        else:
-            if cache_params is not None:
-                hidden_states = cache_params.update_conv_state(
-                    hidden_states, self.layer_idx, conv_kernel_size=self.conv_kernel_size
-                )
-
-            hidden_states = causal_conv1d_fn(
-                hidden_states,
-                self.conv1d.weight.squeeze(1),
-                self.conv1d.bias,
-                activation=self.activation,
-                seq_idx=kwargs.get("seq_idx"),
-            )
-
-            # Drop the additional previous states
-            if cache_params is not None:
-                hidden_states = hidden_states[:, :, -seq_len:]
-
-        return hidden_states
-
-    def cuda_kernels_forward(
-        self,
-        hidden_states: torch.Tensor,
-        cache_params: Cache | None = None,
-        attention_mask: torch.LongTensor | None = None,
-        **kwargs,
-    ):
         # 1. Gated MLP's linear projection
+        hidden_states = apply_mask_to_padding_states(hidden_states, attention_mask)
         projected_states = self.in_proj(hidden_states).transpose(1, 2)
 
-        if self.training and cache_params is None:  # Doesn't support outputting the states -> used for training
-            return mamba_inner_fn(  # noqa
+        A = -torch.exp(self.A_log.float())
+        if self.training and cache_params is None:
+            fused_output = mamba_inner_fn(
                 projected_states,
                 self.conv1d.weight,
                 self.conv1d.bias if self.use_conv_bias else None,
@@ -303,219 +402,142 @@ def cuda_kernels_forward(
                 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()),
+                A,
                 None,  # input-dependent B
                 None,  # input-dependent C
                 self.D.float(),
                 delta_bias=self.dt_proj.bias.float(),
                 delta_softplus=True,
-                b_rms_weight=self.b_c_rms,
-                c_rms_weight=self.b_c_rms,
-                dt_rms_weight=self.dt_rms,
+                # Key difference: norms on B, C, and dt
+                b_rms_weight=self.b_layernorm.weight,
+                c_rms_weight=self.c_layernorm.weight,
+                dt_rms_weight=self.dt_layernorm.weight,
                 b_c_dt_rms_eps=self.rms_eps,
             )
 
-        hidden_states, gate = projected_states.chunk(2, dim=1)
-        use_precomputed_states = cache_params is not None and cache_params.has_previous_state(self.layer_idx)
+            # Only kernels can use this shortcircuit, fallback to normal torch otherwise
+            if fused_output is not None:
+                return fused_output
 
-        # Apply the conv
-        hidden_states = self._convolution(hidden_states, cache_params, attention_mask, **kwargs)
+        hidden_states_B_C, gate = projected_states.chunk(2, dim=1)
 
-        if attention_mask is not None:
-            hidden_states = hidden_states * attention_mask.unsqueeze(1)
+        if use_precomputed_states:
+            conv_state = cache_params.layers[self.layer_idx].conv_states[0]
+            recurrent_state = cache_params.layers[self.layer_idx].recurrent_states[0]
 
-        # 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))
+        # 2. Convolution sequence transformation
+        if use_precomputed_states and seq_len == 1 and not cache_params.layers[self.layer_idx].record_past:
+            hidden_states_B_C = causal_conv1d_update(
+                hidden_states_B_C,
+                conv_state,
+                self.conv1d.weight.squeeze(1),
+                self.conv1d.bias,
+                activation=self.activation,
+            )
+        else:
+            if cache_params is not None:
+                hidden_states_B_C = cache_params.update_conv_state(
+                    hidden_states_B_C,
+                    self.layer_idx,
+                    conv_kernel_size=self.conv_kernel_size,
+                )
+
+            hidden_states_B_C = causal_conv1d_fn(
+                hidden_states_B_C,
+                self.conv1d.weight.squeeze(1),
+                self.conv1d.bias,
+                activation=self.activation,
+                **kwargs,
+            )
+
+            if cache_params is not None:
+                hidden_states_B_C = hidden_states_B_C[:, :, -seq_len:]
+
+        # 3. SSM transformation
+        hidden_states_B_C = apply_mask_to_padding_states(hidden_states_B_C.transpose(1, 2), attention_mask)
         time_step, B, C = torch.split(
-            ssm_parameters, [self.time_step_rank, self.ssm_state_size, self.ssm_state_size], dim=-1
+            self.x_proj(hidden_states_B_C),
+            [self.time_step_rank, self.ssm_state_size, self.ssm_state_size],
+            dim=-1,
         )
 
-        B = rms_forward(B, variance_epsilon=self.rms_eps)
-        C = rms_forward(C, variance_epsilon=self.rms_eps)
-        time_step = rms_forward(time_step, variance_epsilon=self.rms_eps)
+        # Key difference: Additional norms on B, C, and dt
+        time_step = self.dt_layernorm(time_step)
+        B = self.b_layernorm(B)
+        C = self.c_layernorm(C)
 
         # In case the model has been quantized, we need a hack to properly call the `nn.Linear` module
         # at the price of a small overhead.
         if hasattr(self.config, "_is_quantized"):
-            discrete_time_step = (self.dt_proj(time_step) - self.dt_proj.bias).transpose(1, 2)
+            time_step = (self.dt_proj(time_step) - self.dt_proj.bias).transpose(1, 2)
         else:
-            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 use_precomputed_states:
-            scan_outputs = selective_state_update(
-                cache_params.layers[self.layer_idx].recurrent_states[0],
-                hidden_states[..., 0],
-                discrete_time_step[..., 0],
+            time_step = self.dt_proj.weight @ time_step.transpose(1, 2)
+        time_proj_bias = self.dt_proj.bias.float() if self.dt_proj.bias is not None else None
+
+        # Recurrent form
+        if use_precomputed_states and seq_len == 1:
+            scan_output = mamba_selective_state_update(
+                recurrent_state,
+                hidden_states_B_C.transpose(1, 2)[..., 0],
+                time_step[..., 0],
                 A,
                 B[:, 0],
                 C[:, 0],
                 self.D,
-                gate[..., 0],
-                time_proj_bias,
+                z=gate[..., 0],
+                dt_bias=time_proj_bias,
                 dt_softplus=True,
             ).unsqueeze(-1)
+
+        # Full sequence form
         else:
-            scan_outputs, ssm_state = selective_scan_fn(
-                hidden_states,
-                discrete_time_step,
+            output_final_state = cache_params is not None
+            scan_result = mamba_selective_scan(
+                hidden_states_B_C.transpose(1, 2),
+                time_step,
                 A,
                 B.transpose(1, 2),
                 C.transpose(1, 2),
-                self.D.float(),
-                gate,
-                time_proj_bias,
+                D=self.D.float(),
+                z=gate,
+                delta_bias=time_proj_bias,
                 delta_softplus=True,
-                return_last_state=True,
-            )
-            if ssm_state is not None and cache_params is not None:
-                cache_params.update_recurrent_state(ssm_state, self.layer_idx)
-
-        # 4. Final linear projection
-        contextualized_states = self.out_proj(scan_outputs.transpose(1, 2))
-
-        return contextualized_states
-
-    def slow_forward(
-        self,
-        hidden_states: torch.Tensor,
-        cache_params: Cache | None = None,
-        attention_mask: torch.LongTensor | None = None,
-        **kwargs,
-    ):
-        batch_size, seq_len, _ = hidden_states.shape
-        dtype = hidden_states.dtype
-        use_precomputed_states = cache_params is not None and cache_params.has_previous_state(self.layer_idx)
-        # 1. Gated MLP's linear projection
-        projected_states = self.in_proj(hidden_states).transpose(1, 2)  # [batch, 2 * intermediate_size, seq_len]
-        hidden_states, gate = projected_states.chunk(2, dim=1)
-
-        if attention_mask is not None:
-            hidden_states = hidden_states * attention_mask.unsqueeze(1)
-
-        # Apply the convolution
-        hidden_states = self._convolution(hidden_states, cache_params, attention_mask, **kwargs)
-
-        if attention_mask is not None:
-            hidden_states = hidden_states * attention_mask.unsqueeze(1)
-
-        if use_precomputed_states:
-            ssm_state = cache_params.layers[self.layer_idx].recurrent_states[0].clone()
-        else:
-            ssm_state = torch.zeros(
-                (batch_size, self.intermediate_size, self.ssm_state_size), device=hidden_states.device, dtype=dtype
+                return_last_state=output_final_state,
+                # TODO: rename to normal mambapy
+                use_mambapy=self.use_falcon_mambapy,
+                use_associative_scan=self.use_associative_scan,
             )
 
-        # 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
-        )
-
-        B = rms_forward(B, variance_epsilon=self.rms_eps)
-        C = rms_forward(C, variance_epsilon=self.rms_eps)
-        time_step = rms_forward(time_step, variance_epsilon=self.rms_eps)
-
-        discrete_time_step = self.dt_proj(time_step)  # [batch, seq_len, intermediate_size]
-        discrete_time_step = nn.functional.softplus(discrete_time_step).transpose(
-            1, 2
-        )  # [batch, intermediate_size, seq_len]
-
-        # 3.b. Discretization: B and C to [batch, seq_len, intermediate_size, ssm_state_size] (SRAM)
-        A = -torch.exp(self.A_log.float())  # [intermediate_size, 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()
-        )  # [batch, intermediate_size, seq_len, ssm_state_size]
-        deltaB_u = discrete_B * hidden_states[:, :, :, None].float()
-
-        # 3.c perform the recurrence y ← SSM(A, B, C)(x)
-        if self.use_falcon_mambapy and self.training and cache_params is None:
-            hs = pscan(
-                discrete_A.transpose(1, 2), deltaB_u.transpose(1, 2)
-            )  # [batch, seq_len, intermediate_size, ssm_state_size]
-            scan_output = (hs @ C.unsqueeze(-1)).squeeze(3).transpose(1, 2)  # [batch, intermediate_size, seq_len]
-            scan_output = scan_output + hidden_states * self.D[None, :, None]
-            scan_output = scan_output * self.act(gate)
-        else:
-            # Use associative_scan for parallel computation when available
-            if (
-                self.use_associative_scan
-                and associative_scan is not None
-                and is_tracing(hidden_states)
-                and cache_params is None
-            ):
-
-                def combine_fn(left, right):
-                    a_left, b_left = left
-                    a_right, b_right = right
-                    return (a_left * a_right, a_right * b_left + b_right)
-
-                combine_mode = "pointwise" if discrete_A.device.type in ("cuda", "xpu") else "generic"
-                _, all_h = associative_scan(combine_fn, (discrete_A, deltaB_u), dim=2, combine_mode=combine_mode)
-                # all_h: [B, D, S, N] -> output: [B, D, S]
-                scan_output = (
-                    torch.matmul(all_h.permute(0, 2, 1, 3).to(dtype), C.unsqueeze(-1)).squeeze(-1).permute(0, 2, 1)
-                )
-                ssm_state = all_h[:, :, -1, :]
+            if output_final_state:
+                scan_output, final_state = scan_result
+                cache_params.update_recurrent_state(final_state, self.layer_idx)
             else:
-                # Sequential loop for decoding or when associative_scan unavailable
-                scan_outputs = []
-                for i in range(seq_len):
-                    ssm_state = (
-                        discrete_A[:, :, i, :] * ssm_state + deltaB_u[:, :, i, :]
-                    )  # [batch, intermediate_size, ssm_state]
-                    scan_output = torch.matmul(
-                        ssm_state.to(dtype), C[:, i, :].unsqueeze(-1)
-                    )  # [batch, intermediate_size, 1]
-                    scan_outputs.append(scan_output[:, :, 0])
-                scan_output = torch.stack(scan_outputs, dim=-1)  # [batch, intermediate_size, seq_len]
-
-            scan_output = scan_output + (hidden_states * self.D[None, :, None])
-            scan_output = scan_output * self.act(gate)
-
-            if cache_params is not None:
-                cache_params.update_recurrent_state(ssm_state, self.layer_idx)
+                scan_output = scan_result
 
         # 4. Final linear projection
-        contextualized_states = self.out_proj(scan_output.transpose(1, 2))  # [batch, seq_len, hidden_size]
+        contextualized_states = self.out_proj(scan_output.transpose(1, 2).to(dtype))
         return contextualized_states
 
-    @force_accelerate_hooks("conv1d")
-    def forward(
-        self,
-        hidden_states,
-        cache_params: Cache | None = None,
-        attention_mask: torch.LongTensor | None = None,
-        **kwargs,
-    ):
-        if is_fast_path_available and "cuda" in self.x_proj.weight.device.type and not is_tracing(hidden_states):
-            return self.cuda_kernels_forward(hidden_states, cache_params, attention_mask, **kwargs)
-        return self.slow_forward(hidden_states, cache_params, attention_mask, **kwargs)
-
 
+@use_kernel_forward_from_hub("RMSNorm")
 class FalconMambaRMSNorm(nn.Module):
-    def __init__(self, hidden_size, eps=1e-6):
+    def __init__(self, hidden_size, eps: float = 1e-6) -> None:
         """
-        FalconMambaRMSNorm is equivalent to T5LayerNorm and LlamaRMSNorm
+        FalconMambaRMSNorm is equivalent to T5LayerNorm
         """
         super().__init__()
         self.weight = nn.Parameter(torch.ones(hidden_size))
         self.variance_epsilon = eps
 
-    def forward(self, hidden_states):
-        return self.weight.to(hidden_states.device) * rms_forward(
-            hidden_states, variance_epsilon=self.variance_epsilon
-        )
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        input_dtype = hidden_states.dtype
+        hidden_states = hidden_states.to(torch.float32)
+        variance = hidden_states.pow(2).mean(-1, keepdim=True)
+        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
+        return self.weight * hidden_states.to(input_dtype)
 
     def extra_repr(self):
-        return f"{self.weight.shape[0]}, eps={self.variance_epsilon}"
+        return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}"
 
 
 class FalconMambaBlock(GradientCheckpointingLayer):
diff --git a/src/transformers/models/falcon_mamba/modular_falcon_mamba.py b/src/transformers/models/falcon_mamba/modular_falcon_mamba.py
index 16155cc044fb..f65315338b85 100644
--- a/src/transformers/models/falcon_mamba/modular_falcon_mamba.py
+++ b/src/transformers/models/falcon_mamba/modular_falcon_mamba.py
@@ -20,7 +20,7 @@
 from ... import initialization as init
 from ...cache_utils import Cache
 from ...utils import auto_docstring, logging
-from ...utils.import_utils import is_mambapy_available, is_torch_greater_or_equal, is_tracing
+from ..llama.modeling_llama import LlamaRMSNorm
 from ..mamba.configuration_mamba import MambaConfig
 from ..mamba.modeling_mamba import (
     MambaBlock,
@@ -30,30 +30,18 @@
     MambaModel,
     MambaOutput,
     MambaPreTrainedModel,
-    MambaRMSNorm,
+    apply_mask_to_padding_states,
+    causal_conv1d_fn,
+    causal_conv1d_update,
+    mamba_inner_fn,
+    mamba_selective_scan,
+    mamba_selective_state_update,
 )
+from ..nanochat.modeling_nanochat import NanoChatRMSNorm
 
 
 logger = logging.get_logger(__name__)
 
-if is_torch_greater_or_equal("2.9.0"):
-    from torch._higher_order_ops.associative_scan import associative_scan
-else:
-    associative_scan = None
-
-if is_mambapy_available():
-    from mambapy.pscan import pscan
-else:
-    pscan = None
-
-selective_state_update, selective_scan_fn, causal_conv1d_fn, causal_conv1d_update, falcon_mamba_inner_fn = (
-    None,
-    None,
-    None,
-    None,
-    None,
-)
-
 
 @auto_docstring(checkpoint="tiiuae/falcon-mamba-7b")
 @strict
@@ -106,74 +94,46 @@ def layer_types(self):
         return ["linear_attention"] * self.num_hidden_layers
 
 
-def rms_forward(hidden_states, variance_epsilon=1e-6):
-    """
-    Calculates simple RMSNorm with no learnable weights. `MambaRMSNorm` will
-    leverage this in order to multiply the final result with the RMSNorm weight
-
-    Args:
-        hidden_states (`torch.Tensor`):
-            Hidden states to normalize
-        variance_epsilon (`float`):
-            The eps value to add in the square root scaling factor
-    """
-    input_dtype = hidden_states.dtype
-    hidden_states = hidden_states.to(torch.float32)
-
-    variance = hidden_states.pow(2).mean(-1, keepdim=True)
-    hidden_states = hidden_states * torch.rsqrt(variance + variance_epsilon)
-    return hidden_states.to(input_dtype)
+class FalconMambaWeightlessRMSNorm(NanoChatRMSNorm):
+    def __init__(self, hidden_size, eps: float = 1e-6):
+        super().__init__(eps)
+        # Dummy weights that are not used (only for imitating on kernels path)
+        self.register_buffer("weight", torch.ones(hidden_size, requires_grad=False), persistent=False)
 
 
 class FalconMambaMixer(MambaMixer):
     def __init__(self, config: FalconMambaConfig, layer_idx: int, initialize_mixer_weights: bool = True):
-        super().__init__(config, layer_idx)
-        # Triton expects to pass RMS weights even if they are non learnable, thus we need to create these weights here
-        self.register_buffer("b_c_rms", torch.ones(self.ssm_state_size, requires_grad=False), persistent=False)
-        self.register_buffer("dt_rms", torch.ones(self.intermediate_size, requires_grad=False), persistent=False)
+        super().__init__(config, layer_idx, initialize_mixer_weights)
+        self.dt_layernorm = FalconMambaWeightlessRMSNorm(self.intermediate_size, eps=config.mixer_rms_eps)
+        self.b_layernorm = FalconMambaWeightlessRMSNorm(self.ssm_state_size, eps=config.mixer_rms_eps)
+        self.c_layernorm = FalconMambaWeightlessRMSNorm(self.ssm_state_size, eps=config.mixer_rms_eps)
         self.rms_eps = config.mixer_rms_eps
 
     @torch.no_grad()
     def init_falcon_mamba_weights(self):
         super().init_falcon_mamba_weights()
-        init.ones_(self.b_c_rms)
-        init.ones_(self.dt_rms)
-
-    def warn_slow_implementation(self):
-        if not is_fast_path_available:  # noqa
-            if self.use_falcon_mambapy:
-                if is_mambapy_available():
-                    logger.warning_once(
-                        "The fast path is not available because one of `(selective_state_update, selective_scan_fn, causal_conv1d_fn, causal_conv1d_update, mamba_inner_fn)`"
-                        " is None. Falling back to the mamba.py backend. The recommended way to enable the fast path is `pip install kernels`, which provides the"
-                        " FalconMamba kernels (loaded on demand). Alternatively, install mamba-ssm (https://github.com/state-spaces/mamba/#installation) and"
-                        " causal-conv1d (https://github.com/Dao-AILab/causal-conv1d)."
-                    )
-                else:
-                    raise ImportError(
-                        "use_mambapy is set to True but the mambapy package is not installed. To install it follow https://github.com/alxndrTL/mamba.py."
-                    )
-            else:
-                logger.warning_once(
-                    "The fast path is not available because one of `(selective_state_update, selective_scan_fn, causal_conv1d_fn, causal_conv1d_update, mamba_inner_fn)`"
-                    " is None. Falling back to the sequential implementation of Mamba, as use_mambapy is set to False. The recommended way to enable the fast path is"
-                    " `pip install kernels`, which provides the FalconMamba kernels (loaded on demand). Alternatively, install mamba-ssm"
-                    " (https://github.com/state-spaces/mamba/#installation) and causal-conv1d (https://github.com/Dao-AILab/causal-conv1d)."
-                    " For the mamba.py backend, follow https://github.com/alxndrTL/mamba.py."
-                )
+        init.ones_(self.dt_layernorm.weight)
+        init.ones_(self.b_layernorm.weight)
+        init.ones_(self.c_layernorm.weight)
 
-    def cuda_kernels_forward(
+    def forward(
         self,
         hidden_states: torch.Tensor,
         cache_params: Cache | None = None,
-        attention_mask: torch.LongTensor | None = None,
+        attention_mask: torch.Tensor | None = None,
         **kwargs,
     ):
+        seq_len = hidden_states.shape[1]
+        dtype = hidden_states.dtype
+        use_precomputed_states = cache_params is not None and cache_params.has_previous_state(self.layer_idx)
+
         # 1. Gated MLP's linear projection
+        hidden_states = apply_mask_to_padding_states(hidden_states, attention_mask)
         projected_states = self.in_proj(hidden_states).transpose(1, 2)
 
-        if self.training and cache_params is None:  # Doesn't support outputting the states -> used for training
-            return mamba_inner_fn(  # noqa
+        A = -torch.exp(self.A_log.float())
+        if self.training and cache_params is None:
+            fused_output = mamba_inner_fn(
                 projected_states,
                 self.conv1d.weight,
                 self.conv1d.bias if self.use_conv_bias else None,
@@ -181,196 +141,125 @@ def cuda_kernels_forward(
                 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()),
+                A,
                 None,  # input-dependent B
                 None,  # input-dependent C
                 self.D.float(),
                 delta_bias=self.dt_proj.bias.float(),
                 delta_softplus=True,
-                b_rms_weight=self.b_c_rms,
-                c_rms_weight=self.b_c_rms,
-                dt_rms_weight=self.dt_rms,
+                # Key difference: norms on B, C, and dt
+                b_rms_weight=self.b_layernorm.weight,
+                c_rms_weight=self.c_layernorm.weight,
+                dt_rms_weight=self.dt_layernorm.weight,
                 b_c_dt_rms_eps=self.rms_eps,
             )
 
-        hidden_states, gate = projected_states.chunk(2, dim=1)
-        use_precomputed_states = cache_params is not None and cache_params.has_previous_state(self.layer_idx)
+            # Only kernels can use this shortcircuit, fallback to normal torch otherwise
+            if fused_output is not None:
+                return fused_output
+
+        hidden_states_B_C, gate = projected_states.chunk(2, dim=1)
+
+        if use_precomputed_states:
+            conv_state = cache_params.layers[self.layer_idx].conv_states[0]
+            recurrent_state = cache_params.layers[self.layer_idx].recurrent_states[0]
+
+        # 2. Convolution sequence transformation
+        if use_precomputed_states and seq_len == 1 and not cache_params.layers[self.layer_idx].record_past:
+            hidden_states_B_C = causal_conv1d_update(
+                hidden_states_B_C,
+                conv_state,
+                self.conv1d.weight.squeeze(1),
+                self.conv1d.bias,
+                activation=self.activation,
+            )
+        else:
+            if cache_params is not None:
+                hidden_states_B_C = cache_params.update_conv_state(
+                    hidden_states_B_C,
+                    self.layer_idx,
+                    conv_kernel_size=self.conv_kernel_size,
+                )
 
-        # Apply the conv
-        hidden_states = self._convolution(hidden_states, cache_params, attention_mask, **kwargs)
+            hidden_states_B_C = causal_conv1d_fn(
+                hidden_states_B_C,
+                self.conv1d.weight.squeeze(1),
+                self.conv1d.bias,
+                activation=self.activation,
+                **kwargs,
+            )
 
-        if attention_mask is not None:
-            hidden_states = hidden_states * attention_mask.unsqueeze(1)
+            if cache_params is not None:
+                hidden_states_B_C = hidden_states_B_C[:, :, -seq_len:]
 
-        # 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))
+        # 3. SSM transformation
+        hidden_states_B_C = apply_mask_to_padding_states(hidden_states_B_C.transpose(1, 2), attention_mask)
         time_step, B, C = torch.split(
-            ssm_parameters, [self.time_step_rank, self.ssm_state_size, self.ssm_state_size], dim=-1
+            self.x_proj(hidden_states_B_C),
+            [self.time_step_rank, self.ssm_state_size, self.ssm_state_size],
+            dim=-1,
         )
 
-        B = rms_forward(B, variance_epsilon=self.rms_eps)
-        C = rms_forward(C, variance_epsilon=self.rms_eps)
-        time_step = rms_forward(time_step, variance_epsilon=self.rms_eps)
+        # Key difference: Additional norms on B, C, and dt
+        time_step = self.dt_layernorm(time_step)
+        B = self.b_layernorm(B)
+        C = self.c_layernorm(C)
 
         # In case the model has been quantized, we need a hack to properly call the `nn.Linear` module
         # at the price of a small overhead.
         if hasattr(self.config, "_is_quantized"):
-            discrete_time_step = (self.dt_proj(time_step) - self.dt_proj.bias).transpose(1, 2)
+            time_step = (self.dt_proj(time_step) - self.dt_proj.bias).transpose(1, 2)
         else:
-            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 use_precomputed_states:
-            scan_outputs = selective_state_update(
-                cache_params.layers[self.layer_idx].recurrent_states[0],
-                hidden_states[..., 0],
-                discrete_time_step[..., 0],
+            time_step = self.dt_proj.weight @ time_step.transpose(1, 2)
+        time_proj_bias = self.dt_proj.bias.float() if self.dt_proj.bias is not None else None
+
+        # Recurrent form
+        if use_precomputed_states and seq_len == 1:
+            scan_output = mamba_selective_state_update(
+                recurrent_state,
+                hidden_states_B_C.transpose(1, 2)[..., 0],
+                time_step[..., 0],
                 A,
                 B[:, 0],
                 C[:, 0],
                 self.D,
-                gate[..., 0],
-                time_proj_bias,
+                z=gate[..., 0],
+                dt_bias=time_proj_bias,
                 dt_softplus=True,
             ).unsqueeze(-1)
+
+        # Full sequence form
         else:
-            scan_outputs, ssm_state = selective_scan_fn(
-                hidden_states,
-                discrete_time_step,
+            output_final_state = cache_params is not None
+            scan_result = mamba_selective_scan(
+                hidden_states_B_C.transpose(1, 2),
+                time_step,
                 A,
                 B.transpose(1, 2),
                 C.transpose(1, 2),
-                self.D.float(),
-                gate,
-                time_proj_bias,
+                D=self.D.float(),
+                z=gate,
+                delta_bias=time_proj_bias,
                 delta_softplus=True,
-                return_last_state=True,
+                return_last_state=output_final_state,
+                # TODO: rename to normal mambapy
+                use_mambapy=self.use_falcon_mambapy,
+                use_associative_scan=self.use_associative_scan,
             )
-            if ssm_state is not None and cache_params is not None:
-                cache_params.update_recurrent_state(ssm_state, self.layer_idx)
-
-        # 4. Final linear projection
-        contextualized_states = self.out_proj(scan_outputs.transpose(1, 2))
 
-        return contextualized_states
-
-    def slow_forward(
-        self,
-        hidden_states: torch.Tensor,
-        cache_params: Cache | None = None,
-        attention_mask: torch.LongTensor | None = None,
-        **kwargs,
-    ):
-        batch_size, seq_len, _ = hidden_states.shape
-        dtype = hidden_states.dtype
-        use_precomputed_states = cache_params is not None and cache_params.has_previous_state(self.layer_idx)
-        # 1. Gated MLP's linear projection
-        projected_states = self.in_proj(hidden_states).transpose(1, 2)  # [batch, 2 * intermediate_size, seq_len]
-        hidden_states, gate = projected_states.chunk(2, dim=1)
-
-        if attention_mask is not None:
-            hidden_states = hidden_states * attention_mask.unsqueeze(1)
-
-        # Apply the convolution
-        hidden_states = self._convolution(hidden_states, cache_params, attention_mask, **kwargs)
-
-        if attention_mask is not None:
-            hidden_states = hidden_states * attention_mask.unsqueeze(1)
-
-        if use_precomputed_states:
-            ssm_state = cache_params.layers[self.layer_idx].recurrent_states[0].clone()
-        else:
-            ssm_state = torch.zeros(
-                (batch_size, self.intermediate_size, self.ssm_state_size), device=hidden_states.device, dtype=dtype
-            )
-
-        # 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
-        )
-
-        B = rms_forward(B, variance_epsilon=self.rms_eps)
-        C = rms_forward(C, variance_epsilon=self.rms_eps)
-        time_step = rms_forward(time_step, variance_epsilon=self.rms_eps)
-
-        discrete_time_step = self.dt_proj(time_step)  # [batch, seq_len, intermediate_size]
-        discrete_time_step = nn.functional.softplus(discrete_time_step).transpose(
-            1, 2
-        )  # [batch, intermediate_size, seq_len]
-
-        # 3.b. Discretization: B and C to [batch, seq_len, intermediate_size, ssm_state_size] (SRAM)
-        A = -torch.exp(self.A_log.float())  # [intermediate_size, 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()
-        )  # [batch, intermediate_size, seq_len, ssm_state_size]
-        deltaB_u = discrete_B * hidden_states[:, :, :, None].float()
-
-        # 3.c perform the recurrence y ← SSM(A, B, C)(x)
-        if self.use_falcon_mambapy and self.training and cache_params is None:
-            hs = pscan(
-                discrete_A.transpose(1, 2), deltaB_u.transpose(1, 2)
-            )  # [batch, seq_len, intermediate_size, ssm_state_size]
-            scan_output = (hs @ C.unsqueeze(-1)).squeeze(3).transpose(1, 2)  # [batch, intermediate_size, seq_len]
-            scan_output = scan_output + hidden_states * self.D[None, :, None]
-            scan_output = scan_output * self.act(gate)
-        else:
-            # Use associative_scan for parallel computation when available
-            if (
-                self.use_associative_scan
-                and associative_scan is not None
-                and is_tracing(hidden_states)
-                and cache_params is None
-            ):
-
-                def combine_fn(left, right):
-                    a_left, b_left = left
-                    a_right, b_right = right
-                    return (a_left * a_right, a_right * b_left + b_right)
-
-                combine_mode = "pointwise" if discrete_A.device.type in ("cuda", "xpu") else "generic"
-                _, all_h = associative_scan(combine_fn, (discrete_A, deltaB_u), dim=2, combine_mode=combine_mode)
-                # all_h: [B, D, S, N] -> output: [B, D, S]
-                scan_output = (
-                    torch.matmul(all_h.permute(0, 2, 1, 3).to(dtype), C.unsqueeze(-1)).squeeze(-1).permute(0, 2, 1)
-                )
-                ssm_state = all_h[:, :, -1, :]
+            if output_final_state:
+                scan_output, final_state = scan_result
+                cache_params.update_recurrent_state(final_state, self.layer_idx)
             else:
-                # Sequential loop for decoding or when associative_scan unavailable
-                scan_outputs = []
-                for i in range(seq_len):
-                    ssm_state = (
-                        discrete_A[:, :, i, :] * ssm_state + deltaB_u[:, :, i, :]
-                    )  # [batch, intermediate_size, ssm_state]
-                    scan_output = torch.matmul(
-                        ssm_state.to(dtype), C[:, i, :].unsqueeze(-1)
-                    )  # [batch, intermediate_size, 1]
-                    scan_outputs.append(scan_output[:, :, 0])
-                scan_output = torch.stack(scan_outputs, dim=-1)  # [batch, intermediate_size, seq_len]
-
-            scan_output = scan_output + (hidden_states * self.D[None, :, None])
-            scan_output = scan_output * self.act(gate)
-
-            if cache_params is not None:
-                cache_params.update_recurrent_state(ssm_state, self.layer_idx)
+                scan_output = scan_result
 
         # 4. Final linear projection
-        contextualized_states = self.out_proj(scan_output.transpose(1, 2))  # [batch, seq_len, hidden_size]
+        contextualized_states = self.out_proj(scan_output.transpose(1, 2).to(dtype))
         return contextualized_states
 
 
-class FalconMambaRMSNorm(MambaRMSNorm):
-    def forward(self, hidden_states):
-        return self.weight.to(hidden_states.device) * rms_forward(
-            hidden_states, variance_epsilon=self.variance_epsilon
-        )
+class FalconMambaRMSNorm(LlamaRMSNorm):
+    pass
 
 
 class FalconMambaBlock(MambaBlock):
diff --git a/src/transformers/models/jamba/modeling_jamba.py b/src/transformers/models/jamba/modeling_jamba.py
index 18deed22f1d4..04b33eb42805 100755
--- a/src/transformers/models/jamba/modeling_jamba.py
+++ b/src/transformers/models/jamba/modeling_jamba.py
@@ -33,9 +33,9 @@
 from ...cache_utils import Cache, DynamicCache
 from ...generation import GenerationMixin
 from ...integrations import (
-    lazy_load_kernel,
     use_experts_implementation,
     use_kernel_forward_from_hub,
+    use_kernel_func_from_hub_with_fallback,
     use_kernelized_func,
 )
 from ...integrations.accelerate import force_accelerate_hooks
@@ -44,19 +44,13 @@
 from ...modeling_outputs import MoeCausalLMOutputWithPast, MoeModelOutputWithPast
 from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
 from ...processing_utils import Unpack
-from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, logging
+from ...utils import TransformersKwargs, auto_docstring, can_return_tuple
 from ...utils.generic import merge_with_config_defaults
-from ...utils.import_utils import (
-    is_tracing,
-    resolve_internal_import,
-)
+from ...utils.import_utils import is_mambapy_available, is_torch_greater_or_equal, is_tracing
 from ...utils.output_capturing import OutputRecorder, capture_outputs
 from .configuration_jamba import JambaConfig
 
 
-logger = logging.get_logger(__name__)
-
-
 @use_kernel_forward_from_hub("RMSNorm")
 class JambaRMSNorm(nn.Module):
     def __init__(self, hidden_size, eps: float = 1e-6) -> None:
@@ -203,6 +197,19 @@ def forward(
         return attn_output, attn_weights
 
 
+def apply_mask_to_padding_states(hidden_states, attention_mask):
+    """
+    Tunes out the hidden states for padding tokens, see https://github.com/state-spaces/jamba/issues/66
+    """
+    # NOTE: attention mask is a 2D boolean tensor
+    if attention_mask is not None:
+        dtype = hidden_states.dtype
+        hidden_states = (hidden_states * attention_mask[:, :, None]).to(dtype)
+
+    return hidden_states
+
+
+@use_kernel_func_from_hub_with_fallback("causal_conv1d_update", "causal_conv1d")
 def causal_conv1d_update(
     hidden_states: torch.Tensor,
     conv_state: torch.Tensor,
@@ -222,6 +229,7 @@ def causal_conv1d_update(
     return out.to(hidden_states.dtype)
 
 
+@use_kernel_func_from_hub_with_fallback("causal_conv1d_fn", "causal_conv1d")
 def causal_conv1d_fn(
     hidden_states: torch.Tensor,
     weight: nn.Parameter,
@@ -244,6 +252,191 @@ def causal_conv1d_fn(
     return out.to(hidden_states.dtype)
 
 
+@use_kernel_func_from_hub_with_fallback(
+    "mamba_inner_fn",
+    "mamba_ssm",
+)
+def mamba_inner_fn(
+    xz: torch.Tensor,
+    conv1d_weight: torch.Tensor,
+    conv1d_bias: torch.Tensor | None,
+    x_proj_weight: torch.Tensor,
+    delta_proj_weight: torch.Tensor,
+    out_proj_weight: torch.Tensor,
+    out_proj_bias: torch.Tensor | None,
+    A: torch.Tensor,
+    B: torch.Tensor | None = None,
+    C: torch.Tensor | None = None,
+    D: torch.Tensor | None = None,
+    delta_bias: torch.Tensor | None = None,
+    delta_softplus: bool = True,
+    b_rms_weight: torch.Tensor | None = None,
+    c_rms_weight: torch.Tensor | None = None,
+    dt_rms_weight: torch.Tensor | None = None,
+    b_c_dt_rms_eps: float = 1e-6,
+    **kwargs,
+):
+    return None
+
+
+@use_kernel_func_from_hub_with_fallback(
+    "selective_state_update",
+    "mamba_ssm",
+)
+def mamba_selective_state_update(
+    state: torch.Tensor,
+    hidden_states: torch.Tensor,
+    dt: torch.Tensor,
+    A: torch.Tensor,
+    B: torch.Tensor,
+    C: torch.Tensor,
+    D: torch.Tensor | None = None,
+    dt_bias: torch.Tensor | None = None,
+    dt_softplus: bool = False,
+    z: torch.Tensor | None = None,
+    **kwargs,
+):
+    input_dtype = hidden_states.dtype
+
+    if dt_bias is not None:
+        dt = dt + dt_bias.to(dt.dtype)
+    if dt_softplus:
+        dt = F.softplus(dt)
+
+    # Discretize A
+    dA = torch.exp(dt.float()[..., None] * A.float()).to(device=state.device)
+
+    # Discretize B
+    dB = dt.float()[..., None] * B.float()[:, None, :]
+    # Discretize x into dB
+    dBx = dB * hidden_states.float()[..., None]
+
+    # State calculation
+    ssm_state = state.float() * dA + dBx
+    state.copy_(ssm_state.to(state.dtype))
+
+    # Subsequent output
+    out = torch.matmul(ssm_state.to(C.dtype), C.unsqueeze(-1)).squeeze(-1)
+
+    # D skip connection
+    if D is not None:
+        out = out + hidden_states * D
+
+    if z is not None:
+        out = out * F.silu(z)
+
+    return out.to(input_dtype)
+
+
+@use_kernel_func_from_hub_with_fallback(
+    "selective_scan_fn",
+    "mamba_ssm",
+)
+def mamba_selective_scan(
+    hidden_states: torch.Tensor,
+    dt: torch.Tensor,
+    A: torch.Tensor,
+    B: torch.Tensor,
+    C: torch.Tensor,
+    D: torch.Tensor | None = None,
+    z: torch.Tensor | None = None,
+    delta_bias: torch.Tensor | None = None,
+    delta_softplus: bool = False,
+    return_last_state: bool = False,
+    use_jambapy: bool = False,
+    use_associative_scan: bool = False,
+    **kwargs,
+):
+    # Torch only alternatives to the recurrent path
+    if is_torch_greater_or_equal("2.9.0"):
+        from torch._higher_order_ops.associative_scan import associative_scan
+    else:
+        associative_scan = None
+
+    if is_mambapy_available():
+        from mambapy.pscan import pscan
+    else:
+        pscan = None
+
+    batch_size, intermediate_size, seq_len = hidden_states.shape
+    input_dtype = hidden_states.dtype
+
+    if delta_bias is not None:
+        dt = dt + delta_bias.to(dt.dtype)[..., None]
+    if delta_softplus:
+        dt = F.softplus(dt)
+
+    # We need to transpose on the basis of the original kernel layout
+    B = B.transpose(1, 2)
+    C = C.transpose(1, 2)
+
+    # Discretize A and B for the entire sequence
+    discrete_A = torch.exp(A[None, :, None, :] * dt[:, :, :, None])
+    discrete_B = dt[:, :, :, None] * B[:, None, :, :].float()
+    deltaB_u = discrete_B * hidden_states[:, :, :, None].float()
+
+    if use_jambapy and pscan is not None:
+        all_states = pscan(discrete_A.transpose(1, 2), deltaB_u.transpose(1, 2))
+
+        scan_output = (all_states @ C.unsqueeze(-1)).squeeze(3).transpose(1, 2)
+        ssm_state = all_states[:, -1]
+
+    elif use_associative_scan and associative_scan is not None and is_tracing(hidden_states):
+
+        def combine_fn(left, right):
+            a_left, b_left = left
+            a_right, b_right = right
+            return a_left * a_right, a_right * b_left + b_right
+
+        combine_mode = "pointwise" if discrete_A.device.type in ("cuda", "xpu") else "generic"
+        _, all_states = associative_scan(
+            combine_fn,
+            (discrete_A, deltaB_u),
+            dim=2,
+            combine_mode=combine_mode,
+        )
+
+        scan_output = torch.matmul(all_states.transpose(1, 2).to(input_dtype), C.unsqueeze(-1))
+        scan_output = scan_output.squeeze(-1).transpose(1, 2)
+        ssm_state = all_states[:, :, -1]
+
+    # Recurrent iteration
+    else:
+        # "Initial hidden state" is not supported by the kernel path, so use
+        # the same zero initialization as the kernel
+        ssm_state = torch.zeros(
+            batch_size,
+            intermediate_size,
+            A.shape[-1],
+            dtype=input_dtype,
+            device=hidden_states.device,
+        )
+
+        scan_outputs = []
+        for index in range(seq_len):
+            # State calculation
+            ssm_state = discrete_A[:, :, index] * ssm_state + deltaB_u[:, :, index]
+
+            # Subsequent output
+            scan_output = torch.matmul(ssm_state.to(input_dtype), C[:, index, :].unsqueeze(-1))
+            scan_outputs.append(scan_output[:, :, 0])
+        scan_output = torch.stack(scan_outputs, dim=-1)
+
+    if D is not None:
+        scan_output = scan_output + hidden_states * D[None, :, None]
+
+    if z is not None:
+        scan_output = scan_output * F.silu(z)
+
+    if return_last_state:
+        return scan_output, ssm_state
+
+    return scan_output
+
+
+@use_kernelized_func(
+    [mamba_inner_fn, mamba_selective_scan, mamba_selective_state_update, causal_conv1d_fn, causal_conv1d_update]
+)
 class JambaMambaMixer(nn.Module):
     """
     Compute ∆, A, B, C, and D the state space parameters and compute the `contextualized_states`.
@@ -295,236 +488,125 @@ def __init__(self, config: JambaConfig, layer_idx):
         self.b_layernorm = JambaRMSNorm(self.ssm_state_size, eps=config.rms_norm_eps)
         self.c_layernorm = JambaRMSNorm(self.ssm_state_size, eps=config.rms_norm_eps)
 
-        global causal_conv1d, causal_conv1d_update, causal_conv1d_fn
-        causal_conv1d = lazy_load_kernel("causal-conv1d")
-        causal_conv1d_update = getattr(causal_conv1d, "causal_conv1d_update", causal_conv1d_update)
-        causal_conv1d_fn = getattr(causal_conv1d, "causal_conv1d_fn", causal_conv1d_fn)
-
-        global mamba_ssm, selective_state_update, selective_scan_fn
-        mamba_ssm = lazy_load_kernel("mamba-ssm")
-        selective_state_update = resolve_internal_import(
-            mamba_ssm, chained_path="ops.triton.selective_state_update.selective_state_update"
-        )
-        selective_scan_fn = getattr(mamba_ssm, "selective_scan_fn", None)
-
-        global is_fast_path_available
-        is_fast_path_available = (
-            all((selective_state_update, selective_scan_fn))
-            and hasattr(causal_conv1d, "causal_conv1d_update")
-            and hasattr(causal_conv1d, "causal_conv1d_fn")
-        )
-
-        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. To install follow https://github.com/state-spaces/mamba/#installation and https://github.com/Dao-AILab/causal-conv1d."
-            )
-
-        self.layer_type = config.layer_types[layer_idx]
-
-    def _convolution(
+    @force_accelerate_hooks(["conv1d", "dt_proj"])
+    def forward(
         self,
         hidden_states: torch.Tensor,
         cache_params: Cache | None = None,
-        attention_mask: torch.LongTensor | None = None,
+        attention_mask: torch.Tensor | None = None,
         **kwargs,
     ):
-        seq_len = hidden_states.shape[-1]
-        if attention_mask is not None:
-            hidden_states = hidden_states * attention_mask.unsqueeze(1)
-
+        seq_len = hidden_states.shape[1]
+        dtype = hidden_states.dtype
         use_precomputed_states = cache_params is not None and cache_params.has_previous_state(self.layer_idx)
 
-        if use_precomputed_states and seq_len == 1 and not cache_params.layers[self.layer_idx].record_past:
+        # 1. Gated MLP's linear projection
+        hidden_states = apply_mask_to_padding_states(hidden_states, attention_mask)
+        projected_states = self.in_proj(hidden_states).transpose(1, 2)
+
+        # Key difference to falcon mamba: mamba inner fn applies its dt norm at a different time
+        # leading to incompatible implementations
+        A = -torch.exp(self.A_log.float())
+
+        hidden_states_B_C, gate = projected_states.chunk(2, dim=1)
+
+        if use_precomputed_states:
             conv_state = cache_params.layers[self.layer_idx].conv_states[0]
-            hidden_states = causal_conv1d_update(
-                hidden_states,
+            recurrent_state = cache_params.layers[self.layer_idx].recurrent_states[0]
+
+        # 2. Convolution sequence transformation
+        if use_precomputed_states and seq_len == 1 and not cache_params.layers[self.layer_idx].record_past:
+            hidden_states_B_C = causal_conv1d_update(
+                hidden_states_B_C,
                 conv_state,
                 self.conv1d.weight.squeeze(1),
                 self.conv1d.bias,
-                self.activation,
+                activation=self.activation,
             )
         else:
             if cache_params is not None:
-                hidden_states = cache_params.update_conv_state(
-                    hidden_states, self.layer_idx, conv_kernel_size=self.conv_kernel_size
+                hidden_states_B_C = cache_params.update_conv_state(
+                    hidden_states_B_C,
+                    self.layer_idx,
+                    conv_kernel_size=self.conv_kernel_size,
                 )
 
-            hidden_states = causal_conv1d_fn(
-                hidden_states,
+            hidden_states_B_C = causal_conv1d_fn(
+                hidden_states_B_C,
                 self.conv1d.weight.squeeze(1),
                 self.conv1d.bias,
                 activation=self.activation,
-                seq_idx=kwargs.get("seq_idx"),
+                **kwargs,
             )
 
-            # Drop the additional previous states
             if cache_params is not None:
-                hidden_states = hidden_states[:, :, -seq_len:]
-
-        return hidden_states
-
-    def cuda_kernels_forward(
-        self,
-        hidden_states: torch.Tensor,
-        cache_params: Cache | None = None,
-        attention_mask: torch.Tensor | None = None,
-        **kwargs,
-    ):
-        # Note: we cannot use `mamba_inner_fn` as in mamba even if in training and without cache params because we have the
-        # inner layernorms which isn't supported by this fused kernel
-        batch_size, seq_len, _ = hidden_states.shape
-        use_precomputed_states = (
-            cache_params is not None and cache_params.has_previous_state(self.layer_idx) and seq_len == 1
-        )
-
-        # 1. Gated MLP's linear projection
-        projected_states = self.in_proj(hidden_states).transpose(1, 2)
-        hidden_states, gate = projected_states.chunk(2, dim=1)
-
-        # Apply the conv
-        hidden_states = self._convolution(hidden_states, cache_params, attention_mask, **kwargs)
+                hidden_states_B_C = hidden_states_B_C[:, :, -seq_len:]
 
-        if attention_mask is not None:
-            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))
+        # 3. SSM transformation
+        hidden_states_B_C = apply_mask_to_padding_states(hidden_states_B_C.transpose(1, 2), attention_mask)
         time_step, B, C = torch.split(
-            ssm_parameters, [self.time_step_rank, self.ssm_state_size, self.ssm_state_size], dim=-1
+            self.x_proj(hidden_states_B_C),
+            [self.time_step_rank, self.ssm_state_size, self.ssm_state_size],
+            dim=-1,
         )
 
+        # Key difference to (falcon) mamba: Additional norms on B, C, and dt
         time_step = self.dt_layernorm(time_step)
         B = self.b_layernorm(B)
         C = self.c_layernorm(C)
 
-        # Here we need to apply dt_proj without the bias, as the bias is added in the selective scan kernel.
-        # This is a hack to apply dt_proj while still using the forward pass of `torch.nn.Linear`, which is needed
-        # in order to make quantization work. Quantization code replaces `torch.nn.Linear` layers with quantized
-        # linear layers, and requires to call the forward pass directly.
-        # Quantized model can't work with the original code:
-        # ```discrete_time_step = self.dt_proj.weight @ time_step.transpose(1, 2)```
-        time_proj_bias = self.dt_proj.bias.data
-        with torch.no_grad():
-            self.dt_proj.bias.data = torch.zeros_like(self.dt_proj.bias.data)
-        discrete_time_step = self.dt_proj(time_step).transpose(1, 2)
-        with torch.no_grad():
-            self.dt_proj.bias.data = time_proj_bias
-
-        A = -torch.exp(self.A_log.float())
-        # 3.c perform the recurrence y ← SSM(A, B, C)(x)
-        time_proj_bias = time_proj_bias.float() if time_proj_bias is not None else None
-        if use_precomputed_states:
-            scan_outputs = selective_state_update(
-                cache_params.layers[self.layer_idx].recurrent_states[0],
-                hidden_states[..., 0],
-                discrete_time_step[..., 0],
+        # In case the model has been quantized, we need a hack to properly call the `nn.Linear` module
+        # at the price of a small overhead.
+        if hasattr(self.config, "_is_quantized"):
+            time_step = (self.dt_proj(time_step) - self.dt_proj.bias).transpose(1, 2)
+        else:
+            time_step = self.dt_proj.weight @ time_step.transpose(1, 2)
+        time_proj_bias = self.dt_proj.bias.float() if self.dt_proj.bias is not None else None
+
+        # Recurrent form
+        if use_precomputed_states and seq_len == 1:
+            scan_output = mamba_selective_state_update(
+                recurrent_state,
+                hidden_states_B_C.transpose(1, 2)[..., 0],
+                time_step[..., 0],
                 A,
                 B[:, 0],
                 C[:, 0],
                 self.D,
-                gate[..., 0],
-                time_proj_bias,
+                z=gate[..., 0],
+                dt_bias=time_proj_bias,
                 dt_softplus=True,
             ).unsqueeze(-1)
+
+        # Full sequence form
         else:
-            scan_outputs, ssm_state = selective_scan_fn(
-                hidden_states,
-                discrete_time_step,
+            output_final_state = cache_params is not None
+            scan_result = mamba_selective_scan(
+                hidden_states_B_C.transpose(1, 2),
+                time_step,
                 A,
                 B.transpose(1, 2),
                 C.transpose(1, 2),
-                self.D.float(),
-                gate,
-                time_proj_bias,
+                D=self.D.float(),
+                z=gate,
+                delta_bias=time_proj_bias,
                 delta_softplus=True,
-                return_last_state=True,
+                return_last_state=output_final_state,
+                # TODO: No faster alternatives for mamba atm (needs config adjustments)
+                use_mambapy=False,
+                use_associative_scan=False,
             )
-            if ssm_state is not None and cache_params is not None:
-                cache_params.update_recurrent_state(ssm_state, self.layer_idx)
 
-        # 4. Final linear projection
-        contextualized_states = self.out_proj(scan_outputs.transpose(1, 2))
-
-        return contextualized_states
-
-    def slow_forward(
-        self,
-        hidden_states: torch.Tensor,
-        cache_params: Cache | None = None,
-        attention_mask: torch.LongTensor | None = None,
-        **kwargs,
-    ):
-        batch_size, seq_len, _ = hidden_states.shape
-        dtype = hidden_states.dtype
-        use_precomputed_states = cache_params is not None and cache_params.has_previous_state(self.layer_idx)
-        # 1. Gated MLP's linear projection
-        projected_states = self.in_proj(hidden_states).transpose(1, 2)
-        hidden_states, gate = projected_states.chunk(2, dim=1)
-
-        # Apply the convolution
-        hidden_states = self._convolution(hidden_states, cache_params, attention_mask, **kwargs)
-
-        if attention_mask is not None:
-            hidden_states = hidden_states * attention_mask.unsqueeze(1)
-
-        if use_precomputed_states:
-            # In training mode, we don't want to perform in-place operations on ssm_state so we can compute the backwards pass
-            ssm_state = cache_params.layers[self.layer_idx].recurrent_states[0].clone()
-        else:
-            ssm_state = torch.zeros(
-                (batch_size, self.intermediate_size, self.ssm_state_size), device=hidden_states.device, dtype=dtype
-            )
-
-        # 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
-        )
-
-        time_step = self.dt_layernorm(time_step)
-        B = self.b_layernorm(B)
-        C = self.c_layernorm(C)
-
-        discrete_time_step = self.dt_proj(time_step)
-        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)
-        A = -torch.exp(self.A_log.float())
-        discrete_A = torch.exp(A[None, :, None, :] * discrete_time_step[:, :, :, None])
-        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(seq_len):
-            ssm_state = discrete_A[:, :, i, :] * ssm_state + deltaB_u[:, :, i, :]
-            scan_output = torch.matmul(ssm_state.to(dtype), C[:, i, :].unsqueeze(-1))
-            scan_outputs.append(scan_output[:, :, 0])
-        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)
-
-        if cache_params is not None:
-            cache_params.update_recurrent_state(ssm_state, self.layer_idx)
+            if output_final_state:
+                scan_output, final_state = scan_result
+                cache_params.update_recurrent_state(final_state, self.layer_idx)
+            else:
+                scan_output = scan_result
 
         # 4. Final linear projection
-        contextualized_states = self.out_proj(scan_output.transpose(1, 2))
+        contextualized_states = self.out_proj(scan_output.transpose(1, 2).to(dtype))
         return contextualized_states
 
-    @force_accelerate_hooks("conv1d")
-    def forward(
-        self,
-        hidden_states,
-        cache_params: Cache | None = None,
-        attention_mask: torch.LongTensor | None = None,
-        **kwargs,
-    ):
-        if is_fast_path_available and "cuda" in self.x_proj.weight.device.type and not is_tracing(hidden_states):
-            return self.cuda_kernels_forward(hidden_states, cache_params, attention_mask, **kwargs)
-        return self.slow_forward(hidden_states, cache_params, attention_mask, **kwargs)
-
 
 class JambaMLP(nn.Module):
     def __init__(self, config):
diff --git a/src/transformers/models/jamba/modular_jamba.py b/src/transformers/models/jamba/modular_jamba.py
index 076ff55be418..03e8e769986e 100644
--- a/src/transformers/models/jamba/modular_jamba.py
+++ b/src/transformers/models/jamba/modular_jamba.py
@@ -24,7 +24,6 @@
 from ... import initialization as init
 from ...activations import ACT2FN
 from ...cache_utils import Cache, DynamicCache
-from ...integrations import lazy_load_kernel
 from ...masking_utils import create_causal_mask, create_recurrent_attention_mask
 from ...modeling_layers import GenericForSequenceClassification, GradientCheckpointingLayer
 from ...modeling_outputs import MoeCausalLMOutputWithPast, MoeModelOutputWithPast
@@ -32,10 +31,16 @@
 from ...processing_utils import Unpack
 from ...utils import TransformersKwargs, auto_docstring, logging
 from ...utils.generic import merge_with_config_defaults
-from ...utils.import_utils import resolve_internal_import
 from ...utils.output_capturing import OutputRecorder, capture_outputs
+from ..falcon_mamba.modeling_falcon_mamba import FalconMambaMixer
 from ..llama.modeling_llama import LlamaAttention, LlamaRMSNorm, eager_attention_forward
-from ..mamba.modeling_mamba import MambaMixer, causal_conv1d_fn, causal_conv1d_update
+from ..mamba.modeling_mamba import (
+    apply_mask_to_padding_states,
+    causal_conv1d_fn,
+    causal_conv1d_update,
+    mamba_selective_scan,
+    mamba_selective_state_update,
+)
 from ..mistral.modeling_mistral import MistralMLP
 from ..mixtral.modeling_mixtral import MixtralExperts, MixtralForCausalLM
 from .configuration_jamba import JambaConfig
@@ -93,7 +98,7 @@ def forward(
         return attn_output, attn_weights
 
 
-class JambaMambaMixer(MambaMixer):
+class JambaMambaMixer(FalconMambaMixer):
     """
     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)
@@ -144,186 +149,125 @@ def __init__(self, config: JambaConfig, layer_idx):
         self.b_layernorm = JambaRMSNorm(self.ssm_state_size, eps=config.rms_norm_eps)
         self.c_layernorm = JambaRMSNorm(self.ssm_state_size, eps=config.rms_norm_eps)
 
-        global causal_conv1d, causal_conv1d_update, causal_conv1d_fn
-        causal_conv1d = lazy_load_kernel("causal-conv1d")
-        causal_conv1d_update = getattr(causal_conv1d, "causal_conv1d_update", causal_conv1d_update)
-        causal_conv1d_fn = getattr(causal_conv1d, "causal_conv1d_fn", causal_conv1d_fn)
-
-        global mamba_ssm, selective_state_update, selective_scan_fn
-        mamba_ssm = lazy_load_kernel("mamba-ssm")
-        selective_state_update = resolve_internal_import(
-            mamba_ssm, chained_path="ops.triton.selective_state_update.selective_state_update"
-        )
-        selective_scan_fn = getattr(mamba_ssm, "selective_scan_fn", None)
-
-        global is_fast_path_available
-        is_fast_path_available = (
-            all((selective_state_update, selective_scan_fn))
-            and hasattr(causal_conv1d, "causal_conv1d_update")
-            and hasattr(causal_conv1d, "causal_conv1d_fn")
-        )
-
-        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. To install follow https://github.com/state-spaces/mamba/#installation and https://github.com/Dao-AILab/causal-conv1d."
-            )
-
-        self.layer_type = config.layer_types[layer_idx]
-
-    def warn_slow_implementation(self):
+    def init_jamba_weights(self):
         raise NotImplementedError("Not needed for jamba")
 
-    def init_jamba_mamba_weights(self):
-        raise NotImplementedError("Not needed for jamba")
-
-    def cuda_kernels_forward(
+    def forward(
         self,
         hidden_states: torch.Tensor,
         cache_params: Cache | None = None,
         attention_mask: torch.Tensor | None = None,
         **kwargs,
     ):
-        # Note: we cannot use `mamba_inner_fn` as in mamba even if in training and without cache params because we have the
-        # inner layernorms which isn't supported by this fused kernel
-        batch_size, seq_len, _ = hidden_states.shape
-        use_precomputed_states = (
-            cache_params is not None and cache_params.has_previous_state(self.layer_idx) and seq_len == 1
-        )
+        seq_len = hidden_states.shape[1]
+        dtype = hidden_states.dtype
+        use_precomputed_states = cache_params is not None and cache_params.has_previous_state(self.layer_idx)
 
         # 1. Gated MLP's linear projection
+        hidden_states = apply_mask_to_padding_states(hidden_states, attention_mask)
         projected_states = self.in_proj(hidden_states).transpose(1, 2)
-        hidden_states, gate = projected_states.chunk(2, dim=1)
 
-        # Apply the conv
-        hidden_states = self._convolution(hidden_states, cache_params, attention_mask, **kwargs)
+        # Key difference to falcon mamba: mamba inner fn applies its dt norm at a different time
+        # leading to incompatible implementations
+        A = -torch.exp(self.A_log.float())
 
-        if attention_mask is not None:
-            hidden_states = hidden_states * attention_mask.unsqueeze(1)
+        hidden_states_B_C, gate = projected_states.chunk(2, dim=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))
+        if use_precomputed_states:
+            conv_state = cache_params.layers[self.layer_idx].conv_states[0]
+            recurrent_state = cache_params.layers[self.layer_idx].recurrent_states[0]
+
+        # 2. Convolution sequence transformation
+        if use_precomputed_states and seq_len == 1 and not cache_params.layers[self.layer_idx].record_past:
+            hidden_states_B_C = causal_conv1d_update(
+                hidden_states_B_C,
+                conv_state,
+                self.conv1d.weight.squeeze(1),
+                self.conv1d.bias,
+                activation=self.activation,
+            )
+        else:
+            if cache_params is not None:
+                hidden_states_B_C = cache_params.update_conv_state(
+                    hidden_states_B_C,
+                    self.layer_idx,
+                    conv_kernel_size=self.conv_kernel_size,
+                )
+
+            hidden_states_B_C = causal_conv1d_fn(
+                hidden_states_B_C,
+                self.conv1d.weight.squeeze(1),
+                self.conv1d.bias,
+                activation=self.activation,
+                **kwargs,
+            )
+
+            if cache_params is not None:
+                hidden_states_B_C = hidden_states_B_C[:, :, -seq_len:]
+
+        # 3. SSM transformation
+        hidden_states_B_C = apply_mask_to_padding_states(hidden_states_B_C.transpose(1, 2), attention_mask)
         time_step, B, C = torch.split(
-            ssm_parameters, [self.time_step_rank, self.ssm_state_size, self.ssm_state_size], dim=-1
+            self.x_proj(hidden_states_B_C),
+            [self.time_step_rank, self.ssm_state_size, self.ssm_state_size],
+            dim=-1,
         )
 
+        # Key difference to (falcon) mamba: Additional norms on B, C, and dt
         time_step = self.dt_layernorm(time_step)
         B = self.b_layernorm(B)
         C = self.c_layernorm(C)
 
-        # Here we need to apply dt_proj without the bias, as the bias is added in the selective scan kernel.
-        # This is a hack to apply dt_proj while still using the forward pass of `torch.nn.Linear`, which is needed
-        # in order to make quantization work. Quantization code replaces `torch.nn.Linear` layers with quantized
-        # linear layers, and requires to call the forward pass directly.
-        # Quantized model can't work with the original code:
-        # ```discrete_time_step = self.dt_proj.weight @ time_step.transpose(1, 2)```
-        time_proj_bias = self.dt_proj.bias.data
-        with torch.no_grad():
-            self.dt_proj.bias.data = torch.zeros_like(self.dt_proj.bias.data)
-        discrete_time_step = self.dt_proj(time_step).transpose(1, 2)
-        with torch.no_grad():
-            self.dt_proj.bias.data = time_proj_bias
-
-        A = -torch.exp(self.A_log.float())
-        # 3.c perform the recurrence y ← SSM(A, B, C)(x)
-        time_proj_bias = time_proj_bias.float() if time_proj_bias is not None else None
-        if use_precomputed_states:
-            scan_outputs = selective_state_update(
-                cache_params.layers[self.layer_idx].recurrent_states[0],
-                hidden_states[..., 0],
-                discrete_time_step[..., 0],
+        # In case the model has been quantized, we need a hack to properly call the `nn.Linear` module
+        # at the price of a small overhead.
+        if hasattr(self.config, "_is_quantized"):
+            time_step = (self.dt_proj(time_step) - self.dt_proj.bias).transpose(1, 2)
+        else:
+            time_step = self.dt_proj.weight @ time_step.transpose(1, 2)
+        time_proj_bias = self.dt_proj.bias.float() if self.dt_proj.bias is not None else None
+
+        # Recurrent form
+        if use_precomputed_states and seq_len == 1:
+            scan_output = mamba_selective_state_update(
+                recurrent_state,
+                hidden_states_B_C.transpose(1, 2)[..., 0],
+                time_step[..., 0],
                 A,
                 B[:, 0],
                 C[:, 0],
                 self.D,
-                gate[..., 0],
-                time_proj_bias,
+                z=gate[..., 0],
+                dt_bias=time_proj_bias,
                 dt_softplus=True,
             ).unsqueeze(-1)
+
+        # Full sequence form
         else:
-            scan_outputs, ssm_state = selective_scan_fn(
-                hidden_states,
-                discrete_time_step,
+            output_final_state = cache_params is not None
+            scan_result = mamba_selective_scan(
+                hidden_states_B_C.transpose(1, 2),
+                time_step,
                 A,
                 B.transpose(1, 2),
                 C.transpose(1, 2),
-                self.D.float(),
-                gate,
-                time_proj_bias,
+                D=self.D.float(),
+                z=gate,
+                delta_bias=time_proj_bias,
                 delta_softplus=True,
-                return_last_state=True,
+                return_last_state=output_final_state,
+                # TODO: No faster alternatives for mamba atm (needs config adjustments)
+                use_mambapy=False,
+                use_associative_scan=False,
             )
-            if ssm_state is not None and cache_params is not None:
-                cache_params.update_recurrent_state(ssm_state, self.layer_idx)
 
-        # 4. Final linear projection
-        contextualized_states = self.out_proj(scan_outputs.transpose(1, 2))
-
-        return contextualized_states
-
-    def slow_forward(
-        self,
-        hidden_states: torch.Tensor,
-        cache_params: Cache | None = None,
-        attention_mask: torch.LongTensor | None = None,
-        **kwargs,
-    ):
-        batch_size, seq_len, _ = hidden_states.shape
-        dtype = hidden_states.dtype
-        use_precomputed_states = cache_params is not None and cache_params.has_previous_state(self.layer_idx)
-        # 1. Gated MLP's linear projection
-        projected_states = self.in_proj(hidden_states).transpose(1, 2)
-        hidden_states, gate = projected_states.chunk(2, dim=1)
-
-        # Apply the convolution
-        hidden_states = self._convolution(hidden_states, cache_params, attention_mask, **kwargs)
-
-        if attention_mask is not None:
-            hidden_states = hidden_states * attention_mask.unsqueeze(1)
-
-        if use_precomputed_states:
-            # In training mode, we don't want to perform in-place operations on ssm_state so we can compute the backwards pass
-            ssm_state = cache_params.layers[self.layer_idx].recurrent_states[0].clone()
-        else:
-            ssm_state = torch.zeros(
-                (batch_size, self.intermediate_size, self.ssm_state_size), device=hidden_states.device, dtype=dtype
-            )
-
-        # 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
-        )
-
-        time_step = self.dt_layernorm(time_step)
-        B = self.b_layernorm(B)
-        C = self.c_layernorm(C)
-
-        discrete_time_step = self.dt_proj(time_step)
-        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)
-        A = -torch.exp(self.A_log.float())
-        discrete_A = torch.exp(A[None, :, None, :] * discrete_time_step[:, :, :, None])
-        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(seq_len):
-            ssm_state = discrete_A[:, :, i, :] * ssm_state + deltaB_u[:, :, i, :]
-            scan_output = torch.matmul(ssm_state.to(dtype), C[:, i, :].unsqueeze(-1))
-            scan_outputs.append(scan_output[:, :, 0])
-        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)
-
-        if cache_params is not None:
-            cache_params.update_recurrent_state(ssm_state, self.layer_idx)
+            if output_final_state:
+                scan_output, final_state = scan_result
+                cache_params.update_recurrent_state(final_state, self.layer_idx)
+            else:
+                scan_output = scan_result
 
         # 4. Final linear projection
-        contextualized_states = self.out_proj(scan_output.transpose(1, 2))
+        contextualized_states = self.out_proj(scan_output.transpose(1, 2).to(dtype))
         return contextualized_states
 
 
diff --git a/src/transformers/models/mamba/modeling_mamba.py b/src/transformers/models/mamba/modeling_mamba.py
index 595545ab2d9a..25fb5383ea8b 100644
--- a/src/transformers/models/mamba/modeling_mamba.py
+++ b/src/transformers/models/mamba/modeling_mamba.py
@@ -118,6 +118,10 @@ def mamba_inner_fn(
     D: torch.Tensor | None = None,
     delta_bias: torch.Tensor | None = None,
     delta_softplus: bool = True,
+    b_rms_weight: torch.Tensor | None = None,
+    c_rms_weight: torch.Tensor | None = None,
+    dt_rms_weight: torch.Tensor | None = None,
+    b_c_dt_rms_eps: float = 1e-6,
     **kwargs,
 ):
     return None
diff --git a/tests/models/falcon_mamba/test_modeling_falcon_mamba.py b/tests/models/falcon_mamba/test_modeling_falcon_mamba.py
index 06d48a4086cb..de56f843fb7c 100644
--- a/tests/models/falcon_mamba/test_modeling_falcon_mamba.py
+++ b/tests/models/falcon_mamba/test_modeling_falcon_mamba.py
@@ -210,7 +210,11 @@ def create_and_check_falcon_mamba_cached_slow_forward_and_backwards(
         self, config, input_ids, *args, gradient_checkpointing=False
     ):
         model = FalconMambaModel(config)
-        model.to(torch_device)
+
+        # force torch path in any case
+        model.to("cpu")
+        input_ids = input_ids.to("cpu")
+
         if gradient_checkpointing:
             model.gradient_checkpointing_enable()
 
@@ -220,7 +224,7 @@ def create_and_check_falcon_mamba_cached_slow_forward_and_backwards(
 
         # use cache
         token_emb = model.embeddings(input_ids)
-        outputs = model.layers[0].mixer.slow_forward(token_emb, cache)
+        outputs = model.layers[0].mixer(token_emb, cache)
 
         loss = torch.log1p(torch.abs(outputs.sum()))
         self.parent.assertEqual(loss.shape, ())
@@ -513,6 +517,10 @@ def test_batched_generation(self):
                     "Hello today I will be talking about the “Theory of Relativity” by Albert Einstein.\nThe",
                     "Hello my name is Younes and today I will be talking about the importance of the internet in our lives.\nThe internet is a global",
                 ],
+                ("cuda", 9): [
+                    'Hello today I am going to talk about the “Theory of Relativity” by Albert Einstein.\n',
+                    'Hello my name is Younes and today I will be talking about the importance of the internet in our lives.\nThe internet is a global'
+                ]
             }
         )
         EXPECTED_OUTPUT = EXPECTED_OUTPUTS.get_expectation()
@@ -545,6 +553,10 @@ def test_batched_generation(self):
                 ("cuda", (8, 6)): [
                     ' I will be talking about the “Theory of Relativity” by Albert Einstein.\nThe',
                     ' I will be talking about the importance of the internet in our lives.\nThe internet is a global'
+                ],
+                ("cuda", 9): [
+                    ' I am going to talk about the “Theory of Relativity” by Albert Einstein.\n',
+                    ' I will be talking about the importance of the internet in our lives.\nThe internet is a global'
                 ]
             }
         )  # fmt: skip
diff --git a/tests/models/jamba/test_modeling_jamba.py b/tests/models/jamba/test_modeling_jamba.py
index ecbfa916899e..04c6ac8ec49f 100644
--- a/tests/models/jamba/test_modeling_jamba.py
+++ b/tests/models/jamba/test_modeling_jamba.py
@@ -567,6 +567,7 @@ def test_simple_batched_generate_with_padding(self):
             {
                 ("cuda", 7): ["<|startoftext|>Hey how are you doing on this lovely evening? Canyon rins hugaughter glamour Rutgers Singh Hebrew cases Cats", "<|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|startoftext|>Tell me a storyptus Nets Madison El chamadamodern updximVaparsed",],
                 ("cuda", 8): ["<|startoftext|>Hey how are you doing on this lovely evening? I'm so glad you're here.", "<|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|startoftext|>Tell me a story about a woman who was born in the United States",],
+                ("cuda", 9): ["<|startoftext|>Hey how are you doing on this lovely evening? I'm so glad you're here.", "<|startoftext|>Tell me a story<|pad|><|pad|><|pad|><|pad|><|pad|><|pad|>, I'm not sure, but I'",],
                 ("rocm", 9): ["<|startoftext|>Hey how are you doing on this lovely evening? Canyon rins hugaughter glamour Rutgers Singh<|reserved_797|>cw algunas", "<|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|startoftext|>Tell me a storyptus Nets Madison El chamadamodern updximVaparsed",],
                 ("xpu", 3): ["<|startoftext|>Hey how are you doing on this lovely evening? I'm so glad you're here.", "<|startoftext|>Tell me a story<|pad|><|pad|><|pad|><|pad|><|pad|><|pad|>, I'm not sure, but I'"]
             }
diff --git a/utils/modular_model_converter.py b/utils/modular_model_converter.py
index a30cec1e37f4..9d22a7696d8f 100644
--- a/utils/modular_model_converter.py
+++ b/utils/modular_model_converter.py
@@ -83,6 +83,7 @@ def get_module_source_from_name(module_name: str) -> str:
     "mamba_inner_fn",
     "mamba_selective_state_update",
     "mamba_selective_scan",
+    "is_mambapy_available",
     "mamba2_split_conv1d_scan_combined",
     "mamba2_selective_state_update",
     "mamba2_chunk_scan",

From fcf7292f176f3c2957003eecb6d7e5b506064713 Mon Sep 17 00:00:00 2001
From: vasqu <antonprogamer@gmail.com>
Date: Mon, 3 Aug 2026 21:21:16 +0000
Subject: [PATCH 30/43] style

---
 tests/models/falcon_mamba/test_modeling_falcon_mamba.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/tests/models/falcon_mamba/test_modeling_falcon_mamba.py b/tests/models/falcon_mamba/test_modeling_falcon_mamba.py
index de56f843fb7c..2ed54e6b6e79 100644
--- a/tests/models/falcon_mamba/test_modeling_falcon_mamba.py
+++ b/tests/models/falcon_mamba/test_modeling_falcon_mamba.py
@@ -518,9 +518,9 @@ def test_batched_generation(self):
                     "Hello my name is Younes and today I will be talking about the importance of the internet in our lives.\nThe internet is a global",
                 ],
                 ("cuda", 9): [
-                    'Hello today I am going to talk about the “Theory of Relativity” by Albert Einstein.\n',
-                    'Hello my name is Younes and today I will be talking about the importance of the internet in our lives.\nThe internet is a global'
-                ]
+                    "Hello today I am going to talk about the “Theory of Relativity” by Albert Einstein.\n",
+                    "Hello my name is Younes and today I will be talking about the importance of the internet in our lives.\nThe internet is a global",
+                ],
             }
         )
         EXPECTED_OUTPUT = EXPECTED_OUTPUTS.get_expectation()

From 3823bcbd90eba82d3fb5bf26deccfbf616d82392 Mon Sep 17 00:00:00 2001
From: vasqu <antonprogamer@gmail.com>
Date: Mon, 3 Aug 2026 21:29:49 +0000
Subject: [PATCH 31/43] fix mambapy + enable on jamba

---
 .../falcon_mamba/configuration_falcon_mamba.py  | 11 +++++++----
 .../falcon_mamba/modeling_falcon_mamba.py       |  9 ++++-----
 .../models/falcon_mamba/modular_falcon_mamba.py | 17 +++++++++++------
 .../models/jamba/configuration_jamba.py         | 11 +++++++++++
 src/transformers/models/jamba/modeling_jamba.py | 12 +++++++-----
 src/transformers/models/jamba/modular_jamba.py  |  8 +++++---
 utils/modular_model_converter.py                |  1 +
 7 files changed, 46 insertions(+), 23 deletions(-)

diff --git a/src/transformers/models/falcon_mamba/configuration_falcon_mamba.py b/src/transformers/models/falcon_mamba/configuration_falcon_mamba.py
index 49e25e4d3bfd..d3caacc0a0df 100644
--- a/src/transformers/models/falcon_mamba/configuration_falcon_mamba.py
+++ b/src/transformers/models/falcon_mamba/configuration_falcon_mamba.py
@@ -41,9 +41,10 @@ class FalconMambaConfig(PreTrainedConfig):
         Whether or not residuals should be in `float32`. If set to `False` residuals will keep the same `dtype` as the rest of the model
     rescale_prenorm_residual (`bool`, *optional*, defaults to `False`):
         Whether or not to rescale `out_proj` weights when initializing.
-    use_falcon_mambapy (`bool`, *optional*, defaults to `False`):
-        This argument corresponds to `use_mambapy` in MambaConfig.
-        Determines the fallback strategy during training if the CUDA-based official implementation of Mamba is not available. If `True`, the mamba.py implementation is used. If `False`, the naive and slower implementation is used. Consider switching to the naive version if memory is limited.
+    use_mambapy (`bool`, *optional*, defaults to `False`):
+        Determines the fallback strategy during training if the CUDA-based official implementation of Mamba is not available. If `True`,
+        the mamba.py implementation is used. If `False`, the naive and slower implementation is used. Consider switching to the naive
+        version if memory is limited.
     use_associative_scan (`bool`, *optional*, defaults to `True`):
         Whether to use PyTorch's `torch._higher_order_ops.associative_scan` for the parallel scan instead of the naive
         sequential implementation. The associative scan is only active during `torch.compile` tracing and
@@ -93,12 +94,14 @@ class FalconMambaConfig(PreTrainedConfig):
     rescale_prenorm_residual: bool = False
     use_cache: bool = True
 
-    use_falcon_mambapy: bool = False
+    use_mambapy: bool = False
     use_associative_scan: bool = True
     tie_word_embeddings: bool = True
     mixer_rms_eps: float = 1e-6
 
     def __post_init__(self, **kwargs):
+        # BC for rename
+        self.use_mambapy = kwargs.pop("use_falcon_mambapy", False)
         self.intermediate_size = int(self.expand * self.hidden_size)
         self.time_step_rank = (
             math.ceil(self.hidden_size / 16) if self.time_step_rank == "auto" else self.time_step_rank
diff --git a/src/transformers/models/falcon_mamba/modeling_falcon_mamba.py b/src/transformers/models/falcon_mamba/modeling_falcon_mamba.py
index ac833f31db9c..d11425b2f566 100644
--- a/src/transformers/models/falcon_mamba/modeling_falcon_mamba.py
+++ b/src/transformers/models/falcon_mamba/modeling_falcon_mamba.py
@@ -202,7 +202,7 @@ def mamba_selective_scan(
     delta_bias: torch.Tensor | None = None,
     delta_softplus: bool = False,
     return_last_state: bool = False,
-    use_falcon_mambapy: bool = False,
+    use_mambapy: bool = False,
     use_associative_scan: bool = False,
     **kwargs,
 ):
@@ -234,7 +234,7 @@ def mamba_selective_scan(
     discrete_B = dt[:, :, :, None] * B[:, None, :, :].float()
     deltaB_u = discrete_B * hidden_states[:, :, :, None].float()
 
-    if use_falcon_mambapy and pscan is not None:
+    if use_mambapy and pscan is not None:
         all_states = pscan(discrete_A.transpose(1, 2), deltaB_u.transpose(1, 2))
 
         scan_output = (all_states @ C.unsqueeze(-1)).squeeze(3).transpose(1, 2)
@@ -326,7 +326,7 @@ def __init__(self, config: FalconMambaConfig, layer_idx: int, initialize_mixer_w
         self.activation = config.hidden_act
         self.act = ACT2FN[config.hidden_act]
 
-        self.use_falcon_mambapy = config.use_falcon_mambapy
+        self.use_mambapy = config.use_mambapy
         self.use_associative_scan = config.use_associative_scan
 
         # projection of the input hidden states
@@ -503,8 +503,7 @@ def forward(
                 delta_bias=time_proj_bias,
                 delta_softplus=True,
                 return_last_state=output_final_state,
-                # TODO: rename to normal mambapy
-                use_mambapy=self.use_falcon_mambapy,
+                use_mambapy=self.use_mambapy,
                 use_associative_scan=self.use_associative_scan,
             )
 
diff --git a/src/transformers/models/falcon_mamba/modular_falcon_mamba.py b/src/transformers/models/falcon_mamba/modular_falcon_mamba.py
index f65315338b85..006b3a6b8fd5 100644
--- a/src/transformers/models/falcon_mamba/modular_falcon_mamba.py
+++ b/src/transformers/models/falcon_mamba/modular_falcon_mamba.py
@@ -59,9 +59,10 @@ class FalconMambaConfig(MambaConfig):
         Whether or not residuals should be in `float32`. If set to `False` residuals will keep the same `dtype` as the rest of the model
     rescale_prenorm_residual (`bool`, *optional*, defaults to `False`):
         Whether or not to rescale `out_proj` weights when initializing.
-    use_falcon_mambapy (`bool`, *optional*, defaults to `False`):
-        This argument corresponds to `use_mambapy` in MambaConfig.
-        Determines the fallback strategy during training if the CUDA-based official implementation of Mamba is not available. If `True`, the mamba.py implementation is used. If `False`, the naive and slower implementation is used. Consider switching to the naive version if memory is limited.
+    use_mambapy (`bool`, *optional*, defaults to `False`):
+        Determines the fallback strategy during training if the CUDA-based official implementation of Mamba is not available. If `True`,
+        the mamba.py implementation is used. If `False`, the naive and slower implementation is used. Consider switching to the naive
+        version if memory is limited.
     use_associative_scan (`bool`, *optional*, defaults to `True`):
         Whether to use PyTorch's `torch._higher_order_ops.associative_scan` for the parallel scan instead of the naive
         sequential implementation. The associative scan is only active during `torch.compile` tracing and
@@ -85,10 +86,15 @@ class FalconMambaConfig(MambaConfig):
     >>> configuration = model.config
     ```"""
 
-    use_falcon_mambapy: bool = False
+    use_mambapy: bool = False
     use_associative_scan: bool = True
     mixer_rms_eps: float = 1e-6
 
+    def __post_init__(self, **kwargs):
+        # BC for rename
+        self.use_mambapy = kwargs.pop("use_falcon_mambapy", False)
+        super().__post_init__(self, **kwargs)
+
     @property
     def layer_types(self):
         return ["linear_attention"] * self.num_hidden_layers
@@ -242,8 +248,7 @@ def forward(
                 delta_bias=time_proj_bias,
                 delta_softplus=True,
                 return_last_state=output_final_state,
-                # TODO: rename to normal mambapy
-                use_mambapy=self.use_falcon_mambapy,
+                use_mambapy=self.use_mambapy,
                 use_associative_scan=self.use_associative_scan,
             )
 
diff --git a/src/transformers/models/jamba/configuration_jamba.py b/src/transformers/models/jamba/configuration_jamba.py
index cd9e567647c3..cbb14c392085 100644
--- a/src/transformers/models/jamba/configuration_jamba.py
+++ b/src/transformers/models/jamba/configuration_jamba.py
@@ -39,6 +39,15 @@ class JambaConfig(PreTrainedConfig):
         `True` and kernels are not available
     mamba_dt_rank (`Union[int,str]`, *optional*, defaults to `"auto"`):
         Rank of the mamba discretization projection matrix. `"auto"` means that it will default to `math.ceil(self.hidden_size / 16)`
+    use_mambapy (`bool`, *optional*, defaults to `False`):
+        Determines the fallback strategy during training if the CUDA-based official implementation of Mamba is not available. If `True`,
+        the mamba.py implementation is used. If `False`, the naive and slower implementation is used. Consider switching to the naive
+        version if memory is limited.
+    use_associative_scan (`bool`, *optional*, defaults to `True`):
+        Whether to use PyTorch's `torch._higher_order_ops.associative_scan` for the parallel scan instead of the naive
+        sequential implementation. The associative scan is only active during `torch.compile` tracing and
+        requires torch >= 2.9.0. Both paths are tested to produce numerically identical results (see
+        `test_associative_scan_matches_sequential`). Set to `False` to fall back to the sequential loop.
     """
 
     model_type = "jamba"
@@ -78,6 +87,8 @@ class JambaConfig(PreTrainedConfig):
     mamba_dt_rank: int | str = "auto"
     mamba_conv_bias: bool = True
     mamba_proj_bias: bool = False
+    use_mambapy: bool = False
+    use_associative_scan: bool = True
 
     def __post_init__(self, **kwargs):
         if self.num_key_value_heads is None:
diff --git a/src/transformers/models/jamba/modeling_jamba.py b/src/transformers/models/jamba/modeling_jamba.py
index 04b33eb42805..44ba6fbf9bda 100755
--- a/src/transformers/models/jamba/modeling_jamba.py
+++ b/src/transformers/models/jamba/modeling_jamba.py
@@ -343,7 +343,7 @@ def mamba_selective_scan(
     delta_bias: torch.Tensor | None = None,
     delta_softplus: bool = False,
     return_last_state: bool = False,
-    use_jambapy: bool = False,
+    use_mambapy: bool = False,
     use_associative_scan: bool = False,
     **kwargs,
 ):
@@ -375,7 +375,7 @@ def mamba_selective_scan(
     discrete_B = dt[:, :, :, None] * B[:, None, :, :].float()
     deltaB_u = discrete_B * hidden_states[:, :, :, None].float()
 
-    if use_jambapy and pscan is not None:
+    if use_mambapy and pscan is not None:
         all_states = pscan(discrete_A.transpose(1, 2), deltaB_u.transpose(1, 2))
 
         scan_output = (all_states @ C.unsqueeze(-1)).squeeze(3).transpose(1, 2)
@@ -488,6 +488,9 @@ def __init__(self, config: JambaConfig, layer_idx):
         self.b_layernorm = JambaRMSNorm(self.ssm_state_size, eps=config.rms_norm_eps)
         self.c_layernorm = JambaRMSNorm(self.ssm_state_size, eps=config.rms_norm_eps)
 
+        self.use_mambapy = config.use_mambapy
+        self.use_associative_scan = config.use_associative_scan
+
     @force_accelerate_hooks(["conv1d", "dt_proj"])
     def forward(
         self,
@@ -592,9 +595,8 @@ def forward(
                 delta_bias=time_proj_bias,
                 delta_softplus=True,
                 return_last_state=output_final_state,
-                # TODO: No faster alternatives for mamba atm (needs config adjustments)
-                use_mambapy=False,
-                use_associative_scan=False,
+                use_mambapy=self.use_mambapy,
+                use_associative_scan=self.use_associative_scan,
             )
 
             if output_final_state:
diff --git a/src/transformers/models/jamba/modular_jamba.py b/src/transformers/models/jamba/modular_jamba.py
index 03e8e769986e..5d7390729555 100644
--- a/src/transformers/models/jamba/modular_jamba.py
+++ b/src/transformers/models/jamba/modular_jamba.py
@@ -149,6 +149,9 @@ def __init__(self, config: JambaConfig, layer_idx):
         self.b_layernorm = JambaRMSNorm(self.ssm_state_size, eps=config.rms_norm_eps)
         self.c_layernorm = JambaRMSNorm(self.ssm_state_size, eps=config.rms_norm_eps)
 
+        self.use_mambapy = config.use_mambapy
+        self.use_associative_scan = config.use_associative_scan
+
     def init_jamba_weights(self):
         raise NotImplementedError("Not needed for jamba")
 
@@ -255,9 +258,8 @@ def forward(
                 delta_bias=time_proj_bias,
                 delta_softplus=True,
                 return_last_state=output_final_state,
-                # TODO: No faster alternatives for mamba atm (needs config adjustments)
-                use_mambapy=False,
-                use_associative_scan=False,
+                use_mambapy=self.use_mambapy,
+                use_associative_scan=self.use_associative_scan,
             )
 
             if output_final_state:
diff --git a/utils/modular_model_converter.py b/utils/modular_model_converter.py
index 9d22a7696d8f..7ec9b85edc34 100644
--- a/utils/modular_model_converter.py
+++ b/utils/modular_model_converter.py
@@ -80,6 +80,7 @@ def get_module_source_from_name(module_name: str) -> str:
 NAMES_TO_NEVER_REPLACE = (
     "mamba_ssm",
     "mamba-ssm",
+    "use_mambapy",
     "mamba_inner_fn",
     "mamba_selective_state_update",
     "mamba_selective_scan",

From 380e5fc7dda6ff85201a5011bffc43d3567e9de0 Mon Sep 17 00:00:00 2001
From: vasqu <antonprogamer@gmail.com>
Date: Mon, 3 Aug 2026 22:04:56 +0000
Subject: [PATCH 32/43] zamba1

---
 .../models/zamba/configuration_zamba.py       |  32 +-
 .../models/zamba/modeling_zamba.py            | 459 +++++++++++-------
 2 files changed, 293 insertions(+), 198 deletions(-)

diff --git a/src/transformers/models/zamba/configuration_zamba.py b/src/transformers/models/zamba/configuration_zamba.py
index ce57dedc79bc..efaf9fd6bc68 100644
--- a/src/transformers/models/zamba/configuration_zamba.py
+++ b/src/transformers/models/zamba/configuration_zamba.py
@@ -17,7 +17,7 @@
 
 from huggingface_hub.dataclasses import strict
 
-from ...configuration_utils import PreTrainedConfig
+from ...configuration_utils import PreTrainedConfig, remap_legacy_layer_types
 from ...utils import auto_docstring
 
 
@@ -49,6 +49,8 @@ class ZambaConfig(PreTrainedConfig):
         `True` and kernels are not available
     mamba_dt_rank (`Union[int,str]`, *optional*, defaults to `"auto"`):
         Rank of the mamba discretization projection matrix. `"auto"` means that it will default to `math.ceil(self.hidden_size / 16)`
+    layers_block_type (`str`, *optional*):
+        Alias for `layer_types` which is kept for BC purposes.
     """
 
     model_type = "zamba"
@@ -88,14 +90,25 @@ class ZambaConfig(PreTrainedConfig):
     time_step_floor: float = 1e-4
     mamba_conv_bias: bool = True
     mamba_proj_bias: bool = False
+    layers_block_type: list[str] | None = None
 
     def __post_init__(self, **kwargs):
         self.attention_hidden_size = self.attention_hidden_size or 2 * self.hidden_size
         self.attention_head_dim = self.attention_head_dim or 2 * self.hidden_size // self.num_attention_heads
         self.mamba_dt_rank = math.ceil(self.hidden_size / 16) if self.mamba_dt_rank == "auto" else self.mamba_dt_rank
-        self.layers_block_type = self._layers_block_type(
-            self.num_hidden_layers, self.attn_layer_period, self.attn_layer_offset
-        )
+
+        if self.layers_block_type is None:
+            self.layers_block_type = [
+                "linear_attention",
+                "linear_attention",
+                "hybrid",
+            ] + [
+                "hybrid" if i % self.attn_layer_period == self.attn_layer_offset else "linear_attention"
+                for i in range(self.num_hidden_layers - 3)
+            ]
+        else:
+            self.layers_block_type = remap_legacy_layer_types(self.layers_block_type)
+
         super().__post_init__(**kwargs)
 
     def validate_architecture(self):
@@ -103,16 +116,5 @@ def validate_architecture(self):
         if (self.mamba_expand * self.hidden_size) % self.n_mamba_heads != 0:
             raise ValueError("`intermediate_size` should be divisible by `n_mamba_heads`.")
 
-    def _layers_block_type(self, num_hidden_layers, attn_layer_period, attn_layer_offset):
-        layers = [
-            "linear_attention",
-            "linear_attention",
-            "hybrid",
-        ] + [
-            "hybrid" if i % attn_layer_period == attn_layer_offset else "linear_attention"
-            for i in range(num_hidden_layers - 3)
-        ]
-        return layers
-
 
 __all__ = ["ZambaConfig"]
diff --git a/src/transformers/models/zamba/modeling_zamba.py b/src/transformers/models/zamba/modeling_zamba.py
index fa1ebd8ba4ef..ef12466937f7 100644
--- a/src/transformers/models/zamba/modeling_zamba.py
+++ b/src/transformers/models/zamba/modeling_zamba.py
@@ -30,7 +30,8 @@
 from ...activations import ACT2FN
 from ...cache_utils import Cache, DynamicCache
 from ...generation import GenerationMixin
-from ...integrations.hub_kernels import lazy_load_kernel
+from ...integrations import use_kernel_func_from_hub_with_fallback, use_kernelized_func
+from ...integrations.accelerate import force_accelerate_hooks
 from ...masking_utils import create_causal_mask, create_recurrent_attention_mask
 from ...modeling_layers import GradientCheckpointingLayer
 from ...modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast, SequenceClassifierOutputWithPast
@@ -38,7 +39,11 @@
 from ...processing_utils import Unpack
 from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, logging
 from ...utils.generic import merge_with_config_defaults
-from ...utils.import_utils import resolve_internal_import
+from ...utils.import_utils import (
+    is_mambapy_available,
+    is_torch_greater_or_equal,
+    is_tracing,
+)
 from ...utils.output_capturing import capture_outputs
 from .configuration_zamba import ZambaConfig
 
@@ -175,6 +180,19 @@ def forward(
         return attn_output, attn_weights
 
 
+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
+    """
+    # NOTE: attention mask is a 2D boolean tensor
+    if attention_mask is not None:
+        dtype = hidden_states.dtype
+        hidden_states = (hidden_states * attention_mask[:, :, None]).to(dtype)
+
+    return hidden_states
+
+
+@use_kernel_func_from_hub_with_fallback("causal_conv1d_update", "causal_conv1d")
 def causal_conv1d_update(
     hidden_states: torch.Tensor,
     conv_state: torch.Tensor,
@@ -194,6 +212,7 @@ def causal_conv1d_update(
     return out.to(hidden_states.dtype)
 
 
+@use_kernel_func_from_hub_with_fallback("causal_conv1d_fn", "causal_conv1d")
 def causal_conv1d_fn(
     hidden_states: torch.Tensor,
     weight: nn.Parameter,
@@ -216,6 +235,164 @@ def causal_conv1d_fn(
     return out.to(hidden_states.dtype)
 
 
+@use_kernel_func_from_hub_with_fallback(
+    "selective_state_update",
+    "mamba_ssm",
+)
+def mamba_selective_state_update(
+    state: torch.Tensor,
+    hidden_states: torch.Tensor,
+    dt: torch.Tensor,
+    A: torch.Tensor,
+    B: torch.Tensor,
+    C: torch.Tensor,
+    D: torch.Tensor | None = None,
+    dt_bias: torch.Tensor | None = None,
+    dt_softplus: bool = False,
+    z: torch.Tensor | None = None,
+    **kwargs,
+):
+    input_dtype = hidden_states.dtype
+
+    if dt_bias is not None:
+        dt = dt + dt_bias.to(dt.dtype)
+    if dt_softplus:
+        dt = F.softplus(dt)
+
+    # Discretize A
+    dA = torch.exp(dt.float()[..., None] * A.float()).to(device=state.device)
+
+    # Discretize B
+    dB = dt.float()[..., None] * B.float()[:, None, :]
+    # Discretize x into dB
+    dBx = dB * hidden_states.float()[..., None]
+
+    # State calculation
+    ssm_state = state.float() * dA + dBx
+    state.copy_(ssm_state.to(state.dtype))
+
+    # Subsequent output
+    out = torch.matmul(ssm_state.to(C.dtype), C.unsqueeze(-1)).squeeze(-1)
+
+    # D skip connection
+    if D is not None:
+        out = out + hidden_states * D
+
+    if z is not None:
+        out = out * F.silu(z)
+
+    return out.to(input_dtype)
+
+
+@use_kernel_func_from_hub_with_fallback(
+    "selective_scan_fn",
+    "mamba_ssm",
+)
+def mamba_selective_scan(
+    hidden_states: torch.Tensor,
+    dt: torch.Tensor,
+    A: torch.Tensor,
+    B: torch.Tensor,
+    C: torch.Tensor,
+    D: torch.Tensor | None = None,
+    z: torch.Tensor | None = None,
+    delta_bias: torch.Tensor | None = None,
+    delta_softplus: bool = False,
+    return_last_state: bool = False,
+    use_mambapy: bool = False,
+    use_associative_scan: bool = False,
+    **kwargs,
+):
+    # Torch only alternatives to the recurrent path
+    if is_torch_greater_or_equal("2.9.0"):
+        from torch._higher_order_ops.associative_scan import associative_scan
+    else:
+        associative_scan = None
+
+    if is_mambapy_available():
+        from mambapy.pscan import pscan
+    else:
+        pscan = None
+
+    batch_size, intermediate_size, seq_len = hidden_states.shape
+    input_dtype = hidden_states.dtype
+
+    if delta_bias is not None:
+        dt = dt + delta_bias.to(dt.dtype)[..., None]
+    if delta_softplus:
+        dt = F.softplus(dt)
+
+    # We need to transpose on the basis of the original kernel layout
+    B = B.transpose(1, 2)
+    C = C.transpose(1, 2)
+
+    # Discretize A and B for the entire sequence
+    discrete_A = torch.exp(A[None, :, None, :] * dt[:, :, :, None])
+    discrete_B = dt[:, :, :, None] * B[:, None, :, :].float()
+    deltaB_u = discrete_B * hidden_states[:, :, :, None].float()
+
+    if use_mambapy and pscan is not None:
+        all_states = pscan(discrete_A.transpose(1, 2), deltaB_u.transpose(1, 2))
+
+        scan_output = (all_states @ C.unsqueeze(-1)).squeeze(3).transpose(1, 2)
+        ssm_state = all_states[:, -1]
+
+    elif use_associative_scan and associative_scan is not None and is_tracing(hidden_states):
+
+        def combine_fn(left, right):
+            a_left, b_left = left
+            a_right, b_right = right
+            return a_left * a_right, a_right * b_left + b_right
+
+        combine_mode = "pointwise" if discrete_A.device.type in ("cuda", "xpu") else "generic"
+        _, all_states = associative_scan(
+            combine_fn,
+            (discrete_A, deltaB_u),
+            dim=2,
+            combine_mode=combine_mode,
+        )
+
+        scan_output = torch.matmul(all_states.transpose(1, 2).to(input_dtype), C.unsqueeze(-1))
+        scan_output = scan_output.squeeze(-1).transpose(1, 2)
+        ssm_state = all_states[:, :, -1]
+
+    # Recurrent iteration
+    else:
+        # "Initial hidden state" is not supported by the kernel path, so use
+        # the same zero initialization as the kernel
+        ssm_state = torch.zeros(
+            batch_size,
+            intermediate_size,
+            A.shape[-1],
+            dtype=input_dtype,
+            device=hidden_states.device,
+        )
+
+        scan_outputs = []
+        for index in range(seq_len):
+            # State calculation
+            ssm_state = discrete_A[:, :, index] * ssm_state + deltaB_u[:, :, index]
+
+            # Subsequent output
+            scan_output = torch.matmul(ssm_state.to(input_dtype), C[:, index, :].unsqueeze(-1))
+            scan_outputs.append(scan_output[:, :, 0])
+        scan_output = torch.stack(scan_outputs, dim=-1)
+
+    if D is not None:
+        scan_output = scan_output + hidden_states * D[None, :, None]
+
+    if z is not None:
+        scan_output = scan_output * F.silu(z)
+
+    if return_last_state:
+        return scan_output, ssm_state
+
+    return scan_output
+
+
+@use_kernelized_func(
+    [mamba_selective_scan, mamba_selective_state_update, causal_conv1d_fn, causal_conv1d_update]
+)
 class ZambaMambaMixer(nn.Module):
     """
     Compute ∆, A, B, C, and D the state space parameters and compute the `contextualized_states`.
@@ -283,232 +460,148 @@ def __init__(self, config: ZambaConfig, layer_idx):
         self.D = nn.Parameter(torch.ones(self.n_mamba_heads, self.mamba_head_dim))
         self.out_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=self.use_bias)
 
-        global causal_conv1d, causal_conv1d_update, causal_conv1d_fn
-        causal_conv1d = lazy_load_kernel("causal-conv1d")
-        causal_conv1d_update = getattr(causal_conv1d, "causal_conv1d_update", causal_conv1d_update)
-        causal_conv1d_fn = getattr(causal_conv1d, "causal_conv1d_fn", causal_conv1d_fn)
-
-        global mamba_ssm, selective_state_update, selective_scan_fn
-        mamba_ssm = lazy_load_kernel("mamba-ssm")
-        selective_state_update = resolve_internal_import(
-            mamba_ssm, chained_path="ops.triton.selective_state_update.selective_state_update"
-        )
-        selective_scan_fn = getattr(mamba_ssm, "selective_scan_fn", None)
-
-        global is_fast_path_available
-        is_fast_path_available = (
-            all((selective_state_update, selective_scan_fn))
-            and hasattr(causal_conv1d, "causal_conv1d_update")
-            and hasattr(causal_conv1d, "causal_conv1d_fn")
-        )
-
-        if not is_fast_path_available:
-            logger.warning_once(
-                "The fast path is not available because one of `(selective_state_update, selective_scan_fn, causal_conv1d_fn, causal_conv1d_update, mamba_inner_fn)`"
-                " is None. To install follow https://github.com/state-spaces/mamba/#installation and"
-                " https://github.com/Dao-AILab/causal-conv1d. If you want to use the naive implementation, set `use_mamba_kernels=False` in the model config"
-            )
-
         self.layer_type = config.layer_types[layer_idx]
 
-    def _convolution(
+    @force_accelerate_hooks("conv1d")
+    def forward(
         self,
         hidden_states: torch.Tensor,
         cache_params: Cache | None = None,
         attention_mask: torch.LongTensor | None = None,
         **kwargs,
     ):
-        seq_len = hidden_states.shape[-1]
-        if attention_mask is not None:
-            hidden_states = hidden_states * attention_mask.unsqueeze(1)
-
+        batch_size, seq_len, _ = hidden_states.shape
+        dtype = hidden_states.dtype
         use_precomputed_states = cache_params is not None and cache_params.has_previous_state(self.layer_idx)
 
-        if use_precomputed_states and seq_len == 1 and not cache_params.layers[self.layer_idx].record_past:
+        # 1. Gated MLP's linear projection
+        hidden_states = apply_mask_to_padding_states(hidden_states, attention_mask)
+        projected_states = self.in_proj(hidden_states).transpose(1, 2)
+
+        # Key difference: Split the projection into independent Mamba heads
+        hidden_states_B_C, gate = projected_states.view(batch_size, -1, 2, seq_len).chunk(2, dim=2)
+        hidden_states_B_C = hidden_states_B_C.squeeze(2).contiguous()
+        gate = gate.reshape(batch_size, self.n_mamba_heads, self.mamba_head_dim, seq_len).transpose(0, 1)
+
+        if use_precomputed_states:
             conv_state = cache_params.layers[self.layer_idx].conv_states[0]
-            hidden_states = causal_conv1d_update(
-                hidden_states,
+            recurrent_state = cache_params.layers[self.layer_idx].recurrent_states[0]
+
+        # 2. Convolution sequence transformation
+        if use_precomputed_states and seq_len == 1 and not cache_params.layers[self.layer_idx].record_past:
+            hidden_states_B_C = causal_conv1d_update(
+                hidden_states_B_C,
                 conv_state,
                 self.conv1d.weight.squeeze(1),
                 self.conv1d.bias,
-                self.activation,
+                activation=self.activation,
             )
         else:
             if cache_params is not None:
-                hidden_states = cache_params.update_conv_state(
-                    hidden_states, self.layer_idx, conv_kernel_size=self.conv_kernel_size
+                hidden_states_B_C = cache_params.update_conv_state(
+                    hidden_states_B_C,
+                    self.layer_idx,
+                    conv_kernel_size=self.conv_kernel_size,
                 )
 
-            hidden_states = causal_conv1d_fn(
-                hidden_states,
+            hidden_states_B_C = causal_conv1d_fn(
+                hidden_states_B_C,
                 self.conv1d.weight.squeeze(1),
                 self.conv1d.bias,
                 activation=self.activation,
                 seq_idx=kwargs.get("seq_idx"),
             )
 
-            # Drop the additional previous states
             if cache_params is not None:
-                hidden_states = hidden_states[:, :, -seq_len:]
-
-        return hidden_states
-
-    def cuda_kernels_forward(
-        self, hidden_states: torch.Tensor, cache_params: Cache | None = None, attention_mask=None, **kwargs
-    ):
-        batch_size, seq_len, _ = hidden_states.shape
-        use_precomputed_states = cache_params is not None and cache_params.has_previous_state(self.layer_idx)
-
-        # 1. Gated linear projection
-        projected_states = self.in_proj(hidden_states).transpose(1, 2)
-
-        hidden_states, gate = projected_states.view(batch_size, -1, 2, seq_len).chunk(2, dim=2)
-        hidden_states = hidden_states.squeeze(2).contiguous()
-        gate = gate.reshape(batch_size, self.n_mamba_heads, -1, seq_len).transpose(0, 1)
-
-        # Apply the conv
-        hidden_states = self._convolution(hidden_states, cache_params, attention_mask, **kwargs)
-
-        if attention_mask is not None:
-            hidden_states = hidden_states * attention_mask.unsqueeze(1)
-
-        # 3. SSM sequence transformation
-        # 3.a. input varying initialization of time_step, B and C
-
-        hidden_states = hidden_states.reshape(-1, self.n_mamba_heads, self.mamba_head_dim, seq_len).transpose(0, 1)
-        ssm_parameters = (self.x_proj_weight[:, None, :, :] @ hidden_states).transpose(-1, -2)
+                hidden_states_B_C = hidden_states_B_C[:, :, -seq_len:]
 
+        # 3. SSM transformation
+        hidden_states_B_C = apply_mask_to_padding_states(
+            hidden_states_B_C.transpose(1, 2),
+            attention_mask,
+        )
+        hidden_states_B_C = hidden_states_B_C.transpose(1, 2)
+        hidden_states_B_C = hidden_states_B_C.reshape(
+            batch_size,
+            self.n_mamba_heads,
+            self.mamba_head_dim,
+            seq_len,
+        ).transpose(0, 1)
+
+        # Key difference: x_proj and dt_proj have separate weights for each Mamba head
+        ssm_parameters = (self.x_proj_weight[:, None] @ hidden_states_B_C).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
+            ssm_parameters,
+            [self.time_step_rank, self.ssm_state_size, self.ssm_state_size],
+            dim=-1,
         )
 
-        discrete_time_step = self.dt_proj_weight[:, None] @ time_step.transpose(-1, -2)
-
-        A = -torch.exp(self.A_log.float())
-
-        # 3.c perform the recurrence y ← SSM(A, B, C)(x)
+        time_step = self.dt_proj_weight[:, None] @ time_step.transpose(-1, -2)
         time_proj_bias = self.dt_proj_bias.float() if self.dt_proj_bias is not None else None
-        scan_outputs = torch.empty((batch_size, 0, seq_len), device=hidden_states.device, dtype=hidden_states.dtype)
+        A = -torch.exp(self.A_log.float())
 
-        if use_precomputed_states:
-            for n in range(self.n_mamba_heads):
-                scan_outputs_ = selective_state_update(
-                    cache_params.layers[self.layer_idx].recurrent_states[0][:, n],
-                    hidden_states[n, ..., 0],
-                    discrete_time_step[n, ..., 0],
-                    A[n],
-                    B[n, :, 0],
-                    C[n, :, 0],
-                    self.D[n],
-                    gate[n, ..., 0],
-                    time_proj_bias[n],
+        # Key difference: per head scans
+        # Recurrent form
+        if use_precomputed_states and seq_len == 1:
+            scan_outputs = []
+
+            for head_idx in range(self.n_mamba_heads):
+                scan_output = mamba_selective_state_update(
+                    recurrent_state[:, head_idx],
+                    hidden_states_B_C[head_idx, ..., 0],
+                    time_step[head_idx, ..., 0],
+                    A[head_idx],
+                    B[head_idx, :, 0],
+                    C[head_idx, :, 0],
+                    self.D[head_idx],
+                    z=gate[head_idx, ..., 0],
+                    dt_bias=time_proj_bias[head_idx] if time_proj_bias is not None else None,
                     dt_softplus=True,
                 ).unsqueeze(-1)
-                scan_outputs = torch.cat((scan_outputs, scan_outputs_), dim=1)
+                scan_outputs.append(scan_output)
 
+            scan_output = torch.cat(scan_outputs, dim=1)
+
+        # Full sequence form
         else:
-            ssm_state = torch.empty(
-                (batch_size, 0, self.mamba_head_dim, self.ssm_state_size),
-                device=hidden_states.device,
-                dtype=hidden_states.dtype,
-            )
-            for n in range(self.n_mamba_heads):
-                scan_outputs_, ssm_state_ = selective_scan_fn(
-                    hidden_states[n],
-                    discrete_time_step[n],
-                    A[n],
-                    B[n].transpose(1, 2),
-                    C[n].transpose(1, 2),
-                    self.D[n].float(),
-                    gate[n],
-                    time_proj_bias[n],
+            output_final_state = cache_params is not None
+            scan_outputs = []
+            final_states = []
+
+            for head_idx in range(self.n_mamba_heads):
+                scan_result = mamba_selective_scan(
+                    hidden_states_B_C[head_idx],
+                    time_step[head_idx],
+                    A[head_idx],
+                    B[head_idx].transpose(1, 2),
+                    C[head_idx].transpose(1, 2),
+                    D=self.D[head_idx].float(),
+                    z=gate[head_idx],
+                    delta_bias=time_proj_bias[head_idx] if time_proj_bias is not None else None,
                     delta_softplus=True,
-                    return_last_state=True,
+                    return_last_state=output_final_state,
+                    # Old model: only when user request it explicitly
+                    use_mambapy=False,
+                    use_associative_scan=False,
                 )
-                scan_outputs = torch.cat((scan_outputs, scan_outputs_), dim=1).contiguous()
-                ssm_state = torch.cat((ssm_state, ssm_state_.unsqueeze(1)), dim=1)
-            if ssm_state is not None and cache_params is not None:
-                cache_params.update_recurrent_state(ssm_state, self.layer_idx)
-
-        # 4. Final linear projection
-        contextualized_states = self.out_proj(scan_outputs.transpose(1, 2))
-        return contextualized_states
-
-    def slow_forward(self, input_states, cache_params: Cache | None = None, attention_mask=None, **kwargs):
-        batch_size, seq_len, _ = input_states.shape
-        dtype = input_states.dtype
-        use_precomputed_states = cache_params is not None and cache_params.has_previous_state(self.layer_idx)
-        # 1. Gated linear projection
-        projected_states = self.in_proj(input_states).transpose(1, 2)
-
-        hidden_states, gate = projected_states.view(batch_size, -1, 2, seq_len).chunk(2, dim=2)
-        hidden_states = hidden_states.squeeze(2).contiguous()
-        gate = gate.reshape(batch_size, self.n_mamba_heads, -1, seq_len).transpose(0, 1)
-
-        # Apply the convolution
-        hidden_states = self._convolution(hidden_states, cache_params, attention_mask, **kwargs)
-
-        if attention_mask is not None:
-            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]
-        hidden_states = hidden_states.reshape(-1, self.n_mamba_heads, self.mamba_head_dim, seq_len).transpose(0, 1)
-        ssm_parameters = (self.x_proj_weight[:, None, :, :] @ 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[:, None] @ time_step.transpose(-1, -2)) + self.dt_proj_bias[
-            :, None, :, None
-        ]
+                if output_final_state:
+                    head_output, final_state = scan_result
+                    scan_outputs.append(head_output)
+                    final_states.append(final_state)
+                else:
+                    scan_outputs.append(scan_result)
 
-        discrete_time_step = nn.functional.softplus(discrete_time_step)
-        if use_precomputed_states:
-            # In training mode, we don't want to perform in-place operations on ssm_state so we can compute the backwards pass
-            ssm_state = cache_params.layers[self.layer_idx].recurrent_states[0].clone()
-        else:
-            ssm_state = torch.zeros(
-                (batch_size, self.n_mamba_heads, self.mamba_head_dim, self.ssm_state_size),
-                device=hidden_states.device,
-                dtype=dtype,
-            )
+            scan_output = torch.cat(scan_outputs, dim=1)
 
-        # 3.b. Discretization: B and C to [batch, seq_len, intermediate_size, ssm_state_size] (SRAM)
-        A = -torch.exp(self.A_log.float())
-        discrete_A = torch.exp(A[:, None, :, None, :] * discrete_time_step[:, :, :, :, None])
-        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(seq_len):
-            ssm_state = discrete_A[:, :, :, i, :].transpose(0, 1) * ssm_state + deltaB_u[:, :, :, i, :].transpose(0, 1)
-            scan_output = torch.matmul(ssm_state.transpose(0, 1).to(dtype), C[:, :, i, :].unsqueeze(-1))
-            scan_outputs.append(scan_output[:, :, :, 0])
-        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)
-
-        if cache_params is not None:
-            cache_params.update_recurrent_state(ssm_state, self.layer_idx)
+            if output_final_state:
+                final_state = torch.stack(final_states, dim=1)
+                cache_params.update_recurrent_state(final_state, self.layer_idx)
 
         # 4. Final linear projection
-        contextualized_states = self.out_proj(
-            scan_output.transpose(0, 1).reshape(batch_size, -1, seq_len).transpose(1, 2)
-        )
+        contextualized_states = self.out_proj(scan_output.transpose(1, 2).to(dtype))
         return contextualized_states
 
-    def forward(self, hidden_states, cache_params: Cache | None = None, attention_mask=None, **kwargs):
-        if self.use_fast_kernels:
-            if not is_fast_path_available or "cuda" not in self.x_proj_weight.device.type:
-                raise ValueError(
-                    "Fast Mamba kernels are not available. Make sure to they are installed and that "
-                    "the mamba module is on a CUDA device. lease run 'pip install causal-conv1d>=1.2.0' "
-                    "and 'pip install mamba-ssm', or set use_mamba_kernels=False in the model's config."
-                )
-            return self.cuda_kernels_forward(hidden_states, cache_params, attention_mask=attention_mask, **kwargs)
-        return self.slow_forward(hidden_states, cache_params, attention_mask=attention_mask, **kwargs)
-
 
 # Copied from transformers.models.mistral.modeling_mistral.MistralMLP with Mistral->Zamba
 class ZambaMLP(nn.Module):

From ad78e0966840e5094afb370f57bdf4bebf103e13 Mon Sep 17 00:00:00 2001
From: vasqu <antonprogamer@gmail.com>
Date: Mon, 3 Aug 2026 22:05:22 +0000
Subject: [PATCH 33/43] style

---
 src/transformers/models/zamba/modeling_zamba.py | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/src/transformers/models/zamba/modeling_zamba.py b/src/transformers/models/zamba/modeling_zamba.py
index ef12466937f7..d0a55071ad1c 100644
--- a/src/transformers/models/zamba/modeling_zamba.py
+++ b/src/transformers/models/zamba/modeling_zamba.py
@@ -390,9 +390,7 @@ def combine_fn(left, right):
     return scan_output
 
 
-@use_kernelized_func(
-    [mamba_selective_scan, mamba_selective_state_update, causal_conv1d_fn, causal_conv1d_update]
-)
+@use_kernelized_func([mamba_selective_scan, mamba_selective_state_update, causal_conv1d_fn, causal_conv1d_update])
 class ZambaMambaMixer(nn.Module):
     """
     Compute ∆, A, B, C, and D the state space parameters and compute the `contextualized_states`.

From 3c51e11d8f80a2cf8b5ac1f837096ad3dd1fbe2c Mon Sep 17 00:00:00 2001
From: vasqu <antonprogamer@gmail.com>
Date: Mon, 3 Aug 2026 22:39:36 +0000
Subject: [PATCH 34/43] fix early cast (leads to non fp32 norm)

---
 src/transformers/models/bamba/modeling_bamba.py                | 3 ---
 src/transformers/models/falcon_h1/modeling_falcon_h1.py        | 3 ---
 .../models/granitemoehybrid/modeling_granitemoehybrid.py       | 3 ---
 src/transformers/models/mamba2/modeling_mamba2.py              | 3 ---
 src/transformers/models/nemotron_h/modeling_nemotron_h.py      | 3 ---
 src/transformers/models/zamba2/modeling_zamba2.py              | 3 ---
 6 files changed, 18 deletions(-)

diff --git a/src/transformers/models/bamba/modeling_bamba.py b/src/transformers/models/bamba/modeling_bamba.py
index decf58d2bcda..63928fd4bb6a 100644
--- a/src/transformers/models/bamba/modeling_bamba.py
+++ b/src/transformers/models/bamba/modeling_bamba.py
@@ -498,7 +498,6 @@ def mamba2_chunk_scan(
     return_final_states: bool = False,
     **kwargs,
 ):
-    input_dtype = hidden_states.dtype
     batch_size, sequence_length, num_heads, head_dim = hidden_states.shape
     num_groups = B.shape[2]
 
@@ -575,8 +574,6 @@ def mamba2_chunk_scan(
     if pad_size > 0:
         output = output[:, :sequence_length]
 
-    output = output.to(input_dtype)
-
     if return_final_states:
         return output, final_state
 
diff --git a/src/transformers/models/falcon_h1/modeling_falcon_h1.py b/src/transformers/models/falcon_h1/modeling_falcon_h1.py
index 03f7216b034e..b53a2a3fc626 100644
--- a/src/transformers/models/falcon_h1/modeling_falcon_h1.py
+++ b/src/transformers/models/falcon_h1/modeling_falcon_h1.py
@@ -505,7 +505,6 @@ def mamba2_chunk_scan(
     return_final_states: bool = False,
     **kwargs,
 ):
-    input_dtype = hidden_states.dtype
     batch_size, sequence_length, num_heads, head_dim = hidden_states.shape
     num_groups = B.shape[2]
 
@@ -582,8 +581,6 @@ def mamba2_chunk_scan(
     if pad_size > 0:
         output = output[:, :sequence_length]
 
-    output = output.to(input_dtype)
-
     if return_final_states:
         return output, final_state
 
diff --git a/src/transformers/models/granitemoehybrid/modeling_granitemoehybrid.py b/src/transformers/models/granitemoehybrid/modeling_granitemoehybrid.py
index f3685dc728b4..0173ef0df9de 100644
--- a/src/transformers/models/granitemoehybrid/modeling_granitemoehybrid.py
+++ b/src/transformers/models/granitemoehybrid/modeling_granitemoehybrid.py
@@ -409,7 +409,6 @@ def mamba2_chunk_scan(
     return_final_states: bool = False,
     **kwargs,
 ):
-    input_dtype = hidden_states.dtype
     batch_size, sequence_length, num_heads, head_dim = hidden_states.shape
     num_groups = B.shape[2]
 
@@ -486,8 +485,6 @@ def mamba2_chunk_scan(
     if pad_size > 0:
         output = output[:, :sequence_length]
 
-    output = output.to(input_dtype)
-
     if return_final_states:
         return output, final_state
 
diff --git a/src/transformers/models/mamba2/modeling_mamba2.py b/src/transformers/models/mamba2/modeling_mamba2.py
index 5ae3991df372..42472d973ad3 100644
--- a/src/transformers/models/mamba2/modeling_mamba2.py
+++ b/src/transformers/models/mamba2/modeling_mamba2.py
@@ -275,7 +275,6 @@ def mamba2_chunk_scan(
     return_final_states: bool = False,
     **kwargs,
 ):
-    input_dtype = hidden_states.dtype
     batch_size, sequence_length, num_heads, head_dim = hidden_states.shape
     num_groups = B.shape[2]
 
@@ -352,8 +351,6 @@ def mamba2_chunk_scan(
     if pad_size > 0:
         output = output[:, :sequence_length]
 
-    output = output.to(input_dtype)
-
     if return_final_states:
         return output, final_state
 
diff --git a/src/transformers/models/nemotron_h/modeling_nemotron_h.py b/src/transformers/models/nemotron_h/modeling_nemotron_h.py
index 6cba8aa1ff6a..681990420373 100644
--- a/src/transformers/models/nemotron_h/modeling_nemotron_h.py
+++ b/src/transformers/models/nemotron_h/modeling_nemotron_h.py
@@ -271,7 +271,6 @@ def mamba2_chunk_scan(
     return_final_states: bool = False,
     **kwargs,
 ):
-    input_dtype = hidden_states.dtype
     batch_size, sequence_length, num_heads, head_dim = hidden_states.shape
     num_groups = B.shape[2]
 
@@ -348,8 +347,6 @@ def mamba2_chunk_scan(
     if pad_size > 0:
         output = output[:, :sequence_length]
 
-    output = output.to(input_dtype)
-
     if return_final_states:
         return output, final_state
 
diff --git a/src/transformers/models/zamba2/modeling_zamba2.py b/src/transformers/models/zamba2/modeling_zamba2.py
index 8c8634e4b9e2..a42139d0a2f8 100644
--- a/src/transformers/models/zamba2/modeling_zamba2.py
+++ b/src/transformers/models/zamba2/modeling_zamba2.py
@@ -563,7 +563,6 @@ def mamba2_chunk_scan(
     return_final_states: bool = False,
     **kwargs,
 ):
-    input_dtype = hidden_states.dtype
     batch_size, sequence_length, num_heads, head_dim = hidden_states.shape
     num_groups = B.shape[2]
 
@@ -640,8 +639,6 @@ def mamba2_chunk_scan(
     if pad_size > 0:
         output = output[:, :sequence_length]
 
-    output = output.to(input_dtype)
-
     if return_final_states:
         return output, final_state
 

From b05fb5227e2a3f276494a3a4ea3654139d5e04c4 Mon Sep 17 00:00:00 2001
From: vasqu <antonprogamer@gmail.com>
Date: Mon, 3 Aug 2026 23:16:46 +0000
Subject: [PATCH 35/43] fix padding free path

---
 src/transformers/models/olmo_hybrid/modeling_olmo_hybrid.py | 2 ++
 src/transformers/models/olmo_hybrid/modular_olmo_hybrid.py  | 2 ++
 src/transformers/models/qwen3_5/modeling_qwen3_5.py         | 2 ++
 src/transformers/models/qwen3_5/modular_qwen3_5.py          | 2 ++
 src/transformers/models/qwen3_5_moe/modeling_qwen3_5_moe.py | 2 ++
 src/transformers/models/qwen3_next/modeling_qwen3_next.py   | 2 ++
 src/transformers/models/qwen3_next/modular_qwen3_next.py    | 2 ++
 7 files changed, 14 insertions(+)

diff --git a/src/transformers/models/olmo_hybrid/modeling_olmo_hybrid.py b/src/transformers/models/olmo_hybrid/modeling_olmo_hybrid.py
index 4abd5d6232c3..b2ac2b682d8a 100644
--- a/src/transformers/models/olmo_hybrid/modeling_olmo_hybrid.py
+++ b/src/transformers/models/olmo_hybrid/modeling_olmo_hybrid.py
@@ -652,6 +652,7 @@ def forward(
                 initial_state=recurrent_state,
                 output_final_state=use_cache,
                 use_qk_l2norm_in_kernel=True,
+                cu_seqlens=kwargs.pop("cu_seq_lens_q", None),
                 **kwargs,
             )
         else:
@@ -664,6 +665,7 @@ def forward(
                 initial_state=recurrent_state if use_precomputed_states else None,
                 output_final_state=use_cache,
                 use_qk_l2norm_in_kernel=True,
+                cu_seqlens=kwargs.pop("cu_seq_lens_q", None),
                 **kwargs,
             )
 
diff --git a/src/transformers/models/olmo_hybrid/modular_olmo_hybrid.py b/src/transformers/models/olmo_hybrid/modular_olmo_hybrid.py
index d0fab6a6109c..f004471bf301 100644
--- a/src/transformers/models/olmo_hybrid/modular_olmo_hybrid.py
+++ b/src/transformers/models/olmo_hybrid/modular_olmo_hybrid.py
@@ -424,6 +424,7 @@ def forward(
                 initial_state=recurrent_state,
                 output_final_state=use_cache,
                 use_qk_l2norm_in_kernel=True,
+                cu_seqlens=kwargs.pop("cu_seq_lens_q", None),
                 **kwargs,
             )
         else:
@@ -436,6 +437,7 @@ def forward(
                 initial_state=recurrent_state if use_precomputed_states else None,
                 output_final_state=use_cache,
                 use_qk_l2norm_in_kernel=True,
+                cu_seqlens=kwargs.pop("cu_seq_lens_q", None),
                 **kwargs,
             )
 
diff --git a/src/transformers/models/qwen3_5/modeling_qwen3_5.py b/src/transformers/models/qwen3_5/modeling_qwen3_5.py
index 831fa13c4b89..ea8b907ef5e8 100644
--- a/src/transformers/models/qwen3_5/modeling_qwen3_5.py
+++ b/src/transformers/models/qwen3_5/modeling_qwen3_5.py
@@ -516,6 +516,7 @@ def forward(
                 initial_state=recurrent_state,
                 output_final_state=cache_params is not None,
                 use_qk_l2norm_in_kernel=True,
+                cu_seqlens=kwargs.pop("cu_seq_lens_q", None),
                 **kwargs,
             )
         else:
@@ -528,6 +529,7 @@ def forward(
                 initial_state=recurrent_state,
                 output_final_state=cache_params is not None,
                 use_qk_l2norm_in_kernel=True,
+                cu_seqlens=kwargs.pop("cu_seq_lens_q", None),
                 **kwargs,
             )
 
diff --git a/src/transformers/models/qwen3_5/modular_qwen3_5.py b/src/transformers/models/qwen3_5/modular_qwen3_5.py
index 2e1834282c51..3196ee52a4d2 100644
--- a/src/transformers/models/qwen3_5/modular_qwen3_5.py
+++ b/src/transformers/models/qwen3_5/modular_qwen3_5.py
@@ -304,6 +304,7 @@ def forward(
                 initial_state=recurrent_state,
                 output_final_state=cache_params is not None,
                 use_qk_l2norm_in_kernel=True,
+                cu_seqlens=kwargs.pop("cu_seq_lens_q", None),
                 **kwargs,
             )
         else:
@@ -316,6 +317,7 @@ def forward(
                 initial_state=recurrent_state,
                 output_final_state=cache_params is not None,
                 use_qk_l2norm_in_kernel=True,
+                cu_seqlens=kwargs.pop("cu_seq_lens_q", None),
                 **kwargs,
             )
 
diff --git a/src/transformers/models/qwen3_5_moe/modeling_qwen3_5_moe.py b/src/transformers/models/qwen3_5_moe/modeling_qwen3_5_moe.py
index e12ba4482672..e264b9ca70d1 100644
--- a/src/transformers/models/qwen3_5_moe/modeling_qwen3_5_moe.py
+++ b/src/transformers/models/qwen3_5_moe/modeling_qwen3_5_moe.py
@@ -514,6 +514,7 @@ def forward(
                 initial_state=recurrent_state,
                 output_final_state=cache_params is not None,
                 use_qk_l2norm_in_kernel=True,
+                cu_seqlens=kwargs.pop("cu_seq_lens_q", None),
                 **kwargs,
             )
         else:
@@ -526,6 +527,7 @@ def forward(
                 initial_state=recurrent_state,
                 output_final_state=cache_params is not None,
                 use_qk_l2norm_in_kernel=True,
+                cu_seqlens=kwargs.pop("cu_seq_lens_q", None),
                 **kwargs,
             )
 
diff --git a/src/transformers/models/qwen3_next/modeling_qwen3_next.py b/src/transformers/models/qwen3_next/modeling_qwen3_next.py
index 00e88fd9cd9c..8d68766f23fb 100644
--- a/src/transformers/models/qwen3_next/modeling_qwen3_next.py
+++ b/src/transformers/models/qwen3_next/modeling_qwen3_next.py
@@ -667,6 +667,7 @@ def forward(
                 initial_state=recurrent_state,
                 output_final_state=cache_params is not None,
                 use_qk_l2norm_in_kernel=True,
+                cu_seqlens=kwargs.pop("cu_seq_lens_q", None),
                 **kwargs,
             )
         else:
@@ -679,6 +680,7 @@ def forward(
                 initial_state=recurrent_state,
                 output_final_state=cache_params is not None,
                 use_qk_l2norm_in_kernel=True,
+                cu_seqlens=kwargs.pop("cu_seq_lens_q", None),
                 **kwargs,
             )
 
diff --git a/src/transformers/models/qwen3_next/modular_qwen3_next.py b/src/transformers/models/qwen3_next/modular_qwen3_next.py
index 3b10b4a54a67..a858d877c955 100644
--- a/src/transformers/models/qwen3_next/modular_qwen3_next.py
+++ b/src/transformers/models/qwen3_next/modular_qwen3_next.py
@@ -500,6 +500,7 @@ def forward(
                 initial_state=recurrent_state,
                 output_final_state=cache_params is not None,
                 use_qk_l2norm_in_kernel=True,
+                cu_seqlens=kwargs.pop("cu_seq_lens_q", None),
                 **kwargs,
             )
         else:
@@ -512,6 +513,7 @@ def forward(
                 initial_state=recurrent_state,
                 output_final_state=cache_params is not None,
                 use_qk_l2norm_in_kernel=True,
+                cu_seqlens=kwargs.pop("cu_seq_lens_q", None),
                 **kwargs,
             )
 

From 98f62d4128e9f4da65190e88e0f3556c927b30a7 Mon Sep 17 00:00:00 2001
From: vasqu <antonprogamer@gmail.com>
Date: Mon, 3 Aug 2026 23:42:51 +0000
Subject: [PATCH 36/43] fix offload

---
 src/transformers/models/olmo_hybrid/modeling_olmo_hybrid.py | 2 ++
 src/transformers/models/olmo_hybrid/modular_olmo_hybrid.py  | 2 ++
 2 files changed, 4 insertions(+)

diff --git a/src/transformers/models/olmo_hybrid/modeling_olmo_hybrid.py b/src/transformers/models/olmo_hybrid/modeling_olmo_hybrid.py
index b2ac2b682d8a..cf6add73f5b1 100644
--- a/src/transformers/models/olmo_hybrid/modeling_olmo_hybrid.py
+++ b/src/transformers/models/olmo_hybrid/modeling_olmo_hybrid.py
@@ -31,6 +31,7 @@
 from ...cache_utils import Cache, DynamicCache
 from ...generation import GenerationMixin
 from ...integrations import use_kernel_forward_from_hub, use_kernel_func_from_hub_with_fallback, use_kernelized_func
+from ...integrations.accelerate import force_accelerate_hooks
 from ...masking_utils import create_causal_mask, create_recurrent_attention_mask
 from ...modeling_layers import GradientCheckpointingLayer
 from ...modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
@@ -562,6 +563,7 @@ def __init__(self, config: OlmoHybridConfig, layer_idx: int):
 
         self.layer_type = config.layer_types[layer_idx]
 
+    @force_accelerate_hooks("conv1d")
     def forward(
         self,
         hidden_states: torch.Tensor,
diff --git a/src/transformers/models/olmo_hybrid/modular_olmo_hybrid.py b/src/transformers/models/olmo_hybrid/modular_olmo_hybrid.py
index f004471bf301..47f3d7d4292f 100644
--- a/src/transformers/models/olmo_hybrid/modular_olmo_hybrid.py
+++ b/src/transformers/models/olmo_hybrid/modular_olmo_hybrid.py
@@ -26,6 +26,7 @@
 from ...cache_utils import Cache, DynamicCache
 from ...configuration_utils import PreTrainedConfig, remap_legacy_layer_types
 from ...integrations import use_kernelized_func
+from ...integrations.accelerate import force_accelerate_hooks
 from ...masking_utils import create_causal_mask, create_recurrent_attention_mask
 from ...modeling_outputs import BaseModelOutputWithPast
 from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
@@ -334,6 +335,7 @@ def __init__(self, config: OlmoHybridConfig, layer_idx: int):
 
         self.layer_type = config.layer_types[layer_idx]
 
+    @force_accelerate_hooks("conv1d")
     def forward(
         self,
         hidden_states: torch.Tensor,

From 9b266aa3815304a3a7714624ce9aad1c7d3d4c7d Mon Sep 17 00:00:00 2001
From: vasqu <antonprogamer@gmail.com>
Date: Tue, 4 Aug 2026 12:02:48 +0000
Subject: [PATCH 37/43] bump kernels

---
 src/transformers/integrations/hub_kernels.py | 43 ++++++++++++--------
 1 file changed, 26 insertions(+), 17 deletions(-)

diff --git a/src/transformers/integrations/hub_kernels.py b/src/transformers/integrations/hub_kernels.py
index dd5e35f1c414..d05f281532ce 100644
--- a/src/transformers/integrations/hub_kernels.py
+++ b/src/transformers/integrations/hub_kernels.py
@@ -168,12 +168,12 @@ def _build_kernel_mapping() -> dict:
                     Mode.TRAINING: LayerRepository(
                         repo_id="kernels-community/mamba-ssm",
                         layer_name="causal_conv1d_fn",
-                        version=1,
+                        version=2,
                     ),
                     Mode.INFERENCE: LayerRepository(
                         repo_id="kernels-community/mamba-ssm",
                         layer_name="causal_conv1d_fn",
-                        version=1,
+                        version=2,
                     ),
                 },
             },
@@ -182,12 +182,12 @@ def _build_kernel_mapping() -> dict:
                     Mode.TRAINING: LayerRepository(
                         repo_id="kernels-community/mamba-ssm",
                         layer_name="causal_conv1d_update",
-                        version=1,
+                        version=2,
                     ),
                     Mode.INFERENCE: LayerRepository(
                         repo_id="kernels-community/mamba-ssm",
                         layer_name="causal_conv1d_update",
-                        version=1,
+                        version=2,
                     ),
                 },
             },
@@ -224,12 +224,12 @@ def _build_kernel_mapping() -> dict:
                     Mode.TRAINING: LayerRepository(
                         repo_id="kernels-community/mamba-ssm",
                         layer_name="mamba_chunk_scan_combined",
-                        version=1,
+                        version=2,
                     ),
                     Mode.INFERENCE: LayerRepository(
                         repo_id="kernels-community/mamba-ssm",
                         layer_name="mamba_chunk_scan_combined",
-                        version=1,
+                        version=2,
                     ),
                 },
             },
@@ -238,12 +238,12 @@ def _build_kernel_mapping() -> dict:
                     Mode.TRAINING: LayerRepository(
                         repo_id="kernels-community/mamba-ssm",
                         layer_name="mamba_split_conv1d_scan_combined",
-                        version=1,
+                        version=2,
                     ),
                     Mode.INFERENCE: LayerRepository(
                         repo_id="kernels-community/mamba-ssm",
                         layer_name="mamba_split_conv1d_scan_combined",
-                        version=1,
+                        version=2,
                     ),
                 },
             },
@@ -252,12 +252,12 @@ def _build_kernel_mapping() -> dict:
                     Mode.TRAINING: LayerRepository(
                         repo_id="kernels-community/mamba-ssm",
                         layer_name="mamba_inner_fn",
-                        version=1,
+                        version=2,
                     ),
                     Mode.INFERENCE: LayerRepository(
                         repo_id="kernels-community/mamba-ssm",
                         layer_name="mamba_inner_fn",
-                        version=1,
+                        version=2,
                     ),
                 },
             },
@@ -266,12 +266,12 @@ def _build_kernel_mapping() -> dict:
                     Mode.TRAINING: LayerRepository(
                         repo_id="kernels-community/mamba-ssm",
                         layer_name="selective_scan_fn",
-                        version=1,
+                        version=2,
                     ),
                     Mode.INFERENCE: LayerRepository(
                         repo_id="kernels-community/mamba-ssm",
                         layer_name="selective_scan_fn",
-                        version=1,
+                        version=2,
                     ),
                 },
             },
@@ -280,12 +280,12 @@ def _build_kernel_mapping() -> dict:
                     Mode.TRAINING: LayerRepository(
                         repo_id="kernels-community/mamba-ssm",
                         layer_name="selective_state_update",
-                        version=1,
+                        version=2,
                     ),
                     Mode.INFERENCE: LayerRepository(
                         repo_id="kernels-community/mamba-ssm",
                         layer_name="selective_state_update",
-                        version=1,
+                        version=2,
                     ),
                 },
             },
@@ -577,9 +577,6 @@ def register_kernel_mapping_transformers(*args, **kwargs):
 
 
 _HUB_KERNEL_MAPPING: dict[str, dict[str, str]] = {
-    "causal-conv1d": {"repo_id": "kernels-community/causal-conv1d", "version": 1},
-    "mamba-ssm": {"repo_id": "kernels-community/mamba-ssm", "version": 1},
-    "falcon_mamba-ssm": {"repo_id": "kernels-community/mamba-ssm", "version": 1},
     "finegrained-fp8": {"repo_id": "kernels-community/finegrained-fp8", "version": 4},
     "deep-gemm": {"repo_id": "kernels-community/deep-gemm", "version": 2},
     "sonic-moe": {"repo_id": "kernels-community/sonic-moe", "revision": "ep-support"},
@@ -765,6 +762,16 @@ def get_kernel(
 
 
 def use_kernel_func_from_hub_with_fallback(func_name: str, package: str, internal_path: str | None = None):
+    """
+    The same as `use_kernel_forward_from_hub` but with the optional fallback to an original package if it exists, e.g.,
+    FLA for Gated Delta Rule, mamba-ssm for mamba2, etc.
+
+    This combines all options with kernels, enabling kernels on top of the original package if requested as well.
+    The order of priority is
+        1. Hf kernels (if requested)
+        2. Original package
+        3. Torch only path
+    """
     kernel_wrapper_decorator = use_kernel_forward_from_hub(func_name)
 
     # Allow internal path prefix if given to resolve non __init__ imports
@@ -774,6 +781,8 @@ def use_kernel_func_from_hub_with_fallback(func_name: str, package: str, interna
     def decorator(torch_function: Callable) -> Callable:
         implementation = None
         try:
+            # TODO: remove raise to force torch path
+            raise
             module = importlib.import_module(package)
             implementation = resolve_internal_import(module, full_path)
         except Exception:

From c4b785c89f1eeb6ca9669ae9aca55daea1a3d48f Mon Sep 17 00:00:00 2001
From: vasqu <antonprogamer@gmail.com>
Date: Tue, 4 Aug 2026 12:45:32 +0000
Subject: [PATCH 38/43] remove todo

---
 src/transformers/integrations/hub_kernels.py | 2 --
 1 file changed, 2 deletions(-)

diff --git a/src/transformers/integrations/hub_kernels.py b/src/transformers/integrations/hub_kernels.py
index d05f281532ce..1da7f75f585d 100644
--- a/src/transformers/integrations/hub_kernels.py
+++ b/src/transformers/integrations/hub_kernels.py
@@ -781,8 +781,6 @@ def use_kernel_func_from_hub_with_fallback(func_name: str, package: str, interna
     def decorator(torch_function: Callable) -> Callable:
         implementation = None
         try:
-            # TODO: remove raise to force torch path
-            raise
             module = importlib.import_module(package)
             implementation = resolve_internal_import(module, full_path)
         except Exception:

From c85e6220d580ba9c5aaf9bf96c10919a5f154313 Mon Sep 17 00:00:00 2001
From: vasqu <antonprogamer@gmail.com>
Date: Tue, 4 Aug 2026 13:14:25 +0000
Subject: [PATCH 39/43] avoid onnx export and fix mamba2 test

---
 .../models/falcon_mamba/modeling_falcon_mamba.py   | 14 ++++++++++++--
 src/transformers/models/jamba/modeling_jamba.py    | 14 ++++++++++++--
 src/transformers/models/mamba/modeling_mamba.py    | 10 ++++++++--
 tests/models/mamba2/test_modeling_mamba2.py        |  6 ++----
 4 files changed, 34 insertions(+), 10 deletions(-)

diff --git a/src/transformers/models/falcon_mamba/modeling_falcon_mamba.py b/src/transformers/models/falcon_mamba/modeling_falcon_mamba.py
index ea774d23a5fe..823ee4c6f159 100644
--- a/src/transformers/models/falcon_mamba/modeling_falcon_mamba.py
+++ b/src/transformers/models/falcon_mamba/modeling_falcon_mamba.py
@@ -35,7 +35,12 @@
 from ...modeling_layers import GradientCheckpointingLayer
 from ...modeling_utils import PreTrainedModel
 from ...utils import ModelOutput, auto_docstring
-from ...utils.import_utils import is_mambapy_available, is_torch_greater_or_equal, is_tracing
+from ...utils.import_utils import (
+    is_mambapy_available,
+    is_torch_greater_or_equal,
+    is_torchdynamo_compiling,
+    is_torchdynamo_exporting,
+)
 from .configuration_falcon_mamba import FalconMambaConfig
 
 
@@ -240,7 +245,12 @@ def mamba_selective_scan(
         scan_output = (all_states @ C.unsqueeze(-1)).squeeze(3).transpose(1, 2)
         ssm_state = all_states[:, -1]
 
-    elif use_associative_scan and associative_scan is not None and is_tracing(hidden_states):
+    elif (
+        use_associative_scan
+        and associative_scan is not None
+        # There is no onnx translation for this op so we rely on the normal sequential path then
+        and (is_torchdynamo_compiling() and not is_torchdynamo_exporting())
+    ):
 
         def combine_fn(left, right):
             a_left, b_left = left
diff --git a/src/transformers/models/jamba/modeling_jamba.py b/src/transformers/models/jamba/modeling_jamba.py
index 44ba6fbf9bda..ea9ab6389925 100755
--- a/src/transformers/models/jamba/modeling_jamba.py
+++ b/src/transformers/models/jamba/modeling_jamba.py
@@ -46,7 +46,12 @@
 from ...processing_utils import Unpack
 from ...utils import TransformersKwargs, auto_docstring, can_return_tuple
 from ...utils.generic import merge_with_config_defaults
-from ...utils.import_utils import is_mambapy_available, is_torch_greater_or_equal, is_tracing
+from ...utils.import_utils import (
+    is_mambapy_available,
+    is_torch_greater_or_equal,
+    is_torchdynamo_compiling,
+    is_torchdynamo_exporting,
+)
 from ...utils.output_capturing import OutputRecorder, capture_outputs
 from .configuration_jamba import JambaConfig
 
@@ -381,7 +386,12 @@ def mamba_selective_scan(
         scan_output = (all_states @ C.unsqueeze(-1)).squeeze(3).transpose(1, 2)
         ssm_state = all_states[:, -1]
 
-    elif use_associative_scan and associative_scan is not None and is_tracing(hidden_states):
+    elif (
+        use_associative_scan
+        and associative_scan is not None
+        # There is no onnx translation for this op so we rely on the normal sequential path then
+        and (is_torchdynamo_compiling() and not is_torchdynamo_exporting())
+    ):
 
         def combine_fn(left, right):
             a_left, b_left = left
diff --git a/src/transformers/models/mamba/modeling_mamba.py b/src/transformers/models/mamba/modeling_mamba.py
index 25fb5383ea8b..e16b1d853929 100644
--- a/src/transformers/models/mamba/modeling_mamba.py
+++ b/src/transformers/models/mamba/modeling_mamba.py
@@ -37,7 +37,8 @@
 from ...utils.import_utils import (
     is_mambapy_available,
     is_torch_greater_or_equal,
-    is_tracing,
+    is_torchdynamo_compiling,
+    is_torchdynamo_exporting,
 )
 from .configuration_mamba import MambaConfig
 
@@ -229,7 +230,12 @@ def mamba_selective_scan(
         scan_output = (all_states @ C.unsqueeze(-1)).squeeze(3).transpose(1, 2)
         ssm_state = all_states[:, -1]
 
-    elif use_associative_scan and associative_scan is not None and is_tracing(hidden_states):
+    elif (
+        use_associative_scan
+        and associative_scan is not None
+        # There is no onnx translation for this op so we rely on the normal sequential path then
+        and (is_torchdynamo_compiling() and not is_torchdynamo_exporting())
+    ):
 
         def combine_fn(left, right):
             a_left, b_left = left
diff --git a/tests/models/mamba2/test_modeling_mamba2.py b/tests/models/mamba2/test_modeling_mamba2.py
index 664353e6b6c5..329c7754c430 100644
--- a/tests/models/mamba2/test_modeling_mamba2.py
+++ b/tests/models/mamba2/test_modeling_mamba2.py
@@ -423,11 +423,9 @@ def test_simple_generate(self):
 
         model = Mamba2ForCausalLM.from_pretrained(self.model_id, dtype=torch.bfloat16)
         model.to(torch_device)
-        input_ids = tokenizer("[INST]Write a hello world program in C++.[/INST]", return_tensors="pt")["input_ids"].to(
-            torch_device
-        )
+        inputs = tokenizer("[INST]Write a hello world program in C++.[/INST]", return_tensors="pt").to(torch_device)
 
-        out = model.generate(input_ids, do_sample=False, use_cache=True, max_new_tokens=30)
+        out = model.generate(**inputs, do_sample=False, use_cache=True, max_new_tokens=30)
         output_sentence = tokenizer.decode(out[0])
         ground_truth_sentences = Expectations(
             {

From 2b4fd754e8de887e5b1dc368a9210317fac2d8e6 Mon Sep 17 00:00:00 2001
From: vasqu <antonprogamer@gmail.com>
Date: Tue, 4 Aug 2026 13:31:38 +0000
Subject: [PATCH 40/43] fix bamba test

---
 tests/models/bamba/test_modeling_bamba.py | 8 +++-----
 1 file changed, 3 insertions(+), 5 deletions(-)

diff --git a/tests/models/bamba/test_modeling_bamba.py b/tests/models/bamba/test_modeling_bamba.py
index 5580fb16d40c..9c7c0e1640b4 100644
--- a/tests/models/bamba/test_modeling_bamba.py
+++ b/tests/models/bamba/test_modeling_bamba.py
@@ -568,10 +568,8 @@ def test_simple_generate(self):
         )
         # fmt: on
 
-        input_ids = self.tokenizer("Hey how are you doing on this lovely evening?", return_tensors="pt")[
-            "input_ids"
-        ].to(torch_device)
-        out = self.model.generate(input_ids, do_sample=False, max_new_tokens=10)
+        inputs = self.tokenizer("Hey how are you doing on this lovely evening?", return_tensors="pt").to(torch_device)
+        out = self.model.generate(**inputs, do_sample=False, max_new_tokens=10)
         output_sentence = self.tokenizer.decode(out[0, :])
         expected = expectations.get_expectation()
         self.assertEqual(output_sentence, expected)
@@ -579,7 +577,7 @@ def test_simple_generate(self):
         # TODO: there are significant differences in the logits across major cuda versions, which shouldn't exist
         if self.device_properties[0] == "cuda" and self.device_properties[1] == 8:
             with torch.no_grad():
-                logits = self.model(input_ids=input_ids, logits_to_keep=40).logits
+                logits = self.model(**inputs, logits_to_keep=40).logits
 
             EXPECTED_LOGITS_NO_GRAD = torch.tensor(
                 [

From 2e485f3f6e893af8505f76d42f412a1754fe5a4a Mon Sep 17 00:00:00 2001
From: vasqu <antonprogamer@gmail.com>
Date: Tue, 4 Aug 2026 14:05:01 +0000
Subject: [PATCH 41/43] update falcon mamba - aligned with all other devices

---
 tests/models/falcon_mamba/test_modeling_falcon_mamba.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/models/falcon_mamba/test_modeling_falcon_mamba.py b/tests/models/falcon_mamba/test_modeling_falcon_mamba.py
index 2ed54e6b6e79..ddec05c6db20 100644
--- a/tests/models/falcon_mamba/test_modeling_falcon_mamba.py
+++ b/tests/models/falcon_mamba/test_modeling_falcon_mamba.py
@@ -551,7 +551,7 @@ def test_batched_generation(self):
                     ' I will be talking about the importance of the internet in our lives.\nThe internet is a global'
                 ],
                 ("cuda", (8, 6)): [
-                    ' I will be talking about the “Theory of Relativity” by Albert Einstein.\nThe',
+                    ' I am going to talk about the “Theory of Relativity” by Albert Einstein.\n',
                     ' I will be talking about the importance of the internet in our lives.\nThe internet is a global'
                 ],
                 ("cuda", 9): [

From d81829f6b0f50d630e0423589a19565a4738339c Mon Sep 17 00:00:00 2001
From: vasqu <antonprogamer@gmail.com>
Date: Tue, 4 Aug 2026 14:12:53 +0000
Subject: [PATCH 42/43] oops

---
 tests/models/falcon_mamba/test_modeling_falcon_mamba.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/models/falcon_mamba/test_modeling_falcon_mamba.py b/tests/models/falcon_mamba/test_modeling_falcon_mamba.py
index ddec05c6db20..a827ef29c146 100644
--- a/tests/models/falcon_mamba/test_modeling_falcon_mamba.py
+++ b/tests/models/falcon_mamba/test_modeling_falcon_mamba.py
@@ -514,7 +514,7 @@ def test_batched_generation(self):
                     "Hello my name is Younes and today I will be talking about the importance of the internet in our lives.\nThe internet is a global",
                 ],
                 ("cuda", (8, 6)): [
-                    "Hello today I will be talking about the “Theory of Relativity” by Albert Einstein.\nThe",
+                    "Hello today I am going to talk about the “Theory of Relativity” by Albert Einstein.\n",
                     "Hello my name is Younes and today I will be talking about the importance of the internet in our lives.\nThe internet is a global",
                 ],
                 ("cuda", 9): [

From 4c1864de12f0314b6647edc915510eeb4ec4a6bd Mon Sep 17 00:00:00 2001
From: vasqu <antonprogamer@gmail.com>
Date: Wed, 5 Aug 2026 08:28:50 +0000
Subject: [PATCH 43/43] adress review

---
 .../models/bamba/modeling_bamba.py            | 19 ++++-------------
 .../models/falcon_h1/modeling_falcon_h1.py    | 21 +++++--------------
 .../models/falcon_h1/modular_falcon_h1.py     |  6 ++----
 .../falcon_mamba/modeling_falcon_mamba.py     | 18 +++++-----------
 .../falcon_mamba/modular_falcon_mamba.py      |  3 ++-
 .../modeling_granitemoehybrid.py              | 19 ++++-------------
 .../models/jamba/modeling_jamba.py            | 15 +++----------
 .../models/mamba/modeling_mamba.py            | 15 +++----------
 .../models/mamba2/modeling_mamba2.py          | 19 ++++-------------
 .../models/nemotron_h/modeling_nemotron_h.py  | 19 ++++-------------
 .../olmo_hybrid/modeling_olmo_hybrid.py       |  2 +-
 .../models/olmo_hybrid/modular_olmo_hybrid.py |  2 +-
 .../models/zamba2/modeling_zamba2.py          | 19 ++++-------------
 13 files changed, 42 insertions(+), 135 deletions(-)

diff --git a/src/transformers/models/bamba/modeling_bamba.py b/src/transformers/models/bamba/modeling_bamba.py
index 017e64d87322..2a06e67cd53e 100644
--- a/src/transformers/models/bamba/modeling_bamba.py
+++ b/src/transformers/models/bamba/modeling_bamba.py
@@ -384,10 +384,7 @@ def causal_conv1d_fn(
     return out.to(hidden_states.dtype)
 
 
-@use_kernel_func_from_hub_with_fallback(
-    "mamba_split_conv1d_scan_combined",
-    "mamba_ssm",
-)
+@use_kernel_func_from_hub_with_fallback("mamba_split_conv1d_scan_combined", "mamba_ssm")
 def mamba2_split_conv1d_scan_combined(
     zxbcdt: torch.Tensor,
     conv1d_weight: torch.Tensor,
@@ -412,10 +409,7 @@ def mamba2_split_conv1d_scan_combined(
     return None
 
 
-@use_kernel_func_from_hub_with_fallback(
-    "selective_state_update",
-    "mamba_ssm",
-)
+@use_kernel_func_from_hub_with_fallback("selective_state_update", "mamba_ssm")
 def mamba2_selective_state_update(
     state: torch.Tensor,
     hidden_states: torch.Tensor,
@@ -477,10 +471,7 @@ def mamba2_selective_state_update(
     return out.to(hidden_states.dtype)
 
 
-@use_kernel_func_from_hub_with_fallback(
-    "mamba_chunk_scan_combined",
-    "mamba_ssm",
-)
+@use_kernel_func_from_hub_with_fallback("mamba_chunk_scan_combined", "mamba_ssm")
 def mamba2_chunk_scan(
     hidden_states: torch.Tensor,
     dt: torch.Tensor,
@@ -673,9 +664,7 @@ def forward(
         projected_states = self.in_proj(hidden_states)
 
         A = -torch.exp(self.A_log.float())
-        fused_kwargs = (
-            kwargs | {} if self.time_step_limit == (0.0, float("inf")) else kwargs | {"dt_limit": self.time_step_limit}
-        )
+        fused_kwargs = kwargs | {"dt_limit": self.time_step_limit}
         if self.training and cache_params is None:
             fused_output = mamba2_split_conv1d_scan_combined(
                 projected_states,
diff --git a/src/transformers/models/falcon_h1/modeling_falcon_h1.py b/src/transformers/models/falcon_h1/modeling_falcon_h1.py
index c23ed86786ab..9edbd2b59b32 100644
--- a/src/transformers/models/falcon_h1/modeling_falcon_h1.py
+++ b/src/transformers/models/falcon_h1/modeling_falcon_h1.py
@@ -391,10 +391,7 @@ def causal_conv1d_fn(
     return out.to(hidden_states.dtype)
 
 
-@use_kernel_func_from_hub_with_fallback(
-    "mamba_split_conv1d_scan_combined",
-    "mamba_ssm",
-)
+@use_kernel_func_from_hub_with_fallback("mamba_split_conv1d_scan_combined", "mamba_ssm")
 def mamba2_split_conv1d_scan_combined(
     zxbcdt: torch.Tensor,
     conv1d_weight: torch.Tensor,
@@ -419,10 +416,7 @@ def mamba2_split_conv1d_scan_combined(
     return None
 
 
-@use_kernel_func_from_hub_with_fallback(
-    "selective_state_update",
-    "mamba_ssm",
-)
+@use_kernel_func_from_hub_with_fallback("selective_state_update", "mamba_ssm")
 def mamba2_selective_state_update(
     state: torch.Tensor,
     hidden_states: torch.Tensor,
@@ -484,10 +478,7 @@ def mamba2_selective_state_update(
     return out.to(hidden_states.dtype)
 
 
-@use_kernel_func_from_hub_with_fallback(
-    "mamba_chunk_scan_combined",
-    "mamba_ssm",
-)
+@use_kernel_func_from_hub_with_fallback("mamba_chunk_scan_combined", "mamba_ssm")
 def mamba2_chunk_scan(
     hidden_states: torch.Tensor,
     dt: torch.Tensor,
@@ -692,9 +683,7 @@ def forward(
         projected_states = projected_states * self.mup_vector
 
         A = -torch.exp(self.A_log.float())
-        fused_kwargs = (
-            kwargs | {} if self.time_step_limit == (0.0, float("inf")) else kwargs | {"dt_limit": self.time_step_limit}
-        )
+        fused_kwargs = kwargs | {"dt_limit": self.time_step_limit}
         if self.training and cache_params is None:
             fused_output = mamba2_split_conv1d_scan_combined(
                 projected_states,
@@ -730,7 +719,7 @@ def forward(
 
         # 2. Convolution sequence transformation
         hidden_states_B_C = hidden_states_B_C.transpose(1, 2)
-        if use_precomputed_states and seq_len == 1:
+        if use_precomputed_states and seq_len == 1 and not cache_params.layers[self.layer_idx].record_past:
             hidden_states_B_C = causal_conv1d_update(
                 hidden_states_B_C,
                 conv_state,
diff --git a/src/transformers/models/falcon_h1/modular_falcon_h1.py b/src/transformers/models/falcon_h1/modular_falcon_h1.py
index 3a00b3694991..0a8fed7fb230 100644
--- a/src/transformers/models/falcon_h1/modular_falcon_h1.py
+++ b/src/transformers/models/falcon_h1/modular_falcon_h1.py
@@ -195,9 +195,7 @@ def forward(
         projected_states = projected_states * self.mup_vector
 
         A = -torch.exp(self.A_log.float())
-        fused_kwargs = (
-            kwargs | {} if self.time_step_limit == (0.0, float("inf")) else kwargs | {"dt_limit": self.time_step_limit}
-        )
+        fused_kwargs = kwargs | {"dt_limit": self.time_step_limit}
         if self.training and cache_params is None:
             fused_output = mamba2_split_conv1d_scan_combined(
                 projected_states,
@@ -233,7 +231,7 @@ def forward(
 
         # 2. Convolution sequence transformation
         hidden_states_B_C = hidden_states_B_C.transpose(1, 2)
-        if use_precomputed_states and seq_len == 1:
+        if use_precomputed_states and seq_len == 1 and not cache_params.layers[self.layer_idx].record_past:
             hidden_states_B_C = causal_conv1d_update(
                 hidden_states_B_C,
                 conv_state,
diff --git a/src/transformers/models/falcon_mamba/modeling_falcon_mamba.py b/src/transformers/models/falcon_mamba/modeling_falcon_mamba.py
index 823ee4c6f159..eb3b73bb7431 100644
--- a/src/transformers/models/falcon_mamba/modeling_falcon_mamba.py
+++ b/src/transformers/models/falcon_mamba/modeling_falcon_mamba.py
@@ -48,7 +48,7 @@ class FalconMambaWeightlessRMSNorm(torch.nn.Module):
     def __init__(self, hidden_size, eps: float = 1e-6):
         super().__init__()
         self.eps = eps
-        # Dummy weights that are not used (only for imitating on kernels path)
+        # Dummy weights that are not used (only for imitating the kernels path)
         self.weight = nn.Buffer(torch.ones(hidden_size, requires_grad=False), persistent=False)
 
     def _norm(self, x):
@@ -116,10 +116,7 @@ def causal_conv1d_fn(
     return out.to(hidden_states.dtype)
 
 
-@use_kernel_func_from_hub_with_fallback(
-    "mamba_inner_fn",
-    "mamba_ssm",
-)
+@use_kernel_func_from_hub_with_fallback("mamba_inner_fn", "mamba_ssm")
 def mamba_inner_fn(
     xz: torch.Tensor,
     conv1d_weight: torch.Tensor,
@@ -143,10 +140,7 @@ def mamba_inner_fn(
     return None
 
 
-@use_kernel_func_from_hub_with_fallback(
-    "selective_state_update",
-    "mamba_ssm",
-)
+@use_kernel_func_from_hub_with_fallback("selective_state_update", "mamba_ssm")
 def mamba_selective_state_update(
     state: torch.Tensor,
     hidden_states: torch.Tensor,
@@ -192,10 +186,7 @@ def mamba_selective_state_update(
     return out.to(input_dtype)
 
 
-@use_kernel_func_from_hub_with_fallback(
-    "selective_scan_fn",
-    "mamba_ssm",
-)
+@use_kernel_func_from_hub_with_fallback("selective_scan_fn", "mamba_ssm")
 def mamba_selective_scan(
     hidden_states: torch.Tensor,
     dt: torch.Tensor,
@@ -356,6 +347,7 @@ def __init__(self, config: FalconMambaConfig, layer_idx: int, initialize_mixer_w
         self.use_bias = config.use_bias
 
         self.layer_type = config.layer_types[layer_idx]
+        # These include dummy weights that are not used (only for imitating the kernels path)
         self.dt_layernorm = FalconMambaWeightlessRMSNorm(self.intermediate_size, eps=config.mixer_rms_eps)
         self.b_layernorm = FalconMambaWeightlessRMSNorm(self.ssm_state_size, eps=config.mixer_rms_eps)
         self.c_layernorm = FalconMambaWeightlessRMSNorm(self.ssm_state_size, eps=config.mixer_rms_eps)
diff --git a/src/transformers/models/falcon_mamba/modular_falcon_mamba.py b/src/transformers/models/falcon_mamba/modular_falcon_mamba.py
index b901f284ef98..87a92c25dd75 100644
--- a/src/transformers/models/falcon_mamba/modular_falcon_mamba.py
+++ b/src/transformers/models/falcon_mamba/modular_falcon_mamba.py
@@ -103,13 +103,14 @@ def layer_types(self):
 class FalconMambaWeightlessRMSNorm(NanoChatRMSNorm):
     def __init__(self, hidden_size, eps: float = 1e-6):
         super().__init__(eps)
-        # Dummy weights that are not used (only for imitating on kernels path)
+        # Dummy weights that are not used (only for imitating the kernels path)
         self.weight = nn.Buffer(torch.ones(hidden_size, requires_grad=False), persistent=False)
 
 
 class FalconMambaMixer(MambaMixer):
     def __init__(self, config: FalconMambaConfig, layer_idx: int, initialize_mixer_weights: bool = True):
         super().__init__(config, layer_idx, initialize_mixer_weights)
+        # These include dummy weights that are not used (only for imitating the kernels path)
         self.dt_layernorm = FalconMambaWeightlessRMSNorm(self.intermediate_size, eps=config.mixer_rms_eps)
         self.b_layernorm = FalconMambaWeightlessRMSNorm(self.ssm_state_size, eps=config.mixer_rms_eps)
         self.c_layernorm = FalconMambaWeightlessRMSNorm(self.ssm_state_size, eps=config.mixer_rms_eps)
diff --git a/src/transformers/models/granitemoehybrid/modeling_granitemoehybrid.py b/src/transformers/models/granitemoehybrid/modeling_granitemoehybrid.py
index 1f705369c6ac..af16a799c8a5 100644
--- a/src/transformers/models/granitemoehybrid/modeling_granitemoehybrid.py
+++ b/src/transformers/models/granitemoehybrid/modeling_granitemoehybrid.py
@@ -297,10 +297,7 @@ def causal_conv1d_fn(
     return out.to(hidden_states.dtype)
 
 
-@use_kernel_func_from_hub_with_fallback(
-    "mamba_split_conv1d_scan_combined",
-    "mamba_ssm",
-)
+@use_kernel_func_from_hub_with_fallback("mamba_split_conv1d_scan_combined", "mamba_ssm")
 def mamba2_split_conv1d_scan_combined(
     zxbcdt: torch.Tensor,
     conv1d_weight: torch.Tensor,
@@ -325,10 +322,7 @@ def mamba2_split_conv1d_scan_combined(
     return None
 
 
-@use_kernel_func_from_hub_with_fallback(
-    "selective_state_update",
-    "mamba_ssm",
-)
+@use_kernel_func_from_hub_with_fallback("selective_state_update", "mamba_ssm")
 def mamba2_selective_state_update(
     state: torch.Tensor,
     hidden_states: torch.Tensor,
@@ -390,10 +384,7 @@ def mamba2_selective_state_update(
     return out.to(hidden_states.dtype)
 
 
-@use_kernel_func_from_hub_with_fallback(
-    "mamba_chunk_scan_combined",
-    "mamba_ssm",
-)
+@use_kernel_func_from_hub_with_fallback("mamba_chunk_scan_combined", "mamba_ssm")
 def mamba2_chunk_scan(
     hidden_states: torch.Tensor,
     dt: torch.Tensor,
@@ -586,9 +577,7 @@ def forward(
         projected_states = self.in_proj(hidden_states)
 
         A = -torch.exp(self.A_log.float())
-        fused_kwargs = (
-            kwargs | {} if self.time_step_limit == (0.0, float("inf")) else kwargs | {"dt_limit": self.time_step_limit}
-        )
+        fused_kwargs = kwargs | {"dt_limit": self.time_step_limit}
         if self.training and cache_params is None:
             fused_output = mamba2_split_conv1d_scan_combined(
                 projected_states,
diff --git a/src/transformers/models/jamba/modeling_jamba.py b/src/transformers/models/jamba/modeling_jamba.py
index ea9ab6389925..af47f43cb05e 100755
--- a/src/transformers/models/jamba/modeling_jamba.py
+++ b/src/transformers/models/jamba/modeling_jamba.py
@@ -257,10 +257,7 @@ def causal_conv1d_fn(
     return out.to(hidden_states.dtype)
 
 
-@use_kernel_func_from_hub_with_fallback(
-    "mamba_inner_fn",
-    "mamba_ssm",
-)
+@use_kernel_func_from_hub_with_fallback("mamba_inner_fn", "mamba_ssm")
 def mamba_inner_fn(
     xz: torch.Tensor,
     conv1d_weight: torch.Tensor,
@@ -284,10 +281,7 @@ def mamba_inner_fn(
     return None
 
 
-@use_kernel_func_from_hub_with_fallback(
-    "selective_state_update",
-    "mamba_ssm",
-)
+@use_kernel_func_from_hub_with_fallback("selective_state_update", "mamba_ssm")
 def mamba_selective_state_update(
     state: torch.Tensor,
     hidden_states: torch.Tensor,
@@ -333,10 +327,7 @@ def mamba_selective_state_update(
     return out.to(input_dtype)
 
 
-@use_kernel_func_from_hub_with_fallback(
-    "selective_scan_fn",
-    "mamba_ssm",
-)
+@use_kernel_func_from_hub_with_fallback("selective_scan_fn", "mamba_ssm")
 def mamba_selective_scan(
     hidden_states: torch.Tensor,
     dt: torch.Tensor,
diff --git a/src/transformers/models/mamba/modeling_mamba.py b/src/transformers/models/mamba/modeling_mamba.py
index e16b1d853929..6d57a2c6e3f2 100644
--- a/src/transformers/models/mamba/modeling_mamba.py
+++ b/src/transformers/models/mamba/modeling_mamba.py
@@ -101,10 +101,7 @@ def causal_conv1d_fn(
     return out.to(hidden_states.dtype)
 
 
-@use_kernel_func_from_hub_with_fallback(
-    "mamba_inner_fn",
-    "mamba_ssm",
-)
+@use_kernel_func_from_hub_with_fallback("mamba_inner_fn", "mamba_ssm")
 def mamba_inner_fn(
     xz: torch.Tensor,
     conv1d_weight: torch.Tensor,
@@ -128,10 +125,7 @@ def mamba_inner_fn(
     return None
 
 
-@use_kernel_func_from_hub_with_fallback(
-    "selective_state_update",
-    "mamba_ssm",
-)
+@use_kernel_func_from_hub_with_fallback("selective_state_update", "mamba_ssm")
 def mamba_selective_state_update(
     state: torch.Tensor,
     hidden_states: torch.Tensor,
@@ -177,10 +171,7 @@ def mamba_selective_state_update(
     return out.to(input_dtype)
 
 
-@use_kernel_func_from_hub_with_fallback(
-    "selective_scan_fn",
-    "mamba_ssm",
-)
+@use_kernel_func_from_hub_with_fallback("selective_scan_fn", "mamba_ssm")
 def mamba_selective_scan(
     hidden_states: torch.Tensor,
     dt: torch.Tensor,
diff --git a/src/transformers/models/mamba2/modeling_mamba2.py b/src/transformers/models/mamba2/modeling_mamba2.py
index 42472d973ad3..20af9e6d5db1 100644
--- a/src/transformers/models/mamba2/modeling_mamba2.py
+++ b/src/transformers/models/mamba2/modeling_mamba2.py
@@ -163,10 +163,7 @@ def causal_conv1d_fn(
     return out.to(hidden_states.dtype)
 
 
-@use_kernel_func_from_hub_with_fallback(
-    "mamba_split_conv1d_scan_combined",
-    "mamba_ssm",
-)
+@use_kernel_func_from_hub_with_fallback("mamba_split_conv1d_scan_combined", "mamba_ssm")
 def mamba2_split_conv1d_scan_combined(
     zxbcdt: torch.Tensor,
     conv1d_weight: torch.Tensor,
@@ -191,10 +188,7 @@ def mamba2_split_conv1d_scan_combined(
     return None
 
 
-@use_kernel_func_from_hub_with_fallback(
-    "selective_state_update",
-    "mamba_ssm",
-)
+@use_kernel_func_from_hub_with_fallback("selective_state_update", "mamba_ssm")
 def mamba2_selective_state_update(
     state: torch.Tensor,
     hidden_states: torch.Tensor,
@@ -256,10 +250,7 @@ def mamba2_selective_state_update(
     return out.to(hidden_states.dtype)
 
 
-@use_kernel_func_from_hub_with_fallback(
-    "mamba_chunk_scan_combined",
-    "mamba_ssm",
-)
+@use_kernel_func_from_hub_with_fallback("mamba_chunk_scan_combined", "mamba_ssm")
 def mamba2_chunk_scan(
     hidden_states: torch.Tensor,
     dt: torch.Tensor,
@@ -466,9 +457,7 @@ def forward(
         projected_states = self.in_proj(hidden_states)
 
         A = -torch.exp(self.A_log.float())
-        fused_kwargs = (
-            kwargs | {} if self.time_step_limit == (0.0, float("inf")) else kwargs | {"dt_limit": self.time_step_limit}
-        )
+        fused_kwargs = kwargs | {"dt_limit": self.time_step_limit}
         if self.training and cache_params is None:
             fused_output = mamba2_split_conv1d_scan_combined(
                 projected_states,
diff --git a/src/transformers/models/nemotron_h/modeling_nemotron_h.py b/src/transformers/models/nemotron_h/modeling_nemotron_h.py
index d7758035d54d..b0c1f181e575 100644
--- a/src/transformers/models/nemotron_h/modeling_nemotron_h.py
+++ b/src/transformers/models/nemotron_h/modeling_nemotron_h.py
@@ -159,10 +159,7 @@ def causal_conv1d_fn(
     return out.to(hidden_states.dtype)
 
 
-@use_kernel_func_from_hub_with_fallback(
-    "mamba_split_conv1d_scan_combined",
-    "mamba_ssm",
-)
+@use_kernel_func_from_hub_with_fallback("mamba_split_conv1d_scan_combined", "mamba_ssm")
 def mamba2_split_conv1d_scan_combined(
     zxbcdt: torch.Tensor,
     conv1d_weight: torch.Tensor,
@@ -187,10 +184,7 @@ def mamba2_split_conv1d_scan_combined(
     return None
 
 
-@use_kernel_func_from_hub_with_fallback(
-    "selective_state_update",
-    "mamba_ssm",
-)
+@use_kernel_func_from_hub_with_fallback("selective_state_update", "mamba_ssm")
 def mamba2_selective_state_update(
     state: torch.Tensor,
     hidden_states: torch.Tensor,
@@ -252,10 +246,7 @@ def mamba2_selective_state_update(
     return out.to(hidden_states.dtype)
 
 
-@use_kernel_func_from_hub_with_fallback(
-    "mamba_chunk_scan_combined",
-    "mamba_ssm",
-)
+@use_kernel_func_from_hub_with_fallback("mamba_chunk_scan_combined", "mamba_ssm")
 def mamba2_chunk_scan(
     hidden_states: torch.Tensor,
     dt: torch.Tensor,
@@ -449,9 +440,7 @@ def forward(
         projected_states = self.in_proj(hidden_states)
 
         A = -torch.exp(self.A_log.float())
-        fused_kwargs = (
-            kwargs | {} if self.time_step_limit == (0.0, float("inf")) else kwargs | {"dt_limit": self.time_step_limit}
-        )
+        fused_kwargs = kwargs | {"dt_limit": self.time_step_limit}
         if self.training and cache_params is None:
             fused_output = mamba2_split_conv1d_scan_combined(
                 projected_states,
diff --git a/src/transformers/models/olmo_hybrid/modeling_olmo_hybrid.py b/src/transformers/models/olmo_hybrid/modeling_olmo_hybrid.py
index 66373b1b395e..a1c214e75853 100644
--- a/src/transformers/models/olmo_hybrid/modeling_olmo_hybrid.py
+++ b/src/transformers/models/olmo_hybrid/modeling_olmo_hybrid.py
@@ -594,7 +594,7 @@ def forward(
             recurrent_state = cache_params.layers[self.layer_idx].recurrent_states[0]
 
         # Single token decode path
-        if use_precomputed_states and seq_len == 1:
+        if use_precomputed_states and seq_len == 1 and not cache_params.layers[self.layer_idx].record_past:
             mixed_qkv = causal_conv1d_update(
                 mixed_qkv,
                 conv_state,
diff --git a/src/transformers/models/olmo_hybrid/modular_olmo_hybrid.py b/src/transformers/models/olmo_hybrid/modular_olmo_hybrid.py
index 47f3d7d4292f..e55fe21c00f3 100644
--- a/src/transformers/models/olmo_hybrid/modular_olmo_hybrid.py
+++ b/src/transformers/models/olmo_hybrid/modular_olmo_hybrid.py
@@ -368,7 +368,7 @@ def forward(
             recurrent_state = cache_params.layers[self.layer_idx].recurrent_states[0]
 
         # Single token decode path
-        if use_precomputed_states and seq_len == 1:
+        if use_precomputed_states and seq_len == 1 and not cache_params.layers[self.layer_idx].record_past:
             mixed_qkv = causal_conv1d_update(
                 mixed_qkv,
                 conv_state,
diff --git a/src/transformers/models/zamba2/modeling_zamba2.py b/src/transformers/models/zamba2/modeling_zamba2.py
index fbf9a99971e0..5869d3b1a651 100644
--- a/src/transformers/models/zamba2/modeling_zamba2.py
+++ b/src/transformers/models/zamba2/modeling_zamba2.py
@@ -449,10 +449,7 @@ def causal_conv1d_fn(
     return out.to(hidden_states.dtype)
 
 
-@use_kernel_func_from_hub_with_fallback(
-    "mamba_split_conv1d_scan_combined",
-    "mamba_ssm",
-)
+@use_kernel_func_from_hub_with_fallback("mamba_split_conv1d_scan_combined", "mamba_ssm")
 def mamba2_split_conv1d_scan_combined(
     zxbcdt: torch.Tensor,
     conv1d_weight: torch.Tensor,
@@ -477,10 +474,7 @@ def mamba2_split_conv1d_scan_combined(
     return None
 
 
-@use_kernel_func_from_hub_with_fallback(
-    "selective_state_update",
-    "mamba_ssm",
-)
+@use_kernel_func_from_hub_with_fallback("selective_state_update", "mamba_ssm")
 def mamba2_selective_state_update(
     state: torch.Tensor,
     hidden_states: torch.Tensor,
@@ -542,10 +536,7 @@ def mamba2_selective_state_update(
     return out.to(hidden_states.dtype)
 
 
-@use_kernel_func_from_hub_with_fallback(
-    "mamba_chunk_scan_combined",
-    "mamba_ssm",
-)
+@use_kernel_func_from_hub_with_fallback("mamba_chunk_scan_combined", "mamba_ssm")
 def mamba2_chunk_scan(
     hidden_states: torch.Tensor,
     dt: torch.Tensor,
@@ -738,9 +729,7 @@ def forward(
         projected_states = self.in_proj(hidden_states)
 
         A = -torch.exp(self.A_log.float())
-        fused_kwargs = (
-            kwargs | {} if self.time_step_limit == (0.0, float("inf")) else kwargs | {"dt_limit": self.time_step_limit}
-        )
+        fused_kwargs = kwargs | {"dt_limit": self.time_step_limit}
         if self.training and cache_params is None:
             fused_output = mamba2_split_conv1d_scan_combined(
                 projected_states,