| 1 | // SPDX-License-Identifier: Apache-2.0 |
| 2 | // SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | // |
| 4 | // W4A16 dequant primitives for RDNA3 (gfx1100/gfx1101/gfx1102), templated on |
| 5 | // the activation/scale dtype (half or __hip_bfloat16). The fp16 path reuses |
| 6 | // the classic exllamav2 bit-trick: |
| 7 | // |
| 8 | // (qa & 0x000F000F) | 0x64006400 -> half2(1024+q_lo, 1024+q_hi) |
| 9 | // (qa & 0x00F000F0) | 0x64006400 -> half2(1024+q_lo*16, 1024+q_hi*16) |
| 10 | // |
| 11 | // The "*16 then divide by 16 in the FMA" trick for the upper-nibble pairs |
| 12 | // works in fp16 because the mantissa (10 bits) is wide enough to hold a value |
| 13 | // shifted by 4 bits. In bf16 the mantissa is only 7 bits, so shifting an upper |
| 14 | // nibble into bits [7:4] would spill into the exponent. To avoid that, the |
| 15 | // bf16 path shifts each pair of nibbles down to bits [3:0]/[19:16] with a |
| 16 | // single right-shift before the OR with 0x43004300 (= bf162(128, 128)). |
| 17 | |
| 18 | #ifndef _qdq_4_rdna3_cuh |
| 19 | #define _qdq_4_rdna3_cuh |
| 20 | |
| 21 | #include <cstdint> |
| 22 | |
| 23 | #include <hip/hip_bf16.h> |
| 24 | #include <hip/hip_fp16.h> |
| 25 | |
| 26 | namespace vllm { |
| 27 | namespace gptq_rdna3 { |
| 28 | |
| 29 | using bf16_t = __hip_bfloat16; |
| 30 | using bf162_t = __hip_bfloat162; |
| 31 | |
| 32 | // Bit-shuffle for an int32 holding 8 sequential 4-bit weights q[0..7]: |
| 33 | // in: q[7] q[6] q[5] q[4] q[3] q[2] q[1] q[0] (LSB first) |
| 34 | // out: q[7] q[5] q[3] q[1] q[6] q[4] q[2] q[0] (even/odd interleaved) |
| 35 | // |
| 36 | // After shuffle, q[2k] sits at bits [4k : 4k+3] (lower 16) |
| 37 | // q[2k+1] sits at bits [16+4k: 16+4k+3] (upper 16) |
| 38 | // so a single mask 0x000F000F selects the matching even/odd pair, ready to |
| 39 | // bitcast to half2 / bfloat162 after OR-ing with the magic constant. |
| 40 | __forceinline__ __device__ void shuffle_4bit_8(uint32_t* q) { |
| 41 | uint32_t qa = q[0]; |
| 42 | uint32_t qb = 0; |
| 43 | #pragma unroll |
| 44 | for (int i = 0; i < 4; i++) { |
| 45 | uint32_t qa0 = qa & 0x0F; |
| 46 | uint32_t qa1 = (qa & 0xF0) >> 4; |
| 47 | qa >>= 8; |
| 48 | qb |= (qa1 << (i * 4 + 16)); |
| 49 | qb |= (qa0 << (i * 4)); |
| 50 | } |
| 51 | q[0] = qb; |
| 52 | } |
| 53 | |
| 54 | // --------------------------------------------------------------------------- |
| 55 | // fp16 path |
| 56 | // --------------------------------------------------------------------------- |
| 57 | |
| 58 | // Precompute scale-baked constants for a single zero/scale pair. |
| 59 | // z1z16[0] = scale * (-1024 - zero) (used for "low" pairs) |
| 60 | // z1z16[1] = scale * (-64 - zero) (used for "high" pairs) |
| 61 | // y1y16[0] = scale * 1 (low pairs are q + 1024) |
| 62 | // y1y16[1] = scale * (1/16) (high pairs are q*16 + 1024) |
| 63 | __forceinline__ __device__ void prep_zero_scale_fp16(uint32_t zero, half scale, |
| 64 | half2 (&z1z16)[2], |
| 65 | half2 (&y1y16)[2]) { |
| 66 | // half(-1024 - zero) via the exllamav2 bit-trick: |
| 67 | // half bits 0xE400 == -1024.0 ; ORing the zero into mantissa subtracts it. |
| 68 | union { |
| 69 | uint16_t u; |
| 70 | half h; |
| 71 | } z1u; |
| 72 | z1u.u = (uint16_t)(0xE400 | zero); |
| 73 | half z1 = z1u.h; |
| 74 | half z16 = __hsub(__int2half_rn(-64), __int2half_rn((int)zero)); |
| 75 | |
| 76 | half2 scale2 = __half2half2(scale); |
| 77 | z1z16[0] = __hmul2(scale2, __half2half2(z1)); |
| 78 | z1z16[1] = __hmul2(scale2, __half2half2(z16)); |
| 79 | |
| 80 | half y1 = __float2half_rn(1.0f); |
| 81 | half y16 = __float2half_rn(1.0f / 16.0f); |
| 82 | y1y16[0] = __hmul2(scale2, __half2half2(y1)); |
| 83 | y1y16[1] = __hmul2(scale2, __half2half2(y16)); |
| 84 | } |
| 85 | |
| 86 | // Dequantize one int32 (8 shuffled 4-bit weights) into 4 half2 pairs: |
| 87 | // dq[0] = (q[0], q[1]) * scale - zero*scale |
| 88 | // dq[1] = (q[2], q[3]) * scale - zero*scale |
| 89 | // dq[2] = (q[4], q[5]) * scale - zero*scale |
| 90 | // dq[3] = (q[6], q[7]) * scale - zero*scale |
| 91 | __forceinline__ __device__ void dequant_4bit_8_fp16(uint32_t qa, half2 (&dq)[4], |
| 92 | half2 (&z1z16)[2], |
| 93 | half2 (&y1y16)[2]) { |
| 94 | const uint32_t c0 = 0x64006400; |
| 95 | |
| 96 | union { |
| 97 | uint32_t u; |
| 98 | half2 h2; |
| 99 | } q0, q1, q2, q3; |
| 100 | q0.u = (qa & 0x000F000F) | c0; // half2(q[0]+1024, q[1]+1024) |
| 101 | q1.u = (qa & 0x00F000F0) | c0; // half2(q[2]*16+1024, q[3]*16+1024) |
| 102 | uint32_t qa_hi = qa >> 8; |
| 103 | q2.u = (qa_hi & 0x000F000F) | c0; // half2(q[4]+1024, q[5]+1024) |
| 104 | q3.u = (qa_hi & 0x00F000F0) | c0; // half2(q[6]*16+1024, q[7]*16+1024) |
| 105 | |
| 106 | dq[0] = __hfma2(q0.h2, y1y16[0], z1z16[0]); |
| 107 | dq[1] = __hfma2(q1.h2, y1y16[1], z1z16[1]); |
| 108 | dq[2] = __hfma2(q2.h2, y1y16[0], z1z16[0]); |
| 109 | dq[3] = __hfma2(q3.h2, y1y16[1], z1z16[1]); |
| 110 | } |
| 111 | |
| 112 | // --------------------------------------------------------------------------- |
| 113 | // bf16 path |
| 114 | // --------------------------------------------------------------------------- |
| 115 | |
| 116 | __forceinline__ __device__ void prep_zero_scale_bf16_f32(uint32_t zero, |
| 117 | bf16_t scale, |
| 118 | float& z_prep, |
| 119 | float& y_prep) { |
| 120 | float scale_f = __bfloat162float(scale); |
| 121 | z_prep = -(128.0f + (float)zero) * scale_f; |
| 122 | y_prep = scale_f; |
| 123 | } |
| 124 | |
| 125 | } // namespace gptq_rdna3 |
| 126 | } // namespace vllm |
| 127 | |
| 128 | #endif // _qdq_4_rdna3_cuh |