Use LinearAttention to implement Mamba2 and other cleanup - #163
Merged
Conversation
- Replace op.Constant calls with raw Python values in _ssm.py to avoid potential constant caching issues across dtype builds - Create L2 YAML test case (nemotron-h-nano-4b.yaml) - Add nemotron_h to L3 synthetic parity _SKIP_REASONS (Mamba2BlockSingle only handles seq_len=1) - Reorder example script CLI args (--ci before --repetition-penalty) Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: G Ramalingam <grama@microsoft.com>
These shapes are reused for both B and C group-to-head expansion, so drop the b_ prefix to reflect shared usage. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: G Ramalingam <grama@microsoft.com>
Replace the three Mamba2Block variants (Single, Scan, ChunkedSSD) with a single unified Mamba2Block that uses com.microsoft.LinearAttention (update_rule="gated") for the SSM recurrence and CausalConvWithState for the depthwise Conv1D. This supports both T=1 decode and T>1 prefill in one code path. Key changes: - Rewrite _mamba_block.py: single Mamba2Block class using LinearAttention - Delete _mamba_block_scan.py and _mamba_block_chunked.py (3 variants → 1) - Remove mamba_scan flag from _flags.py (no longer needed) - Fix GatedRMSNorm to handle 3D input (B, T, H) for grouped normalization - Update SSM2CausalLMTask: new state shape (H, d_state, d_head) matching LinearAttention convention, register function ops - Add _register_linear_attention_functions_for_ssm2 for pure Mamba2 models - Add Mamba2-using models to _CHECKER_SKIP_MODELS (custom op shape limits) The Mamba2 SSD recurrence maps to LinearAttention as: query=C, key=B, value=dt*x, decay=A*dt (log-space) with kv_num_heads=num_heads (each head has its own state/decay). All 2390 tests pass. Lint clean. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: G Ramalingam <grama@microsoft.com>
Replace the three Mamba2Block variants (Single, Scan, ChunkedSSD) with a single unified Mamba2Block that uses com.microsoft.LinearAttention (update_rule="gated") for the SSM recurrence and CausalConvWithState for the depthwise Conv1D. This supports both T=1 decode and T>1 prefill in one code path. Key changes: - Rewrite _mamba_block.py: single Mamba2Block class using LinearAttention - Delete _mamba_block_scan.py and _mamba_block_chunked.py (3 variants → 1) - Remove mamba_scan flag from _flags.py (no longer needed) - Fix GatedRMSNorm to handle 3D input (B, T, H) for grouped normalization - Update SSM2CausalLMTask: new state shape (H, d_state, d_head) matching LinearAttention convention, register function ops - Add _register_linear_attention_functions_for_ssm2 for pure Mamba2 models - Add Mamba2-using models to _CHECKER_SKIP_MODELS (custom op shape limits) The Mamba2 SSD recurrence maps to LinearAttention as: query=C, key=B, value=dt*x, decay=A*dt (log-space) with kv_num_heads=num_heads (each head has its own state/decay). All 2390 tests pass. Lint clean. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: G Ramalingam <grama@microsoft.com>
…tant conflict - Fix SSM state shape: (B, H, d_head, d_state) → (B, H, d_state, d_head) to match LinearAttention convention - Add --ep flag to example script (cpu/cuda/onnx-standard) - Fix op.Clip constant conflict: use op.Constant(value_float=...) to avoid initializer name collision across layers Verified: ONNX output matches HuggingFace exactly on f32/cpu/onnx-standard. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: G Ramalingam <grama@microsoft.com>
Performance Comparison
|
|
The author of this PR, gramalingam, is not an activated member of this organization on Codecov. |
🏗️ Architecture Diff
mamba (ssm-text-generation) / model — 22 change(s)Op summary: 99 → 98 nodes --- base
+++ head
@@ -14,7 +14,6 @@
Transpose
Transpose
MatMul
-Constant
Split
Transpose
MatMulRemoved nodes:
Modified attributes:
Connectivity changes:
Initializer changes:
Legend: ⚪ No change · 🔵 Minor (attrs/inits) · 🟡 Moderate (nodes added/removed) · 🔴 Major (interface changed) |
Contributor
There was a problem hiding this comment.
Pull request overview
This PR updates the Mamba2-based model stack (notably NemotronH) to express the Mamba2 recurrence via com.microsoft.LinearAttention + com.microsoft.CausalConvWithState, and cleans up related weight preprocessing, cache tensor layouts, and flags.
Changes:
- Replaces prior Mamba2 multi-token implementations (Scan / chunked SSD) with a unified LinearAttention-based
Mamba2Block, and removes themamba_scanflag + deleted implementations. - Updates cache state tensor layout for Mamba2 to
(B, H, d_state, d_head)and registers LinearAttention/CausalConv function bodies in tasks. - Aligns HF↔ONNX weight naming for Mamba2 SSM params by removing the
.ssm.nesting renames; updates tests/skips and adds a NemotronH Nano 4B test case stub.
Reviewed changes
Copilot reviewed 18 out of 18 changed files in this pull request and generated 3 comments.
Show a summary per file
| File | Description |
|---|---|
| tests/synthetic_parity_test.py | Skips NemotronH synthetic parity due to LinearAttention op availability assumptions. |
| tests/build_graph_test.py | Updates preprocess-weight expectations (no .ssm. nesting) and expands shape-check skip list for custom ops. |
| testdata/cases/causal-lm/nemotron-h-nano-4b.yaml | Adds a NemotronH Nano 4B golden-generation case (skipped due to size). |
| src/mobius/tasks/_ssm_causal_lm.py | Updates SSM2 state layout + registers LinearAttention/CausalConv function ops for pure Mamba2 tasks. |
| src/mobius/tasks/_cache_utils.py | Updates hybrid cache Mamba2 state shape + extends function-op registration to include Mamba2. |
| src/mobius/models/nemotron_h.py | Removes Mamba2 SSM .ssm. nesting from weight renaming rules and docs. |
| src/mobius/models/mamba.py | Removes Mamba2 SSM param renaming (HF/ONNX naming now matches). |
| src/mobius/models/granitemoehybrid.py | Removes Mamba2 SSM .ssm. nesting rename logic. |
| src/mobius/models/bamba.py | Simplifies preprocess_weights to tie embeddings only; no renaming. |
| src/mobius/components/_ssm.py | Minor builder-API cleanups (Split/Reshape/Clip constant handling). |
| src/mobius/components/_rms_norm.py | Fixes grouped RMSNorm reshaping to support 2D/3D inputs by flattening/restoring leading dims. |
| src/mobius/components/_mamba_block.py | Reimplements Mamba2Block using LinearAttention + CausalConvWithState and removes older variant machinery. |
| src/mobius/components/_mamba_block_scan.py | Deleted (old Mamba2 Scan path). |
| src/mobius/components/_mamba_block_chunked.py | Deleted (old chunked SSD path). |
| src/mobius/components/init.py | Removes Mamba2Scan from public exports. |
| src/mobius/_flags.py | Removes the mamba_scan flag and _env_str helper. |
| pyproject.toml | Removes now-unneeded Ruff per-file ignores for deleted chunked SSD file. |
| examples/nemotron_3_nano_text_generation.py | Updates hybrid cache init shape, adds --ep, and tweaks CLI help/arg ordering. |
gramalingam
changed the title
Address PR #163 review comments: 1. Upcast Mamba2 SSM recurrence to fp32: Cast dt_raw, dt_bias, A_log, x_hidden, B, C, ssm_state to float32 before computing dt/decay/value and calling LinearAttention. This matches HuggingFace which computes the SSM recurrence in float32. Cast outputs back to model dtype. Also set stash_type=ir.DataType.FLOAT for the LinearAttention Scan body so the internal recurrence runs in fp32. 2. Fix _FUNCTIONS_DOMAIN: change from 'pkg.mobius' to 'com.microsoft' to match the actual function domain. The com.microsoft opset is already imported by _make_graph(). 3. Fix --device help text to mention both ORT and HF usage. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: G Ramalingam <grama@microsoft.com>
The granitemoehybrid L3 test was never passing on main — the old Mamba2Scan
code crashed at ORT runtime (Squeeze shape error) and was caught by the
pytest.skip('ONNX inference failed...') handler. The LinearAttention rewrite
fixed the crash, but now the test actually runs and reveals the model
produces wrong output due to ORT function inlining dropping initializers
inside function ops. Skip with an explicit reason.
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Signed-off-by: G Ramalingam <grama@microsoft.com>
gramalingam
enabled auto-merge (squash)
April 13, 2026 02:37
justinchuby
approved these changes
Apr 13, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Update the NemotronH/Mamba2 implementations to use LinearAttention.