[ROCm] [Bugfix] [Critical]: Fix mamba compilation bug #20883

tjtanaa · 2025-07-13T16:41:41Z

Essential Elements of an Effective PR Description Checklist

The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
The test plan, such as providing test command.
The test results, such as pasting the results comparison before and after, or e2e results
(Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.

Purpose

Fix this bug #20882 which breaks vLLM package installation on ROCm
Break introduced in commit 2c11a73

Bugfixes is a referenced to (#18565)

FIX #20882

Test Plan

Able to install on ROCm.
Run mamba unit tests

tests/kernels/mamba/test_causal_conv1d.py
tests/kernels/mamba/test_mamba_mixer2.py
tests/kernels/mamba/test_mamba_ssm_ssd.py
tests/kernels/mamba/test_mamba_ssm.py

Test Result

Compilable
tests/kernels/mamba/test_causal_conv1d.py: PASSED ==== 244 passed in 95.18s (0:01:35) ====
tests/kernels/mamba/test_mamba_mixer2.py: SKIPPED ==== 4 skipped in 13.36s ====
tests/kernels/mamba/test_mamba_ssm.py: PASSED ==== 462 passed, 2 skipped in 151.27s (0:02:31) ====
tests/kernels/mamba/test_mamba_ssm_ssd.py::test_mamba_chunk_scan_single_example: PASSED
tests/kernels/mamba/test_mamba_ssm_ssd.py::test_mamba_chunk_scan_cont_batch: This test has been failing before the commit that breaks the compilation (commit: 2c11a73)

Validated the same unit tests with commit b639327ad94b3aa16022ebea49f8e525660b736b (this commit is before commit: 2c11a73)

For fixing this critical issue, it is suggested that we accept a partial fix. The core dump in unitt test is tracked at here. #20885

lm_eval score of microsoft/Phi-4-mini-reasoning V1 Engine
vllm (pretrained=microsoft/Phi-4-mini-reasoning,tensor_parallel_size=1,max_model_len=32768,trust_remote_code=True), gen_kwargs: (None), limit: None, num_fewshot: 5, batch_size: auto

Tasks	Version	Filter	n-shot	Metric		Value		Stderr
gsm8k	3	flexible-extract	5	exact_match	↑	0.4761	±	0.0138
		strict-match	5	exact_match	↑	0.4738	±	0.0138

(Optional) Documentation Update

Signed-off-by: tjtanaa <[email protected]>

github-actions · 2025-07-13T16:41:50Z

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run fastcheck CI which starts running only a small and essential subset of CI tests to quickly catch errors. You can run other CI tests on top of those by going to your fastcheck build on Buildkite UI (linked in the PR checks section) and unblock them. If you do not have permission to unblock, ping simon-mo or khluu to add you in our Buildkite org.

Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging.

To run CI, PR reviewers can either: Add ready label to the PR or enable auto-merge.

🚀

gemini-code-assist

Summary of Changes

Hello @tjtanaa, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request delivers a critical bug fix that addresses compilation issues for the Mamba kernel on ROCm, which was blocking vLLM package installation. It also extends functionality by unblocking speculative decoding for vLLM V1 on ROCm, enhancing the platform's capabilities on AMD hardware.

Highlights

ROCm Compatibility Fixes: Implemented conditional compilation within the Mamba SSM kernel (selective_scan_fwd.cu) to correctly use HIP-specific error handling macros (C10_HIP_CHECK, C10_HIP_KERNEL_LAUNCH_CHECK) and function attribute setting (hipFuncSetAttribute) when building for ROCm. This ensures the kernel compiles and runs correctly on AMD GPUs, resolving a critical bug that prevented vLLM installation on ROCm.
Speculative Decoding Unblocked for vLLM V1 on ROCm: Removed a NotImplementedError in vllm/platforms/rocm.py that previously prevented speculative decoding for vLLM V1 on ROCm. The configuration now correctly assigns the vllm.v1.worker.gpu_worker.Worker class, enabling this feature on ROCm platforms.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.

Feature	Command	Description
Code Review	`/gemini review`	Performs a code review for the current pull request in its current state.
Pull Request Summary	`/gemini summary`	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in issue comments and review comments.
Help	`/gemini help`	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist is currently in preview and may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments to provide feedback.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

gemini-code-assist

Code Review

This pull request aims to fix a Mamba compilation bug on ROCm. The changes include using conditional compilation for CUDA/HIP-specific headers and function calls in selective_scan_fwd.cu. Additionally, it enables speculative decoding for vLLM V1 on ROCm by removing a NotImplementedError in rocm.py.

The changes are generally good and move towards better ROCm support. However, I've identified a critical issue in selective_scan_fwd.cu where a CUDA-specific macro is used unconditionally, which will lead to compilation failures on ROCm. I've provided a comment with details on how to fix this.

gemini-code-assist · 2025-07-13T16:43:35Z

csrc/mamba/mamba_ssm/selective_scan_fwd.cu

+#ifdef USE_ROCM
+                    C10_HIP_CHECK(hipFuncSetAttribute(
+                        reinterpret_cast<const void*>(kernel), hipFuncAttributeMaxDynamicSharedMemorySize, kSmemSize));
+#else
                    C10_CUDA_CHECK(cudaFuncSetAttribute(
                        kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, kSmemSize));
+#endif


While you've correctly handled cudaFuncSetAttribute with a conditional for ROCm, the C10_CUDA_KERNEL_LAUNCH_CHECK() on line 336 is still CUDA-specific and will cause a compilation error on ROCm builds. This should also be wrapped in an #ifdef USE_ROCM block to call C10_HIP_KERNEL_LAUNCH_CHECK() for ROCm.

Signed-off-by: tjtanaa <[email protected]>

mgoin

Looks reasonable to me, thanks for the context

…-bug

…0883) Signed-off-by: tjtanaa <[email protected]> Co-authored-by: vllmellm <[email protected]> Signed-off-by: x22x22 <[email protected]>

…0883) Signed-off-by: tjtanaa <[email protected]> Co-authored-by: vllmellm <[email protected]>

…0883) Signed-off-by: tjtanaa <[email protected]> Co-authored-by: vllmellm <[email protected]> Signed-off-by: Jinzhen Lin <[email protected]>

…0883) Signed-off-by: tjtanaa <[email protected]> Co-authored-by: vllmellm <[email protected]> Signed-off-by: Paul Pak <[email protected]>

…0883) Signed-off-by: tjtanaa <[email protected]> Co-authored-by: vllmellm <[email protected]> Signed-off-by: Diego-Castan <[email protected]>

…0883) Signed-off-by: tjtanaa <[email protected]> Co-authored-by: vllmellm <[email protected]>

tjtanaa added 4 commits July 4, 2025 16:32

enable MTP V1 ROCm

f5d3b26

Signed-off-by: tjtanaa <[email protected]>

Merge remote-tracking branch 'origin/main' into mtp-v1

d2e01bf

Merge remote-tracking branch 'origin/main' into mtp-v1

f184e89

fix mamba_ssm compilation bug on ROCm

161e909

Signed-off-by: tjtanaa <[email protected]>

gemini-code-assist bot reviewed Jul 13, 2025

View reviewed changes

mergify bot added the rocm Related to AMD ROCm label Jul 13, 2025

gemini-code-assist bot reviewed Jul 13, 2025

View reviewed changes

revert changes to platforms/rocm.py

fd70bc9

Signed-off-by: tjtanaa <[email protected]>

tjtanaa marked this pull request as ready for review July 13, 2025 17:46

mgoin approved these changes Jul 13, 2025

View reviewed changes

mgoin added bug Something isn't working ready ONLY add when PR is ready to merge/full CI is needed labels Jul 13, 2025

Merge remote-tracking branch 'origin/main' into fix-mamba-compilation…

835c662

…-bug

DarkLight1337 merged commit c488b92 into vllm-project:main Jul 14, 2025
98 checks passed

DarkLight1337 mentioned this pull request Jul 14, 2025

[Bug] [Critical] [ROCm]: Compilation error during vLLM package installation #20882

Closed

1 task

DarkLight1337 added this to the v0.10.0 milestone Jul 14, 2025

Pradyun92 pushed a commit to Pradyun92/vllm that referenced this pull request Aug 6, 2025

[ROCm] [Bugfix] [Critical]: Fix mamba compilation bug (vllm-project#2…

ec7efd7

…0883) Signed-off-by: tjtanaa <[email protected]> Co-authored-by: vllmellm <[email protected]>

npanpaliya pushed a commit to odh-on-pz/vllm-upstream that referenced this pull request Aug 6, 2025

[ROCm] [Bugfix] [Critical]: Fix mamba compilation bug (vllm-project#2…

c288e8d

…0883) Signed-off-by: tjtanaa <[email protected]> Co-authored-by: vllmellm <[email protected]>

epwalsh pushed a commit to epwalsh/vllm that referenced this pull request Aug 27, 2025

[ROCm] [Bugfix] [Critical]: Fix mamba compilation bug (vllm-project#2…

f506c2d

…0883) Signed-off-by: tjtanaa <[email protected]> Co-authored-by: vllmellm <[email protected]>

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Uh oh!

[ROCm] [Bugfix] [Critical]: Fix mamba compilation bug #20883

[ROCm] [Bugfix] [Critical]: Fix mamba compilation bug #20883

Uh oh!

tjtanaa commented Jul 13, 2025 •

edited by github-actions bot

Loading

Uh oh!

github-actions bot commented Jul 13, 2025

Uh oh!

gemini-code-assist bot left a comment

Uh oh!

gemini-code-assist bot left a comment

Uh oh!

gemini-code-assist bot Jul 13, 2025

Uh oh!

mgoin left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

[ROCm] [Bugfix] [Critical]: Fix mamba compilation bug #20883

[ROCm] [Bugfix] [Critical]: Fix mamba compilation bug #20883

Uh oh!

Conversation

tjtanaa commented Jul 13, 2025 • edited by github-actions bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Essential Elements of an Effective PR Description Checklist

Purpose

Test Plan

Test Result

(Optional) Documentation Update

Uh oh!

github-actions bot commented Jul 13, 2025

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Summary of Changes

Highlights

Footnotes

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

gemini-code-assist bot Jul 13, 2025

Choose a reason for hiding this comment

Uh oh!

mgoin left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

tjtanaa commented Jul 13, 2025 •

edited by github-actions bot

Loading