| """Step 4.10b (D-016): INT8 weight-only matmul for decode, reading the int8 weights directly. |
| |
| The PyTorch path in tinyllm.quant (Int8Linear, backend "torch") runs three kernels per layer: |
| cast the int8 weight to bf16 (reads 1 byte and writes 2 per weight), matmul (reads those 2 bytes |
| again), then scale. Here one kernel reads each int8 weight once, converts it in registers (exact: |
| |q| <= 127), accumulates in fp32 and applies the per-row scale at the end: |
| |
| y[m, n] = scale[n] * sum_k x[m, k] * q[n, k] x: (M, K) bf16/fp16, q: (N, K) int8 |
| |
| At decode, M (tokens in the step = batch) is small, so the step is bound by reading q. The tile |
| over M is padded to 16 rows for tl.dot (tensor cores); the wasted flops are cheap next to the |
| bytes. Each program owns BLOCK_N output columns and loops over K. Block sizes are autotuned per |
| (N, K, M bucket) on first use, which must happen outside CUDA-graph capture (GraphDecoder warms up |
| first). M above MAX_M (prefill) stays on the PyTorch path; see Int8Linear.forward. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import torch |
| import triton |
| import triton.language as tl |
| from torch import Tensor |
|
|
| MAX_M = 64 |
|
|
| _CONFIGS = [triton.Config({"BLOCK_N": bn, "BLOCK_K": bk}, num_warps=w, num_stages=s) |
| for bn in (16, 32, 64) for bk in (128, 256) for w in (4,) for s in (3, 4)] |
| _CONFIGS += [triton.Config({"BLOCK_N": 32, "BLOCK_K": 512}, num_warps=8, num_stages=3), |
| triton.Config({"BLOCK_N": 64, "BLOCK_K": 128}, num_warps=2, num_stages=4)] |
|
|
|
|
| @triton.autotune(configs=_CONFIGS, key=["N", "K", "BLOCK_M"]) |
| @triton.jit |
| def _int8_matmul_kernel(x_ptr, q_ptr, s_ptr, y_ptr, M, N, K, stride_xm, stride_qn, stride_ym, |
| BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr): |
| pid = tl.program_id(0) |
| offs_m = tl.arange(0, BLOCK_M) |
| offs_n = pid * BLOCK_N + tl.arange(0, BLOCK_N) |
| m_ok, n_ok = offs_m < M, offs_n < N |
| acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) |
| for k0 in range(0, K, BLOCK_K): |
| offs_k = k0 + tl.arange(0, BLOCK_K) |
| k_ok = offs_k < K |
| x = tl.load(x_ptr + offs_m[:, None] * stride_xm + offs_k[None, :], |
| mask=m_ok[:, None] & k_ok[None, :], other=0.0) |
| q = tl.load(q_ptr + offs_n[:, None] * stride_qn + offs_k[None, :], |
| mask=n_ok[:, None] & k_ok[None, :], other=0) |
| acc = tl.dot(x, tl.trans(q.to(x.dtype)), acc) |
| s = tl.load(s_ptr + offs_n, mask=n_ok, other=0.0) |
| y = acc * s[None, :] |
| tl.store(y_ptr + offs_m[:, None] * stride_ym + offs_n[None, :], y, mask=m_ok[:, None] & n_ok[None, :]) |
|
|
|
|
| def int8_matmul(x: Tensor, q: Tensor, scale: Tensor) -> Tensor: |
| """x (M, K) bf16/fp16 CUDA, M <= MAX_M; q (N, K) int8; scale (N,) fp32 -> y (M, N) fp32. |
| |
| Same values as F.linear(x, q.to(x.dtype)) * scale, except that the matmul result isn't rounded |
| to bf16 before scaling (closer to the fp32 dequantize-then-matmul reference, not further).""" |
| M, K = x.shape |
| N = q.shape[0] |
| if M > MAX_M or q.shape[1] != K or x.dtype not in (torch.bfloat16, torch.float16): |
| raise ValueError(f"int8_matmul: x {tuple(x.shape)} {x.dtype}, q {tuple(q.shape)}; needs M <= {MAX_M}, " |
| "matching K, bf16/fp16 x") |
| x = x if x.stride(1) == 1 else x.contiguous() |
| q = q if q.stride(1) == 1 else q.contiguous() |
| y = torch.empty((M, N), device=x.device, dtype=torch.float32) |
| block_m = max(16, triton.next_power_of_2(M)) |
| grid = lambda meta: (triton.cdiv(N, meta["BLOCK_N"]),) |
| _int8_matmul_kernel[grid](x, q, scale, y, M, N, K, x.stride(0), q.stride(0), y.stride(0), BLOCK_M=block_m) |
| return y |
|
|
|
|
| def best_configs() -> dict[str, str]: |
| """The autotuned config per (N, K, BLOCK_M) key so far, for the benchmark JSON.""" |
| return {str(k): str(v) for k, v in _int8_matmul_kernel.cache.items()} |
|
|
