KV Cache INT8量化

/**
 * KV Cache INT8量化实现
 * 
 * 功能:
 * 1. FP16 → INT8量化
 * 2. INT8 → FP16反量化
 * 3. 精度测试和内存对比
 * 
 * 编译: nvcc -o kv_quant kv_quant.cu
 */

#include <stdio.h>
#include <cuda_runtime.h>
#include <cuda_fp16.h>

#define CUDA_CHECK(call) do { cudaError_t e = call; if(e != cudaSuccess) { \
    fprintf(stderr, "CUDA Error: %s\n", cudaGetErrorString(e)); exit(1); } } while(0)

// ==================== 对称量化Kernel ====================
// FP16 → INT8: x_int8 = round(x_fp16 / scale)
// scale = max(|x|) / 127
__global__ void quantize_symmetric(
    const __half* input,   // FP16输入
    int8_t* output,         // INT8输出
    float* scale_out,       // 输出scale值
    int n
) {
    // 1. 找最大绝对值
    float max_abs = 0.0f;
    for (int i = threadIdx.x; i < n; i += blockDim.x) {
        float val = __half2float(input[i]);
        max_abs = fmaxf(max_abs, fabsf(val));
    }
    
    // Warp归约找全局最大值
    for (int offset = 16; offset > 0; offset >>= 1) {
        float other = __shfl_down_sync(0xFFFFFFFF, max_abs, offset);
        max_abs = fmaxf(max_abs, other);
    }
    if (threadIdx.x == 0) {
        scale_out[blockIdx.x] = max_abs / 127.0f;
    }
    __syncthreads();
    
    // 2. 量化
    float scale = scale_out[blockIdx.x];
    for (int i = threadIdx.x; i < n; i += blockDim.x) {
        float val = __half2float(input[i]);
        output[i] = (int8_t)round(val / scale);
    }
}

// ==================== 反量化Kernel ====================
// INT8 → FP16: x_fp16 = x_int8 * scale
__global__ void dequantize_symmetric(
    const int8_t* input,   // INT8输入
    __half* output,         // FP16输出
    const float* scale,     // scale值
    int n
) {
    float s = scale[blockIdx.x];
    for (int i = threadIdx.x; i < n; i += blockDim.x) {
        output[i] = __float2half((float)input[i] * s);
    }
}

// ==================== 主函数 ====================
int main() {
    const int seq_len = 2048;
    const int head_dim = 128;
    const int n = seq_len * head_dim;
    
    printf("KV Cache INT8量化演示\n");
    printf("序列长度: %d, 头维度: %d\n", seq_len, head_dim);
    printf("原始大小 (FP16): %.2f KB\n", n * 2 / 1024.0);
    printf("量化后 (INT8): %.2f KB\n", n * 1 / 1024.0);
    printf("内存节省: 50%%\n\n");
    
    // 分配内存
    __half* h_fp16 = (__half*)malloc(n * sizeof(__half));
    int8_t* h_int8 = (int8_t*)malloc(n * sizeof(int8_t));
    __half* h_dequant = (__half*)malloc(n * sizeof(__half));
    
    // 初始化FP16数据
    for (int i = 0; i < n; i++) {
        h_fp16[i] = __float2half((float)(rand() % 1000) / 1000.0f - 0.5f);
    }
    
    // GPU量化
    __half *d_fp16, *d_dequant;
    int8_t *d_int8;
    float *d_scale;
    
    CUDA_CHECK(cudaMalloc(&d_fp16, n * sizeof(__half)));
    CUDA_CHECK(cudaMalloc(&d_int8, n * sizeof(int8_t)));
    CUDA_CHECK(cudaMalloc(&d_dequant, n * sizeof(__half)));
    CUDA_CHECK(cudaMalloc(&d_scale, 256 * sizeof(float)));  // scale数组
    
    CUDA_CHECK(cudaMemcpy(d_fp16, h_fp16, n * sizeof(__half), cudaMemcpyHostToDevice));
    
    // 量化
    int block_size = 256;
    int grid_size = (n + block_size - 1) / block_size;
    quantize_symmetric<<<grid_size, block_size>>>(d_fp16, d_int8, d_scale, n);
    
    // 反量化
    dequantize_symmetric<<<grid_size, block_size>>>(d_int8, d_dequant, d_scale, n);
    
    CUDA_CHECK(cudaDeviceSynchronize());
    
    // 计算量化误差
    CUDA_CHECK(cudaMemcpy(h_fp16, d_fp16, n * sizeof(__half), cudaMemcpyDeviceToHost));
    CUDA_CHECK(cudaMemcpy(h_dequant, d_dequant, n * sizeof(__half), cudaMemcpyDeviceToHost));
    
    float mse = 0, max_err = 0;
    for (int i = 0; i < n; i++) {
        float diff = __half2float(h_fp16[i]) - __half2float(h_dequant[i]);
        mse += diff * diff;
        max_err = fmaxf(max_err, fabsf(diff));
    }
    mse /= n;
    
    printf("量化误差 (MSE): %e\n", mse);
    printf("最大误差: %e\n", max_err);
    printf(mse < 1e-4 ? "✅ 精度可接受\n" : "❌ 精度损失过大\n");
    
    // 清理
    cudaFree(d_fp16); cudaFree(d_int8); cudaFree(d_dequant); cudaFree(d_scale);
    free(h_fp16); free(h_int8); free(h_dequant);
    
    return 0;
}