Kernel for computing the matmul C = A x qw a: (M, K) qw: (K // pack_num, N) scales: (K // group_size, N) qzeros: (K // group_size // pack_num, N)
(
# Pointers to matrices
a_ptr, qw_ptr, c_ptr, scales_ptr, zeros_ptr,
# Matrix dimensions
M, N, K,
pack_num, w_bit,
# Quantization parameters
group_size, offset,
# Meta-parameters
BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_N: tl.constexpr, BLOCK_SIZE_K: tl.constexpr,
GROUP_SIZE_M: tl.constexpr,
)
| 21 | ) |
| 22 | @triton.jit |
| 23 | def quant_matmul_kernel( |
| 24 | # Pointers to matrices |
| 25 | a_ptr, qw_ptr, c_ptr, scales_ptr, zeros_ptr, |
| 26 | # Matrix dimensions |
| 27 | M, N, K, |
| 28 | pack_num, w_bit, |
| 29 | # Quantization parameters |
| 30 | group_size, offset, |
| 31 | # Meta-parameters |
| 32 | BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_N: tl.constexpr, BLOCK_SIZE_K: tl.constexpr, |
| 33 | GROUP_SIZE_M: tl.constexpr, |
| 34 | ): |
| 35 | """ |
| 36 | Kernel for computing the matmul C = A x qw |
| 37 | |
| 38 | a: (M, K) |
| 39 | qw: (K // pack_num, N) |
| 40 | scales: (K // group_size, N) |
| 41 | qzeros: (K // group_size // pack_num, N) |
| 42 | """ |
| 43 | |
| 44 | stride_zeros_k = N |
| 45 | stride_scales_k = N |
| 46 | stride_a_m = K |
| 47 | stride_qw_k = N |
| 48 | |
| 49 | pid = tl.program_id(axis=0) |
| 50 | num_pid_m = tl.cdiv(M, BLOCK_SIZE_M) |
| 51 | num_pid_n = tl.cdiv(N, BLOCK_SIZE_N) |
| 52 | num_pid_in_group = GROUP_SIZE_M * num_pid_n |
| 53 | group_id = pid // num_pid_in_group |
| 54 | first_pid_m = group_id * GROUP_SIZE_M |
| 55 | group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M) |
| 56 | |
| 57 | pid_m = first_pid_m + (pid % group_size_m) |
| 58 | pid_n = (pid % num_pid_in_group) // group_size_m |
| 59 | |
| 60 | offs_am = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M |
| 61 | offs_bn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)) |
| 62 | offs_k = tl.arange(0, BLOCK_SIZE_K) # (K,) |
| 63 | qw_shifter = (offs_k % pack_num) * w_bit |
| 64 | |
| 65 | accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32) |
| 66 | for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)): |
| 67 | a_offs = (k * BLOCK_SIZE_K) + (offs_am[:, None] * stride_a_m + offs_k[None, :]) # (M, K) |
| 68 | a = tl.load(a_ptr + a_offs) |
| 69 | |
| 70 | # load weight |
| 71 | qw_offs = (((k * BLOCK_SIZE_K) + offs_k[:, None]) // pack_num) * stride_qw_k + offs_bn[ |
| 72 | None, : |
| 73 | ] # (K, N) |
| 74 | qw_packed = tl.load(qw_ptr + qw_offs) # (K, N) |
| 75 | qw_unpacked = (qw_packed >> qw_shifter[:, None]) & offset |
| 76 | |
| 77 | # load sacle |
| 78 | k_iters_per_quant_group = group_size // BLOCK_SIZE_K |
| 79 | grp_idx = k // k_iters_per_quant_group |
| 80 | col_offs = offs_bn |