Hotdry.
归档
compiler-design

RL驱动CUDA L2 GEMM内核自动调优：A100/H100 tile size与融合优化

基于强化学习动态优化GEMM tile sizes、融合顺序与数据布局，针对A100/H100 FP16/INT8实现超cuBLAS吞吐，支持L2缓存高效利用与多精度自适应调度。

在深度学习训练与推理中，通用矩阵乘法（GEMM）占据计算量的 70% 以上，其性能直接决定整体效率。传统 cuBLAS 虽高效，但针对特定矩阵尺寸（M,N,K）和硬件如 A100/H100 的 L2 缓存特性，手动调优 tile sizes、融合顺序与数据布局耗时巨大。deepreinforce-ai/CUDA-L2 项目引入强化学习（RL）代理，结合大语言模型自动生成并优化 HGEMM（半精度 GEMM）CUDA 内核，实现系统性超越 torch.matmul 与 cuBLAS 性能。

RL autotune 的核心在于构建高效的状态 - 动作 - 奖励框架。状态空间包括当前 tile 参数（如 block_M、block_N、block_K，通常为 16x8x16 基于 SM80 Tensor Core）、共享内存布局（swizzle 优化 L2 命中）和融合策略（e.g., async copy 与 double buffering）。动作空间通过 LLM 生成候选 CUDA 代码变体，调整参数如 warp specialization、TMA store 与 cluster shape。奖励函数融合多维指标：首要为 TFLOPS 吞吐，其次 L2 缓存命中率（目标 > 90%）、SM 利用率（>85%）与内存带宽饱和度。通过 PPO 或类似算法迭代，代理快速收敛至 Pareto 最优配置。在 A100 上，针对 1000 种 (M,N,K) 配置训练的 kernels 平均加速 cuBLAS 1.2-1.5x，峰值达 1.8x。

针对 A100/H100 的 L2 缓存（40MB/50MB），tile size 优化至关重要。大 tile 提升重用但降低 occupancy，小 tile 反之。CUDA-L2 RL 代理探索：A100 FP16 下 block_M=128、block_N=64、block_K=32（H100 相应 128/128/32），匹配 L2 扇入，提升命中率 27%。融合顺序优先 cp.async 全局→共享→ldmatrix 共享→register，再 mma.sync Tensor Core 计算，最后 store 回全局。数据布局采用 row-major A 与 col-major B，swizzle 重排避 bank 冲突。多精度自适应：FP16/INT8 切换 via 动态精度累加器（F16F16F16F16TN），INT8 下 tile_K 扩展至 64 利用更高峰值。

落地参数与监控要点如下：

  • Tile 参数推荐（A100 FP16）：block_M=128, block_N=64, block_K=32; threads=128/warp=4; stages=2-3 pipeline。
  • H100 调整：block_M/N=128, block_K=64，利用 TMA 加速 store，num_stages=3。
  • 编译旗帜：--ptxas-options=--register-usage-level=10，maxrregcount=255。
  • 监控清单：
    1. Nsight Compute：sm__throughput.avg.pct_of_peak_sustained_elapsed>85%，l1tex__data_bank_conflicts<5%。
    2. nvidia-smi dmon：L2 读 / 写命中率 > 90%，DRAM 利用 < 80%。
    3. nvprof：occupancy>50%，IPC>2。
  • 回滚策略：若 L2 miss>20%，降 tile_K 16，回退 cuBLASLt-heuristic。

集成步骤简明：

  1. git clone https://github.com/deepreinforce-ai/CUDA-L2 && git clone -b v4.2.1 https://github.com/NVIDIA/cutlass.git cutlass。
  2. export CUTLASS_DIR=./cutlass; export TORCH_CUDA_ARCH_LIST="8.0"。
  3. ./eval_one_file.sh --mnk 64_4096_64 --mode offline --gpu_device_id 0。 服务器模式加 --target_qps 100 模拟推理负载。正确性验证：randn_correctness_check.py 对比 cuBLAS 相对误差 < 1e-3。

实测 A100 上 (M=64,N=4096,K=64)，CUDA-L2 offline 模式 0.85ms/batch（cuBLAS 1.05ms），server QPS=120（cuBLAS 95）。L2 hit 率 92% vs cuBLAS 81%。该方法不限于 GEMM，可扩展 FlashAttention 等，未来支持 Hopper/Blackwell。

资料来源：deepreinforce-ai/CUDA-L2 GitHub 仓库，“CUDA-L2 systematically outperforms major matmul baselines on A100 across 1000 configurations。”

查看归档