[kashif]

Related issue

n/a

What does this PR do?

Adds a ROCm build of mamba-ssm (selective scan + causal conv1d).

Motivation

The sources already carry upstream's USE_ROCM paths, so the only thing missing was the build config. This lets transformers' Mamba-family models use the hub kernel on AMD GPUs.

Changes

• Add the rocm backend and cpp-rocm kernels for selective scan and causal conv1d.
• reverse_scan.cuh: newer hipCUB dropped HIPCUB_WARP_THREADS (rocprim::warp_size() is gone), use HIPCUB_DEVICE_WARP_THREADS when it exists.
• reverse_scan.cuh: lane id was hard-coded to 32-wide warps (threadIdx.x & 0x1f), which breaks the backward reverse scan on wave64 (gfx9/MI300). Mask with the warp size instead; no change on CUDA.
• test_transformers_compile.py: map layers to the rocm device on HIP, otherwise kernelize silently falls back.

Testing

• Radeon 890M (gfx1150), torch 2.15 + ROCm: pytest -m kernels_ci gives 646 passed, 3 failed.
• The selective scan, causal conv1d and transformers compile tests all pass.
• The 3 failures are bf16/fp16 tolerance misses in the Triton SSD/state-update tests (max diff ~0.009 vs atol 3e-3), not in the compiled ops.
• The wave64 fix is only tested on wave32 (gfx1150); it'd be good if someone could run the selective scan bwd tests on MI300.

Checklist

• This PR is linked to an issue that was discussed and approved
• I have tested these changes locally
• New/changed functionality has test coverage
• LLM disclosure:
  • I did not use an LLM to create this PR.
  • I used and LLM for assistance while creating this PR.
  • This PR was mostly or completely generated by an LLM.

[github-actions]

This PR changes source code for the following kernel(s): mamba-ssm.

If you have the kernel-bot build permissions, then trigger /kernel-bot build mamba-ssm yourself. Otherwise, a maintainer will review the PR and will take an appropriate action. Thanks for your patience.

[drbh]

/kernel-bot build mamba-ssm

[github-actions]

Build request processed.

Command: /kernel-bot build mamba-ssm
Mode: build only
Target branch: pr-1308
PR head SHA: 8f8e4f30b86e49902b5f9523b9d11e1aa8822769
Workflows: build.yaml, build-mac.yaml, build-windows.yaml

Dispatched (1):

• mamba-ssm (build.yaml): https://github.com/huggingface/kernels-community/actions/runs/42638268582

[drbh]

/kernel-bot build mamba-ssm

[github-actions]

Build request processed.

Command: /kernel-bot build mamba-ssm
Mode: build only
Target branch: pr-1308
PR head SHA: 16ecce8c060d1f6dfff50cb837c3ea25bdd60aa8
Workflows: build.yaml, build-mac.yaml, build-windows.yaml

Dispatched (1):

• mamba-ssm (build.yaml): https://github.com/huggingface/kernels-community/actions/runs/42662488168

[drbh]

lgtm! the ci error should now be resolved by huggingface/kernels#984 and builds look good

[drbh]

/kernel-bot merge-and-upload mamba-ssm

[github-actions]

Build request processed.

Command: /kernel-bot merge-and-upload mamba-ssm
Mode: merge, build and upload
Target branch: ``
PR head SHA: 16ecce8c060d1f6dfff50cb837c3ea25bdd60aa8
Workflows: `build.yaml, build-mac.yaml, build-windows.yaml`

Merge result: Pull Request successfully merged

Dispatched (1):

• mamba-ssm (build.yaml): https://github.com/huggingface/kernels-community/actions/runs/42882568010

Hub uploads:

• Hub repo: https://huggingface.co/kernels/kernels-community/mamba-ssm