@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
cuh 61 lines 1.88 KB
Raw
1 #pragma once
2
3 /**
4 * Quantization utilities including:
5 * Adjusted maximum values for qtypes.
6 * Minimum scaling factors for qtypes.
7 */
8
9 #include <cmath>
10 #include <torch/headeronly/macros/Macros.h>
11
12 #ifndef USE_ROCM
13 #include <torch/headeronly/util/Float8_e4m3fn.h>
14 #define MAYBE_HOST_DEVICE C10_HOST_DEVICE
15 #else
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,
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 };
30
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<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
41 template <typename T>
42 MAYBE_HOST_DEVICE static constexpr T quant_type_max_v =
43 quant_type_max<T>::val();
44
45 template <typename 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);
53 }
54 };
55
56 template <>
57 struct min_scaling_factor<int8_t> {
58 C10_DEVICE C10_ALWAYS_INLINE static float val() {
59 return std::numeric_limits<float>::epsilon();
60 }
61 };