@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
h 192 lines 6.25 KB
Raw
1
2 #include <cuda_fp16.h>
3 #include <cuda_bf16.h>
4
5 template <typename scalar_t>
6 class ScalarType {};
7
8 template <>
9 class ScalarType<half> {
10 public:
11 using scalar_t = half;
12 using scalar_t2 = half2;
13
14 static __device__ float inline num2float(const half x) {
15 return __half2float(x);
16 }
17
18 static __device__ half2 inline num2num2(const half x) {
19 return __half2half2(x);
20 }
21
22 static __device__ half2 inline nums2num2(const half x1, const half x2) {
23 return __halves2half2(x1, x2);
24 }
25
26 static __host__ __device__ half inline float2num(const float x) {
27 return __float2half(x);
28 }
29
30 static __host__ __device__ half inline int2num(const float x) {
31 return __int2half_rn(x);
32 }
33
34 static __host__ __device__ float2 inline num22float2(const half2 x) {
35 return __half22float2(x);
36 }
37 };
38
39 template <>
40 class ScalarType<nv_bfloat16> {
41 public:
42 using scalar_t = nv_bfloat16;
43 using scalar_t2 = nv_bfloat162;
44
45 #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800
46 static __device__ float inline num2float(const nv_bfloat16 x) {
47 return __bfloat162float(x);
48 }
49
50 static __device__ nv_bfloat162 inline num2num2(const nv_bfloat16 x) {
51 return __bfloat162bfloat162(x);
52 }
53
54 static __device__ nv_bfloat162 inline nums2num2(const nv_bfloat16 x1,
55 const nv_bfloat16 x2) {
56 return __halves2bfloat162(x1, x2);
57 }
58
59 static __host__ __device__ nv_bfloat16 inline float2num(const float x) {
60 return __float2bfloat16(x);
61 }
62
63 static __host__ __device__ nv_bfloat16 inline int2num(const float x) {
64 return __int2bfloat16_rn(x);
65 }
66
67 static __host__ __device__ float2 inline num22float2(const nv_bfloat162 x) {
68 return __bfloat1622float2(x);
69 }
70 #endif
71 };
72
73 template <int lut>
74 __device__ inline int lop3(int a, int b, int c) {
75 int res;
76 asm volatile("lop3.b32 %0, %1, %2, %3, %4;\n"
77 : "=r"(res)
78 : "r"(a), "r"(b), "r"(c), "n"(lut));
79 return res;
80 }
81
82 template <int start_byte, int mask>
83 __device__ inline uint32_t prmt(uint32_t a) {
84 uint32_t res;
85 asm volatile("prmt.b32 %0, %1, %2, %3;\n"
86 : "=r"(res)
87 : "r"(a), "n"(start_byte), "n"(mask));
88 return res;
89 }
90
91 template <typename scalar_t2, int bit>
92 __device__ inline void dequant(int q, scalar_t2* res) {}
93
94 template <>
95 __device__ inline void dequant<half2, 4>(int q, half2* res) {
96 const int LO = 0x000f000f;
97 const int HI = 0x00f000f0;
98 const int EX = 0x64006400;
99 const int SUB = 0x64006400;
100 const int MUL = 0x2c002c00;
101 const int ADD = 0xd400d400;
102
103 int lo0 = lop3<(0xf0 & 0xcc) | 0xaa>(q, LO, EX);
104 int hi0 = lop3<(0xf0 & 0xcc) | 0xaa>(q, HI, EX);
105 q >>= 8;
106 int lo1 = lop3<(0xf0 & 0xcc) | 0xaa>(q, LO, EX);
107 int hi1 = lop3<(0xf0 & 0xcc) | 0xaa>(q, HI, EX);
108
109 res[0] = __hsub2(*reinterpret_cast<half2*>(&lo0),
110 *reinterpret_cast<const half2*>(&SUB));
111 res[1] = __hfma2(*reinterpret_cast<half2*>(&hi0),
112 *reinterpret_cast<const half2*>(&MUL),
113 *reinterpret_cast<const half2*>(&ADD));
114 res[2] = __hsub2(*reinterpret_cast<half2*>(&lo1),
115 *reinterpret_cast<const half2*>(&SUB));
116 res[3] = __hfma2(*reinterpret_cast<half2*>(&hi1),
117 *reinterpret_cast<const half2*>(&MUL),
118 *reinterpret_cast<const half2*>(&ADD));
119 }
120
121 template <>
122 __device__ inline void dequant<half2, 8>(int q, half2* res) {
123 static constexpr uint32_t mask_for_elt_01 = 0x5250;
124 static constexpr uint32_t mask_for_elt_23 = 0x5351;
125 static constexpr uint32_t start_byte_for_fp16 = 0x64646464;
126
127 uint32_t lo = prmt<start_byte_for_fp16, mask_for_elt_01>(q);
128 uint32_t hi = prmt<start_byte_for_fp16, mask_for_elt_23>(q);
129
130 static constexpr uint32_t I8s_TO_F16s_MAGIC_NUM = 0x64006400;
131
132 res[0] = __hsub2(*reinterpret_cast<half2*>(&lo),
133 *reinterpret_cast<const half2*>(&I8s_TO_F16s_MAGIC_NUM));
134 res[1] = __hsub2(*reinterpret_cast<half2*>(&hi),
135 *reinterpret_cast<const half2*>(&I8s_TO_F16s_MAGIC_NUM));
136 }
137
138 #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800
139 template <>
140 __device__ inline void dequant<nv_bfloat162, 4>(int q, nv_bfloat162* res) {
141 static constexpr uint32_t MASK = 0x000f000f;
142 static constexpr uint32_t EX = 0x43004300;
143
144 int lo0 = lop3<(0xf0 & 0xcc) | 0xaa>(q, MASK, EX);
145 q >>= 4;
146 int hi0 = lop3<(0xf0 & 0xcc) | 0xaa>(q, MASK, EX);
147 q >>= 4;
148 int lo1 = lop3<(0xf0 & 0xcc) | 0xaa>(q, MASK, EX);
149 q >>= 4;
150 int hi1 = lop3<(0xf0 & 0xcc) | 0xaa>(q, MASK, EX);
151
152 static constexpr uint32_t MUL = 0x3F803F80;
153 static constexpr uint32_t ADD = 0xC300C300;
154
155 res[0] = __hfma2(*reinterpret_cast<nv_bfloat162*>(&lo0),
156 *reinterpret_cast<const nv_bfloat162*>(&MUL),
157 *reinterpret_cast<const nv_bfloat162*>(&ADD));
158 res[1] = __hfma2(*reinterpret_cast<nv_bfloat162*>(&hi0),
159 *reinterpret_cast<const nv_bfloat162*>(&MUL),
160 *reinterpret_cast<const nv_bfloat162*>(&ADD));
161 res[2] = __hfma2(*reinterpret_cast<nv_bfloat162*>(&lo1),
162 *reinterpret_cast<const nv_bfloat162*>(&MUL),
163 *reinterpret_cast<const nv_bfloat162*>(&ADD));
164 res[3] = __hfma2(*reinterpret_cast<nv_bfloat162*>(&hi1),
165 *reinterpret_cast<const nv_bfloat162*>(&MUL),
166 *reinterpret_cast<const nv_bfloat162*>(&ADD));
167 }
168
169 template <>
170 __device__ inline void dequant<nv_bfloat162, 8>(int q, nv_bfloat162* res) {
171 float fp32_intermediates[4];
172 uint32_t* fp32_intermediates_casted =
173 reinterpret_cast<uint32_t*>(fp32_intermediates);
174
175 static constexpr uint32_t fp32_base = 0x4B000000;
176 fp32_intermediates_casted[0] = __byte_perm(q, fp32_base, 0x7650);
177 fp32_intermediates_casted[1] = __byte_perm(q, fp32_base, 0x7652);
178 fp32_intermediates_casted[2] = __byte_perm(q, fp32_base, 0x7651);
179 fp32_intermediates_casted[3] = __byte_perm(q, fp32_base, 0x7653);
180
181 fp32_intermediates[0] -= 8388608.f;
182 fp32_intermediates[1] -= 8388608.f;
183 fp32_intermediates[2] -= 8388608.f;
184 fp32_intermediates[3] -= 8388608.f;
185
186 uint32_t* bf16_result_ptr = reinterpret_cast<uint32_t*>(res);
187 bf16_result_ptr[0] = __byte_perm(fp32_intermediates_casted[0],
188 fp32_intermediates_casted[1], 0x7632);
189 bf16_result_ptr[1] = __byte_perm(fp32_intermediates_casted[2],
190 fp32_intermediates_casted[3], 0x7632);
191 }
192 #endif