INT8 Quantized MatMul

原始题目：LeetGPU - INT8 Quantized MatMul

题目描述

编写一个 8 位有符号整数矩阵的量化矩阵乘法程序。给定输入矩阵 $A$（$M \times K$）和 $B$（$K \times N$），量化缩放因子 $s_A, s_B, s_C$ 和零点 $z_A, z_B, z_C$，计算：

$$C_{quant}(i, j) = \text{clamp}\!\left(\text{round}\!\left(\sum_{k=0}^{K-1} (A_{ik} - z_A)(B_{kj} - z_B) \cdot \frac{s_A s_B}{s_C}\right) + z_C,\ -128,\ 127\right)$$

累加使用 int32，缩放使用 float32，最后四舍五入、加零点、钳位到 $[-128, 127]$。

实现要求

• 不允许使用外部库。
• solve 函数签名必须保持不变。
• 最终结果必须以 int8 类型存储在 $C$ 中。

示例

示例 1

Input:  A = [[1,2],[3,4]], B = [[5,6],[7,8]], M=N=K=2
        s_A=0.1, s_B=0.2, s_C=0.05, z_A=z_B=z_C=0
Output: C = [[19,22],[43,50]]

示例 2

Input:  A = [[1,2]], B = [[3],[4]], M=1, N=1, K=2
        s_A=s_B=s_C=1.0, z_A=1, z_B=3, z_C=5
Output: C = [[6]]

约束条件

• $1 \le M, N, K \le 4{,}096$。
• $s_A, s_B, s_C > 0$。
• $-128 \le z_A, z_B, z_C \le 127$。
• 性能测试在 $K = 2{,}048,\ M = 8{,}192,\ N = 4{,}096$ 的规模下进行。

解题思路

INT8 量化推理是现代 LLM 部署的核心技术。核心思路是用 int8 矩阵乘法替代 float 乘法，利用 GPU 的 int8 Tensor Core 指令（如 mma.sync.aligned.m16n8k32.row.col.s32.s8.s8.s32）获得远高于浮点的吞吐。去零点操作可以提前或融合到乘法中。缩放因子 $s_A s_B / s_C$ 是标量，可以在所有累加完成后统一应用。

代码实现

CUDA

#include <cuda_runtime.h>
__global__ void int8_matmul(const int8_t* A, const int8_t* B, int8_t* C,
    int M, int K, int N, float sA, float sB, float sC, int zA, int zB, int zC) {
    int row=blockIdx.y*blockDim.y+threadIdx.y, col=blockIdx.x*blockDim.x+threadIdx.x;
    if(row<M && col<N) {
        int acc=0;
        for(int k=0;k<K;k++) acc+=(A[row*K+k]-zA)*(B[k*N+col]-zB);
        float val = acc * sA * sB / sC + zC;
        C[row*N+col]=(int8_t)min(max(roundf(val),-128.0f),127.0f);
    }
}
extern "C" void solve(const int8_t* A, const int8_t* B, int8_t* C,
    int M, int K, int N, float sA, float sB, float sC, int zA, int zB, int zC) {
    dim3 t(16,16), b((N+15)/16,(M+15)/16);
    int8_matmul<<<b,t>>>(A,B,C,M,K,N,sA,sB,sC,zA,zB,zC);
    cudaDeviceSynchronize();
}

Triton

import triton, triton.language as tl
@triton.jit
def int8_matmul(A_ptr,B_ptr,C_ptr, M,N,K, sA,sB,sC,zA,zB,zC, BLOCK_M:tl.constexpr,BLOCK_N:tl.constexpr,BLOCK_K:tl.constexpr):
    pm= tl.program_id(0); pn= tl.program_id(1)
    rm= pm*BLOCK_M+tl.arange(0,BLOCK_M); rn= pn*BLOCK_N+tl.arange(0,BLOCK_N); rk= tl.arange(0,BLOCK_K)
    acc= tl.zeros((BLOCK_M,BLOCK_N),tl.int32)
    for k in range(0,K,BLOCK_K):
        a= tl.load(A_ptr+rm[:,None]*K+(k+rk)[None,:]) - zA
        b= tl.load(B_ptr+(k+rk)[:,None]*N+rn[None,:]) - zB
        acc+= tl.dot(a.to(tl.int32), b.to(tl.int32))
    val= acc.to(tl.float32)*sA*sB/sC + zC
    tl.store(C_ptr+rm[:,None]*N+rn[None,:], tl.clamp(tl.math.round(val),-128,127).to(tl.int8))