Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 6 commits into from
Jul 14, 2025
Merged

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

Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
f5d3b26
enable MTP V1 ROCm
tjtanaa Jul 4, 2025
d2e01bf
Merge remote-tracking branch 'origin/main' into mtp-v1
tjtanaa Jul 7, 2025
f184e89
Merge remote-tracking branch 'origin/main' into mtp-v1
tjtanaa Jul 13, 2025
161e909
fix mamba_ssm compilation bug on ROCm
tjtanaa Jul 13, 2025
fd70bc9
revert changes to platforms/rocm.py
tjtanaa Jul 13, 2025
835c662
Merge remote-tracking branch 'origin/main' into fix-mamba-compilation…
vllmellm Jul 14, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 10 additions & 1 deletion csrc/mamba/mamba_ssm/selective_scan_fwd.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,11 @@

#include <c10/util/BFloat16.h>
#include <c10/util/Half.h>
#include <c10/cuda/CUDAException.h> // For C10_CUDA_CHECK and C10_CUDA_KERNEL_LAUNCH_CHECK
#ifdef USE_ROCM
#include <c10/hip/HIPException.h> // For C10_HIP_CHECK and C10_HIP_KERNEL_LAUNCH_CHECK
#else
#include <c10/cuda/CUDAException.h> // For C10_CUDA_CHECK and C10_CUDA_KERNEL_LAUNCH_CHECK
#endif

#ifndef USE_ROCM
#include <cub/block/block_load.cuh>
Expand Down Expand Up @@ -320,8 +324,13 @@ void selective_scan_fwd_launch(SSMParamsBase &params, cudaStream_t stream) {
dim3 grid(params.batch, params.dim / kNRows);
auto kernel = &selective_scan_fwd_kernel<Ktraits>;
if (kSmemSize >= 48 * 1024) {
#ifdef USE_ROCM
C10_HIP_CHECK(hipFuncSetAttribute(
reinterpret_cast<const void*>(kernel), hipFuncAttributeMaxDynamicSharedMemorySize, kSmemSize));
#else
C10_CUDA_CHECK(cudaFuncSetAttribute(
kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, kSmemSize));
#endif
Comment on lines +327 to +333
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

critical

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.

}
kernel<<<grid, Ktraits::kNThreads, kSmemSize, stream>>>(params);
C10_CUDA_KERNEL_LAUNCH_CHECK();
Expand Down