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$ 是标量，可以在所有累加完成后统一应用。欢迎在 GitHub Discussions 分享你的解法。