A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 1, 20:00 UTC
[Bugfix] [ROCm] [Critical] fallback to regular abi for ROCm (#44648)
Signed-off-by: tjtanaa <tunjian.tan@embeddedllm.com>
TJian committed
Jun 5, 2026 at 23:51 UTC
aa6fb8a329bf092a3fd29cfb399aa22feed7f5a1
6 files changed
+127
-25
CMakeLists.txt
+16
-8
@@ -590,8 +590,11 @@ endif()
590
591
if (VLLM_GPU_LANG STREQUAL "HIP")
592
# Add QuickReduce kernels (ROCm-only; not part of stable ABI migration).
593
+ # TODO: Remove the cuda_view when ROCm upgrade to torch 2.11.
594
list(APPEND VLLM_EXT_SRC
595
"csrc/custom_quickreduce.cu"
596
+ "csrc/cuda_view.cu"
597
+ "csrc/libtorch_stable/cuda_utils_kernels.cu"
598
)
599
# if ROCM endif
600
endif()
@@ -621,7 +624,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
624
#
625
set(VLLM_STABLE_EXT_SRC
626
"csrc/libtorch_stable/torch_bindings.cpp"
624
- "csrc/libtorch_stable/cuda_view.cu"
627
"csrc/libtorch_stable/activation_kernels.cu"
628
"csrc/libtorch_stable/quantization/w8a8/int8/scaled_quant.cu"
629
"csrc/libtorch_stable/quantization/w8a8/fp8/common.cu"
@@ -649,6 +651,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
651
652
if(VLLM_GPU_LANG STREQUAL "CUDA")
653
list(APPEND VLLM_STABLE_EXT_SRC
654
+ "csrc/libtorch_stable/cuda_view.cu"
655
"csrc/libtorch_stable/cuda_utils_kernels.cu"
656
"csrc/libtorch_stable/cutlass_extensions/common.cpp"
657
"csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_entry.cu"
@@ -1071,20 +1074,25 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
1074
USE_SABI 3
1075
WITH_SOABI)
1076
1074
- # Set TORCH_TARGET_VERSION for stable ABI compatibility.
1075
- # This ensures we only use C-shim APIs available in PyTorch 2.11.
1076
- # _C_stable_libtorch is abi compatible with PyTorch >= TORCH_TARGET_VERSION
1077
- # which is currently set to 2.11.
1078
- target_compile_definitions(_C_stable_libtorch PRIVATE
1079
- TORCH_TARGET_VERSION=0x020B000000000000ULL)
1080
-
1077
# Needed to use cuda/hip APIs from C-shim
1078
if(VLLM_GPU_LANG STREQUAL "CUDA")
1079
+ # Set TORCH_TARGET_VERSION for stable ABI compatibility.
1080
+ # This ensures we only use C-shim APIs available in PyTorch 2.11.
1081
+ # _C_stable_libtorch is abi compatible with PyTorch >= TORCH_TARGET_VERSION
1082
+ # which is currently set to 2.11.
1083
+ target_compile_definitions(_C_stable_libtorch PRIVATE
1084
+ TORCH_TARGET_VERSION=0x020B000000000000ULL)
1085
target_compile_definitions(_C_stable_libtorch PRIVATE USE_CUDA)
1086
# Needed by CUTLASS kernels
1087
target_compile_definitions(_C_stable_libtorch PRIVATE
1088
CUTLASS_ENABLE_DIRECT_CUDA_DRIVER_CALL=1)
1089
elseif(VLLM_GPU_LANG STREQUAL "HIP")
1090
+ # Set TORCH_TARGET_VERSION for stable ABI compatibility.
1091
+ # This ensures we only use C-shim APIs available in PyTorch 2.10.
1092
+ # _C_stable_libtorch is abi compatible with PyTorch >= TORCH_TARGET_VERSION
1093
+ # which is currently set to 2.10.
1094
+ target_compile_definitions(_C_stable_libtorch PRIVATE
1095
+ TORCH_TARGET_VERSION=0x020A000000000000ULL)
1096
target_compile_definitions(_C_stable_libtorch PRIVATE USE_ROCM)
1097
endif()
1098
csrc/cuda_view.cu
new
+60
@@ -0,0 +1,60 @@
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
+4
-3
@@ -162,12 +162,13 @@ torch::stable::Tensor awq_dequantize(torch::stable::Tensor _kernel,
162
// AllSpark ops: declarations are in the source files
163
// (allspark_repack.cu and allspark_qgemm_w8a16.cu)
164
165
-#endif
166
-
167
-// CPU tensor -> CUDA UVA view (shared CUDA/ROCm)
165
+// TODO: Move this out once ROCm upgrade their torch to 2.11.
166
+// CPU tensor -> CUDA UVA view (shared CUDA)
167
torch::stable::Tensor get_cuda_view_from_cpu_tensor(
168
torch::stable::Tensor& cpu_tensor);
169
170
+#endif
171
+
172
// Attention kernels (shared CUDA/ROCm)
173
void merge_attn_states(
174
torch::stable::Tensor& output,
csrc/libtorch_stable/torch_bindings.cpp
+20
-14
@@ -28,9 +28,11 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
28
"output_s, int group_size, float eps, float int8_min, float int8_max) -> "
29
"()");
30
31
+#ifndef USE_ROCM
32
+
33
+ // TODO: Remove this once ROCm upgrade to torch 2.11.
34
ops.def("get_cuda_view_from_cpu_tensor(Tensor cpu_tensor) -> Tensor");
35
33
-#ifndef USE_ROCM
36
ops.def("permute_cols(Tensor A, Tensor perm) -> Tensor");
37
#endif
38
@@ -676,11 +678,28 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) {
678
ops.impl("paged_attention_v2", TORCH_BOX(&paged_attention_v2));
679
}
680
681
+// TODO: Remove this once ROCm upgrade to torch 2.11.
682
+#ifndef USE_ROCM
683
STABLE_TORCH_LIBRARY_IMPL(_C, CPU, ops) {
684
ops.impl("get_cuda_view_from_cpu_tensor",
685
TORCH_BOX(&get_cuda_view_from_cpu_tensor));
686
}
687
688
+STABLE_TORCH_LIBRARY_FRAGMENT(_C_cuda_utils, cuda_utils) {
689
+ cuda_utils.def("get_device_attribute(int attribute, int device_id) -> int");
690
+ cuda_utils.def(
691
+ "get_max_shared_memory_per_block_device_attribute(int device_id) -> int");
692
+}
693
+
694
+STABLE_TORCH_LIBRARY_IMPL(_C_cuda_utils, CompositeExplicitAutograd,
695
+ cuda_utils) {
696
+ cuda_utils.impl("get_device_attribute", TORCH_BOX(&get_device_attribute));
697
+ cuda_utils.impl("get_max_shared_memory_per_block_device_attribute",
698
+ TORCH_BOX(&get_max_shared_memory_per_block_device_attribute));
699
+}
700
+
701
+#endif
702
+
703
// These capability-check functions take only primitive args (no tensors), so
704
// there is no device to dispatch on. CompositeExplicitAutograd makes them
705
// available for all backends. This is the stable ABI equivalent of calling
@@ -701,19 +720,6 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CompositeExplicitAutograd, ops) {
720
ops.impl("ggml_moe_get_block_size", TORCH_BOX(&ggml_moe_get_block_size));
721
}
722
704
-STABLE_TORCH_LIBRARY_FRAGMENT(_C_cuda_utils, cuda_utils) {
705
- cuda_utils.def("get_device_attribute(int attribute, int device_id) -> int");
706
- cuda_utils.def(
707
- "get_max_shared_memory_per_block_device_attribute(int device_id) -> int");
708
-}
709
-
710
-STABLE_TORCH_LIBRARY_IMPL(_C_cuda_utils, CompositeExplicitAutograd,
711
- cuda_utils) {
712
- cuda_utils.impl("get_device_attribute", TORCH_BOX(&get_device_attribute));
713
- cuda_utils.impl("get_max_shared_memory_per_block_device_attribute",
714
- TORCH_BOX(&get_max_shared_memory_per_block_device_attribute));
715
-}
716
-
723
// Cache ops
724
STABLE_TORCH_LIBRARY_FRAGMENT(_C_cache_ops, ops) {
725
// Swap in (out) the cache blocks from src to dst.
csrc/ops.h
+3
@@ -102,4 +102,7 @@ void qr_open_handles(fptr_t _fa, const std::vector<torch::Tensor>& handles);
102
void qr_all_reduce(fptr_t _fa, torch::Tensor& inp, torch::Tensor& out,
103
int64_t quant_level, bool cast_bf2half = false);
104
int64_t qr_max_size();
105
+
106
+// TODO: Remove this once ROCm upgrade to torch 2.11.
107
+torch::Tensor get_cuda_view_from_cpu_tensor(torch::Tensor& cpu_tensor);
108
#endif
csrc/torch_bindings.cpp
+24
@@ -2,6 +2,7 @@
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"
6
#include "ops.h"
7
#include "core/registration.h"
8
#include <torch/library.h>
@@ -31,6 +32,15 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
32
ops.def("weak_ref_tensor(Tensor input) -> Tensor");
33
ops.impl("weak_ref_tensor", torch::kCUDA, &weak_ref_tensor);
34
35
+#ifdef USE_ROCM
36
+ // TODO: Remove this once we upgrade to torch 2.11.
37
+ // ROCm still uses torch 2.10,
38
+ // So we still need to use unstable torch ABI for now.
39
+ ops.def("get_cuda_view_from_cpu_tensor(Tensor cpu_tensor) -> Tensor");
40
+ ops.impl("get_cuda_view_from_cpu_tensor", torch::kCPU,
41
+ &get_cuda_view_from_cpu_tensor);
42
+#endif
43
+
44
// Activation ops (quantized only — basic ops moved to _C_stable_libtorch)
45
ops.def(
46
"silu_and_mul_quant(Tensor! result, Tensor input, Tensor scale) -> ()");
@@ -159,6 +169,20 @@ TORCH_LIBRARY_FRAGMENT(CONCAT(TORCH_EXTENSION_NAME, _custom_ar), custom_ar) {
169
170
custom_ar.def("qr_max_size", &qr_max_size);
171
}
172
+
173
+// TODO: Remove this once ROCm upgrade to torch 2.11.
174
+TORCH_LIBRARY_EXPAND(CONCAT(TORCH_EXTENSION_NAME, _cuda_utils), cuda_utils) {
175
+ // Cuda utils
176
+ // Gets the specified device attribute.
177
+ cuda_utils.def("get_device_attribute(int attribute, int device_id) -> int");
178
+ cuda_utils.impl("get_device_attribute", &get_device_attribute);
179
+
180
+ // Gets the maximum shared memory per block device attribute.
181
+ cuda_utils.def(
182
+ "get_max_shared_memory_per_block_device_attribute(int device_id) -> int");
183
+ cuda_utils.impl("get_max_shared_memory_per_block_device_attribute",
184
+ &get_max_shared_memory_per_block_device_attribute);
185
+}
186
#endif
187
188
REGISTER_EXTENSION(TORCH_EXTENSION_NAME)