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, 22:35 UTC
[Perf] Using `__nv_fp8_e4m3` instead of `c10::e4m3` for `per_token_group_quant` (#21867)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
Wentao Ye committed
Jul 29, 2025 at 23:50 UTC
1b0a15553420e5459d9a8512a3f1bd7d4117db08
1 file changed
+2
-4
csrc/quantization/fp8/per_token_group_quant.cu
+2
-4
@@ -1,12 +1,10 @@
1
#include <ATen/cuda/CUDAContext.h>
2
-#include <c10/util/Float8_e4m3fn.h>
2
3
#include "../per_token_group_quant_8bit.h"
4
5
#include <cmath>
6
8
-#include <cuda_fp16.h>
9
-#include <cuda_bf16.h>
7
+#include <cuda_fp8.h>
8
9
#include <torch/all.h>
10
@@ -199,7 +197,7 @@ void per_token_group_quant_8bit(const torch::Tensor& input,
197
VLLM_DISPATCH_FLOATING_TYPES(
198
input.scalar_type(), "per_token_group_quant_8bit", ([&] {
199
if (dst_type == at::ScalarType::Float8_e4m3fn) {
202
- LAUNCH_KERNEL(scalar_t, c10::Float8_e4m3fn);
200
+ LAUNCH_KERNEL(scalar_t, __nv_fp8_e4m3);
201
} else if (dst_type == at::ScalarType::Char) {
202
LAUNCH_KERNEL(scalar_t, int8_t);
203
}