Skip to content

[Atlas Inference] Strix Halo (gfx1151) Mamba-2 + conv1d Hub kernel path - #49709

Open
AzeezIsh wants to merge 1 commit into
huggingface:mainfrom
AzeezIsh:gdn-rocm-mamba2-gfx1151
Open

AzeezIsh wants to merge 1 commit into
huggingface:mainfrom
AzeezIsh:gdn-rocm-mamba2-gfx1151

Conversation

@AzeezIsh

@AzeezIsh AzeezIsh commented Oct 18, 2026 •

Copy link
Copy Markdown
Contributor

🤗 Transformers CI

On AMD Strix Halo (Ryzen AI Max, gfx1151, RDNA 3.5) the mamba_ssm and causal_conv1d packages have no ROCm wheel, and _KERNEL_MAPPING only defines "cuda" and "xpu" entries for the Mamba-2 and causal-conv1d ops, so NemotronH and Mamba2 models fall back to the slow PyTorch einsum segment-sum reference (mamba2_chunk_scan / mamba2_selective_state_update). On NVIDIA GB10 (SM121, aarch64, torch 2.12 + CUDA 13.0), kernels-community/mamba-ssm v4 now backfills the torch212-cu130-aarch64-linux cell.

This PR:

  • Bumps the kernels-community/mamba-ssm entries in integrations/hub_kernels.py from version=3 to version=4 so Torch 2.12 aarch64 (SM121 / GH200 / GB200) users resolve the working v4 build (_mamba_ssm_cuda_4a743b8.abi3.so).
  • Adds Device(type="rocm", properties=ROCMProperties(min_capability=115, max_capability=115)) inference entries for causal_conv1d_fn, causal_conv1d_update, mamba_chunk_scan_combined, mamba_split_conv1d_scan_combined, and selective_state_update pointing at the newly published gfx1151 ROCm revision of Atlas-Inference/gdn.
  • Updates the Qwen3_5GatedDeltaNet gfx1151 pin to the same unified revision and adds a usage note with measured numbers to docs/source/en/model_doc/nemotron_h.md.

On gfx1151, the new Mamba-2 ops run the register-resident STATE=128 HIP kernels from Atlas Inference (mamba2_ssm_scan_reg with TB=16 LDS staging for prefill and mamba2_ssm_decode_reg for single-step decode) alongside the existing causal_conv1d_fwd and causal_conv1d_update kernels.

What does this PR do?

The gated rocm (gfx1151) entries for the Mamba-2 and causal_conv1d ops plus the mamba-ssm version=4 bump in integrations/hub_kernels.py and the docs note in docs/source/en/model_doc/nemotron_h.md. NemotronHMamba2Mixer and Mamba2Mixer are already decorated with @use_kernelized_func, so no modeling changes are needed.

Verified on Strix Halo (gfx1151, torch.cuda.get_device_capability() == (11, 5)), bf16, torch 2.14.1 + ROCm 7.2, and on DGX Spark GB10 (SM121), torch 2.12.1 + CUDA 13.0, kernels 0.17.2, transformers main:

  • Op-level parity against the reference implementations across ragged prefill lengths (1, 33, 65, 128, 257, 1024) at NVIDIA-Nemotron-3.5-Super-VL-120B Mamba-2 dimensions (H=128, P=64, N=128, G=8, D_conv=10240, K=4)]
  • Per Nemotron-3.5-Super-VL-120B Mamba-2 mixer layer on gfx1151, prefill is 3.56x faster at 4,096 tokens (718.9 ms -> 202.1 ms) and 3.07x faster at 8,192 tokens (1,439.6 ms -> 468.4 ms).
Prompt Tokens use_kernels=False TTFT use_kernels=True TTFT Prefill Speedup use_kernels=False Decode use_kernels=True Decode
1,024 0.64 s 0.38 s 1.67x 27.84 tok/s 31.54 tok/s (+13.3%)
4,096 2.71 s 1.65 s 1.64x 24.70 tok/s 29.22 tok/s (+18.3%)
8,192 6.36 s 3.85 s 1.65x 24.98 tok/s 31.22 tok/s (+25.0%)
16,384 18.32 s 9.38 s 1.95x 20.88 tok/s 23.02 tok/s (+10.2%)

Follow-up to #46423, #49127, #49452, and #49453. Just bumping the trust_remote_code=True note until Atlas-Inference is on the kernels trusted-publisher allowlist!

Code Agent Policy

  • I confirm that this is not a pure code agent PR.

Before submitting

Who can review?

@vasqu (merged #46423, #49127, #49452, #49453), @danieldk (mamba-ssm v4), @stevhliu (docs)

@github-actions

Copy link
Copy Markdown
Contributor

CI recap

Dashboard: View test results in Grafana
Latest run: 39096059815:2
Result: failure | Jobs: 15 | Tests: 205,390 | Failures: 1 | Duration: 17h 23m

@vasqu

vasqu commented Oct 19, 2026

Copy link
Copy Markdown
Collaborator

I'd like to avoid adding new kernels until #47773 is merged. Is that ok or do we need to update the existing kernels for the atlas kernels? Don't hesitate to ping me tho

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants