
\documentclass[a4paper]{article}
\usepackage[fontset=fandol]{ctex}
\usepackage{amsmath,amssymb}
\usepackage{xcolor}
\usepackage{graphicx}
\usepackage[margin=2.15cm]{geometry}
\usepackage[most]{tcolorbox}
\usepackage{etoolbox}
\usepackage{listings}
\usepackage{hyperref}
\usepackage{booktabs}
\usepackage{float}
\usepackage{array}
\usepackage{longtable}
\usepackage{needspace}
\graphicspath{{./}{slides-images/}}
\hypersetup{colorlinks=true,linkcolor=blue!55!black,urlcolor=blue!65!black,citecolor=blue!55!black}
\newtcolorbox{knowledgebox}[1]{enhanced, colback=blue!5!white, colframe=blue!75!black, colbacktitle=blue!75!black,coltitle=white, fonttitle=\bfseries, title=#1, boxrule=1pt, sharp corners}
\newtcolorbox{importantbox}[1]{enhanced, colback=yellow!10!white, colframe=yellow!80!black, colbacktitle=yellow!80!black,coltitle=black, fonttitle=\bfseries, title=#1, sharp corners}
\newtcolorbox{warningbox}[1]{enhanced, colback=red!5!white, colframe=red!75!black, colbacktitle=red!75!black,coltitle=white, fonttitle=\bfseries, title=#1, sharp corners}
\lstset{language=Python,basicstyle=\ttfamily\small,keywordstyle=\color{blue!70!black},stringstyle=\color{red!60!black},commentstyle=\color{green!50!black},breaklines=true,frame=single,numbers=left,numberstyle=\tiny\color{gray},captionpos=t,showstringspaces=false,columns=fullflexible,keepspaces=true,extendedchars=false}
\BeforeBeginEnvironment{lstlisting}{\Needspace{12\baselineskip}}
\newcommand{\notetitle}{CS336 2026 Lecture 5：GPUs、Roofline 与 FlashAttention}
\newcommand{\noteauthors}{基于 Stanford CS336 Spring 2026 官方讲义整理}
\newcommand{\notedate}{2026 年春季}
\newcommand{\videochannel}{Stanford CS336}
\newcommand{\videopublishdate}{Spring 2026}
\newcommand{\videoduration}{官方幻灯片 lecture05-slides.pdf}
\newcommand{\videourl}{https://cs336.stanford.edu/}
\newcommand{\videocoverpath}{slides-images/slide-001.jpg}
\newcommand{\repourl}{https://github.com/hqhq1025/ai-course-notes}
\newcommand{\figsource}[1]{\par{\footnotesize\emph{来源：#1}}\normalsize}
\usepackage{fancyhdr}
\pagestyle{fancy}
\fancyhf{}
\fancyfoot[C]{\thepage}
\fancyfoot[R]{\footnotesize\color{gray}\href{\repourl}{AI Course Notes}}
\renewcommand{\headrulewidth}{0pt}
\renewcommand{\footrulewidth}{0pt}
\begin{document}
\begin{titlepage}\centering
{\Large 课程笔记\par}\vspace{1.2cm}{\huge\bfseries \notetitle\par}\vspace{0.8cm}{\large \noteauthors\par}\vspace{0.3cm}{\large \notedate\par}\vspace{1.2cm}
\includegraphics[width=0.82\textwidth,height=0.45\textheight,keepaspectratio]{\videocoverpath}\par
\vfill\begin{tcolorbox}[width=0.9\textwidth, colback=black!2!white, colframe=black!60, sharp corners]
\textbf{课程/频道}：\videochannel\par\textbf{发布日期}：\videopublishdate\par\textbf{材料类型}：\videoduration\par\textbf{课程主页}：\href{\videourl}{\nolinkurl{\videourl}}\par\textbf{项目主页}：\href{\repourl}{\nolinkurl{\repourl}}\par
\end{tcolorbox}\end{titlepage}\tableofcontents\newpage

\section{本讲的问题：GPU 为什么既快又难用？}
Lecture 5 把 Lecture 2 的 resource accounting 落到硬件：为什么 GPU 能支撑 scaling？为什么同样是 FLOPs，有的操作飞快，有的操作很慢？为什么 FlashAttention 这种算法能改变 attention 的实际速度？
\begin{knowledgebox}{术语消化：GPU memory hierarchy}
HBM 是 High Bandwidth Memory，GPU 上的大容量高带宽显存，容量大但离计算单元远；SRAM 是 Static Random-Access Memory，常对应片上 shared memory/cache，容量小但速度快；register 是每个 thread 私有的最快存储。GPU 优化的核心之一，就是让数据从 HBM 读进来以后尽量在 SRAM/register 附近复用，而不是频繁写回 HBM。
\end{knowledgebox}
\begin{importantbox}{本讲主线}
GPU performance 的核心不是“GPU 很快”四个字，而是让数据移动、并行执行和矩阵硬件匹配。低精度、fusion、recomputation、coalescing、tiling、FlashAttention 都是在减少 HBM 往返或提高数据复用。
\end{importantbox}
\section{GPU 为什么适合大规模深度学习}

本节从 compute scaling 的历史讲到 GPU 与 CPU 的区别，再拆 SM、SP、warp、block、register、shared memory 和 HBM。


\begin{knowledgebox}{术语消化：SM、SP、warp、block}
SM 是 Streaming Multiprocessor，是 GPU 上执行 thread blocks 的主要单元；SP/CUDA core 是 SM 内的标量执行单元；warp 是一组通常 32 个 lockstep 执行的 threads；thread block 是被调度到一个 SM 上的线程组。HBM 是 GPU 的高带宽显存，容量大但离计算单元远；SRAM/shared memory/register 更靠近 SM，容量小但快。
\end{knowledgebox}

\begin{figure}[H]
\centering
\includegraphics[width=0.88\textwidth]{slide-001.jpg}
\caption{Slide 1: Lecture 5}
\figsource{CS336 Spring 2026 Lecture 5 官方幻灯片。}
\end{figure}

\noindent\textbf{展开说明：}GPUS


\begin{figure}[H]
\centering
\includegraphics[width=0.88\textwidth]{slide-002.jpg}
\caption{Slide 2: Outline and goals}
\figsource{CS336 Spring 2026 Lecture 5 官方幻灯片。}
\end{figure}

\noindent\textbf{展开说明：} Make CUDA and GPUs less magic


\begin{figure}[H]
\centering
\includegraphics[width=0.88\textwidth]{slide-003.jpg}
\caption{Slide 3: Before we start..}
\figsource{CS336 Spring 2026 Lecture 5 官方幻灯片。}
\end{figure}

\noindent\textbf{展开说明：}Substantial credit goes to a few sources that I’d like to highlight..


\begin{figure}[H]
\centering
\includegraphics[width=0.88\textwidth]{slide-004.jpg}
\caption{Slide 4: Organization today:}
\figsource{CS336 Spring 2026 Lecture 5 官方幻灯片。}
\end{figure}

\noindent\textbf{展开说明：} Part 1: GPUs in depth – how they work and important parts


\begin{figure}[H]
\centering
\includegraphics[width=0.88\textwidth]{slide-005.jpg}
\caption{Slide 5: Setting the stage: compute leads to predictable perf}
\figsource{CS336 Spring 2026 Lecture 5 官方幻灯片。}
\end{figure}

\noindent\textbf{展开说明：}Often times, compute leads to predictable performance gains for language models


\begin{figure}[H]
\centering
\includegraphics[width=0.88\textwidth]{slide-006.jpg}
\caption{Slide 6: How do we get compute scaling? Early on – Dennard scaing}
\figsource{CS336 Spring 2026 Lecture 5 官方幻灯片。}
\end{figure}

\noindent\textbf{展开说明：}But the traditional form of scaling (Dennard scaling) from 1980-2000s has tapped out.


\begin{figure}[H]
\centering
\includegraphics[width=0.88\textwidth]{slide-007.jpg}
\caption{Slide 7: Parallel scaling continues}
\figsource{CS336 Spring 2026 Lecture 5 官方幻灯片。}
\end{figure}

\noindent\textbf{展开说明：}Bill dally, HotChips keynote

\begin{knowledgebox}{读图：Slide 7 应该怎么看}
这页是硬件机制或性能证据页。先看数据从哪里来、在哪里复用、是否写回 HBM；再判断瓶颈是 compute、memory bandwidth、thread scheduling 还是 alignment。GPU 优化不是让公式更漂亮，而是让数据在正确层级停留更久、让 tensor cores 少等数据。
\end{knowledgebox}


\begin{figure}[H]
\centering
\includegraphics[width=0.88\textwidth]{slide-008.jpg}
\caption{Slide 8: How is a GPU different from a CPU?}
\figsource{CS336 Spring 2026 Lecture 5 官方幻灯片。}
\end{figure}

\noindent\textbf{展开说明：}CPUs optimize for a few, fast threads while GPUs optimize for many many threads

\begin{knowledgebox}{读图：Slide 8 应该怎么看}
这页是硬件机制或性能证据页。先看数据从哪里来、在哪里复用、是否写回 HBM；再判断瓶颈是 compute、memory bandwidth、thread scheduling 还是 alignment。GPU 优化不是让公式更漂亮，而是让数据在正确层级停留更久、让 tensor cores 少等数据。
\end{knowledgebox}


\begin{figure}[H]
\centering
\includegraphics[width=0.88\textwidth]{slide-009.jpg}
\caption{Slide 9: Anatomy of a GPU (execution units)}
\figsource{CS336 Spring 2026 Lecture 5 官方幻灯片。}
\end{figure}

\noindent\textbf{展开说明：}Each SM further contains many SPs   GPUs have many SM (streaming

\begin{knowledgebox}{读图：Slide 9 应该怎么看}
这页是硬件机制或性能证据页。先看数据从哪里来、在哪里复用、是否写回 HBM；再判断瓶颈是 compute、memory bandwidth、thread scheduling 还是 alignment。GPU 优化不是让公式更漂亮，而是让数据在正确层级停留更久、让 tensor cores 少等数据。
\end{knowledgebox}


\begin{figure}[H]
\centering
\includegraphics[width=0.88\textwidth]{slide-010.jpg}
\caption{Slide 10: Anatomy of a GPU (memory)}
\figsource{CS336 Spring 2026 Lecture 5 官方幻灯片。}
\end{figure}

\noindent\textbf{展开说明：}The closer the memory to the SM, the faster it is – L1 and shared memory is inside

\begin{knowledgebox}{读图：Slide 10 应该怎么看}
这页是硬件机制或性能证据页。先看数据从哪里来、在哪里复用、是否写回 HBM；再判断瓶颈是 compute、memory bandwidth、thread scheduling 还是 alignment。GPU 优化不是让公式更漂亮，而是让数据在正确层级停留更久、让 tensor cores 少等数据。
\end{knowledgebox}


\begin{figure}[H]
\centering
\includegraphics[width=0.88\textwidth]{slide-011.jpg}
\caption{Slide 11: Execution model of a GPU}
\figsource{CS336 Spring 2026 Lecture 5 官方幻灯片。}
\end{figure}

\noindent\textbf{展开说明：}There are 3 important players in the execution model

\begin{knowledgebox}{读图：Slide 11 应该怎么看}
这页是硬件机制或性能证据页。先看数据从哪里来、在哪里复用、是否写回 HBM；再判断瓶颈是 compute、memory bandwidth、thread scheduling 还是 alignment。GPU 优化不是让公式更漂亮，而是让数据在正确层级停留更久、让 tensor cores 少等数据。
\end{knowledgebox}


\begin{figure}[H]
\centering
\includegraphics[width=0.88\textwidth]{slide-012.jpg}
\caption{Slide 12: Memory model of a GPU}
\figsource{CS336 Spring 2026 Lecture 5 官方幻灯片。}
\end{figure}

\noindent\textbf{展开说明：}Each thread can access its own register, and shared memory within the block.


\begin{figure}[H]
\centering
\includegraphics[width=0.88\textwidth]{slide-013.jpg}
\caption{Slide 13: Side thread – What about TPUs?}
\figsource{CS336 Spring 2026 Lecture 5 官方幻灯片。}
\end{figure}

\noindent\textbf{展开说明：}GPUs, TPUs, and many other accelerators are at a high level, similar


\begin{figure}[H]
\centering
\includegraphics[width=0.88\textwidth]{slide-014.jpg}
\caption{Slide 14: Side thread – What about TPUs?}
\figsource{CS336 Spring 2026 Lecture 5 官方幻灯片。}
\end{figure}

\noindent\textbf{展开说明：}Core structure – lightweight control, fast (big) matmul unit, fast memory.


\begin{figure}[H]
\centering
\includegraphics[width=0.88\textwidth]{slide-015.jpg}
\caption{Slide 15: Strengths of the GPU model}
\figsource{CS336 Spring 2026 Lecture 5 官方幻灯片。}
\end{figure}

\noindent\textbf{展开说明：} Easily scales up hard workloads (by adding more SMs)


\begin{figure}[H]
\centering
\includegraphics[width=0.88\textwidth]{slide-016.jpg}
\caption{Slide 16: GPUs as fast matrix multipliers}
\figsource{CS336 Spring 2026 Lecture 5 官方幻灯片。}
\end{figure}

\noindent\textbf{展开说明：}Early days of NVIDIA GPUs – programmable shaders. Researchers hacked this to do matmuls

\begin{knowledgebox}{读图：Slide 16 应该怎么看}
这页是硬件机制或性能证据页。先看数据从哪里来、在哪里复用、是否写回 HBM；再判断瓶颈是 compute、memory bandwidth、thread scheduling 还是 alignment。GPU 优化不是让公式更漂亮，而是让数据在正确层级停留更久、让 tensor cores 少等数据。
\end{knowledgebox}


\begin{figure}[H]
\centering
\includegraphics[width=0.88\textwidth]{slide-017.jpg}
\caption{Slide 17: New matmul hardware means matmuls are fast and special}
\figsource{CS336 Spring 2026 Lecture 5 官方幻灯片。}
\end{figure}

\noindent\textbf{展开说明：}Tensor cores (introduced in V, T series) are specialized matrix multiplication circuits.

\begin{knowledgebox}{读图：Slide 17 应该怎么看}
这页是硬件机制或性能证据页。先看数据从哪里来、在哪里复用、是否写回 HBM；再判断瓶颈是 compute、memory bandwidth、thread scheduling 还是 alignment。GPU 优化不是让公式更漂亮，而是让数据在正确层级停留更久、让 tensor cores 少等数据。
\end{knowledgebox}


\begin{figure}[H]
\centering
\includegraphics[width=0.88\textwidth]{slide-018.jpg}
\caption{Slide 18: Compute scaling is faster than memory scaling}
\figsource{CS336 Spring 2026 Lecture 5 官方幻灯片。}
\end{figure}

\noindent\textbf{展开说明：}https://medium.com/riselab/ai-and-memory-wall-2cb4265cb0b8

\subsection{本章小结}
本节的共同问题是如何减少昂贵的数据移动并提高硬件利用率。GPU 的速度来自海量并行和专用矩阵硬件，但只有当 memory access、shape、precision 和 kernel 组织匹配时，这些速度才会真正出现。

\section{Roofline 与低精度：让 ML workload 跑快}

本节解释 GPU 性能模型、divergence、低精度如何提高 arithmetic intensity。随后再看 FP8、MXFP8、MXFP4 的实际含义。


\begin{importantbox}{性能公式：roofline 心智模型}
\[
\text{arithmetic intensity}
=
\frac{\text{FLOPs}}{\text{bytes moved}},
\]
\[
\text{attainable FLOP/s}
\le
\min(\text{peak FLOP/s},\text{bandwidth}\times I).
\]
低精度减少 bytes moved，也可能启用 tensor cores，因此同时影响 memory traffic 和 peak compute。
\end{importantbox}

\begin{figure}[H]
\centering
\includegraphics[width=0.88\textwidth]{slide-019.jpg}
\caption{Slide 19: Recap: GPUs – what are they and how do they work}
\figsource{CS336 Spring 2026 Lecture 5 官方幻灯片。}
\end{figure}

\noindent\textbf{展开说明：} GPUs are massively parallel – same instructions


\begin{figure}[H]
\centering
\includegraphics[width=0.88\textwidth]{slide-020.jpg}
\caption{Slide 20: Part 2: Making ML workloads fast on a GPU}
\figsource{CS336 Spring 2026 Lecture 5 官方幻灯片。}
\end{figure}

\noindent\textbf{展开说明：}Performance on a GPU can be complex, even for something as simple as a square matmul

\begin{knowledgebox}{读图：Slide 20 应该怎么看}
这页是硬件机制或性能证据页。先看数据从哪里来、在哪里复用、是否写回 HBM；再判断瓶颈是 compute、memory bandwidth、thread scheduling 还是 alignment。GPU 优化不是让公式更漂亮，而是让数据在正确层级停留更久、让 tensor cores 少等数据。
\end{knowledgebox}


\begin{figure}[H]
\centering
\includegraphics[width=0.88\textwidth]{slide-021.jpg}
\caption{Slide 21: What makes ML workloads fast?}
\figsource{CS336 Spring 2026 Lecture 5 官方幻灯片。}
\end{figure}

\noindent\textbf{展开说明：}The roofline model

\begin{knowledgebox}{读图：Slide 21 应该怎么看}
这页是硬件机制或性能证据页。先看数据从哪里来、在哪里复用、是否写回 HBM；再判断瓶颈是 compute、memory bandwidth、thread scheduling 还是 alignment。GPU 优化不是让公式更漂亮，而是让数据在正确层级停留更久、让 tensor cores 少等数据。
\end{knowledgebox}


\begin{figure}[H]
\centering
\includegraphics[width=0.88\textwidth]{slide-022.jpg}
\caption{Slide 22: How do we make GPUs go fast?}
\figsource{CS336 Spring 2026 Lecture 5 官方幻灯片。}
\end{figure}

\noindent\textbf{展开说明：}1. Control divergence (not a memory bottleneck..)


\begin{figure}[H]
\centering
\includegraphics[width=0.88\textwidth]{slide-023.jpg}
\caption{Slide 23: Control divergence (not a memory issue)}
\figsource{CS336 Spring 2026 Lecture 5 官方幻灯片。}
\end{figure}

\noindent\textbf{展开说明：}GPUs operate in a SIMT model – every thread in a warp is executing the same instruction

\begin{knowledgebox}{读图：Slide 23 应该怎么看}
这页是硬件机制或性能证据页。先看数据从哪里来、在哪里复用、是否写回 HBM；再判断瓶颈是 compute、memory bandwidth、thread scheduling 还是 alignment。GPU 优化不是让公式更漂亮，而是让数据在正确层级停留更久、让 tensor cores 少等数据。
\end{knowledgebox}


\begin{figure}[H]
\centering
\includegraphics[width=0.88\textwidth]{slide-024.jpg}
\caption{Slide 24: Trick 1: Low precision computation}
\figsource{CS336 Spring 2026 Lecture 5 官方幻灯片。}
\end{figure}

\noindent\textbf{展开说明：}If you have fewer bits, you have fewer bits to move

\begin{knowledgebox}{读图：Slide 24 应该怎么看}
这页是硬件机制或性能证据页。先看数据从哪里来、在哪里复用、是否写回 HBM；再判断瓶颈是 compute、memory bandwidth、thread scheduling 还是 alignment。GPU 优化不是让公式更漂亮，而是让数据在正确层级停留更久、让 tensor cores 少等数据。
\end{knowledgebox}


\begin{figure}[H]
\centering
\includegraphics[width=0.88\textwidth]{slide-025.jpg}
\caption{Slide 25: Low precision improves arithmetic intensity}
\figsource{CS336 Spring 2026 Lecture 5 官方幻灯片。}
\end{figure}

\noindent\textbf{展开说明：}Example: elementwise ReLU ($x$ = max(0, $x$)) on a vector of size $n$.

\begin{knowledgebox}{读图：Slide 25 应该怎么看}
这页是硬件机制或性能证据页。先看数据从哪里来、在哪里复用、是否写回 HBM；再判断瓶颈是 compute、memory bandwidth、thread scheduling 还是 alignment。GPU 优化不是让公式更漂亮，而是让数据在正确层级停留更久、让 tensor cores 少等数据。
\end{knowledgebox}


\begin{figure}[H]
\centering
\includegraphics[width=0.88\textwidth]{slide-026.jpg}
\caption{Slide 26: Low precision drives faster matrix multiplies}
\figsource{CS336 Spring 2026 Lecture 5 官方幻灯片。}
\end{figure}

\noindent\textbf{展开说明：}Lots of operations in modern GPUs are accelerated via low / mixed precision operations

\begin{knowledgebox}{读图：Slide 26 应该怎么看}
这页是硬件机制或性能证据页。先看数据从哪里来、在哪里复用、是否写回 HBM；再判断瓶颈是 compute、memory bandwidth、thread scheduling 还是 alignment。GPU 优化不是让公式更漂亮，而是让数据在正确层级停留更久、让 tensor cores 少等数据。
\end{knowledgebox}


\begin{figure}[H]
\centering
\includegraphics[width=0.88\textwidth]{slide-027.jpg}
\caption{Slide 27: Frontiers in low precision}
\figsource{CS336 Spring 2026 Lecture 5 官方幻灯片。}
\end{figure}

\noindent\textbf{展开说明：}Very low precision (FP8) with different tradeoffs                 Multiple scaling factors MXFP8 (Blackwell)


\begin{figure}[H]
\centering
\includegraphics[width=0.88\textwidth]{slide-028.jpg}
\caption{Slide 28: MXFP8 training in practice}
\figsource{CS336 Spring 2026 Lecture 5 官方幻灯片。}
\end{figure}

\noindent\textbf{展开说明：}Notice – not all weights in MXFP8, transposes also separately quantized

\subsection{本章小结}
本节的共同问题是如何减少昂贵的数据移动并提高硬件利用率。GPU 的速度来自海量并行和专用矩阵硬件，但只有当 memory access、shape、precision 和 kernel 组织匹配时，这些速度才会真正出现。

\section{Operator fusion 与 recomputation：少搬数据}

本节用工厂/仓库类比说明 HBM 往返为什么贵，再解释 fusion 和 recomputation 如何减少 memory access。


\begin{knowledgebox}{first-use glossary：operator fusion 与 recomputation}
Operator fusion 把多个小 kernel 合成一个 kernel，让中间结果留在寄存器/shared memory 中，减少 HBM 读写。Recomputation 是丢弃某些中间 activations，在需要时重算，用额外 compute 换更低 memory traffic 或 peak memory。
\end{knowledgebox}

\begin{figure}[H]
\centering
\includegraphics[width=0.88\textwidth]{slide-029.jpg}
\caption{Slide 29: Frontiers in low precision}
\figsource{CS336 Spring 2026 Lecture 5 官方幻灯片。}
\end{figure}

\noindent\textbf{展开说明：}MXFP4..

\begin{knowledgebox}{读图：Slide 29 应该怎么看}
这页是硬件机制或性能证据页。先看数据从哪里来、在哪里复用、是否写回 HBM；再判断瓶颈是 compute、memory bandwidth、thread scheduling 还是 alignment。GPU 优化不是让公式更漂亮，而是让数据在正确层级停留更久、让 tensor cores 少等数据。
\end{knowledgebox}


\begin{figure}[H]
\centering
\includegraphics[width=0.88\textwidth]{slide-030.jpg}
\caption{Slide 30: Trick 2: Operator fusion}
\figsource{CS336 Spring 2026 Lecture 5 官方幻灯片。}
\end{figure}

\noindent\textbf{展开说明：}Think of a GPU like a factory – inputs come from a warehouse (memory) and is

\begin{knowledgebox}{读图：Slide 30 应该怎么看}
这页是硬件机制或性能证据页。先看数据从哪里来、在哪里复用、是否写回 HBM；再判断瓶颈是 compute、memory bandwidth、thread scheduling 还是 alignment。GPU 优化不是让公式更漂亮，而是让数据在正确层级停留更久、让 tensor cores 少等数据。
\end{knowledgebox}


\begin{figure}[H]
\centering
\includegraphics[width=0.88\textwidth]{slide-031.jpg}
\caption{Slide 31: Operator fusion to minimize memory access}
\figsource{CS336 Spring 2026 Lecture 5 官方幻灯片。}
\end{figure}

\noindent\textbf{展开说明：}What if we have to do many operations? Shipping back and forth is somewhat silly

\begin{knowledgebox}{读图：Slide 31 应该怎么看}
这页是硬件机制或性能证据页。先看数据从哪里来、在哪里复用、是否写回 HBM；再判断瓶颈是 compute、memory bandwidth、thread scheduling 还是 alignment。GPU 优化不是让公式更漂亮，而是让数据在正确层级停留更久、让 tensor cores 少等数据。
\end{knowledgebox}


\begin{figure}[H]
\centering
\includegraphics[width=0.88\textwidth]{slide-032.jpg}
\caption{Slide 32: Example – sines and cosines}
\figsource{CS336 Spring 2026 Lecture 5 官方幻灯片。}
\end{figure}

\noindent\textbf{展开说明：}Computing sin 2 $x$ + cos 2 $x$ naively launches 5 CUDA kernels (back and forth)

\begin{knowledgebox}{读图：Slide 32 应该怎么看}
这页是硬件机制或性能证据页。先看数据从哪里来、在哪里复用、是否写回 HBM；再判断瓶颈是 compute、memory bandwidth、thread scheduling 还是 alignment。GPU 优化不是让公式更漂亮，而是让数据在正确层级停留更久、让 tensor cores 少等数据。
\end{knowledgebox}


\begin{figure}[H]
\centering
\includegraphics[width=0.88\textwidth]{slide-033.jpg}
\caption{Slide 33: Fusion example}
\figsource{CS336 Spring 2026 Lecture 5 官方幻灯片。}
\end{figure}

\noindent\textbf{展开说明：}All 5 pointwise operations can be fused into a single CUDA kernel call.


\begin{figure}[H]
\centering
\includegraphics[width=0.88\textwidth]{slide-034.jpg}
\caption{Slide 34: Trick 3: recomputation}
\figsource{CS336 Spring 2026 Lecture 5 官方幻灯片。}
\end{figure}

\noindent\textbf{展开说明：}[From cs221]

\begin{knowledgebox}{读图：Slide 34 应该怎么看}
这页是硬件机制或性能证据页。先看数据从哪里来、在哪里复用、是否写回 HBM；再判断瓶颈是 compute、memory bandwidth、thread scheduling 还是 alignment。GPU 优化不是让公式更漂亮，而是让数据在正确层级停留更久、让 tensor cores 少等数据。
\end{knowledgebox}


\begin{figure}[H]
\centering
\includegraphics[width=0.88\textwidth]{slide-035.jpg}
\caption{Slide 35: Storing (and retrieving) activations can be expensive!}
\figsource{CS336 Spring 2026 Lecture 5 官方幻灯片。}
\end{figure}

\noindent\textbf{展开说明：}Let’s say we stack 3 sigmoids on top of each other.

\begin{knowledgebox}{读图：Slide 35 应该怎么看}
这页是硬件机制或性能证据页。先看数据从哪里来、在哪里复用、是否写回 HBM；再判断瓶颈是 compute、memory bandwidth、thread scheduling 还是 alignment。GPU 优化不是让公式更漂亮，而是让数据在正确层级停留更久、让 tensor cores 少等数据。
\end{knowledgebox}

\subsection{本章小结}
本节的共同问题是如何减少昂贵的数据移动并提高硬件利用率。GPU 的速度来自海量并行和专用矩阵硬件，但只有当 memory access、shape、precision 和 kernel 组织匹配时，这些速度才会真正出现。

\section{Coalescing、tiling 与矩阵性能异常}

本节解释 DRAM burst、memory coalescing、row-major matmul、tiling、alignment、wave quantization 和 matrix mystery。


\begin{knowledgebox}{first-use glossary：coalescing、tiling、wave quantization}
Memory coalescing 指同一个 warp 的 memory accesses 落在连续 burst 中，从而合并成高效访问。Tiling 把大矩阵切成小块放入 shared memory 重复使用。Wave quantization 指 thread blocks 数量不能完美填满 SM waves 时产生的周期性利用率波动。
\end{knowledgebox}

\begin{figure}[H]
\centering
\includegraphics[width=0.88\textwidth]{slide-036.jpg}
\caption{Slide 36: Throw away the activations, re-compute them!}
\figsource{CS336 Spring 2026 Lecture 5 官方幻灯片。}
\end{figure}

\noindent\textbf{展开说明：}Throwing away computation can actually be optimal, w/ 5/8 th the memory accesses!

\begin{knowledgebox}{读图：Slide 36 应该怎么看}
这页是硬件机制或性能证据页。先看数据从哪里来、在哪里复用、是否写回 HBM；再判断瓶颈是 compute、memory bandwidth、thread scheduling 还是 alignment。GPU 优化不是让公式更漂亮，而是让数据在正确层级停留更久、让 tensor cores 少等数据。
\end{knowledgebox}


\begin{figure}[H]
\centering
\includegraphics[width=0.88\textwidth]{slide-037.jpg}
\caption{Slide 37: Trick (?) 4: Memory coalescing and DRAM}
\figsource{CS336 Spring 2026 Lecture 5 官方幻灯片。}
\end{figure}

\noindent\textbf{展开说明：}DRAM (global memory) is read in ‘burst mode’ – each read gives you many bytes!

\begin{knowledgebox}{读图：Slide 37 应该怎么看}
这页是硬件机制或性能证据页。先看数据从哪里来、在哪里复用、是否写回 HBM；再判断瓶颈是 compute、memory bandwidth、thread scheduling 还是 alignment。GPU 优化不是让公式更漂亮，而是让数据在正确层级停留更久、让 tensor cores 少等数据。
\end{knowledgebox}


\begin{figure}[H]
\centering
\includegraphics[width=0.88\textwidth]{slide-038.jpg}
\caption{Slide 38: Memory coalescing}
\figsource{CS336 Spring 2026 Lecture 5 官方幻灯片。}
\end{figure}

\noindent\textbf{展开说明：}Memory accesses are coalesced if all the threads (in a warp) fall within the same burst

\begin{knowledgebox}{读图：Slide 38 应该怎么看}
这页是硬件机制或性能证据页。先看数据从哪里来、在哪里复用、是否写回 HBM；再判断瓶颈是 compute、memory bandwidth、thread scheduling 还是 alignment。GPU 优化不是让公式更漂亮，而是让数据在正确层级停留更久、让 tensor cores 少等数据。
\end{knowledgebox}


\begin{figure}[H]
\centering
\includegraphics[width=0.88\textwidth]{slide-039.jpg}
\caption{Slide 39: Coalescing for matrix multiplication}
\figsource{CS336 Spring 2026 Lecture 5 官方幻灯片。}
\end{figure}

\noindent\textbf{展开说明：}For row-major matrices – threads that move along rows are not coalesced

\begin{knowledgebox}{读图：Slide 39 应该怎么看}
这页是硬件机制或性能证据页。先看数据从哪里来、在哪里复用、是否写回 HBM；再判断瓶颈是 compute、memory bandwidth、thread scheduling 还是 alignment。GPU 优化不是让公式更漂亮，而是让数据在正确层级停留更久、让 tensor cores 少等数据。
\end{knowledgebox}


\begin{figure}[H]
\centering
\includegraphics[width=0.88\textwidth]{slide-040.jpg}
\caption{Slide 40: Trick 5 (the big one): tiling}
\figsource{CS336 Spring 2026 Lecture 5 官方幻灯片。}
\end{figure}

\noindent\textbf{展开说明：}Tiling is the idea of grouping and ordering threads to minimize global memory access.

\begin{knowledgebox}{读图：Slide 40 应该怎么看}
这页是硬件机制或性能证据页。先看数据从哪里来、在哪里复用、是否写回 HBM；再判断瓶颈是 compute、memory bandwidth、thread scheduling 还是 alignment。GPU 优化不是让公式更漂亮，而是让数据在正确层级停留更久、让 tensor cores 少等数据。
\end{knowledgebox}


\begin{figure}[H]
\centering
\includegraphics[width=0.88\textwidth]{slide-041.jpg}
\caption{Slide 41: Tiling – store and reuse information in shared memory}
\figsource{CS336 Spring 2026 Lecture 5 官方幻灯片。}
\end{figure}

\noindent\textbf{展开说明：}Cut up the matrix into smaller ‘tiles’, and load this into shared memory

\begin{knowledgebox}{读图：Slide 41 应该怎么看}
这页是硬件机制或性能证据页。先看数据从哪里来、在哪里复用、是否写回 HBM；再判断瓶颈是 compute、memory bandwidth、thread scheduling 还是 alignment。GPU 优化不是让公式更漂亮，而是让数据在正确层级停留更久、让 tensor cores 少等数据。
\end{knowledgebox}


\begin{figure}[H]
\centering
\includegraphics[width=0.88\textwidth]{slide-042.jpg}
\caption{Slide 42: Tiling math}
\figsource{CS336 Spring 2026 Lecture 5 官方幻灯片。}
\end{figure}

\noindent\textbf{展开说明：}Non-tiled matrix multiply: each input is read $N$ times from global memory

\begin{knowledgebox}{读图：Slide 42 应该怎么看}
这页是硬件机制或性能证据页。先看数据从哪里来、在哪里复用、是否写回 HBM；再判断瓶颈是 compute、memory bandwidth、thread scheduling 还是 alignment。GPU 优化不是让公式更漂亮，而是让数据在正确层级停留更久、让 tensor cores 少等数据。
\end{knowledgebox}


\begin{figure}[H]
\centering
\includegraphics[width=0.88\textwidth]{slide-043.jpg}
\caption{Slide 43: Complexities with tiling}
\figsource{CS336 Spring 2026 Lecture 5 官方幻灯片。}
\end{figure}

\noindent\textbf{展开说明：}Tile sizes may not divide the matrix size and lead to low utilization

\begin{knowledgebox}{读图：Slide 43 应该怎么看}
这页是硬件机制或性能证据页。先看数据从哪里来、在哪里复用、是否写回 HBM；再判断瓶颈是 compute、memory bandwidth、thread scheduling 还是 alignment。GPU 优化不是让公式更漂亮，而是让数据在正确层级停留更久、让 tensor cores 少等数据。
\end{knowledgebox}


\begin{figure}[H]
\centering
\includegraphics[width=0.88\textwidth]{slide-044.jpg}
\caption{Slide 44: Complexities with tiling 2 – memory alignment}
\figsource{CS336 Spring 2026 Lecture 5 官方幻灯片。}
\end{figure}

\noindent\textbf{展开说明：}Memory comes in bursts

\begin{knowledgebox}{读图：Slide 44 应该怎么看}
这页是硬件机制或性能证据页。先看数据从哪里来、在哪里复用、是否写回 HBM；再判断瓶颈是 compute、memory bandwidth、thread scheduling 还是 alignment。GPU 优化不是让公式更漂亮，而是让数据在正确层级停留更久、让 tensor cores 少等数据。
\end{knowledgebox}


\begin{figure}[H]
\centering
\includegraphics[width=0.88\textwidth]{slide-045.jpg}
\caption{Slide 45: Putting it together: understanding a matrix mystery}
\figsource{CS336 Spring 2026 Lecture 5 官方幻灯片。}
\end{figure}

\noindent\textbf{展开说明：}Why is it faster to have bigger matrices?

\begin{knowledgebox}{读图：Slide 45 应该怎么看}
这页是硬件机制或性能证据页。先看数据从哪里来、在哪里复用、是否写回 HBM；再判断瓶颈是 compute、memory bandwidth、thread scheduling 还是 alignment。GPU 优化不是让公式更漂亮，而是让数据在正确层级停留更久、让 tensor cores 少等数据。
\end{knowledgebox}


\begin{figure}[H]
\centering
\includegraphics[width=0.88\textwidth]{slide-046.jpg}
\caption{Slide 46: Matrix mystery}
\figsource{CS336 Spring 2026 Lecture 5 官方幻灯片。}
\end{figure}

\noindent\textbf{展开说明：}We understand some of this (compute intensity, tiling). Let’s take a closer look..

\begin{knowledgebox}{读图：Slide 46 应该怎么看}
这页是硬件机制或性能证据页。先看数据从哪里来、在哪里复用、是否写回 HBM；再判断瓶颈是 compute、memory bandwidth、thread scheduling 还是 alignment。GPU 优化不是让公式更漂亮，而是让数据在正确层级停留更久、让 tensor cores 少等数据。
\end{knowledgebox}


\begin{figure}[H]
\centering
\includegraphics[width=0.88\textwidth]{slide-047.jpg}
\caption{Slide 47: Part 1: tiling}
\figsource{CS336 Spring 2026 Lecture 5 官方幻灯片。}
\end{figure}

\noindent\textbf{展开说明：}Tiling has a major impact through alignment.

\begin{knowledgebox}{读图：Slide 47 应该怎么看}
这页是硬件机制或性能证据页。先看数据从哪里来、在哪里复用、是否写回 HBM；再判断瓶颈是 compute、memory bandwidth、thread scheduling 还是 alignment。GPU 优化不是让公式更漂亮，而是让数据在正确层级停留更久、让 tensor cores 少等数据。
\end{knowledgebox}


\begin{figure}[H]
\centering
\includegraphics[width=0.88\textwidth]{slide-048.jpg}
\caption{Slide 48: Part 2: wave quantization}
\figsource{CS336 Spring 2026 Lecture 5 官方幻灯片。}
\end{figure}

\noindent\textbf{展开说明：}What’s with the periodic behavior?

\subsection{本章小结}
本节的共同问题是如何减少昂贵的数据移动并提高硬件利用率。GPU 的速度来自海量并行和专用矩阵硬件，但只有当 memory access、shape、precision 和 kernel 组织匹配时，这些速度才会真正出现。

\section{FlashAttention：把 attention 写成 IO-aware 算法}

本节把前面所有硬件技巧汇总到 FlashAttention：KQV tiling、online softmax、forward pass 分块。


\begin{importantbox}{FlashAttention 的本质}
FlashAttention 不是改变 attention 数学，而是改变计算顺序：把 Q/K/V 分块载入 SRAM，分块计算 softmax 统计量并在线更新输出，避免把完整 \(n\times n\) attention matrix 写回 HBM。
\end{importantbox}

\begin{figure}[H]
\centering
\includegraphics[width=0.88\textwidth]{slide-049.jpg}
\caption{Slide 49: Recap of part 2: making ML workloads go fast}
\figsource{CS336 Spring 2026 Lecture 5 官方幻灯片。}
\end{figure}

\noindent\textbf{展开说明：} Reduce memory accesses

\begin{knowledgebox}{读图：Slide 49 应该怎么看}
这页是硬件机制或性能证据页。先看数据从哪里来、在哪里复用、是否写回 HBM；再判断瓶颈是 compute、memory bandwidth、thread scheduling 还是 alignment。GPU 优化不是让公式更漂亮，而是让数据在正确层级停留更久、让 tensor cores 少等数据。
\end{knowledgebox}


\begin{figure}[H]
\centering
\includegraphics[width=0.88\textwidth]{slide-050.jpg}
\caption{Slide 50: Part 3: Using what we know to understand Flash Attention}
\figsource{CS336 Spring 2026 Lecture 5 官方幻灯片。}
\end{figure}

\noindent\textbf{展开说明：}Flash attention [Dao et al] dramatically accelerates attention.. But how?

\begin{knowledgebox}{读图：Slide 50 应该怎么看}
这页是硬件机制或性能证据页。先看数据从哪里来、在哪里复用、是否写回 HBM；再判断瓶颈是 compute、memory bandwidth、thread scheduling 还是 alignment。GPU 优化不是让公式更漂亮，而是让数据在正确层级停留更久、让 tensor cores 少等数据。
\end{knowledgebox}


\begin{figure}[H]
\centering
\includegraphics[width=0.88\textwidth]{slide-051.jpg}
\caption{Slide 51: Recap of attention computation}
\figsource{CS336 Spring 2026 Lecture 5 官方幻灯片。}
\end{figure}

\noindent\textbf{展开说明：}Attention computation: 3 matrix multiplies (K, Q, V) with a softmax in between

\begin{knowledgebox}{读图：Slide 51 应该怎么看}
这页是硬件机制或性能证据页。先看数据从哪里来、在哪里复用、是否写回 HBM；再判断瓶颈是 compute、memory bandwidth、thread scheduling 还是 alignment。GPU 优化不是让公式更漂亮，而是让数据在正确层级停留更久、让 tensor cores 少等数据。
\end{knowledgebox}


\begin{figure}[H]
\centering
\includegraphics[width=0.88\textwidth]{slide-052.jpg}
\caption{Slide 52: Tiling part 1: tiling for the KQV matrix multiply}
\figsource{CS336 Spring 2026 Lecture 5 官方幻灯片。}
\end{figure}

\noindent\textbf{展开说明：}This figure 1 from the paper is literally just tiling for a KQV matrix multiply..

\begin{knowledgebox}{读图：Slide 52 应该怎么看}
这页是硬件机制或性能证据页。先看数据从哪里来、在哪里复用、是否写回 HBM；再判断瓶颈是 compute、memory bandwidth、thread scheduling 还是 alignment。GPU 优化不是让公式更漂亮，而是让数据在正确层级停留更久、让 tensor cores 少等数据。
\end{knowledgebox}


\begin{figure}[H]
\centering
\includegraphics[width=0.88\textwidth]{slide-053.jpg}
\caption{Slide 53: Tiling part 2: incremental computation of the softmax}
\figsource{CS336 Spring 2026 Lecture 5 官方幻灯片。}
\end{figure}

\noindent\textbf{展开说明：}From Mikailov and Gimelshein 2018,

\begin{knowledgebox}{读图：Slide 53 应该怎么看}
这页是硬件机制或性能证据页。先看数据从哪里来、在哪里复用、是否写回 HBM；再判断瓶颈是 compute、memory bandwidth、thread scheduling 还是 alignment。GPU 优化不是让公式更漂亮，而是让数据在正确层级停留更久、让 tensor cores 少等数据。
\end{knowledgebox}


\begin{figure}[H]
\centering
\includegraphics[width=0.88\textwidth]{slide-054.jpg}
\caption{Slide 54: Putting it all together – the forward pass of flash attention}
\figsource{CS336 Spring 2026 Lecture 5 官方幻灯片。}
\end{figure}

\noindent\textbf{展开说明：}From Dao 2023, we see

\subsection{本章小结}
本节的共同问题是如何减少昂贵的数据移动并提高硬件利用率。GPU 的速度来自海量并行和专用矩阵硬件，但只有当 memory access、shape、precision 和 kernel 组织匹配时，这些速度才会真正出现。

\section{总结}

最后综合 GPU 低层细节如何支撑模型 scaling。


\begin{figure}[H]
\centering
\includegraphics[width=0.88\textwidth]{slide-055.jpg}
\caption{Slide 55: Recap for the whole lecture}
\figsource{CS336 Spring 2026 Lecture 5 官方幻灯片。}
\end{figure}

\noindent\textbf{展开说明：} Hardware powers scale, and low-level details

\subsection{本章小结}
本节的共同问题是如何减少昂贵的数据移动并提高硬件利用率。GPU 的速度来自海量并行和专用矩阵硬件，但只有当 memory access、shape、precision 和 kernel 组织匹配时，这些速度才会真正出现。


\section{总结与延伸}
Lecture 5 把 GPU 性能拆成三个层级：硬件结构、算子组织、算法重排。硬件层给出 SM/warp/shared memory/HBM/tensor cores；算子层使用 low precision、fusion、coalescing、tiling；算法层用 FlashAttention 这类 IO-aware 设计避免不必要 HBM 写回。
\begin{importantbox}{最终 takeaway}
大模型训练不是“把 PyTorch 放到 GPU 上”这么简单。真正的 GPU-aware thinking 是：每个 byte 从 HBM 进来后要被复用多少次？中间结果是否必须写回？thread/warp/block 是否对齐硬件？算法能否重排成更少 IO 的形式？
\end{importantbox}
\subsection{拓展阅读}
\begin{itemize}
    \item NVIDIA CUDA Programming Guide and H100/B200 architecture material.
    \item JAX Scaling Book roofline and inference chapters.
    \item Dao et al. FlashAttention papers.
    \item NVIDIA Transformer Engine low precision documentation.
\end{itemize}
\end{document}
