@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
cuh 102 lines 4.15 KB
Raw
1 /*
2 Adapted from https://github.com/mit-han-lab/llm-awq
3 Modified from NVIDIA FasterTransformer:
4 https://github.com/NVIDIA/FasterTransformer/blob/main/src/fastertransformer/cutlass_extensions/include/cutlass_extensions/interleaved_numeric_conversion.h
5 @article{lin2023awq,
6 title={AWQ: Activation-aware Weight Quantization for LLM Compression and
7 Acceleration}, author={Lin, Ji and Tang, Jiaming and Tang, Haotian and Yang,
8 Shang and Dang, Xingyu and Han, Song}, journal={arXiv}, year={2023}
9 }
10 */
11
12 #pragma once
13
14 namespace vllm {
15 namespace awq {
16
17 __device__ uint4 dequantize_s4_to_fp16x2(uint32_t const& source) {
18 #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 750
19 assert(false);
20 #else
21 uint4 result;
22
23 uint32_t* h = reinterpret_cast<uint32_t*>(&result);
24 uint32_t const i4s = reinterpret_cast<uint32_t const&>(source);
25
26 // First, we extract the i4s and construct an intermediate fp16 number.
27 static constexpr uint32_t immLut = (0xf0 & 0xcc) | 0xaa;
28 static constexpr uint32_t BOTTOM_MASK = 0x000f000f;
29 static constexpr uint32_t TOP_MASK = 0x00f000f0;
30 static constexpr uint32_t I4s_TO_F16s_MAGIC_NUM = 0x64006400;
31
32 // Note that the entire sequence only requires 1 shift instruction. This is
33 // thanks to the register packing format and the fact that we force our
34 // integers to be unsigned, and account for this in the fp16 subtractions. In
35 // addition, I exploit the fact that sub and fma have the same throughput in
36 // order to convert elt_23 and elt_67 to fp16 without having to shift them to
37 // the bottom bits before hand.
38
39 // Shift right by 8 to now consider elt_45 and elt_67. Issue first to hide RAW
40 // dependency if we issue immediately before required.
41 const uint32_t top_i4s = i4s >> 8;
42 // Extract elt_01 - (i4s & 0x000f000f) | 0x64006400
43 asm volatile("lop3.b32 %0, %1, %2, %3, %4;\n"
44 : "=r"(h[0])
45 : "r"(i4s), "n"(BOTTOM_MASK), "n"(I4s_TO_F16s_MAGIC_NUM),
46 "n"(immLut));
47 // Extract elt_23 (i4s & 0x00f000f0) | 0x64006400
48 asm volatile("lop3.b32 %0, %1, %2, %3, %4;\n"
49 : "=r"(h[1])
50 : "r"(i4s), "n"(TOP_MASK), "n"(I4s_TO_F16s_MAGIC_NUM),
51 "n"(immLut));
52 // Extract elt_45 (top_i4s & 0x000f000f) | 0x64006400
53 asm volatile("lop3.b32 %0, %1, %2, %3, %4;\n"
54 : "=r"(h[2])
55 : "r"(top_i4s), "n"(BOTTOM_MASK), "n"(I4s_TO_F16s_MAGIC_NUM),
56 "n"(immLut));
57 // Extract elt_67 (top_i4s & 0x00f000f0) | 0x64006400
58 asm volatile("lop3.b32 %0, %1, %2, %3, %4;\n"
59 : "=r"(h[3])
60 : "r"(top_i4s), "n"(TOP_MASK), "n"(I4s_TO_F16s_MAGIC_NUM),
61 "n"(immLut));
62
63 // I use inline PTX below because I am not sure if the compiler will emit
64 // float2half instructions if I use the half2 ctor. In this case, I chose
65 // performance reliability over code readability.
66
67 // This is the half2 {1032, 1032} represented as an integer.
68 // static constexpr uint32_t FP16_TOP_MAGIC_NUM = 0x64086408;
69 // Haotian: subtract {1024, 1024} instead, we do not need to map to [-8, 7]
70 static constexpr uint32_t FP16_TOP_MAGIC_NUM = 0x64006400;
71 // This is the half2 {1 / 16, 1 / 16} represented as an integer.
72 static constexpr uint32_t ONE_SIXTEENTH = 0x2c002c00;
73 // This is the half2 {-72, -72} represented as an integer.
74 // static constexpr uint32_t NEG_72 = 0xd480d480;
75 // Haotian: Let's use {-64, -64}.
76 static constexpr uint32_t NEG_64 = 0xd400d400;
77
78 // Finally, we construct the output numbers.
79 // Convert elt_01
80 asm volatile("sub.f16x2 %0, %1, %2;\n"
81 : "=r"(h[0])
82 : "r"(h[0]), "r"(FP16_TOP_MAGIC_NUM));
83 // Convert elt_23
84 asm volatile("fma.rn.f16x2 %0, %1, %2, %3;\n"
85 : "=r"(h[1])
86 : "r"(h[1]), "r"(ONE_SIXTEENTH), "r"(NEG_64));
87 // Convert elt_45
88 asm volatile("sub.f16x2 %0, %1, %2;\n"
89 : "=r"(h[2])
90 : "r"(h[2]), "r"(FP16_TOP_MAGIC_NUM));
91 // Convert elt_67
92 asm volatile("fma.rn.f16x2 %0, %1, %2, %3;\n"
93 : "=r"(h[3])
94 : "r"(h[3]), "r"(ONE_SIXTEENTH), "r"(NEG_64));
95
96 return result;
97 #endif
98 __builtin_unreachable(); // Suppress missing return statement warning
99 }
100
101 } // namespace awq
102 } // namespace vllm