| 1 | /* |
| 2 | Copied from https://github.com/turboderp/exllamav2 |
| 3 | */ |
| 4 | |
| 5 | #ifndef _qdq_2_cuh |
| 6 | #define _qdq_2_cuh |
| 7 | |
| 8 | #include "qdq_util.cuh" |
| 9 | |
| 10 | namespace vllm { |
| 11 | namespace gptq { |
| 12 | |
| 13 | // Permutation: |
| 14 | // |
| 15 | // ffddbb99 77553311 eeccaa88 66442200 |
| 16 | |
| 17 | __forceinline__ __device__ void shuffle_2bit_16(uint32_t* q, int stride) { |
| 18 | uint32_t qa = q[0]; |
| 19 | uint32_t qb = 0; |
| 20 | |
| 21 | #pragma unroll |
| 22 | for (int i = 0; i < 8; i++) { |
| 23 | uint32_t qa0 = qa & 0x03; |
| 24 | uint32_t qa1 = (qa & 0x0c) >> 2; |
| 25 | qa >>= 4; |
| 26 | qb |= (qa1 << (i * 2 + 16)); |
| 27 | qb |= (qa0 << (i * 2)); |
| 28 | } |
| 29 | q[0] = qb; |
| 30 | } |
| 31 | |
| 32 | __forceinline__ __device__ void dequant_2bit_16(const uint32_t q_0, |
| 33 | half2 (&dq)[8], int stride, |
| 34 | const uint32_t zero) { |
| 35 | const uint32_t c0 = 0x64006400; |
| 36 | const half y4_ = __float2half_rn(1.0f / 4.0f); |
| 37 | const half y16_ = __float2half_rn(1.0f / 16.0f); |
| 38 | const half y64_ = __float2half_rn(1.0f / 64.0f); |
| 39 | const half2 y4 = __halves2half2(y4_, y4_); |
| 40 | const half2 y16 = __halves2half2(y16_, y16_); |
| 41 | const half2 y64 = __halves2half2(y64_, y64_); |
| 42 | |
| 43 | const half_uint16 z1_(0xe400 | zero); // half(-1024.0f - zero); |
| 44 | const half z4_ = __hsub(__int2half_rn(-256), __int2half_rn(zero)); |
| 45 | const half z16_ = __hsub(__int2half_rn(-64), __int2half_rn(zero)); |
| 46 | const half z64_ = __hsub(__int2half_rn(-16), __int2half_rn(zero)); |
| 47 | const half2 z1 = __half2half2(z1_.as_half); |
| 48 | const half2 z4 = __half2half2(z4_); |
| 49 | const half2 z16 = __half2half2(z16_); |
| 50 | const half2 z64 = __half2half2(z64_); |
| 51 | |
| 52 | uint32_t qa = q_0; |
| 53 | half2_uint32 q0((qa & 0x00030003) | c0); // half2(q[ 0], q[ 1]) + 1024 |
| 54 | half2_uint32 q1((qa & 0x000c000c) | c0); // half2(q[ 2], q[ 3]) * 4 + 1024 |
| 55 | half2_uint32 q2((qa & 0x00300030) | c0); // half2(q[ 4], q[ 5]) * 16 + 1024 |
| 56 | half2_uint32 q3((qa & 0x00c000c0) | c0); // half2(q[ 6], q[ 7]) * 64 + 1024 |
| 57 | qa >>= 8; |
| 58 | half2_uint32 q4((qa & 0x00030003) | c0); // half2(q[ 8], q[ 8]) + 1024 |
| 59 | half2_uint32 q5((qa & 0x000c000c) | c0); // half2(q[10], q[11]) * 4 + 1024 |
| 60 | half2_uint32 q6((qa & 0x00300030) | c0); // half2(q[12], q[13]) * 16 + 1024 |
| 61 | half2_uint32 q7((qa & 0x00c000c0) | c0); // half2(q[14], q[15]) * 64 + 1024 |
| 62 | |
| 63 | dq[0] = __hadd2(q0.as_half2, z1); |
| 64 | dq[1] = __hfma2(q1.as_half2, y4, z4); |
| 65 | dq[2] = __hfma2(q2.as_half2, y16, z16); |
| 66 | dq[3] = __hfma2(q3.as_half2, y64, z64); |
| 67 | dq[4] = __hadd2(q4.as_half2, z1); |
| 68 | dq[5] = __hfma2(q5.as_half2, y4, z4); |
| 69 | dq[6] = __hfma2(q6.as_half2, y16, z16); |
| 70 | dq[7] = __hfma2(q7.as_half2, y64, z64); |
| 71 | } |
| 72 | |
| 73 | } // namespace gptq |
| 74 | } // namespace vllm |
| 75 | |
| 76 | #endif |