Kernel for computing GEMV C = qw @ a - Input a: (Batch_size, N) - Weight qw: (M, N // pack_num) - scales: (M, N // group_size) - qzeros: (M, N // group_size // pack_num) - Output Y: (Batch_size, M)
(
# Pointers to matrices
a_ptr, qw_ptr, c_ptr, scales_ptr, zeros_ptr,
# Matrix dimensions
M, N,
pack_num, group_size,
# Quantization parameters
w_bit, 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_qw, stride_scale, stride_zeros,
# Quantization parameters
# Meta-parameters
BATCHSIZE: tl.constexpr,
BLOCK_M: tl.constexpr,
BLOCK_N: tl.constexpr,
EVEN_N: tl.constexpr,
)
| 345 | ) |
| 346 | @triton.jit |
| 347 | def gemv_kernel_v3( |
| 348 | # Pointers to matrices |
| 349 | a_ptr, qw_ptr, c_ptr, scales_ptr, zeros_ptr, |
| 350 | # Matrix dimensions |
| 351 | M, N, |
| 352 | pack_num, group_size, |
| 353 | # Quantization parameters |
| 354 | w_bit, offset, |
| 355 | CACHE_KEY_M, |
| 356 | CACHE_KEY_N, |
| 357 | # The stride variables represent how much to increase the ptr by when moving by 1 |
| 358 | # element in a particular dimension. E.g. stride_am is how much to increase a_ptr |
| 359 | # by to get the element one row down (A has M rows) |
| 360 | stride_qw, stride_scale, stride_zeros, |
| 361 | # Quantization parameters |
| 362 | # Meta-parameters |
| 363 | BATCHSIZE: tl.constexpr, |
| 364 | BLOCK_M: tl.constexpr, |
| 365 | BLOCK_N: tl.constexpr, |
| 366 | EVEN_N: tl.constexpr, |
| 367 | ): |
| 368 | """ |
| 369 | Kernel for computing GEMV C = qw @ a |
| 370 | - Input a: (Batch_size, N) |
| 371 | - Weight qw: (M, N // pack_num) |
| 372 | - scales: (M, N // group_size) |
| 373 | - qzeros: (M, N // group_size // pack_num) |
| 374 | - Output Y: (Batch_size, M) |
| 375 | """ |
| 376 | start_m = tl.program_id(0) |
| 377 | |
| 378 | rm = start_m * BLOCK_M + tl.arange(0, BLOCK_M) |
| 379 | rn = tl.arange(0, BLOCK_N) |
| 380 | |
| 381 | a_ptr = a_ptr + rn |
| 382 | |
| 383 | # load weight |
| 384 | qw_shifter = (rn % pack_num) * w_bit |
| 385 | qw_off = rm[:, None] * stride_qw + rn[None, :] // pack_num |
| 386 | qw_ptr = qw_ptr + qw_off |
| 387 | |
| 388 | acc0 = tl.zeros((BLOCK_M,), dtype=tl.float32) |
| 389 | |
| 390 | for n in range(0, N, BLOCK_N): |
| 391 | # load activation |
| 392 | a0 = tl.load(a_ptr) if EVEN_N else tl.load(a_ptr, mask=rn[None, :] < N, other=0.0) |
| 393 | |
| 394 | # unpack weight |
| 395 | qw_packed = tl.load(qw_ptr) |
| 396 | qw_unpacked = (qw_packed >> qw_shifter[None, :]) & offset |
| 397 | |
| 398 | # load scale |
| 399 | grp_idx = rn[None, :] // group_size; |
| 400 | scales = tl.load(scales_ptr + rm[:, None] * stride_scale + grp_idx) |
| 401 | |
| 402 | # load zero |
| 403 | packed_zeros = tl.load( |
| 404 | zeros_ptr + rm[:, None] * stride_zeros + (grp_idx // pack_num) |