| 1 | /* |
| 2 | Copied from https://github.com/turboderp/exllamav2 |
| 3 | */ |
| 4 | |
| 5 | #ifndef _qdq_4_cuh |
| 6 | #define _qdq_4_cuh |
| 7 | |
| 8 | #include "qdq_util.cuh" |
| 9 | |
| 10 | namespace vllm { |
| 11 | namespace gptq { |
| 12 | // Permutation: |
| 13 | // |
| 14 | // 77775555 33331111 66664444 22220000 |
| 15 | |
| 16 | __forceinline__ __device__ void shuffle_4bit_8(uint32_t* q, int stride) { |
| 17 | uint32_t qa = q[0]; |
| 18 | uint32_t qb = 0; |
| 19 | |
| 20 | #pragma unroll |
| 21 | for (int i = 0; i < 4; i++) { |
| 22 | uint32_t qa0 = qa & 0x0f; |
| 23 | uint32_t qa1 = (qa & 0xf0) >> 4; |
| 24 | qa >>= 8; |
| 25 | qb |= (qa1 << (i * 4 + 16)); |
| 26 | qb |= (qa0 << (i * 4)); |
| 27 | } |
| 28 | q[0] = qb; |
| 29 | } |
| 30 | |
| 31 | __forceinline__ __device__ void dequant_4bit_8_prep_zero(const uint32_t zero, |
| 32 | half2 (&z1z16)[2], |
| 33 | half2 (&y1y16)[2]) { |
| 34 | half_uint16 z1(0xe400 | zero); // half(-1024.0f - zero); |
| 35 | half z16 = __hsub(__int2half_rn(-64), __int2half_rn(zero)); |
| 36 | |
| 37 | z1z16[0] = __half2half2(z1.as_half); |
| 38 | z1z16[1] = __half2half2(z16); |
| 39 | |
| 40 | const half y1 = __float2half_rn(1.0f); |
| 41 | const half y16 = __float2half_rn(1.0f / 16.0f); |
| 42 | |
| 43 | y1y16[0] = __half2half2(y1); |
| 44 | y1y16[1] = __half2half2(y16); |
| 45 | } |
| 46 | |
| 47 | __forceinline__ __device__ void dequant_4bit_8_gptq(const uint32_t q_0, |
| 48 | half2 (&dq)[4], |
| 49 | half2 (&z1z16)[2], |
| 50 | half2 (&y1y16)[2], |
| 51 | int stride, bool scaled) { |
| 52 | const uint32_t c0 = 0x64006400; |
| 53 | |
| 54 | uint32_t qa = q_0; |
| 55 | half2_uint32 q0((qa & 0x000f000f) | |
| 56 | c0); // half2( q[0] + 1024, q[1] + 1024 ) |
| 57 | half2_uint32 q1((qa & 0x00f000f0) | |
| 58 | c0); // half2( q[2] * 16 + 1024, q[3] * 16 + 1024 ) |
| 59 | qa >>= 8; |
| 60 | half2_uint32 q2((qa & 0x000f000f) | |
| 61 | c0); // half2( q[4] + 1024, q[5] + 1024 ) |
| 62 | half2_uint32 q3((qa & 0x00f000f0) | |
| 63 | c0); // half2( q[6] * 16 + 1024, q[7] * 16 + 1024 ) |
| 64 | |
| 65 | if (scaled) { |
| 66 | dq[0] = __hfma2(q0.as_half2, y1y16[0], |
| 67 | z1z16[0]); // half2( q[0] * s - z * s, q[1] * s - z * s) |
| 68 | dq[1] = __hfma2(q1.as_half2, y1y16[1], |
| 69 | z1z16[1]); // half2( q[2] * s - z * s, q[3] * s - z * s) |
| 70 | dq[2] = __hfma2(q2.as_half2, y1y16[0], z1z16[0]); |
| 71 | dq[3] = __hfma2(q3.as_half2, y1y16[1], z1z16[1]); |
| 72 | } else { |
| 73 | dq[0] = __hadd2(q0.as_half2, z1z16[0]); // half2( q[0] - z, q[1] - z ) |
| 74 | dq[1] = __hfma2(q1.as_half2, y1y16[1], |
| 75 | z1z16[1]); // half2( q[2] - z, q[3] - z ) |
| 76 | dq[2] = __hadd2(q2.as_half2, z1z16[0]); // half2( q[4] - z, q[5] - z ) |
| 77 | dq[3] = __hfma2(q3.as_half2, y1y16[1], |
| 78 | z1z16[1]); // half2( q[6] - z, q[7] - z ) |
| 79 | } |
| 80 | } |
| 81 | } // namespace gptq |
| 82 | } // namespace vllm |
| 83 | |
| 84 | #endif |