iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[12/n] final _C library kernel migration (#45415)

Browse Source
This commit is contained in:
Chris Leonard
2026-06-19 06:57:26 -07:00
committed by GitHub
parent 69bdd34542
commit b9a7cd464c
17 changed files with 239 additions and 212 deletions
Download Patch File Download Diff File Expand all files Collapse all files
CMakeLists.txt
+56 -72
View File
@@ -319,82 +319,35 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
endif()
#
# _C extension
# Legacy _C extension (ROCm only — CUDA ops migrated to _C_stable_libtorch)
#
set(VLLM_EXT_SRC
"csrc/quantization/activation_kernels.cu"
"csrc/torch_bindings.cpp")
if(VLLM_GPU_LANG STREQUAL "CUDA")
SET(CUTLASS_ENABLE_HEADERS_ONLY ON CACHE BOOL "Enable only the header library")
# Set CUTLASS_REVISION. Used for FetchContent. Also fixes some bogus messages when building.
set(CUTLASS_REVISION "v4.4.2")
# Use the specified CUTLASS source directory for compilation if VLLM_CUTLASS_SRC_DIR is provided
if (DEFINED ENV{VLLM_CUTLASS_SRC_DIR})
set(VLLM_CUTLASS_SRC_DIR $ENV{VLLM_CUTLASS_SRC_DIR})
endif()
if(VLLM_CUTLASS_SRC_DIR)
if(NOT IS_ABSOLUTE VLLM_CUTLASS_SRC_DIR)
get_filename_component(VLLM_CUTLASS_SRC_DIR "${VLLM_CUTLASS_SRC_DIR}" ABSOLUTE)
endif()
message(STATUS "The VLLM_CUTLASS_SRC_DIR is set, using ${VLLM_CUTLASS_SRC_DIR} for compilation")
FetchContent_Declare(cutlass SOURCE_DIR ${VLLM_CUTLASS_SRC_DIR})
else()
FetchContent_Declare(
cutlass
GIT_REPOSITORY https://github.com/nvidia/cutlass.git
# Please keep this in sync with CUTLASS_REVISION line above.
GIT_TAG ${CUTLASS_REVISION}
GIT_PROGRESS TRUE
# Speed up CUTLASS download by retrieving only the specified GIT_TAG instead of the history.
# Important: If GIT_SHALLOW is enabled then GIT_TAG works only with branch names and tags.
# So if the GIT_TAG above is updated to a commit hash, GIT_SHALLOW must be set to FALSE
GIT_SHALLOW TRUE
)
endif()
FetchContent_MakeAvailable(cutlass)
set_gencode_flags_for_srcs(
SRCS "${VLLM_EXT_SRC}"
CUDA_ARCHS "${CUDA_ARCHS}")
# if CUDA endif
endif()
if (VLLM_GPU_LANG STREQUAL "HIP")
# Add QuickReduce kernels (ROCm-only; not part of stable ABI migration).
# TODO: Remove the cuda_view when ROCm upgrade to torch 2.11.
list(APPEND VLLM_EXT_SRC
if(VLLM_GPU_LANG STREQUAL "HIP")
set(VLLM_EXT_SRC
"csrc/torch_bindings.cpp"
"csrc/custom_quickreduce.cu"
"csrc/cuda_view.cu"
"csrc/libtorch_stable/cuda_utils_kernels.cu"
)
# if ROCM endif
endif()
"csrc/libtorch_stable/cuda_utils_kernels.cu")
message(STATUS "Enabling C extension.")
define_extension_target(
_C
DESTINATION vllm
LANGUAGE ${VLLM_GPU_LANG}
SOURCES ${VLLM_EXT_SRC}
COMPILE_FLAGS ${VLLM_GPU_FLAGS}
ARCHITECTURES ${VLLM_GPU_ARCHES}
INCLUDE_DIRECTORIES ${CUTLASS_INCLUDE_DIR}
INCLUDE_DIRECTORIES ${CUTLASS_TOOLS_UTIL_INCLUDE_DIR}
USE_SABI 3
WITH_SOABI)
message(STATUS "Enabling C extension.")
define_extension_target(
_C
DESTINATION vllm
LANGUAGE ${VLLM_GPU_LANG}
SOURCES ${VLLM_EXT_SRC}
COMPILE_FLAGS ${VLLM_GPU_FLAGS}
ARCHITECTURES ${VLLM_GPU_ARCHES}
INCLUDE_DIRECTORIES ${CUTLASS_INCLUDE_DIR}
INCLUDE_DIRECTORIES ${CUTLASS_TOOLS_UTIL_INCLUDE_DIR}
USE_SABI 3
WITH_SOABI)
# If CUTLASS is compiled on NVCC >= 12.5, it by default uses
# cudaGetDriverEntryPointByVersion as a wrapper to avoid directly calling the
# driver API. This causes problems when linking with earlier versions of CUDA.
# Setting this variable sidesteps the issue by calling the driver directly.
target_compile_definitions(_C PRIVATE CUTLASS_ENABLE_DIRECT_CUDA_DRIVER_CALL=1)
# If CUTLASS is compiled on NVCC >= 12.5, it by default uses
# cudaGetDriverEntryPointByVersion as a wrapper to avoid directly calling the
# driver API. This causes problems when linking with earlier versions of CUDA.
# Setting this variable sidesteps the issue by calling the driver directly.
target_compile_definitions(_C PRIVATE CUTLASS_ENABLE_DIRECT_CUDA_DRIVER_CALL=1)
endif() # _C HIP endif
if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
#
@@ -403,6 +356,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
set(VLLM_STABLE_EXT_SRC
"csrc/libtorch_stable/torch_bindings.cpp"
"csrc/libtorch_stable/activation_kernels.cu"
"csrc/libtorch_stable/quantization/activation_kernels.cu"
"csrc/libtorch_stable/quantization/w8a8/int8/scaled_quant.cu"
"csrc/libtorch_stable/quantization/w8a8/fp8/common.cu"
"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")
"csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu")
if(VLLM_GPU_LANG STREQUAL "CUDA")
SET(CUTLASS_ENABLE_HEADERS_ONLY ON CACHE BOOL "Enable only the header library")
# Set CUTLASS_REVISION. Used for FetchContent. Also fixes some bogus messages when building.
set(CUTLASS_REVISION "v4.4.2")
# Use the specified CUTLASS source directory for compilation if VLLM_CUTLASS_SRC_DIR is provided
if (DEFINED ENV{VLLM_CUTLASS_SRC_DIR})
set(VLLM_CUTLASS_SRC_DIR $ENV{VLLM_CUTLASS_SRC_DIR})
endif()
if(VLLM_CUTLASS_SRC_DIR)
if(NOT IS_ABSOLUTE VLLM_CUTLASS_SRC_DIR)
get_filename_component(VLLM_CUTLASS_SRC_DIR "${VLLM_CUTLASS_SRC_DIR}" ABSOLUTE)
endif()
message(STATUS "The VLLM_CUTLASS_SRC_DIR is set, using ${VLLM_CUTLASS_SRC_DIR} for compilation")
FetchContent_Declare(cutlass SOURCE_DIR ${VLLM_CUTLASS_SRC_DIR})
else()
FetchContent_Declare(
cutlass
GIT_REPOSITORY https://github.com/nvidia/cutlass.git
# Please keep this in sync with CUTLASS_REVISION line above.
GIT_TAG ${CUTLASS_REVISION}
GIT_PROGRESS TRUE
# Speed up CUTLASS download by retrieving only the specified GIT_TAG instead of the history.
# Important: If GIT_SHALLOW is enabled then GIT_TAG works only with branch names and tags.
# So if the GIT_TAG above is updated to a commit hash, GIT_SHALLOW must be set to FALSE
GIT_SHALLOW TRUE
)
endif()
FetchContent_MakeAvailable(cutlass)
list(APPEND VLLM_STABLE_EXT_SRC
"csrc/libtorch_stable/cuda_view.cu"
"csrc/libtorch_stable/cuda_utils_kernels.cu"
@@ -929,7 +915,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
SRCS "${FP4_SM120_SRCS}"
CUDA_ARCHS "${FP4_SM120_ARCHS}")
list(APPEND VLLM_STABLE_EXT_SRC "${FP4_SM120_SRCS}")
target_compile_definitions(_C PRIVATE ENABLE_NVFP4_SM120=1)
list(APPEND VLLM_GPU_FLAGS "-DENABLE_NVFP4_SM120=1")
list(APPEND VLLM_GPU_FLAGS "-DENABLE_CUTLASS_MOE_SM120=1")
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")
SRCS "${FP4_SM100_SRCS}"
CUDA_ARCHS "${FP4_SM100_ARCHS}")
list(APPEND VLLM_STABLE_EXT_SRC "${FP4_SM100_SRCS}")
target_compile_definitions(_C PRIVATE ENABLE_NVFP4_SM100=1)
list(APPEND VLLM_GPU_FLAGS "-DENABLE_NVFP4_SM100=1")
list(APPEND VLLM_GPU_FLAGS "-DENABLE_CUTLASS_MOE_SM100=1")
message(STATUS "Building SM10x/11x NVFP4/MXFP4 for archs: ${FP4_SM100_ARCHS}")
cmake/external_projects/qutlass.cmake
+29 -5
View File
@@ -60,6 +60,7 @@ endif()
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND QUTLASS_ARCHS)
set(QUTLASS_SOURCES
csrc/qutlass_registration.cpp
${qutlass_SOURCE_DIR}/qutlass/csrc/bindings.cpp
${qutlass_SOURCE_DIR}/qutlass/csrc/gemm.cu
${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)
if(CUTLASS_INCLUDE_DIR AND EXISTS "${CUTLASS_INCLUDE_DIR}/cutlass/cutlass.h")
list(APPEND QUTLASS_INCLUDES "${CUTLASS_INCLUDE_DIR}")
if(CUTLASS_TOOLS_UTIL_INCLUDE_DIR AND
EXISTS "${CUTLASS_TOOLS_UTIL_INCLUDE_DIR}/cutlass/util/packed_stride.hpp")
list(APPEND QUTLASS_INCLUDES "${CUTLASS_TOOLS_UTIL_INCLUDE_DIR}")
else()
get_filename_component(_qutlass_cutlass_root "${CUTLASS_INCLUDE_DIR}" DIRECTORY)
if(EXISTS "${_qutlass_cutlass_root}/tools/util/include/cutlass/util/packed_stride.hpp")
list(APPEND QUTLASS_INCLUDES "${_qutlass_cutlass_root}/tools/util/include")
endif()
endif()
elseif(EXISTS "${qutlass_SOURCE_DIR}/qutlass/third_party/cutlass/include/cutlass/cutlass.h")
list(APPEND QUTLASS_INCLUDES "${qutlass_SOURCE_DIR}/qutlass/third_party/cutlass/include")
list(APPEND QUTLASS_INCLUDES
"${qutlass_SOURCE_DIR}/qutlass/third_party/cutlass/include"
"${qutlass_SOURCE_DIR}/qutlass/third_party/cutlass/tools/util/include")
message(STATUS "[QUTLASS] Using QuTLASS vendored CUTLASS headers (no vLLM CUTLASS detected).")
else()
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)
CUDA_ARCHS "${QUTLASS_ARCHS}"
)
target_sources(_C PRIVATE ${QUTLASS_SOURCES})
target_include_directories(_C PRIVATE ${QUTLASS_INCLUDES})
target_compile_definitions(_C PRIVATE
# QuTLASS uses legacy ATen headers and cannot be built with TORCH_TARGET_VERSION.
# Keep it as its own extension (registers torch.ops._qutlass_C).
define_extension_target(
_qutlass_C
DESTINATION vllm
LANGUAGE ${VLLM_GPU_LANG}
SOURCES ${QUTLASS_SOURCES}
COMPILE_FLAGS ${VLLM_GPU_FLAGS}
ARCHITECTURES ${VLLM_GPU_ARCHES}
INCLUDE_DIRECTORIES ${QUTLASS_INCLUDES}
USE_SABI 3
WITH_SOABI)
target_compile_definitions(_qutlass_C PRIVATE
QUTLASS_DISABLE_PYBIND=1
TARGET_CUDA_ARCH=${QUTLASS_TARGET_CC}
)
CUTLASS_ENABLE_DIRECT_CUDA_DRIVER_CALL=1)
set_property(SOURCE ${QUTLASS_SOURCES} APPEND PROPERTY COMPILE_OPTIONS
$<$<COMPILE_LANGUAGE:CUDA>:--expt-relaxed-constexpr --use_fast_math -O3>
@@ -111,4 +134,5 @@ else()
"[QUTLASS] Skipping build: no supported arch (12.0f / 10.0f) found in "
"CUDA_ARCHS='${CUDA_ARCHS}'.")
endif()
add_custom_target(_qutlass_C)
endif()
csrc/core/math.hpp → csrc/libtorch_stable/core/math.hpp
View File
csrc/libtorch_stable/moe/moe_align_sum_kernels.cu
+1 -1
View File
@@ -9,7 +9,7 @@
#include <torch/headeronly/core/ScalarType.h>
#include "../../cuda_compat.h"
#include "core/math.hpp"
#include "libtorch_stable/core/math.hpp"
#include "libtorch_stable/dispatch_utils.h"
#include "libtorch_stable/torch_utils.h"
csrc/libtorch_stable/ops.h
+28
View File
@@ -2,9 +2,25 @@
#include <torch/csrc/stable/library.h>
#include <torch/csrc/stable/tensor.h>
#include <torch/headeronly/util/Exception.h>
#include <optional>
#include <string>
#include <vector>
#include <torch/csrc/stable/ops.h>
inline torch::stable::Tensor weak_ref_tensor(torch::stable::Tensor& tensor) {
// Ensure tensor is on CUDA
STD_TORCH_CHECK(tensor.device().is_cuda(), "Tensor must be on CUDA device");
// Get the raw data pointer
void* data_ptr = tensor.mutable_data_ptr();
/// Create a new tensor from the raw data pointer
return torch::stable::from_blob(data_ptr, tensor.sizes(), tensor.strides(),
tensor.device(), tensor.scalar_type());
}
void per_token_group_quant_fp8(const torch::stable::Tensor& input,
torch::stable::Tensor& output_q,
@@ -371,6 +387,18 @@ void silu_and_mul(torch::stable::Tensor& out, torch::stable::Tensor& input);
void silu_and_mul_clamp(torch::stable::Tensor& out,
torch::stable::Tensor& input, double limit,
double alpha = 1.0, double beta = 0.0);
void silu_and_mul_quant(torch::stable::Tensor& out,
torch::stable::Tensor& input,
torch::stable::Tensor& scale);
void persistent_masked_m_silu_mul_quant(
const torch::stable::Tensor& input, // (E, T, 2*H)
const torch::stable::Tensor& tokens_per_expert, // (E)
torch::stable::Tensor& y_q, // (E, T, H) [OUT]
torch::stable::Tensor& y_s, // (E, T, H//group_size) [OUT]
bool use_ue8m0);
void mul_and_silu(torch::stable::Tensor& out, torch::stable::Tensor& input);
void gelu_and_mul(torch::stable::Tensor& out, torch::stable::Tensor& input);
void gelu_tanh_and_mul(torch::stable::Tensor& out,
csrc/quantization/activation_kernels.cu → csrc/libtorch_stable/quantization/activation_kernels.cu
+63 -55
View File
@@ -1,16 +1,12 @@
#include <ATen/cuda/CUDAContext.h>
#include <torch/all.h>
#include <c10/cuda/CUDAGuard.h>
#include "libtorch_stable/torch_utils.h"
#include <cmath>
#include "core/math.hpp"
#include "../cuda_compat.h"
#include "dispatch_utils.h"
#include "libtorch_stable/core/math.hpp"
#include "cuda_compat.h"
#include "libtorch_stable/dispatch_utils.h"
#include "quantization/w8a8/fp8/common.cuh"
#include <c10/util/Float8_e4m3fn.h>
#ifndef USE_ROCM
#include <cuda_bf16.h>
#include <cuda_fp16.h>
@@ -33,7 +29,6 @@ typedef __hip_fp8x4_e4m3_fnuz __nv_fp8x4_e4m3;
#endif
#endif
#include "core/registration.h"
namespace vllm {
template <typename T>
@@ -564,41 +559,47 @@ __global__ void silu_mul_fp8_quant_deep_gemm_kernel(
} // namespace vllm
// Launch activation, gating, and quantize kernel.
#define LAUNCH_ACTIVATION_GATE_KERNEL(KERNEL) \
int d = input.size(-1) / 2; \
int64_t num_tokens = input.numel() / input.size(-1); \
dim3 grid(num_tokens, num_tokens > 16 ? num_tokens > 32 ? 1 : 2 : 4); \
dim3 block(std::min(d, 512)); \
const at::cuda::OptionalCUDAGuard device_guard(device_of(input)); \
const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); \
VLLM_DISPATCH_FLOATING_TYPES( \
input.scalar_type(), "act_and_mul_kernel", [&] { \
VLLM_DISPATCH_FP8_TYPES( \
out.scalar_type(), "fused_add_rms_norm_kernel_fp8_type", [&] { \
vllm::act_and_mul_quant_kernel<scalar_t, KERNEL<scalar_t>, \
fp8_t> \
<<<grid, block, 0, stream>>>(out.data_ptr<fp8_t>(), \
input.data_ptr<scalar_t>(), \
scale.data_ptr<float>(), d); \
}); \
#define LAUNCH_ACTIVATION_GATE_KERNEL(KERNEL) \
int d = input.size(-1) / 2; \
int64_t num_tokens = input.numel() / input.size(-1); \
dim3 grid(num_tokens, num_tokens > 16 ? num_tokens > 32 ? 1 : 2 : 4); \
dim3 block(std::min(d, 512)); \
const torch::stable::accelerator::DeviceGuard device_guard( \
input.get_device_index()); \
const cudaStream_t stream = \
get_current_cuda_stream(input.get_device_index()); \
VLLM_STABLE_DISPATCH_FLOATING_TYPES( \
input.scalar_type(), "act_and_mul_kernel", [&] { \
VLLM_STABLE_DISPATCH_FP8_TYPES( \
out.scalar_type(), "act_and_mul_quant_kernel_fp8_type", [&] { \
vllm::act_and_mul_quant_kernel<scalar_t, KERNEL<scalar_t>, \
fp8_t> \
<<<grid, block, 0, stream>>>( \
out.mutable_data_ptr<fp8_t>(), \
input.const_data_ptr<scalar_t>(), \
scale.const_data_ptr<float>(), d); \
}); \
});
void silu_and_mul_quant(torch::Tensor& out, // [..., d]
torch::Tensor& input, // [..., 2 * d]
torch::Tensor& scale) {
TORCH_CHECK(out.dtype() == torch::kFloat8_e4m3fn ||
out.dtype() == torch::kFloat8_e4m3fnuz);
TORCH_CHECK(input.dtype() == torch::kFloat16 ||
input.dtype() == torch::kBFloat16);
TORCH_CHECK(input.size(-1) % 2 == 0);
void silu_and_mul_quant(torch::stable::Tensor& out, // [..., d]
torch::stable::Tensor& input, // [..., 2 * d]
torch::stable::Tensor& scale) {
STD_TORCH_CHECK(
out.scalar_type() == torch::headeronly::ScalarType::Float8_e4m3fn ||
out.scalar_type() == torch::headeronly::ScalarType::Float8_e4m3fnuz);
STD_TORCH_CHECK(
input.scalar_type() == torch::headeronly::ScalarType::Half ||
input.scalar_type() == torch::headeronly::ScalarType::BFloat16,
"Input must be FP16 or BF16");
STD_TORCH_CHECK(input.size(-1) % 2 == 0);
LAUNCH_ACTIVATION_GATE_KERNEL(vllm::silu_kernel);
}
void persistent_masked_m_silu_mul_quant(
const at::Tensor& input, // (E, T, 2*H)
const at::Tensor& tokens_per_expert, // (E)
at::Tensor& y_q, // (E, T, H) [OUT]
at::Tensor& y_s, // (E, T, H//group_size) [OUT]
const torch::stable::Tensor& input, // (E, T, 2*H)
const torch::stable::Tensor& tokens_per_expert, // (E)
torch::stable::Tensor& y_q, // (E, T, H) [OUT]
torch::stable::Tensor& y_s, // (E, T, H//group_size) [OUT]
bool cast_scale_ue8m0) {
#ifndef USE_ROCM
@@ -606,14 +607,18 @@ void persistent_masked_m_silu_mul_quant(
// fixed GROUP_SIZE of 128.
static constexpr int GROUP_SIZE = 128;
TORCH_CHECK(input.dtype() == torch::kBFloat16);
TORCH_CHECK(y_q.dtype() == torch::kFloat8_e4m3fn ||
y_q.dtype() == torch::kFloat8_e4m3fnuz);
TORCH_CHECK(input.size(-1) % (GROUP_SIZE * 2) == 0);
STD_TORCH_CHECK(input.scalar_type() ==
torch::headeronly::ScalarType::BFloat16);
STD_TORCH_CHECK(
y_q.scalar_type() == torch::headeronly::ScalarType::Float8_e4m3fn ||
y_q.scalar_type() == torch::headeronly::ScalarType::Float8_e4m3fnuz);
STD_TORCH_CHECK(input.size(-1) % (GROUP_SIZE * 2) == 0);
bool const is_packed_ue8m0 =
(y_s.dtype() == torch::kInt32 && cast_scale_ue8m0);
TORCH_CHECK(y_s.dtype() == torch::kFloat32 || is_packed_ue8m0);
(y_s.scalar_type() == torch::headeronly::ScalarType::Int &&
cast_scale_ue8m0);
STD_TORCH_CHECK(y_s.scalar_type() == torch::headeronly::ScalarType::Float ||
is_packed_ue8m0);
using Idx_t = int64_t;
@@ -631,7 +636,7 @@ void persistent_masked_m_silu_mul_quant(
int const NUM_GROUPS = H / GROUP_SIZE;
const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
const cudaStream_t stream = get_current_cuda_stream(input.get_device_index());
// TODO: Get this from cuda_arch ?
static constexpr int SILU_V2_BLOCK_COUNT = 132 * 32;
@@ -643,18 +648,21 @@ void persistent_masked_m_silu_mul_quant(
static constexpr int max_shared_mem_bytes = \
GROUP_SIZE * 2 * STAGES * NUM_WARPS * 2; \
dim3 grid(sms), block(THREAD_COUNT); \
const at::cuda::OptionalCUDAGuard device_guard(device_of(input)); \
VLLM_DISPATCH_FP8_TYPES( \
const torch::stable::accelerator::DeviceGuard device_guard( \
input.get_device_index()); \
VLLM_STABLE_DISPATCH_FP8_TYPES( \
y_q.scalar_type(), "silu_mul_fp8_quant_deep_gemm_kernel", [&] { \
vllm::silu_mul_fp8_quant_deep_gemm_kernel< \
BLOCK_COUNT, max_shared_mem_bytes, fp8_t, scale_t, THREAD_COUNT, \
Idx_t, CEIL_UE8M0, GROUP_SIZE, STAGES> \
<<<grid, block, max_shared_mem_bytes + (E + 1) * 16, stream>>>( \
reinterpret_cast<__nv_bfloat16*>(input.data_ptr()), \
(fp8_t*)y_q.data_ptr(), \
reinterpret_cast<scale_t*>(y_s.data_ptr()), \
reinterpret_cast<int32_t*>(tokens_per_expert.data_ptr()), E, \
T, H, stride_i_e, stride_i_t, stride_i_h, stride_yq_e, \
reinterpret_cast<const __nv_bfloat16*>( \
input.const_data_ptr()), \
y_q.mutable_data_ptr<fp8_t>(), \
reinterpret_cast<scale_t*>(y_s.mutable_data_ptr()), \
reinterpret_cast<const int32_t*>( \
tokens_per_expert.const_data_ptr()), \
E, T, H, stride_i_e, stride_i_t, stride_i_h, stride_yq_e, \
stride_yq_t, stride_yq_h, STRIDE_YS_E, STRIDE_YS_T, \
STRIDE_YS_G, STRIDE_YS_P, stride_counts_e); \
});
@@ -679,7 +687,7 @@ void persistent_masked_m_silu_mul_quant(
Idx_t stride_ys_g = y_s.stride(2);
Idx_t stride_ys_p = 0;
if (!cast_scale_ue8m0) {
TORCH_CHECK(!is_packed_ue8m0);
STD_TORCH_CHECK(!is_packed_ue8m0);
LAUNCH_ON_H(float, stride_ys_e, stride_ys_t, stride_ys_g, stride_ys_p,
false);
return;
@@ -692,8 +700,8 @@ void persistent_masked_m_silu_mul_quant(
return;
}
TORCH_CHECK(cast_scale_ue8m0 && is_packed_ue8m0);
TORCH_CHECK(y_s.dtype() == torch::kInt32);
STD_TORCH_CHECK(cast_scale_ue8m0 && is_packed_ue8m0);
STD_TORCH_CHECK(y_s.scalar_type() == torch::headeronly::ScalarType::Int);
// Int32 packed ue8m0 scales tensor.
// 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
View File
@@ -31,7 +31,7 @@
#include "cutlass/util/packed_stride.hpp"
#include "core/math.hpp"
#include "libtorch_stable/core/math.hpp"
#include "core/batch_invariant.hpp"
using namespace cute;
csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_sm120_kernels.cu
+1 -1
View File
@@ -31,7 +31,7 @@
#include "cutlass/util/packed_stride.hpp"
#include "core/math.hpp"
#include "libtorch_stable/core/math.hpp"
#include "core/batch_invariant.hpp"
using namespace cute;
csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/cutlass_gemm_caller.cuh
+1 -1
View File
@@ -19,7 +19,7 @@
#include "cutlass/gemm/collective/collective_builder.hpp"
#include "cutlass/util/packed_stride.hpp"
#include "core/math.hpp"
#include "libtorch_stable/core/math.hpp"
#include "libtorch_stable/cutlass_extensions/common.hpp"
// clang-format on
csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm.cuh
+1 -1
View File
@@ -14,7 +14,7 @@
#include "cutlass/epilogue/collective/collective_builder.hpp"
#include "cutlass/gemm/collective/collective_builder.hpp"
#include "core/math.hpp"
#include "libtorch_stable/core/math.hpp"
#include "libtorch_stable/cutlass_extensions/common.hpp"
// clang-format on
csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_c2x.cuh
+1 -1
View File
@@ -22,7 +22,7 @@
#include "cutlass/epilogue/threadblock/fusion/visitors.hpp"
#include "cutlass/gemm/kernel/default_gemm_universal_with_visitor.h"
#include "core/math.hpp"
#include "libtorch_stable/core/math.hpp"
#include "libtorch_stable/cutlass_extensions/common.hpp"
// clang-format on
csrc/libtorch_stable/torch_bindings.cpp
+27
View File
@@ -34,6 +34,20 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
// TODO: Remove this once ROCm upgrade to torch 2.11.
ops.def("get_cuda_view_from_cpu_tensor(Tensor cpu_tensor) -> Tensor");
// Note about marlin kernel 'workspace' arguments:
// Technically these should be mutable since they are modified by the kernel.
// But since they are set back to zero once the kernel is finished we can
// hand wave and say that they have no net effect.
//
// The reason to mark 'workspace' as immutable is so that they don't interfere
// with using ScalarType arguments in the ops. If they are marked as mutable,
// pytorch throws an assert in
// 'torch._higher_order_ops._register_effectful_op' that prevents these
// kernels from being torch.compile'd.
// See the following document for more info on custom types and ops that use
// custom types:
// https://docs.google.com/document/d/18fBMPuOJ0fY5ZQ6YyrHUppw9FA332CpNtgB6SOIgyuA
// Machete (Dense) Optimized Mixed Precision GEMM for Hopper.
ops.def(
"machete_supported_schedules("
@@ -480,6 +494,11 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
"Tensor workspace, int k, int max_seq_len) -> ()");
// Activation ops
ops.def(
"persistent_masked_m_silu_mul_quant(Tensor input, Tensor counts, Tensor! "
"y_q, Tensor! y_s, bool use_ue8m0) -> ()");
ops.def("weak_ref_tensor(Tensor input) -> Tensor");
// Activation function used in SwiGLU.
ops.def("silu_and_mul(Tensor! result, Tensor input) -> ()");
@@ -492,6 +511,10 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
"silu_and_mul_with_clamp(Tensor! result, Tensor input, float limit, "
"float alpha=1.0, float beta=0.0) -> ()");
// SwiGLU activation with FP8 quantization.
ops.def(
"silu_and_mul_quant(Tensor! result, Tensor input, Tensor scale) -> ()");
// Activation function used in GeGLU with `none` approximation.
ops.def("gelu_and_mul(Tensor! out, Tensor input) -> ()");
@@ -690,6 +713,10 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) {
ops.impl("persistent_topk", TORCH_BOX(&persistent_topk));
// Activation kernels (shared CUDA/ROCm)
ops.impl("persistent_masked_m_silu_mul_quant",
TORCH_BOX(&persistent_masked_m_silu_mul_quant));
ops.impl("weak_ref_tensor", TORCH_BOX(&weak_ref_tensor));
ops.impl("silu_and_mul_quant", TORCH_BOX(&silu_and_mul_quant));
ops.impl("silu_and_mul", TORCH_BOX(&silu_and_mul));
ops.impl("mul_and_silu", TORCH_BOX(&mul_and_silu));
ops.impl("gelu_and_mul", TORCH_BOX(&gelu_and_mul));
csrc/ops.h
-32
View File
@@ -9,28 +9,6 @@
#include <vector>
torch::Tensor weak_ref_tensor(torch::Tensor& tensor) {
// Ensure tensor is on CUDA
if (!tensor.is_cuda()) {
throw std::runtime_error("Tensor must be on CUDA device");
}
// Get the raw data pointer
void* data_ptr = tensor.data_ptr();
// Get tensor sizes and strides
std::vector<int64_t> sizes = tensor.sizes().vec();
std::vector<int64_t> strides = tensor.strides().vec();
// Get tensor options (dtype, device)
auto options = tensor.options();
// Create a new tensor from the raw data pointer
auto new_tensor = torch::from_blob(data_ptr, sizes, strides, options);
return new_tensor;
}
// rms_norm and fused_add_rms_norm declarations also exist in
// csrc/libtorch_stable/ops.h (torch::stable ABI for CUDA). They remain here
// because the CPU build still uses these torch::Tensor declarations.
@@ -53,16 +31,6 @@ void silu_and_mul(torch::Tensor& out, torch::Tensor& input);
void silu_and_mul_clamp(torch::Tensor& out, torch::Tensor& input, double limit,
double alpha = 1.0, double beta = 0.0);
void silu_and_mul_quant(torch::Tensor& out, torch::Tensor& input,
torch::Tensor& scale);
void persistent_masked_m_silu_mul_quant(
const at::Tensor& input, // (E, T, 2*H)
const at::Tensor& counts, // (E)
at::Tensor& y_q, // (E, T, H) [OUT]
at::Tensor& y_s, // (E, T, H//group_size) [OUT]
bool use_ue8m0);
void gelu_and_mul(torch::Tensor& out, torch::Tensor& input);
void gelu_tanh_and_mul(torch::Tensor& out, torch::Tensor& input);
csrc/qutlass_registration.cpp
+5
View File
@@ -0,0 +1,5 @@
#include "core/registration.h"
// QuTLASS registers torch.ops._qutlass_C via TORCH_LIBRARY in bindings.cpp.
// This stub lets Python import vllm._qutlass_C to trigger op registration.
REGISTER_EXTENSION(_qutlass_C)
csrc/torch_bindings.cpp
-40
View File
@@ -20,17 +20,6 @@
TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
// vLLM custom ops
//
ops.def(
"persistent_masked_m_silu_mul_quant(Tensor input, Tensor counts, Tensor! "
"y_q, Tensor! y_s,"
"bool use_ue8m0) -> ()");
ops.impl("persistent_masked_m_silu_mul_quant", torch::kCUDA,
&persistent_masked_m_silu_mul_quant);
ops.def("weak_ref_tensor(Tensor input) -> Tensor");
ops.impl("weak_ref_tensor", torch::kCUDA, &weak_ref_tensor);
#ifdef USE_ROCM
// TODO: Remove this once we upgrade to torch 2.11.
@@ -39,35 +28,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
ops.def("get_cuda_view_from_cpu_tensor(Tensor cpu_tensor) -> Tensor");
ops.impl("get_cuda_view_from_cpu_tensor", torch::kCPU,
&get_cuda_view_from_cpu_tensor);
#endif
// Activation ops (quantized only — basic ops moved to _C_stable_libtorch)
ops.def(
"silu_and_mul_quant(Tensor! result, Tensor input, Tensor scale) -> ()");
ops.impl("silu_and_mul_quant", torch::kCUDA, &silu_and_mul_quant);
// Horizontally-fused DeepseekV4-MLA: per-head RMSNorm + GPT-J RoPE for Q, and
// GPT-J RoPE + UE8M0 FP8 quant + paged cache insert for KV, all in one
// kernel launch. Registered in _C_stable_libtorch (incl. the FlashInfer V4
// full-cache bf16/fp8 variants).
// Quantization ops
#ifndef USE_ROCM
// Note about marlin kernel 'workspace' arguments:
// Technically these should be mutable since they are modified by the kernel.
// But since they are set back to zero once the kernel is finished we can
// hand wave and say that they have no net effect.
//
// The reason to mark 'workspace' as immutable is so that they don't interfere
// with using ScalarType arguments in the ops. If they are marked as mutable,
// pytorch throws an assert in
// 'torch._higher_order_ops._register_effectful_op' that prevents these
// kernels from being torch.compile'd.
// See the following document for more info on custom types and ops that use
// custom types:
// https://docs.google.com/document/d/18fBMPuOJ0fY5ZQ6YyrHUppw9FA332CpNtgB6SOIgyuA
#endif
}
setup.py
+4 -1
View File
@@ -769,6 +769,7 @@ class precompiled_wheel_utils:
"vllm/_C.abi3.so",
"vllm/_C_stable_libtorch.abi3.so",
"vllm/_moe_C_stable_libtorch.abi3.so",
"vllm/_qutlass_C.abi3.so",
"vllm/_flashmla_C.abi3.so",
"vllm/_flashmla_extension_C.abi3.so",
"vllm/_sparse_flashmla_C.abi3.so",
@@ -1135,6 +1136,7 @@ if _is_cuda():
# DeepGEMM requires CUDA 12.3+ (SM90/SM100)
# Optional since it won't build on unsupported architectures
ext_modules.append(CMakeExtension(name="vllm._deep_gemm_C", optional=True))
ext_modules.append(CMakeExtension(name="vllm._qutlass_C", optional=True))
# fmha_sm100 is a Python/CuTe-DSL package installed into vllm.third_party.
ext_modules.append(CMakeExtension(name="vllm.fmha_sm100", optional=True))
@@ -1149,7 +1151,8 @@ if _is_cpu():
ext_modules.append(CMakeExtension(name="vllm._C"))
if _build_custom_ops():
ext_modules.append(CMakeExtension(name="vllm._C"))
if _is_hip():
ext_modules.append(CMakeExtension(name="vllm._C"))
if _is_cuda() or _is_hip():
ext_modules.append(CMakeExtension(name="vllm._C_stable_libtorch"))
ext_modules.append(CMakeExtension(name="vllm._moe_C_stable_libtorch"))
vllm/platforms/cuda.py
+21 -1
View File
@@ -19,7 +19,6 @@ from torch.distributed.distributed_c10d import is_nccl_available
from typing_extensions import ParamSpec
# import custom ops, trigger op registration
import vllm._C # noqa
import vllm._C_stable_libtorch # noqa
import vllm.envs as envs
from vllm.logger import init_logger
@@ -40,6 +39,11 @@ else:
logger = init_logger(__name__)
try:
import vllm._qutlass_C # noqa: F401
except ImportError as e:
logger.warning("Failed to import from vllm._qutlass_C: %r", e)
_P = ParamSpec("_P")
_R = TypeVar("_R")
@@ -187,6 +191,22 @@ class CudaPlatformBase(Platform):
"RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES",
]
@classmethod
def import_kernels(cls) -> None:
"""Import CUDA kernel extensions (_C_stable_libtorch, optional _qutlass_C)."""
try:
import vllm._C_stable_libtorch # noqa: F401
except ImportError as e:
logger.warning("Failed to import from vllm._C_stable_libtorch: %r", e)
try:
import vllm._moe_C_stable_libtorch # noqa: F401
except ImportError as e:
logger.warning("Failed to import from vllm._moe_C_stable_libtorch: %r", e)
try:
import vllm._qutlass_C # noqa: F401
except ImportError as e:
logger.warning("Failed to import from vllm._qutlass_C: %r", e)
@property
def supported_dtypes(self) -> list[torch.dtype]:
if self.has_device_capability(80):