Compute the matrix multiplication C = A x B. A is of shape (M, K) float16 B is of shape (K//8, N) int32 C is of shape (M, N) float16 scales is of shape (G, N) float16 zeros is of shape (G, N) float16 g_ptr is of shape (K) int32
(a_ptr, b_ptr, c_ptr,
scales_ptr, zeros_ptr, g_ptr,
M, N, K, bits, maxq,
stride_am, stride_ak,
stride_bk, stride_bn,
stride_cm, stride_cn,
stride_scales, stride_zeros,
BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_N: tl.constexpr, BLOCK_SIZE_K: tl.constexpr,
GROUP_SIZE_M: tl.constexpr)
| 55 | ) |
| 56 | @triton.jit |
| 57 | def matmul_248_kernel(a_ptr, b_ptr, c_ptr, |
| 58 | scales_ptr, zeros_ptr, g_ptr, |
| 59 | M, N, K, bits, maxq, |
| 60 | stride_am, stride_ak, |
| 61 | stride_bk, stride_bn, |
| 62 | stride_cm, stride_cn, |
| 63 | stride_scales, stride_zeros, |
| 64 | BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_N: tl.constexpr, BLOCK_SIZE_K: tl.constexpr, |
| 65 | GROUP_SIZE_M: tl.constexpr): |
| 66 | """ |
| 67 | Compute the matrix multiplication C = A x B. |
| 68 | A is of shape (M, K) float16 |
| 69 | B is of shape (K//8, N) int32 |
| 70 | C is of shape (M, N) float16 |
| 71 | scales is of shape (G, N) float16 |
| 72 | zeros is of shape (G, N) float16 |
| 73 | g_ptr is of shape (K) int32 |
| 74 | """ |
| 75 | infearure_per_bits = 32 // bits |
| 76 | |
| 77 | pid = tl.program_id(axis=0) |
| 78 | num_pid_m = tl.cdiv(M, BLOCK_SIZE_M) |
| 79 | num_pid_n = tl.cdiv(N, BLOCK_SIZE_N) |
| 80 | num_pid_k = tl.cdiv(K, BLOCK_SIZE_K) |
| 81 | num_pid_in_group = GROUP_SIZE_M * num_pid_n |
| 82 | group_id = pid // num_pid_in_group |
| 83 | first_pid_m = group_id * GROUP_SIZE_M |
| 84 | group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M) |
| 85 | pid_m = first_pid_m + (pid % group_size_m) |
| 86 | pid_n = (pid % num_pid_in_group) // group_size_m |
| 87 | |
| 88 | offs_am = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) |
| 89 | offs_bn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) |
| 90 | offs_k = tl.arange(0, BLOCK_SIZE_K) |
| 91 | a_ptrs = a_ptr + (offs_am[:, None] * stride_am + offs_k[None, :] * stride_ak) # (BLOCK_SIZE_M, BLOCK_SIZE_K) |
| 92 | a_mask = (offs_am[:, None] < M) |
| 93 | # b_ptrs is set up such that it repeats elements along the K axis 8 times |
| 94 | b_ptrs = b_ptr + ((offs_k[:, None] // infearure_per_bits) * stride_bk + offs_bn[None, |
| 95 | :] * stride_bn) # (BLOCK_SIZE_K, BLOCK_SIZE_N) |
| 96 | g_ptrs = g_ptr + offs_k |
| 97 | # shifter is used to extract the N bits of each element in the 32-bit word from B |
| 98 | scales_ptrs = scales_ptr + offs_bn[None, :] |
| 99 | zeros_ptrs = zeros_ptr + (offs_bn[None, :] // infearure_per_bits) |
| 100 | |
| 101 | shifter = (offs_k % infearure_per_bits) * bits |
| 102 | zeros_shifter = (offs_bn % infearure_per_bits) * bits |
| 103 | accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32) |
| 104 | |
| 105 | for k in range(0, num_pid_k): |
| 106 | g_idx = tl.load(g_ptrs) |
| 107 | |
| 108 | # Fetch scales and zeros; these are per-outfeature and thus reused in the inner loop |
| 109 | scales = tl.load(scales_ptrs + g_idx[:, None] * stride_scales) # (BLOCK_SIZE_K, BLOCK_SIZE_N,) |
| 110 | zeros = tl.load(zeros_ptrs + g_idx[:, None] * stride_zeros) # (BLOCK_SIZE_K, BLOCK_SIZE_N,) |
| 111 | |
| 112 | zeros = (zeros >> zeros_shifter[None, :]) & maxq |
| 113 | zeros = (zeros + 1) |
| 114 |
nothing calls this directly
no outgoing calls
no test coverage detected