A high-throughput and memory-efficient inference and serving engine for LLMs
[ROCm][Kernel] Enable permute_cols for ROCm (#44674)
Signed-off-by: charlifu <charlifu@amd.com>
Charlie Fu committed
Jun 7, 2026 at 04:50 UTC
228bcc436b0f09cf1824f00fd44f8f9c94060940
3 files changed
+3
-8
CMakeLists.txt
+1
-3
@@ -629,6 +629,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
629
"csrc/libtorch_stable/quantization/w8a8/fp8/common.cu"
630
"csrc/libtorch_stable/quantization/w8a8/fp8/per_token_group_quant.cu"
631
"csrc/libtorch_stable/quantization/w8a8/int8/per_token_group_quant.cu"
632
+ "csrc/libtorch_stable/permute_cols.cu"
633
"csrc/libtorch_stable/quantization/gptq/q_gemm.cu"
634
"csrc/libtorch_stable/quantization/gguf/gguf_kernel.cu"
635
"csrc/libtorch_stable/pos_encoding_kernels.cu"
@@ -657,9 +658,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
658
"csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_entry.cu"
659
"csrc/libtorch_stable/quantization/fp4/nvfp4_quant_entry.cu"
660
"csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_entry.cu"
660
- "csrc/libtorch_stable/permute_cols.cu"
661
- "csrc/libtorch_stable/quantization/w8a8/fp8/per_token_group_quant.cu"
662
- "csrc/libtorch_stable/quantization/w8a8/int8/per_token_group_quant.cu"
661
"csrc/libtorch_stable/quantization/awq/gemm_kernels.cu"
662
"csrc/libtorch_stable/minimax_reduce_rms_kernel.cu")
663
csrc/libtorch_stable/ops.h
+1
-1
@@ -24,10 +24,10 @@ void per_token_group_quant_int8(const torch::stable::Tensor& input,
24
int64_t group_size, double eps, double int8_min,
25
double int8_max);
26
27
-#ifndef USE_ROCM
27
torch::stable::Tensor permute_cols(torch::stable::Tensor const& A,
28
torch::stable::Tensor const& perm);
29
30
+#ifndef USE_ROCM
31
bool cutlass_scaled_mm_supports_fp8(int64_t cuda_device_capability);
32
bool cutlass_scaled_mm_supports_block_fp8(int64_t cuda_device_capability);
33
bool cutlass_group_gemm_supported(int64_t cuda_device_capability);
csrc/libtorch_stable/torch_bindings.cpp
+1
-4
@@ -27,13 +27,12 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
27
"per_token_group_quant_int8(Tensor input, Tensor! output_q, Tensor! "
28
"output_s, int group_size, float eps, float int8_min, float int8_max) -> "
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.
35
ops.def("get_cuda_view_from_cpu_tensor(Tensor cpu_tensor) -> Tensor");
35
-
36
- ops.def("permute_cols(Tensor A, Tensor perm) -> Tensor");
36
#endif
37
38
#ifndef USE_ROCM
@@ -562,9 +561,7 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) {
561
ops.impl("per_token_group_quant_int8",
562
TORCH_BOX(&per_token_group_quant_int8));
563
565
-#ifndef USE_ROCM
564
ops.impl("permute_cols", TORCH_BOX(&permute_cols));
567
-#endif
565
566
#ifndef USE_ROCM
567
// CUTLASS scaled_mm ops