From de2a8fc042b6e9ca834ff20746734ed32b5ca219 Mon Sep 17 00:00:00 2001 From: TJian Date: Thu, 2 Jul 2026 18:34:07 +0800 Subject: [PATCH] [ROCm] [PyTorch] Move to stable abi since ROCm upgraded to torch 2.11 (#47128) Signed-off-by: tjtanaa --- CMakeLists.txt | 46 +++++++------------ csrc/cuda_view.cu | 60 ------------------------- csrc/libtorch_stable/ops.h | 7 ++- csrc/libtorch_stable/torch_bindings.cpp | 9 +--- csrc/ops.h | 3 -- csrc/torch_bindings.cpp | 28 ------------ 6 files changed, 22 insertions(+), 131 deletions(-) delete mode 100644 csrc/cuda_view.cu diff --git a/CMakeLists.txt b/CMakeLists.txt index 1ef9d596aec..cf215a6f249 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -350,9 +350,7 @@ endif() 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") + "csrc/custom_quickreduce.cu") message(STATUS "Enabling C extension.") define_extension_target( @@ -380,6 +378,8 @@ 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/cuda_view.cu" + "csrc/libtorch_stable/cuda_utils_kernels.cu" "csrc/libtorch_stable/activation_kernels.cu" "csrc/libtorch_stable/quantization/activation_kernels.cu" "csrc/libtorch_stable/quantization/w8a8/int8/scaled_quant.cu" @@ -459,8 +459,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") FetchContent_MakeAvailable(cutlass) list(APPEND VLLM_STABLE_EXT_SRC - "csrc/libtorch_stable/cuda_view.cu" - "csrc/libtorch_stable/cuda_utils_kernels.cu" "csrc/libtorch_stable/cutlass_extensions/common.cpp" "csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_entry.cu" "csrc/libtorch_stable/quantization/fp4/nvfp4_quant_entry.cu" @@ -1091,14 +1089,15 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") USE_SABI 3 WITH_SOABI) + # Set TORCH_TARGET_VERSION for stable ABI compatibility. + # This ensures we only use C-shim APIs available in PyTorch 2.11. + # _C_stable_libtorch is abi compatible with PyTorch >= TORCH_TARGET_VERSION + # which is currently set to 2.11. + target_compile_definitions(_C_stable_libtorch PRIVATE + TORCH_TARGET_VERSION=0x020B000000000000ULL) + # Needed to use cuda/hip APIs from C-shim if(VLLM_GPU_LANG STREQUAL "CUDA") - # Set TORCH_TARGET_VERSION for stable ABI compatibility. - # This ensures we only use C-shim APIs available in PyTorch 2.11. - # _C_stable_libtorch is abi compatible with PyTorch >= TORCH_TARGET_VERSION - # which is currently set to 2.11. - target_compile_definitions(_C_stable_libtorch PRIVATE - TORCH_TARGET_VERSION=0x020B000000000000ULL) target_compile_definitions(_C_stable_libtorch PRIVATE USE_CUDA) if(COOPERATIVE_TOPK_ARCHS) target_compile_definitions(_C_stable_libtorch PRIVATE @@ -1108,12 +1107,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") target_compile_definitions(_C_stable_libtorch PRIVATE CUTLASS_ENABLE_DIRECT_CUDA_DRIVER_CALL=1) elseif(VLLM_GPU_LANG STREQUAL "HIP") - # Set TORCH_TARGET_VERSION for stable ABI compatibility. - # This ensures we only use C-shim APIs available in PyTorch 2.10. - # _C_stable_libtorch is abi compatible with PyTorch >= TORCH_TARGET_VERSION - # which is currently set to 2.10. - target_compile_definitions(_C_stable_libtorch PRIVATE - TORCH_TARGET_VERSION=0x020A000000000000ULL) target_compile_definitions(_C_stable_libtorch PRIVATE USE_ROCM) endif() @@ -1321,25 +1314,20 @@ define_extension_target( USE_SABI 3 WITH_SOABI) +# Set TORCH_TARGET_VERSION for stable ABI compatibility. +# This ensures we only use C-shim APIs available in PyTorch 2.11. +# _moe_C_stable_libtorch is abi compatible with PyTorch >= TORCH_TARGET_VERSION +# which is currently set to 2.11. +target_compile_definitions(_moe_C_stable_libtorch PRIVATE + TORCH_TARGET_VERSION=0x020B000000000000ULL) + # Needed to use cuda/hip APIs from C-shim if(VLLM_GPU_LANG STREQUAL "CUDA") - # Set TORCH_TARGET_VERSION for stable ABI compatibility. - # This ensures we only use C-shim APIs available in PyTorch 2.11. - # _moe_C_stable_libtorch is abi compatible with PyTorch >= TORCH_TARGET_VERSION - # which is currently set to 2.11. - target_compile_definitions(_moe_C_stable_libtorch PRIVATE - TORCH_TARGET_VERSION=0x020B000000000000ULL) target_compile_definitions(_moe_C_stable_libtorch PRIVATE USE_CUDA) # Needed by CUTLASS kernels target_compile_definitions(_moe_C_stable_libtorch PRIVATE CUTLASS_ENABLE_DIRECT_CUDA_DRIVER_CALL=1) elseif(VLLM_GPU_LANG STREQUAL "HIP") - # Set TORCH_TARGET_VERSION for stable ABI compatibility. - # This ensures we only use C-shim APIs available in PyTorch 2.10. - # _moe_C_stable_libtorch is abi compatible with PyTorch >= TORCH_TARGET_VERSION - # which is currently set to 2.10. - target_compile_definitions(_moe_C_stable_libtorch PRIVATE - TORCH_TARGET_VERSION=0x020A000000000000ULL) target_compile_definitions(_moe_C_stable_libtorch PRIVATE USE_ROCM) endif() diff --git a/csrc/cuda_view.cu b/csrc/cuda_view.cu deleted file mode 100644 index 00e06a9329f..00000000000 --- a/csrc/cuda_view.cu +++ /dev/null @@ -1,60 +0,0 @@ -// TODO: Remove this once ROCm upgrade to torch 2.11. -#include -#include -#include - -// This function assumes that `cpu_tensor` is a CPU tensor, -// and that UVA (Unified Virtual Addressing) is enabled. -torch::Tensor get_cuda_view_from_cpu_tensor(torch::Tensor& cpu_tensor) { - TORCH_CHECK(cpu_tensor.device().is_cpu(), "Input tensor must be on CPU"); - - // handle empty tensor - if (cpu_tensor.numel() == 0) { - return torch::empty(cpu_tensor.sizes(), - cpu_tensor.options().device(torch::kCUDA)); - } - - if (cpu_tensor.is_pinned()) { - // If CPU tensor is pinned, directly get the device pointer. - void* host_ptr = const_cast(cpu_tensor.data_ptr()); - void* device_ptr = nullptr; - cudaError_t err = cudaHostGetDevicePointer(&device_ptr, host_ptr, 0); - TORCH_CHECK(err == cudaSuccess, - "cudaHostGetDevicePointer failed: ", cudaGetErrorString(err)); - - return torch::from_blob( - device_ptr, cpu_tensor.sizes(), cpu_tensor.strides(), - [base = cpu_tensor](void*) {}, // keep cpu tensor alive - cpu_tensor.options().device(torch::kCUDA)); - } - - // If CPU tensor is not pinned, allocate a new pinned memory buffer. - torch::Tensor contiguous_cpu = cpu_tensor.contiguous(); - size_t nbytes = contiguous_cpu.nbytes(); - - void* host_ptr = nullptr; - cudaError_t err = cudaHostAlloc(&host_ptr, nbytes, cudaHostAllocMapped); - if (err != cudaSuccess) { - AT_ERROR("cudaHostAlloc failed: ", cudaGetErrorString(err)); - } - - err = cudaMemcpy(host_ptr, contiguous_cpu.data_ptr(), nbytes, - cudaMemcpyDefault); - if (err != cudaSuccess) { - cudaFreeHost(host_ptr); - AT_ERROR("cudaMemcpy failed: ", cudaGetErrorString(err)); - } - - void* device_ptr = nullptr; - err = cudaHostGetDevicePointer(&device_ptr, host_ptr, 0); - if (err != cudaSuccess) { - cudaFreeHost(host_ptr); - AT_ERROR("cudaHostGetDevicePointer failed: ", cudaGetErrorString(err)); - } - - auto deleter = [host_ptr](void*) { cudaFreeHost(host_ptr); }; - - return torch::from_blob(device_ptr, contiguous_cpu.sizes(), - contiguous_cpu.strides(), deleter, - contiguous_cpu.options().device(torch::kCUDA)); -} \ No newline at end of file diff --git a/csrc/libtorch_stable/ops.h b/csrc/libtorch_stable/ops.h index 7cf34d8b03a..6438daefcab 100644 --- a/csrc/libtorch_stable/ops.h +++ b/csrc/libtorch_stable/ops.h @@ -181,13 +181,12 @@ torch::stable::Tensor awq_dequantize(torch::stable::Tensor _kernel, // AllSpark ops: declarations are in the source files // (allspark_repack.cu and allspark_qgemm_w8a16.cu) -// TODO: Move this out once ROCm upgrade their torch to 2.11. -// CPU tensor -> CUDA UVA view (shared CUDA) +#endif + +// CPU tensor -> CUDA UVA view (shared CUDA/ROCm) torch::stable::Tensor get_cuda_view_from_cpu_tensor( torch::stable::Tensor& cpu_tensor); -#endif - // Attention kernels (shared CUDA/ROCm) void merge_attn_states( torch::stable::Tensor& output, diff --git a/csrc/libtorch_stable/torch_bindings.cpp b/csrc/libtorch_stable/torch_bindings.cpp index 158999a6633..a760df260ac 100644 --- a/csrc/libtorch_stable/torch_bindings.cpp +++ b/csrc/libtorch_stable/torch_bindings.cpp @@ -29,11 +29,10 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) { "()"); ops.def("permute_cols(Tensor A, Tensor perm) -> Tensor"); -#ifndef USE_ROCM - - // TODO: Remove this once ROCm upgrade to torch 2.11. ops.def("get_cuda_view_from_cpu_tensor(Tensor cpu_tensor) -> Tensor"); +#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 @@ -753,8 +752,6 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) { ops.impl("paged_attention_v2", TORCH_BOX(&paged_attention_v2)); } -// TODO: Remove this once ROCm upgrade to torch 2.11. -#ifndef USE_ROCM STABLE_TORCH_LIBRARY_IMPL(_C, CPU, ops) { ops.impl("get_cuda_view_from_cpu_tensor", TORCH_BOX(&get_cuda_view_from_cpu_tensor)); @@ -773,8 +770,6 @@ STABLE_TORCH_LIBRARY_IMPL(_C_cuda_utils, CompositeExplicitAutograd, TORCH_BOX(&get_max_shared_memory_per_block_device_attribute)); } -#endif - // These capability-check functions take only primitive args (no tensors), so // there is no device to dispatch on. CompositeExplicitAutograd makes them // available for all backends. This is the stable ABI equivalent of calling diff --git a/csrc/ops.h b/csrc/ops.h index 0cf73f6bfb3..cd18b1e5e0d 100644 --- a/csrc/ops.h +++ b/csrc/ops.h @@ -67,7 +67,4 @@ void qr_open_handles(fptr_t _fa, const std::vector& handles); void qr_all_reduce(fptr_t _fa, torch::Tensor& inp, torch::Tensor& out, int64_t quant_level, bool cast_bf2half = false); int64_t qr_max_size(); - -// TODO: Remove this once ROCm upgrade to torch 2.11. -torch::Tensor get_cuda_view_from_cpu_tensor(torch::Tensor& cpu_tensor); #endif diff --git a/csrc/torch_bindings.cpp b/csrc/torch_bindings.cpp index e1430c08d3a..bcf0ce3e1c4 100644 --- a/csrc/torch_bindings.cpp +++ b/csrc/torch_bindings.cpp @@ -2,7 +2,6 @@ // cache.h, which is no longer included here after cache ops moved to // _C_stable_libtorch). #include -#include "cuda_utils.h" #include "ops.h" #include "core/registration.h" #include @@ -18,19 +17,6 @@ // https://docs.google.com/document/d/1_W62p8WJOQQUzPsJYa7s701JXt0qf2OfLub2sbkHOaU/edit#heading=h.ptttacy8y1u9 // https://github.com/pytorch/pytorch/blob/main/aten/src/ATen/native/README.md#annotations -TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { - // vLLM custom ops - -#ifdef USE_ROCM - // TODO: Remove this once we upgrade to torch 2.11. - // ROCm still uses torch 2.10, - // So we still need to use unstable torch ABI for now. - 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 -} - #ifdef USE_ROCM TORCH_LIBRARY_FRAGMENT(CONCAT(TORCH_EXTENSION_NAME, _custom_ar), custom_ar) { // Quick Reduce all-reduce kernels (ROCm-only; stays on legacy _C). @@ -48,20 +34,6 @@ TORCH_LIBRARY_FRAGMENT(CONCAT(TORCH_EXTENSION_NAME, _custom_ar), custom_ar) { custom_ar.def("qr_max_size", &qr_max_size); } - -// TODO: Remove this once ROCm upgrade to torch 2.11. -TORCH_LIBRARY_EXPAND(CONCAT(TORCH_EXTENSION_NAME, _cuda_utils), cuda_utils) { - // Cuda utils - // Gets the specified device attribute. - cuda_utils.def("get_device_attribute(int attribute, int device_id) -> int"); - cuda_utils.impl("get_device_attribute", &get_device_attribute); - - // Gets the maximum shared memory per block device attribute. - cuda_utils.def( - "get_max_shared_memory_per_block_device_attribute(int device_id) -> int"); - cuda_utils.impl("get_max_shared_memory_per_block_device_attribute", - &get_max_shared_memory_per_block_device_attribute); -} #endif REGISTER_EXTENSION(TORCH_EXTENSION_NAME)