From b9a7cd464c9ae9b1b450f8982b76d7be4de73724 Mon Sep 17 00:00:00 2001 From: Chris Leonard Date: Fri, 19 Jun 2026 09:57:26 -0400 Subject: [PATCH] [12/n] final _C library kernel migration (#45415) --- CMakeLists.txt | 128 ++++++++---------- cmake/external_projects/qutlass.cmake | 34 ++++- csrc/{ => libtorch_stable}/core/math.hpp | 0 .../moe/moe_align_sum_kernels.cu | 2 +- csrc/libtorch_stable/ops.h | 28 ++++ .../quantization/activation_kernels.cu | 118 ++++++++-------- .../fp4/nvfp4_scaled_mm_kernels.cu | 2 +- .../fp4/nvfp4_scaled_mm_sm120_kernels.cu | 2 +- .../w8a8/cutlass/c3x/cutlass_gemm_caller.cuh | 2 +- .../w8a8/cutlass/c3x/scaled_mm.cuh | 2 +- .../w8a8/cutlass/scaled_mm_c2x.cuh | 2 +- csrc/libtorch_stable/torch_bindings.cpp | 27 ++++ csrc/ops.h | 32 ----- csrc/qutlass_registration.cpp | 5 + csrc/torch_bindings.cpp | 40 ------ setup.py | 5 +- vllm/platforms/cuda.py | 22 ++- 17 files changed, 239 insertions(+), 212 deletions(-) rename csrc/{ => libtorch_stable}/core/math.hpp (100%) rename csrc/{ => libtorch_stable}/quantization/activation_kernels.cu (87%) create mode 100644 csrc/qutlass_registration.cpp diff --git a/CMakeLists.txt b/CMakeLists.txt index a2651ab344..e95fe38d32 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -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}") diff --git a/cmake/external_projects/qutlass.cmake b/cmake/external_projects/qutlass.cmake index 66c001919b..b653bbfce7 100644 --- a/cmake/external_projects/qutlass.cmake +++ b/cmake/external_projects/qutlass.cmake @@ -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 $<$:--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() diff --git a/csrc/core/math.hpp b/csrc/libtorch_stable/core/math.hpp similarity index 100% rename from csrc/core/math.hpp rename to csrc/libtorch_stable/core/math.hpp diff --git a/csrc/libtorch_stable/moe/moe_align_sum_kernels.cu b/csrc/libtorch_stable/moe/moe_align_sum_kernels.cu index d7c68ff25a..1e84238134 100644 --- a/csrc/libtorch_stable/moe/moe_align_sum_kernels.cu +++ b/csrc/libtorch_stable/moe/moe_align_sum_kernels.cu @@ -9,7 +9,7 @@ #include #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" diff --git a/csrc/libtorch_stable/ops.h b/csrc/libtorch_stable/ops.h index 9efc12e9f4..1cc8e8167a 100644 --- a/csrc/libtorch_stable/ops.h +++ b/csrc/libtorch_stable/ops.h @@ -2,9 +2,25 @@ #include #include +#include #include #include +#include + +#include + +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, diff --git a/csrc/quantization/activation_kernels.cu b/csrc/libtorch_stable/quantization/activation_kernels.cu similarity index 87% rename from csrc/quantization/activation_kernels.cu rename to csrc/libtorch_stable/quantization/activation_kernels.cu index 8cc645c33e..822a41969e 100644 --- a/csrc/quantization/activation_kernels.cu +++ b/csrc/libtorch_stable/quantization/activation_kernels.cu @@ -1,16 +1,12 @@ -#include -#include -#include +#include "libtorch_stable/torch_utils.h" #include -#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 - #ifndef USE_ROCM #include #include @@ -33,7 +29,6 @@ typedef __hip_fp8x4_e4m3_fnuz __nv_fp8x4_e4m3; #endif #endif -#include "core/registration.h" namespace vllm { template @@ -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, \ - fp8_t> \ - <<>>(out.data_ptr(), \ - input.data_ptr(), \ - scale.data_ptr(), 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, \ + fp8_t> \ + <<>>( \ + out.mutable_data_ptr(), \ + input.const_data_ptr(), \ + scale.const_data_ptr(), 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> \ <<>>( \ - reinterpret_cast<__nv_bfloat16*>(input.data_ptr()), \ - (fp8_t*)y_q.data_ptr(), \ - reinterpret_cast(y_s.data_ptr()), \ - reinterpret_cast(tokens_per_expert.data_ptr()), E, \ - T, H, stride_i_e, stride_i_t, stride_i_h, stride_yq_e, \ + reinterpret_cast( \ + input.const_data_ptr()), \ + y_q.mutable_data_ptr(), \ + reinterpret_cast(y_s.mutable_data_ptr()), \ + reinterpret_cast( \ + 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 diff --git a/csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_kernels.cu b/csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_kernels.cu index 86355bf706..af9f24a70e 100644 --- a/csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_kernels.cu +++ b/csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_kernels.cu @@ -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; diff --git a/csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_sm120_kernels.cu b/csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_sm120_kernels.cu index 7adba6308f..3a45ede8df 100644 --- a/csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_sm120_kernels.cu +++ b/csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_sm120_kernels.cu @@ -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; diff --git a/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/cutlass_gemm_caller.cuh b/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/cutlass_gemm_caller.cuh index 1eed757992..1d9023484f 100644 --- a/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/cutlass_gemm_caller.cuh +++ b/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/cutlass_gemm_caller.cuh @@ -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 diff --git a/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm.cuh b/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm.cuh index 4cb591be05..7b7d4d7147 100644 --- a/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm.cuh +++ b/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm.cuh @@ -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 diff --git a/csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_c2x.cuh b/csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_c2x.cuh index 7846e609fe..d2b54cb911 100644 --- a/csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_c2x.cuh +++ b/csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_c2x.cuh @@ -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 diff --git a/csrc/libtorch_stable/torch_bindings.cpp b/csrc/libtorch_stable/torch_bindings.cpp index c1d2d26fcd..d55c12d382 100644 --- a/csrc/libtorch_stable/torch_bindings.cpp +++ b/csrc/libtorch_stable/torch_bindings.cpp @@ -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)); diff --git a/csrc/ops.h b/csrc/ops.h index ec3f5e187c..398ae1016f 100644 --- a/csrc/ops.h +++ b/csrc/ops.h @@ -9,28 +9,6 @@ #include -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 sizes = tensor.sizes().vec(); - std::vector 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); diff --git a/csrc/qutlass_registration.cpp b/csrc/qutlass_registration.cpp new file mode 100644 index 0000000000..effb440413 --- /dev/null +++ b/csrc/qutlass_registration.cpp @@ -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) diff --git a/csrc/torch_bindings.cpp b/csrc/torch_bindings.cpp index cfd185394a..e1430c08d3 100644 --- a/csrc/torch_bindings.cpp +++ b/csrc/torch_bindings.cpp @@ -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 } diff --git a/setup.py b/setup.py index 2aaa7dfc49..b807b2215d 100644 --- a/setup.py +++ b/setup.py @@ -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")) diff --git a/vllm/platforms/cuda.py b/vllm/platforms/cuda.py index 49181eaec6..30a16e2746 100644 --- a/vllm/platforms/cuda.py +++ b/vllm/platforms/cuda.py @@ -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):