@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
md 168 lines 5.55 KB
Rendered Raw
1 # CUTLASS Epilogues
2
3 ## Introduction
4
5 This document describes the various CUTLASS epilogues implemented for fusing de-quantization operations onto GEMMs.
6
7 Currently, we only support symmetric quantization for weights,
8 and symmetric and asymmetric quantization for activations.
9 Both can be quantized per-tensor or per-channel (weights) / per-token (activations).
10
11 There are 4 epilogues:
12
13 1. `ScaledEpilogue`: symmetric quantization for activations, no bias.
14 1. `ScaledEpilogueBias`: symmetric quantization for activations, supports bias.
15 1. `ScaledEpilogueAzp`: asymmetric per-tensor quantization for activations, supports bias.
16 1. `ScaledEpilogueAzpPerToken`: asymmetric per-token quantization for activations, supports bias.
17
18 We do not have epilogues for asymmetric quantization of activations without bias in order to reduce final binary size.
19 Instead, if no bias is passed, the epilogue will use 0 as the bias.
20 That induces a redundant addition operation (and runtime check), but the performance impact is minor.
21
22 ## Underlying Linear Algebra
23
24 More details available in the [Activation Quantization RFC](https://github.com/vllm-project/vllm/issues/3975).
25
26 If $` \widehat X `$ is the quantized $` X `$, our matrices become the following
27
28 ```math
29 A = s_a (\widehat A - J_a z_a)
30 ```
31
32 ```math
33 B = s_b \widehat B
34 ```
35
36 ```math
37 D = A B + C
38 ```
39
40 ```math
41 D = s_a s_b \widehat D + C
42 ```
43
44 Here, D is the output of the GEMM, and C is the bias.
45 A is the activations and supports asymmetric quantization,
46 and B is the weights and only supports symmetric quantization.
47 $ s_a $ and $s_b$ are the scales for activations and weights, respectively.
48 $ z_a $ is the zero-point for activations, and $ J_a $ is the matrix of all ones with dimensions of A.
49 Additional epilogues would be required to support asymmetric quantization for weights.
50
51 Expanding further, we can calculate $` \widehat D `$ as follows:
52
53 ```math
54 A B = s_a ( \widehat A - J_a z_a ) s_b \widehat B
55 ```
56
57 ```math
58 A B = s_a s_b \left( \widehat A \widehat B - J_a z_a \widehat B \right)
59 ```
60
61 ```math
62 \widehat D = \widehat A \widehat B - z_a J_a \widehat B
63 ```
64
65 Note that $` \widehat A \widehat B `$ is the raw output of the GEMM,
66 and $` J_a \widehat B `$ is known ahead of time.
67 Each row of it is equal to $` \mathbf 1 \widehat B `$, which is a row-vector of column sums of $` \widehat B `$.
68
69 ## Epilogues
70
71 ### `ScaledEpilogue`
72
73 This epilogue computes the symmetric quantization for activations without bias, meaning $` C = 0 `$ and $` z_a = 0 `$.
74 The output of the GEMM is:
75
76 ```math
77 \widehat D = \widehat A \widehat B
78 ```
79
80 ```math
81 D = s_a s_b \widehat D
82 ```
83
84 ```math
85 D = s_a s_b \widehat A \widehat B
86 ```
87
88 Epilogue parameters:
89
90 - `scale_a` is the scale for activations, can be per-tensor (scalar) or per-token (column-vector).
91 - `scale_b` is the scale for weights, can be per-tensor (scalar) or per-channel (row-vector).
92
93 ### `ScaledEpilogueBias`
94
95 This epilogue computes the symmetric quantization for activations with bias, meaning $` z_a = 0 `$.
96 The output of the GEMM is:
97
98 ```math
99 \widehat D = \widehat A \widehat B
100 ```
101
102 ```math
103 D = s_a s_b \widehat D + C
104 ```
105
106 ```math
107 D = s_a s_b \widehat A \widehat B + C
108 ```
109
110 Epilogue parameters:
111
112 - `scale_a` is the scale for activations, can be per-tensor (scalar) or per-token (column-vector).
113 - `scale_b` is the scale for weights, can be per-tensor (scalar) or per-channel (row-vector).
114 - `bias` is the bias, is always per-channel (row-vector).
115
116 ### `ScaledEpilogueAzp`
117
118 This epilogue computes the asymmetric per-tensor quantization for activations with bias.
119 The output of the GEMM is:
120
121 ```math
122 \widehat D = \widehat A \widehat B - z_a J_a \widehat B
123 ```
124
125 ```math
126 D = s_a s_b \widehat D + C
127 ```
128
129 ```math
130 D = s_a s_b \left( \widehat A \widehat B - z_a J_a \widehat B \right) + C
131 ```
132
133 Because $` z_a `$ is a scalar, the zero-point term $` z_a J_a \widehat B `$ has every row equal to $` z_a \mathbf 1 B `$.
134 That is precomputed and stored in `azp_with_adj` as a row-vector.
135
136 Epilogue parameters:
137
138 - `scale_a` is the scale for activations, can be per-tensor (scalar) or per-token (column-vector).
139 - Generally this will be per-tensor as the zero-points are per-tensor.
140 - `scale_b` is the scale for weights, can be per-tensor (scalar) or per-channel (row-vector).
141 - `azp_with_adj` is the precomputed zero-point term ($` z_a J_a \widehat B `$), is per-channel (row-vector).
142 - `bias` is the bias, is always per-channel (row-vector).
143
144 To use these kernels efficiently, users must precompute the `azp_with_adj` term offline and pass it to the kernel.
145
146 ### `ScaledEpilogueAzpPerToken`
147
148 This epilogue computes the asymmetric per-token quantization for activations with bias.
149
150 The output of the GEMM is the same as above, but the $` z_a `$ is a column-vector.
151 That means the zero-point term $` z_a J_a \widehat B `$ becomes an outer product of $` z_a `$ and $` \mathbf 1 \widehat B `$.
152
153 Epilogue parameters:
154
155 - `scale_a` is the scale for activations, can be per-tensor (scalar) or per-token (column-vector).
156 - Generally this will be per-token as the zero-points are per-token.
157 - `scale_b` is the scale for weights, can be per-tensor (scalar) or per-channel (row-vector).
158 - `azp_adj` is the precomputed zero-point adjustment term ($` \mathbf 1 \widehat B `$), is per-channel (row-vector).
159 - `azp` is the zero-point (`z_a`), is per-token (column-vector).
160 - `bias` is the bias, is always per-channel (row-vector).
161
162 To use these kernels efficiently, users must precompute the `azp_adj` term offline and pass it to the kernel.
163
164 The epilogue performs the following computation (where `Dq` is the raw quantized output of the GEMM):
165
166 ```math
167 out = scale_a * scale_b * (Dq - azp_adj * azp) + bias
168 ```