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 6, 10:45 UTC
[1/n] Migrate permute_cols to libtorch stable ABI (#31509)
Signed-off-by: Mikayla Gawarecki <mikaylagawarecki@gmail.com>
mikaylagawarecki committed
Mar 19, 2026 at 11:27 UTC
8b10e4fb316c14cfdb3109ac6f87722ec2a6c3c8
9 files changed
+114
-22
CMakeLists.txt
+42
-1
@@ -340,7 +340,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
340
341
list(APPEND VLLM_EXT_SRC
342
"csrc/quantization/awq/gemm_kernels.cu"
343
- "csrc/permute_cols.cu"
343
"csrc/quantization/w8a8/cutlass/scaled_mm_entry.cu"
344
"csrc/quantization/fp4/nvfp4_quant_entry.cu"
345
"csrc/quantization/fp4/nvfp4_scaled_mm_entry.cu"
@@ -986,6 +985,48 @@ define_extension_target(
985
# Setting this variable sidesteps the issue by calling the driver directly.
986
target_compile_definitions(_C PRIVATE CUTLASS_ENABLE_DIRECT_CUDA_DRIVER_CALL=1)
987
988
+# add OR VLLM_GPU_LANG STREQUAL "HIP" here once
989
+# https://github.com/vllm-project/vllm/issues/35163 is resolved
990
+if(VLLM_GPU_LANG STREQUAL "CUDA")
991
+ #
992
+ # _C_stable_libtorch extension (ops registered via STABLE_TORCH_LIBRARY)
993
+ #
994
+ set(VLLM_STABLE_EXT_SRC
995
+ "csrc/libtorch_stable/torch_bindings.cpp")
996
+
997
+ if(VLLM_GPU_LANG STREQUAL "CUDA")
998
+ list(APPEND VLLM_STABLE_EXT_SRC "csrc/libtorch_stable/permute_cols.cu")
999
+ endif()
1000
+
1001
+ if(VLLM_GPU_LANG STREQUAL "CUDA")
1002
+ set_gencode_flags_for_srcs(
1003
+ SRCS "${VLLM_STABLE_EXT_SRC}"
1004
+ CUDA_ARCHS "${CUDA_ARCHS}")
1005
+ endif()
1006
+
1007
+ message(STATUS "Enabling C_stable extension.")
1008
+ define_extension_target(
1009
+ _C_stable_libtorch
1010
+ DESTINATION vllm
1011
+ LANGUAGE ${VLLM_GPU_LANG}
1012
+ SOURCES ${VLLM_STABLE_EXT_SRC}
1013
+ COMPILE_FLAGS ${VLLM_GPU_FLAGS}
1014
+ ARCHITECTURES ${VLLM_GPU_ARCHES}
1015
+ USE_SABI 3
1016
+ WITH_SOABI)
1017
+
1018
+ # Set TORCH_TARGET_VERSION for stable ABI compatibility.
1019
+ # This ensures we only use C-shim APIs available in PyTorch 2.10.
1020
+ # _C_stable_libtorch is abi compatible with PyTorch >= TORCH_TARGET_VERSION
1021
+ # which is currently set to 2.10.
1022
+ target_compile_definitions(_C_stable_libtorch PRIVATE
1023
+ TORCH_TARGET_VERSION=0x020A000000000000ULL)
1024
+
1025
+ # Needed to use cuda APIs from C-shim
1026
+ target_compile_definitions(_C_stable_libtorch PRIVATE
1027
+ USE_CUDA)
1028
+endif()
1029
+
1030
#
1031
# _moe_C extension
1032
#
csrc/libtorch_stable/ops.h
new
+9
@@ -0,0 +1,9 @@
1
+#pragma once
2
+
3
+#include <torch/csrc/stable/library.h>
4
+#include <torch/csrc/stable/tensor.h>
5
+
6
+#ifndef USE_ROCM
7
+torch::stable::Tensor permute_cols(torch::stable::Tensor const& A,
8
+ torch::stable::Tensor const& perm);
9
+#endif
csrc/libtorch_stable/permute_cols.cu
renamed
+23
-17
@@ -1,10 +1,13 @@
1
-#include <torch/all.h>
2
-
3
-#include <ATen/cuda/CUDAContext.h>
4
-#include <c10/cuda/CUDAGuard.h>
1
+#include <torch/csrc/stable/library.h>
2
+#include <torch/csrc/stable/tensor.h>
3
+#include <torch/csrc/stable/accelerator.h>
4
+#include <torch/csrc/stable/ops.h>
5
+#include <torch/headeronly/core/ScalarType.h>
6
7
#include <cuda_fp16.h>
8
9
+#include "torch_utils.h"
10
+
11
static constexpr int default_threads = 256;
12
static constexpr int div_ceil(int a, int b) { return (a + b - 1) / b; }
13
@@ -64,19 +67,22 @@ __global__ void permute_cols_kernel(int4 const* __restrict__ a_int4_ptr,
67
68
// More efficient version of A[..., perm]
69
// taken from gptq_marlin.cu
67
-torch::Tensor permute_cols(torch::Tensor const& A, torch::Tensor const& perm) {
68
- const at::cuda::OptionalCUDAGuard device_guard(device_of(A));
69
- auto dev = A.get_device();
70
- auto stream = at::cuda::getCurrentCUDAStream(dev);
71
-
72
- TORCH_CHECK(A.scalar_type() == at::kHalf || A.scalar_type() == at::kBFloat16,
73
- "Currently only 16bit types are supported");
74
- TORCH_CHECK(A.is_contiguous(), "A must be contiguous");
75
- TORCH_CHECK(A.size(-1) % 8 == 0,
76
- "A columns must be a multiple of 8 (128bits)");
77
- auto A_2d = A.view({-1, A.size(-1)});
78
-
79
- torch::Tensor D = torch::empty_like(A);
70
+torch::stable::Tensor permute_cols(torch::stable::Tensor const& A,
71
+ torch::stable::Tensor const& perm) {
72
+ const int32_t dev = A.get_device_index();
73
+ const torch::stable::accelerator::DeviceGuard device_guard(dev);
74
+ const auto stream = get_current_cuda_stream(dev);
75
+
76
+ STD_TORCH_CHECK(
77
+ A.scalar_type() == torch::headeronly::ScalarType::Half ||
78
+ A.scalar_type() == torch::headeronly::ScalarType::BFloat16,
79
+ "Currently only 16bit types are supported");
80
+ STD_TORCH_CHECK(A.is_contiguous(), "A must be contiguous");
81
+ STD_TORCH_CHECK(A.size(-1) % 8 == 0,
82
+ "A columns must be a multiple of 8 (128bits)");
83
+ auto A_2d = torch::stable::view(A, {-1, A.size(-1)});
84
+
85
+ torch::stable::Tensor D = torch::stable::empty_like(A);
86
int sms;
87
cudaDeviceGetAttribute(&sms, cudaDevAttrMultiProcessorCount, dev);
88
int block_rows = div_ceil(A_2d.size(0), sms);
csrc/libtorch_stable/torch_bindings.cpp
new
+21
@@ -0,0 +1,21 @@
1
+#include "ops.h"
2
+#include "core/registration.h"
3
+
4
+#include <torch/csrc/stable/library.h>
5
+
6
+// Register ops with STABLE_TORCH_LIBRARY for libtorch stable ABI compatibility.
7
+// Note: We register under namespace "_C" so ops are accessible as
8
+// torch.ops._C.<op_name> for compatibility with existing code.
9
+STABLE_TORCH_LIBRARY_FRAGMENT(_C, m) {
10
+#ifndef USE_ROCM
11
+ m.def("permute_cols(Tensor A, Tensor perm) -> Tensor");
12
+#endif
13
+}
14
+
15
+STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, m) {
16
+#ifndef USE_ROCM
17
+ m.impl("permute_cols", TORCH_BOX(&permute_cols));
18
+#endif
19
+}
20
+
21
+REGISTER_EXTENSION(_C_stable_libtorch)
csrc/libtorch_stable/torch_utils.h
new
+13
@@ -0,0 +1,13 @@
1
+#pragma once
2
+
3
+#include <torch/csrc/inductor/aoti_torch/c/shim.h>
4
+#include <cuda_runtime.h>
5
+
6
+// Utility to get the current CUDA stream for a given device using stable APIs.
7
+// Returns a cudaStream_t for use in kernel launches.
8
+inline cudaStream_t get_current_cuda_stream(int32_t device_index) {
9
+ void* stream_ptr = nullptr;
10
+ TORCH_ERROR_CODE_CHECK(
11
+ aoti_torch_get_current_cuda_stream(device_index, &stream_ptr));
12
+ return reinterpret_cast<cudaStream_t>(stream_ptr);
13
+}
csrc/ops.h
-1
@@ -201,7 +201,6 @@ torch::Tensor awq_dequantize(torch::Tensor _kernel,
201
torch::Tensor _zeros, int64_t split_k_iters,
202
int64_t thx, int64_t thy);
203
204
-torch::Tensor permute_cols(torch::Tensor const& A, torch::Tensor const& perm);
204
#endif
205
206
torch::Tensor ggml_dequantize(torch::Tensor W, int64_t type, int64_t m,
csrc/torch_bindings.cpp
-3
@@ -303,9 +303,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
303
") -> Tensor");
304
// conditionally compiled so impl registration is in source file
305
306
- ops.def("permute_cols(Tensor A, Tensor perm) -> Tensor");
307
- ops.impl("permute_cols", torch::kCUDA, &permute_cols);
308
-
306
// Marlin Optimized Quantized GEMM (supports GPTQ, AWQ, FP8, NVFP4, MXFP4).
307
ops.def(
308
"marlin_gemm(Tensor a, Tensor? c_or_none, Tensor b_q_weight, "
setup.py
+5
@@ -597,6 +597,7 @@ class precompiled_wheel_utils:
597
with zipfile.ZipFile(wheel_path) as wheel:
598
files_to_copy = [
599
"vllm/_C.abi3.so",
600
+ "vllm/_C_stable_libtorch.abi3.so",
601
"vllm/_moe_C.abi3.so",
602
"vllm/_flashmla_C.abi3.so",
603
"vllm/_flashmla_extension_C.abi3.so",
@@ -932,6 +933,10 @@ if _is_cpu():
933
934
if _build_custom_ops():
935
ext_modules.append(CMakeExtension(name="vllm._C"))
936
+ # also _is_hip() once https://github.com/vllm-project/vllm/issues/35163 is
937
+ # fixed
938
+ if _is_cuda():
939
+ ext_modules.append(CMakeExtension(name="vllm._C_stable_libtorch"))
940
941
package_data = {
942
"vllm": [
vllm/platforms/cuda.py
+1
@@ -19,6 +19,7 @@ from typing_extensions import ParamSpec
19
20
# import custom ops, trigger op registration
21
import vllm._C # noqa
22
+import vllm._C_stable_libtorch # noqa
23
from vllm.logger import init_logger
24
from vllm.utils.import_utils import import_pynvml
25
from vllm.utils.torch_utils import cuda_device_count_stateless