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 14, 23:45 UTC
[Perf] Cuda Kernel for Int8 Per Token Group Quant (#21476)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
Wentao Ye committed
Jul 25, 2025 at 20:07 UTC
75d29cf4e1d7e950c2308b12e944b507fb3e1916
6 files changed
+47
-3
csrc/ops.h
+5
@@ -292,6 +292,11 @@ void per_token_group_quant_fp8(const torch::Tensor& input,
292
torch::Tensor& output_q, torch::Tensor& output_s,
293
int64_t group_size, double eps, double fp8_min,
294
double fp8_max, bool scale_ue8m0);
295
+
296
+void per_token_group_quant_int8(const torch::Tensor& input,
297
+ torch::Tensor& output_q,
298
+ torch::Tensor& output_s, int64_t group_size,
299
+ double eps, double int8_min, double int8_max);
300
#endif
301
302
void static_scaled_int8_quant(torch::Tensor& out, torch::Tensor const& input,
csrc/quantization/compressed_tensors/int8_quant_kernels.cu
+10
@@ -1,6 +1,8 @@
1
#include <ATen/cuda/CUDAContext.h>
2
#include <torch/all.h>
3
4
+#include "../per_token_group_quant_8bit.h"
5
+
6
#include <cmath>
7
8
#include "../../dispatch_utils.h"
@@ -336,3 +338,11 @@ void dynamic_scaled_int8_quant(
338
}
339
});
340
}
341
+
342
+void per_token_group_quant_int8(const torch::Tensor& input,
343
+ torch::Tensor& output_q,
344
+ torch::Tensor& output_s, int64_t group_size,
345
+ double eps, double int8_min, double int8_max) {
346
+ per_token_group_quant_8bit(input, output_q, output_s, group_size, eps,
347
+ int8_min, int8_max);
348
+}
\ No newline at end of file
csrc/quantization/fp8/per_token_group_quant.cu
+5
-1
@@ -1,6 +1,8 @@
1
#include <ATen/cuda/CUDAContext.h>
2
#include <c10/util/Float8_e4m3fn.h>
3
4
+#include "../per_token_group_quant_8bit.h"
5
+
6
#include <cmath>
7
8
#include <cuda_fp16.h>
@@ -120,7 +122,7 @@ void per_token_group_quant_8bit(const torch::Tensor& input,
122
torch::Tensor& output_q,
123
torch::Tensor& output_s, int64_t group_size,
124
double eps, double min_8bit, double max_8bit,
123
- bool scale_ue8m0 = false) {
125
+ bool scale_ue8m0) {
126
TORCH_CHECK(input.is_contiguous());
127
TORCH_CHECK(output_q.is_contiguous());
128
@@ -198,6 +200,8 @@ void per_token_group_quant_8bit(const torch::Tensor& input,
200
input.scalar_type(), "per_token_group_quant_8bit", ([&] {
201
if (dst_type == at::ScalarType::Float8_e4m3fn) {
202
LAUNCH_KERNEL(scalar_t, c10::Float8_e4m3fn);
203
+ } else if (dst_type == at::ScalarType::Char) {
204
+ LAUNCH_KERNEL(scalar_t, int8_t);
205
}
206
}));
207
csrc/quantization/per_token_group_quant_8bit.h
new
+10
@@ -0,0 +1,10 @@
1
+#pragma once
2
+#include <torch/all.h>
3
+
4
+// TODO(wentao): refactor the folder to 8bit, then includes fp8 and int8 folders
5
+// 8-bit per-token-group quantization helper used by both FP8 and INT8
6
+void per_token_group_quant_8bit(const torch::Tensor& input,
7
+ torch::Tensor& output_q,
8
+ torch::Tensor& output_s, int64_t group_size,
9
+ double eps, double min_8bit, double max_8bit,
10
+ bool scale_ue8m0 = false);
\ No newline at end of file
csrc/torch_bindings.cpp
+8
@@ -624,6 +624,14 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
624
ops.impl("per_token_group_fp8_quant", torch::kCUDA,
625
&per_token_group_quant_fp8);
626
627
+ // Compute per-token-group INT8 quantized tensor and scaling factor.
628
+ ops.def(
629
+ "per_token_group_quant_int8(Tensor input, Tensor! output_q, Tensor! "
630
+ "output_s, int group_size, float eps, float int8_min, float int8_max) -> "
631
+ "()");
632
+ ops.impl("per_token_group_quant_int8", torch::kCUDA,
633
+ &per_token_group_quant_int8);
634
+
635
// reorder weight for AllSpark Ampere W8A16 Fused Gemm kernel
636
ops.def(
637
"rearrange_kn_weight_as_n32k16_order(Tensor b_qweight, Tensor b_scales, "
vllm/model_executor/layers/quantization/utils/int8_utils.py
+9
-2
@@ -238,13 +238,20 @@ def per_token_group_quant_int8(
238
int8_min = iinfo.min
239
240
x_q = torch.empty_like(x, device=x.device, dtype=dtype)
241
- M = x.numel() // group_size
242
- N = group_size
241
x_s = torch.empty(
242
x.shape[:-1] + (x.shape[-1] // group_size, ),
243
device=x.device,
244
dtype=torch.float32,
245
)
246
+ # prefer CUDA kernel if available
247
+ if current_platform.is_cuda():
248
+ torch.ops._C.per_token_group_quant_int8(x, x_q, x_s, group_size, eps,
249
+ float(int8_min),
250
+ float(int8_max))
251
+ return x_q, x_s
252
+
253
+ M = x.numel() // group_size
254
+ N = group_size
255
256
BLOCK = triton.next_power_of_2(N)
257
# heuristics for number of warps