[ROCm] [PyTorch] Move to stable abi since ROCm upgraded to torch 2.11 (#47128)
Signed-off-by: tjtanaa <tunjian.tan@embeddedllm.com>
TJian committed
Jul 2, 2026 at 18:34 UTC
de2a8fc042b6e9ca834ff20746734ed32b5ca219
6 files changed
+22
-131
CMakeLists.txt
+17
-29
@@ -350,9 +350,7 @@ endif()
350
if(VLLM_GPU_LANG STREQUAL "HIP")
351
set(VLLM_EXT_SRC
352
"csrc/torch_bindings.cpp"
353
- "csrc/custom_quickreduce.cu"
354
- "csrc/cuda_view.cu"
355
- "csrc/libtorch_stable/cuda_utils_kernels.cu")
353
+ "csrc/custom_quickreduce.cu")
354
355
message(STATUS "Enabling C extension.")
356
define_extension_target(
@@ -380,6 +378,8 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
378
#
379
set(VLLM_STABLE_EXT_SRC
380
"csrc/libtorch_stable/torch_bindings.cpp"
381
+ "csrc/libtorch_stable/cuda_view.cu"
382
+ "csrc/libtorch_stable/cuda_utils_kernels.cu"
383
"csrc/libtorch_stable/activation_kernels.cu"
384
"csrc/libtorch_stable/quantization/activation_kernels.cu"
385
"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")
459
FetchContent_MakeAvailable(cutlass)
460
461
list(APPEND VLLM_STABLE_EXT_SRC
462
- "csrc/libtorch_stable/cuda_view.cu"
463
- "csrc/libtorch_stable/cuda_utils_kernels.cu"
462
"csrc/libtorch_stable/cutlass_extensions/common.cpp"
463
"csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_entry.cu"
464
"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")
1089
USE_SABI 3
1090
WITH_SOABI)
1091
1092
+ # Set TORCH_TARGET_VERSION for stable ABI compatibility.
1093
+ # This ensures we only use C-shim APIs available in PyTorch 2.11.
1094
+ # _C_stable_libtorch is abi compatible with PyTorch >= TORCH_TARGET_VERSION
1095
+ # which is currently set to 2.11.
1096
+ target_compile_definitions(_C_stable_libtorch PRIVATE
1097
+ TORCH_TARGET_VERSION=0x020B000000000000ULL)
1098
+
1099
# Needed to use cuda/hip APIs from C-shim
1100
if(VLLM_GPU_LANG STREQUAL "CUDA")
1096
- # Set TORCH_TARGET_VERSION for stable ABI compatibility.
1097
- # This ensures we only use C-shim APIs available in PyTorch 2.11.
1098
- # _C_stable_libtorch is abi compatible with PyTorch >= TORCH_TARGET_VERSION
1099
- # which is currently set to 2.11.
1100
- target_compile_definitions(_C_stable_libtorch PRIVATE
1101
- TORCH_TARGET_VERSION=0x020B000000000000ULL)
1101
target_compile_definitions(_C_stable_libtorch PRIVATE USE_CUDA)
1102
if(COOPERATIVE_TOPK_ARCHS)
1103
target_compile_definitions(_C_stable_libtorch PRIVATE
@@ -1108,12 +1107,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
1107
target_compile_definitions(_C_stable_libtorch PRIVATE
1108
CUTLASS_ENABLE_DIRECT_CUDA_DRIVER_CALL=1)
1109
elseif(VLLM_GPU_LANG STREQUAL "HIP")
1111
- # Set TORCH_TARGET_VERSION for stable ABI compatibility.
1112
- # This ensures we only use C-shim APIs available in PyTorch 2.10.
1113
- # _C_stable_libtorch is abi compatible with PyTorch >= TORCH_TARGET_VERSION
1114
- # which is currently set to 2.10.
1115
- target_compile_definitions(_C_stable_libtorch PRIVATE
1116
- TORCH_TARGET_VERSION=0x020A000000000000ULL)
1110
target_compile_definitions(_C_stable_libtorch PRIVATE USE_ROCM)
1111
endif()
1112
@@ -1321,25 +1314,20 @@ define_extension_target(
1314
USE_SABI 3
1315
WITH_SOABI)
1316
1317
+# Set TORCH_TARGET_VERSION for stable ABI compatibility.
1318
+# This ensures we only use C-shim APIs available in PyTorch 2.11.
1319
+# _moe_C_stable_libtorch is abi compatible with PyTorch >= TORCH_TARGET_VERSION
1320
+# which is currently set to 2.11.
1321
+target_compile_definitions(_moe_C_stable_libtorch PRIVATE
1322
+ TORCH_TARGET_VERSION=0x020B000000000000ULL)
1323
+
1324
# Needed to use cuda/hip APIs from C-shim
1325
if(VLLM_GPU_LANG STREQUAL "CUDA")
1326
- # Set TORCH_TARGET_VERSION for stable ABI compatibility.
1327
- # This ensures we only use C-shim APIs available in PyTorch 2.11.
1328
- # _moe_C_stable_libtorch is abi compatible with PyTorch >= TORCH_TARGET_VERSION
1329
- # which is currently set to 2.11.
1330
- target_compile_definitions(_moe_C_stable_libtorch PRIVATE
1331
- TORCH_TARGET_VERSION=0x020B000000000000ULL)
1326
target_compile_definitions(_moe_C_stable_libtorch PRIVATE USE_CUDA)
1327
# Needed by CUTLASS kernels
1328
target_compile_definitions(_moe_C_stable_libtorch PRIVATE
1329
CUTLASS_ENABLE_DIRECT_CUDA_DRIVER_CALL=1)
1330
elseif(VLLM_GPU_LANG STREQUAL "HIP")
1337
- # Set TORCH_TARGET_VERSION for stable ABI compatibility.
1338
- # This ensures we only use C-shim APIs available in PyTorch 2.10.
1339
- # _moe_C_stable_libtorch is abi compatible with PyTorch >= TORCH_TARGET_VERSION
1340
- # which is currently set to 2.10.
1341
- target_compile_definitions(_moe_C_stable_libtorch PRIVATE
1342
- TORCH_TARGET_VERSION=0x020A000000000000ULL)
1331
target_compile_definitions(_moe_C_stable_libtorch PRIVATE USE_ROCM)
1332
endif()
1333
csrc/cuda_view.cu
deleted
-60
@@ -1,60 +0,0 @@
1
-// TODO: Remove this once ROCm upgrade to torch 2.11.
2
-#include <torch/all.h>
3
-#include <torch/cuda.h>
4
-#include <cuda_runtime.h>
5
-
6
-// This function assumes that `cpu_tensor` is a CPU tensor,
7
-// and that UVA (Unified Virtual Addressing) is enabled.
8
-torch::Tensor get_cuda_view_from_cpu_tensor(torch::Tensor& cpu_tensor) {
9
- TORCH_CHECK(cpu_tensor.device().is_cpu(), "Input tensor must be on CPU");
10
-
11
- // handle empty tensor
12
- if (cpu_tensor.numel() == 0) {
13
- return torch::empty(cpu_tensor.sizes(),
14
- cpu_tensor.options().device(torch::kCUDA));
15
- }
16
-
17
- if (cpu_tensor.is_pinned()) {
18
- // If CPU tensor is pinned, directly get the device pointer.
19
- void* host_ptr = const_cast<void*>(cpu_tensor.data_ptr());
20
- void* device_ptr = nullptr;
21
- cudaError_t err = cudaHostGetDevicePointer(&device_ptr, host_ptr, 0);
22
- TORCH_CHECK(err == cudaSuccess,
23
- "cudaHostGetDevicePointer failed: ", cudaGetErrorString(err));
24
-
25
- return torch::from_blob(
26
- device_ptr, cpu_tensor.sizes(), cpu_tensor.strides(),
27
- [base = cpu_tensor](void*) {}, // keep cpu tensor alive
28
- cpu_tensor.options().device(torch::kCUDA));
29
- }
30
-
31
- // If CPU tensor is not pinned, allocate a new pinned memory buffer.
32
- torch::Tensor contiguous_cpu = cpu_tensor.contiguous();
33
- size_t nbytes = contiguous_cpu.nbytes();
34
-
35
- void* host_ptr = nullptr;
36
- cudaError_t err = cudaHostAlloc(&host_ptr, nbytes, cudaHostAllocMapped);
37
- if (err != cudaSuccess) {
38
- AT_ERROR("cudaHostAlloc failed: ", cudaGetErrorString(err));
39
- }
40
-
41
- err = cudaMemcpy(host_ptr, contiguous_cpu.data_ptr(), nbytes,
42
- cudaMemcpyDefault);
43
- if (err != cudaSuccess) {
44
- cudaFreeHost(host_ptr);
45
- AT_ERROR("cudaMemcpy failed: ", cudaGetErrorString(err));
46
- }
47
-
48
- void* device_ptr = nullptr;
49
- err = cudaHostGetDevicePointer(&device_ptr, host_ptr, 0);
50
- if (err != cudaSuccess) {
51
- cudaFreeHost(host_ptr);
52
- AT_ERROR("cudaHostGetDevicePointer failed: ", cudaGetErrorString(err));
53
- }
54
-
55
- auto deleter = [host_ptr](void*) { cudaFreeHost(host_ptr); };
56
-
57
- return torch::from_blob(device_ptr, contiguous_cpu.sizes(),
58
- contiguous_cpu.strides(), deleter,
59
- contiguous_cpu.options().device(torch::kCUDA));
60
-}
\ No newline at end of file
csrc/libtorch_stable/ops.h
+3
-4
@@ -181,13 +181,12 @@ torch::stable::Tensor awq_dequantize(torch::stable::Tensor _kernel,
181
// AllSpark ops: declarations are in the source files
182
// (allspark_repack.cu and allspark_qgemm_w8a16.cu)
183
184
-// TODO: Move this out once ROCm upgrade their torch to 2.11.
185
-// CPU tensor -> CUDA UVA view (shared CUDA)
184
+#endif
185
+
186
+// CPU tensor -> CUDA UVA view (shared CUDA/ROCm)
187
torch::stable::Tensor get_cuda_view_from_cpu_tensor(
188
torch::stable::Tensor& cpu_tensor);
189
189
-#endif
190
-
190
// Attention kernels (shared CUDA/ROCm)
191
void merge_attn_states(
192
torch::stable::Tensor& output,
csrc/libtorch_stable/torch_bindings.cpp
+2
-7
@@ -29,11 +29,10 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
29
"()");
30
ops.def("permute_cols(Tensor A, Tensor perm) -> Tensor");
31
32
-#ifndef USE_ROCM
33
-
34
- // TODO: Remove this once ROCm upgrade to torch 2.11.
32
ops.def("get_cuda_view_from_cpu_tensor(Tensor cpu_tensor) -> Tensor");
33
34
+#ifndef USE_ROCM
35
+
36
// Note about marlin kernel 'workspace' arguments:
37
// Technically these should be mutable since they are modified by the kernel.
38
// 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) {
752
ops.impl("paged_attention_v2", TORCH_BOX(&paged_attention_v2));
753
}
754
756
-// TODO: Remove this once ROCm upgrade to torch 2.11.
757
-#ifndef USE_ROCM
755
STABLE_TORCH_LIBRARY_IMPL(_C, CPU, ops) {
756
ops.impl("get_cuda_view_from_cpu_tensor",
757
TORCH_BOX(&get_cuda_view_from_cpu_tensor));
@@ -773,8 +770,6 @@ STABLE_TORCH_LIBRARY_IMPL(_C_cuda_utils, CompositeExplicitAutograd,
770
TORCH_BOX(&get_max_shared_memory_per_block_device_attribute));
771
}
772
776
-#endif
777
-
773
// These capability-check functions take only primitive args (no tensors), so
774
// there is no device to dispatch on. CompositeExplicitAutograd makes them
775
// available for all backends. This is the stable ABI equivalent of calling
csrc/ops.h
-3
@@ -67,7 +67,4 @@ void qr_open_handles(fptr_t _fa, const std::vector<torch::Tensor>& handles);
67
void qr_all_reduce(fptr_t _fa, torch::Tensor& inp, torch::Tensor& out,
68
int64_t quant_level, bool cast_bf2half = false);
69
int64_t qr_max_size();
70
-
71
-// TODO: Remove this once ROCm upgrade to torch 2.11.
72
-torch::Tensor get_cuda_view_from_cpu_tensor(torch::Tensor& cpu_tensor);
70
#endif
csrc/torch_bindings.cpp
-28
@@ -2,7 +2,6 @@
2
// cache.h, which is no longer included here after cache ops moved to
3
// _C_stable_libtorch).
4
#include <torch/all.h>
5
-#include "cuda_utils.h"
5
#include "ops.h"
6
#include "core/registration.h"
7
#include <torch/library.h>
@@ -18,19 +17,6 @@
17
// https://docs.google.com/document/d/1_W62p8WJOQQUzPsJYa7s701JXt0qf2OfLub2sbkHOaU/edit#heading=h.ptttacy8y1u9
18
// https://github.com/pytorch/pytorch/blob/main/aten/src/ATen/native/README.md#annotations
19
21
-TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
22
- // vLLM custom ops
23
-
24
-#ifdef USE_ROCM
25
- // TODO: Remove this once we upgrade to torch 2.11.
26
- // ROCm still uses torch 2.10,
27
- // So we still need to use unstable torch ABI for now.
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);
31
-#endif
32
-}
33
-
20
#ifdef USE_ROCM
21
TORCH_LIBRARY_FRAGMENT(CONCAT(TORCH_EXTENSION_NAME, _custom_ar), custom_ar) {
22
// 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) {
34
35
custom_ar.def("qr_max_size", &qr_max_size);
36
}
51
-
52
-// TODO: Remove this once ROCm upgrade to torch 2.11.
53
-TORCH_LIBRARY_EXPAND(CONCAT(TORCH_EXTENSION_NAME, _cuda_utils), cuda_utils) {
54
- // Cuda utils
55
- // Gets the specified device attribute.
56
- cuda_utils.def("get_device_attribute(int attribute, int device_id) -> int");
57
- cuda_utils.impl("get_device_attribute", &get_device_attribute);
58
-
59
- // Gets the maximum shared memory per block device attribute.
60
- cuda_utils.def(
61
- "get_max_shared_memory_per_block_device_attribute(int device_id) -> int");
62
- cuda_utils.impl("get_max_shared_memory_per_block_device_attribute",
63
- &get_max_shared_memory_per_block_device_attribute);
64
-}
37
#endif
38
39
REGISTER_EXTENSION(TORCH_EXTENSION_NAME)