INT4 Weight-Only Quantized MatMul

原始题目：LeetGPU - INT4 Weight-Only Quantized MatMul

题目描述

实现仅权重的 INT4 量化矩阵乘法（W4A16），这是现代 LLM 推理中使用的核心 kernel。给定 float16 激活矩阵 $x$（$M \times K$）和以 packed INT4 格式存储的权重矩阵，计算输出矩阵 $y = x \times W^T$（$M \times N$），其中 $W$ 是解量化后的 float16 权重矩阵（$N \times K$）。

打包格式：w_q 的每个字节存储两个 INT4 权重。高半字节（bit 7–4）存储 $w[n, 2i]$，低半字节（bit 3–0）存储 $w[n, 2i+1]$。INT4 值以无符号形式存储在 $[0, 15]$ 范围内，偏移量为 8，因此有符号权重为 $\text{nibble} - 8$，范围为 $[-8, 7]$。

解量化：权重按组解量化。沿 $K$ 维度的每 $group\_size$ 个连续权重共享一个 float16 缩放因子：

$$W[n, k] = (w\_q\_nibble[n, k] - 8) \times scales[n, \lfloor k / group\_size \rfloor]$$

实现要求

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

示例

Input:  M=2, N=4, K=4, group_size=2
        x (float16, 2×4): [[1,0,1,0],[0,1,0,1]]
        w_q (uint8, 4×2): packed, 带符号INT4值为 [[1,1],[1,1],[2,2],[2,2],[-1,-1],[-1,-1],[0,0],[0,0]]
        scales (float16, 4×2): 全 0.5
        W_dequant = (nibble-8)*0.5
Output: y = x × W^T (float16, 2×4): [[1,2,-1,0],[1,2,-1,0]]

约束条件

• $1 \le M, N, K \le 8{,}192$。
• $K$ 可被 2 和 $group\_size$ 整除。
• $group\_size \in \{2, 4, 8, 16, 32, 64, 128\}$。
• x 和 scales 为 float16，w_q 为 uint8，y 为 float16。
• 性能测试在 $M=N=K=4{,}096$，$group\_size=128$ 下进行。

解题思路

INT4 量化将权重存储需求减少到原来的 1/4（相比 FP16）。解量化需要在矩阵乘法之前或融合到乘法之中完成。高效的实现会将 nibble 解包、减偏移、乘缩放因子与矩阵乘法融合为一个 kernel，避免中间解量化矩阵的显存分配。注意利用 dp4a（点积 4x8bit）指令加速 INT4 计算。欢迎在 GitHub Discussions 分享你的解法。