diff --git a/examples/nemotron_3_nano_text_generation.py b/examples/nemotron_3_nano_text_generation.py
index 1fe73fbd..838e8bc9 100644
--- a/examples/nemotron_3_nano_text_generation.py
+++ b/examples/nemotron_3_nano_text_generation.py
@@ -242,9 +242,9 @@ def init_hybrid_states(config, dtype: np.dtype = np.float32) -> dict[str, np.nda
             states[f"past_key_values.{i}.conv_state"] = np.zeros(
                 (batch_size, conv_dim, d_conv - 1), dtype=dtype
             )
-            # ssm_state: (batch, n_heads, d_head, d_state)
+            # ssm_state: (batch, n_heads, d_state, d_head) — LinearAttention convention
             states[f"past_key_values.{i}.ssm_state"] = np.zeros(
-                (batch_size, n_heads, d_head, d_state),
+                (batch_size, n_heads, d_state, d_head),
                 dtype=dtype,
             )
         elif ltype in ("attention", "full_attention"):
@@ -593,24 +593,39 @@ def main():
         "--device",
         choices=["cpu", "cuda"],
         default="cpu",
-        help="Device for ONNX Runtime and PyTorch inference (default: %(default)s).",
+        help=(
+            "Device for inference (used for ONNX Runtime and for "
+            "HuggingFace comparison when --compare-hf is set) "
+            "(default: %(default)s)."
+        ),
+    )
+    parser.add_argument(
+        "--ep",
+        choices=["cpu", "cuda", "onnx-standard"],
+        default=None,
+        help=(
+            "Execution provider for ONNX model build. "
+            "'onnx-standard' inlines custom ops (LinearAttention, etc.) "
+            "into standard ONNX ops, runnable on any ORT version. "
+            "Defaults to matching --device."
+        ),
     )
     parser.add_argument(
         "--no-chat",
         action="store_true",
         help="Disable chat template (send raw text).",
     )
+    parser.add_argument(
+        "--ci",
+        action="store_true",
+        help="Exit with non-zero code on failure (for CI pipelines).",
+    )
     parser.add_argument(
         "--repetition-penalty",
         type=float,
         default=REPETITION_PENALTY,
         help="Repetition penalty (1.0 = none, default: %(default)s).",
     )
-    parser.add_argument(
-        "--ci",
-        action="store_true",
-        help="Exit with non-zero code on failure (for CI pipelines).",
-    )
     args = parser.parse_args()
 
     use_chat = not args.no_chat
@@ -632,7 +647,7 @@ def main():
         build_flags = {}
         if args.device == "cuda":
             build_flags["ort_cuda_grouped_rmsnorm_workaround"] = True
-        ep = "cuda" if args.device == "cuda" else "cpu"
+        ep = args.ep or ("cuda" if args.device == "cuda" else "cpu")
         print(f"Building model for {args.model!r} (dtype={args.dtype}, ep={ep}) ...")
         with override_flags(**build_flags):
             pkg = build(
diff --git a/pyproject.toml b/pyproject.toml
index 24154874..acbb2e1b 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -135,8 +135,6 @@ ignore = [
 ]
 
 [tool.ruff.lint.per-file-ignores]
-# SSM math notation uses uppercase variable names (A, B, C, D, H, N, etc.)
-"src/mobius/components/_mamba_block_chunked.py" = ["N803", "N806"]
 
 [tool.ruff.lint.flake8-tidy-imports]
 ban-relative-imports = "all"
diff --git a/src/mobius/_flags.py b/src/mobius/_flags.py
index 7f3d3693..d11b2938 100644
--- a/src/mobius/_flags.py
+++ b/src/mobius/_flags.py
@@ -55,28 +55,6 @@ def _env_bool(name: str, default: bool) -> bool:
     return default
 
 
-def _env_str(name: str, default: str, choices: tuple[str, ...]) -> str:
-    """Read a string from an environment variable.
-
-    Returns *default* if the variable is unset or has an unrecognised value.
-    For backwards compatibility, ``"1"``/``"true"``/``"yes"`` map to the
-    first choice, and ``"0"``/``"false"``/``"no"`` map to the last choice.
-    """
-    val = os.environ.get(name, "").strip().lower()
-    if not val:
-        return default
-    # Direct match against choices
-    for c in choices:
-        if val == c.lower():
-            return c
-    # Boolean-style aliases: truthy → first choice, falsy → last choice
-    if val in ("1", "true", "yes"):
-        return choices[0]
-    if val in ("0", "false", "no"):
-        return choices[-1]
-    return default
-
-
 @dataclasses.dataclass
 class _Flags:
     """Runtime feature flags singleton.
@@ -120,28 +98,6 @@ class _Flags:
     Set ``MOBIUS_ORT_CUDA_GROUPED_RMSNORM_WORKAROUND=1`` when targeting CUDA.
     """
 
-    mamba_scan: str = dataclasses.field(
-        default_factory=lambda: _env_str(
-            "MOBIUS_MAMBA_SCAN",
-            "single",
-            ("chunked_ssd", "scan", "single"),
-        )
-    )
-    """Multi-token Mamba2 forward strategy.
-
-    - ``"single"`` (default): single-token-only path (seq_len must
-      be 1).  The simplest and most debuggable mode.
-    - ``"chunked_ssd"``: chunked SSD algorithm — processes the full
-      sequence in parallel within chunks, with cross-chunk state
-      propagation.  Matches HF ``torch_forward``.
-    - ``"scan"``: ONNX Scan op that iterates token-by-token.  Supports
-      arbitrary seq_len but is sequential.
-
-    Set via ``MOBIUS_MAMBA_SCAN=chunked_ssd|scan|single``.
-    For backwards compatibility, ``1``/``true`` → ``chunked_ssd``,
-    ``0``/``false`` → ``single``.
-    """
-
 
 # Global singleton — import and use this directly.
 flags = _Flags()
diff --git a/src/mobius/components/__init__.py b/src/mobius/components/__init__.py
index 3c1e3ff9..45e52f41 100644
--- a/src/mobius/components/__init__.py
+++ b/src/mobius/components/__init__.py
@@ -52,7 +52,6 @@
     "Mistral3MultiModalProjector",
     "MLPMultiModalProjector",
     "Mamba2Block",
-    "Mamba2Scan",
     "MambaBlock",
     "MoELayer",
     "OffsetRMSNorm",
@@ -233,7 +232,6 @@
 from mobius.components._rotary_embedding import initialize_rope
 from mobius.components._ssm import (
     JambaSelectiveScan,
-    Mamba2Scan,
     SelectiveScan,
 )
 from mobius.components._vision import (
diff --git a/src/mobius/components/_mamba_block.py b/src/mobius/components/_mamba_block.py
index e5babd91..59721e2f 100644
--- a/src/mobius/components/_mamba_block.py
+++ b/src/mobius/components/_mamba_block.py
@@ -7,17 +7,11 @@
 
 - **MambaBlock** (Mamba1): standard Mamba layer for Mamba, Jamba,
   FalconMamba, etc.
-- **Mamba2BlockBase**: shared ``__init__`` and helpers for all Mamba2
-  multi-token modes.
-- **Mamba2BlockSingle**: single-token path (seq_len must be 1).
-- **Mamba2Block**: factory function that instantiates the correct
-  subclass based on ``flags.mamba_scan``.
-
-The full set of Mamba2 subclasses is:
-
-- ``Mamba2BlockSingle``  — this file
-- ``Mamba2BlockScan``    — ``_mamba_block_scan.py``
-- ``Mamba2BlockChunkedSSD`` — ``_mamba_block_chunked.py``
+- **Mamba2Block**: Mamba2 layer using ``com.microsoft.LinearAttention``
+  with ``update_rule="gated"`` for the SSM recurrence and
+  ``com.microsoft.CausalConvWithState`` for depthwise Conv1D.
+  Supports both single-token decode (T=1) and multi-token prefill
+  (T>1) in a single code path.
 
 HuggingFace reference: ``MambaMixer``, ``BambaMixer``,
 ``NemotronHMamba2Mixer``.
@@ -29,17 +23,16 @@
 from onnxscript import nn
 from onnxscript._internal import builder
 
-from mobius._flags import flags
 from mobius.components._common import INT64_MAX, Linear
 from mobius.components._rms_norm import GatedRMSNorm
-from mobius.components._ssm import Mamba2Scan, SelectiveScan
+from mobius.components._ssm import SelectiveScan
 
 
 class _DepthwiseConv1d(nn.Module):
     """Depthwise 1D convolution with optional bias.
 
     Each input channel is convolved with its own kernel (groups=channels).
-    Used for causal convolution in the Mamba block.
+    Used for causal convolution in the Mamba1 block.
     """
 
     def __init__(self, channels: int, kernel_size: int, bias: bool = True):
@@ -180,24 +173,111 @@ def forward(
 
 
 # =====================================================================
-# Mamba2 base class and subclasses
+# Mamba2 block using LinearAttention
 # =====================================================================
 
 
-class Mamba2BlockBase(nn.Module):
-    """Base class for all Mamba2 block variants.
+class _Mamba2DepthwiseConv1d(nn.Module):
+    """Depthwise 1D convolution via CausalConvWithState function op.
+
+    Wraps ``weight`` and optional ``bias`` parameters so that
+    HuggingFace weight names (``conv1d.weight``, ``conv1d.bias``)
+    automatically align with ONNX initializer names.
+
+    The ``forward()`` method calls the ``CausalConvWithState``
+    function op in the ``com.microsoft`` domain.
+    """
+
+    def __init__(self, channels: int, kernel_size: int, bias: bool = True):
+        super().__init__()
+        self.weight = nn.Parameter([channels, 1, kernel_size])
+        self.bias = nn.Parameter([channels]) if bias else None
+        self._channels = channels
 
-    Shared ``__init__`` defines the parameters (in_proj, conv1d, ssm,
-    norm, out_proj) so that all subclasses produce identical ONNX
-    weight paths.  Subclasses override ``forward()`` with the specific
-    multi-token algorithm.
+    def forward(
+        self,
+        op: builder.OpBuilder,
+        input_val: ir.Value,
+        conv_state: ir.Value,
+    ):
+        """Run CausalConvWithState function op.
+
+        Args:
+            op: ONNX op builder.
+            input_val: (B, D, T) — channels-first input.
+            conv_state: (B, D, K-1) — carry state.
+
+        Returns:
+            output: (B, D, T) — convolution output with SiLU.
+            present_state: (B, D, K-1) — updated carry state.
+        """
+        if self.bias is not None:
+            conv_bias = self.bias
+        else:
+            # Zero bias — the function requires a bias input.
+            conv_bias = op.Expand(
+                op.CastLike(op.Constant(value_float=0.0), self.weight),
+                op.Constant(value_ints=[self._channels]),
+            )
+        return op.CausalConvWithState(
+            input_val,
+            self.weight,
+            conv_bias,
+            conv_state,
+            activation="silu",
+            _domain="com.microsoft",
+            _outputs=2,
+        )
+
+
+class Mamba2Block(nn.Module):
+    """Mamba2 block using LinearAttention for the SSM recurrence.
+
+    Uses ``com.microsoft.LinearAttention`` with ``update_rule="gated"``
+    to express the Mamba2 SSD recurrence, and
+    ``com.microsoft.CausalConvWithState`` for the depthwise Conv1D.
+    Supports both single-token decode (T=1) and multi-token prefill
+    (T>1) in a single unified code path.
+
+    The Mamba2 SSD recurrence maps to LinearAttention as:
+
+    - **query** = C (readout matrix), ``(B, T, num_heads * d_state)``
+    - **key** = B (input matrix), ``(B, T, num_heads * d_state)``
+    - **value** = dt * x (discretized input), ``(B, T, num_heads * d_head)``
+    - **decay** = A * dt (log-decay), ``(B, T, num_heads)``
+
+    B and C are expanded from ``n_groups`` to ``num_heads`` before
+    passing to LinearAttention (each group's B/C is shared across
+    ``heads_per_group`` heads).
 
     Key differences from MambaBlock (Mamba1):
+
     - in_proj outputs [gate, xBC, dt] instead of [x, z]
     - Conv1D on wider xBC (conv_dim channels)
     - Multi-head SSM with grouped B/C
     - GatedRMSNorm instead of SiLU gating
     - dt direct from in_proj (no rank reduction), just bias
+
+    Args:
+        d_model: Model hidden dimension.
+        d_inner: Expanded inner dimension (``num_heads * d_head``).
+        num_heads: Number of SSM heads.
+        d_head: Per-head hidden dimension.
+        d_state: SSM state dimension.
+        n_groups: Number of B/C groups (``num_heads // n_groups``
+            heads share the same B/C).
+        chunk_size: Chunk size hint for LinearAttention (does not
+            affect correctness).
+        conv_kernel: Causal Conv1D kernel size (typically 4).
+        conv_bias: Whether conv1d has a bias.
+        proj_bias: Whether in_proj/out_proj have biases.
+        eps: RMSNorm epsilon.
+        norm_group_size: If set, GatedRMSNorm normalizes within
+            groups of this size.
+        time_step_min: Minimum dt clamp value (0 = no clamp).
+
+    HuggingFace reference: ``Mamba2Mixer``, ``BambaMixer``,
+    ``NemotronHMamba2Mixer``.
     """
 
     def __init__(
@@ -232,21 +312,17 @@ def __init__(
 
         proj_size = d_inner + self.conv_dim + num_heads
         self.in_proj = Linear(d_model, proj_size, bias=proj_bias)
-        self.conv1d = _DepthwiseConv1d(
+        self.conv1d = _Mamba2DepthwiseConv1d(
             self.conv_dim,
             conv_kernel,
             bias=conv_bias,
         )
-        # SSM parameters live under self.ssm so that ONNX weight paths
-        # stay as ``mamba.ssm.{A_log,D,dt_bias}`` — compatible with all
-        # existing preprocess_weights rename rules.
-        self.ssm = Mamba2Scan(
-            num_heads,
-            d_head,
-            d_state,
-            n_groups,
-            time_step_min=time_step_min,
-        )
+        # SSM parameters directly on this module so they appear as
+        # graph initializers. Weight paths: mamba.{A_log,D,dt_bias}
+        # matching HuggingFace naming (no extra nesting needed).
+        self.A_log = nn.Parameter([num_heads])
+        self.D = nn.Parameter([num_heads])
+        self.dt_bias = nn.Parameter([num_heads])
         self.norm = GatedRMSNorm(
             d_inner,
             eps=eps,
@@ -254,35 +330,6 @@ def __init__(
         )
         self.out_proj = Linear(d_inner, d_model, bias=proj_bias)
 
-    @staticmethod
-    def _realize_submodule(
-        parent_builder: builder.GraphBuilder,
-        submodule: nn.Module,
-    ) -> None:
-        """Register a submodule's parameters as parent-graph initializers.
-
-        When a forward path references ``self.ssm.A_log`` etc. directly
-        (without calling ``self.ssm(...)``), or when Scan body graphs
-        need parameters as implicit inputs, this helper pushes the
-        naming context and calls ``_realize()`` on every parameter so
-        they appear as graph initializers.
-        """
-        name = submodule._name or ""
-        parent_builder.push_module(name, type(submodule).__qualname__)
-        for param in submodule._parameters.values():
-            param._realize(parent_builder)
-        for child in submodule._modules.values():
-            Mamba2BlockBase._realize_submodule(parent_builder, child)
-        parent_builder.pop_module()
-
-
-class Mamba2BlockSingle(Mamba2BlockBase):
-    """Mamba2 block: single-token path (seq_len must be 1).
-
-    Uses the Mamba2Scan recurrence directly with no chunking or Scan.
-    Useful for debugging numerical issues.
-    """
-
     def forward(
         self,
         op: builder.OpBuilder,
@@ -290,112 +337,149 @@ def forward(
         conv_state: ir.Value,
         ssm_state: ir.Value,
     ):
-        """Single-token forward pass (seq_len must be 1).
+        """Forward pass for the Mamba2 block.
 
         Args:
             op: ONNX op builder.
-            hidden_states: (batch, 1, d_model)
+            hidden_states: (batch, seq_len, d_model)
             conv_state: (batch, conv_dim, conv_kernel-1)
-            ssm_state: (batch, num_heads, d_head, d_state)
+            ssm_state: (batch, num_heads, d_state, d_head) — matches
+                LinearAttention state layout (B, H, d_k, d_v).
 
         Returns:
-            output: (batch, 1, d_model)
+            output: (batch, seq_len, d_model)
             new_conv_state: (batch, conv_dim, conv_kernel-1)
-            new_ssm_state: (batch, num_heads, d_head, d_state)
+            new_ssm_state: (batch, num_heads, d_state, d_head)
         """
-        # Step 1: Input projection -> gate, xBC, dt
+        # Step 1: Input projection → gate, xBC, dt
+        # projected: (B, T, d_inner + conv_dim + num_heads)
         projected = self.in_proj(op, hidden_states)
-        gate, x_bc, dt = op.Split(
+        gate, x_bc, dt_raw = op.Split(
             projected,
             [self.d_inner, self.conv_dim, self.num_heads],
             axis=-1,
             _outputs=3,
         )
+        # gate: (B, T, d_inner) — gating signal for GatedRMSNorm
+        # x_bc: (B, T, conv_dim) — input to conv1d
+        # dt_raw: (B, T, num_heads) — raw time step
 
-        # Step 2: Causal Conv1D with state update
+        # Step 2: CausalConvWithState + SiLU
+        # Transpose to channels-first: (B, conv_dim, T)
         x_bc_t = op.Transpose(x_bc, perm=[0, 2, 1])
-        conv_input = op.Concat(conv_state, x_bc_t, axis=2)
-        new_conv_state = op.Slice(
-            conv_input,
-            starts=[1],
-            ends=[INT64_MAX],
-            axes=[2],
-        )
-        conv_out = self.conv1d(op, conv_input)
-
-        # Step 3: SiLU activation
-        conv_out = op.Mul(conv_out, op.Sigmoid(conv_out))
+        conv_out, new_conv_state = self.conv1d(op, x_bc_t, conv_state)
+        # Transpose back: (B, T, conv_dim)
         x_bc_activated = op.Transpose(conv_out, perm=[0, 2, 1])
 
-        # Step 4: Split xBC -> hidden, B, C
+        # Step 3: Split xBC → x, B, C
         gs = self.n_groups * self.d_state
-        hidden_x, b_mat, c_mat = op.Split(
+        x_hidden, b_mat, c_mat = op.Split(
             x_bc_activated,
             [self.d_inner, gs, gs],
             axis=-1,
             _outputs=3,
         )
-
-        # Squeeze seq dim for SSM: (B, 1, D) → (B, D)
-        hidden_flat = op.Squeeze(hidden_x, [1])
-        dt_flat = op.Squeeze(dt, [1])
-        b_flat = op.Squeeze(b_mat, [1])
-        c_flat = op.Squeeze(c_mat, [1])
-
-        # Step 5: Multi-head selective scan
-        y, new_ssm_state = self.ssm(
-            op,
-            hidden_flat,
-            dt_flat,
-            b_flat,
-            c_flat,
-            ssm_state,
+        # x_hidden: (B, T, d_inner = num_heads * d_head)
+        # b_mat: (B, T, n_groups * d_state)
+        # c_mat: (B, T, n_groups * d_state)
+
+        # Step 4: Compute dt and decay for LinearAttention.
+        # Upcast to fp32 for softplus/exp to match HuggingFace which
+        # computes the SSM recurrence in float32.
+        dt_raw_f32 = op.Cast(dt_raw, to=ir.DataType.FLOAT)
+        dt_bias_f32 = op.Cast(self.dt_bias, to=ir.DataType.FLOAT)
+        a_log_f32 = op.Cast(self.A_log, to=ir.DataType.FLOAT)
+
+        # dt = softplus(dt_raw + dt_bias): (B, T, num_heads)
+        dt = op.Softplus(op.Add(dt_raw_f32, dt_bias_f32))
+        if self.time_step_min > 0.0:
+            dt = op.Clip(dt, op.Constant(value_float=self.time_step_min))
+
+        # decay = A * dt in log-space: g_t where exp(g_t) is the decay
+        # A = -exp(A_log), so decay = -exp(A_log) * dt
+        a_neg = op.Neg(op.Exp(a_log_f32))  # (num_heads,)
+        decay = op.Mul(a_neg, dt)  # (B, T, num_heads) in f32
+
+        # Step 5: Prepare value = dt * x (absorb dt into input, in f32)
+        # dt: (B, T, num_heads) → (B, T, num_heads, 1)
+        dt_4d = op.Unsqueeze(dt, [-1])
+        # x: (B, T, d_inner) → (B, T, num_heads, d_head), cast to f32
+        x_f32 = op.Cast(x_hidden, to=ir.DataType.FLOAT)
+        x_4d = op.Reshape(x_f32, [0, 0, self.num_heads, self.d_head])
+        # value = dt * x: (B, T, num_heads, d_head)
+        value_4d = op.Mul(dt_4d, x_4d)
+        # Pack back: (B, T, num_heads * d_head)
+        value = op.Reshape(value_4d, [0, 0, self.num_heads * self.d_head])
+
+        # Step 6: Expand B and C from n_groups to num_heads (in f32)
+        # Each group's B/C vector is shared across heads_per_group heads.
+        b_expanded = self._expand_groups(op, op.Cast(b_mat, to=ir.DataType.FLOAT))
+        c_expanded = self._expand_groups(op, op.Cast(c_mat, to=ir.DataType.FLOAT))
+
+        # Step 7: Call LinearAttention (gated mode, all inputs in f32)
+        # query = C: (B, T, num_heads * d_state) — d_k = d_state
+        # key = B: (B, T, num_heads * d_state)
+        # value = dt*x: (B, T, num_heads * d_head) — d_v = d_head
+        # decay: (B, T, num_heads) — per-head scalar in log-space
+        # state: (B, num_heads, d_state, d_head) = (B, H, d_k, d_v)
+        ssm_state_f32 = op.Cast(ssm_state, to=ir.DataType.FLOAT)
+        la_output, new_ssm_state = op.LinearAttention(
+            c_expanded,
+            b_expanded,
+            value,
+            ssm_state_f32,
+            decay,
+            scale=1.0,
+            q_num_heads=self.num_heads,
+            kv_num_heads=self.num_heads,
+            update_rule="gated",
+            _domain="com.microsoft",
+            _outputs=2,
         )
+        # la_output: (B, T, num_heads * d_head) in f32
+        # new_ssm_state: (B, num_heads, d_state, d_head) in f32
+        # Cast back to model dtype for downstream ops and state output
+        la_output = op.CastLike(la_output, hidden_states)
+        new_ssm_state = op.CastLike(new_ssm_state, ssm_state)
+
+        # Step 8: D skip connection — y += D * x (per-head broadcast)
+        # D: (num_heads,) → (1, 1, num_heads, 1) for broadcast
+        # x_4d is still f32; cast D to f32 for the multiply, result cast
+        # back via la_output's dtype.
+        d_f32 = op.Cast(self.D, to=ir.DataType.FLOAT)
+        d_4d = op.Reshape(d_f32, [1, 1, self.num_heads, 1])
+        d_skip = op.CastLike(
+            op.Reshape(
+                op.Mul(d_4d, x_4d),
+                [0, 0, self.num_heads * self.d_head],
+            ),
+            hidden_states,
+        )
+        y = op.Add(la_output, d_skip)
 
-        # Step 6: Gated RMSNorm
-        gate_flat = op.Squeeze(gate, [1])
-        y_normed = self.norm(op, y, gate_flat)
-
-        # Restore seq dim: (B, d_inner) → (B, 1, d_inner)
-        y_3d = op.Unsqueeze(y_normed, [1])
+        # Step 9: GatedRMSNorm
+        y_normed = self.norm(op, y, gate)
 
-        # Step 7: Output projection
-        output = self.out_proj(op, y_3d)
+        # Step 10: Output projection
+        output = self.out_proj(op, y_normed)
 
         return output, new_conv_state, new_ssm_state
 
+    def _expand_groups(self, op: builder.OpBuilder, x: ir.Value) -> ir.Value:
+        """Expand grouped B or C from n_groups to num_heads.
 
-# =====================================================================
-# Factory function
-# =====================================================================
-
-
-def Mamba2Block(*args, **kwargs) -> Mamba2BlockBase:  # noqa: N802
-    """Instantiate the Mamba2 block variant for the current flag.
-
-    Reads ``flags.mamba_scan`` at construction time and returns the
-    matching subclass.  Callers use this exactly like the old
-    ``Mamba2Block`` class — no API change.
-
-    Available modes:
-
-    - ``"single"`` → ``Mamba2BlockSingle`` (this file)
-    - ``"scan"`` → ``Mamba2BlockScan`` (``_mamba_block_scan.py``)
-    - ``"chunked_ssd"`` → ``Mamba2BlockChunkedSSD``
-      (``_mamba_block_chunked.py``)
-    """
-    mode = flags.mamba_scan
-    if mode == "single":
-        return Mamba2BlockSingle(*args, **kwargs)
-    if mode == "scan":
-        from mobius.components._mamba_block_scan import Mamba2BlockScan
-
-        return Mamba2BlockScan(*args, **kwargs)
-    if mode == "chunked_ssd":
-        from mobius.components._mamba_block_chunked import (
-            Mamba2BlockChunkedSSD,
-        )
+        Args:
+            x: (B, T, n_groups * d_state)
 
-        return Mamba2BlockChunkedSSD(*args, **kwargs)
-    msg = f"Unknown mamba_scan mode {mode!r}. Expected 'single', 'scan', or 'chunked_ssd'."
-    raise ValueError(msg)
+        Returns:
+            (B, T, num_heads * d_state) with each group's d_state
+            vector replicated ``heads_per_group`` times.
+        """
+        if self.n_groups == self.num_heads:
+            return x  # No expansion needed
+        # (B, T, n_groups, 1, d_state)
+        x_5d = op.Reshape(x, [0, 0, self.n_groups, 1, self.d_state])
+        # Expand → (B, T, n_groups, heads_per_group, d_state)
+        x_expanded = op.Expand(x_5d, [1, 1, 1, self.heads_per_group, 1])
+        # Flatten → (B, T, num_heads * d_state)
+        return op.Reshape(x_expanded, [0, 0, self.num_heads * self.d_state])
diff --git a/src/mobius/components/_mamba_block_chunked.py b/src/mobius/components/_mamba_block_chunked.py
deleted file mode 100644
index 474a05bb..00000000
--- a/src/mobius/components/_mamba_block_chunked.py
+++ /dev/null
@@ -1,608 +0,0 @@
-# Copyright (c) ONNX Project Contributors
-# SPDX-License-Identifier: Apache-2.0
-
-"""Mamba2 chunked SSD (Structured State Space Duality) implementation.
-
-Processes all tokens in parallel within chunks of ``chunk_size``, then
-propagates SSM state across chunk boundaries.  Matches HuggingFace's
-``torch_forward`` multi-token path.
-
-This module contains the ``Mamba2BlockChunkedSSD`` subclass and its
-helper functions (``_segment_sum``, ``_segment_sum_dynamic``).
-"""
-
-from __future__ import annotations
-
-import numpy as np
-import onnx_ir as ir
-from onnxscript._internal import builder
-
-from mobius.components._common import INT64_MAX
-from mobius.components._mamba_block import Mamba2BlockBase
-
-# -----------------------------------------------------------------------
-# Chunked SSD helpers — ONNX equivalents of the PyTorch helper functions
-# -----------------------------------------------------------------------
-
-
-def _segment_sum(
-    op: builder.OpBuilder,
-    x: ir.Value,
-    chunk_size: int,
-):
-    """Stable segment sum via cumulative sums and masking.
-
-    Input:  x with last dim = chunk_size  (..., chunk_size)
-    Output: (..., chunk_size, chunk_size) lower-triangular cumsum.
-
-    Matches HF's ``segment_sum()`` function.  ``chunk_size`` must be a
-    compile-time constant so we can build the fixed triangular masks.
-    """
-    # Expand: (..., chunk_size) → (..., chunk_size, chunk_size)
-    x_expanded = op.Unsqueeze(x, [-1])
-    x_tiled = op.Expand(
-        x_expanded,
-        op.Concat(
-            op.Shape(x),
-            op.Constant(value_ints=[chunk_size]),
-            axis=0,
-        ),
-    )
-
-    # Strict lower-triangular mask (diagonal=-1)
-    mask_strict = op.Constant(
-        value=ir.tensor(
-            np.tril(
-                np.ones((chunk_size, chunk_size), dtype=np.float32),
-                k=-1,
-            )
-        )
-    )
-    x_masked = op.Mul(x_tiled, mask_strict)
-
-    # CumSum along the second-to-last axis
-    cumsum = op.CumSum(x_masked, op.Constant(value_int=-2))
-
-    # Full lower-triangular mask (including diagonal)
-    mask_full_bool = op.Constant(
-        value=ir.tensor(
-            np.tril(
-                np.ones((chunk_size, chunk_size), dtype=np.bool_),
-                k=0,
-            )
-        )
-    )
-    neg_inf = op.Constant(
-        value=ir.tensor(
-            np.full(
-                (chunk_size, chunk_size),
-                float("-inf"),
-                dtype=np.float32,
-            )
-        )
-    )
-    result = op.Where(mask_full_bool, cumsum, neg_inf)
-    return result
-
-
-def _segment_sum_dynamic(op: builder.OpBuilder, x: ir.Value):
-    """Segment sum for dynamic-sized last dimension.
-
-    Same logic as ``_segment_sum`` but uses ONNX ops (``Trilu``) to
-    build masks dynamically, since nc+1 (number of chunks + 1) is not
-    known at graph-build time.
-
-    Input:  x with shape (..., K) where K is dynamic.
-    Output: (..., K, K) lower-triangular cumsum.
-    """
-    K = op.Shape(x, start=-1, end=None)  # [K] as 1-d tensor
-
-    # Expand: (..., K) → (..., K, K)
-    x_expanded = op.Unsqueeze(x, [-1])
-    new_shape = op.Concat(op.Shape(x), K, axis=0)
-    x_tiled = op.Expand(x_expanded, new_shape)
-
-    # Build KxK ones, then lower-triangular masks via Trilu
-    KK_shape = op.Concat(K, K, axis=0)
-    ones_KK = op.ConstantOfShape(
-        KK_shape,
-        value=ir.tensor(np.array([1.0], dtype=np.float32)),
-    )
-    # Strict lower triangle (diagonal=-1)
-    lower_strict = op.Trilu(
-        ones_KK,
-        op.Constant(value_int=-1),
-        upper=0,
-    )
-    x_masked = op.Mul(x_tiled, lower_strict)
-
-    # CumSum along second-to-last axis
-    cumsum = op.CumSum(x_masked, op.Constant(value_int=-2))
-
-    # Full lower triangle (including diagonal)
-    lower_full = op.Trilu(ones_KK, upper=0)
-    lower_bool = op.Cast(lower_full, to=ir.DataType.BOOL)
-
-    # Fill non-lower-triangle with -inf
-    neg_inf_KK = op.ConstantOfShape(
-        KK_shape,
-        value=ir.tensor(np.array([float("-inf")], dtype=np.float32)),
-    )
-    result = op.Where(lower_bool, cumsum, neg_inf_KK)
-    return result
-
-
-class Mamba2BlockChunkedSSD(Mamba2BlockBase):
-    """Mamba2 block using the chunked SSD algorithm.
-
-    Processes the full sequence in parallel within chunks of
-    ``chunk_size`` tokens, then propagates SSM state across chunk
-    boundaries.  Matches HF's ``torch_forward`` multi-token path.
-
-    Note: seq_len must be divisible by chunk_size (the task layer
-    is responsible for padding if necessary).
-    """
-
-    def _chunked_ssd(
-        self,
-        op: builder.OpBuilder,
-        x: ir.Value,
-        dt: ir.Value,
-        B_mat: ir.Value,
-        C_mat: ir.Value,
-        ssm_state_in: ir.Value,
-    ):
-        """Chunked SSD computation in ONNX ops.
-
-        Implements the "ssd naive implementation without einsums" from
-        HuggingFace's ``torch_forward``, translated to ONNX operations.
-
-        All 5 stages of the SSD algorithm:
-          1. Intra-chunk diagonal blocks (attention-like within each
-             chunk).
-          2. Inter-chunk state computation (B terms: how each chunk
-             contributes to the running state).
-          3. Inter-chunk SSM recurrence (A terms: decay across chunk
-             boundaries).
-          4. State-to-output per chunk (C terms: readout from
-             propagated state).
-          5. Combine intra-chunk and inter-chunk contributions.
-
-        Args:
-            op: ONNX op builder.
-            x: (B, T, H, D) float32 — activated hidden (discretised).
-            dt: (B, T, H) float32 — softplus(dt_raw + dt_bias).
-            B_mat: (B, T, H, N) float32.
-            C_mat: (B, T, H, N) float32.
-            ssm_state_in: (B, H, D, N) — carry state.
-
-        Returns:
-            y: (B, T, H, D) float32 — output.
-            new_ssm_state: (B, H, D, N).
-        """
-        CS = self.chunk_size
-        H = self.num_heads
-        D = self.d_head
-        N = self.d_state
-
-        # --- D residual: D[..., None] * x (before discretisation) ---
-        D_param = op.Cast(self.ssm.D, to=ir.DataType.FLOAT)
-        D_4d = op.Unsqueeze(D_param, [0, 1, 3])  # (1, 1, H, 1)
-        D_residual = op.Mul(D_4d, x)  # (B, T, H, D)
-
-        # --- Discretise ---
-        # x = x * dt[..., None]: (B, T, H, D) * (B, T, H, 1)
-        dt_4d = op.Unsqueeze(dt, [-1])
-        x_disc = op.Mul(x, dt_4d)
-
-        # A = -exp(A_log) in float32
-        A_neg = op.Neg(op.Exp(op.Cast(self.ssm.A_log, to=ir.DataType.FLOAT)))
-        # A_dt = A * dt: (H,) * (B, T, H) → (B, T, H)
-        A_2d = op.Unsqueeze(A_neg, [0, 1])
-        A_dt = op.Mul(A_2d, dt)
-
-        # --- Reshape into chunks ---
-        # Requires T divisible by CS (caller pads).
-        x_chunked = op.Reshape(
-            x_disc,
-            op.Constant(value_ints=[0, -1, CS, H, D]),
-        )  # (B, nc, CS, H, D)
-        A_chunked = op.Reshape(
-            A_dt,
-            op.Constant(value_ints=[0, -1, CS, H]),
-        )  # (B, nc, CS, H)
-        B_chunked = op.Reshape(
-            B_mat,
-            op.Constant(value_ints=[0, -1, CS, H, N]),
-        )
-        C_chunked = op.Reshape(
-            C_mat,
-            op.Constant(value_ints=[0, -1, CS, H, N]),
-        )
-
-        # A_cumsum: cumsum of A within each chunk
-        # (B, nc, CS, H) → permute to (B, H, nc, CS) for cumsum
-        A_perm = op.Transpose(A_chunked, perm=[0, 3, 1, 2])
-        A_cumsum = op.CumSum(A_perm, op.Constant(value_int=-1))
-
-        # =============================================
-        # 1. Intra-chunk (diagonal blocks)
-        # =============================================
-        # L = exp(segment_sum(A)): causal decay within chunk
-        L = op.Exp(_segment_sum(op, A_perm, CS))  # (B,H,nc,CS,CS)
-
-        # G = sum_n(C[l,n] * B[s,n]) — contraction over state dim
-        C_exp = op.Unsqueeze(C_chunked, [3])  # (B,nc,CS,1,H,N)
-        B_exp = op.Unsqueeze(B_chunked, [2])  # (B,nc,1,CS,H,N)
-        G = op.ReduceSum(
-            op.Mul(C_exp, B_exp),
-            [-1],
-            keepdims=False,
-        )  # (B,nc,CS,CS,H)
-
-        # M = G * L (permuted to match G layout)
-        L_perm = op.Transpose(L, perm=[0, 2, 3, 4, 1])
-        M = op.Mul(G, L_perm)  # (B, nc, CS, CS, H)
-
-        # Y_diag = (M[...,None] * x_chunked[:,None]).sum(dim=3)
-        M_exp = op.Unsqueeze(M, [-1])  # (B,nc,l,s,H,1)
-        x_exp = op.Unsqueeze(x_chunked, [2])  # (B,nc,1,s,H,D)
-        Y_diag = op.ReduceSum(
-            op.Mul(M_exp, x_exp),
-            [3],
-            keepdims=False,
-        )  # (B, nc, CS, H, D)
-
-        # =============================================
-        # 2. Inter-chunk state computation (B terms)
-        # =============================================
-        # decay_states = exp(A_last - A_cumsum)
-        A_last = op.Slice(
-            A_cumsum,
-            starts=[-1],
-            ends=[INT64_MAX],
-            axes=[-1],
-        )  # (B, H, nc, 1)
-        decay_states = op.Exp(op.Sub(A_last, A_cumsum))
-
-        # B_decay = B * decay_states (permuted)
-        decay_perm = op.Transpose(
-            decay_states,
-            perm=[0, 2, 3, 1],
-        )  # (B, nc, CS, H)
-        decay_exp = op.Unsqueeze(decay_perm, [-1])
-        B_decay = op.Mul(B_chunked, decay_exp)  # (B, nc, CS, H, N)
-
-        # states = sum_over_chunk(B_decay * x_disc)
-        B_decay_exp = op.Unsqueeze(B_decay, [-2])  # (B,nc,CS,H,1,N)
-        x_disc_exp = op.Unsqueeze(x_chunked, [-1])  # (B,nc,CS,H,D,1)
-        states = op.ReduceSum(
-            op.Mul(B_decay_exp, x_disc_exp),
-            [2],
-            keepdims=False,
-        )  # (B, nc, H, D, N)
-
-        # =============================================
-        # 3. Inter-chunk SSM recurrence (A terms)
-        # =============================================
-        # Prepend previous state along chunk dim
-        prev = op.Unsqueeze(
-            op.Cast(ssm_state_in, to=ir.DataType.FLOAT),
-            [1],
-        )  # (B, 1, H, D, N)
-        states_cat = op.Concat(
-            prev,
-            states,
-            axis=1,
-        )  # (B, nc+1, H, D, N)
-
-        # decay_chunk = exp(segment_sum(pad(A_ends, (1,0))))
-        A_ends = op.Squeeze(
-            op.Slice(
-                A_cumsum,
-                starts=[-1],
-                ends=[INT64_MAX],
-                axes=[-1],
-            ),
-            [-1],
-        )  # (B, H, nc)
-        A_ends_padded = op.Pad(
-            A_ends,
-            op.Constant(value_ints=[0, 0, 1, 0, 0, 0]),
-            op.Constant(value_float=0.0),
-        )  # (B, H, nc+1)
-        decay_chunk = op.Exp(
-            _segment_sum_dynamic(op, A_ends_padded),
-        )  # (B, H, nc+1, nc+1)
-
-        # Propagate state across chunks
-        decay_chunk_t = op.Transpose(
-            decay_chunk,
-            perm=[0, 2, 3, 1],
-        )  # (B, nc+1, nc+1, H)
-        decay_exp2 = op.Unsqueeze(decay_chunk_t, [-1, -2])
-        states_exp = op.Unsqueeze(states_cat, [1])
-        new_states = op.ReduceSum(
-            op.Mul(decay_exp2, states_exp),
-            [2],
-            keepdims=False,
-        )  # (B, nc+1, H, D, N)
-
-        # Split: first nc for output, last for carry
-        states_out = op.Slice(
-            new_states,
-            starts=[0],
-            ends=[-1],
-            axes=[1],
-        )  # (B, nc, H, D, N)
-        new_ssm_state = op.Squeeze(
-            op.Slice(
-                new_states,
-                starts=[-1],
-                ends=[INT64_MAX],
-                axes=[1],
-            ),
-            [1],
-        )  # (B, H, D, N)
-
-        # =============================================
-        # 4. State → output per chunk (C terms)
-        # =============================================
-        state_decay_out = op.Exp(A_cumsum)  # (B, H, nc, CS)
-
-        C_exp2 = op.Unsqueeze(C_chunked, [-2])  # (B,nc,CS,H,1,N)
-        states_exp2 = op.Unsqueeze(states_out, [2])  # (B,nc,1,H,D,N)
-        C_states_sum = op.ReduceSum(
-            op.Mul(C_exp2, states_exp2),
-            [-1],
-            keepdims=False,
-        )  # (B, nc, CS, H, D)
-
-        sdo_perm = op.Transpose(
-            state_decay_out,
-            perm=[0, 2, 3, 1],
-        )  # (B, nc, CS, H)
-        sdo_exp = op.Unsqueeze(sdo_perm, [-1])
-        Y_off = op.Mul(C_states_sum, sdo_exp)  # (B, nc, CS, H, D)
-
-        # =============================================
-        # 5. Combine
-        # =============================================
-        y = op.Add(Y_diag, Y_off)  # (B, nc, CS, H, D)
-
-        # Reshape back to (B, T, H, D)
-        bt_shape = op.Shape(x, start=0, end=2)
-        y_shape = op.Concat(
-            bt_shape,
-            op.Constant(value_ints=[H, D]),
-            axis=0,
-        )
-        y_reshaped = op.Reshape(y, y_shape)
-
-        y_out = op.Add(y_reshaped, D_residual)  # (B, T, H, D)
-
-        return y_out, new_ssm_state
-
-    def forward(
-        self,
-        op: builder.OpBuilder,
-        hidden_states: ir.Value,
-        conv_state: ir.Value,
-        ssm_state: ir.Value,
-    ):
-        """Multi-token forward using the chunked SSD algorithm.
-
-        Processes the full sequence in parallel within chunks of
-        ``chunk_size`` tokens, then propagates SSM state across chunk
-        boundaries.  Matching HF's ``torch_forward`` multi-token path.
-
-        Note: seq_len must be divisible by chunk_size (the task layer
-        is responsible for padding if necessary).
-        """
-        # Realize SSM parameters so they are visible as graph
-        # initializers when referenced directly (not via self.ssm()).
-        self._realize_submodule(op.builder, self.ssm)
-        H = self.num_heads
-        D = self.d_head
-        N = self.d_state
-
-        # Dtype-matched zero for Pad ops (avoids f32/f16 mismatch)
-        pad_zero = op.CastLike(op.Constant(value_float=0.0), hidden_states)
-
-        # Step 1: Batch-project all tokens at once
-        projected = self.in_proj(op, hidden_states)  # (B, T, proj)
-        gate, x_bc, dt_raw = op.Split(
-            projected,
-            [self.d_inner, self.conv_dim, self.num_heads],
-            axis=-1,
-            _outputs=3,
-        )
-
-        # Step 2: Causal Conv1D over full sequence
-        x_bc_t = op.Transpose(x_bc, perm=[0, 2, 1])
-        # Pad left with K-1 zeros for causal convolution
-        padded = op.Pad(
-            x_bc_t,
-            op.Constant(
-                value_ints=[0, 0, self.conv_kernel - 1, 0, 0, 0],
-            ),
-            pad_zero,
-        )  # (B, conv_dim, T+K-1)
-        conv_out_raw = self.conv1d(op, padded)
-        # SiLU activation
-        conv_activated = op.Mul(
-            conv_out_raw,
-            op.Sigmoid(conv_out_raw),
-        )
-        hidden_B_C = op.Transpose(conv_activated, perm=[0, 2, 1])
-
-        # Extract new conv_state: last K-1 positions from the combined
-        # old state + new tokens.  For T >= K-1 we could just slice
-        # x_bc_t, but for T < K-1 (e.g. T=1 during decode) we need
-        # to include history from the old conv_state.
-        conv_combined = op.Concat(
-            conv_state,
-            x_bc_t,
-            axis=2,
-        )  # (B, conv_dim, K-1+T)
-        new_conv_state = op.Slice(
-            conv_combined,
-            starts=[-(self.conv_kernel - 1)],
-            ends=[INT64_MAX],
-            axes=[2],
-        )  # (B, conv_dim, K-1)
-
-        # Step 3: Split xBC → x, B, C
-        gs = self.n_groups * self.d_state
-        x_split, B_split, C_split = op.Split(
-            hidden_B_C,
-            [self.d_inner, gs, gs],
-            axis=-1,
-            _outputs=3,
-        )
-
-        # Step 4: Prepare SSM inputs
-        dt_bias = op.Cast(self.ssm.dt_bias, to=ir.DataType.FLOAT)
-        dt_bias_3d = op.Unsqueeze(dt_bias, [0, 1])
-        dt = op.Softplus(
-            op.Add(
-                op.Cast(dt_raw, to=ir.DataType.FLOAT),
-                dt_bias_3d,
-            )
-        )
-        # Clamp dt to time_step_min (matches HF torch.clamp(dt, min=...))
-        if self.time_step_min > 0.0:
-            dt = op.Clip(dt, op.Constant(value_float=self.time_step_min))
-
-        # Reshape x to (B, T, H, D)
-        x_4d = op.Reshape(
-            op.Cast(x_split, to=ir.DataType.FLOAT),
-            op.Constant(value_ints=[0, 0, H, D]),
-        )
-        # Reshape B, C to (B, T, n_groups, N) then expand to heads
-        B_grouped = op.Reshape(
-            op.Cast(B_split, to=ir.DataType.FLOAT),
-            op.Constant(value_ints=[0, 0, self.n_groups, N]),
-        )
-        C_grouped = op.Reshape(
-            op.Cast(C_split, to=ir.DataType.FLOAT),
-            op.Constant(value_ints=[0, 0, self.n_groups, N]),
-        )
-        # Expand groups → heads
-        B_exp = op.Unsqueeze(B_grouped, [3])
-        B_expand = op.Expand(
-            B_exp,
-            op.Concat(
-                op.Shape(B_exp, start=0, end=3),
-                op.Constant(value_ints=[self.heads_per_group]),
-                op.Shape(B_exp, start=4, end=5),
-                axis=0,
-            ),
-        )
-        B_heads = op.Reshape(
-            B_expand,
-            op.Constant(value_ints=[0, 0, H, N]),
-        )
-
-        C_exp = op.Unsqueeze(C_grouped, [3])
-        C_expand = op.Expand(
-            C_exp,
-            op.Concat(
-                op.Shape(C_exp, start=0, end=3),
-                op.Constant(value_ints=[self.heads_per_group]),
-                op.Shape(C_exp, start=4, end=5),
-                axis=0,
-            ),
-        )
-        C_heads = op.Reshape(
-            C_expand,
-            op.Constant(value_ints=[0, 0, H, N]),
-        )
-
-        # Step 5: Pad seq_len to a multiple of chunk_size
-        # pad_size = (CS - T % CS) % CS
-        CS_c = op.Constant(value_int=self.chunk_size)
-        T_val = op.Shape(x_4d, start=1, end=2)  # [T]
-        T_scalar = op.Squeeze(T_val, [0])
-        pad_size = op.Mod(
-            op.Sub(CS_c, op.Mod(T_scalar, CS_c)),
-            CS_c,
-        )
-        # Pad along dim=1: pads = [0, 0, 0, pad_size, 0...0]
-        # For 4D: (B, T, H, D) → pad format [d0b,d1b,d2b,d3b,
-        #                                     d0e,d1e,d2e,d3e]
-        pad_size_1d = op.Reshape(pad_size, op.Constant(value_ints=[1]))
-        zero_1d = op.Constant(value_ints=[0])
-        pads_4d = op.Concat(
-            zero_1d,
-            zero_1d,
-            zero_1d,
-            zero_1d,  # begins
-            zero_1d,
-            pad_size_1d,
-            zero_1d,
-            zero_1d,  # ends
-            axis=0,
-        )
-        pads_3d = op.Concat(
-            zero_1d,
-            zero_1d,
-            zero_1d,
-            zero_1d,
-            pad_size_1d,
-            zero_1d,
-            axis=0,
-        )
-        pad_zero_f32 = op.Constant(value_float=0.0)
-        x_4d = op.Pad(x_4d, pads_4d, pad_zero_f32)
-        dt = op.Pad(dt, pads_3d, pad_zero_f32)
-        B_heads = op.Pad(B_heads, pads_4d, pad_zero_f32)
-        C_heads = op.Pad(C_heads, pads_4d, pad_zero_f32)
-
-        # Step 6: Chunked SSD
-        y, new_ssm_state = self._chunked_ssd(
-            op,
-            x_4d,
-            dt,
-            B_heads,
-            C_heads,
-            ssm_state,
-        )
-
-        # Trim padding: (B, T_padded, H, D) → (B, T, H, D)
-        y = op.Slice(
-            y,
-            starts=[0],
-            ends=T_val,
-            axes=[1],
-        )
-
-        # Step 7: GatedRMSNorm (expects 2D — flatten B*T)
-        y_flat = op.Reshape(
-            y,
-            op.Constant(value_ints=[0, 0, self.d_inner]),
-        )
-        y_2d = op.Reshape(
-            y_flat,
-            op.Constant(value_ints=[-1, self.d_inner]),
-        )
-        gate_2d = op.Reshape(
-            gate,
-            op.Constant(value_ints=[-1, self.d_inner]),
-        )
-        y_normed = self.norm(op, y_2d, gate_2d)
-
-        # Reshape back to (B, T, d_inner)
-        bt_shape = op.Shape(hidden_states, start=0, end=2)
-        y_shape = op.Concat(
-            bt_shape,
-            op.Constant(value_ints=[self.d_inner]),
-            axis=0,
-        )
-        y_3d = op.Reshape(y_normed, y_shape)
-
-        # Step 7: Output projection — cast back to input dtype
-        y_proj = op.CastLike(y_3d, hidden_states)
-        output = self.out_proj(op, y_proj)
-
-        return output, new_conv_state, new_ssm_state
diff --git a/src/mobius/components/_mamba_block_scan.py b/src/mobius/components/_mamba_block_scan.py
deleted file mode 100644
index 1b1b6d8f..00000000
--- a/src/mobius/components/_mamba_block_scan.py
+++ /dev/null
@@ -1,159 +0,0 @@
-# Copyright (c) ONNX Project Contributors
-# SPDX-License-Identifier: Apache-2.0
-
-"""Mamba2 ONNX Scan-based multi-token implementation.
-
-Uses an ONNX Scan op to iterate over the sequence token-by-token,
-carrying conv and SSM states across tokens.  Supports arbitrary
-seq_len but is sequential (no intra-sequence parallelism).
-
-This module contains the ``Mamba2BlockScan`` subclass.
-"""
-
-from __future__ import annotations
-
-import onnx_ir as ir
-from onnxscript._internal import builder
-from onnxscript.onnx_types import FLOAT
-
-from mobius.components._common import INT64_MAX
-from mobius.components._mamba_block import Mamba2BlockBase
-
-
-class Mamba2BlockScan(Mamba2BlockBase):
-    """Mamba2 block using ONNX Scan (token-by-token iteration).
-
-    Multi-token path that supports arbitrary seq_len but is sequential.
-    The Scan body performs one conv state update + SSM step per token.
-    """
-
-    def _scan_body(
-        self,
-        op: builder.OpBuilder,
-        conv_state: ir.Value,
-        ssm_state: ir.Value,
-        xbc_t: ir.Value,
-        dt_t: ir.Value,
-    ):
-        """Scan body: per-token conv state update + SSM step.
-
-        Carry states (updated each iteration):
-            conv_state: (B, conv_dim, conv_kernel-1)
-            ssm_state:  (B, num_heads, d_head, d_state)
-
-        Scan inputs (per-token, sliced along axis 1):
-            xbc_t: (B, conv_dim) -- projected xBC for this token
-            dt_t:  (B, num_heads) -- time step for this token
-
-        Scan outputs (per-token, stacked along axis 1):
-            y_t: (B, num_heads * d_head) -- SSM output before norm
-        """
-        # --- Conv state update (shift register) ---
-        xbc_3d = op.Unsqueeze(xbc_t, [2])
-        conv_cat = op.Concat(conv_state, xbc_3d, axis=2)
-        new_conv_state = op.Slice(
-            conv_cat,
-            starts=[1],
-            ends=[INT64_MAX],
-            axes=[2],
-        )
-
-        # --- Apply conv1d (params already realized as implicit inputs) ---
-        conv_out = self.conv1d(op, conv_cat)
-
-        # SiLU activation, then squeeze: (B, conv_dim, 1) -> (B, conv_dim)
-        conv_out = op.Mul(conv_out, op.Sigmoid(conv_out))
-        x_bc_act = op.Squeeze(conv_out, [2])
-
-        # --- Split xBC -> hidden_x, B, C ---
-        gs = self.n_groups * self.d_state
-        hidden_x, b_mat, c_mat = op.Split(
-            x_bc_act,
-            [self.d_inner, gs, gs],
-            axis=-1,
-            _outputs=3,
-        )
-
-        # --- SSM step (params already realized as implicit inputs) ---
-        y_flat, new_ssm = self.ssm(
-            op,
-            hidden_x,
-            dt_t,
-            b_mat,
-            c_mat,
-            ssm_state,
-        )
-
-        return new_conv_state, new_ssm, y_flat
-
-    def forward(
-        self,
-        op: builder.OpBuilder,
-        hidden_states: ir.Value,
-        conv_state: ir.Value,
-        ssm_state: ir.Value,
-    ):
-        """Multi-token forward via ONNX Scan (token-by-token iteration).
-
-        Uses an ONNX Scan op to iterate over the sequence, carrying conv
-        and SSM states across tokens.  Supports arbitrary seq_len but is
-        sequential (no intra-sequence parallelism).
-        """
-        # Realize conv1d and ssm parameters in the parent graph so they
-        # are visible as implicit inputs to the Scan body.
-        parent_builder = op.builder
-        self._realize_submodule(parent_builder, self.conv1d)
-        self._realize_submodule(parent_builder, self.ssm)
-
-        # Step 1: Batch-project all tokens at once
-        projected = self.in_proj(op, hidden_states)  # (B, T, proj_size)
-        gate, x_bc, dt = op.Split(
-            projected,
-            [self.d_inner, self.conv_dim, self.num_heads],
-            axis=-1,
-            _outputs=3,
-        )
-
-        # Step 2: Scan over axis 1 (time)
-        body = parent_builder.subgraph(
-            self._scan_body,
-            inputs={
-                "conv_state": FLOAT[...],
-                "ssm_state": FLOAT[...],
-                "xbc_t": FLOAT[...],
-                "dt_t": FLOAT[...],
-            },
-            outputs={
-                "new_conv_state": FLOAT[...],
-                "new_ssm_state": FLOAT[...],
-                "y_t": FLOAT[...],
-            },
-            name="mamba2_recurrence",
-        )
-        new_conv_state, new_ssm_state, y_all = op.Scan(
-            conv_state,
-            ssm_state,
-            x_bc,
-            dt,
-            body=body,
-            num_scan_inputs=2,
-            scan_input_axes=[1, 1],
-            scan_output_axes=[1],
-            _outputs=3,
-        )
-        # y_all: (B, T, d_inner)
-
-        # Step 3: GatedRMSNorm (expects 2D -- flatten B*T)
-        y_flat = op.Reshape(y_all, [-1, self.d_inner])
-        gate_flat = op.Reshape(gate, [-1, self.d_inner])
-        y_normed = self.norm(op, y_flat, gate_flat)
-
-        # Reshape back to (B, T, d_inner)
-        bt_shape = op.Shape(hidden_states, start=0, end=2)
-        y_shape = op.Concat(bt_shape, [self.d_inner], axis=0)
-        y_3d = op.Reshape(y_normed, y_shape)
-
-        # Step 4: Batch output projection
-        output = self.out_proj(op, y_3d)
-
-        return output, new_conv_state, new_ssm_state
diff --git a/src/mobius/components/_rms_norm.py b/src/mobius/components/_rms_norm.py
index bae405dc..95eb1f00 100644
--- a/src/mobius/components/_rms_norm.py
+++ b/src/mobius/components/_rms_norm.py
@@ -81,16 +81,18 @@ def forward(self, op: builder.OpBuilder, hidden_states: ir.Value, gate: ir.Value
         gated = op.Mul(h_f32, gate_activated)
 
         if self.group_size is not None and self.group_size < self.hidden_size:
-            # Grouped RMSNorm: reshape to (batch, n_groups, group_size),
-            # normalize within each group, then reshape back.
+            # Grouped RMSNorm: normalize within each group of size group_size.
+            # Input may be 2D (B, H) or 3D (B, T, H). Flatten leading dims
+            # so that Reshape always sees (..., n_groups, group_size).
             n_groups = self.hidden_size // self.group_size
+            orig_shape = op.Shape(gated)
             if flags.ort_cuda_grouped_rmsnorm_workaround:
                 # ORT ≤1.24.4 CUDA kernel for RMSNormalization produces
                 # wrong results when scale is 2D. Decompose into basic
                 # ops as a workaround.
                 grouped = op.Reshape(
                     gated,
-                    op.Constant(value_ints=[0, n_groups, self.group_size]),
+                    op.Constant(value_ints=[-1, n_groups, self.group_size]),
                 )
                 variance = op.ReduceMean(
                     op.Mul(grouped, grouped),
@@ -101,10 +103,7 @@ def forward(self, op: builder.OpBuilder, hidden_states: ir.Value, gate: ir.Value
                     op.Sqrt(op.Add(variance, self.variance_epsilon)),
                 )
                 normed = op.Mul(grouped, rnorm)
-                normed = op.Reshape(
-                    normed,
-                    op.Constant(value_ints=[0, self.hidden_size]),
-                )
+                normed = op.Reshape(normed, orig_shape)
                 normed = op.Mul(
                     normed,
                     op.Cast(self.weight, to=ir.DataType.FLOAT),
@@ -115,7 +114,7 @@ def forward(self, op: builder.OpBuilder, hidden_states: ir.Value, gate: ir.Value
                 gated = op.CastLike(gated, hidden_states)
                 grouped = op.Reshape(
                     gated,
-                    op.Constant(value_ints=[0, n_groups, self.group_size]),
+                    op.Constant(value_ints=[-1, n_groups, self.group_size]),
                 )
                 weight_grouped = op.Reshape(
                     self.weight,
@@ -127,10 +126,7 @@ def forward(self, op: builder.OpBuilder, hidden_states: ir.Value, gate: ir.Value
                     epsilon=self.variance_epsilon,
                     axis=-1,
                 )
-                normed = op.Reshape(
-                    normed,
-                    op.Constant(value_ints=[0, self.hidden_size]),
-                )
+                normed = op.Reshape(normed, orig_shape)
         else:
             # Standard RMSNorm over the full dimension.
             # Cast gated back to native dtype; stash_type=1 handles fp32.
diff --git a/src/mobius/components/_ssm.py b/src/mobius/components/_ssm.py
index f9a45a84..6bca946e 100644
--- a/src/mobius/components/_ssm.py
+++ b/src/mobius/components/_ssm.py
@@ -75,7 +75,7 @@ def _project_ssm_params(self, op: builder.OpBuilder, x_db):
         """
         dt_raw, b_mat, c_mat = op.Split(
             x_db,
-            op.Constant(value_ints=[self.dt_rank, self.d_state, self.d_state]),
+            [self.dt_rank, self.d_state, self.d_state],
             axis=-1,
             _outputs=3,
         )
@@ -180,7 +180,7 @@ def _project_ssm_params(self, op, x_db):
         """Split + layernorm on dt, B, C."""
         dt_raw, b_mat, c_mat = op.Split(
             x_db,
-            op.Constant(value_ints=[self.dt_rank, self.d_state, self.d_state]),
+            [self.dt_rank, self.d_state, self.d_state],
             axis=-1,
             _outputs=3,
         )
@@ -262,7 +262,7 @@ def forward(
         )
         # Clamp dt to time_step_min (matches HF torch.clamp(dt, min=...))
         if self.time_step_min > 0.0:
-            dt = op.Clip(dt, op.Constant(value_float=self.time_step_min))
+            dt = op.Clip(dt, self.time_step_min)
 
         # A = -exp(A_log) in fp32: (num_heads,)
         a_neg = op.Neg(op.Exp(op.Cast(self.A_log, to=ir.DataType.FLOAT)))
@@ -274,16 +274,18 @@ def forward(
         da = op.Exp(op.Mul(dt_4d, a_4d))
 
         # Reshape hidden: (batch, num_heads, d_head)
-        hidden_shape = op.Constant(value_ints=[0, self.num_heads, self.d_head])
-        hidden_3d = op.Cast(op.Reshape(hidden_states, hidden_shape), to=ir.DataType.FLOAT)
+        hidden_3d = op.Cast(
+            op.Reshape(hidden_states, [0, self.num_heads, self.d_head]),
+            to=ir.DataType.FLOAT,
+        )
 
         # Expand B from groups to heads (in fp32)
-        b_shape = op.Constant(value_ints=[0, self.n_groups, 1, self.d_state])
+        b_shape = [0, self.n_groups, 1, self.d_state]
         b_4d = op.Reshape(op.Cast(b_mat, to=ir.DataType.FLOAT), b_shape)
-        b_expand_shape = op.Constant(value_ints=[1, 1, self.heads_per_group, 1])
-        b_expanded = op.Expand(b_4d, b_expand_shape)
-        b_heads_shape = op.Constant(value_ints=[0, self.num_heads, self.d_state])
-        b_heads = op.Reshape(b_expanded, b_heads_shape)
+        expand_shape = [1, 1, self.heads_per_group, 1]
+        b_expanded = op.Expand(b_4d, expand_shape)
+        heads_shape = [0, self.num_heads, self.d_state]
+        b_heads = op.Reshape(b_expanded, heads_shape)
         b_ssm = op.Unsqueeze(b_heads, [2])
 
         # dBx: dt * B * x (in fp32)
@@ -296,8 +298,8 @@ def forward(
 
         # Readout: y = C . h + D * x (in fp32)
         c_4d = op.Reshape(op.Cast(c_mat, to=ir.DataType.FLOAT), b_shape)
-        c_expanded = op.Expand(c_4d, b_expand_shape)
-        c_heads = op.Reshape(c_expanded, b_heads_shape)
+        c_expanded = op.Expand(c_4d, expand_shape)
+        c_heads = op.Reshape(c_expanded, heads_shape)
         c_ssm = op.Unsqueeze(c_heads, [2])
 
         y = op.ReduceSum(
@@ -309,8 +311,10 @@ def forward(
         y = op.Add(y, op.Mul(d_3d, hidden_3d))
 
         # Flatten and cast back to input dtype: (batch, num_heads * d_head)
-        flat_shape = op.Constant(value_ints=[0, self.num_heads * self.d_head])
-        y = op.CastLike(op.Reshape(y, flat_shape), hidden_states)
+        y = op.CastLike(
+            op.Reshape(y, [0, self.num_heads * self.d_head]),
+            hidden_states,
+        )
 
         return y, op.CastLike(new_ssm_state, ssm_state)
 
@@ -329,7 +333,7 @@ def forward(self, op: builder.OpBuilder, x: ir.Value):
         variance = op.ReduceMean(op.Mul(x_f32, x_f32), [-1], keepdims=True)
         x_normed = op.Div(
             x_f32,
-            op.Sqrt(op.Add(variance, op.Constant(value_float=self._eps))),
+            op.Sqrt(op.Add(variance, self._eps)),
         )
         result = op.Mul(x_normed, op.Cast(self.weight, to=ir.DataType.FLOAT))
         return op.CastLike(result, x)
diff --git a/src/mobius/models/bamba.py b/src/mobius/models/bamba.py
index 1c852d47..57ae7a0e 100644
--- a/src/mobius/models/bamba.py
+++ b/src/mobius/models/bamba.py
@@ -273,41 +273,9 @@ def preprocess_weights(
 
         Handles:
         1. Weight tying (embed_tokens ↔ lm_head)
-        2. Mamba2 SSM params: A_log, D, dt_bias stay under mamba.ssm
-        3. Norm rename: mamba.norm → mamba.norm (matches HF naming)
-        4. MLP rename: feed_forward → feed_forward (matches HF naming)
         """
         if self.config.tie_word_embeddings:
             tie_word_embeddings(state_dict)
 
-        new_state_dict: dict[str, torch.Tensor] = {}
-        for key, value in state_dict.items():
-            new_key = _rename_bamba_weight(key)
-            new_state_dict[new_key] = value
-
-        return new_state_dict
-
-
-def _rename_bamba_weight(key: str) -> str:
-    """Rename a single HF weight key to match ONNX module structure.
-
-    HF BambaForCausalLM weight naming:
-        model.layers.N.mamba.{in_proj, conv1d, out_proj, norm, A_log, D, dt_bias}
-        model.layers.N.self_attn.{q,k,v,o}_proj
-        model.layers.N.feed_forward.{gate,up,down}_proj
-        model.layers.N.{input_layernorm, pre_ff_layernorm}
-        model.{embed_tokens, final_layernorm}
-        lm_head
-
-    ONNX parameter naming:
-        Same as HF, except SSM params are nested under mamba.ssm:
-        model.layers.N.mamba.ssm.{A_log, D, dt_bias}
-    """
-    # SSM params: nest A_log, D, dt_bias under mamba.ssm
-    ssm_params = (".mamba.A_log", ".mamba.D", ".mamba.dt_bias")
-    for param in ssm_params:
-        if key.endswith(param):
-            # e.g. "model.layers.0.mamba.A_log" → "model.layers.0.mamba.ssm.A_log"
-            return key.replace(".mamba.", ".mamba.ssm.")
-
-    return key
+        # HF and ONNX naming match directly — no renaming needed.
+        return state_dict
diff --git a/src/mobius/models/granitemoehybrid.py b/src/mobius/models/granitemoehybrid.py
index eda9d04b..e5a68705 100644
--- a/src/mobius/models/granitemoehybrid.py
+++ b/src/mobius/models/granitemoehybrid.py
@@ -319,15 +319,14 @@ def preprocess_weights(
 
         Handles:
         1. Weight tying (embed_tokens ↔ lm_head)
-        2. Mamba2 SSM params: A_log, D, dt_bias nested under mamba.ssm
-        3. MoE gate: block_sparse_moe.router.layer.weight → block_sparse_moe.gate.weight
-        4. MoE fused input: block_sparse_moe.input_linear [n_experts, 2*mid, hidden]
+        2. MoE gate: block_sparse_moe.router.layer.weight → block_sparse_moe.gate.weight
+        3. MoE fused input: block_sparse_moe.input_linear [n_experts, 2*mid, hidden]
            → per-expert block_sparse_moe.experts.{e}.{gate,up}_proj.weight
-        5. MoE fused output: block_sparse_moe.output_linear [n_experts, hidden, mid]
+        4. MoE fused output: block_sparse_moe.output_linear [n_experts, hidden, mid]
            → per-expert block_sparse_moe.experts.{e}.down_proj.weight
-        6. Shared MLP fused gate+up: shared_mlp.input_linear [2*shared_mid, hidden]
+        5. Shared MLP fused gate+up: shared_mlp.input_linear [2*shared_mid, hidden]
            → shared_mlp.gate_proj.weight + shared_mlp.up_proj.weight
-        7. Shared MLP down proj: shared_mlp.output_linear → shared_mlp.down_proj
+        6. Shared MLP down proj: shared_mlp.output_linear → shared_mlp.down_proj
         """
         if self.config.tie_word_embeddings:
             tie_word_embeddings(state_dict)
@@ -345,9 +344,6 @@ def preprocess_weights(
 # Weight name mapping
 # ---------------------------------------------------------------------------
 
-# Mamba2 SSM params stored flat on HF "mamba" that we nest under "mamba.ssm"
-_MAMBA2_SSM_PARAMS = (".mamba.A_log", ".mamba.D", ".mamba.dt_bias")
-
 
 def _rename_granitemoehybrid_weight(
     key: str,
@@ -359,12 +355,6 @@ def _rename_granitemoehybrid_weight(
     Returns the new key, or None if the weight was handled inline
     (fused tensors split into multiple per-expert outputs).
     """
-    # SSM params: nest A_log, D, dt_bias under mamba.ssm
-    # e.g. "model.layers.0.mamba.A_log" → "model.layers.0.mamba.ssm.A_log"
-    for param in _MAMBA2_SSM_PARAMS:
-        if key.endswith(param):
-            return key.replace(".mamba.", ".mamba.ssm.")
-
     # MoE gate: router.layer.weight → gate.weight
     # e.g. "…block_sparse_moe.router.layer.weight" → "…block_sparse_moe.gate.weight"
     key = key.replace(".block_sparse_moe.router.layer.", ".block_sparse_moe.gate.")
diff --git a/src/mobius/models/mamba.py b/src/mobius/models/mamba.py
index 42471711..61d98454 100644
--- a/src/mobius/models/mamba.py
+++ b/src/mobius/models/mamba.py
@@ -431,35 +431,14 @@ def preprocess_weights(
     ) -> dict[str, torch.Tensor]:
         """Map HuggingFace Mamba2ForCausalLM weights to ONNX names.
 
-        HF naming:
+        HF and ONNX naming match directly — no SSM param renaming needed:
             backbone.layers.{i}.norm.weight (same)
-            backbone.layers.{i}.mixer.{A_log,D,dt_bias}
-                -> backbone.layers.{i}.mixer.ssm.{A_log,D,dt_bias}
-            backbone.layers.{i}.mixer.{in_proj,conv1d,norm,out_proj}
+            backbone.layers.{i}.mixer.{A_log,D,dt_bias,in_proj,conv1d,norm,out_proj}
                 (same)
             backbone.embeddings.weight (same)
             backbone.norm_f.weight (same)
             lm_head.weight (same)
         """
-        renames = {}
-        _ssm_params = ("A_log", "D", "dt_bias")
-
-        for key in list(state_dict):
-            new_key = key
-            # mixer.{ssm_param} -> mixer.ssm.{ssm_param}
-            for param in _ssm_params:
-                old_seg = f".mixer.{param}"
-                new_seg = f".mixer.ssm.{param}"
-                if new_key.endswith(old_seg):
-                    new_key = new_key.replace(old_seg, new_seg)
-                    break
-
-            if new_key != key:
-                renames[key] = new_key
-
-        for old_key, new_key in renames.items():
-            state_dict[new_key] = state_dict.pop(old_key)
-
         # Tied embeddings (Mamba2 uses self.backbone, not self.model)
         if self.config.tie_word_embeddings:
             if "lm_head.weight" in state_dict:
diff --git a/src/mobius/models/nemotron_h.py b/src/mobius/models/nemotron_h.py
index b56f6c9a..30b30437 100644
--- a/src/mobius/models/nemotron_h.py
+++ b/src/mobius/models/nemotron_h.py
@@ -311,7 +311,7 @@ def preprocess_weights(
         3. ``backbone.embeddings.`` → ``model.embed_tokens.``
         4. ``backbone.norm_f.`` → ``model.norm.``
         5. Per-layer ``mixer.`` rename based on layer type:
-           - mamba: ``mixer.`` → ``mamba.`` (SSM params nested under ``mamba.ssm.``)
+           - mamba: ``mixer.`` → ``mamba.``
            - attention: ``mixer.`` → ``self_attn.``
            - mlp: ``mixer.`` → ``mlp.``
         """
@@ -336,9 +336,6 @@ def preprocess_weights(
 # Layer index regex: backbone.layers.<N>.<rest>
 _LAYER_RE = re.compile(r"^backbone\.layers\.(\d+)\.(.+)$")
 
-# Mamba SSM params that need to be nested under mamba.ssm
-_MAMBA_SSM_PARAMS = ("A_log", "D", "dt_bias")
-
 
 def _rename_nemotron_h_weight(key: str, layer_types: list[str]) -> str:
     """Rename a single HF weight key to match ONNX module structure.
@@ -356,8 +353,7 @@ def _rename_nemotron_h_weight(key: str, layer_types: list[str]) -> str:
         model.embed_tokens.weight
         model.norm.weight
         model.layers.N.norm.weight
-        model.layers.N.mamba.{in_proj, conv1d, out_proj, norm}   (mamba direct)
-        model.layers.N.mamba.ssm.{A_log, D, dt_bias}            (mamba SSM nested)
+        model.layers.N.mamba.{in_proj, conv1d, out_proj, norm, A_log, D, dt_bias}
         model.layers.N.self_attn.{q_proj, k_proj, v_proj, o_proj}.weight
         model.layers.N.mlp.{up_proj, down_proj}.weight
         lm_head.weight
@@ -378,10 +374,6 @@ def _rename_nemotron_h_weight(key: str, layer_types: list[str]) -> str:
         if rest.startswith("mixer."):
             mixer_rest = rest[len("mixer.") :]
             if ltype == "mamba2":
-                # Check if this is an SSM param that needs nesting
-                param_name = mixer_rest.split(".")[0]
-                if param_name in _MAMBA_SSM_PARAMS:
-                    return f"model.layers.{layer_idx}.mamba.ssm.{mixer_rest}"
                 return f"model.layers.{layer_idx}.mamba.{mixer_rest}"
             elif ltype == "full_attention":
                 return f"model.layers.{layer_idx}.self_attn.{mixer_rest}"
diff --git a/src/mobius/tasks/_cache_utils.py b/src/mobius/tasks/_cache_utils.py
index 395d62f2..8049aaf1 100644
--- a/src/mobius/tasks/_cache_utils.py
+++ b/src/mobius/tasks/_cache_utils.py
@@ -17,7 +17,7 @@
 
 from mobius._configs import BaseModelConfig
 
-_FUNCTIONS_DOMAIN = "pkg.mobius"
+_FUNCTIONS_DOMAIN = "com.microsoft"
 
 # Cache state pair: (key, value) or (conv_state, ssm_state) for stateful
 # layers; (rec_state,) for lightning/conv attention (single state);
@@ -244,7 +244,7 @@ def _make_hybrid_cache_inputs(
             )
             ssm_state = ir.Value(
                 name=f"{prefix}.{i}.ssm_state",
-                shape=ir.Shape([batch, mamba2_n_heads, mamba2_d_head, mamba2_d_state]),
+                shape=ir.Shape([batch, mamba2_n_heads, mamba2_d_state, mamba2_d_head]),
                 type=ir.TensorType(dtype),
             )
             flat.extend([conv_state, ssm_state])
@@ -324,15 +324,17 @@ def _register_linear_attention_functions(
 ) -> None:
     """Register CausalConvWithState and LinearAttention functions.
 
-    Registers functions for DeltaNet (``linear_attention`` layers) and/or
-    Lightning Attention (``lightning_attention`` layers) as needed.
-    Adds the ``pkg.mobius`` opset import to the graph.
+    Registers functions for DeltaNet (``linear_attention`` layers),
+    Lightning Attention (``lightning_attention`` layers), and/or
+    Mamba2 (``mamba2`` layers) as needed.
+    Adds the ``com.microsoft`` opset import to the graph.
     """
     layer_types = getattr(config, "layer_types", None) or []
     has_deltanet = "linear_attention" in layer_types
     has_lightning = "lightning_attention" in layer_types
+    has_mamba2 = "mamba2" in layer_types
 
-    if not has_deltanet and not has_lightning:
+    if not has_deltanet and not has_lightning and not has_mamba2:
         return
 
     from mobius.functions import (
@@ -369,4 +371,76 @@ def _register_linear_attention_functions(
         )
         model.functions[attn_func_gated.identifier()] = attn_func_gated
 
+    if has_mamba2:
+        mamba2_n_heads = getattr(config, "mamba_n_heads", 0)
+        mamba2_d_head = getattr(config, "mamba_d_head", 0)
+        mamba2_d_state = getattr(config, "mamba_d_state", 0)
+        mamba2_n_groups = getattr(config, "mamba_n_groups", 1)
+        mamba2_d_conv = getattr(config, "mamba_d_conv", 4)
+        mamba_expand = getattr(config, "mamba_expand", 2)
+        mamba2_d_inner = (
+            mamba2_n_heads * mamba2_d_head
+            if mamba2_n_heads and mamba2_d_head
+            else config.hidden_size * mamba_expand
+        )
+        mamba2_conv_dim = mamba2_d_inner + 2 * mamba2_n_groups * mamba2_d_state
+        conv_func = causal_conv_nd_with_state(
+            kernel_size=mamba2_d_conv,
+            channels=mamba2_conv_dim,
+            ndim=1,
+            activation="silu",
+        )
+        attn_func = linear_attention(
+            q_num_heads=mamba2_n_heads,
+            kv_num_heads=mamba2_n_heads,
+            update_rule="gated",
+            scale=1.0,
+            stash_type=ir.DataType.FLOAT,
+        )
+        model.functions[conv_func.identifier()] = conv_func
+        model.functions[attn_func.identifier()] = attn_func
+
+    model.graph.opset_imports[_FUNCTIONS_DOMAIN] = 1
+
+
+def _register_linear_attention_functions_for_ssm2(
+    model: ir.Model,
+    config: BaseModelConfig,
+) -> None:
+    """Register CausalConvWithState and LinearAttention for pure Mamba2 models.
+
+    Unlike :func:`_register_linear_attention_functions` (which inspects
+    ``layer_types``), this always registers the Mamba2 function ops.
+    Called by :class:`SSM2CausalLMTask` for pure Mamba2 models that don't
+    have a ``layer_types`` attribute.
+    """
+    from mobius._configs import Mamba2Config
+    from mobius.functions import (
+        causal_conv_nd_with_state,
+        linear_attention,
+    )
+
+    assert isinstance(config, Mamba2Config)
+    n_heads = config.num_heads
+    d_state = config.state_size
+    n_groups = config.n_groups
+    d_inner = config.intermediate_size
+    d_conv = config.conv_kernel
+    conv_dim = d_inner + 2 * n_groups * d_state
+
+    conv_func = causal_conv_nd_with_state(
+        kernel_size=d_conv,
+        channels=conv_dim,
+        ndim=1,
+        activation="silu",
+    )
+    attn_func = linear_attention(
+        q_num_heads=n_heads,
+        kv_num_heads=n_heads,
+        update_rule="gated",
+        scale=1.0,
+        stash_type=ir.DataType.FLOAT,
+    )
+    model.functions[conv_func.identifier()] = conv_func
+    model.functions[attn_func.identifier()] = attn_func
     model.graph.opset_imports[_FUNCTIONS_DOMAIN] = 1
diff --git a/src/mobius/tasks/_ssm_causal_lm.py b/src/mobius/tasks/_ssm_causal_lm.py
index 1df88f1e..63f6f504 100644
--- a/src/mobius/tasks/_ssm_causal_lm.py
+++ b/src/mobius/tasks/_ssm_causal_lm.py
@@ -130,17 +130,18 @@ class SSM2CausalLMTask(ModelTask):
     """Causal language model with Mamba2/SSD state carry.
 
     Like SSMCausalLMTask but with 4D SSM state for Mamba2 multi-head
-    architecture and wider conv_dim.
+    architecture and wider conv_dim.  Uses LinearAttention + CausalConvWithState
+    function ops.
 
     Inputs:
         - input_ids: [batch, sequence_len] INT64
         - past_states.{i}.conv_state: [batch, conv_dim, conv_kernel-1]
-        - past_states.{i}.ssm_state: [batch, num_heads, head_dim, state_size]
+        - past_states.{i}.ssm_state: [batch, num_heads, state_size, head_dim]
 
     Outputs:
         - logits: FLOAT
         - present.{i}.conv_state: [batch, conv_dim, conv_kernel-1]
-        - present.{i}.ssm_state: [batch, num_heads, head_dim, state_size]
+        - present.{i}.ssm_state: [batch, num_heads, state_size, head_dim]
     """
 
     model_roles: ClassVar[dict[str, str]] = {"model": "decoder"}
@@ -156,9 +157,17 @@ def build(
         n_groups = config.n_groups
         state_size = config.state_size
         conv_dim = config.intermediate_size + 2 * n_groups * state_size
-        return _build_ssm_task(
+        pkg = _build_ssm_task(
             module,
             config,
             conv_state_shape=[conv_dim, config.conv_kernel - 1],
-            ssm_state_shape=[config.num_heads, config.head_dim, state_size],
+            # LinearAttention convention: (H, d_k, d_v) = (H, d_state, d_head)
+            ssm_state_shape=[config.num_heads, state_size, config.head_dim],
+        )
+        # Register CausalConvWithState and LinearAttention function ops.
+        from mobius.tasks._cache_utils import (
+            _register_linear_attention_functions_for_ssm2,
         )
+
+        _register_linear_attention_functions_for_ssm2(pkg["model"], config)
+        return pkg
diff --git a/testdata/cases/causal-lm/nemotron-h-nano-4b.yaml b/testdata/cases/causal-lm/nemotron-h-nano-4b.yaml
new file mode 100644
index 00000000..404b1f6a
--- /dev/null
+++ b/testdata/cases/causal-lm/nemotron-h-nano-4b.yaml
@@ -0,0 +1,17 @@
+model_id: "nvidia/NVIDIA-Nemotron-3-Nano-4B-BF16"
+revision: "main"
+task_type: "text-generation"
+dtype: "float32"
+
+inputs:
+  prompts:
+    - "Here is my poem:"
+
+level: "L4+L5"
+
+generation:
+  max_new_tokens: 20
+  do_sample: false
+
+skip_reason: "4B parameter model — too large for CI golden generation."
+notes: "NemotronH Nano 4B. Hybrid Mamba2 + Attention + MLP architecture from NVIDIA."
diff --git a/tests/build_graph_test.py b/tests/build_graph_test.py
index efd0bbbe..5475ea39 100644
--- a/tests/build_graph_test.py
+++ b/tests/build_graph_test.py
@@ -76,6 +76,12 @@
     "qwen3_5_text",
     "qwen3_5_moe",
     "qwen3_next",
+    # Models using LinearAttention / CausalConvWithState custom ops
+    # prevent full shape/type propagation through com.microsoft domain.
+    "bamba",
+    "granitemoehybrid",
+    "mamba2",
+    "nemotron_h",
     # VL/Speech models with value_info/shape checker issues
     "qwen2_vl",
     "qwen2_5_vl",
@@ -2967,7 +2973,7 @@ def test_mamba2_state_io(self):
             assert f"present.{i}.ssm_state" in output_names
 
     def test_mamba2_preprocess_weights(self):
-        """Verify SSM param nesting: mixer.A_log -> mixer.ssm.A_log."""
+        """Verify SSM params stay at mixer level (no nesting)."""
         from mobius.models.mamba import Mamba2CausalLMModel
 
         config = self._mamba2_config()
@@ -2984,9 +2990,10 @@ def test_mamba2_preprocess_weights(self):
         }
         result = module.preprocess_weights(state_dict)
 
-        assert "backbone.layers.0.mixer.ssm.A_log" in result
-        assert "backbone.layers.0.mixer.ssm.D" in result
-        assert "backbone.layers.0.mixer.ssm.dt_bias" in result
+        # SSM params stay directly on mixer (no .ssm. nesting)
+        assert "backbone.layers.0.mixer.A_log" in result
+        assert "backbone.layers.0.mixer.D" in result
+        assert "backbone.layers.0.mixer.dt_bias" in result
         # Non-SSM params stay as-is
         assert "backbone.layers.0.mixer.in_proj.weight" in result
         assert "backbone.layers.0.norm.weight" in result
@@ -3078,7 +3085,7 @@ def test_bamba_registry_lookup(self):
         assert _default_task_for_model("bamba") == "hybrid-text-generation"
 
     def test_bamba_preprocess_weights(self):
-        """Verify preprocess_weights nests SSM params under mamba.ssm."""
+        """Verify preprocess_weights passes SSM params through unchanged."""
         import torch
 
         from mobius.models.bamba import BambaCausalLMModel
@@ -3100,9 +3107,10 @@ def test_bamba_preprocess_weights(self):
         }
         result = module.preprocess_weights(state_dict)
 
-        assert "model.layers.0.mamba.ssm.A_log" in result
-        assert "model.layers.0.mamba.ssm.D" in result
-        assert "model.layers.0.mamba.ssm.dt_bias" in result
+        # SSM params stay directly on mamba (no .ssm. nesting)
+        assert "model.layers.0.mamba.A_log" in result
+        assert "model.layers.0.mamba.D" in result
+        assert "model.layers.0.mamba.dt_bias" in result
         assert "model.layers.0.mamba.in_proj.weight" in result
         assert "model.layers.1.self_attn.q_proj.weight" in result
 
@@ -3185,10 +3193,10 @@ def test_nemotron_h_preprocess_weights(self):
         assert "model.embed_tokens.weight" in result
         assert "model.norm.weight" in result
 
-        # Layer 0 (mamba2): SSM params nested under mamba.ssm
-        assert "model.layers.0.mamba.ssm.A_log" in result
-        assert "model.layers.0.mamba.ssm.D" in result
-        assert "model.layers.0.mamba.ssm.dt_bias" in result
+        # Layer 0 (mamba2): SSM params directly under mamba (no nesting)
+        assert "model.layers.0.mamba.A_log" in result
+        assert "model.layers.0.mamba.D" in result
+        assert "model.layers.0.mamba.dt_bias" in result
         # Non-SSM mamba params stay under mamba.*
         assert "model.layers.0.mamba.in_proj.weight" in result
         assert "model.layers.0.mamba.conv1d.weight" in result
diff --git a/tests/synthetic_parity_test.py b/tests/synthetic_parity_test.py
index c1023d19..77f4d860 100644
--- a/tests/synthetic_parity_test.py
+++ b/tests/synthetic_parity_test.py
@@ -87,6 +87,12 @@
     # Zamba weight-tying references layers.2.shared_transf (the third layer) but
     # the tiny config only has 2 layers — HF tie_weights validation crashes.
     "zamba": "Zamba weight-tying requires num_layers > 2; tiny 2-layer config causes HF tie_weights error",
+    # NemotronH: uses com.microsoft.LinearAttention which requires ORT nightly
+    "nemotron_h": "NemotronH uses LinearAttention contrib op not available in stable ORT",
+    # GraniteMoeHybrid: Mamba2+Attention hybrid uses LinearAttention; ORT function
+    # inlining drops initializers inside function ops (pre-existing on main where
+    # the old Mamba2Scan code crashed at runtime with a Squeeze shape error).
+    "granitemoehybrid": "Mamba2 hybrid model — LinearAttention function op inlining drops initializers",
 }
 
 # Per-model atol overrides for L3 synthetic parity.