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 7, 21:30 UTC
[Bugfix] Fix DSV3 kernels breaking _C and _moe_C on unsupported arches (#35123)
Signed-off-by: mgoin <mgoin64@gmail.com>
Michael Goin committed
Feb 23, 2026 at 20:11 UTC
3ef9fd0f989d7dfc5e3633c0dcfc64acbaf2f8c5
5 files changed
+12
-3
CMakeLists.txt
-1
@@ -783,7 +783,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
783
SRCS "${DSV3_FUSED_A_GEMM_SRC}"
784
CUDA_ARCHS "${DSV3_FUSED_A_GEMM_ARCHS}")
785
list(APPEND VLLM_EXT_SRC ${DSV3_FUSED_A_GEMM_SRC})
786
- list(APPEND VLLM_GPU_FLAGS "-DENABLE_DSV3_FUSED_A_GEMM=1")
786
message(STATUS "Building dsv3_fused_a_gemm for archs: ${DSV3_FUSED_A_GEMM_ARCHS}")
787
else()
788
message(STATUS "Not building dsv3_fused_a_gemm as no compatible archs found "
csrc/dsv3_fused_a_gemm.cu
+4
@@ -745,3 +745,7 @@ void dsv3_fused_a_gemm(torch::Tensor& output, torch::Tensor const& mat_a,
745
stream);
746
}
747
}
748
+
749
+TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, CUDA, m) {
750
+ m.impl("dsv3_fused_a_gemm", &dsv3_fused_a_gemm);
751
+}
csrc/moe/dsv3_router_gemm_entry.cu
+6
@@ -20,10 +20,12 @@
20
21
#include <ATen/ATen.h>
22
#include <ATen/cuda/CUDAContext.h>
23
+#include <torch/all.h>
24
25
#include <cuda_bf16.h>
26
#include <cuda_runtime.h>
27
28
+#include "core/registration.h"
29
#include "dsv3_router_gemm_utils.h"
30
31
static constexpr int DEFAULT_NUM_EXPERTS = 256;
@@ -161,3 +163,7 @@ void dsv3_router_gemm(at::Tensor& output, // [num_tokens, num_experts]
163
}
164
}
165
}
166
+
167
+TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, CUDA, m) {
168
+ m.impl("dsv3_router_gemm", &dsv3_router_gemm);
169
+}
csrc/moe/torch_bindings.cpp
+1
-1
@@ -127,7 +127,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, m) {
127
128
// DeepSeek V3 optimized router GEMM for SM90+
129
m.def("dsv3_router_gemm(Tensor! output, Tensor mat_a, Tensor mat_b) -> ()");
130
- m.impl("dsv3_router_gemm", torch::kCUDA, &dsv3_router_gemm);
130
+ // conditionally compiled so impl registration is in source file
131
#endif
132
}
133
csrc/torch_bindings.cpp
+1
-1
@@ -242,7 +242,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
242
// DeepSeek V3 fused A GEMM (SM 9.0+, bf16 only, 1-16 tokens).
243
ops.def(
244
"dsv3_fused_a_gemm(Tensor! output, Tensor mat_a, Tensor mat_b) -> ()");
245
- ops.impl("dsv3_fused_a_gemm", torch::kCUDA, &dsv3_fused_a_gemm);
245
+ // conditionally compiled so impl registration is in source file
246
247
// Quantized GEMM for AWQ.
248
ops.def(