@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
cuh 228 lines 7.23 KB
Raw
1 /*
2 Adapted from https://github.com/turboderp/exllamav2 and
3 https://github.com/turboderp/exllama
4 */
5
6 #ifndef _matrix_view_cuh
7 #define _matrix_view_cuh
8
9 #include <cuda_runtime.h>
10 #include <cuda_fp16.h>
11
12 #include "qdq_util.cuh"
13
14 namespace vllm {
15 namespace gptq {
16
17 class MatrixView_half {
18 public:
19 const half* data;
20 const int height;
21 const int width;
22
23 __device__ __forceinline__ MatrixView_half(const half* data, const int height,
24 const int width)
25 : data(data), height(height), width(width) {}
26
27 __device__ __forceinline__ half item(int row, int column) const {
28 return data[row * width + column];
29 }
30 __device__ __forceinline__ const half* item_ptr(int row, int column) const {
31 return &data[row * width + column];
32 }
33
34 __device__ __forceinline__ void item4(half (&items)[4], int row,
35 int column) const {
36 half2* ptr = (half2*)item_ptr(row, column);
37 half2 i01 = ptr[0];
38 half2 i23 = ptr[1];
39 items[0] = __low2half(i01);
40 items[1] = __high2half(i01);
41 items[2] = __low2half(i23);
42 items[3] = __high2half(i23);
43 }
44 __device__ __forceinline__ void item4_f(float (&items)[4], int row,
45 int column) const {
46 half2* ptr = (half2*)item_ptr(row, column);
47 half2 i01 = ptr[0];
48 half2 i23 = ptr[1];
49 items[0] = __half2float(__low2half(i01));
50 items[1] = __half2float(__high2half(i01));
51 items[2] = __half2float(__low2half(i23));
52 items[3] = __half2float(__high2half(i23));
53 }
54
55 __device__ __forceinline__ void item4_h2(half2 (&items)[4], int row,
56 int column) const {
57 half2* ptr = (half2*)item_ptr(row, column);
58 half2 i01 = ptr[0];
59 half2 i23 = ptr[1];
60 items[0] = __half2half2(__low2half(i01));
61 items[1] = __half2half2(__high2half(i01));
62 items[2] = __half2half2(__low2half(i23));
63 items[3] = __half2half2(__high2half(i23));
64 }
65 };
66
67 class MatrixView_half_rw {
68 public:
69 half* data;
70 const int height;
71 const int width;
72
73 __device__ __forceinline__ MatrixView_half_rw(half* data, const int height,
74 const int width)
75 : data(data), height(height), width(width) {}
76
77 __device__ __forceinline__ half* item_ptr(int row, int column) {
78 return &data[row * width + column];
79 }
80
81 __device__ __forceinline__ void set4(int row, int column, half v0, half v1,
82 half v2, half v3) {
83 half2 v01 = __halves2half2(v0, v1);
84 half2 v23 = __halves2half2(v2, v3);
85 half2* ptr = (half2*)item_ptr(row, column);
86 ptr[0] = v01;
87 ptr[1] = v23;
88 }
89 };
90
91 class MatrixView_q4_row {
92 public:
93 const uint32_t* data;
94 const int height;
95 const int width;
96
97 __device__ __forceinline__ MatrixView_q4_row(const uint32_t* data,
98 const int height,
99 const int width)
100 : data(data), height(height), width(width) {}
101
102 __device__ __forceinline__ int item(int row, int column) const {
103 int shift = (column & 0x07) * 4;
104 return (data[row * width / 8 + column / 8] >> shift) & 0x0f;
105 }
106
107 __device__ __forceinline__ void item4(int (&items)[4], int row,
108 int column) const {
109 int shift = (column & 0x07) * 4;
110 uint32_t d = data[row * width / 8 + column / 8] >> shift;
111 items[0] = d & 0x0f;
112 items[1] = (d >> 4) & 0x0f;
113 items[2] = (d >> 8) & 0x0f;
114 items[3] = (d >> 12) & 0x0f;
115 }
116 };
117
118 class MatrixView_q2_row {
119 public:
120 const uint32_t* data;
121 const int height;
122 const int width;
123
124 __device__ __forceinline__ MatrixView_q2_row(const uint32_t* data,
125 const int height,
126 const int width)
127 : data(data), height(height), width(width) {}
128
129 __device__ __forceinline__ int item(int row, int column) const {
130 int shift = (column & 0x0f) * 2;
131 return (data[row * width / 16 + column / 16] >> shift) & 0x03;
132 }
133
134 __device__ __forceinline__ void item4(int (&items)[4], int row,
135 int column) const {
136 int shift = (column & 0x0f) * 2;
137 uint32_t d = data[row * width / 16 + column / 16] >> shift;
138 items[0] = d & 0x03;
139 items[1] = (d >> 2) & 0x03;
140 items[2] = (d >> 4) & 0x03;
141 items[3] = (d >> 6) & 0x03;
142 }
143 };
144
145 class MatrixView_q3_row {
146 public:
147 const uint32_t* data;
148 const int height;
149 const int width;
150
151 __device__ __forceinline__ MatrixView_q3_row(const uint32_t* data,
152 const int height,
153 const int width)
154 : data(data), height(height), width(width) {}
155
156 __device__ __forceinline__ int item(int row, int column) const {
157 int z_w = column * 3 / 32;
158 int z_mod = column & 0x1f;
159
160 if (z_mod == 10) {
161 return (data[row * width * 3 / 32 + z_w] >> 30) |
162 ((data[row * width * 3 / 32 + (z_w + 1)] << 2) & 0x4);
163 } else if (z_mod == 21) {
164 return (data[row * width * 3 / 32 + z_w] >> 31) |
165 ((data[row * width * 3 / 32 + (z_w + 1)] << 1) & 0x6);
166 } else if (z_mod < 10) {
167 return (data[row * width * 3 / 32 + z_w] >> (z_mod * 3)) & 0x07;
168 } else if (z_mod < 21) {
169 return (data[row * width * 3 / 32 + z_w] >> (z_mod * 3 - 32)) & 0x07;
170 } else {
171 return (data[row * width * 3 / 32 + z_w] >> (z_mod * 3 - 64)) & 0x07;
172 }
173 }
174
175 __device__ __forceinline__ void item4(int (&items)[4], int row,
176 int column) const {
177 int shift = (column & 0x1f);
178 uint32_t d;
179 if (shift <= 4) {
180 d = data[row * width / 32 * 3 + column * 3 / 32] >> (shift * 3);
181 } else if (shift == 8) {
182 d = (data[row * width / 32 * 3 + column * 3 / 32] >> 24) |
183 ((data[row * width / 32 * 3 + column * 3 / 32 + 1] & 0x0f) << 8);
184 } else if (shift <= 16) {
185 d = data[row * width / 32 * 3 + column * 3 / 32] >> (shift * 3 - 32);
186 } else if (shift == 20) {
187 d = (data[row * width / 32 * 3 + column * 3 / 32] >> 28) |
188 ((data[row * width / 32 * 3 + column * 3 / 32 + 1] & 0xff) << 4);
189 } else {
190 d = data[row * width / 32 * 3 + column * 3 / 32] >> (shift * 3 - 64);
191 }
192 items[0] = d & 0x07;
193 items[1] = (d >> 3) & 0x07;
194 items[2] = (d >> 6) & 0x07;
195 items[3] = (d >> 9) & 0x07;
196 }
197 };
198
199 class MatrixView_q8_row {
200 public:
201 const uint32_t* data;
202 const int height;
203 const int width;
204
205 __device__ __forceinline__ MatrixView_q8_row(const uint32_t* data,
206 const int height,
207 const int width)
208 : data(data), height(height), width(width) {}
209
210 __device__ __forceinline__ int item(int row, int column) const {
211 int shift = (column & 0x03) * 8;
212 return (data[row * width / 4 + column / 4] >> shift) & 0xff;
213 }
214
215 __device__ __forceinline__ void item4(int (&items)[4], int row,
216 int column) const {
217 int shift = (column & 0x03) * 2;
218 uint32_t d = data[row * width / 4 + column / 4] >> shift;
219 items[0] = d & 0xff;
220 items[1] = (d >> 8) & 0xff;
221 items[2] = (d >> 16) & 0xff;
222 items[3] = (d >> 24) & 0xff;
223 }
224 };
225
226 } // namespace gptq
227 } // namespace vllm
228 #endif