[12/n] final _C library kernel migration (#45415)
Chris Leonard committed
Jun 19, 2026 at 09:57 UTC
b9a7cd464c9ae9b1b450f8982b76d7be4de73724
17 files changed
+239
-212
CMakeLists.txt
+56
-72
@@ -319,82 +319,35 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
319
endif()
320
321
#
322
-# _C extension
322
+# Legacy _C extension (ROCm only — CUDA ops migrated to _C_stable_libtorch)
323
#
324
325
-set(VLLM_EXT_SRC
326
- "csrc/quantization/activation_kernels.cu"
327
- "csrc/torch_bindings.cpp")
328
-
329
-if(VLLM_GPU_LANG STREQUAL "CUDA")
330
- SET(CUTLASS_ENABLE_HEADERS_ONLY ON CACHE BOOL "Enable only the header library")
331
-
332
- # Set CUTLASS_REVISION. Used for FetchContent. Also fixes some bogus messages when building.
333
- set(CUTLASS_REVISION "v4.4.2")
334
-
335
- # Use the specified CUTLASS source directory for compilation if VLLM_CUTLASS_SRC_DIR is provided
336
- if (DEFINED ENV{VLLM_CUTLASS_SRC_DIR})
337
- set(VLLM_CUTLASS_SRC_DIR $ENV{VLLM_CUTLASS_SRC_DIR})
338
- endif()
339
-
340
- if(VLLM_CUTLASS_SRC_DIR)
341
- if(NOT IS_ABSOLUTE VLLM_CUTLASS_SRC_DIR)
342
- get_filename_component(VLLM_CUTLASS_SRC_DIR "${VLLM_CUTLASS_SRC_DIR}" ABSOLUTE)
343
- endif()
344
- message(STATUS "The VLLM_CUTLASS_SRC_DIR is set, using ${VLLM_CUTLASS_SRC_DIR} for compilation")
345
- FetchContent_Declare(cutlass SOURCE_DIR ${VLLM_CUTLASS_SRC_DIR})
346
- else()
347
- FetchContent_Declare(
348
- cutlass
349
- GIT_REPOSITORY https://github.com/nvidia/cutlass.git
350
- # Please keep this in sync with CUTLASS_REVISION line above.
351
- GIT_TAG ${CUTLASS_REVISION}
352
- GIT_PROGRESS TRUE
353
-
354
- # Speed up CUTLASS download by retrieving only the specified GIT_TAG instead of the history.
355
- # Important: If GIT_SHALLOW is enabled then GIT_TAG works only with branch names and tags.
356
- # So if the GIT_TAG above is updated to a commit hash, GIT_SHALLOW must be set to FALSE
357
- GIT_SHALLOW TRUE
358
- )
359
- endif()
360
- FetchContent_MakeAvailable(cutlass)
361
-
362
- set_gencode_flags_for_srcs(
363
- SRCS "${VLLM_EXT_SRC}"
364
- CUDA_ARCHS "${CUDA_ARCHS}")
365
-
366
-# if CUDA endif
367
-endif()
368
-
369
-if (VLLM_GPU_LANG STREQUAL "HIP")
370
- # Add QuickReduce kernels (ROCm-only; not part of stable ABI migration).
371
- # TODO: Remove the cuda_view when ROCm upgrade to torch 2.11.
372
- list(APPEND VLLM_EXT_SRC
325
+if(VLLM_GPU_LANG STREQUAL "HIP")
326
+ set(VLLM_EXT_SRC
327
+ "csrc/torch_bindings.cpp"
328
"csrc/custom_quickreduce.cu"
329
"csrc/cuda_view.cu"
375
- "csrc/libtorch_stable/cuda_utils_kernels.cu"
376
- )
377
-# if ROCM endif
378
-endif()
330
+ "csrc/libtorch_stable/cuda_utils_kernels.cu")
331
380
-message(STATUS "Enabling C extension.")
381
-define_extension_target(
382
- _C
383
- DESTINATION vllm
384
- LANGUAGE ${VLLM_GPU_LANG}
385
- SOURCES ${VLLM_EXT_SRC}
386
- COMPILE_FLAGS ${VLLM_GPU_FLAGS}
387
- ARCHITECTURES ${VLLM_GPU_ARCHES}
388
- INCLUDE_DIRECTORIES ${CUTLASS_INCLUDE_DIR}
389
- INCLUDE_DIRECTORIES ${CUTLASS_TOOLS_UTIL_INCLUDE_DIR}
390
- USE_SABI 3
391
- WITH_SOABI)
332
+ message(STATUS "Enabling C extension.")
333
+ define_extension_target(
334
+ _C
335
+ DESTINATION vllm
336
+ LANGUAGE ${VLLM_GPU_LANG}
337
+ SOURCES ${VLLM_EXT_SRC}
338
+ COMPILE_FLAGS ${VLLM_GPU_FLAGS}
339
+ ARCHITECTURES ${VLLM_GPU_ARCHES}
340
+ INCLUDE_DIRECTORIES ${CUTLASS_INCLUDE_DIR}
341
+ INCLUDE_DIRECTORIES ${CUTLASS_TOOLS_UTIL_INCLUDE_DIR}
342
+ USE_SABI 3
343
+ WITH_SOABI)
344
393
-# If CUTLASS is compiled on NVCC >= 12.5, it by default uses
394
-# cudaGetDriverEntryPointByVersion as a wrapper to avoid directly calling the
395
-# driver API. This causes problems when linking with earlier versions of CUDA.
396
-# Setting this variable sidesteps the issue by calling the driver directly.
397
-target_compile_definitions(_C PRIVATE CUTLASS_ENABLE_DIRECT_CUDA_DRIVER_CALL=1)
345
+ # If CUTLASS is compiled on NVCC >= 12.5, it by default uses
346
+ # cudaGetDriverEntryPointByVersion as a wrapper to avoid directly calling the
347
+ # driver API. This causes problems when linking with earlier versions of CUDA.
348
+ # Setting this variable sidesteps the issue by calling the driver directly.
349
+ target_compile_definitions(_C PRIVATE CUTLASS_ENABLE_DIRECT_CUDA_DRIVER_CALL=1)
350
+endif() # _C HIP endif
351
352
if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
353
#
@@ -403,6 +356,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
356
set(VLLM_STABLE_EXT_SRC
357
"csrc/libtorch_stable/torch_bindings.cpp"
358
"csrc/libtorch_stable/activation_kernels.cu"
359
+ "csrc/libtorch_stable/quantization/activation_kernels.cu"
360
"csrc/libtorch_stable/quantization/w8a8/int8/scaled_quant.cu"
361
"csrc/libtorch_stable/quantization/w8a8/fp8/common.cu"
362
"csrc/libtorch_stable/quantization/w8a8/fp8/per_token_group_quant.cu"
@@ -429,6 +383,38 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
383
"csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu")
384
385
if(VLLM_GPU_LANG STREQUAL "CUDA")
386
+ SET(CUTLASS_ENABLE_HEADERS_ONLY ON CACHE BOOL "Enable only the header library")
387
+
388
+ # Set CUTLASS_REVISION. Used for FetchContent. Also fixes some bogus messages when building.
389
+ set(CUTLASS_REVISION "v4.4.2")
390
+
391
+ # Use the specified CUTLASS source directory for compilation if VLLM_CUTLASS_SRC_DIR is provided
392
+ if (DEFINED ENV{VLLM_CUTLASS_SRC_DIR})
393
+ set(VLLM_CUTLASS_SRC_DIR $ENV{VLLM_CUTLASS_SRC_DIR})
394
+ endif()
395
+
396
+ if(VLLM_CUTLASS_SRC_DIR)
397
+ if(NOT IS_ABSOLUTE VLLM_CUTLASS_SRC_DIR)
398
+ get_filename_component(VLLM_CUTLASS_SRC_DIR "${VLLM_CUTLASS_SRC_DIR}" ABSOLUTE)
399
+ endif()
400
+ message(STATUS "The VLLM_CUTLASS_SRC_DIR is set, using ${VLLM_CUTLASS_SRC_DIR} for compilation")
401
+ FetchContent_Declare(cutlass SOURCE_DIR ${VLLM_CUTLASS_SRC_DIR})
402
+ else()
403
+ FetchContent_Declare(
404
+ cutlass
405
+ GIT_REPOSITORY https://github.com/nvidia/cutlass.git
406
+ # Please keep this in sync with CUTLASS_REVISION line above.
407
+ GIT_TAG ${CUTLASS_REVISION}
408
+ GIT_PROGRESS TRUE
409
+
410
+ # Speed up CUTLASS download by retrieving only the specified GIT_TAG instead of the history.
411
+ # Important: If GIT_SHALLOW is enabled then GIT_TAG works only with branch names and tags.
412
+ # So if the GIT_TAG above is updated to a commit hash, GIT_SHALLOW must be set to FALSE
413
+ GIT_SHALLOW TRUE
414
+ )
415
+ endif()
416
+ FetchContent_MakeAvailable(cutlass)
417
+
418
list(APPEND VLLM_STABLE_EXT_SRC
419
"csrc/libtorch_stable/cuda_view.cu"
420
"csrc/libtorch_stable/cuda_utils_kernels.cu"
@@ -929,7 +915,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
915
SRCS "${FP4_SM120_SRCS}"
916
CUDA_ARCHS "${FP4_SM120_ARCHS}")
917
list(APPEND VLLM_STABLE_EXT_SRC "${FP4_SM120_SRCS}")
932
- target_compile_definitions(_C PRIVATE ENABLE_NVFP4_SM120=1)
918
list(APPEND VLLM_GPU_FLAGS "-DENABLE_NVFP4_SM120=1")
919
list(APPEND VLLM_GPU_FLAGS "-DENABLE_CUTLASS_MOE_SM120=1")
920
message(STATUS "Building SM12x NVFP4 for archs: ${FP4_SM120_ARCHS}")
@@ -962,7 +947,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
947
SRCS "${FP4_SM100_SRCS}"
948
CUDA_ARCHS "${FP4_SM100_ARCHS}")
949
list(APPEND VLLM_STABLE_EXT_SRC "${FP4_SM100_SRCS}")
965
- target_compile_definitions(_C PRIVATE ENABLE_NVFP4_SM100=1)
950
list(APPEND VLLM_GPU_FLAGS "-DENABLE_NVFP4_SM100=1")
951
list(APPEND VLLM_GPU_FLAGS "-DENABLE_CUTLASS_MOE_SM100=1")
952
message(STATUS "Building SM10x/11x NVFP4/MXFP4 for archs: ${FP4_SM100_ARCHS}")
cmake/external_projects/qutlass.cmake
+29
-5
@@ -60,6 +60,7 @@ endif()
60
61
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND QUTLASS_ARCHS)
62
set(QUTLASS_SOURCES
63
+ csrc/qutlass_registration.cpp
64
${qutlass_SOURCE_DIR}/qutlass/csrc/bindings.cpp
65
${qutlass_SOURCE_DIR}/qutlass/csrc/gemm.cu
66
${qutlass_SOURCE_DIR}/qutlass/csrc/gemm_ada.cu
@@ -78,8 +79,19 @@ if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND QUTLASS_ARCHS)
79
80
if(CUTLASS_INCLUDE_DIR AND EXISTS "${CUTLASS_INCLUDE_DIR}/cutlass/cutlass.h")
81
list(APPEND QUTLASS_INCLUDES "${CUTLASS_INCLUDE_DIR}")
82
+ if(CUTLASS_TOOLS_UTIL_INCLUDE_DIR AND
83
+ EXISTS "${CUTLASS_TOOLS_UTIL_INCLUDE_DIR}/cutlass/util/packed_stride.hpp")
84
+ list(APPEND QUTLASS_INCLUDES "${CUTLASS_TOOLS_UTIL_INCLUDE_DIR}")
85
+ else()
86
+ get_filename_component(_qutlass_cutlass_root "${CUTLASS_INCLUDE_DIR}" DIRECTORY)
87
+ if(EXISTS "${_qutlass_cutlass_root}/tools/util/include/cutlass/util/packed_stride.hpp")
88
+ list(APPEND QUTLASS_INCLUDES "${_qutlass_cutlass_root}/tools/util/include")
89
+ endif()
90
+ endif()
91
elseif(EXISTS "${qutlass_SOURCE_DIR}/qutlass/third_party/cutlass/include/cutlass/cutlass.h")
82
- list(APPEND QUTLASS_INCLUDES "${qutlass_SOURCE_DIR}/qutlass/third_party/cutlass/include")
92
+ list(APPEND QUTLASS_INCLUDES
93
+ "${qutlass_SOURCE_DIR}/qutlass/third_party/cutlass/include"
94
+ "${qutlass_SOURCE_DIR}/qutlass/third_party/cutlass/tools/util/include")
95
message(STATUS "[QUTLASS] Using QuTLASS vendored CUTLASS headers (no vLLM CUTLASS detected).")
96
else()
97
message(FATAL_ERROR "[QUTLASS] CUTLASS headers not found. "
@@ -91,12 +103,23 @@ if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND QUTLASS_ARCHS)
103
CUDA_ARCHS "${QUTLASS_ARCHS}"
104
)
105
94
- target_sources(_C PRIVATE ${QUTLASS_SOURCES})
95
- target_include_directories(_C PRIVATE ${QUTLASS_INCLUDES})
96
- target_compile_definitions(_C PRIVATE
106
+ # QuTLASS uses legacy ATen headers and cannot be built with TORCH_TARGET_VERSION.
107
+ # Keep it as its own extension (registers torch.ops._qutlass_C).
108
+ define_extension_target(
109
+ _qutlass_C
110
+ DESTINATION vllm
111
+ LANGUAGE ${VLLM_GPU_LANG}
112
+ SOURCES ${QUTLASS_SOURCES}
113
+ COMPILE_FLAGS ${VLLM_GPU_FLAGS}
114
+ ARCHITECTURES ${VLLM_GPU_ARCHES}
115
+ INCLUDE_DIRECTORIES ${QUTLASS_INCLUDES}
116
+ USE_SABI 3
117
+ WITH_SOABI)
118
+
119
+ target_compile_definitions(_qutlass_C PRIVATE
120
QUTLASS_DISABLE_PYBIND=1
121
TARGET_CUDA_ARCH=${QUTLASS_TARGET_CC}
99
- )
122
+ CUTLASS_ENABLE_DIRECT_CUDA_DRIVER_CALL=1)
123
124
set_property(SOURCE ${QUTLASS_SOURCES} APPEND PROPERTY COMPILE_OPTIONS
125
$<$<COMPILE_LANGUAGE:CUDA>:--expt-relaxed-constexpr --use_fast_math -O3>
@@ -111,4 +134,5 @@ else()
134
"[QUTLASS] Skipping build: no supported arch (12.0f / 10.0f) found in "
135
"CUDA_ARCHS='${CUDA_ARCHS}'.")
136
endif()
137
+ add_custom_target(_qutlass_C)
138
endif()
csrc/libtorch_stable/core/math.hpp
renamed
csrc/libtorch_stable/moe/moe_align_sum_kernels.cu
+1
-1
@@ -9,7 +9,7 @@
9
#include <torch/headeronly/core/ScalarType.h>
10
11
#include "../../cuda_compat.h"
12
-#include "core/math.hpp"
12
+#include "libtorch_stable/core/math.hpp"
13
#include "libtorch_stable/dispatch_utils.h"
14
#include "libtorch_stable/torch_utils.h"
15
csrc/libtorch_stable/ops.h
+28
@@ -2,9 +2,25 @@
2
3
#include <torch/csrc/stable/library.h>
4
#include <torch/csrc/stable/tensor.h>
5
+#include <torch/headeronly/util/Exception.h>
6
7
#include <optional>
8
#include <string>
9
+#include <vector>
10
+
11
+#include <torch/csrc/stable/ops.h>
12
+
13
+inline torch::stable::Tensor weak_ref_tensor(torch::stable::Tensor& tensor) {
14
+ // Ensure tensor is on CUDA
15
+ STD_TORCH_CHECK(tensor.device().is_cuda(), "Tensor must be on CUDA device");
16
+
17
+ // Get the raw data pointer
18
+ void* data_ptr = tensor.mutable_data_ptr();
19
+
20
+ /// Create a new tensor from the raw data pointer
21
+ return torch::stable::from_blob(data_ptr, tensor.sizes(), tensor.strides(),
22
+ tensor.device(), tensor.scalar_type());
23
+}
24
25
void per_token_group_quant_fp8(const torch::stable::Tensor& input,
26
torch::stable::Tensor& output_q,
@@ -371,6 +387,18 @@ void silu_and_mul(torch::stable::Tensor& out, torch::stable::Tensor& input);
387
void silu_and_mul_clamp(torch::stable::Tensor& out,
388
torch::stable::Tensor& input, double limit,
389
double alpha = 1.0, double beta = 0.0);
390
+
391
+void silu_and_mul_quant(torch::stable::Tensor& out,
392
+ torch::stable::Tensor& input,
393
+ torch::stable::Tensor& scale);
394
+
395
+void persistent_masked_m_silu_mul_quant(
396
+ const torch::stable::Tensor& input, // (E, T, 2*H)
397
+ const torch::stable::Tensor& tokens_per_expert, // (E)
398
+ torch::stable::Tensor& y_q, // (E, T, H) [OUT]
399
+ torch::stable::Tensor& y_s, // (E, T, H//group_size) [OUT]
400
+ bool use_ue8m0);
401
+
402
void mul_and_silu(torch::stable::Tensor& out, torch::stable::Tensor& input);
403
void gelu_and_mul(torch::stable::Tensor& out, torch::stable::Tensor& input);
404
void gelu_tanh_and_mul(torch::stable::Tensor& out,
csrc/libtorch_stable/quantization/activation_kernels.cu
renamed
+63
-55
@@ -1,16 +1,12 @@
1
-#include <ATen/cuda/CUDAContext.h>
2
-#include <torch/all.h>
3
-#include <c10/cuda/CUDAGuard.h>
1
+#include "libtorch_stable/torch_utils.h"
2
3
#include <cmath>
6
-#include "core/math.hpp"
7
-#include "../cuda_compat.h"
8
-#include "dispatch_utils.h"
4
5
+#include "libtorch_stable/core/math.hpp"
6
+#include "cuda_compat.h"
7
+#include "libtorch_stable/dispatch_utils.h"
8
#include "quantization/w8a8/fp8/common.cuh"
9
12
-#include <c10/util/Float8_e4m3fn.h>
13
-
10
#ifndef USE_ROCM
11
#include <cuda_bf16.h>
12
#include <cuda_fp16.h>
@@ -33,7 +29,6 @@ typedef __hip_fp8x4_e4m3_fnuz __nv_fp8x4_e4m3;
29
#endif
30
#endif
31
36
-#include "core/registration.h"
32
namespace vllm {
33
34
template <typename T>
@@ -564,41 +559,47 @@ __global__ void silu_mul_fp8_quant_deep_gemm_kernel(
559
} // namespace vllm
560
561
// Launch activation, gating, and quantize kernel.
567
-#define LAUNCH_ACTIVATION_GATE_KERNEL(KERNEL) \
568
- int d = input.size(-1) / 2; \
569
- int64_t num_tokens = input.numel() / input.size(-1); \
570
- dim3 grid(num_tokens, num_tokens > 16 ? num_tokens > 32 ? 1 : 2 : 4); \
571
- dim3 block(std::min(d, 512)); \
572
- const at::cuda::OptionalCUDAGuard device_guard(device_of(input)); \
573
- const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); \
574
- VLLM_DISPATCH_FLOATING_TYPES( \
575
- input.scalar_type(), "act_and_mul_kernel", [&] { \
576
- VLLM_DISPATCH_FP8_TYPES( \
577
- out.scalar_type(), "fused_add_rms_norm_kernel_fp8_type", [&] { \
578
- vllm::act_and_mul_quant_kernel<scalar_t, KERNEL<scalar_t>, \
579
- fp8_t> \
580
- <<<grid, block, 0, stream>>>(out.data_ptr<fp8_t>(), \
581
- input.data_ptr<scalar_t>(), \
582
- scale.data_ptr<float>(), d); \
583
- }); \
562
+#define LAUNCH_ACTIVATION_GATE_KERNEL(KERNEL) \
563
+ int d = input.size(-1) / 2; \
564
+ int64_t num_tokens = input.numel() / input.size(-1); \
565
+ dim3 grid(num_tokens, num_tokens > 16 ? num_tokens > 32 ? 1 : 2 : 4); \
566
+ dim3 block(std::min(d, 512)); \
567
+ const torch::stable::accelerator::DeviceGuard device_guard( \
568
+ input.get_device_index()); \
569
+ const cudaStream_t stream = \
570
+ get_current_cuda_stream(input.get_device_index()); \
571
+ VLLM_STABLE_DISPATCH_FLOATING_TYPES( \
572
+ input.scalar_type(), "act_and_mul_kernel", [&] { \
573
+ VLLM_STABLE_DISPATCH_FP8_TYPES( \
574
+ out.scalar_type(), "act_and_mul_quant_kernel_fp8_type", [&] { \
575
+ vllm::act_and_mul_quant_kernel<scalar_t, KERNEL<scalar_t>, \
576
+ fp8_t> \
577
+ <<<grid, block, 0, stream>>>( \
578
+ out.mutable_data_ptr<fp8_t>(), \
579
+ input.const_data_ptr<scalar_t>(), \
580
+ scale.const_data_ptr<float>(), d); \
581
+ }); \
582
});
583
586
-void silu_and_mul_quant(torch::Tensor& out, // [..., d]
587
- torch::Tensor& input, // [..., 2 * d]
588
- torch::Tensor& scale) {
589
- TORCH_CHECK(out.dtype() == torch::kFloat8_e4m3fn ||
590
- out.dtype() == torch::kFloat8_e4m3fnuz);
591
- TORCH_CHECK(input.dtype() == torch::kFloat16 ||
592
- input.dtype() == torch::kBFloat16);
593
- TORCH_CHECK(input.size(-1) % 2 == 0);
584
+void silu_and_mul_quant(torch::stable::Tensor& out, // [..., d]
585
+ torch::stable::Tensor& input, // [..., 2 * d]
586
+ torch::stable::Tensor& scale) {
587
+ STD_TORCH_CHECK(
588
+ out.scalar_type() == torch::headeronly::ScalarType::Float8_e4m3fn ||
589
+ out.scalar_type() == torch::headeronly::ScalarType::Float8_e4m3fnuz);
590
+ STD_TORCH_CHECK(
591
+ input.scalar_type() == torch::headeronly::ScalarType::Half ||
592
+ input.scalar_type() == torch::headeronly::ScalarType::BFloat16,
593
+ "Input must be FP16 or BF16");
594
+ STD_TORCH_CHECK(input.size(-1) % 2 == 0);
595
LAUNCH_ACTIVATION_GATE_KERNEL(vllm::silu_kernel);
596
}
597
598
void persistent_masked_m_silu_mul_quant(
598
- const at::Tensor& input, // (E, T, 2*H)
599
- const at::Tensor& tokens_per_expert, // (E)
600
- at::Tensor& y_q, // (E, T, H) [OUT]
601
- at::Tensor& y_s, // (E, T, H//group_size) [OUT]
599
+ const torch::stable::Tensor& input, // (E, T, 2*H)
600
+ const torch::stable::Tensor& tokens_per_expert, // (E)
601
+ torch::stable::Tensor& y_q, // (E, T, H) [OUT]
602
+ torch::stable::Tensor& y_s, // (E, T, H//group_size) [OUT]
603
bool cast_scale_ue8m0) {
604
#ifndef USE_ROCM
605
@@ -606,14 +607,18 @@ void persistent_masked_m_silu_mul_quant(
607
// fixed GROUP_SIZE of 128.
608
static constexpr int GROUP_SIZE = 128;
609
609
- TORCH_CHECK(input.dtype() == torch::kBFloat16);
610
- TORCH_CHECK(y_q.dtype() == torch::kFloat8_e4m3fn ||
611
- y_q.dtype() == torch::kFloat8_e4m3fnuz);
612
- TORCH_CHECK(input.size(-1) % (GROUP_SIZE * 2) == 0);
610
+ STD_TORCH_CHECK(input.scalar_type() ==
611
+ torch::headeronly::ScalarType::BFloat16);
612
+ STD_TORCH_CHECK(
613
+ y_q.scalar_type() == torch::headeronly::ScalarType::Float8_e4m3fn ||
614
+ y_q.scalar_type() == torch::headeronly::ScalarType::Float8_e4m3fnuz);
615
+ STD_TORCH_CHECK(input.size(-1) % (GROUP_SIZE * 2) == 0);
616
617
bool const is_packed_ue8m0 =
615
- (y_s.dtype() == torch::kInt32 && cast_scale_ue8m0);
616
- TORCH_CHECK(y_s.dtype() == torch::kFloat32 || is_packed_ue8m0);
618
+ (y_s.scalar_type() == torch::headeronly::ScalarType::Int &&
619
+ cast_scale_ue8m0);
620
+ STD_TORCH_CHECK(y_s.scalar_type() == torch::headeronly::ScalarType::Float ||
621
+ is_packed_ue8m0);
622
623
using Idx_t = int64_t;
624
@@ -631,7 +636,7 @@ void persistent_masked_m_silu_mul_quant(
636
637
int const NUM_GROUPS = H / GROUP_SIZE;
638
634
- const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
639
+ const cudaStream_t stream = get_current_cuda_stream(input.get_device_index());
640
641
// TODO: Get this from cuda_arch ?
642
static constexpr int SILU_V2_BLOCK_COUNT = 132 * 32;
@@ -643,18 +648,21 @@ void persistent_masked_m_silu_mul_quant(
648
static constexpr int max_shared_mem_bytes = \
649
GROUP_SIZE * 2 * STAGES * NUM_WARPS * 2; \
650
dim3 grid(sms), block(THREAD_COUNT); \
646
- const at::cuda::OptionalCUDAGuard device_guard(device_of(input)); \
647
- VLLM_DISPATCH_FP8_TYPES( \
651
+ const torch::stable::accelerator::DeviceGuard device_guard( \
652
+ input.get_device_index()); \
653
+ VLLM_STABLE_DISPATCH_FP8_TYPES( \
654
y_q.scalar_type(), "silu_mul_fp8_quant_deep_gemm_kernel", [&] { \
655
vllm::silu_mul_fp8_quant_deep_gemm_kernel< \
656
BLOCK_COUNT, max_shared_mem_bytes, fp8_t, scale_t, THREAD_COUNT, \
657
Idx_t, CEIL_UE8M0, GROUP_SIZE, STAGES> \
658
<<<grid, block, max_shared_mem_bytes + (E + 1) * 16, stream>>>( \
653
- reinterpret_cast<__nv_bfloat16*>(input.data_ptr()), \
654
- (fp8_t*)y_q.data_ptr(), \
655
- reinterpret_cast<scale_t*>(y_s.data_ptr()), \
656
- reinterpret_cast<int32_t*>(tokens_per_expert.data_ptr()), E, \
657
- T, H, stride_i_e, stride_i_t, stride_i_h, stride_yq_e, \
659
+ reinterpret_cast<const __nv_bfloat16*>( \
660
+ input.const_data_ptr()), \
661
+ y_q.mutable_data_ptr<fp8_t>(), \
662
+ reinterpret_cast<scale_t*>(y_s.mutable_data_ptr()), \
663
+ reinterpret_cast<const int32_t*>( \
664
+ tokens_per_expert.const_data_ptr()), \
665
+ E, T, H, stride_i_e, stride_i_t, stride_i_h, stride_yq_e, \
666
stride_yq_t, stride_yq_h, STRIDE_YS_E, STRIDE_YS_T, \
667
STRIDE_YS_G, STRIDE_YS_P, stride_counts_e); \
668
});
@@ -679,7 +687,7 @@ void persistent_masked_m_silu_mul_quant(
687
Idx_t stride_ys_g = y_s.stride(2);
688
Idx_t stride_ys_p = 0;
689
if (!cast_scale_ue8m0) {
682
- TORCH_CHECK(!is_packed_ue8m0);
690
+ STD_TORCH_CHECK(!is_packed_ue8m0);
691
LAUNCH_ON_H(float, stride_ys_e, stride_ys_t, stride_ys_g, stride_ys_p,
692
false);
693
return;
@@ -692,8 +700,8 @@ void persistent_masked_m_silu_mul_quant(
700
return;
701
}
702
695
- TORCH_CHECK(cast_scale_ue8m0 && is_packed_ue8m0);
696
- TORCH_CHECK(y_s.dtype() == torch::kInt32);
703
+ STD_TORCH_CHECK(cast_scale_ue8m0 && is_packed_ue8m0);
704
+ STD_TORCH_CHECK(y_s.scalar_type() == torch::headeronly::ScalarType::Int);
705
706
// Int32 packed ue8m0 scales tensor.
707
// Let E, T, G be the number to experts, number of tokens and number of groups
csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_kernels.cu
+1
-1
@@ -31,7 +31,7 @@
31
32
#include "cutlass/util/packed_stride.hpp"
33
34
-#include "core/math.hpp"
34
+#include "libtorch_stable/core/math.hpp"
35
#include "core/batch_invariant.hpp"
36
37
using namespace cute;
csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_sm120_kernels.cu
+1
-1
@@ -31,7 +31,7 @@
31
32
#include "cutlass/util/packed_stride.hpp"
33
34
-#include "core/math.hpp"
34
+#include "libtorch_stable/core/math.hpp"
35
#include "core/batch_invariant.hpp"
36
37
using namespace cute;
csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/cutlass_gemm_caller.cuh
+1
-1
@@ -19,7 +19,7 @@
19
#include "cutlass/gemm/collective/collective_builder.hpp"
20
#include "cutlass/util/packed_stride.hpp"
21
22
-#include "core/math.hpp"
22
+#include "libtorch_stable/core/math.hpp"
23
#include "libtorch_stable/cutlass_extensions/common.hpp"
24
// clang-format on
25
csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm.cuh
+1
-1
@@ -14,7 +14,7 @@
14
#include "cutlass/epilogue/collective/collective_builder.hpp"
15
#include "cutlass/gemm/collective/collective_builder.hpp"
16
17
-#include "core/math.hpp"
17
+#include "libtorch_stable/core/math.hpp"
18
#include "libtorch_stable/cutlass_extensions/common.hpp"
19
// clang-format on
20
csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_c2x.cuh
+1
-1
@@ -22,7 +22,7 @@
22
#include "cutlass/epilogue/threadblock/fusion/visitors.hpp"
23
#include "cutlass/gemm/kernel/default_gemm_universal_with_visitor.h"
24
25
-#include "core/math.hpp"
25
+#include "libtorch_stable/core/math.hpp"
26
#include "libtorch_stable/cutlass_extensions/common.hpp"
27
// clang-format on
28
csrc/libtorch_stable/torch_bindings.cpp
+27
@@ -34,6 +34,20 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
34
// TODO: Remove this once ROCm upgrade to torch 2.11.
35
ops.def("get_cuda_view_from_cpu_tensor(Tensor cpu_tensor) -> Tensor");
36
37
+ // Note about marlin kernel 'workspace' arguments:
38
+ // Technically these should be mutable since they are modified by the kernel.
39
+ // But since they are set back to zero once the kernel is finished we can
40
+ // hand wave and say that they have no net effect.
41
+ //
42
+ // The reason to mark 'workspace' as immutable is so that they don't interfere
43
+ // with using ScalarType arguments in the ops. If they are marked as mutable,
44
+ // pytorch throws an assert in
45
+ // 'torch._higher_order_ops._register_effectful_op' that prevents these
46
+ // kernels from being torch.compile'd.
47
+ // See the following document for more info on custom types and ops that use
48
+ // custom types:
49
+ // https://docs.google.com/document/d/18fBMPuOJ0fY5ZQ6YyrHUppw9FA332CpNtgB6SOIgyuA
50
+
51
// Machete (Dense) Optimized Mixed Precision GEMM for Hopper.
52
ops.def(
53
"machete_supported_schedules("
@@ -480,6 +494,11 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
494
"Tensor workspace, int k, int max_seq_len) -> ()");
495
496
// Activation ops
497
+ ops.def(
498
+ "persistent_masked_m_silu_mul_quant(Tensor input, Tensor counts, Tensor! "
499
+ "y_q, Tensor! y_s, bool use_ue8m0) -> ()");
500
+ ops.def("weak_ref_tensor(Tensor input) -> Tensor");
501
+
502
// Activation function used in SwiGLU.
503
ops.def("silu_and_mul(Tensor! result, Tensor input) -> ()");
504
@@ -492,6 +511,10 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
511
"silu_and_mul_with_clamp(Tensor! result, Tensor input, float limit, "
512
"float alpha=1.0, float beta=0.0) -> ()");
513
514
+ // SwiGLU activation with FP8 quantization.
515
+ ops.def(
516
+ "silu_and_mul_quant(Tensor! result, Tensor input, Tensor scale) -> ()");
517
+
518
// Activation function used in GeGLU with `none` approximation.
519
ops.def("gelu_and_mul(Tensor! out, Tensor input) -> ()");
520
@@ -690,6 +713,10 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) {
713
ops.impl("persistent_topk", TORCH_BOX(&persistent_topk));
714
715
// Activation kernels (shared CUDA/ROCm)
716
+ ops.impl("persistent_masked_m_silu_mul_quant",
717
+ TORCH_BOX(&persistent_masked_m_silu_mul_quant));
718
+ ops.impl("weak_ref_tensor", TORCH_BOX(&weak_ref_tensor));
719
+ ops.impl("silu_and_mul_quant", TORCH_BOX(&silu_and_mul_quant));
720
ops.impl("silu_and_mul", TORCH_BOX(&silu_and_mul));
721
ops.impl("mul_and_silu", TORCH_BOX(&mul_and_silu));
722
ops.impl("gelu_and_mul", TORCH_BOX(&gelu_and_mul));
csrc/ops.h
-32
@@ -9,28 +9,6 @@
9
10
#include <vector>
11
12
-torch::Tensor weak_ref_tensor(torch::Tensor& tensor) {
13
- // Ensure tensor is on CUDA
14
- if (!tensor.is_cuda()) {
15
- throw std::runtime_error("Tensor must be on CUDA device");
16
- }
17
-
18
- // Get the raw data pointer
19
- void* data_ptr = tensor.data_ptr();
20
-
21
- // Get tensor sizes and strides
22
- std::vector<int64_t> sizes = tensor.sizes().vec();
23
- std::vector<int64_t> strides = tensor.strides().vec();
24
-
25
- // Get tensor options (dtype, device)
26
- auto options = tensor.options();
27
-
28
- // Create a new tensor from the raw data pointer
29
- auto new_tensor = torch::from_blob(data_ptr, sizes, strides, options);
30
-
31
- return new_tensor;
32
-}
33
-
12
// rms_norm and fused_add_rms_norm declarations also exist in
13
// csrc/libtorch_stable/ops.h (torch::stable ABI for CUDA). They remain here
14
// because the CPU build still uses these torch::Tensor declarations.
@@ -53,16 +31,6 @@ void silu_and_mul(torch::Tensor& out, torch::Tensor& input);
31
void silu_and_mul_clamp(torch::Tensor& out, torch::Tensor& input, double limit,
32
double alpha = 1.0, double beta = 0.0);
33
56
-void silu_and_mul_quant(torch::Tensor& out, torch::Tensor& input,
57
- torch::Tensor& scale);
58
-
59
-void persistent_masked_m_silu_mul_quant(
60
- const at::Tensor& input, // (E, T, 2*H)
61
- const at::Tensor& counts, // (E)
62
- at::Tensor& y_q, // (E, T, H) [OUT]
63
- at::Tensor& y_s, // (E, T, H//group_size) [OUT]
64
- bool use_ue8m0);
65
-
34
void gelu_and_mul(torch::Tensor& out, torch::Tensor& input);
35
36
void gelu_tanh_and_mul(torch::Tensor& out, torch::Tensor& input);
csrc/qutlass_registration.cpp
new
+5
@@ -0,0 +1,5 @@
1
+#include "core/registration.h"
2
+
3
+// QuTLASS registers torch.ops._qutlass_C via TORCH_LIBRARY in bindings.cpp.
4
+// This stub lets Python import vllm._qutlass_C to trigger op registration.
5
+REGISTER_EXTENSION(_qutlass_C)
csrc/torch_bindings.cpp
-40
@@ -20,17 +20,6 @@
20
21
TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
22
// vLLM custom ops
23
- //
24
-
25
- ops.def(
26
- "persistent_masked_m_silu_mul_quant(Tensor input, Tensor counts, Tensor! "
27
- "y_q, Tensor! y_s,"
28
- "bool use_ue8m0) -> ()");
29
- ops.impl("persistent_masked_m_silu_mul_quant", torch::kCUDA,
30
- &persistent_masked_m_silu_mul_quant);
31
-
32
- ops.def("weak_ref_tensor(Tensor input) -> Tensor");
33
- ops.impl("weak_ref_tensor", torch::kCUDA, &weak_ref_tensor);
23
24
#ifdef USE_ROCM
25
// TODO: Remove this once we upgrade to torch 2.11.
@@ -39,35 +28,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
28
ops.def("get_cuda_view_from_cpu_tensor(Tensor cpu_tensor) -> Tensor");
29
ops.impl("get_cuda_view_from_cpu_tensor", torch::kCPU,
30
&get_cuda_view_from_cpu_tensor);
42
-#endif
43
-
44
- // Activation ops (quantized only — basic ops moved to _C_stable_libtorch)
45
- ops.def(
46
- "silu_and_mul_quant(Tensor! result, Tensor input, Tensor scale) -> ()");
47
- ops.impl("silu_and_mul_quant", torch::kCUDA, &silu_and_mul_quant);
48
-
49
- // Horizontally-fused DeepseekV4-MLA: per-head RMSNorm + GPT-J RoPE for Q, and
50
- // GPT-J RoPE + UE8M0 FP8 quant + paged cache insert for KV, all in one
51
- // kernel launch. Registered in _C_stable_libtorch (incl. the FlashInfer V4
52
- // full-cache bf16/fp8 variants).
53
-
54
- // Quantization ops
55
-#ifndef USE_ROCM
56
-
57
- // Note about marlin kernel 'workspace' arguments:
58
- // Technically these should be mutable since they are modified by the kernel.
59
- // But since they are set back to zero once the kernel is finished we can
60
- // hand wave and say that they have no net effect.
61
- //
62
- // The reason to mark 'workspace' as immutable is so that they don't interfere
63
- // with using ScalarType arguments in the ops. If they are marked as mutable,
64
- // pytorch throws an assert in
65
- // 'torch._higher_order_ops._register_effectful_op' that prevents these
66
- // kernels from being torch.compile'd.
67
- // See the following document for more info on custom types and ops that use
68
- // custom types:
69
- // https://docs.google.com/document/d/18fBMPuOJ0fY5ZQ6YyrHUppw9FA332CpNtgB6SOIgyuA
70
-
31
#endif
32
}
33
setup.py
+4
-1
@@ -769,6 +769,7 @@ class precompiled_wheel_utils:
769
"vllm/_C.abi3.so",
770
"vllm/_C_stable_libtorch.abi3.so",
771
"vllm/_moe_C_stable_libtorch.abi3.so",
772
+ "vllm/_qutlass_C.abi3.so",
773
"vllm/_flashmla_C.abi3.so",
774
"vllm/_flashmla_extension_C.abi3.so",
775
"vllm/_sparse_flashmla_C.abi3.so",
@@ -1135,6 +1136,7 @@ if _is_cuda():
1136
# DeepGEMM requires CUDA 12.3+ (SM90/SM100)
1137
# Optional since it won't build on unsupported architectures
1138
ext_modules.append(CMakeExtension(name="vllm._deep_gemm_C", optional=True))
1139
+ ext_modules.append(CMakeExtension(name="vllm._qutlass_C", optional=True))
1140
# fmha_sm100 is a Python/CuTe-DSL package installed into vllm.third_party.
1141
ext_modules.append(CMakeExtension(name="vllm.fmha_sm100", optional=True))
1142
@@ -1149,7 +1151,8 @@ if _is_cpu():
1151
ext_modules.append(CMakeExtension(name="vllm._C"))
1152
1153
if _build_custom_ops():
1152
- ext_modules.append(CMakeExtension(name="vllm._C"))
1154
+ if _is_hip():
1155
+ ext_modules.append(CMakeExtension(name="vllm._C"))
1156
if _is_cuda() or _is_hip():
1157
ext_modules.append(CMakeExtension(name="vllm._C_stable_libtorch"))
1158
ext_modules.append(CMakeExtension(name="vllm._moe_C_stable_libtorch"))
vllm/platforms/cuda.py
+21
-1
@@ -19,7 +19,6 @@ from torch.distributed.distributed_c10d import is_nccl_available
19
from typing_extensions import ParamSpec
20
21
# import custom ops, trigger op registration
22
-import vllm._C # noqa
22
import vllm._C_stable_libtorch # noqa
23
import vllm.envs as envs
24
from vllm.logger import init_logger
@@ -40,6 +39,11 @@ else:
39
40
logger = init_logger(__name__)
41
42
+try:
43
+ import vllm._qutlass_C # noqa: F401
44
+except ImportError as e:
45
+ logger.warning("Failed to import from vllm._qutlass_C: %r", e)
46
+
47
_P = ParamSpec("_P")
48
_R = TypeVar("_R")
49
@@ -187,6 +191,22 @@ class CudaPlatformBase(Platform):
191
"RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES",
192
]
193
194
+ @classmethod
195
+ def import_kernels(cls) -> None:
196
+ """Import CUDA kernel extensions (_C_stable_libtorch, optional _qutlass_C)."""
197
+ try:
198
+ import vllm._C_stable_libtorch # noqa: F401
199
+ except ImportError as e:
200
+ logger.warning("Failed to import from vllm._C_stable_libtorch: %r", e)
201
+ try:
202
+ import vllm._moe_C_stable_libtorch # noqa: F401
203
+ except ImportError as e:
204
+ logger.warning("Failed to import from vllm._moe_C_stable_libtorch: %r", e)
205
+ try:
206
+ import vllm._qutlass_C # noqa: F401
207
+ except ImportError as e:
208
+ logger.warning("Failed to import from vllm._qutlass_C: %r", e)
209
+
210
@property
211
def supported_dtypes(self) -> list[torch.dtype]:
212
if self.has_device_capability(80):