cuda编程笔记（37）--逐行量化的kernel分析

什么是 Per-row 量化？

在深度学习中，常见的量化是将 float32 转换为 int8。

全局量化 (Per-tensor): 整个矩阵共用一个 Scale。如果矩阵中某个值特别大（Outlier），会导致其他小值在量化后全部变成 0，精度损失严重。
每行量化 (Per-row): 矩阵的每一行 $i$ 都有一个专属的 $s_i$ 。

数学公式：

对于矩阵中的元素，其量化过程为：

其中通常定义为该行绝对值的最大值缩放到 int8 范围（-128 到 127）：

写法分析

由于处理单位是每行，只用考虑每行交给谁处理即可。每行一个thread处理，该方案舒适太离谱，不做考虑。所以我们要对比block处理一行和warp处理一行的差距

cpp 复制代码

//简单的int8 Per-row 量化Kernel
//一个block负责一行
__global__ void per_row_quantize_kernel(float *input,int8_t *output,float* scales, int cols){
    int row=blockIdx.x;//行id
    int tid=threadIdx.x;//本线程id
    float row_max=0.f;//本线程负责的元素的最大值
    //每个线程block-stride迭代
    for(int i=tid;i<cols;i+=blockDim.x){
        row_max=fmaxf(row_max,fabsf(input[row*cols+i]));
    }
    //开共享内存，在block内准备max规约
    extern __shared__ float shared_data[];
    shared_data[tid]=row_max;
    __syncthreads();
    //block内规约最大值
    for(int s=blockDim.x/2;s>0;s>>=1){
        if(tid<s){
            shared_data[tid]=fmaxf(shared_data[tid],shared_data[tid+s]);
        }
        __syncthreads();
    }
    //最终的本行最大值
    float final_max=shared_data[0];
    //计算缩放系数
    float scale=final_max/127.f;
    //对每个元素进行缩放并以int8存储
    for(int i=tid;i<cols;i+=blockDim.x){
        float val=input[row*cols+i];
        output[row*cols+i]=(int8_t)roundf(val/scale);
    }
    //顺便存一下每行的缩放系数，以供反量化
    if(tid==0){
        scales[row]=scale;
    }
}
//一个warp处理一行
__global__ void per_row_quant_warp_kernel(float *input,int8_t *output,float* scales, int cols){
    //定位当前行
    int warp_id=(blockDim.x*blockIdx.x+threadIdx.x)/warpSize;
    int lane_id=threadIdx.x%32;
    int row=warp_id;
    float local_max=0.0f;
    //依旧每个线程循环处理长行
    for(int j=lane_id;j<cols;j+=warpSize){
        local_max=fmaxf(local_max,input[row*cols+j]);
    }
    //Warp Shuffle 规约，找出一个Warp内的最大值
    for(int offset=warpSize/2;offset>0;offset>>=1){
        local_max=fmaxf(local_max,__shfl_down_sync(0xFFFFFFFF,local_max,offset));
    }
    //此时lane_id=0线程持有整行的最大值
    float final_max=__shfl_sync(0xFFFFFFFF,local_max,0);
    float scale=final_max/127;
    //执行量化
    for (int j = lane_id; j < cols; j += 32) {
        float val = input[row * cols + j];
        output[row * cols + j] = (int8_t)roundf(val / scale);
    }
    // 4. 存储 scale
    if (lane_id == 0) {
        scales[row] = scale;
    }
}

Profile分析

先用cuda事件简单测一下

cpp 复制代码

int main() {
    cudaEvent_t start,end;
    float WarpPerRow=0,BlockPerRow=0;
    cudaEventCreate(&start);
    cudaEventCreate(&end);
    const int M=1024,N=1024;
    std::vector<float> host_data(M * N, 1.0f);
    float *d_data;
    int8_t *d_out;
    float *d_scales;
    cudaMalloc(&d_data,M*N*sizeof(float));
    cudaMalloc(&d_out,M*N*sizeof(int8_t));
    cudaMalloc(&d_scales,M*sizeof(float));
    cudaMemcpy(d_data,host_data.data(),M*N*sizeof(float),cudaMemcpyHostToDevice);
    //给一个block一行的配置
    dim3 block_B(256),grid_B(M);
    size_t shared_mem_size=block_B.x*sizeof(float);
    //给一个warp一行的配置
    dim3 block_W(256);
    int warps_per_block = block_W.x / 32;
    dim3 grid_W((M + warps_per_block - 1) / warps_per_block);


    int warmup=10;
    for(int i=0;i<warmup;i++){
        per_row_quantize_kernel<<<grid_B,block_B,shared_mem_size>>>(d_data,d_out,d_scales,N);
    }
    cudaEventRecord(start);
    per_row_quantize_kernel<<<grid_B,block_B,shared_mem_size>>>(d_data,d_out,d_scales,N);
    cudaEventRecord(end);
    cudaEventSynchronize(end);
    cudaEventElapsedTime(&BlockPerRow,start,end);
 
    cudaEventRecord(start);
    per_row_quant_warp_kernel<<<grid_W,block_W>>>(d_data,d_out,d_scales,N);
    cudaEventRecord(end);
    cudaEventSynchronize(end);
    cudaEventElapsedTime(&WarpPerRow,start,end);
    // 检查是否有异常抛出
    cudaError_t err = cudaGetLastError();
    if (err != cudaSuccess) printf("Kernel Error: %s\n", cudaGetErrorString(err));
 
    std::cout<<"block per row time is: "<<BlockPerRow<<"(ms)"<<std::endl;
    std::cout<<"warp per row time is: "<<WarpPerRow<<"(ms)"<<std::endl;
    cudaFree(d_data);
    cudaFree(d_out);
    cudaFree(d_scales);
    cudaEventDestroy(start);
    cudaEventDestroy(end);
}

cpp 复制代码

block per row time is: 48.2633(ms)
warp per row time is: 49.6557(ms)

居然是block更快，那具体看看Warp瓶颈在哪

由于我也不太会用Nsight，分析都是AI做的，仅供参考

指标	Block 版本 (Kernel 10)	Warp 版本 (Kernel 11)	结论
Compute (计算)	57.14%	20.65%	Block 版在忙着做计算/同步，Warp 版计算闲置严重。
Memory (带宽)	57.14%	59.16%	两者都受限于访存。

Warp 版本的访存利用率更高 (59.16%) ：这说明 __shfl_down_sync 节省下来的指令发射时间，让 GPU 能腾出更多精力去搬运数据。
计算吞吐量暴跌：Warp 版本的计算利用率只有 20%，说明它的指令发射压力非常小，大部分时间都在等内存回传数据。

猜测是每行元素数量N太多，Warp需要迭代32次才能全部访问完，而block只用迭代4次。

由于量化算子是 访存密集型（Memory-bound） 而非计算密集型，谁能以最少的指令、最整齐的队列把数据从显存搬进寄存器，谁就是赢家。

如果想要优化，可以尝试向量化读取。