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 3, 06:45 UTC
[6/n] Migrate activation kernels, gptq, gguf, non cutlass w8a8 to libtorch stable ABI (continued) (#42663)
Signed-off-by: Mikayla Gawarecki <mikaylagawarecki@gmail.com> Signed-off-by: Chris Leonard <chleonar@redhat.com> Co-authored-by: Mikayla Gawarecki <mikaylagawarecki@gmail.com> Co-authored-by: Shengqi Chen <harry-chen@outlook.com>
Chris Leonard committed
May 20, 2026 at 03:18 UTC
07aeaf9d4df870a76d5a0dc19d6a7e74b4be5d3b
28 files changed
+811
-640
CMakeLists.txt
+41
-22
@@ -312,19 +312,14 @@ set(VLLM_EXT_SRC
312
"csrc/attention/paged_attention_v2.cu"
313
"csrc/attention/merge_attn_states.cu"
314
"csrc/pos_encoding_kernels.cu"
315
- "csrc/activation_kernels.cu"
315
"csrc/layernorm_kernels.cu"
316
"csrc/fused_qknorm_rope_kernel.cu"
317
"csrc/layernorm_quant_kernels.cu"
318
"csrc/sampler.cu"
319
"csrc/topk.cu"
320
"csrc/cuda_view.cu"
322
- "csrc/quantization/gptq/q_gemm.cu"
323
- "csrc/quantization/w8a8/int8/scaled_quant.cu"
324
- "csrc/quantization/w8a8/fp8/common.cu"
321
"csrc/quantization/fused_kernels/fused_layernorm_dynamic_per_token_quant.cu"
322
"csrc/quantization/fused_kernels/fused_silu_mul_block_quant.cu"
327
- "csrc/quantization/gguf/gguf_kernel.cu"
323
"csrc/quantization/activation_kernels.cu"
324
"csrc/cuda_utils_kernels.cu"
325
"csrc/custom_all_reduce.cu"
@@ -628,33 +623,33 @@ define_extension_target(
623
# Setting this variable sidesteps the issue by calling the driver directly.
624
target_compile_definitions(_C PRIVATE CUTLASS_ENABLE_DIRECT_CUDA_DRIVER_CALL=1)
625
631
-# add OR VLLM_GPU_LANG STREQUAL "HIP" here once
632
-# https://github.com/vllm-project/vllm/issues/35163 is resolved
633
-if(VLLM_GPU_LANG STREQUAL "CUDA")
626
+if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
627
#
628
# _C_stable_libtorch extension (ops registered via STABLE_TORCH_LIBRARY)
629
#
630
set(VLLM_STABLE_EXT_SRC
631
"csrc/libtorch_stable/torch_bindings.cpp"
639
- "csrc/cutlass_extensions/common.cpp"
640
- "csrc/cuda_utils_kernels.cu"
641
- "csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_entry.cu"
642
- "csrc/libtorch_stable/quantization/fp4/nvfp4_quant_entry.cu"
643
- "csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_entry.cu")
632
+ "csrc/libtorch_stable/activation_kernels.cu"
633
+ "csrc/libtorch_stable/quantization/w8a8/int8/scaled_quant.cu"
634
+ "csrc/libtorch_stable/quantization/w8a8/fp8/common.cu"
635
+ "csrc/libtorch_stable/quantization/gptq/q_gemm.cu"
636
+ "csrc/libtorch_stable/quantization/gguf/gguf_kernel.cu")
637
638
if(VLLM_GPU_LANG STREQUAL "CUDA")
639
list(APPEND VLLM_STABLE_EXT_SRC
640
+ "csrc/cuda_utils_kernels.cu"
641
+ "csrc/cutlass_extensions/common.cpp"
642
+ "csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_entry.cu"
643
+ "csrc/libtorch_stable/quantization/fp4/nvfp4_quant_entry.cu"
644
+ "csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_entry.cu"
645
"csrc/libtorch_stable/permute_cols.cu"
646
"csrc/libtorch_stable/quantization/w8a8/fp8/per_token_group_quant.cu"
647
"csrc/libtorch_stable/quantization/w8a8/int8/per_token_group_quant.cu"
648
"csrc/libtorch_stable/quantization/awq/gemm_kernels.cu")
651
- endif()
649
653
- if(VLLM_GPU_LANG STREQUAL "CUDA")
650
set_gencode_flags_for_srcs(
651
SRCS "${VLLM_STABLE_EXT_SRC}"
652
CUDA_ARCHS "${CUDA_ARCHS}")
657
- endif()
653
654
# DeepSeek V3 fused A GEMM kernel (requires SM 9.0+, Hopper and later)
655
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
@@ -1034,6 +1029,9 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
1029
message(STATUS "Building hadacore")
1030
endif()
1031
1032
+ # if CUDA endif
1033
+ endif()
1034
+
1035
message(STATUS "Enabling C_stable extension.")
1036
define_extension_target(
1037
_C_stable_libtorch
@@ -1053,13 +1051,34 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
1051
target_compile_definitions(_C_stable_libtorch PRIVATE
1052
TORCH_TARGET_VERSION=0x020A000000000000ULL)
1053
1056
- # Needed to use cuda APIs from C-shim
1057
- target_compile_definitions(_C_stable_libtorch PRIVATE
1058
- USE_CUDA)
1054
+ # Needed to use cuda/hip APIs from C-shim
1055
+ if(VLLM_GPU_LANG STREQUAL "CUDA")
1056
+ target_compile_definitions(_C_stable_libtorch PRIVATE USE_CUDA)
1057
+ # Needed by CUTLASS kernels
1058
+ target_compile_definitions(_C_stable_libtorch PRIVATE
1059
+ CUTLASS_ENABLE_DIRECT_CUDA_DRIVER_CALL=1)
1060
+ elseif(VLLM_GPU_LANG STREQUAL "HIP")
1061
+ target_compile_definitions(_C_stable_libtorch PRIVATE USE_ROCM)
1062
+ endif()
1063
1060
- # Needed by CUTLASS kernels
1061
- target_compile_definitions(_C_stable_libtorch PRIVATE
1062
- CUTLASS_ENABLE_DIRECT_CUDA_DRIVER_CALL=1)
1064
+ # On ROCm, _C_stable_libtorch calls raw HIP APIs (e.g. hipGetDevice in
1065
+ # get_device_prop()) which must resolve to the same libamdhip64.so that
1066
+ # PyTorch uses. When PyTorch bundles its own copy (pip/conda wheels),
1067
+ # the raw HIP calls would otherwise resolve to the system ROCm copy,
1068
+ # initializing a second HIP runtime that corrupts device state (wrong
1069
+ # device on DeviceGuard, core dumps on multi-GPU tests).
1070
+ #
1071
+ # If PyTorch doesn't bundle libamdhip64 (built from source against system
1072
+ # ROCm), there is only one copy in the process and no action is needed —
1073
+ # the HIP compiler already links the system libamdhip64 automatically.
1074
+ if(VLLM_GPU_LANG STREQUAL "HIP")
1075
+ find_library(_STABLE_TORCH_AMDHIP64 amdhip64
1076
+ PATHS "${TORCH_INSTALL_PREFIX}/lib" NO_DEFAULT_PATH)
1077
+ if(_STABLE_TORCH_AMDHIP64)
1078
+ message(STATUS "Found PyTorch-bundled libamdhip64 at ${_STABLE_TORCH_AMDHIP64}")
1079
+ target_link_libraries(_C_stable_libtorch PRIVATE ${_STABLE_TORCH_AMDHIP64})
1080
+ endif()
1081
+ endif()
1082
endif()
1083
1084
#
csrc/attention/dtype_fp8.cuh
+2
-1
@@ -1,6 +1,7 @@
1
#pragma once
2
3
#include "attention_generic.cuh"
4
+#include "torch_utils.h"
5
6
#include <stdint.h>
7
#ifdef ENABLE_FP8
@@ -30,7 +31,7 @@ inline Fp8KVCacheDataType get_fp8_kv_cache_data_type(
31
} else if (dtype_str == "fp8_e5m2") {
32
return Fp8KVCacheDataType::kFp8E5M2;
33
}
33
- TORCH_CHECK(false, "Unsupported fp8 kv cache data type: ", dtype_str);
34
+ TORCH_UTILS_CHECK(false, "Unsupported fp8 kv cache data type: ", dtype_str);
35
}
36
37
// fp8 vector types for quantization of kv cache
csrc/cuda_vec_utils.cuh
+2
@@ -9,6 +9,8 @@
9
10
#ifdef USE_ROCM
11
#include <hip/hip_runtime.h>
12
+ #include <hip/hip_bf16.h>
13
+ #include <hip/hip_fp16.h>
14
#else
15
#include <cuda_bf16.h>
16
#include <cuda_fp16.h>
csrc/cutlass_extensions/torch_utils.hpp
+2
-4
@@ -1,5 +1,7 @@
1
#pragma once
2
3
+#include "torch_utils.h"
4
+
5
// This header is shared between _C (unstable ABI, used by machete) and
6
// _C_stable_libtorch (stable ABI, used by W4A8/sparse). TORCH_TARGET_VERSION
7
// is defined only for the stable target, so we switch includes and types
@@ -8,13 +10,9 @@
10
#include <torch/csrc/stable/tensor.h>
11
#include <torch/headeronly/util/BFloat16.h>
12
#include <torch/headeronly/util/Half.h>
11
- #include <torch/headeronly/util/shim_utils.h> // for STD_TORCH_CHECK
13
using TorchTensor = torch::stable::Tensor;
13
- #define TORCH_UTILS_CHECK STD_TORCH_CHECK
14
#else
15
- #include <torch/all.h>
15
using TorchTensor = torch::Tensor;
17
- #define TORCH_UTILS_CHECK TORCH_CHECK
16
#endif
17
18
#include "cute/layout.hpp"
csrc/libtorch_stable/activation_kernels.cu
renamed
+115
-105
@@ -1,12 +1,12 @@
1
-#include <ATen/cuda/CUDAContext.h>
2
-#include <torch/all.h>
3
-#include <c10/cuda/CUDAGuard.h>
1
+#include <cuda.h>
2
+#include <torch/csrc/stable/tensor.h>
3
4
#include <cmath>
5
7
-#include "cuda_compat.h"
8
-#include "cuda_vec_utils.cuh"
6
+#include "../cuda_compat.h"
7
+#include "../cuda_vec_utils.cuh"
8
#include "dispatch_utils.h"
9
+#include "torch_utils.h"
10
11
namespace vllm {
12
@@ -210,64 +210,68 @@ packed_gelu_tanh_kernel(const packed_t& val) {
210
return; \
211
} \
212
dim3 grid(num_tokens); \
213
- int cc_major = at::cuda::getCurrentDeviceProperties()->major; \
213
+ int cc_major = get_device_prop()->major; \
214
int support_vec = \
215
(CUDA_VERSION >= 12090 && cc_major >= 10 && num_tokens > 128) \
216
? vllm::VecTraits<true>::ARCH_MAX_VEC_SIZE \
217
: vllm::VecTraits<false>::ARCH_MAX_VEC_SIZE; \
218
- int vec_size = support_vec / at::elementSize(dtype); \
218
+ int vec_size = support_vec / input.element_size(); \
219
const bool use_vec = (d % vec_size == 0); \
220
- const at::cuda::OptionalCUDAGuard device_guard(device_of(input)); \
221
- const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); \
220
+ const torch::stable::accelerator::DeviceGuard device_guard( \
221
+ input.get_device_index()); \
222
+ const cudaStream_t stream = get_current_cuda_stream(); \
223
if (use_vec) { \
224
dim3 block(std::min(d / vec_size, 1024)); \
225
if (CUDA_VERSION >= 12090 && cc_major >= 10 && num_tokens > 128) { \
225
- VLLM_DISPATCH_FLOATING_TYPES(dtype, "act_and_mul_kernel", [&] { \
226
+ VLLM_STABLE_DISPATCH_FLOATING_TYPES(dtype, "act_and_mul_kernel", [&] { \
227
vllm::act_and_mul_kernel< \
228
scalar_t, typename vllm::PackedTypeConverter<scalar_t>::Type, \
229
KERNEL<scalar_t>, \
230
PACKED_KERNEL<typename vllm::PackedTypeConverter<scalar_t>::Type>, \
231
ACT_FIRST, true, HAS_CLAMP, true><<<grid, block, 0, stream>>>( \
231
- out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d, LIMIT); \
232
+ out.mutable_data_ptr<scalar_t>(), \
233
+ input.const_data_ptr<scalar_t>(), d, LIMIT); \
234
}); \
235
} else { \
234
- VLLM_DISPATCH_FLOATING_TYPES(dtype, "act_and_mul_kernel", [&] { \
236
+ VLLM_STABLE_DISPATCH_FLOATING_TYPES(dtype, "act_and_mul_kernel", [&] { \
237
vllm::act_and_mul_kernel< \
238
scalar_t, typename vllm::PackedTypeConverter<scalar_t>::Type, \
239
KERNEL<scalar_t>, \
240
PACKED_KERNEL<typename vllm::PackedTypeConverter<scalar_t>::Type>, \
241
ACT_FIRST, true, HAS_CLAMP, false><<<grid, block, 0, stream>>>( \
240
- out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d, LIMIT); \
242
+ out.mutable_data_ptr<scalar_t>(), \
243
+ input.const_data_ptr<scalar_t>(), d, LIMIT); \
244
}); \
245
} \
246
} else { \
247
dim3 block(std::min(d, 1024)); \
245
- VLLM_DISPATCH_FLOATING_TYPES(dtype, "act_and_mul_kernel", [&] { \
248
+ VLLM_STABLE_DISPATCH_FLOATING_TYPES(dtype, "act_and_mul_kernel", [&] { \
249
vllm::act_and_mul_kernel< \
250
scalar_t, typename vllm::PackedTypeConverter<scalar_t>::Type, \
251
KERNEL<scalar_t>, \
252
PACKED_KERNEL<typename vllm::PackedTypeConverter<scalar_t>::Type>, \
253
ACT_FIRST, false, HAS_CLAMP><<<grid, block, 0, stream>>>( \
251
- out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d, LIMIT); \
254
+ out.mutable_data_ptr<scalar_t>(), input.const_data_ptr<scalar_t>(), \
255
+ d, LIMIT); \
256
}); \
257
}
258
255
-void silu_and_mul(torch::Tensor& out, // [..., d]
256
- torch::Tensor& input) // [..., 2 * d]
259
+void silu_and_mul(torch::stable::Tensor& out, // [..., d]
260
+ torch::stable::Tensor& input) // [..., 2 * d]
261
{
262
LAUNCH_ACTIVATION_GATE_KERNEL(vllm::silu_kernel, vllm::packed_silu_kernel,
263
true, false, 0.0f);
264
}
265
262
-void silu_and_mul_clamp(torch::Tensor& out, // [..., d]
263
- torch::Tensor& input, // [..., 2 * d]
266
+void silu_and_mul_clamp(torch::stable::Tensor& out, // [..., d]
267
+ torch::stable::Tensor& input, // [..., 2 * d]
268
double limit) {
269
LAUNCH_ACTIVATION_GATE_KERNEL(vllm::silu_kernel, vllm::packed_silu_kernel,
270
true, true, (float)limit);
271
}
272
269
-void mul_and_silu(torch::Tensor& out, // [..., d]
270
- torch::Tensor& input) // [..., 2 * d]
273
+void mul_and_silu(torch::stable::Tensor& out, // [..., d]
274
+ torch::stable::Tensor& input) // [..., 2 * d]
275
{
276
// The difference between mul_and_silu and silu_and_mul is that mul_and_silu
277
// applies the silu to the latter half of the input.
@@ -275,15 +279,15 @@ void mul_and_silu(torch::Tensor& out, // [..., d]
279
false, false, 0.0f);
280
}
281
278
-void gelu_and_mul(torch::Tensor& out, // [..., d]
279
- torch::Tensor& input) // [..., 2 * d]
282
+void gelu_and_mul(torch::stable::Tensor& out, // [..., d]
283
+ torch::stable::Tensor& input) // [..., 2 * d]
284
{
285
LAUNCH_ACTIVATION_GATE_KERNEL(vllm::gelu_kernel, vllm::packed_gelu_kernel,
286
true, false, 0.0f);
287
}
288
285
-void gelu_tanh_and_mul(torch::Tensor& out, // [..., d]
286
- torch::Tensor& input) // [..., 2 * d]
289
+void gelu_tanh_and_mul(torch::stable::Tensor& out, // [..., d]
290
+ torch::stable::Tensor& input) // [..., 2 * d]
291
{
292
LAUNCH_ACTIVATION_GATE_KERNEL(
293
vllm::gelu_tanh_kernel, vllm::packed_gelu_tanh_kernel, true, false, 0.0f);
@@ -434,19 +438,20 @@ __global__ void swigluoai_and_mul_kernel(
438
return; \
439
} \
440
dim3 grid(num_tokens); \
437
- int cc_major = at::cuda::getCurrentDeviceProperties()->major; \
441
+ int cc_major = get_device_prop()->major; \
442
int support_vec = \
443
(CUDA_VERSION >= 12090 && cc_major >= 10 && num_tokens > 128) \
444
? vllm::VecTraits<true>::ARCH_MAX_VEC_SIZE \
445
: vllm::VecTraits<false>::ARCH_MAX_VEC_SIZE; \
442
- int vec_size = support_vec / at::elementSize(dtype); \
446
+ int vec_size = support_vec / input.element_size(); \
447
const bool use_vec = (d % vec_size == 0); \
444
- const at::cuda::OptionalCUDAGuard device_guard(device_of(input)); \
445
- const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); \
448
+ const torch::stable::accelerator::DeviceGuard device_guard( \
449
+ input.get_device_index()); \
450
+ const cudaStream_t stream = get_current_cuda_stream(); \
451
if (use_vec) { \
452
dim3 block(std::min(d / vec_size, 1024)); \
453
if (CUDA_VERSION >= 12090 && cc_major >= 10 && num_tokens > 128) { \
449
- VLLM_DISPATCH_FLOATING_TYPES( \
454
+ VLLM_STABLE_DISPATCH_FLOATING_TYPES( \
455
dtype, "act_and_mul_kernel_with_param", [&] { \
456
vllm::act_and_mul_kernel_with_param< \
457
scalar_t, typename vllm::PackedTypeConverter<scalar_t>::Type, \
@@ -454,11 +459,11 @@ __global__ void swigluoai_and_mul_kernel(
459
PACKED_KERNEL< \
460
typename vllm::PackedTypeConverter<scalar_t>::Type>, \
461
true, true><<<grid, block, 0, stream>>>( \
457
- out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d, \
458
- PARAM); \
462
+ out.mutable_data_ptr<scalar_t>(), \
463
+ input.const_data_ptr<scalar_t>(), d, PARAM); \
464
}); \
465
} else { \
461
- VLLM_DISPATCH_FLOATING_TYPES( \
466
+ VLLM_STABLE_DISPATCH_FLOATING_TYPES( \
467
dtype, "act_and_mul_kernel_with_param", [&] { \
468
vllm::act_and_mul_kernel_with_param< \
469
scalar_t, typename vllm::PackedTypeConverter<scalar_t>::Type, \
@@ -466,45 +471,49 @@ __global__ void swigluoai_and_mul_kernel(
471
PACKED_KERNEL< \
472
typename vllm::PackedTypeConverter<scalar_t>::Type>, \
473
true, false><<<grid, block, 0, stream>>>( \
469
- out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d, \
470
- PARAM); \
474
+ out.mutable_data_ptr<scalar_t>(), \
475
+ input.const_data_ptr<scalar_t>(), d, PARAM); \
476
}); \
477
} \
478
} else { \
479
dim3 block(std::min(d, 1024)); \
475
- VLLM_DISPATCH_FLOATING_TYPES(dtype, "act_and_mul_kernel_with_param", [&] { \
476
- vllm::act_and_mul_kernel_with_param< \
477
- scalar_t, typename vllm::PackedTypeConverter<scalar_t>::Type, \
478
- KERNEL<scalar_t>, \
479
- PACKED_KERNEL<typename vllm::PackedTypeConverter<scalar_t>::Type>, \
480
- false><<<grid, block, 0, stream>>>( \
481
- out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d, PARAM); \
482
- }); \
480
+ VLLM_STABLE_DISPATCH_FLOATING_TYPES( \
481
+ dtype, "act_and_mul_kernel_with_param", [&] { \
482
+ vllm::act_and_mul_kernel_with_param< \
483
+ scalar_t, typename vllm::PackedTypeConverter<scalar_t>::Type, \
484
+ KERNEL<scalar_t>, \
485
+ PACKED_KERNEL< \
486
+ typename vllm::PackedTypeConverter<scalar_t>::Type>, \
487
+ false><<<grid, block, 0, stream>>>( \
488
+ out.mutable_data_ptr<scalar_t>(), \
489
+ input.const_data_ptr<scalar_t>(), d, PARAM); \
490
+ }); \
491
}
492
485
-#define LAUNCH_SIGLUOAI_AND_MUL(KERNEL, ALPHA, LIMIT) \
486
- int d = input.size(-1) / 2; \
487
- int64_t num_tokens = input.numel() / input.size(-1); \
488
- dim3 grid(num_tokens); \
489
- dim3 block(std::min(d, 1024)); \
490
- const at::cuda::OptionalCUDAGuard device_guard(device_of(input)); \
491
- const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); \
492
- VLLM_DISPATCH_FLOATING_TYPES( \
493
- input.scalar_type(), "clamp_swiglu_kernel_with_params", [&] { \
494
- vllm::swigluoai_and_mul_kernel<scalar_t, KERNEL<scalar_t>> \
495
- <<<grid, block, 0, stream>>>(out.data_ptr<scalar_t>(), \
496
- input.data_ptr<scalar_t>(), d, ALPHA, \
497
- LIMIT); \
493
+#define LAUNCH_SIGLUOAI_AND_MUL(KERNEL, ALPHA, LIMIT) \
494
+ int d = input.size(-1) / 2; \
495
+ int64_t num_tokens = input.numel() / input.size(-1); \
496
+ dim3 grid(num_tokens); \
497
+ dim3 block(std::min(d, 1024)); \
498
+ const torch::stable::accelerator::DeviceGuard device_guard( \
499
+ input.get_device_index()); \
500
+ const cudaStream_t stream = get_current_cuda_stream(); \
501
+ VLLM_STABLE_DISPATCH_FLOATING_TYPES( \
502
+ input.scalar_type(), "clamp_swiglu_kernel_with_params", [&] { \
503
+ vllm::swigluoai_and_mul_kernel<scalar_t, KERNEL<scalar_t>> \
504
+ <<<grid, block, 0, stream>>>(out.mutable_data_ptr<scalar_t>(), \
505
+ input.const_data_ptr<scalar_t>(), d, \
506
+ ALPHA, LIMIT); \
507
});
508
500
-void fatrelu_and_mul(torch::Tensor& out, // [..., d],
501
- torch::Tensor& input, // [..., 2 * d]
509
+void fatrelu_and_mul(torch::stable::Tensor& out, // [..., d],
510
+ torch::stable::Tensor& input, // [..., 2 * d]
511
double threshold) {
512
LAUNCH_ACTIVATION_GATE_KERNEL_WITH_PARAM(
513
vllm::fatrelu_kernel, vllm::packed_fatrelu_kernel, threshold);
514
}
506
-void swigluoai_and_mul(torch::Tensor& out, // [..., d]
507
- torch::Tensor& input, // [..., 2 * d]
515
+void swigluoai_and_mul(torch::stable::Tensor& out, // [..., d]
516
+ torch::stable::Tensor& input, // [..., 2 * d]
517
double alpha, double limit) {
518
LAUNCH_SIGLUOAI_AND_MUL(vllm::swigluoai_and_mul, alpha, limit);
519
}
@@ -559,45 +568,46 @@ __global__ void activation_kernel(
568
} // namespace vllm
569
570
// Launch element-wise activation kernel.
562
-#define LAUNCH_ACTIVATION_KERNEL(KERNEL) \
563
- auto dtype = input.scalar_type(); \
564
- int d = input.size(-1); \
565
- int64_t num_tokens = input.numel() / input.size(-1); \
566
- if (num_tokens == 0) { \
567
- return; \
568
- } \
569
- dim3 grid(num_tokens); \
570
- int cc_major = at::cuda::getCurrentDeviceProperties()->major; \
571
- int support_vec = \
572
- (CUDA_VERSION >= 12090 && cc_major >= 10 && num_tokens > 128) \
573
- ? vllm::VecTraits<true>::ARCH_MAX_VEC_SIZE \
574
- : vllm::VecTraits<false>::ARCH_MAX_VEC_SIZE; \
575
- int vec_size = support_vec / at::elementSize(dtype); \
576
- const bool use_vec = (d % vec_size == 0); \
577
- const at::cuda::OptionalCUDAGuard device_guard(device_of(input)); \
578
- const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); \
579
- if (use_vec) { \
580
- dim3 block(std::min(d / vec_size, 1024)); \
581
- if (CUDA_VERSION >= 12090 && cc_major >= 10 && num_tokens > 128) { \
582
- VLLM_DISPATCH_FLOATING_TYPES(dtype, "activation_kernel", [&] { \
583
- vllm::activation_kernel<scalar_t, KERNEL<scalar_t>, true, true> \
584
- <<<grid, block, 0, stream>>>(out.data_ptr<scalar_t>(), \
585
- input.data_ptr<scalar_t>(), d); \
586
- }); \
587
- } else { \
588
- VLLM_DISPATCH_FLOATING_TYPES(dtype, "activation_kernel", [&] { \
589
- vllm::activation_kernel<scalar_t, KERNEL<scalar_t>, true, false> \
590
- <<<grid, block, 0, stream>>>(out.data_ptr<scalar_t>(), \
591
- input.data_ptr<scalar_t>(), d); \
592
- }); \
593
- } \
594
- } else { \
595
- dim3 block(std::min(d, 1024)); \
596
- VLLM_DISPATCH_FLOATING_TYPES(dtype, "activation_kernel", [&] { \
597
- vllm::activation_kernel<scalar_t, KERNEL<scalar_t>, false> \
598
- <<<grid, block, 0, stream>>>(out.data_ptr<scalar_t>(), \
599
- input.data_ptr<scalar_t>(), d); \
600
- }); \
571
+#define LAUNCH_ACTIVATION_KERNEL(KERNEL) \
572
+ auto dtype = input.scalar_type(); \
573
+ int d = input.size(-1); \
574
+ int64_t num_tokens = input.numel() / input.size(-1); \
575
+ if (num_tokens == 0) { \
576
+ return; \
577
+ } \
578
+ dim3 grid(num_tokens); \
579
+ int cc_major = get_device_prop()->major; \
580
+ int support_vec = \
581
+ (CUDA_VERSION >= 12090 && cc_major >= 10 && num_tokens > 128) \
582
+ ? vllm::VecTraits<true>::ARCH_MAX_VEC_SIZE \
583
+ : vllm::VecTraits<false>::ARCH_MAX_VEC_SIZE; \
584
+ int vec_size = support_vec / input.element_size(); \
585
+ const bool use_vec = (d % vec_size == 0); \
586
+ const torch::stable::accelerator::DeviceGuard device_guard( \
587
+ input.get_device_index()); \
588
+ const cudaStream_t stream = get_current_cuda_stream(); \
589
+ if (use_vec) { \
590
+ dim3 block(std::min(d / vec_size, 1024)); \
591
+ if (CUDA_VERSION >= 12090 && cc_major >= 10 && num_tokens > 128) { \
592
+ VLLM_STABLE_DISPATCH_FLOATING_TYPES(dtype, "activation_kernel", [&] { \
593
+ vllm::activation_kernel<scalar_t, KERNEL<scalar_t>, true, true> \
594
+ <<<grid, block, 0, stream>>>(out.mutable_data_ptr<scalar_t>(), \
595
+ input.const_data_ptr<scalar_t>(), d); \
596
+ }); \
597
+ } else { \
598
+ VLLM_STABLE_DISPATCH_FLOATING_TYPES(dtype, "activation_kernel", [&] { \
599
+ vllm::activation_kernel<scalar_t, KERNEL<scalar_t>, true, false> \
600
+ <<<grid, block, 0, stream>>>(out.mutable_data_ptr<scalar_t>(), \
601
+ input.const_data_ptr<scalar_t>(), d); \
602
+ }); \
603
+ } \
604
+ } else { \
605
+ dim3 block(std::min(d, 1024)); \
606
+ VLLM_STABLE_DISPATCH_FLOATING_TYPES(dtype, "activation_kernel", [&] { \
607
+ vllm::activation_kernel<scalar_t, KERNEL<scalar_t>, false> \
608
+ <<<grid, block, 0, stream>>>(out.mutable_data_ptr<scalar_t>(), \
609
+ input.const_data_ptr<scalar_t>(), d); \
610
+ }); \
611
}
612
613
namespace vllm {
@@ -625,20 +635,20 @@ __device__ __forceinline__ T gelu_quick_kernel(const T& x) {
635
636
} // namespace vllm
637
628
-void gelu_new(torch::Tensor& out, // [..., d]
629
- torch::Tensor& input) // [..., d]
638
+void gelu_new(torch::stable::Tensor& out, // [..., d]
639
+ torch::stable::Tensor& input) // [..., d]
640
{
641
LAUNCH_ACTIVATION_KERNEL(vllm::gelu_new_kernel);
642
}
643
634
-void gelu_fast(torch::Tensor& out, // [..., d]
635
- torch::Tensor& input) // [..., d]
644
+void gelu_fast(torch::stable::Tensor& out, // [..., d]
645
+ torch::stable::Tensor& input) // [..., d]
646
{
647
LAUNCH_ACTIVATION_KERNEL(vllm::gelu_fast_kernel);
648
}
649
640
-void gelu_quick(torch::Tensor& out, // [..., d]
641
- torch::Tensor& input) // [..., d]
650
+void gelu_quick(torch::stable::Tensor& out, // [..., d]
651
+ torch::stable::Tensor& input) // [..., d]
652
{
653
LAUNCH_ACTIVATION_KERNEL(vllm::gelu_quick_kernel);
654
}
csrc/libtorch_stable/ops.h
+83
@@ -166,3 +166,86 @@ torch::stable::Tensor awq_dequantize(torch::stable::Tensor _kernel,
166
167
torch::stable::Tensor hadacore_transform(torch::stable::Tensor& x,
168
bool inplace);
169
+
170
+// Activation kernels (shared CUDA/ROCm)
171
+void silu_and_mul(torch::stable::Tensor& out, torch::stable::Tensor& input);
172
+void silu_and_mul_clamp(torch::stable::Tensor& out,
173
+ torch::stable::Tensor& input, double limit);
174
+void mul_and_silu(torch::stable::Tensor& out, torch::stable::Tensor& input);
175
+void gelu_and_mul(torch::stable::Tensor& out, torch::stable::Tensor& input);
176
+void gelu_tanh_and_mul(torch::stable::Tensor& out,
177
+ torch::stable::Tensor& input);
178
+void fatrelu_and_mul(torch::stable::Tensor& out, torch::stable::Tensor& input,
179
+ double threshold);
180
+void swigluoai_and_mul(torch::stable::Tensor& out, torch::stable::Tensor& input,
181
+ double alpha = 1.702, double limit = 7.0);
182
+void gelu_new(torch::stable::Tensor& out, torch::stable::Tensor& input);
183
+void gelu_fast(torch::stable::Tensor& out, torch::stable::Tensor& input);
184
+void gelu_quick(torch::stable::Tensor& out, torch::stable::Tensor& input);
185
+
186
+// INT8 quantization kernels (shared CUDA/ROCm)
187
+void static_scaled_int8_quant(torch::stable::Tensor& out,
188
+ torch::stable::Tensor const& input,
189
+ torch::stable::Tensor const& scale,
190
+ std::optional<torch::stable::Tensor> const& azp);
191
+
192
+void dynamic_scaled_int8_quant(torch::stable::Tensor& out,
193
+ torch::stable::Tensor const& input,
194
+ torch::stable::Tensor& scales,
195
+ std::optional<torch::stable::Tensor> const& azp);
196
+
197
+// FP8 quantization kernels (shared CUDA/ROCm)
198
+void static_scaled_fp8_quant(
199
+ torch::stable::Tensor& out, torch::stable::Tensor const& input,
200
+ torch::stable::Tensor const& scale,
201
+ std::optional<torch::headeronly::IntHeaderOnlyArrayRef> group_shape =
202
+ std::nullopt);
203
+
204
+void dynamic_scaled_fp8_quant(torch::stable::Tensor& out,
205
+ torch::stable::Tensor const& input,
206
+ torch::stable::Tensor& scale);
207
+
208
+void dynamic_per_token_scaled_fp8_quant(
209
+ torch::stable::Tensor& out, torch::stable::Tensor const& input,
210
+ torch::stable::Tensor& scale,
211
+ std::optional<torch::stable::Tensor> const& scale_ub);
212
+
213
+// GPTQ kernels (shared CUDA/ROCm)
214
+torch::stable::Tensor gptq_gemm(torch::stable::Tensor a,
215
+ torch::stable::Tensor b_q_weight,
216
+ torch::stable::Tensor b_gptq_qzeros,
217
+ torch::stable::Tensor b_gptq_scales,
218
+ torch::stable::Tensor b_g_idx, bool use_exllama,
219
+ bool use_v2_format, int64_t bit);
220
+
221
+void gptq_shuffle(torch::stable::Tensor q_weight, torch::stable::Tensor q_perm,
222
+ int64_t bit);
223
+
224
+// GGML kernels (shared CUDA/ROCm)
225
+torch::stable::Tensor ggml_dequantize(
226
+ torch::stable::Tensor W, int64_t type, int64_t m, int64_t n,
227
+ std::optional<torch::headeronly::ScalarType> const& dtype);
228
+
229
+torch::stable::Tensor ggml_mul_mat_vec_a8(torch::stable::Tensor W,
230
+ torch::stable::Tensor X, int64_t type,
231
+ int64_t row);
232
+
233
+torch::stable::Tensor ggml_mul_mat_a8(torch::stable::Tensor W,
234
+ torch::stable::Tensor X, int64_t type,
235
+ int64_t row);
236
+
237
+torch::stable::Tensor ggml_moe_a8(torch::stable::Tensor X,
238
+ torch::stable::Tensor W,
239
+ torch::stable::Tensor sorted_token_ids,
240
+ torch::stable::Tensor expert_ids,
241
+ torch::stable::Tensor num_tokens_post_padded,
242
+ int64_t type, int64_t row, int64_t top_k,
243
+ int64_t tokens);
244
+
245
+torch::stable::Tensor ggml_moe_a8_vec(torch::stable::Tensor X,
246
+ torch::stable::Tensor W,
247
+ torch::stable::Tensor topk_ids,
248
+ int64_t top_k, int64_t type, int64_t row,
249
+ int64_t tokens);
250
+
251
+int64_t ggml_moe_get_block_size(int64_t type);
csrc/libtorch_stable/quantization/gguf/gguf_kernel.cu
renamed
+180
-162
@@ -1,17 +1,20 @@
1
#include <cuda_fp16.h>
2
#include <cuda_runtime.h>
3
4
-#include <torch/all.h>
5
-#include <c10/cuda/CUDAGuard.h>
4
+#include "../../../cuda_compat.h"
5
+#include "../../dispatch_utils.h"
6
+#include "../../torch_utils.h"
7
7
-#include "../../cuda_compat.h"
8
-#include "dispatch_utils.h"
8
+#include <torch/csrc/stable/ops.h>
9
10
-#include "ggml-common.h"
11
-#include "vecdotq.cuh"
12
-#include "dequantize.cuh"
13
-#include "mmvq.cuh"
14
-#include "mmq.cuh"
10
+// NOTE: These headers are intentionally kept in csrc/quantization/gguf/ (not
11
+// moved to libtorch_stable) to avoid unnecessary reformatting that would break
12
+// git rename detection and pollute blame history.
13
+#include "../../../quantization/gguf/ggml-common.h"
14
+#include "../../../quantization/gguf/vecdotq.cuh"
15
+#include "../../../quantization/gguf/dequantize.cuh"
16
+#include "../../../quantization/gguf/mmvq.cuh"
17
+#include "../../../quantization/gguf/mmq.cuh"
18
#include "moe.cuh"
19
#include "moe_vec.cuh"
20
@@ -71,16 +74,17 @@ static void quantize_row_q8_1_cuda(const scalar_t* x, void* vy, const int kx,
74
}
75
}
76
74
-torch::Tensor ggml_dequantize(torch::Tensor W, // quant weight
75
- int64_t type, int64_t m, int64_t n,
76
- std::optional<at::ScalarType> const& dtype) {
77
- const at::cuda::OptionalCUDAGuard device_guard(device_of(W));
78
- auto dtype_ = dtype.value_or(torch::kFloat16);
79
- auto options = torch::TensorOptions().dtype(dtype_).device(W.device());
80
- at::Tensor DW = torch::empty({m, n}, options);
81
- cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream();
77
+torch::stable::Tensor ggml_dequantize(
78
+ torch::stable::Tensor W, // quant weight
79
+ int64_t type, int64_t m, int64_t n,
80
+ std::optional<torch::headeronly::ScalarType> const& dtype) {
81
+ const torch::stable::accelerator::DeviceGuard device_guard(
82
+ W.get_device_index());
83
+ auto dtype_ = dtype.value_or(torch::headeronly::ScalarType::Half);
84
+ auto DW = torch::stable::empty({m, n}, dtype_, std::nullopt, W.device());
85
+ cudaStream_t stream = get_current_cuda_stream();
86
83
- VLLM_DISPATCH_FLOATING_TYPES(DW.scalar_type(), "ggml_dequantize", [&] {
87
+ VLLM_STABLE_DISPATCH_FLOATING_TYPES(DW.scalar_type(), "ggml_dequantize", [&] {
88
auto to_cuda = ggml_get_to_cuda<scalar_t>(type);
89
to_cuda((void*)W.data_ptr(), (scalar_t*)DW.data_ptr(), m * n, stream);
90
});
@@ -88,135 +92,142 @@ torch::Tensor ggml_dequantize(torch::Tensor W, // quant weight
92
return DW;
93
}
94
91
-torch::Tensor ggml_mul_mat_vec_a8(torch::Tensor W, // quant weight
92
- torch::Tensor X, // input
93
- int64_t type, int64_t row) {
95
+torch::stable::Tensor ggml_mul_mat_vec_a8(
96
+ torch::stable::Tensor W, // quant weight
97
+ torch::stable::Tensor X, // input
98
+ int64_t type, int64_t row) {
99
int col = X.sizes()[1];
100
int vecs = X.sizes()[0];
101
const int padded = (col + 512 - 1) / 512 * 512;
97
- const at::cuda::OptionalCUDAGuard device_guard(device_of(X));
98
- auto options = torch::TensorOptions().dtype(X.dtype()).device(W.device());
99
- at::Tensor Y = torch::empty({vecs, row}, options);
100
- cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream();
101
- options = torch::TensorOptions().dtype(torch::kInt32).device(W.device());
102
- at::Tensor quant_X = torch::empty({vecs, padded / 32 * 9}, options);
103
- VLLM_DISPATCH_FLOATING_TYPES(X.scalar_type(), "ggml_mul_mat_vec_a8", [&] {
104
- quantize_row_q8_1_cuda<scalar_t>(
105
- (scalar_t*)X.data_ptr(), (void*)quant_X.data_ptr(), col, vecs, stream);
106
- switch (type) {
107
- case 2:
108
- mul_mat_vec_q4_0_q8_1_cuda<scalar_t>(
109
- (void*)W.data_ptr(), (void*)quant_X.data_ptr(),
110
- (scalar_t*)Y.data_ptr(), col, row, vecs, stream);
111
- break;
112
- case 3:
113
- mul_mat_vec_q4_1_q8_1_cuda<scalar_t>(
114
- (void*)W.data_ptr(), (void*)quant_X.data_ptr(),
115
- (scalar_t*)Y.data_ptr(), col, row, vecs, stream);
116
- break;
117
- case 6:
118
- mul_mat_vec_q5_0_q8_1_cuda<scalar_t>(
119
- (void*)W.data_ptr(), (void*)quant_X.data_ptr(),
120
- (scalar_t*)Y.data_ptr(), col, row, vecs, stream);
121
- break;
122
- case 7:
123
- mul_mat_vec_q5_1_q8_1_cuda<scalar_t>(
124
- (void*)W.data_ptr(), (void*)quant_X.data_ptr(),
125
- (scalar_t*)Y.data_ptr(), col, row, vecs, stream);
126
- break;
127
- case 8:
128
- mul_mat_vec_q8_0_q8_1_cuda<scalar_t>(
129
- (void*)W.data_ptr(), (void*)quant_X.data_ptr(),
130
- (scalar_t*)Y.data_ptr(), col, row, vecs, stream);
131
- break;
132
- case 10:
133
- mul_mat_vec_q2_K_q8_1_cuda<scalar_t>(
134
- (void*)W.data_ptr(), (void*)quant_X.data_ptr(),
135
- (scalar_t*)Y.data_ptr(), col, row, vecs, stream);
136
- break;
137
- case 11:
138
- mul_mat_vec_q3_K_q8_1_cuda<scalar_t>(
139
- (void*)W.data_ptr(), (void*)quant_X.data_ptr(),
140
- (scalar_t*)Y.data_ptr(), col, row, vecs, stream);
141
- break;
142
- case 12:
143
- mul_mat_vec_q4_K_q8_1_cuda<scalar_t>(
144
- (void*)W.data_ptr(), (void*)quant_X.data_ptr(),
145
- (scalar_t*)Y.data_ptr(), col, row, vecs, stream);
146
- break;
147
- case 13:
148
- mul_mat_vec_q5_K_q8_1_cuda<scalar_t>(
149
- (void*)W.data_ptr(), (void*)quant_X.data_ptr(),
150
- (scalar_t*)Y.data_ptr(), col, row, vecs, stream);
151
- break;
152
- case 14:
153
- mul_mat_vec_q6_K_q8_1_cuda<scalar_t>(
154
- (void*)W.data_ptr(), (void*)quant_X.data_ptr(),
155
- (scalar_t*)Y.data_ptr(), col, row, vecs, stream);
156
- break;
157
- case 16:
158
- mul_mat_vec_iq2_xxs_q8_1_cuda<scalar_t>(
159
- (void*)W.data_ptr(), (void*)quant_X.data_ptr(),
160
- (scalar_t*)Y.data_ptr(), col, row, vecs, stream);
161
- break;
162
- case 17:
163
- mul_mat_vec_iq2_xs_q8_1_cuda<scalar_t>(
164
- (void*)W.data_ptr(), (void*)quant_X.data_ptr(),
165
- (scalar_t*)Y.data_ptr(), col, row, vecs, stream);
166
- break;
167
- case 18:
168
- mul_mat_vec_iq3_xxs_q8_1_cuda<scalar_t>(
169
- (void*)W.data_ptr(), (void*)quant_X.data_ptr(),
170
- (scalar_t*)Y.data_ptr(), col, row, vecs, stream);
171
- break;
172
- case 19:
173
- mul_mat_vec_iq1_s_q8_1_cuda<scalar_t>(
174
- (void*)W.data_ptr(), (void*)quant_X.data_ptr(),
175
- (scalar_t*)Y.data_ptr(), col, row, vecs, stream);
176
- break;
177
- case 20:
178
- mul_mat_vec_iq4_nl_q8_1_cuda<scalar_t>(
179
- (void*)W.data_ptr(), (void*)quant_X.data_ptr(),
180
- (scalar_t*)Y.data_ptr(), col, row, vecs, stream);
181
- break;
182
- case 21:
183
- mul_mat_vec_iq3_s_q8_1_cuda<scalar_t>(
184
- (void*)W.data_ptr(), (void*)quant_X.data_ptr(),
185
- (scalar_t*)Y.data_ptr(), col, row, vecs, stream);
186
- break;
187
- case 22:
188
- mul_mat_vec_iq2_s_q8_1_cuda<scalar_t>(
189
- (void*)W.data_ptr(), (void*)quant_X.data_ptr(),
190
- (scalar_t*)Y.data_ptr(), col, row, vecs, stream);
191
- break;
192
- case 23:
193
- mul_mat_vec_iq4_xs_q8_1_cuda<scalar_t>(
194
- (void*)W.data_ptr(), (void*)quant_X.data_ptr(),
195
- (scalar_t*)Y.data_ptr(), col, row, vecs, stream);
196
- break;
197
- case 29:
198
- mul_mat_vec_iq1_m_q8_1_cuda<scalar_t>(
199
- (void*)W.data_ptr(), (void*)quant_X.data_ptr(),
200
- (scalar_t*)Y.data_ptr(), col, row, vecs, stream);
201
- break;
202
- }
203
- });
102
+ const torch::stable::accelerator::DeviceGuard device_guard(
103
+ X.get_device_index());
104
+ auto Y = torch::stable::empty({vecs, row}, X.scalar_type(), std::nullopt,
105
+ W.device());
106
+ cudaStream_t stream = get_current_cuda_stream();
107
+ auto quant_X = torch::stable::empty({vecs, padded / 32 * 9},
108
+ torch::headeronly::ScalarType::Int,
109
+ std::nullopt, W.device());
110
+ VLLM_STABLE_DISPATCH_FLOATING_TYPES(
111
+ X.scalar_type(), "ggml_mul_mat_vec_a8", [&] {
112
+ quantize_row_q8_1_cuda<scalar_t>((scalar_t*)X.data_ptr(),
113
+ (void*)quant_X.data_ptr(), col, vecs,
114
+ stream);
115
+ switch (type) {
116
+ case 2:
117
+ mul_mat_vec_q4_0_q8_1_cuda<scalar_t>(
118
+ (void*)W.data_ptr(), (void*)quant_X.data_ptr(),
119
+ (scalar_t*)Y.data_ptr(), col, row, vecs, stream);
120
+ break;
121
+ case 3:
122
+ mul_mat_vec_q4_1_q8_1_cuda<scalar_t>(
123
+ (void*)W.data_ptr(), (void*)quant_X.data_ptr(),
124
+ (scalar_t*)Y.data_ptr(), col, row, vecs, stream);
125
+ break;
126
+ case 6:
127
+ mul_mat_vec_q5_0_q8_1_cuda<scalar_t>(
128
+ (void*)W.data_ptr(), (void*)quant_X.data_ptr(),
129
+ (scalar_t*)Y.data_ptr(), col, row, vecs, stream);
130
+ break;
131
+ case 7:
132
+ mul_mat_vec_q5_1_q8_1_cuda<scalar_t>(
133
+ (void*)W.data_ptr(), (void*)quant_X.data_ptr(),
134
+ (scalar_t*)Y.data_ptr(), col, row, vecs, stream);
135
+ break;
136
+ case 8:
137
+ mul_mat_vec_q8_0_q8_1_cuda<scalar_t>(
138
+ (void*)W.data_ptr(), (void*)quant_X.data_ptr(),
139
+ (scalar_t*)Y.data_ptr(), col, row, vecs, stream);
140
+ break;
141
+ case 10:
142
+ mul_mat_vec_q2_K_q8_1_cuda<scalar_t>(
143
+ (void*)W.data_ptr(), (void*)quant_X.data_ptr(),
144
+ (scalar_t*)Y.data_ptr(), col, row, vecs, stream);
145
+ break;
146
+ case 11:
147
+ mul_mat_vec_q3_K_q8_1_cuda<scalar_t>(
148
+ (void*)W.data_ptr(), (void*)quant_X.data_ptr(),
149
+ (scalar_t*)Y.data_ptr(), col, row, vecs, stream);
150
+ break;
151
+ case 12:
152
+ mul_mat_vec_q4_K_q8_1_cuda<scalar_t>(
153
+ (void*)W.data_ptr(), (void*)quant_X.data_ptr(),
154
+ (scalar_t*)Y.data_ptr(), col, row, vecs, stream);
155
+ break;
156
+ case 13:
157
+ mul_mat_vec_q5_K_q8_1_cuda<scalar_t>(
158
+ (void*)W.data_ptr(), (void*)quant_X.data_ptr(),
159
+ (scalar_t*)Y.data_ptr(), col, row, vecs, stream);
160
+ break;
161
+ case 14:
162
+ mul_mat_vec_q6_K_q8_1_cuda<scalar_t>(
163
+ (void*)W.data_ptr(), (void*)quant_X.data_ptr(),
164
+ (scalar_t*)Y.data_ptr(), col, row, vecs, stream);
165
+ break;
166
+ case 16:
167
+ mul_mat_vec_iq2_xxs_q8_1_cuda<scalar_t>(
168
+ (void*)W.data_ptr(), (void*)quant_X.data_ptr(),
169
+ (scalar_t*)Y.data_ptr(), col, row, vecs, stream);
170
+ break;
171
+ case 17:
172
+ mul_mat_vec_iq2_xs_q8_1_cuda<scalar_t>(
173
+ (void*)W.data_ptr(), (void*)quant_X.data_ptr(),
174
+ (scalar_t*)Y.data_ptr(), col, row, vecs, stream);
175
+ break;
176
+ case 18:
177
+ mul_mat_vec_iq3_xxs_q8_1_cuda<scalar_t>(
178
+ (void*)W.data_ptr(), (void*)quant_X.data_ptr(),
179
+ (scalar_t*)Y.data_ptr(), col, row, vecs, stream);
180
+ break;
181
+ case 19:
182
+ mul_mat_vec_iq1_s_q8_1_cuda<scalar_t>(
183
+ (void*)W.data_ptr(), (void*)quant_X.data_ptr(),
184
+ (scalar_t*)Y.data_ptr(), col, row, vecs, stream);
185
+ break;
186
+ case 20:
187
+ mul_mat_vec_iq4_nl_q8_1_cuda<scalar_t>(
188
+ (void*)W.data_ptr(), (void*)quant_X.data_ptr(),
189
+ (scalar_t*)Y.data_ptr(), col, row, vecs, stream);
190
+ break;
191
+ case 21:
192
+ mul_mat_vec_iq3_s_q8_1_cuda<scalar_t>(
193
+ (void*)W.data_ptr(), (void*)quant_X.data_ptr(),
194
+ (scalar_t*)Y.data_ptr(), col, row, vecs, stream);
195
+ break;
196
+ case 22:
197
+ mul_mat_vec_iq2_s_q8_1_cuda<scalar_t>(
198
+ (void*)W.data_ptr(), (void*)quant_X.data_ptr(),
199
+ (scalar_t*)Y.data_ptr(), col, row, vecs, stream);
200
+ break;
201
+ case 23:
202
+ mul_mat_vec_iq4_xs_q8_1_cuda<scalar_t>(
203
+ (void*)W.data_ptr(), (void*)quant_X.data_ptr(),
204
+ (scalar_t*)Y.data_ptr(), col, row, vecs, stream);
205
+ break;
206
+ case 29:
207
+ mul_mat_vec_iq1_m_q8_1_cuda<scalar_t>(
208
+ (void*)W.data_ptr(), (void*)quant_X.data_ptr(),
209
+ (scalar_t*)Y.data_ptr(), col, row, vecs, stream);
210
+ break;
211
+ }
212
+ });
213
return Y;
214
}
215
207
-torch::Tensor ggml_mul_mat_a8(torch::Tensor W, // quant weight
208
- torch::Tensor X, // input
209
- int64_t type, int64_t row) {
216
+torch::stable::Tensor ggml_mul_mat_a8(torch::stable::Tensor W, // quant weight
217
+ torch::stable::Tensor X, // input
218
+ int64_t type, int64_t row) {
219
int col = X.sizes()[1];
220
int padded = (col + 512 - 1) / 512 * 512;
221
int batch = X.sizes()[0];
213
- const at::cuda::OptionalCUDAGuard device_guard(device_of(X));
214
- auto options = torch::TensorOptions().dtype(X.dtype()).device(W.device());
215
- at::Tensor Y = torch::empty({batch, row}, options);
216
- cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream();
217
- options = torch::TensorOptions().dtype(torch::kInt32).device(W.device());
218
- at::Tensor quant_X = torch::empty({batch, padded / 32 * 9}, options);
219
- VLLM_DISPATCH_FLOATING_TYPES(X.scalar_type(), "ggml_mul_mat_a8", [&] {
222
+ const torch::stable::accelerator::DeviceGuard device_guard(
223
+ X.get_device_index());
224
+ auto Y = torch::stable::empty({batch, row}, X.scalar_type(), std::nullopt,
225
+ W.device());
226
+ cudaStream_t stream = get_current_cuda_stream();
227
+ auto quant_X = torch::stable::empty({batch, padded / 32 * 9},
228
+ torch::headeronly::ScalarType::Int,
229
+ std::nullopt, W.device());
230
+ VLLM_STABLE_DISPATCH_FLOATING_TYPES(X.scalar_type(), "ggml_mul_mat_a8", [&] {
231
quantize_row_q8_1_cuda((scalar_t*)X.data_ptr(), (void*)quant_X.data_ptr(),
232
col, batch, stream);
233
@@ -276,21 +287,24 @@ torch::Tensor ggml_mul_mat_a8(torch::Tensor W, // quant weight
287
return Y;
288
}
289
279
-torch::Tensor ggml_moe_a8(torch::Tensor X, // input
280
- torch::Tensor W, // expert weights
281
- torch::Tensor sorted_token_ids,
282
- torch::Tensor expert_ids,
283
- torch::Tensor num_tokens_post_padded, int64_t type,
284
- int64_t row, int64_t top_k, int64_t tokens) {
290
+torch::stable::Tensor ggml_moe_a8(torch::stable::Tensor X, // input
291
+ torch::stable::Tensor W, // expert weights
292
+ torch::stable::Tensor sorted_token_ids,
293
+ torch::stable::Tensor expert_ids,
294
+ torch::stable::Tensor num_tokens_post_padded,
295
+ int64_t type, int64_t row, int64_t top_k,
296
+ int64_t tokens) {
297
int col = X.sizes()[1];
298
int padded = (col + 512 - 1) / 512 * 512;
287
- const at::cuda::OptionalCUDAGuard device_guard(device_of(X));
288
- auto options = torch::TensorOptions().dtype(X.dtype()).device(W.device());
289
- at::Tensor Y = torch::empty({tokens * top_k, row}, options);
290
- cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream();
291
- options = torch::TensorOptions().dtype(torch::kInt32).device(W.device());
292
- at::Tensor quant_X = torch::empty({tokens, padded / 32 * 9}, options);
293
- VLLM_DISPATCH_FLOATING_TYPES(X.scalar_type(), "ggml_moe_a8", [&] {
299
+ const torch::stable::accelerator::DeviceGuard device_guard(
300
+ X.get_device_index());
301
+ auto Y = torch::stable::empty({tokens * top_k, row}, X.scalar_type(),
302
+ std::nullopt, W.device());
303
+ cudaStream_t stream = get_current_cuda_stream();
304
+ auto quant_X = torch::stable::empty({tokens, padded / 32 * 9},
305
+ torch::headeronly::ScalarType::Int,
306
+ std::nullopt, W.device());
307
+ VLLM_STABLE_DISPATCH_FLOATING_TYPES(X.scalar_type(), "ggml_moe_a8", [&] {
308
quantize_row_q8_1_cuda((scalar_t*)X.data_ptr(), (void*)quant_X.data_ptr(),
309
col, tokens, stream);
310
switch (type) {
@@ -379,19 +393,23 @@ torch::Tensor ggml_moe_a8(torch::Tensor X, // input
393
return Y;
394
}
395
382
-torch::Tensor ggml_moe_a8_vec(torch::Tensor X, // input
383
- torch::Tensor W, // expert weights
384
- torch::Tensor topk_ids, int64_t top_k,
385
- int64_t type, int64_t row, int64_t tokens) {
396
+torch::stable::Tensor ggml_moe_a8_vec(
397
+ torch::stable::Tensor X, // input
398
+ torch::stable::Tensor W, // expert weights
399
+ torch::stable::Tensor topk_ids, int64_t top_k, int64_t type, int64_t row,
400
+ int64_t tokens) {
401
int col = X.sizes()[1];
402
const int padded = (col + 512 - 1) / 512 * 512;
388
- const at::cuda::OptionalCUDAGuard device_guard(device_of(X));
389
- auto options = torch::TensorOptions().dtype(X.dtype()).device(W.device());
390
- at::Tensor Y = torch::zeros({tokens * top_k, row}, options);
391
- cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream();
392
- options = torch::TensorOptions().dtype(torch::kInt32).device(W.device());
393
- at::Tensor quant_X = torch::empty({tokens, padded / 32 * 9}, options);
394
- VLLM_DISPATCH_FLOATING_TYPES(X.scalar_type(), "ggml_moe_vec_a8", [&] {
403
+ const torch::stable::accelerator::DeviceGuard device_guard(
404
+ X.get_device_index());
405
+ auto Y = torch::stable::empty({tokens * top_k, row}, X.scalar_type(),
406
+ std::nullopt, W.device());
407
+ torch::stable::fill_(Y, 0.0);
408
+ cudaStream_t stream = get_current_cuda_stream();
409
+ auto quant_X = torch::stable::empty({tokens, padded / 32 * 9},
410
+ torch::headeronly::ScalarType::Int,
411
+ std::nullopt, W.device());
412
+ VLLM_STABLE_DISPATCH_FLOATING_TYPES(X.scalar_type(), "ggml_moe_vec_a8", [&] {
413
quantize_row_q8_1_cuda<scalar_t>((scalar_t*)X.data_ptr(),
414
(void*)quant_X.data_ptr(), col, tokens,
415
stream);
csrc/libtorch_stable/quantization/gguf/moe.cuh
renamed
csrc/libtorch_stable/quantization/gguf/moe_vec.cuh
renamed
csrc/libtorch_stable/quantization/gptq/compat.cuh
renamed
csrc/libtorch_stable/quantization/gptq/matrix_view.cuh
renamed
csrc/libtorch_stable/quantization/gptq/q_gemm.cu
renamed
+30
-23
@@ -6,9 +6,8 @@ https://github.com/qwopqwop200/GPTQ-for-LLaMa
6
#include <cstdint>
7
#include <cstdio>
8
9
-#include <torch/all.h>
10
-#include <c10/cuda/CUDAGuard.h>
11
-#include <ATen/cuda/CUDAContext.h>
9
+#include "../../torch_utils.h"
10
+#include <torch/csrc/stable/ops.h>
11
#include <cuda_runtime.h>
12
#include <cuda_fp16.h>
13
@@ -735,7 +734,7 @@ void gemm_half_q_half_cuda_part(const half* a, const uint32_t* b_q_weight,
734
fp_gemm_half_q_half_gptq_kernel kernel =
735
pick_gemm_half_q_half_gptq_kernel(true, m_count, bit);
736
738
- const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
737
+ const cudaStream_t stream = get_current_cuda_stream();
738
kernel<<<gridDim, blockDim, 0, stream>>>(
739
a, b_q_weight, b_gptq_qzeros, b_gptq_scales, c, size_m, size_n, size_k,
740
groups, use_v2_format, b_q_perm);
@@ -1164,7 +1163,7 @@ void reconstruct_exllama(const uint32_t* b_q_weight,
1163
reconstruct_exllama_kernel = reconstruct_exllama_8bit_kernel;
1164
}
1165
1167
- const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
1166
+ const cudaStream_t stream = get_current_cuda_stream();
1167
reconstruct_exllama_kernel<<<gridDim, blockDim, 0, stream>>>(
1168
b_q_weight, b_q_perm, b_gptq_qzeros, b_gptq_scales, height, width, groups,
1169
use_v2_format, out);
@@ -1376,7 +1375,7 @@ void gemm_half_q_half_alt(const half* a, const uint32_t* b_q_weight,
1375
kernel = gemm_half_q_half_alt_8bit_kernel;
1376
}
1377
1379
- const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
1378
+ const cudaStream_t stream = get_current_cuda_stream();
1379
kernel<<<gridDim, blockDim, 0, stream>>>(
1380
(const half2*)a, b_q_weight, c, b_gptq_scales, b_gptq_qzeros, b_g_idx,
1381
size_m, size_k / 32 * bit, size_n, use_v2_format);
@@ -1485,7 +1484,7 @@ void reconstruct_gptq(const uint32_t* b_q_weight, const uint32_t* b_gptq_qzeros,
1484
gridDim.y = DIVIDE(height, 32);
1485
}
1486
1488
- const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
1487
+ const cudaStream_t stream = get_current_cuda_stream();
1488
kernel<<<gridDim, blockDim, 0, stream>>>(b_q_weight, b_gptq_scales,
1489
b_gptq_qzeros, b_g_idx, height,
1490
width, groups, use_v2_format, out);
@@ -1794,7 +1793,7 @@ void shuffle_exllama_weight(uint32_t* q_weight, int* q_perm, int height,
1793
} else if (bit == 8) {
1794
kernel = make_sequential_8bit_kernel;
1795
}
1797
- const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
1796
+ const cudaStream_t stream = get_current_cuda_stream();
1797
kernel<<<gridDim, blockDim, 0, stream>>>(q_weight, new_qweight, q_perm,
1798
width);
1799
// Replace qweights
@@ -1818,29 +1817,34 @@ void shuffle_exllama_weight(uint32_t* q_weight, int* q_perm, int height,
1817
} else if (bit == 8) {
1818
shuffle_kernel = shuffle_8bit_kernel;
1819
}
1821
- const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
1820
+ const cudaStream_t stream = get_current_cuda_stream();
1821
shuffle_kernel<<<gridDim, blockDim, 0, stream>>>(q_weight, height, width);
1822
}
1823
1824
} // namespace gptq
1825
} // namespace vllm
1826
1828
-torch::Tensor gptq_gemm(torch::Tensor a, torch::Tensor b_q_weight,
1829
- torch::Tensor b_gptq_qzeros,
1830
- torch::Tensor b_gptq_scales, torch::Tensor b_g_idx,
1831
- bool use_exllama, bool use_v2_format, int64_t bit) {
1832
- const at::cuda::OptionalCUDAGuard device_guard(device_of(a));
1833
- auto options = torch::TensorOptions().dtype(a.dtype()).device(a.device());
1834
- at::Tensor c = torch::zeros({a.size(0), b_q_weight.size(1)}, options);
1835
- at::Tensor temp_dq = torch::empty(
1836
- {b_q_weight.size(0) * 32 / bit, b_q_weight.size(1)}, options);
1827
+torch::stable::Tensor gptq_gemm(torch::stable::Tensor a,
1828
+ torch::stable::Tensor b_q_weight,
1829
+ torch::stable::Tensor b_gptq_qzeros,
1830
+ torch::stable::Tensor b_gptq_scales,
1831
+ torch::stable::Tensor b_g_idx, bool use_exllama,
1832
+ bool use_v2_format, int64_t bit) {
1833
+ const torch::stable::accelerator::DeviceGuard device_guard(
1834
+ a.get_device_index());
1835
+ auto c = torch::stable::new_zeros(a, {a.size(0), b_q_weight.size(1)});
1836
+ auto temp_dq =
1837
+ torch::stable::empty({b_q_weight.size(0) * 32 / bit, b_q_weight.size(1)},
1838
+ a.scalar_type(), std::nullopt, a.device());
1839
1840
vllm::gptq::gemm_half_q_half_cuda(
1839
- at::cuda::getCurrentCUDABlasHandle(), (const half*)a.data_ptr(),
1841
+ get_current_cuda_blas_handle(), (const half*)a.data_ptr(),
1842
(const uint32_t*)b_q_weight.data_ptr(),
1843
(const uint32_t*)b_gptq_qzeros.data_ptr(),
1844
(const half*)b_gptq_scales.data_ptr(),
1843
- b_g_idx.device().is_meta() ? NULL : (const int*)b_g_idx.data_ptr(),
1845
+ b_g_idx.device().type() == torch::stable::DeviceType::Meta
1846
+ ? NULL
1847
+ : (const int*)b_g_idx.data_ptr(),
1848
(half*)c.data_ptr(), (half*)temp_dq.data_ptr(),
1849
c.size(0), // m
1850
c.size(1), // n
@@ -1850,11 +1854,14 @@ torch::Tensor gptq_gemm(torch::Tensor a, torch::Tensor b_q_weight,
1854
return c;
1855
}
1856
1853
-void gptq_shuffle(torch::Tensor q_weight, torch::Tensor q_perm, int64_t bit) {
1854
- const at::cuda::OptionalCUDAGuard device_guard(device_of(q_weight));
1857
+void gptq_shuffle(torch::stable::Tensor q_weight, torch::stable::Tensor q_perm,
1858
+ int64_t bit) {
1859
+ const torch::stable::accelerator::DeviceGuard device_guard(
1860
+ q_weight.get_device_index());
1861
vllm::gptq::shuffle_exllama_weight(
1862
(uint32_t*)q_weight.data_ptr(),
1857
- q_perm.device().is_meta() || q_perm.numel() == 0
1863
+ q_perm.device().type() == torch::stable::DeviceType::Meta ||
1864
+ q_perm.numel() == 0
1865
? NULL
1866
: (int*)q_perm.data_ptr(),
1867
q_weight.size(0) * 32 / bit, q_weight.size(1), bit);
csrc/libtorch_stable/quantization/gptq/qdq_2.cuh
renamed
csrc/libtorch_stable/quantization/gptq/qdq_3.cuh
renamed
csrc/libtorch_stable/quantization/gptq/qdq_4.cuh
renamed
csrc/libtorch_stable/quantization/gptq/qdq_8.cuh
renamed
csrc/libtorch_stable/quantization/gptq/qdq_util.cuh
renamed
csrc/libtorch_stable/quantization/w8a8/fp8/common.cu
renamed
+108
-89
@@ -1,11 +1,9 @@
1
-#include "common.cuh"
2
-#include "dispatch_utils.h"
3
-#include "cub_helpers.h"
4
-#include "libtorch_stable/quantization/vectorization_utils.cuh"
5
-#include <c10/cuda/CUDAGuard.h>
6
-#include <ATen/cuda/Exceptions.h>
7
-#include <tuple>
8
-
1
+#include "../../../../quantization/w8a8/fp8/common.cuh"
2
+#include "../../../dispatch_utils.h"
3
+#include "../../../../cub_helpers.h"
4
+#include "../../vectorization_utils.cuh"
5
+#include "../../../torch_utils.h"
6
+#include <torch/csrc/stable/macros.h>
7
namespace vllm {
8
9
// STRIDE_I_ZERO: true if scale_stride_i == 0 (per-tensor or per-channel)
@@ -183,16 +181,16 @@ __global__ void dynamic_per_token_scaled_fp8_quant_kernel_strided(
181
} // namespace vllm
182
183
void static_scaled_fp8_quant(
186
- torch::Tensor& out, // [..., d]
187
- torch::Tensor const& input, // [..., d]
188
- torch::Tensor const& scale, // various shapes
189
- std::optional<std::tuple<int64_t, int64_t>>
190
- opt_group_shape) // optional explicit (group_m, group_n)
184
+ torch::stable::Tensor& out, // [..., d]
185
+ torch::stable::Tensor const& input, // [..., d]
186
+ torch::stable::Tensor const& scale, // various shapes
187
+ std::optional<torch::headeronly::IntHeaderOnlyArrayRef>
188
+ opt_group_shape) // optional explicit [group_m, group_n]
189
{
192
- TORCH_CHECK(input.stride(-1) == 1,
193
- "last dimension of input must be contiguous");
194
- TORCH_CHECK(out.stride(-1) == 1,
195
- "last dimension of output must be contiguous");
190
+ STD_TORCH_CHECK(input.stride(-1) == 1,
191
+ "last dimension of input must be contiguous");
192
+ STD_TORCH_CHECK(out.stride(-1) == 1,
193
+ "last dimension of output must be contiguous");
194
195
const int hidden_size = input.size(-1); // N (columns)
196
const int num_tokens = input.numel() / hidden_size; // M (rows)
@@ -212,13 +210,18 @@ void static_scaled_fp8_quant(
210
} else if (scale.dim() == 1) {
211
// 1D scale: require explicit group_shape to disambiguate per-channel vs
212
// per-token (avoids edge case where num_tokens == hidden_size)
215
- TORCH_CHECK(opt_group_shape.has_value(),
216
- "1D scale requires explicit group_shape to disambiguate "
217
- "per-channel vs per-token quantization. "
218
- "Use group_shape=(-1, 1) for per-channel or group_shape=(1, "
219
- "-1) for per-token.");
220
-
221
- const auto& [opt_group_m, opt_group_n] = opt_group_shape.value();
213
+ STD_TORCH_CHECK(
214
+ opt_group_shape.has_value(),
215
+ "1D scale requires explicit group_shape to disambiguate "
216
+ "per-channel vs per-token quantization. "
217
+ "Use group_shape=(-1, 1) for per-channel or group_shape=(1, "
218
+ "-1) for per-token.");
219
+ STD_TORCH_CHECK(opt_group_shape->size() == 2,
220
+ "group_shape must have exactly 2 elements, got ",
221
+ opt_group_shape->size());
222
+
223
+ const auto opt_group_m = (*opt_group_shape)[0];
224
+ const auto opt_group_n = (*opt_group_shape)[1];
225
group_m = opt_group_m == -1 ? num_tokens : static_cast<int>(opt_group_m);
226
group_n = opt_group_n == -1 ? hidden_size : static_cast<int>(opt_group_n);
227
@@ -228,11 +231,11 @@ void static_scaled_fp8_quant(
231
const int64_t expected_scale_n = hidden_size / group_n;
232
const int64_t expected_scale_numel = expected_scale_m * expected_scale_n;
233
231
- TORCH_CHECK(scale_len == expected_scale_numel, "1D scale length (",
232
- scale_len, ") does not match expected size (",
233
- expected_scale_numel, ") for group_shape (", opt_group_m, ", ",
234
- opt_group_n, ") with input shape (", num_tokens, ", ",
235
- hidden_size, ")");
234
+ STD_TORCH_CHECK(scale_len == expected_scale_numel, "1D scale length (",
235
+ scale_len, ") does not match expected size (",
236
+ expected_scale_numel, ") for group_shape (", opt_group_m,
237
+ ", ", opt_group_n, ") with input shape (", num_tokens, ", ",
238
+ hidden_size, ")");
239
240
// For 1D scale, determine strides based on which dim is trivial
241
// Scale indexing: scale[gi * scale_stride_i + gj * scale_stride_j]
@@ -248,7 +251,7 @@ void static_scaled_fp8_quant(
251
scale_stride_i = scale.stride(0);
252
scale_stride_j = 0;
253
} else {
251
- TORCH_CHECK(
254
+ STD_TORCH_CHECK(
255
false,
256
"1D scale can only be used when one of the scale dimensions is 1. "
257
"For 2D group scaling, use a 2D scale tensor.");
@@ -259,10 +262,12 @@ void static_scaled_fp8_quant(
262
const int64_t scale_size_0 = scale.size(0);
263
const int64_t scale_size_1 = scale.size(1);
264
262
- TORCH_CHECK(num_tokens % scale_size_0 == 0, "num_tokens (", num_tokens,
263
- ") must be divisible by scale.size(0) (", scale_size_0, ")");
264
- TORCH_CHECK(hidden_size % scale_size_1 == 0, "hidden_size (", hidden_size,
265
- ") must be divisible by scale.size(1) (", scale_size_1, ")");
265
+ STD_TORCH_CHECK(num_tokens % scale_size_0 == 0, "num_tokens (", num_tokens,
266
+ ") must be divisible by scale.size(0) (", scale_size_0,
267
+ ")");
268
+ STD_TORCH_CHECK(hidden_size % scale_size_1 == 0, "hidden_size (",
269
+ hidden_size, ") must be divisible by scale.size(1) (",
270
+ scale_size_1, ")");
271
272
// Infer from 2D scale shape
273
int inferred_group_m = num_tokens / scale_size_0;
@@ -270,16 +275,21 @@ void static_scaled_fp8_quant(
275
276
// Use explicit if provided, otherwise use inferred
277
if (opt_group_shape.has_value()) {
273
- const auto& [opt_group_m, opt_group_n] = opt_group_shape.value();
278
+ STD_TORCH_CHECK(opt_group_shape->size() == 2,
279
+ "group_shape must have exactly 2 elements, got ",
280
+ opt_group_shape->size());
281
+ const auto opt_group_m = (*opt_group_shape)[0];
282
+ const auto opt_group_n = (*opt_group_shape)[1];
283
group_m = opt_group_m == -1 ? num_tokens : static_cast<int>(opt_group_m);
284
group_n = opt_group_n == -1 ? hidden_size : static_cast<int>(opt_group_n);
285
286
// Validate explicit matches inferred
278
- TORCH_CHECK(group_m == inferred_group_m && group_n == inferred_group_n,
279
- "Explicit group_shape (", opt_group_m, ", ", opt_group_n,
280
- ") does not match inferred group shape (", inferred_group_m,
281
- ", ", inferred_group_n, ") from 2D scale tensor shape (",
282
- scale_size_0, ", ", scale_size_1, ")");
287
+ STD_TORCH_CHECK(
288
+ group_m == inferred_group_m && group_n == inferred_group_n,
289
+ "Explicit group_shape (", opt_group_m, ", ", opt_group_n,
290
+ ") does not match inferred group shape (", inferred_group_m, ", ",
291
+ inferred_group_n, ") from 2D scale tensor shape (", scale_size_0,
292
+ ", ", scale_size_1, ")");
293
} else {
294
group_m = inferred_group_m;
295
group_n = inferred_group_n;
@@ -288,8 +298,8 @@ void static_scaled_fp8_quant(
298
scale_stride_i = scale.stride(0);
299
scale_stride_j = scale.stride(1);
300
} else {
291
- TORCH_CHECK(false, "scale must be 0D, 1D, or 2D tensor, but got ",
292
- scale.dim(), "D");
301
+ STD_TORCH_CHECK(false, "scale must be 0D, 1D, or 2D tensor, but got ",
302
+ scale.dim(), "D");
303
}
304
305
const int block_size = 256;
@@ -299,37 +309,39 @@ void static_scaled_fp8_quant(
309
const int64_t in_row_stride = input.stride(-2);
310
const int64_t out_row_stride = out.stride(-2);
311
302
- const at::cuda::OptionalCUDAGuard device_guard(device_of(input));
303
- const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
312
+ const torch::stable::accelerator::DeviceGuard device_guard(
313
+ input.get_device_index());
314
+ const cudaStream_t stream = get_current_cuda_stream();
315
316
// Dispatch to template-specialized kernel based on stride pattern
306
- VLLM_DISPATCH_FLOATING_TYPES(
317
+ VLLM_STABLE_DISPATCH_FLOATING_TYPES(
318
input.scalar_type(), "scaled_fp8_quant_kernel_scalar_type", [&] {
308
- VLLM_DISPATCH_FP8_TYPES(
319
+ VLLM_STABLE_DISPATCH_FP8_TYPES(
320
out.scalar_type(), "scaled_fp8_quant_kernel_fp8_type", [&] {
310
- VLLM_DISPATCH_BOOL(scale_stride_i == 0, S0_ZERO, [&] {
311
- VLLM_DISPATCH_BOOL(scale_stride_j == 0, S1_ZERO, [&] {
321
+ VLLM_STABLE_DISPATCH_BOOL(scale_stride_i == 0, S0_ZERO, [&] {
322
+ VLLM_STABLE_DISPATCH_BOOL(scale_stride_j == 0, S1_ZERO, [&] {
323
vllm::scaled_fp8_quant_kernel_strided_group_shape<
324
scalar_t, fp8_t, S0_ZERO, S1_ZERO>
325
<<<grid, block, 0, stream>>>(
315
- out.data_ptr<fp8_t>(), input.data_ptr<scalar_t>(),
316
- scale.data_ptr<float>(), hidden_size, in_row_stride,
317
- out_row_stride, group_m, group_n, scale_stride_i,
318
- scale_stride_j);
326
+ out.mutable_data_ptr<fp8_t>(),
327
+ input.const_data_ptr<scalar_t>(),
328
+ scale.const_data_ptr<float>(), hidden_size,
329
+ in_row_stride, out_row_stride, group_m, group_n,
330
+ scale_stride_i, scale_stride_j);
331
});
332
});
333
});
334
});
335
}
336
325
-void dynamic_scaled_fp8_quant(torch::Tensor& out, // [..., d]
326
- torch::Tensor const& input, // [..., d]
327
- torch::Tensor& scale) // [1]
337
+void dynamic_scaled_fp8_quant(torch::stable::Tensor& out, // [..., d]
338
+ torch::stable::Tensor const& input, // [..., d]
339
+ torch::stable::Tensor& scale) // [1]
340
{
329
- TORCH_CHECK(input.stride(-1) == 1,
330
- "last dimension of input must be contiguous");
331
- TORCH_CHECK(out.stride(-1) == 1,
332
- "last dimension of output must be contiguous");
341
+ STD_TORCH_CHECK(input.stride(-1) == 1,
342
+ "last dimension of input must be contiguous");
343
+ STD_TORCH_CHECK(out.stride(-1) == 1,
344
+ "last dimension of output must be contiguous");
345
346
const int hidden_size = input.size(-1);
347
const int num_tokens = input.numel() / hidden_size;
@@ -340,40 +352,43 @@ void dynamic_scaled_fp8_quant(torch::Tensor& out, // [..., d]
352
const int64_t in_row_stride = input.stride(-2);
353
const int64_t out_row_stride = out.stride(-2);
354
343
- const at::cuda::OptionalCUDAGuard device_guard(device_of(input));
344
- const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
355
+ const torch::stable::accelerator::DeviceGuard device_guard(
356
+ input.get_device_index());
357
+ const cudaStream_t stream = get_current_cuda_stream();
358
359
// scale tensor should be initialised to <=0 before reduction
347
- AT_CUDA_CHECK(
348
- cudaMemsetAsync(scale.data_ptr<float>(), 0, sizeof(float), stream));
360
+ STD_CUDA_CHECK(cudaMemsetAsync(scale.mutable_data_ptr<float>(), 0,
361
+ sizeof(float), stream));
362
350
- VLLM_DISPATCH_FLOATING_TYPES(
363
+ VLLM_STABLE_DISPATCH_FLOATING_TYPES(
364
input.scalar_type(), "scaled_fp8_quant_kernel_scalar_type", [&] {
352
- VLLM_DISPATCH_FP8_TYPES(
365
+ VLLM_STABLE_DISPATCH_FP8_TYPES(
366
out.scalar_type(), "scaled_fp8_quant_kernel_fp8_type", [&] {
367
vllm::segmented_max_reduction_strided<scalar_t, fp8_t>
368
<<<grid, block, 0, stream>>>(
356
- scale.data_ptr<float>(), input.data_ptr<scalar_t>(),
357
- hidden_size, in_row_stride,
358
- static_cast<int64_t>(num_tokens));
369
+ scale.mutable_data_ptr<float>(),
370
+ input.const_data_ptr<scalar_t>(), hidden_size,
371
+ in_row_stride, static_cast<int64_t>(num_tokens));
372
373
vllm::scaled_fp8_quant_kernel_strided_dynamic<scalar_t, fp8_t>
361
- <<<grid, block, 0, stream>>>(
362
- out.data_ptr<fp8_t>(), input.data_ptr<scalar_t>(),
363
- scale.data_ptr<float>(), hidden_size, in_row_stride,
364
- out_row_stride);
374
+ <<<grid, block, 0, stream>>>(out.mutable_data_ptr<fp8_t>(),
375
+ input.const_data_ptr<scalar_t>(),
376
+ scale.const_data_ptr<float>(),
377
+ hidden_size, in_row_stride,
378
+ out_row_stride);
379
});
380
});
381
}
382
383
void dynamic_per_token_scaled_fp8_quant(
370
- torch::Tensor& out, // [..., d]
371
- torch::Tensor const& input, // [..., d]
372
- torch::Tensor& scales, std::optional<at::Tensor> const& scale_ub) {
373
- TORCH_CHECK(input.stride(-1) == 1,
374
- "last dimension of input must be contiguous");
375
- TORCH_CHECK(out.stride(-1) == 1,
376
- "last dimension of output must be contiguous");
384
+ torch::stable::Tensor& out, // [..., d]
385
+ torch::stable::Tensor const& input, // [..., d]
386
+ torch::stable::Tensor& scales,
387
+ std::optional<torch::stable::Tensor> const& scale_ub) {
388
+ STD_TORCH_CHECK(input.stride(-1) == 1,
389
+ "last dimension of input must be contiguous");
390
+ STD_TORCH_CHECK(out.stride(-1) == 1,
391
+ "last dimension of output must be contiguous");
392
393
const int hidden_size = input.size(-1);
394
const int num_tokens = input.numel() / hidden_size;
@@ -384,20 +399,24 @@ void dynamic_per_token_scaled_fp8_quant(
399
const int64_t in_row_stride = input.stride(-2);
400
const int64_t out_row_stride = out.stride(-2);
401
387
- const at::cuda::OptionalCUDAGuard device_guard(device_of(input));
388
- const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
389
- VLLM_DISPATCH_FLOATING_TYPES(
402
+ const torch::stable::accelerator::DeviceGuard device_guard(
403
+ input.get_device_index());
404
+ const cudaStream_t stream = get_current_cuda_stream();
405
+ VLLM_STABLE_DISPATCH_FLOATING_TYPES(
406
input.scalar_type(),
407
"dynamic_per_token_scaled_fp8_quant_kernel_scalar_type", [&] {
392
- VLLM_DISPATCH_FP8_TYPES(
408
+ VLLM_STABLE_DISPATCH_FP8_TYPES(
409
out.scalar_type(),
410
"dynamic_per_token_scaled_fp8_quant_kernel_fp8_type", [&] {
395
- vllm::dynamic_per_token_scaled_fp8_quant_kernel_strided<
396
- scalar_t, fp8_t><<<grid, block, 0, stream>>>(
397
- out.data_ptr<fp8_t>(), scales.data_ptr<float>(),
398
- input.data_ptr<scalar_t>(),
399
- scale_ub.has_value() ? scale_ub->data_ptr<float>() : nullptr,
400
- hidden_size, in_row_stride, out_row_stride);
411
+ vllm::dynamic_per_token_scaled_fp8_quant_kernel_strided<scalar_t,
412
+ fp8_t>
413
+ <<<grid, block, 0, stream>>>(
414
+ out.mutable_data_ptr<fp8_t>(),
415
+ scales.mutable_data_ptr<float>(),
416
+ input.const_data_ptr<scalar_t>(),
417
+ scale_ub.has_value() ? scale_ub->const_data_ptr<float>()
418
+ : nullptr,
419
+ hidden_size, in_row_stride, out_row_stride);
420
});
421
});
422
}
csrc/libtorch_stable/quantization/w8a8/int8/scaled_quant.cu
renamed
+46
-40
@@ -1,12 +1,11 @@
1
-#include <ATen/cuda/CUDAContext.h>
2
-#include <torch/all.h>
3
-#include <c10/cuda/CUDAGuard.h>
1
+#include <torch/csrc/stable/tensor.h>
2
3
#include <cmath>
4
7
-#include "dispatch_utils.h"
8
-#include "libtorch_stable/quantization/vectorization_utils.cuh"
9
-#include "cub_helpers.h"
5
+#include "../../../dispatch_utils.h"
6
+#include "../../../torch_utils.h"
7
+#include "../../vectorization_utils.cuh"
8
+#include "../../../../cub_helpers.h"
9
10
static inline __device__ int8_t float_to_int8_rn(float x) {
11
#ifdef USE_ROCM
@@ -263,66 +262,73 @@ __global__ void dynamic_scaled_int8_azp_quant_kernel(
262
263
} // namespace vllm
264
266
-void static_scaled_int8_quant(torch::Tensor& out, // [..., hidden_size]
267
- torch::Tensor const& input, // [..., hidden_size]
268
- torch::Tensor const& scale,
269
- std::optional<torch::Tensor> const& azp) {
270
- TORCH_CHECK(input.is_contiguous());
271
- TORCH_CHECK(out.is_contiguous());
272
- TORCH_CHECK(scale.numel() == 1);
273
- TORCH_CHECK(!azp || azp->numel() == 1);
265
+void static_scaled_int8_quant(
266
+ torch::stable::Tensor& out, // [..., hidden_size]
267
+ torch::stable::Tensor const& input, // [..., hidden_size]
268
+ torch::stable::Tensor const& scale,
269
+ std::optional<torch::stable::Tensor> const& azp) {
270
+ STD_TORCH_CHECK(input.is_contiguous());
271
+ STD_TORCH_CHECK(out.is_contiguous());
272
+ STD_TORCH_CHECK(scale.numel() == 1);
273
+ STD_TORCH_CHECK(!azp || azp->numel() == 1);
274
275
int const hidden_size = input.size(-1);
276
int const num_tokens = input.numel() / hidden_size;
277
dim3 const grid(num_tokens);
278
dim3 const block(std::min(hidden_size, 256));
279
- const at::cuda::OptionalCUDAGuard device_guard(device_of(input));
280
- const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
281
- VLLM_DISPATCH_FLOATING_TYPES(
279
+ const torch::stable::accelerator::DeviceGuard device_guard(
280
+ input.get_device_index());
281
+ const cudaStream_t stream = get_current_cuda_stream();
282
+ VLLM_STABLE_DISPATCH_FLOATING_TYPES(
283
input.scalar_type(), "static_scaled_int8_quant_kernel", [&] {
284
if (!azp) {
285
vllm::static_scaled_int8_quant_kernel<scalar_t, float>
285
- <<<grid, block, 0, stream>>>(
286
- input.data_ptr<scalar_t>(), out.data_ptr<int8_t>(),
287
- scale.data_ptr<float>(), hidden_size);
286
+ <<<grid, block, 0, stream>>>(input.const_data_ptr<scalar_t>(),
287
+ out.mutable_data_ptr<int8_t>(),
288
+ scale.const_data_ptr<float>(),
289
+ hidden_size);
290
} else {
291
vllm::static_scaled_int8_azp_quant_kernel<scalar_t, float, int32_t>
292
<<<grid, block, 0, stream>>>(
291
- input.data_ptr<scalar_t>(), out.data_ptr<int8_t>(),
292
- scale.data_ptr<float>(), azp->data_ptr<int32_t>(),
293
- hidden_size);
293
+ input.const_data_ptr<scalar_t>(),
294
+ out.mutable_data_ptr<int8_t>(), scale.const_data_ptr<float>(),
295
+ azp->const_data_ptr<int32_t>(), hidden_size);
296
}
297
});
298
}
299
300
void dynamic_scaled_int8_quant(
299
- torch::Tensor& out, // [..., hidden_size]
300
- torch::Tensor const& input, // [..., hidden_size]
301
- torch::Tensor& scales, std::optional<torch::Tensor> const& azp) {
302
- TORCH_CHECK(input.is_contiguous());
303
- TORCH_CHECK(out.is_contiguous());
304
- TORCH_CHECK(scales.is_contiguous());
305
- TORCH_CHECK(!azp || azp->is_contiguous());
301
+ torch::stable::Tensor& out, // [..., hidden_size]
302
+ torch::stable::Tensor const& input, // [..., hidden_size]
303
+ torch::stable::Tensor& scales,
304
+ std::optional<torch::stable::Tensor> const& azp) {
305
+ STD_TORCH_CHECK(input.is_contiguous());
306
+ STD_TORCH_CHECK(out.is_contiguous());
307
+ STD_TORCH_CHECK(scales.is_contiguous());
308
+ STD_TORCH_CHECK(!azp || azp->is_contiguous());
309
310
int const hidden_size = input.size(-1);
311
int const num_tokens = input.numel() / hidden_size;
312
dim3 const grid(num_tokens);
313
dim3 const block(std::min(hidden_size, 256));
311
- const at::cuda::OptionalCUDAGuard device_guard(device_of(input));
312
- const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
313
- VLLM_DISPATCH_FLOATING_TYPES(
314
+ const torch::stable::accelerator::DeviceGuard device_guard(
315
+ input.get_device_index());
316
+ const cudaStream_t stream = get_current_cuda_stream();
317
+ VLLM_STABLE_DISPATCH_FLOATING_TYPES(
318
input.scalar_type(), "dynamic_scaled_int8_quant_kernel", [&] {
319
if (!azp) {
320
vllm::dynamic_scaled_int8_quant_kernel<scalar_t, float>
317
- <<<grid, block, 0, stream>>>(
318
- input.data_ptr<scalar_t>(), out.data_ptr<int8_t>(),
319
- scales.data_ptr<float>(), hidden_size);
321
+ <<<grid, block, 0, stream>>>(input.const_data_ptr<scalar_t>(),
322
+ out.mutable_data_ptr<int8_t>(),
323
+ scales.mutable_data_ptr<float>(),
324
+ hidden_size);
325
} else {
326
vllm::dynamic_scaled_int8_azp_quant_kernel<scalar_t, float, int32_t>
322
- <<<grid, block, 0, stream>>>(
323
- input.data_ptr<scalar_t>(), out.data_ptr<int8_t>(),
324
- scales.data_ptr<float>(), azp->data_ptr<int32_t>(),
325
- hidden_size);
327
+ <<<grid, block, 0, stream>>>(input.const_data_ptr<scalar_t>(),
328
+ out.mutable_data_ptr<int8_t>(),
329
+ scales.mutable_data_ptr<float>(),
330
+ azp->mutable_data_ptr<int32_t>(),
331
+ hidden_size);
332
}
333
});
334
}
\ No newline at end of file
csrc/libtorch_stable/torch_bindings.cpp
+139
@@ -266,6 +266,109 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
266
// Hadamard transforms
267
// conditionally compiled so impl registration is in source file
268
ops.def("hadacore_transform(Tensor! x, bool inplace) -> Tensor");
269
+
270
+ // Activation ops
271
+ // Activation function used in SwiGLU.
272
+ ops.def("silu_and_mul(Tensor! result, Tensor input) -> ()");
273
+ ops.def("mul_and_silu(Tensor! out, Tensor input) -> ()");
274
+
275
+ // SwiGLU activation with input clamping.
276
+ ops.def(
277
+ "silu_and_mul_with_clamp(Tensor! result, Tensor input, float limit) "
278
+ "-> ()");
279
+
280
+ // Activation function used in GeGLU with `none` approximation.
281
+ ops.def("gelu_and_mul(Tensor! out, Tensor input) -> ()");
282
+
283
+ // Activation function used in GeGLU with `tanh` approximation.
284
+ ops.def("gelu_tanh_and_mul(Tensor! out, Tensor input) -> ()");
285
+
286
+ // FATReLU implementation.
287
+ ops.def("fatrelu_and_mul(Tensor! out, Tensor input, float threshold) -> ()");
288
+
289
+ ops.def(
290
+ "swigluoai_and_mul(Tensor! out, Tensor input, float alpha=1.702, float "
291
+ "limit=7.0) "
292
+ "-> ()");
293
+
294
+ // GELU implementation used in GPT-2.
295
+ ops.def("gelu_new(Tensor! out, Tensor input) -> ()");
296
+
297
+ // Approximate GELU implementation.
298
+ ops.def("gelu_fast(Tensor! out, Tensor input) -> ()");
299
+
300
+ // Quick GELU implementation.
301
+ ops.def("gelu_quick(Tensor! out, Tensor input) -> ()");
302
+
303
+ // Compute int8 quantized tensor for given scaling factor.
304
+ ops.def(
305
+ "static_scaled_int8_quant(Tensor! result, Tensor input, Tensor scale,"
306
+ "Tensor? azp) -> ()");
307
+
308
+ // Compute int8 quantized tensor and scaling factor
309
+ ops.def(
310
+ "dynamic_scaled_int8_quant(Tensor! result, Tensor input, Tensor! scale, "
311
+ "Tensor!? azp) -> ()");
312
+
313
+ // Compute FP8 quantized tensor for given scaling factor.
314
+ // Supports per-tensor, per-channel, per-token, and arbitrary 2D group
315
+ // scaling. Optional group_m/group_n specify the group shape explicitly;
316
+ // required for 1D scales to disambiguate per-channel vs per-token.
317
+ ops.def(
318
+ "static_scaled_fp8_quant(Tensor! result, Tensor input, Tensor scale, "
319
+ "int[]? group_shape=None) -> ()");
320
+
321
+ // Compute dynamic-per-tensor FP8 quantized tensor and scaling factor.
322
+ ops.def(
323
+ "dynamic_scaled_fp8_quant(Tensor! result, Tensor input, Tensor! scale) "
324
+ "-> "
325
+ "()");
326
+
327
+ // Compute dynamic-per-token FP8 quantized tensor and scaling factor.
328
+ ops.def(
329
+ "dynamic_per_token_scaled_fp8_quant(Tensor! result, Tensor input, "
330
+ "Tensor! scale, Tensor? scale_ub) -> "
331
+ "()");
332
+
333
+ // Quantized GEMM for GPTQ.
334
+ // Note: even though the C++ inferred schema is correct for this op, it seems
335
+ // to prevent the meta function registry.
336
+ ops.def(
337
+ "gptq_gemm(Tensor a, Tensor b_q_weight, Tensor b_gptq_qzeros, "
338
+ "Tensor b_gptq_scales, Tensor b_g_idx, bool use_exllama, bool "
339
+ "use_v2_format, int bit) "
340
+ "-> Tensor");
341
+
342
+ // Post processing for GPTQ.
343
+ ops.def("gptq_shuffle(Tensor! q_weight, Tensor q_perm, int bit) -> ()");
344
+
345
+ // Dequantization for GGML.
346
+ ops.def(
347
+ "ggml_dequantize(Tensor W, int type, SymInt m, SymInt n, ScalarType? "
348
+ "dtype) -> Tensor");
349
+
350
+ // mmvq kernel for GGML.
351
+ ops.def(
352
+ "ggml_mul_mat_vec_a8(Tensor W, Tensor X, int type, SymInt row) "
353
+ "-> Tensor");
354
+
355
+ // mmq kernel for GGML.
356
+ ops.def(
357
+ "ggml_mul_mat_a8(Tensor W, Tensor X, int type, SymInt row) -> Tensor");
358
+
359
+ // moe kernel for GGML.
360
+ ops.def(
361
+ "ggml_moe_a8(Tensor X, Tensor W, "
362
+ "Tensor sorted_token_ids, Tensor expert_ids, Tensor "
363
+ "num_tokens_post_padded, "
364
+ "int type, SymInt row, SymInt top_k, SymInt tokens) -> Tensor");
365
+
366
+ ops.def(
367
+ "ggml_moe_a8_vec(Tensor X, Tensor W, "
368
+ "Tensor topk_ids, int top_k, "
369
+ "int type, SymInt row, SymInt tokens) -> Tensor");
370
+
371
+ ops.def("ggml_moe_get_block_size(int type) -> int");
372
}
373
374
STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) {
@@ -312,6 +415,39 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) {
415
// AllSpark ops: conditionally compiled so impl registrations are in source
416
// files (allspark_repack.cu and allspark_qgemm_w8a16.cu)
417
#endif
418
+
419
+ // Activation kernels (shared CUDA/ROCm)
420
+ ops.impl("silu_and_mul", TORCH_BOX(&silu_and_mul));
421
+ ops.impl("mul_and_silu", TORCH_BOX(&mul_and_silu));
422
+ ops.impl("gelu_and_mul", TORCH_BOX(&gelu_and_mul));
423
+ ops.impl("gelu_tanh_and_mul", TORCH_BOX(&gelu_tanh_and_mul));
424
+ ops.impl("fatrelu_and_mul", TORCH_BOX(&fatrelu_and_mul));
425
+ ops.impl("swigluoai_and_mul", TORCH_BOX(&swigluoai_and_mul));
426
+ ops.impl("gelu_new", TORCH_BOX(&gelu_new));
427
+ ops.impl("gelu_fast", TORCH_BOX(&gelu_fast));
428
+ ops.impl("gelu_quick", TORCH_BOX(&gelu_quick));
429
+ ops.impl("silu_and_mul_with_clamp", TORCH_BOX(&silu_and_mul_clamp));
430
+
431
+ // INT8 quantization kernels
432
+ ops.impl("static_scaled_int8_quant", TORCH_BOX(&static_scaled_int8_quant));
433
+ ops.impl("dynamic_scaled_int8_quant", TORCH_BOX(&dynamic_scaled_int8_quant));
434
+
435
+ // FP8 quantization kernels
436
+ ops.impl("static_scaled_fp8_quant", TORCH_BOX(&static_scaled_fp8_quant));
437
+ ops.impl("dynamic_scaled_fp8_quant", TORCH_BOX(&dynamic_scaled_fp8_quant));
438
+ ops.impl("dynamic_per_token_scaled_fp8_quant",
439
+ TORCH_BOX(&dynamic_per_token_scaled_fp8_quant));
440
+
441
+ // GPTQ kernels
442
+ ops.impl("gptq_gemm", TORCH_BOX(&gptq_gemm));
443
+ ops.impl("gptq_shuffle", TORCH_BOX(&gptq_shuffle));
444
+
445
+ // GGML kernels
446
+ ops.impl("ggml_dequantize", TORCH_BOX(&ggml_dequantize));
447
+ ops.impl("ggml_mul_mat_vec_a8", TORCH_BOX(&ggml_mul_mat_vec_a8));
448
+ ops.impl("ggml_mul_mat_a8", TORCH_BOX(&ggml_mul_mat_a8));
449
+ ops.impl("ggml_moe_a8", TORCH_BOX(&ggml_moe_a8));
450
+ ops.impl("ggml_moe_a8_vec", TORCH_BOX(&ggml_moe_a8_vec));
451
}
452
453
// These capability-check functions take only primitive args (no tensors), so
@@ -329,6 +465,9 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CompositeExplicitAutograd, ops) {
465
ops.impl("cutlass_scaled_mm_supports_fp4",
466
TORCH_BOX(&cutlass_scaled_mm_supports_fp4));
467
#endif
468
+
469
+ // GGML block size lookup (no tensor args)
470
+ ops.impl("ggml_moe_get_block_size", TORCH_BOX(&ggml_moe_get_block_size));
471
}
472
473
REGISTER_EXTENSION(_C_stable_libtorch)
csrc/libtorch_stable/torch_utils.h
+5
-1
@@ -6,8 +6,12 @@
6
#include <torch/csrc/stable/tensor.h>
7
#include <torch/headeronly/util/shim_utils.h>
8
9
+#ifndef USE_ROCM
10
+ #include <cuda_runtime.h>
11
+#else
12
+ #include <hip/hip_runtime.h>
13
+#endif
14
#include <cublas_v2.h>
10
-#include <cuda_runtime.h>
15
16
#include <deque>
17
#include <mutex>
csrc/ops.h
-47
@@ -149,17 +149,10 @@ void persistent_masked_m_silu_mul_quant(
149
at::Tensor& y_s, // (E, T, H//group_size) [OUT]
150
bool use_ue8m0);
151
152
-void mul_and_silu(torch::Tensor& out, torch::Tensor& input);
153
-
152
void gelu_and_mul(torch::Tensor& out, torch::Tensor& input);
153
154
void gelu_tanh_and_mul(torch::Tensor& out, torch::Tensor& input);
155
158
-void fatrelu_and_mul(torch::Tensor& out, torch::Tensor& input,
159
- double threshold);
160
-void swigluoai_and_mul(torch::Tensor& out, torch::Tensor& input,
161
- double alpha = 1.702, double limit = 7.0);
162
-
156
void gelu_new(torch::Tensor& out, torch::Tensor& input);
157
158
void gelu_fast(torch::Tensor& out, torch::Tensor& input);
@@ -174,28 +167,6 @@ void cutlass_mla_decode(torch::Tensor const& out, torch::Tensor const& q_nope,
167
168
torch::Tensor get_cuda_view_from_cpu_tensor(torch::Tensor& cpu_tensor);
169
177
-torch::Tensor ggml_dequantize(torch::Tensor W, int64_t type, int64_t m,
178
- int64_t n,
179
- std::optional<at::ScalarType> const& dtype);
180
-
181
-torch::Tensor ggml_mul_mat_vec_a8(torch::Tensor W, torch::Tensor X,
182
- int64_t type, int64_t row);
183
-
184
-torch::Tensor ggml_mul_mat_a8(torch::Tensor W, torch::Tensor X, int64_t type,
185
- int64_t row);
186
-
187
-torch::Tensor ggml_moe_a8(torch::Tensor X, torch::Tensor W,
188
- torch::Tensor sorted_token_ids,
189
- torch::Tensor expert_ids,
190
- torch::Tensor num_tokens_post_padded, int64_t type,
191
- int64_t row, int64_t top_k, int64_t tokens);
192
-
193
-torch::Tensor ggml_moe_a8_vec(torch::Tensor X, torch::Tensor W,
194
- torch::Tensor topk_ids, int64_t top_k,
195
- int64_t type, int64_t row, int64_t tokens);
196
-
197
-int64_t ggml_moe_get_block_size(int64_t type);
198
-
170
void static_scaled_int8_quant(torch::Tensor& out, torch::Tensor const& input,
171
torch::Tensor const& scale,
172
std::optional<torch::Tensor> const& azp);
@@ -204,24 +175,6 @@ void dynamic_scaled_int8_quant(torch::Tensor& out, torch::Tensor const& input,
175
torch::Tensor& scales,
176
std::optional<torch::Tensor> const& azp);
177
207
-torch::Tensor gptq_gemm(torch::Tensor a, torch::Tensor b_q_weight,
208
- torch::Tensor b_gptq_qzeros,
209
- torch::Tensor b_gptq_scales, torch::Tensor b_g_idx,
210
- bool use_exllama, bool use_v2_format, int64_t bit);
211
-
212
-void gptq_shuffle(torch::Tensor q_weight, torch::Tensor q_perm, int64_t bit);
213
-
214
-void static_scaled_fp8_quant(
215
- torch::Tensor& out, torch::Tensor const& input, torch::Tensor const& scale,
216
- std::optional<std::tuple<int64_t, int64_t>> group_shape = std::nullopt);
217
-
218
-void dynamic_scaled_fp8_quant(torch::Tensor& out, torch::Tensor const& input,
219
- torch::Tensor& scale);
220
-
221
-void dynamic_per_token_scaled_fp8_quant(
222
- torch::Tensor& out, torch::Tensor const& input, torch::Tensor& scale,
223
- std::optional<torch::Tensor> const& scale_ub);
224
-
178
void selective_scan_fwd(
179
const torch::Tensor& u, const torch::Tensor& delta, const torch::Tensor& A,
180
const torch::Tensor& B, const torch::Tensor& C,
csrc/quantization/utils.cuh
+16
-14
@@ -7,23 +7,23 @@
7
*/
8
9
#include <cmath>
10
-#include <torch/types.h>
10
+#include <torch/headeronly/macros/Macros.h>
11
12
#ifndef USE_ROCM
13
- #include <c10/util/Float8_e4m3fn.h>
13
+ #include <torch/headeronly/util/Float8_e4m3fn.h>
14
#define MAYBE_HOST_DEVICE C10_HOST_DEVICE
15
#else
16
- #include <ATen/hip/HIPContext.h>
17
- #include <c10/util/Float8_e4m3fn.h>
18
- #include <c10/util/Float8_e4m3fnuz.h>
16
+ #include <torch/headeronly/util/Float8_e4m3fn.h>
17
+ #include <torch/headeronly/util/Float8_e4m3fnuz.h>
18
// ROCm doesn't seem to need C10_HOST_DEVICE for static constexpr
19
#define MAYBE_HOST_DEVICE
20
#endif
21
22
template <typename T,
24
- typename = std::enable_if_t<std::is_same_v<T, c10::Float8_e4m3fn> ||
25
- std::is_same_v<T, c10::Float8_e4m3fnuz> ||
26
- std::is_same_v<T, int8_t>>>
23
+ typename = std::enable_if_t<
24
+ std::is_same_v<T, torch::headeronly::Float8_e4m3fn> ||
25
+ std::is_same_v<T, torch::headeronly::Float8_e4m3fnuz> ||
26
+ std::is_same_v<T, int8_t>>>
27
struct quant_type_max {
28
static constexpr T val() { return std::numeric_limits<T>::max(); }
29
};
@@ -31,9 +31,10 @@ struct quant_type_max {
31
// Using the default max value from pytorch (240.0 0x7F) will cause accuracy
32
// issues when running dynamic quantization. Here use 224.0 0x7E for rocm.
33
template <>
34
-struct quant_type_max<c10::Float8_e4m3fnuz> {
35
- static constexpr c10::Float8_e4m3fnuz val() {
36
- return c10::Float8_e4m3fnuz(0x7E, c10::Float8_e4m3fnuz::from_bits());
34
+struct quant_type_max<torch::headeronly::Float8_e4m3fnuz> {
35
+ static constexpr torch::headeronly::Float8_e4m3fnuz val() {
36
+ return torch::headeronly::Float8_e4m3fnuz(
37
+ 0x7E, torch::headeronly::Float8_e4m3fnuz::from_bits());
38
}
39
};
40
@@ -42,9 +43,10 @@ MAYBE_HOST_DEVICE static constexpr T quant_type_max_v =
43
quant_type_max<T>::val();
44
45
template <typename T,
45
- typename = std::enable_if_t<std::is_same_v<T, c10::Float8_e4m3fn> ||
46
- std::is_same_v<T, c10::Float8_e4m3fnuz> ||
47
- std::is_same_v<T, int8_t>>>
46
+ typename = std::enable_if_t<
47
+ std::is_same_v<T, torch::headeronly::Float8_e4m3fn> ||
48
+ std::is_same_v<T, torch::headeronly::Float8_e4m3fnuz> ||
49
+ std::is_same_v<T, int8_t>>>
50
struct min_scaling_factor {
51
C10_DEVICE C10_ALWAYS_INLINE static float val() {
52
return 1.0f / (quant_type_max_v<T> * 512.0f);
csrc/quantization/w8a8/fp8/common.cuh
+18
-1
@@ -5,6 +5,19 @@
5
6
#include <cmath>
7
8
+// This header is shared between _C and _C_stable_libtorch targets.
9
+// torch_utils.h provides get_device_prop(). We need to pass USE_CUDA
10
+// to the .so to expose some of the shims used by torch_utils.h. For now
11
+// this is only done for _C_stable_libtorch and not for _C, so we use the
12
+// non stable at::cuda::getCurrentDeviceProperties for _C for now.
13
+#ifdef TORCH_TARGET_VERSION
14
+ #include "../../../libtorch_stable/torch_utils.h"
15
+#else
16
+ #ifdef USE_ROCM
17
+ #include <ATen/hip/HIPContext.h>
18
+ #endif
19
+#endif
20
+
21
#ifndef USE_ROCM
22
#include "nvidia/quant_utils.cuh"
23
#else
@@ -18,7 +31,11 @@ static bool is_fp8_ocp() {
31
#ifndef USE_ROCM
32
return true;
33
#else
21
- auto dprops = at::cuda::getCurrentDeviceProperties();
34
+ #ifdef TORCH_TARGET_VERSION
35
+ auto* dprops = get_device_prop();
36
+ #else
37
+ auto* dprops = at::cuda::getCurrentDeviceProperties();
38
+ #endif
39
std::string device_arch = dprops->gcnArchName;
40
size_t substring = device_arch.find("gfx94");
41
return substring == std::string::npos;
csrc/torch_bindings.cpp
+1
-128
@@ -77,17 +77,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
77
" Tensor? output_scale=None) -> ()");
78
ops.impl("merge_attn_states", torch::kCUDA, &merge_attn_states);
79
80
- // Activation ops
81
- // Activation function used in SwiGLU.
82
- ops.def("silu_and_mul(Tensor! result, Tensor input) -> ()");
83
- ops.impl("silu_and_mul", torch::kCUDA, &silu_and_mul);
84
-
85
- // SwiGLU activation with input clamping.
86
- ops.def(
87
- "silu_and_mul_with_clamp(Tensor! result, Tensor input, float limit) "
88
- "-> ()");
89
- ops.impl("silu_and_mul_with_clamp", torch::kCUDA, &silu_and_mul_clamp);
90
-
80
+ // Activation ops (quantized only — basic ops moved to _C_stable_libtorch)
81
ops.def(
82
"silu_and_mul_quant(Tensor! result, Tensor input, Tensor scale) -> ()");
83
ops.impl("silu_and_mul_quant", torch::kCUDA, &silu_and_mul_quant);
@@ -104,39 +94,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
94
ops.impl("silu_and_mul_per_block_quant", torch::kCUDA,
95
&silu_and_mul_per_block_quant);
96
107
- ops.def("mul_and_silu(Tensor! out, Tensor input) -> ()");
108
- ops.impl("mul_and_silu", torch::kCUDA, &mul_and_silu);
109
-
110
- // Activation function used in GeGLU with `none` approximation.
111
- ops.def("gelu_and_mul(Tensor! out, Tensor input) -> ()");
112
- ops.impl("gelu_and_mul", torch::kCUDA, &gelu_and_mul);
113
-
114
- // Activation function used in GeGLU with `tanh` approximation.
115
- ops.def("gelu_tanh_and_mul(Tensor! out, Tensor input) -> ()");
116
- ops.impl("gelu_tanh_and_mul", torch::kCUDA, &gelu_tanh_and_mul);
117
-
118
- // FATReLU implementation.
119
- ops.def("fatrelu_and_mul(Tensor! out, Tensor input, float threshold) -> ()");
120
- ops.impl("fatrelu_and_mul", torch::kCUDA, &fatrelu_and_mul);
121
-
122
- ops.def(
123
- "swigluoai_and_mul(Tensor! out, Tensor input, float alpha=1.702, float "
124
- "limit=7.0) "
125
- "-> ()");
126
- ops.impl("swigluoai_and_mul", torch::kCUDA, &swigluoai_and_mul);
127
-
128
- // GELU implementation used in GPT-2.
129
- ops.def("gelu_new(Tensor! out, Tensor input) -> ()");
130
- ops.impl("gelu_new", torch::kCUDA, &gelu_new);
131
-
132
- // Approximate GELU implementation.
133
- ops.def("gelu_fast(Tensor! out, Tensor input) -> ()");
134
- ops.impl("gelu_fast", torch::kCUDA, &gelu_fast);
135
-
136
- // Quick GELU implementation.
137
- ops.def("gelu_quick(Tensor! out, Tensor input) -> ()");
138
- ops.impl("gelu_quick", torch::kCUDA, &gelu_quick);
139
-
97
// Layernorm
98
// Apply Root Mean Square (RMS) Normalization to the input tensor.
99
ops.def(
@@ -318,39 +275,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
275
276
#endif
277
321
- // Dequantization for GGML.
322
- ops.def(
323
- "ggml_dequantize(Tensor W, int type, SymInt m, SymInt n, ScalarType? "
324
- "dtype) -> Tensor");
325
- ops.impl("ggml_dequantize", torch::kCUDA, &ggml_dequantize);
326
-
327
- // mmvq kernel for GGML.
328
- ops.def(
329
- "ggml_mul_mat_vec_a8(Tensor W, Tensor X, int type, SymInt row) "
330
- "-> Tensor");
331
- ops.impl("ggml_mul_mat_vec_a8", torch::kCUDA, &ggml_mul_mat_vec_a8);
332
-
333
- // mmq kernel for GGML.
334
- ops.def(
335
- "ggml_mul_mat_a8(Tensor W, Tensor X, int type, SymInt row) -> Tensor");
336
- ops.impl("ggml_mul_mat_a8", torch::kCUDA, &ggml_mul_mat_a8);
337
-
338
- // moe kernel for GGML.
339
- ops.def(
340
- "ggml_moe_a8(Tensor X, Tensor W, "
341
- "Tensor sorted_token_ids, Tensor expert_ids, Tensor "
342
- "num_tokens_post_padded, "
343
- "int type, SymInt row, SymInt top_k, SymInt tokens) -> Tensor");
344
- ops.impl("ggml_moe_a8", torch::kCUDA, &ggml_moe_a8);
345
-
346
- ops.def(
347
- "ggml_moe_a8_vec(Tensor X, Tensor W, "
348
- "Tensor topk_ids, int top_k, "
349
- "int type, SymInt row, SymInt tokens) -> Tensor");
350
- ops.impl("ggml_moe_a8_vec", torch::kCUDA, &ggml_moe_a8_vec);
351
-
352
- ops.def("ggml_moe_get_block_size", &ggml_moe_get_block_size);
353
-
278
#ifndef USE_ROCM
279
// Expert-specialization mxfp8 blockscaled grouped quantization (SM100+).
280
ops.def(
@@ -370,57 +294,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
294
295
#endif
296
373
- // Quantized GEMM for GPTQ.
374
- // Note: even though the C++ inferred schema is correct for this op, it seems
375
- // to prevent the meta function registry.
376
- ops.def(
377
- "gptq_gemm(Tensor a, Tensor b_q_weight, Tensor b_gptq_qzeros, "
378
- "Tensor b_gptq_scales, Tensor b_g_idx, bool use_exllama, bool "
379
- "use_v2_format, int bit) "
380
- "-> Tensor");
381
- ops.impl("gptq_gemm", torch::kCUDA, &gptq_gemm);
382
-
383
- // Post processing for GPTQ.
384
- ops.def("gptq_shuffle(Tensor! q_weight, Tensor q_perm, int bit) -> ()");
385
- ops.impl("gptq_shuffle", torch::kCUDA, &gptq_shuffle);
386
-
387
- // Compute FP8 quantized tensor for given scaling factor.
388
- // Supports per-tensor, per-channel, per-token, and arbitrary 2D group
389
- // scaling. Optional group_m/group_n specify the group shape explicitly;
390
- // required for 1D scales to disambiguate per-channel vs per-token.
391
- ops.def(
392
- "static_scaled_fp8_quant(Tensor! result, Tensor input, Tensor scale, "
393
- "(int, int)? group_shape=None) -> ()");
394
- ops.impl("static_scaled_fp8_quant", torch::kCUDA, &static_scaled_fp8_quant);
395
-
396
- // Compute dynamic-per-tensor FP8 quantized tensor and scaling factor.
397
- ops.def(
398
- "dynamic_scaled_fp8_quant(Tensor! result, Tensor input, Tensor! scale) "
399
- "-> "
400
- "()");
401
- ops.impl("dynamic_scaled_fp8_quant", torch::kCUDA, &dynamic_scaled_fp8_quant);
402
-
403
- // Compute dynamic-per-token FP8 quantized tensor and scaling factor.
404
- ops.def(
405
- "dynamic_per_token_scaled_fp8_quant(Tensor! result, Tensor input, "
406
- "Tensor! scale, Tensor? scale_ub) -> "
407
- "()");
408
- ops.impl("dynamic_per_token_scaled_fp8_quant", torch::kCUDA,
409
- &dynamic_per_token_scaled_fp8_quant);
410
-
411
- // Compute int8 quantized tensor for given scaling factor.
412
- ops.def(
413
- "static_scaled_int8_quant(Tensor! result, Tensor input, Tensor scale,"
414
- "Tensor? azp) -> ()");
415
- ops.impl("static_scaled_int8_quant", torch::kCUDA, &static_scaled_int8_quant);
416
-
417
- // Compute int8 quantized tensor and scaling factor
418
- ops.def(
419
- "dynamic_scaled_int8_quant(Tensor! result, Tensor input, Tensor! scale, "
420
- "Tensor!? azp) -> ()");
421
- ops.impl("dynamic_scaled_int8_quant", torch::kCUDA,
422
- &dynamic_scaled_int8_quant);
423
-
297
// Mamba selective scan kernel
298
ops.def(
299
"selective_scan_fwd(Tensor! u, Tensor! delta,"
csrc/torch_utils.h
new
+17
@@ -0,0 +1,17 @@
1
+#pragma once
2
+
3
+// Shared TORCH_UTILS_CHECK across both libtorch stable and unstable source
4
+// files. Keep this header free of CUTLASS/CUTE so attention/quant headers can
5
+// use it.
6
+//
7
+// If TORCH_TARGET_VERSION is defined, we are building _C_stable_libtorch.so so
8
+// use STD_TORCH_CHECK via header-only.
9
+// Otherwise, use TORCH_CHECK via torch/all.h.
10
+
11
+#ifdef TORCH_TARGET_VERSION
12
+ #include <torch/headeronly/util/Exception.h>
13
+ #define TORCH_UTILS_CHECK STD_TORCH_CHECK
14
+#else
15
+ #include <torch/all.h>
16
+ #define TORCH_UTILS_CHECK TORCH_CHECK
17
+#endif
setup.py
+1
-3
@@ -1045,9 +1045,7 @@ if _is_cpu():
1045
1046
if _build_custom_ops():
1047
ext_modules.append(CMakeExtension(name="vllm._C"))
1048
- # also _is_hip() once https://github.com/vllm-project/vllm/issues/35163 is
1049
- # fixed
1050
- if _is_cuda():
1048
+ if _is_cuda() or _is_hip():
1049
ext_modules.append(CMakeExtension(name="vllm._C_stable_libtorch"))
1050
1051
package_data = {
vllm/platforms/rocm.py
+5
@@ -44,6 +44,11 @@ except ImportError as e:
44
logger.warning("Failed to import from vllm._C with %r", e)
45
46
# import custom ops, trigger op registration
47
+try:
48
+ import vllm._C_stable_libtorch # noqa: F401
49
+except ImportError as e:
50
+ logger.warning("Failed to import from vllm._C_stable_libtorch with %r", e)
51
+
52
try:
53
import vllm._rocm_C # noqa: F401
54
except ImportError as e: