Weight Dequantization

原始题目：LeetGPU - Weight Dequantization

题目描述

编写一个 GPU 程序，在 GPU 上对权重矩阵进行"解量化"。给定一个包含量化值的输入矩阵 $X$（$M \times N$）和一个缩放矩阵 $S$（$\lceil M/T \rceil \times \lceil N/T \rceil$），其中 $T$ 为 tile 大小。

对于每个元素 $X_{i,j}$，对应的缩放因子为 $S_{row, col}$，其中：

$$row = \lfloor i / T \rfloor, \quad col = \lfloor j / T \rfloor$$

输出 $Y_{i,j}$ 计算为：

$$Y_{i,j} = X_{i,j} \times S_{row, col}$$

实现要求

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

示例

Input:  M=4, N=4, T=2
        X = [[10,10,5,5],[10,10,5,5],[2,2,8,8],[2,2,8,8]]
        S = [[0.5, 2.0],[4.0, 0.25]]
Output: Y = [[5,5,10,10],[5,5,10,10],[8,8,2,2],[8,8,2,2]]
        (Tile (0,0)×0.5, (0,1)×2.0, (1,0)×4.0, (1,1)×0.25)

约束条件

• $1 \le M, N \le 8{,}192$。
• $TILE\_SIZE \in \{16, 32, 64, 128\}$。
• 性能测试在 $M = N = 8{,}192,\ T = 128$ 的规模下进行。

解题思路

权重解量化是量化推理的逆操作——将低精度权重乘以每个 tile 的缩放因子恢复为浮点值。每 $T \times T$ block 内的所有元素共享同一个缩放因子。通过先加载缩放因子到寄存器，然后复用计算整个 tile，可以显著减少对缩放矩阵的全局内存访问。每个线程块处理一个或多个 tile 时可以利用共享内存广播缩放因子。

代码实现

CUDA

#include <cuda_runtime.h>
__global__ void dequant_kernel(const float* X, const float* S, float* Y, int M, int N, int T) {
    int row=blockIdx.y*blockDim.y+threadIdx.y, col=blockIdx.x*blockDim.x+threadIdx.x;
    if(row<M&&col<N) {
        int sr=row/T, sc=col/T, SC=(N+T-1)/T;
        Y[row*N+col]=X[row*N+col]*S[sr*SC+sc];
    }
}
extern "C" void solve(const float* X, const float* S, float* Y, int M, int N, int T) {
    dim3 t(16,16), b((N+15)/16,(M+15)/16);
    dequant_kernel<<<b,t>>>(X,S,Y,M,N,T); cudaDeviceSynchronize();
}

Triton

import triton, triton.language as tl
@triton.jit
def dequant_kernel(X_ptr,S_ptr,Y_ptr, M:tl.constexpr,N:tl.constexpr,T:tl.constexpr, BLOCK:tl.constexpr):
    col=tl.program_id(0)*BLOCK+tl.arange(0,BLOCK); row=tl.program_id(1)*BLOCK+tl.arange(0,BLOCK)
    mask=(row[:,None]<M)&(col[None,:]<N)
    sr=row[:,None]//T; sc=col[None,:]//T; SC=(N+T-1)//T
    s=tl.load(S_ptr+sr*SC+sc,mask=mask); x=tl.load(X_ptr+row[:,None]*N+col[None,:],mask=mask)
    tl.store(Y_ptr+row[:,None]*N+col[None,:],x*s,mask=mask)