Study note

B200 Attention Kernel from Scratch to Near-SOTA in 60 Diagrams

Properties

Type: Blogs
Status: 待读
Domain: AI / ML
Category: GPU、CUDA 与内核优化
Source: iaroslavelistratov.github.io
Vault note: library/articles/ai_ml/B200-Attention-Kernel-from-Scratch-to-Near-SOTA-in-60-Diagrams-f9c632706ec23af8.md

Summary

以 NVIDIA Blackwell B200 为目标，从一个可正确运行的 dense BF16 attention 基线开始，逐步构建 14 个 CUDA/PTX kernel，并在每一章只引入一类主要优化。内容覆盖 QK/online softmax/PV 数据流、TMA、TMEM、tcgen05 MMA、shared-memory layout 与 swizzling、warp specialization、加载和计算流水线、近似计算、寄存器缓存以及 persistent kernel。

Highlights

用 60 幅图把 B200 attention 的逻辑 tile、物理内存布局、CTA 分工和流水线重叠关系可视化；最终实现针对 4K、8K、16K 形状达到同次测试 FlashAttention-4 约 94.4% 的性能，并提供完整源码、benchmark 和接入视频生成模型的实践项目。适合已有基础 CUDA 与 attention/online softmax 概念后，用来学习 Blackwell 内核性能工程。

Comments

Loading comments...