专栏名称: GiantPandaLLM

专注于机器学习、深度学习、计算机视觉、图像处理等多个方向技术分享。团队由一群热爱技术且热衷于分享的小伙伴组成。我们坚持原创，每天一到两篇原创技术分享。希望在传播知识、分享知识的同时能够启发你，大家一起共同进步(･ω<)☆

[CUDA基础]CUDA-Learn-Notes: v3.0 大升级-面试刷题不迷路

GiantPandaLLM · 公众号 · 3D · 2025-04-04 09:00

正文

请到「今天看啥」查看全文

直接用NV的官方镜像吧，省事，推荐： http:// https://nvcr.io/nvidia/pytorch:24.10-py3 24.10-py3 ，CUDA-Learn-Notes中所有的Kernel均在该环境中测试通过。特别不推荐CUDA 11的环境，FP8以及最新的一些MMA指令，在CUDA 11没法玩。都2025年了，用CUDA 12+吧，兄弟们~

0x04 PyTorch Python bindings

本小节讲讲CUDA-Learn-Notes的workflow。CUDA-Learn-Notes将kernels按照主题进行划分，并且对每个主题的kernel实现，都进行了PyTorch python binding，可以直接使用python脚本进行性能和数值验证。比如以下这个block all reduce的示例：

//packed_type,acc_type,th_type,element_type,n_elements_per_pack,out_typeTORCH_BINDING_REDUCE(f32,              f32,torch::kFloat32,       float,              1,float)TORCH_BINDING_REDUCE(f32x4,            f32,torch::kFloat32,       float,              4,float)TORCH_BINDING_REDUCE(f16,              f16,torch::kHalf,          half,               1,float)TORCH_BINDING_REDUCE(f16,              f32,torch::kHalf,          half,               1,float)TORCH_BINDING_REDUCE(f16x2,            f16,torch::kHalf,          half,               2,float)TORCH_BINDING_REDUCE(f16x2,            f32,torch::kHalf,          half,               2,float)TORCH_BINDING_REDUCE(f16x8_pack,       f16,torch::kHalf,          half,               8,float)TORCH_BINDING_REDUCE(f16x8_pack,       f32,torch::kHalf,          half,               8,float)TORCH_BINDING_REDUCE(bf16,             bf16,torch::kBFloat16,      __nv_bfloat16,      1,float)TORCH_BINDING_REDUCE(bf16,             f32,torch::kBFloat16,      __nv_bfloat16,      1,float)TORCH_BINDING_REDUCE(bf16x2,           bf16,torch::kBFloat16,      __nv_bfloat16,      2,float)TORCH_BINDING_REDUCE(bf16x2,           f32,torch::kBFloat16,      __nv_bfloat16,      2,float)TORCH_BINDING_REDUCE(bf16x8_pack,      bf16,torch::kBFloat16,      __nv_bfloat16,      8,float)TORCH_BINDING_REDUCE(bf16x8_pack,      f32,torch::kBFloat16,      __nv_bfloat16,      8,float)TORCH_BINDING_REDUCE(fp8_e4m3,         f16,torch::kFloat8_e4m3fn,__nv_fp8_storage_t,1,float)TORCH_BINDING_REDUCE(fp8_e4m3x16_pack,f16,torch::kFloat8_e4m3fn,__nv_fp8_storage_t,16,float)TORCH_BINDING_REDUCE(fp8_e5m2,         f16,torch::kFloat8_e5m2,   __nv_fp8_storage_t,1,float)TORCH_BINDING_REDUCE(fp8_e5m2x16_pack,f16,torch::kFloat8_e5m2,   __nv_fp8_storage_t,16,float)TORCH_BINDING_REDUCE(i8,               i32,torch::kInt8,          int8_t,             1,int32_t)TORCH_BINDING_REDUCE(i8x16_pack,       i32,torch::kInt8,          int8_t,             16,int32_t)

这个实例中，包含了大部分常见的数据类型。不同于已有的CUDA笔记几乎不谈FP16和Tensor Cores的情况，CUDA-Learn-Notes在FP16和Tensor Cores上花了大量的精力，提供了如HGEMM, FlashAttention-MMA等大量的案例。并且，对于BF16/FP8等数据类型，也提供了应用案例，如本示例中的block all reduce，支持FP32/FP16/BF16/FP8/INT8等数据类型。经过了PyTorch binding后，数值结果和性能的验证非常简单，直接跑python脚本就行，不需要去配置C++编译环境。

# 只测试Ada架构 不指定默认编译所有架构 耗时较长: Volta, Ampere, Ada, Hopper, ...exportTORCH_CUDA_ARCH_LIST=Ada 
python3 block_all_reduce.py

日志输出如下：（block all reduce示例）

--------------------------------------------------------------------------------                                        S=4096, K=4096
               out_f32f32: -2295.19458008 , time:0.10227132ms
             out_f32x4f32: -2295.19702148 , time:0.03361320ms
            out_f32f32_th: -2295.19946289 , time:0.02290916ms
--------------------------------------------------------------------------------
               out_f16f16: -2293.83764648 , time:0.10097337ms
               out_f16f32: -2296.36425781 , time:0.10095334ms
             out_f16x2f32: -2297.93896484 , time:0.03533483ms
             out_f16x2f16: -2297.96386719 , time:0.03572583ms
         out_f16x8packf16: -2299.68701172 , time:0.01311255ms
         out_f16x8packf32: -2296.36645508 , time:0.01308966ms
            out_f16f16_th: -2296.00000000 , time:0.01445580ms
--------------------------------------------------------------------------------