mirror of
https://github.com/vllm-project/vllm.git
synced 2026-08-13 09:18:12 +00:00
[ROCm] [PyTorch] Move to stable abi since ROCm upgraded to torch 2.11 (#47128)
Signed-off-by: tjtanaa <[email protected]>
This commit is contained in:
TJian
+17
-29
@@ -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()
|
||||
|
||||
|
||||
@@ -1,60 +0,0 @@
|
||||
// TODO: Remove this once ROCm upgrade to torch 2.11.
|
||||
#include <torch/all.h>
|
||||
#include <torch/cuda.h>
|
||||
#include <cuda_runtime.h>
|
||||
|
||||
// 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<void*>(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));
|
||||
}
|
||||
@@ -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,
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -67,7 +67,4 @@ void qr_open_handles(fptr_t _fa, const std::vector<torch::Tensor>& 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
|
||||
|
||||
@@ -2,7 +2,6 @@
|
||||
// cache.h, which is no longer included here after cache ops moved to
|
||||
// _C_stable_libtorch).
|
||||
#include <torch/all.h>
|
||||
#include "cuda_utils.h"
|
||||
#include "ops.h"
|
||||
#include "core/registration.h"
|
||||
#include <torch/library.h>
|
||||
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user