Use LinearAttention to implement Mamba2 and other cleanup (#163)

Update the NemotronH/Mamba2 implementations to use LinearAttention.

---------

Signed-off-by: G Ramalingam <grama@microsoft.com>
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Author: gramalingam

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(

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"

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()

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 (