A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 15, 02:15 UTC
[ROCm] [Bugfix] [Critical]: Fix mamba compilation bug (#20883)
Signed-off-by: tjtanaa <tunjian.tan@embeddedllm.com> Co-authored-by: vllmellm <vllm.ellm@embeddedllm.com>
TJian committed
Jul 14, 2025 at 00:23 UTC
c488b928a736109cc0a0c824340b928a4f118b2f
1 file changed
+10
-1
csrc/mamba/mamba_ssm/selective_scan_fwd.cu
+10
-1
@@ -7,7 +7,11 @@
7
8
#include <c10/util/BFloat16.h>
9
#include <c10/util/Half.h>
10
-#include <c10/cuda/CUDAException.h> // For C10_CUDA_CHECK and C10_CUDA_KERNEL_LAUNCH_CHECK
10
+#ifdef USE_ROCM
11
+ #include <c10/hip/HIPException.h> // For C10_HIP_CHECK and C10_HIP_KERNEL_LAUNCH_CHECK
12
+#else
13
+ #include <c10/cuda/CUDAException.h> // For C10_CUDA_CHECK and C10_CUDA_KERNEL_LAUNCH_CHECK
14
+#endif
15
16
#ifndef USE_ROCM
17
#include <cub/block/block_load.cuh>
@@ -320,8 +324,13 @@ void selective_scan_fwd_launch(SSMParamsBase ¶ms, cudaStream_t stream) {
324
dim3 grid(params.batch, params.dim / kNRows);
325
auto kernel = &selective_scan_fwd_kernel<Ktraits>;
326
if (kSmemSize >= 48 * 1024) {
327
+#ifdef USE_ROCM
328
+ C10_HIP_CHECK(hipFuncSetAttribute(
329
+ reinterpret_cast<const void*>(kernel), hipFuncAttributeMaxDynamicSharedMemorySize, kSmemSize));
330
+#else
331
C10_CUDA_CHECK(cudaFuncSetAttribute(
332
kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, kSmemSize));
333
+#endif
334
}
335
kernel<<<grid, Ktraits::kNThreads, kSmemSize, stream>>>(params);
336
C10_CUDA_KERNEL_LAUNCH_CHECK();