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, 19:45 UTC
[ROCm][Bugfix] Ensure that the moe_wna16_gemm kernel is not built on ROCm platforms. (#14629)
Signed-off-by: Sage Moore <sage@neuralmagic.com>
Sage Moore committed
Mar 12, 2025 at 05:00 UTC
45f3f3f59e4898a11baf9bcb8d6ec5db2581af34
4 files changed
+8
-3
CMakeLists.txt
+1
-1
@@ -559,7 +559,6 @@ target_compile_definitions(_C PRIVATE CUTLASS_ENABLE_DIRECT_CUDA_DRIVER_CALL=1)
559
set(VLLM_MOE_EXT_SRC
560
"csrc/moe/torch_bindings.cpp"
561
"csrc/moe/moe_align_sum_kernels.cu"
562
- "csrc/moe/moe_wna16.cu"
562
"csrc/moe/topk_softmax_kernels.cu")
563
564
set_gencode_flags_for_srcs(
@@ -574,6 +573,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
573
SRCS "${VLLM_MOE_WNA16_SRC}"
574
CUDA_ARCHS "${CUDA_ARCHS}")
575
576
+ list(APPEND VLLM_MOE_EXT_SRC "${VLLM_MOE_WNA16_SRC}")
577
cuda_archs_loose_intersection(MARLIN_MOE_ARCHS "8.0;8.6;8.7;8.9;9.0;10.0;10.1;12.0" "${CUDA_ARCHS}")
578
if (MARLIN_MOE_ARCHS)
579
set(MARLIN_MOE_SRC
csrc/moe/moe_ops.h
+2
-1
@@ -18,7 +18,7 @@ void sgl_moe_align_block_size(torch::Tensor topk_ids, int64_t num_experts,
18
torch::Tensor sorted_token_ids,
19
torch::Tensor experts_ids,
20
torch::Tensor num_tokens_post_pad);
21
-
21
+#ifndef USE_ROCM
22
torch::Tensor moe_wna16_gemm(torch::Tensor input, torch::Tensor output,
23
torch::Tensor b_qweight, torch::Tensor b_scales,
24
std::optional<torch::Tensor> b_qzeros,
@@ -28,3 +28,4 @@ torch::Tensor moe_wna16_gemm(torch::Tensor input, torch::Tensor output,
28
torch::Tensor num_tokens_post_pad, int64_t top_k,
29
int64_t BLOCK_SIZE_M, int64_t BLOCK_SIZE_N,
30
int64_t BLOCK_SIZE_K, int64_t bit);
31
+#endif
\ No newline at end of file
csrc/moe/torch_bindings.cpp
+1
-1
@@ -31,6 +31,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, m) {
31
" Tensor! num_tokens_post_pad) -> ()");
32
m.impl("sgl_moe_align_block_size", torch::kCUDA, &sgl_moe_align_block_size);
33
34
+#ifndef USE_ROCM
35
m.def(
36
"moe_wna16_gemm(Tensor input, Tensor! output, Tensor b_qweight, "
37
"Tensor b_scales, Tensor? b_qzeros, "
@@ -41,7 +42,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, m) {
42
43
m.impl("moe_wna16_gemm", torch::kCUDA, &moe_wna16_gemm);
44
44
-#ifndef USE_ROCM
45
m.def(
46
"marlin_gemm_moe(Tensor! a, Tensor! b_q_weights, Tensor! sorted_ids, "
47
"Tensor! topk_weights, Tensor! topk_ids, Tensor! b_scales, Tensor! "
vllm/_custom_ops.py
+4
@@ -1146,6 +1146,10 @@ def moe_wna16_gemm(input: torch.Tensor, output: torch.Tensor,
1146
num_tokens_post_pad: torch.Tensor, top_k: int,
1147
BLOCK_SIZE_M: int, BLOCK_SIZE_N: int, BLOCK_SIZE_K: int,
1148
bit: int) -> torch.Tensor:
1149
+ if not current_platform.is_cuda():
1150
+ raise NotImplementedError(
1151
+ "The optimized moe_wna16_gemm kernel is only "
1152
+ "available on CUDA platforms")
1153
torch.ops._moe_C.moe_wna16_gemm(input, output, b_qweight, b_scales,
1154
b_qzeros, topk_weights, sorted_token_ids,
1155
experts_ids, num_tokens_post_pad, top_k,