Kernel for computing GEMV C = qw.T @ a - Input a: (Batch_size, M) - Weight qw: (M // pack_num, N) - scales: (M // group_size, N) - qzeros: (M // group_size // pack_num, N) - Output Y: (Batch_size, N)
(
# Pointers to matrices
a_ptr, qw_ptr, c_ptr, scales_ptr, zeros_ptr,
# Matrix dimensions
M, N,
pack_num, w_bit,
# Quantization parameters
group_size, offset,
CACHE_KEY_M,
CACHE_KEY_N,
# The stride variables represent how much to increase the ptr by when moving by 1
# element in a particular dimension. E.g. stride_am is how much to increase a_ptr
# by to get the element one row down (A has M rows)
stride_am,
# Quantization parameters
# Meta-parameters
BATCHSIZE: tl.constexpr,
BLOCK_M: tl.constexpr,
BLOCK_N: tl.constexpr,
EVEN_N: tl.constexpr,
)
| 174 | ) |
| 175 | @triton.jit |
| 176 | def transposed_gemv_atomicadd_kernel( |
| 177 | # Pointers to matrices |
| 178 | a_ptr, qw_ptr, c_ptr, scales_ptr, zeros_ptr, |
| 179 | # Matrix dimensions |
| 180 | M, N, |
| 181 | pack_num, w_bit, |
| 182 | # Quantization parameters |
| 183 | group_size, offset, |
| 184 | CACHE_KEY_M, |
| 185 | CACHE_KEY_N, |
| 186 | # The stride variables represent how much to increase the ptr by when moving by 1 |
| 187 | # element in a particular dimension. E.g. stride_am is how much to increase a_ptr |
| 188 | # by to get the element one row down (A has M rows) |
| 189 | stride_am, |
| 190 | # Quantization parameters |
| 191 | # Meta-parameters |
| 192 | BATCHSIZE: tl.constexpr, |
| 193 | BLOCK_M: tl.constexpr, |
| 194 | BLOCK_N: tl.constexpr, |
| 195 | EVEN_N: tl.constexpr, |
| 196 | ): |
| 197 | """ |
| 198 | Kernel for computing GEMV C = qw.T @ a |
| 199 | - Input a: (Batch_size, M) |
| 200 | - Weight qw: (M // pack_num, N) |
| 201 | - scales: (M // group_size, N) |
| 202 | - qzeros: (M // group_size // pack_num, N) |
| 203 | - Output Y: (Batch_size, N) |
| 204 | """ |
| 205 | start_m = tl.program_id(0) |
| 206 | start_n = tl.program_id(1) |
| 207 | offs_m = tl.arange(0, BLOCK_M) |
| 208 | rm = start_m * BLOCK_M + tl.arange(0, BLOCK_M) |
| 209 | rn = start_n * BLOCK_N + tl.arange(0, BLOCK_N) |
| 210 | |
| 211 | a_ptr = a_ptr + rm |
| 212 | c_ptr = c_ptr + rn |
| 213 | |
| 214 | # load weight |
| 215 | qw_shifter = (offs_m % pack_num) * w_bit |
| 216 | qw_off = (rm[:, None] // pack_num) * stride_am + rn[None, :] |
| 217 | qw_ptr = qw_ptr + qw_off |
| 218 | qw_packed = tl.load(qw_ptr) if EVEN_N else tl.load(qw_ptr, mask=rn[None, :] < N, other=0.0) |
| 219 | qw_unpacked = (qw_packed >> qw_shifter[:, None]) & offset |
| 220 | |
| 221 | # load sacle |
| 222 | m_iters_per_quant_group = group_size // BLOCK_M |
| 223 | grp_idx = start_m // m_iters_per_quant_group |
| 224 | col_offs = rn |
| 225 | scales = tl.load(scales_ptr + (stride_am * grp_idx) + col_offs) # (N,) |
| 226 | |
| 227 | # load zeros |
| 228 | packed_zeros = tl.load( |
| 229 | zeros_ptr + stride_am * (grp_idx // pack_num) + col_offs |
| 230 | ) |
| 231 | unpacked_zeros = (packed_zeros >> ((grp_idx % pack_num) * w_bit)) & offset |
| 232 | |
| 233 | # dequant w |