Hotdry.
归档
ai-systems

用 RL 自动调优 CUDA GEMM 内核：超越 cuBLAS 在 A100/H100 性能

CUDA-L2 通过强化学习搜索 GEMM 调度与平铺策略，在 A100 上 1000 个配置中多数超越 cuBLAS。给出工程部署参数、基准脚本与 QPS 监控要点。

在大型语言模型训练与推理中，通用矩阵乘法（GEMM）是计算密集型核心算子，占总 FLOPs 的 80% 以上。NVIDIA cuBLAS 作为基准库已高度优化，但针对特定矩阵形状（M, N, K），手工调优或搜索策略仍可带来 10-50% 加速。CUDA-L2 项目创新性地结合大语言模型（LLM）与强化学习（RL），自动生成并优化半精度（HGEMM）CUDA 内核，在 A100 GPU 上针对 1000 个典型配置实现超越 cuBLAS 的性能。该方法的核心在于 RL 代理学习最优的内核调度参数（如线程块尺寸、tiling 策略、共享内存分配），避免人工枚举的指数复杂度。

证据显示，在 A100 上评估的 1000 个 (M,N,K) 配置中，CUDA-L2 内核平均加速 torch.matmul 达 2x，cuBLASLt-heuristic 达 1.2x，甚至 cuBLASLt-AutoTuning 也落后 5-15%。例如，对于形状 64x4096x64 的 HGEMM，项目基准图示其 TFLOPS 高于 cuBLAS 约 20%。当前内核基于 CUTLASS v4.2.1，支持 SM80 架构的 F16F16F16F16（16 位累加器），已开源预编译 ptx 文件，覆盖稠密矩阵典型 workload。

要落地部署此类 RL 调优内核，需关注以下参数与清单，确保生产级稳定性：

1. 环境准备清单

  • GPU 兼容：A100 (Ampere SM80)，TORCH_CUDA_ARCH_LIST="8.0"。
  • 依赖安装： 
    git clone -b v4.2.1 https://github.com/NVIDIA/cutlass.git cutlass
export CUTLASS_DIR=/path/to/cutlass
pip install torch>=2.6.0
  • 克隆项目：git clone https://github.com/deepreinforce-ai/CUDA-L2
  • 风险阈值：若非 A100，性能降级 >20% 时回滚至 cuBLAS；仅支持列主序矩阵。

2. 基准评估参数

使用 ./eval_one_file.sh 脚本验证性能，支持离线（batch）与服务器（QPS）模式：

  • 核心参数：

    参数 值示例 说明
    --mnk 64_4096_64 M_N_K 形状，pad 零填充邻近大配置
    --warmup_seconds 5 预热时长，避免冷启动
    --benchmark_seconds 10 基准时长，>5s 以平滑波动
    --base_dir ./results 输出目录，存 ptx 与 perf 日志
    --gpu_device_id 0 多卡指定 ID
    --mode offline/server 离线 vs. QPS 负载
    --target_qps 100 服务器模式目标吞吐，监控延迟 P99 <50ms

  • 离线模式示例：

    ./eval_one_file.sh --mnk 64_4096_64 --warmup_seconds 5 --benchmark_seconds 10 --base_dir ./results --gpu_device_id 0 --mode offline

    输出：TFLOPS、相对 cuBLAS speedup、正确性检查（randn/zero-one 测试）。

  • 服务器模式：模拟推理负载，target_qps 从 50 渐增至 500，观察吞吐饱和点。若 QPS 降 >10%，检查共享内存 bank conflict 或 L2 命中率。

3. 生产监控与调优清单

  • 性能指标：
    • TFLOPS > cuBLAS 1.1x 为阈值。
    • 内存带宽利用：nsight-compute 监控 L2 读 / 写命中率 >90%。
    • 占用率：>50% SMs 活跃，warp 效率 >70%。
  • 回滚策略：若形状未覆盖，pad 零至最近大配置（e.g., M=64 pad to 128）；若精度需求 32-bit acc，暂用 cuBLASLt。
  • 扩展参数：
    • 线程块：16x8x16（CUTLASS 默认，RL 搜索变体）。
    • 共享内存：128KB/block 上限，tiling K=64-256。
    • 多流：cublasSetStream 异步，batch_size=8 以藏延迟。

在 H100 等新架构上，预期需重训 RL 代理适应 Hopper SM90 tensor cores，tiling 策略偏向更大 block (32x16x32)。项目 To-do 包括 32-bit acc 支持与更多 GPU，落地后可集成 Triton 或 TVM 作为动态调度 fallback。

此方法证明 RL 在内核调优的潜力，尤其对长尾形状，提供 >10% 平均加速。未来结合 LLM 代码生成，可进一步自动化全栈优化。

资料来源：
CUDA-L2 GitHub
项目论文 arXiv:2512.02551

查看归档