@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
cuh 149 lines 4.95 KB
Raw
1 #ifndef _qdq_3_cuh
2 #define _qdq_3_cuh
3
4 #include "qdq_util.cuh"
5
6 namespace vllm {
7 namespace gptq {
8 // Permutation:
9 //
10 // v9997775 55333111 u8886664 44222000 (u, v lsb)
11 // vjjjhhhf ffdddbbb uiiiggge eecccaaa
12 // vtttrrrp ppnnnlll usssqqqo oommmkkk
13
14 __forceinline__ __device__ void shuffle_3bit_32(uint32_t* q, int stride) {
15 uint32_t qa = q[0 * stride];
16 uint32_t qb = q[1 * stride];
17 uint32_t qc = q[2 * stride];
18
19 // qa: aa999888 77766655 54443332 22111000
20 // qb: lkkkjjji iihhhggg fffeeedd dcccbbba
21 // qc: vvvuuutt tsssrrrq qqpppooo nnnmmmll
22
23 uint32_t qd = qc >> 26;
24 qc <<= 4;
25 qc |= qb >> 28;
26 qb <<= 2;
27 qb |= qa >> 30;
28
29 // qa: ..999888 77766655 54443332 22111000
30 // qb: ..jjjiii hhhgggff feeedddc ccbbbaaa
31 // qc: ..tttsss rrrqqqpp pooonnnm mmlllkkk
32 // qd: vvvuuu
33
34 uint32_t za = 0;
35 uint32_t zb = 0;
36 uint32_t zc = 0;
37
38 for (int i = 0; i < 5; i++) {
39 uint32_t t0 = qa & 0x07;
40 uint32_t t1 = (qa & 0x38) >> 3;
41 qa >>= 6;
42 za |= (t0 << (i * 3));
43 za |= (t1 << (i * 3 + 16));
44 }
45 for (int i = 0; i < 5; i++) {
46 uint32_t t0 = qb & 0x07;
47 uint32_t t1 = (qb & 0x38) >> 3;
48 qb >>= 6;
49 zb |= (t0 << (i * 3));
50 zb |= (t1 << (i * 3 + 16));
51 }
52 for (int i = 0; i < 5; i++) {
53 uint32_t t0 = qc & 0x07;
54 uint32_t t1 = (qc & 0x38) >> 3;
55 qc >>= 6;
56 zc |= (t0 << (i * 3));
57 zc |= (t1 << (i * 3 + 16));
58 }
59
60 // za: 9997775 55333111 8886664 44222000
61 // zb: jjjhhhf ffdddbbb iiiggge eecccaaa
62 // zc: tttrrrp ppnnnlll sssqqqo oommmkkk
63 // qd: vvvuuu
64
65 za |= ((qd & 0x01) >> 0) << 15;
66 zb |= ((qd & 0x02) >> 1) << 15;
67 zc |= ((qd & 0x04) >> 2) << 15;
68 za |= ((qd & 0x08) >> 3) << 31;
69 zb |= ((qd & 0x10) >> 4) << 31;
70 zc |= ((qd & 0x20) >> 5) << 31;
71
72 // za: v9997775 55333111 u8886664 44222000 (u, v lsb)
73 // zb: vjjjhhhf ffdddbbb uiiiggge eecccaaa
74 // zc: vtttrrrp ppnnnlll usssqqqo oommmkkk
75
76 q[0 * stride] = za;
77 q[1 * stride] = zb;
78 q[2 * stride] = zc;
79 }
80
81 __forceinline__ __device__ void dequant_3bit_32(const uint32_t q_0,
82 const uint32_t q_1,
83 const uint32_t q_2,
84 half2 (&dq)[16], int stride,
85 const uint32_t zero) {
86 const uint32_t c0 = 0x64006400;
87 const half y8_ = __float2half_rn(1.0f / 8.0f);
88 const half y64_ = __float2half_rn(1.0f / 64.0f);
89 const half2 y8 = __halves2half2(y8_, y8_);
90 const half2 y64 = __halves2half2(y64_, y64_);
91 const half_uint16 z1_(0xe400 | zero); // half(-1024.0f - zero);
92 const half z8_ = __hsub(__int2half_rn(-128), __int2half_rn(zero));
93 const half z64_ = __hsub(__int2half_rn(-16), __int2half_rn(zero));
94 const half2 z1 = __halves2half2(z1_.as_half, z1_.as_half);
95 const half2 z8 = __halves2half2(z8_, z8_);
96 const half2 z64 = __halves2half2(z64_, z64_);
97
98 uint32_t qa = q_0;
99 uint32_t qb = q_1;
100 uint32_t qc = q_2;
101
102 half2_uint32 q0((qa & 0x00070007) | c0); // half2(q[ 0], q[ 1]) + 1024
103 half2_uint32 q1((qa & 0x00380038) | c0); // half2(q[ 2], q[ 3]) * 8 + 1024
104 qa >>= 6;
105 half2_uint32 q2((qa & 0x00070007) | c0); // half2(q[ 4], q[ 5]) + 1024
106 half2_uint32 q3((qa & 0x00380038) | c0); // half2(q[ 6], q[ 7]) * 8 + 1024
107 half2_uint32 q4((qa & 0x01c001c0) | c0); // half2(q[ 8], q[ 9]) * 64 + 1024
108 qa >>= 9;
109 qa &= 0x00010001;
110 half2_uint32 q5((qb & 0x00070007) | c0); // half2(q[10], q[11]) + 1024
111 half2_uint32 q6((qb & 0x00380038) | c0); // half2(q[12], q[13]) * 8 + 1024
112 qb >>= 6;
113 half2_uint32 q7((qb & 0x00070007) | c0); // half2(q[14], q[15]) + 1024
114 half2_uint32 q8((qb & 0x00380038) | c0); // half2(q[16], q[17]) * 8 + 1024
115 half2_uint32 q9((qb & 0x01c001c0) | c0); // half2(q[18], q[19]) * 64 + 1024
116 qb >>= 8;
117 qb &= 0x00020002;
118 half2_uint32 q10((qc & 0x00070007) | c0); // half2(q[20], q[21]) + 1024
119 half2_uint32 q11((qc & 0x00380038) | c0); // half2(q[22], q[23]) * 8 + 1024
120 qc >>= 6;
121 half2_uint32 q12((qc & 0x00070007) | c0); // half2(q[24], q[25]) + 1024
122 half2_uint32 q13((qc & 0x00380038) | c0); // half2(q[26], q[27]) * 8 + 1024
123 half2_uint32 q14((qc & 0x01c001c0) | c0); // half2(q[28], q[29]) * 64 + 1024
124 qc >>= 7;
125 qc &= 0x00040004;
126 half2_uint32 q15((qa | qb | qc) | c0);
127
128 dq[0] = __hadd2(q0.as_half2, z1);
129 dq[1] = __hfma2(q1.as_half2, y8, z8);
130 dq[2] = __hadd2(q2.as_half2, z1);
131 dq[3] = __hfma2(q3.as_half2, y8, z8);
132 dq[4] = __hfma2(q4.as_half2, y64, z64);
133 dq[5] = __hadd2(q5.as_half2, z1);
134 dq[6] = __hfma2(q6.as_half2, y8, z8);
135 dq[7] = __hadd2(q7.as_half2, z1);
136 dq[8] = __hfma2(q8.as_half2, y8, z8);
137 dq[9] = __hfma2(q9.as_half2, y64, z64);
138 dq[10] = __hadd2(q10.as_half2, z1);
139 dq[11] = __hfma2(q11.as_half2, y8, z8);
140 dq[12] = __hadd2(q12.as_half2, z1);
141 dq[13] = __hfma2(q13.as_half2, y8, z8);
142 dq[14] = __hfma2(q14.as_half2, y64, z64);
143 dq[15] = __hadd2(q15.as_half2, z1);
144 }
145
146 } // namespace gptq
147 } // namespace vllm
148
149 #endif