GPU Transformer Kernel Lab

这是一个面向 AI Infra 岗位的 GPGPU 学习与项目实践仓库。目标不是堆零散 CUDA demo，而是围绕 Transformer 推理/训练中的核心算子，完成一套可 benchmark、可 profile、可接入 PyTorch 的 GPU kernel 优化项目。

项目定位

最终项目建议聚焦为：

从零实现并优化 Transformer 常用 GPU kernels，包括 reduction、RMSNorm/LayerNorm、softmax、GEMM、FlashAttention 前向或 KV-cache attention，并通过 CUDA/Triton + PyTorch extension 完成端到端 benchmark 与性能分析。

这个方向适合 AI Infra 简历，因为它能同时展示：

• GPU memory hierarchy、warp/block/thread 协作、shared memory、bank conflict、occupancy 等底层理解。
• 对 Transformer 算子的工程拆解能力。
• 性能评估能力：latency、bandwidth、TFLOPS、roofline、Nsight profiling。
• 与 PyTorch/LLM inference runtime 生态衔接的能力。

推荐目录

.
├── cuda/             # CUDA/C++ kernels and PyTorch extension sources
├── triton/           # Triton reference and optimized kernels
├── python/           # Python package, correctness checks, wrappers
├── benchmarks/       # Benchmark harness and profiling scripts
├── tests/            # Unit tests and numerical correctness checks
├── docs/             # Learning notes, profiling reports, design docs
└── third_party/      # Optional external references, kept minimal

里程碑

1. 建立 benchmark 基础设施：正确性测试、计时、带宽/吞吐计算、profiling 脚本。
2. 实现 memory-bound kernels：vector add、reduce、softmax、RMSNorm/LayerNorm。
3. 实现 compute-bound kernels：naive GEMM、tiled GEMM、register blocking、Tensor Core/WMMA 版本。
4. 实现 Transformer 关键融合算子：fused bias+activation、fused RMSNorm、FlashAttention 或 KV-cache attention。
5. 接入 PyTorch extension/Triton，对比 PyTorch、cuBLAS、FlashAttention 等 baseline。
6. 输出 profiling report 和 resume-ready 项目总结。

目标成果

完成后应当能在简历中写成类似：

Built a CUDA/Triton Transformer kernel optimization lab, implementing optimized GEMM, RMSNorm, softmax and attention kernels with PyTorch extension integration. Achieved measurable speedups over naive baselines through tiling, shared-memory reuse, warp-level reductions and Tensor Core utilization; analyzed bottlenecks with Nsight Compute and roofline modeling.

具体百分比和吞吐数据需要在真实 GPU 上跑 benchmark 后填写。

下一步

从 docs/ROADMAP.md 开始。第一阶段建议先完成环境、benchmark harness 和第一个 memory bandwidth kernel，因为后续所有优化都依赖可靠的测量方法。