FlashAttention-3: gradients, PDF page 3

Scoring unit: science-fa3-gradients. Joined pages use the same full output.

where softmax is applied row-wise and one typically sets $\alpha = 1/\sqrt{d}$ as the scaling factor. In practice, we subtract rowmax($\mathbf{S}$) from $\mathbf{S}$ to prevent numerical instability with the exponential function. For multi-head attention (MHA), each head has its own set of query, key and value projections, and this computation parallelizes across multiple heads and batches to produce the full output tensor.

Now let $\phi$ be a scalar loss function and let $\mathbf{d}(-) = \partial\phi/\partial(-)$ be notation for the gradient. Given the output gradient $\mathbf{dO} \in \mathbb{R}^{N\times d}$, we compute $\mathbf{dQ}$, $\mathbf{dK}$, and $\mathbf{dV}$ according to the chain rule as follows:

$$
\begin{aligned}
\mathbf{dV} &= \mathbf{P}^\top \mathbf{dO} \in \mathbb{R}^{N\times d} \\
\mathbf{dP} &= \mathbf{dO}\mathbf{V}^\top \in \mathbb{R}^{N\times N} \\
\mathbf{dS} &= \mathrm{dsoftmax}(\mathbf{dP}) \in \mathbb{R}^{N\times N} \\
\mathbf{dQ} &= \alpha \mathbf{dS}\mathbf{K} \in \mathbb{R}^{N\times d} \\
\mathbf{dK} &= \alpha \mathbf{dS}^\top \mathbf{Q} \in \mathbb{R}^{N\times d},
\end{aligned}
$$

Here, we have that $\mathbf{d}s = (\mathrm{diag}(p) - pp^\top)\mathbf{d}p$ for $p = \mathrm{softmax}(s)$ as a function of a vector $s$, and we write dsoftmax($\mathbf{dP}$) for this formula applied row-wise. Finally, this computation again parallelizes across the number of heads and batches for the backward pass of MHA.

## 2.2 GPU hardware characteristics and execution model

We describe the aspects of the GPU's execution model relevant for FlashAttention-3, with a focus on the NVIDIA Hopper architecture as a concrete instantiation of this model.

**Memory hierarchy:** The GPU's memories are organized as a hierarchy of data locales, with capacity inversely related to bandwidth (Table 1)<sup>4</sup>. Global memory (GMEM), also known as HBM, is the off-chip DRAM accessible to all streaming multiprocessors (SMs). Data from GMEM gets transparently cached into an on-chip L2 cache. Next, each SM contains a small on-chip, programmer-managed highly banked cache called shared memory (SMEM). Lastly, there is the register file within each SM.

**Thread hierarchy:** The GPU's programming model is organized around logical groupings of execution units called threads. From the finest to coarsest level, the thread hierarchy is comprised of threads, warps (32 threads), warpgroups (4 contiguous warps), threadblocks (i.e., cooperative thread arrays or CTAs), threadblock clusters (in Hopper), and grids.

These two hierarchies are closely interlinked. Threads in the same CTA are co-scheduled on the same SM, and CTAs in the same cluster are co-scheduled on the same GPC. SMEM is directly addressable by all threads within a CTA, whereas each thread has at most 256 registers (RMEM) private to itself.

Table 1: Thread-Memory hierarchy for the NVIDIA Hopper H100 SXM5 GPU.

| Hardware Level | Parallel Agent | Data Locale | Capacity @ Bandwidth |
|---:|---|---|---|
| Chip | Grid | GMEM | 80 GiB @ 3.35 TB/s |
| GPC | Threadblock Clusters | L2 | 50 MiB @ 12 TB/s |
| SM | Threadblock (CTA) | SMEM | 228 KiB per SM, 31TB/s per GPU |
| Thread | Thread | RMEM | 256 KiB per SM |

**Asynchrony and warp-specialization:** GPUs are throughput processors that rely on concurrency and asynchrony to hide memory and execution latencies. For async memory copy between GMEM and SMEM, Hopper has the Tensor Memory Accelerator (TMA) as a dedicated hardware unit [38, §7.29]. Furthermore, unlike prior architectures such as Ampere, the Tensor Core of Hopper, exposed via the warpgroup-wide WGMMA instruction [40, §9.7.14], is also asynchronous and can source its inputs directly from shared memory.

---

<sup>4</sup>Luo et al. [34] reports shared memory bandwidth of 128 bytes per clock cycle per SM, and we multiply that by 132 SMs and the boost clock of 1830 MHz.

3

Scoring unit: science-fa3-gradients. Joined pages use the same full output.

where softmax is applied row-wise and one typically sets \(\alpha = 1/\sqrt{d}\) as the scaling factor. In practice, we subtract rowmax(\(\mathbf{S}\)) from \(\mathbf{S}\) to prevent numerical instability with the exponential function. For multi-head attention (MHA), each head has its own set of query, key and value projections, and this computation parallelizes across multiple heads and batches to produce the full output tensor.

Now let \(\phi\) be a scalar loss function and let \(\mathbf{d}(-) = \partial\phi/\partial(-)\) be notation for the gradient. Given the output gradient \(\mathbf{dO} \in \mathbb{R}^{N\times d}\), we compute \(\mathbf{dQ}\), \(\mathbf{dK}\), and \(\mathbf{dV}\) according to the chain rule as follows:

\[
\begin{aligned}
\mathbf{dV} &= \mathbf{P}^{\mathsf{T}}\mathbf{dO} \in \mathbb{R}^{N\times d}\\
\mathbf{dP} &= \mathbf{dO}\mathbf{V}^{\mathsf{T}} \in \mathbb{R}^{N\times N}\\
\mathbf{dS} &= \operatorname{dsoftmax}(\mathbf{dP}) \in \mathbb{R}^{N\times N}\\
\mathbf{dQ} &= \alpha\mathbf{dS}\mathbf{K} \in \mathbb{R}^{N\times d}\\
\mathbf{dK} &= \alpha\mathbf{dS}^{\mathsf{T}}\mathbf{Q} \in \mathbb{R}^{N\times d},
\end{aligned}
\]

Here, we have that \(ds = (\operatorname{diag}(p) - pp^{\mathsf{T}})dp\) for \(p = \operatorname{softmax}(s)\) as a function of a vector \(s\), and we write dsoftmax(\(\mathbf{dP}\)) for this formula applied row-wise. Finally, this computation again parallelizes across the number of heads and batches for the backward pass of MHA.

## 2.2 GPU hardware characteristics and execution model

We describe the aspects of the GPU’s execution model relevant for FlashAttention-3, with a focus on the NVIDIA Hopper architecture as a concrete instantiation of this model.

**Memory hierarchy:** The GPU’s memories are organized as a hierarchy of data locales, with capacity inversely related to bandwidth (Table 1)<sup>4</sup>. Global memory (GMEM), also known as HBM, is the off-chip DRAM accessible to all streaming multiprocessors (SMs). Data from GMEM gets transparently cached into an on-chip L2 cache. Next, each SM contains a small on-chip, programmer-managed highly banked cache called shared memory (SMEM). Lastly, there is the register file within each SM.

**Thread hierarchy:** The GPU’s programming model is organized around logical groupings of execution units called threads. From the finest to coarsest level, the thread hierarchy is comprised of threads, warps (32 threads), warpgroups (4 contiguous warps), threadblocks (i.e., cooperative thread arrays or CTAs), threadblock clusters (in Hopper), and grids.

These two hierarchies are closely interlinked. Threads in the same CTA are co-scheduled on the same SM, and CTAs in the same cluster are co-scheduled on the same GPC. SMEM is directly addressable by all threads within a CTA, whereas each thread has at most 256 registers (RMEM) private to itself.

Table 1: Thread-Memory hierarchy for the NVIDIA Hopper H100 SXM5 GPU.

| Hardware Level | Parallel Agent | Data Locale | Capacity @ Bandwidth |
|---|---|---|---|
| Chip | Grid | GMEM | 80 GiB @ 3.35 TB/s |
| GPC | Threadblock Clusters | L2 | 50 MiB @ 12 TB/s |
| SM | Threadblock (CTA) | SMEM | 228 KiB per SM, 31TB/s per GPU |
| Thread | Thread | RMEM | 256 KiB per SM |

**Asynchrony and warp-specialization:** GPUs are throughput processors that rely on concurrency and asynchrony to hide memory and execution latencies. For async memory copy between GMEM and SMEM, Hopper has the Tensor Memory Accelerator (TMA) as a dedicated hardware unit [38 §7.29]. Furthermore, unlike prior architectures such as Ampere, the Tensor Core of Hopper, exposed via the warpgroup-wide WGMMA instruction [40 §9.7.14], is also asynchronous and can source its inputs directly from shared memory.

---

<sup>4</sup>Luo et al. [34] reports shared memory bandwidth of 128 bytes per clock cycle per SM, and we multiply that by 132 SMs and the boost clock of 1830 MHz.

3

Scoring unit: science-fa3-gradients. Joined pages use the same full output.

where softmax is applied row-wise and one typically sets $\alpha = 1/\sqrt{d}$ as the scaling factor. In practice, we subtract rowmax($\mathbf{S}$) from $\mathbf{S}$ to prevent numerical instability with the exponential function. For multi-head attention (MHA), each head has its own set of query, key and value projections, and this computation parallelizes across multiple heads and batches to produce the full output tensor.

Now let $\phi$ be a scalar loss function and let $\mathbf{d}(-) = \partial\phi/\partial(-)$ be notation for the gradient. Given the output gradient $\mathbf{dO} \in \mathbb{R}^{N\times d}$, we compute $\mathbf{dQ}$, $\mathbf{dK}$, and $\mathbf{dV}$ according to the chain rule as follows:

$$\mathbf{dV} = \mathbf{P}^\top \mathbf{dO} \in \mathbb{R}^{N\times d}$$

$$\mathbf{dP} = \mathbf{dOV}^\top \in \mathbb{R}^{N\times N}$$

$$\mathbf{dS} = \text{dsoftmax}(\mathbf{dP}) \in \mathbb{R}^{N\times N}$$

$$\mathbf{dQ} = \alpha\mathbf{dSK} \in \mathbb{R}^{N\times d}$$

$$\mathbf{dK} = \alpha\mathbf{dS}^\top\mathbf{Q} \in \mathbb{R}^{N\times d},$$

Here, we have that $\mathbf{d}s = (\text{diag}(p) - pp^\top)\mathbf{d}p$ for $p = \text{softmax}(s)$ as a function of a vector $s$, and we write dsoftmax($\mathbf{dP}$) for this formula applied row-wise. Finally, this computation again parallelizes across the number of heads and batches for the backward pass of MHA.

## 2.2 GPU hardware characteristics and execution model

We describe the aspects of the GPU's execution model relevant for FlashAttention-3, with a focus on the NVIDIA Hopper architecture as a concrete instantiation of this model.

**Memory hierarchy:** The GPU's memories are organized as a hierarchy of data locales, with capacity inversely related to bandwidth (Table 1)$^4$. Global memory (GMEM), also known as HBM, is the off-chip DRAM accessible to all streaming multiprocessors (SMs). Data from GMEM gets transparently cached into an on-chip L2 cache. Next, each SM contains a small on-chip, programmer-managed highly banked cache called shared memory (SMEM). Lastly, there is the register file within each SM.

**Thread hierarchy:** The GPU's programming model is organized around logical groupings of execution units called threads. From the finest to coarsest level, the thread hierarchy is comprised of threads, warps (32 threads), warpgroups (4 contiguous warps), threadblocks (i.e., cooperative thread arrays or CTAs), threadblock clusters (in Hopper), and grids.

These two hierarchies are closely interlinked. Threads in the same CTA are co-scheduled on the same SM, and CTAs in the same cluster are co-scheduled on the same GPC. SMEM is directly addressable by all threads within a CTA, whereas each thread has at most 256 registers (RMEM) private to itself.

Table 1: Thread-Memory hierarchy for the NVIDIA Hopper H100 SXM5 GPU.

| Hardware Level | Parallel Agent | Data Locale | Capacity @ Bandwidth |
|---|---|---|---|
| Chip | Grid | GMEM | 80 GiB @ 3.35 TB/s |
| GPC | Threadblock Clusters | L2 | 50 MiB @ 12 TB/s |
| SM | Threadblock (CTA) | SMEM | 228 KiB per SM, 31TB/s per GPU |
| Thread | Thread | RMEM | 256 KiB per SM |

**Asynchrony and warp-specialization:** GPUs are throughput processors that rely on concurrency and asynchrony to hide memory and execution latencies. For async memory copy between GMEM and SMEM, Hopper has the Tensor Memory Accelerator (TMA) as a dedicated hardware unit [38, §7.29]. Furthermore, unlike prior architectures such as Ampere, the Tensor Core of Hopper, exposed via the warpgroup-wide WGMMA instruction [40, §9.7.14], is also asynchronous and can source its inputs directly from shared memory.

$^4$Luo et al. [34] reports shared memory bandwidth of 128 bytes per clock cycle per SM, and we multiply that by 132 SMs and the boost clock of 1830 MHz.

3

Scoring unit: science-fa3-gradients. Joined pages use the same full output.

where $\mathrm{softmax}$ is applied row-wise and one typically sets $\alpha = 1/\sqrt{d}$ as the scaling factor. In practice, we subtract $\mathrm{rowmax}(\mathbf{S})$ from $\mathbf{S}$ to prevent numerical instability with the exponential function. For multi-head attention (MHA), each head has its own set of query, key and value projections, and this computation parallelizes across multiple heads and batches to produce the full output tensor.

Now let $\phi$ be a scalar loss function and let $\mathbf{d}(-) = \partial\phi/\partial(-)$ be notation for the gradient. Given the output gradient $\mathbf{dO} \in \mathbb{R}^{N \times d}$, we compute $\mathbf{dQ}$, $\mathbf{dK}$, and $\mathbf{dV}$ according to the chain rule as follows:
$$
\begin{aligned}
\mathbf{dV} &= \mathbf{P}^\top\mathbf{dO} \in \mathbb{R}^{N \times d} \\
\mathbf{dP} &= \mathbf{dOV}^\top \in \mathbb{R}^{N \times N} \\
\mathbf{dS} &= \mathrm{dsoftmax}(\mathbf{dP}) \in \mathbb{R}^{N \times N} \\
\mathbf{dQ} &= \alpha\mathbf{dSK} \in \mathbb{R}^{N \times d} \\
\mathbf{dK} &= \alpha\mathbf{dS}^\top\mathbf{Q} \in \mathbb{R}^{N \times d},
\end{aligned}
$$
Here, we have that $\mathbf{ds} = (\mathrm{diag}(p) - pp^\top)\mathbf{dp}$ for $p = \mathrm{softmax}(s)$ as a function of a vector $s$, and we write $\mathrm{dsoftmax}(\mathbf{dP})$ for this formula applied row-wise. Finally, this computation again parallelizes across the number of heads and batches for the backward pass of MHA.

### 2.2 GPU hardware characteristics and execution model

We describe the aspects of the GPU's execution model relevant for FLASHATTENTION-3, with a focus on the NVIDIA Hopper architecture as a concrete instantiation of this model.

**Memory hierarchy:** The GPU's memories are organized as a hierarchy of data locales, with capacity inversely related to bandwidth (Table 1)$^4$. Global memory (GMEM), also known as HBM, is the off-chip DRAM accessible to all streaming multiprocessors (SMs). Data from GMEM gets transparently cached into an on-chip L2 cache. Next, each SM contains a small on-chip, programmer-managed highly banked cache called shared memory (SMEM). Lastly, there is the register file within each SM.

**Thread hierarchy:** The GPU's programming model is organized around logical groupings of execution units called threads. From the finest to coarsest level, the thread hierarchy is comprised of threads, warps (32 threads), warpgroups (4 contiguous warps), threadblocks (i.e., cooperative thread arrays or CTAs), threadblock clusters (in Hopper), and grids.

These two hierarchies are closely interlinked. Threads in the same CTA are co-scheduled on the same SM, and CTAs in the same cluster are co-scheduled on the same GPC. SMEM is directly addressable by all threads within a CTA, whereas each thread has at most 256 registers (RMEM) private to itself.

Table 1: Thread-Memory hierarchy for the NVIDIA Hopper H100 SXM5 GPU.

| Hardware Level | Parallel Agent | Data Locale | Capacity @ Bandwidth |
| :--- | :--- | :--- | :--- |
| Chip | Grid | GMEM | 80 GiB @ 3.35 TB/s |
| GPC | Threadblock Clusters | L2 | 50 MiB @ 12 TB/s |
| SM | Threadblock (CTA) | SMEM | 228 KiB per SM, 31TB/s per GPU |
| Thread | Thread | RMEM | 256 KiB per SM |

**Asynchrony and warp-specialization:** GPUs are throughput processors that rely on concurrency and asynchrony to hide memory and execution latencies. For async memory copy between GMEM and SMEM, Hopper has the Tensor Memory Accelerator (TMA) as a dedicated hardware unit [38, §7.29]. Furthermore, unlike prior architectures such as Ampere, the Tensor Core of Hopper, exposed via the warpgroup-wide WGMMA instruction [40, §9.7.14], is also asynchronous and can source its inputs directly from shared memory.

---
$^4$Luo et al. [34] reports shared memory bandwidth of 128 bytes per clock cycle per SM, and we multiply that by 132 SMs and the boost clock of 1830 MHz.

3

Scoring unit: science-fa3-gradients. Joined pages use the same full output.

where softmax is applied row-wise and one typically sets $\alpha = 1/\sqrt{d}$ as the scaling factor. In practice, we subtract $\text{rowmax}(\mathbf{S})$ from $\mathbf{S}$ to prevent numerical instability with the exponential function. For multi-head attention (MHA), each head has its own set of query, key and value projections, and this computation parallelizes across multiple heads and batches to produce the full output tensor.

Now let $\phi$ be a scalar loss function and let $\mathbf{d}(-) = \partial \phi / \partial (-)$ be notation for the gradient. Given the output gradient $\mathbf{dO} \in \mathbb{R}^{N \times d}$, we compute $\mathbf{dQ}, \mathbf{dK}$, and $\mathbf{dV}$ according to the chain rule as follows:

$$\begin{aligned}
\mathbf{dV} &= \mathbf{P}^\top \mathbf{dO} \in \mathbb{R}^{N \times d} \\
\mathbf{dP} &= \mathbf{dO}\mathbf{V}^\top \in \mathbb{R}^{N \times N} \\
\mathbf{dS} &= \text{dsoftmax}(\mathbf{dP}) \in \mathbb{R}^{N \times N} \\
\mathbf{dQ} &= \alpha \mathbf{dS}\mathbf{K} \in \mathbb{R}^{N \times d} \\
\mathbf{dK} &= \alpha \mathbf{dS}^\top \mathbf{Q} \in \mathbb{R}^{N \times d},
\end{aligned}$$

Here, we have that $\mathbf{d}s = (\text{diag}(p) - pp^\top)\mathbf{d}p$ for $p = \text{softmax}(s)$ as a function of a vector $s$, and we write $\text{dsoftmax}(\mathbf{dP})$ for this formula applied row-wise. Finally, this computation again parallelizes across the number of heads and batches for the backward pass of MHA.

### 2.2 GPU hardware characteristics and execution model

We describe the aspects of the GPU’s execution model relevant for FLASHATTENTION-3, with a focus on the NVIDIA Hopper architecture as a concrete instantiation of this model.

**Memory hierarchy:** The GPU’s memories are organized as a hierarchy of data locales, with capacity inversely related to bandwidth (Table 1).$^4$ Global memory (GMEM), also known as HBM, is the off-chip DRAM accessible to all streaming multiprocessors (SMs). Data from GMEM gets transparently cached into an on-chip L2 cache. Next, each SM contains a small on-chip, programmer-managed highly banked cache called shared memory (SMEM). Lastly, there is the register file within each SM.

**Thread hierarchy:** The GPU’s programming model is organized around logical groupings of execution units called threads. From the finest to coarsest level, the thread hierarchy is comprised of threads, warps (32 threads), warpgroups (4 contiguous warps), threadblocks (i.e., cooperative thread arrays or CTAs), threadblock clusters (in Hopper), and grids.

These two hierarchies are closely interlinked. Threads in the same CTA are co-scheduled on the same SM, and CTAs in the same cluster are co-scheduled on the same GPC. SMEM is directly addressable by all threads within a CTA, whereas each thread has at most 256 registers (RMEM) private to itself.

Table 1: Thread-Memory hierarchy for the NVIDIA Hopper H100 SXM5 GPU.

| Hardware Level | Parallel Agent | Data Locale | Capacity @ Bandwidth |
| :--- | :--- | :--- | :--- |
| Chip | Grid | GMEM | 80 GiB @ 3.35 TB/s |
| GPC | Threadblock Clusters | L2 | 50 MiB @ 12 TB/s |
| SM | Threadblock (CTA) | SMEM | 228 KiB per SM, 31TB/s per GPU |
| Thread | Thread | RMEM | 256 KiB per SM |

**Asynchrony and warp-specialization:** GPUs are throughput processors that rely on concurrency and asynchrony to hide memory and execution latencies. For async memory copy between GMEM and SMEM, Hopper has the Tensor Memory Accelerator (TMA) as a dedicated hardware unit [38, §7.29]. Furthermore, unlike prior architectures such as Ampere, the Tensor Core of Hopper, exposed via the warpgroup-wide WGMMA instruction [40, §9.7.14], is also asynchronous and can source its inputs directly from shared memory.

---
$^4$Luo et al. [34] reports shared memory bandwidth of 128 bytes per clock cycle per SM, and we multiply that by 132 SMs and the boost clock of 1830 MHz.

3

Scoring unit: science-fa3-gradients. Joined pages use the same full output.

<div data-bbox="113 93 942 152" data-label="Text">
<p>where softmax is applied row-wise and one typically sets <math>\alpha = 1/\sqrt{d}</math> as the scaling factor. In practice, we subtract <math>\text{rowmax}(\mathbf{S})</math> from <math>\mathbf{S}</math> to prevent numerical instability with the exponential function. For multi-head attention (MHA), each head has its own set of query, key and value projections, and this computation parallelizes across multiple heads and batches to produce the full output tensor.</p>
</div>
<div data-bbox="113 153 942 183" data-label="Text">
<p>Now let <math>\phi</math> be a scalar loss function and let <math>\mathbf{d}(-) = \partial\phi/\partial(-)</math> be notation for the gradient. Given the output gradient <math>\mathbf{dO} \in \mathbb{R}^{N \times d}</math>, we compute <math>\mathbf{dQ}</math>, <math>\mathbf{dK}</math>, and <math>\mathbf{dV}</math> according to the chain rule as follows:</p>
</div>
<div data-bbox="417 193 623 292" data-label="Equation-Block">
<math display="block">\begin{aligned}\mathbf{dV} &amp;= \mathbf{P}^\top \mathbf{dO} \in \mathbb{R}^{N \times d} \\ \mathbf{dP} &amp;= \mathbf{dO} \mathbf{V}^\top \in \mathbb{R}^{N \times N} \\ \mathbf{dS} &amp;= \text{dsoftmax}(\mathbf{dP}) \in \mathbb{R}^{N \times N} \\ \mathbf{dQ} &amp;= \alpha \mathbf{dS} \mathbf{K} \in \mathbb{R}^{N \times d} \\ \mathbf{dK} &amp;= \alpha \mathbf{dS}^\top \mathbf{Q} \in \mathbb{R}^{N \times d},\end{aligned}</math>
</div>
<div data-bbox="113 306 942 350" data-label="Text">
<p>Here, we have that <math>\mathbf{ds} = (\text{diag}(p) - pp^\top) \mathbf{dP}</math> for <math>p = \text{softmax}(s)</math> as a function of a vector <math>s</math>, and we write <math>\text{dsoftmax}(\mathbf{dP})</math> for this formula applied row-wise. Finally, this computation again parallelizes across the number of heads and batches for the backward pass of MHA.</p>
</div>
<div data-bbox="113 370 666 385" data-label="Section-Header">
<h2>2.2 GPU hardware characteristics and execution model</h2>
</div>
<div data-bbox="113 395 942 424" data-label="Text">
<p>We describe the aspects of the GPU’s execution model relevant for FLASHATTENTION-3, with a focus on the NVIDIA Hopper architecture as a concrete instantiation of this model.</p>
</div>
<div data-bbox="113 442 942 517" data-label="Text">
<p><b>Memory hierarchy:</b> The GPU’s memories are organized as a hierarchy of data locales, with capacity inversely related to bandwidth (Table 1<sup>4</sup>). Global memory (GMEM), also known as HBM, is the off-chip DRAM accessible to all streaming multiprocessors (SMs). Data from GMEM gets transparently cached into an on-chip L2 cache. Next, each SM contains a small on-chip, programmer-managed highly banked cache called shared memory (SMEM). Lastly, there is the register file within each SM.</p>
</div>
<div data-bbox="113 535 942 595" data-label="Text">
<p><b>Thread hierarchy:</b> The GPU’s programming model is organized around logical groupings of execution units called threads. From the finest to coarsest level, the thread hierarchy is comprised of threads, warps (32 threads), warpgroups (4 contiguous warps), threadblocks (i.e., cooperative thread arrays or CTAs), threadblock clusters (in Hopper), and grids.</p>
</div>
<div data-bbox="113 596 942 640" data-label="Text">
<p>These two hierarchies are closely interlinked. Threads in the same CTA are co-scheduled on the same SM, and CTAs in the same cluster are co-scheduled on the same GPC. SMEM is directly addressable by all threads within a CTA, whereas each thread has at most 256 registers (RMEM) private to itself.</p>
</div>
<div data-bbox="236 654 808 669" data-label="Caption">
<p>Table 1: Thread-Memory hierarchy for the NVIDIA Hopper H100 SXM5 GPU.</p>
</div>
<div data-bbox="189 682 852 753" data-label="Table">
<table border="1">
<thead>
<tr>
<th>Hardware Level</th>
<th>Parallel Agent</th>
<th>Data Locale</th>
<th>Capacity @ Bandwidth</th>
</tr>
</thead>
<tbody>
<tr>
<td>Chip</td>
<td>Grid</td>
<td>GMEM</td>
<td>80 GiB @ 3.35 TB/s</td>
</tr>
<tr>
<td>GPC</td>
<td>Threadblock Clusters</td>
<td>L2</td>
<td>50 MiB @ 12 TB/s</td>
</tr>
<tr>
<td>SM</td>
<td>Threadblock (CTA)</td>
<td>SMEM</td>
<td>228 KiB per SM, 31TB/s per GPU</td>
</tr>
<tr>
<td>Thread</td>
<td>Thread</td>
<td>RMEM</td>
<td>256 KiB per SM</td>
</tr>
</tbody>
</table>
</div>
<div data-bbox="113 787 942 861" data-label="Text">
<p><b>Asynchrony and warp-specialization:</b> GPUs are throughput processors that rely on concurrency and asynchrony to hide memory and execution latencies. For async memory copy between GMEM and SMEM, Hopper has the Tensor Memory Accelerator (TMA) as a dedicated hardware unit [38] §7.29]. Furthermore, unlike prior architectures such as Ampere, the Tensor Core of Hopper, exposed via the warpgroup-wide WGMMA instruction [40] §9.7.14], is also asynchronous and can source its inputs directly from shared memory.</p>
</div>
<div data-bbox="113 870 942 895" data-label="Footnote">
<p><sup>4</sup>Luo et al. [34] reports shared memory bandwidth of 128 bytes per clock cycle per SM, and we multiply that by 132 SMs and the boost clock of 1830 MHz.</p>
</div>
<div data-bbox="515 934 529 948" data-label="Page-Footer">
<p>3</p>
</div>

Scoring unit: science-fa3-gradients. Joined pages use the same full output.

<file name="science-fa3-gradients.pdf">
where softmax is applied row-wise and one typically sets  $\alpha = 1 / \sqrt{d}$  as the scaling factor. In practice, we subtract rowmax(S) from S to prevent numerical instability with the exponential function. For multi-head attention (MHA), each head has its own set of query, key and value projections, and this computation parallelizes across multiple heads and batches to produce the full output tensor.

Now let  $\phi$  be a scalar loss function and let  $\mathbf{d}(-) = \partial \phi / \partial (-)$  be notation for the gradient. Given the output gradient  $\mathbf{d}\mathbf{O} \in \mathbb{R}^{N \times d}$ , we compute  $\mathbf{d}\mathbf{Q}$ ,  $\mathbf{d}\mathbf{K}$ , and  $\mathbf{dV}$  according to the chain rule as follows:

$\mathbf{d}\mathbf{V} = \mathbf{P}^{\top}\mathbf{d}\mathbf{O}\in \mathbb{R}^{N\times d}$

$\mathbf{d}\mathbf{P} = \mathbf{d}\mathbf{O}\mathbf{V}^{\top}\in \mathbb{R}^{N\times N}$

$\mathbf{d}\mathbf{S} = \mathrm{dsoftmax}(\mathbf{dP})\in \mathbb{R}^{N\times N}$

$\mathbf{d}\mathbf{Q} = \alpha \mathbf{d}\mathbf{S}\mathbf{K}\in \mathbb{R}^{N\times d}$

$\mathbf{d}\mathbf{K} = \alpha \mathbf{d}\mathbf{S}^{\top}\mathbf{Q}\in \mathbb{R}^{N\times d},$

Here, we have that  $\mathbf{d}s = (\mathrm{diag}(p) - pp^{\top})\mathbf{d}p$  for  $p = \mathrm{softmax}(s)$  as a function of a vector  $s$ , and we write  $\mathrm{dsoftmax}(\mathbf{dP})$  for this formula applied row-wise. Finally, this computation again parallelizes across the number of heads and batches for the backward pass of MHA.

# 2.2 GPU hardware characteristics and execution model

We describe the aspects of the GPU's execution model relevant for FLASHATTENTION-3, with a focus on the NVIDIA Hopper architecture as a concrete instantiation of this model.

Memory hierarchy: The GPU's memories are organized as a hierarchy of data locales, with capacity inversely related to bandwidth (Table 1). Global memory (GMEM), also known as HBM, is the off-chip DRAM accessible to all streaming multiprocessors (SMs). Data from GMEM gets transparently cached into an on-chip L2 cache. Next, each SM contains a small on-chip, programmer-managed highly banked cache called shared memory (SMEM). Lastly, there is the register file within each SM.

Thread hierarchy: The GPU's programming model is organized around logical groupings of execution units called threads. From the finest to coarsest level, the thread hierarchy is comprised of threads, warps (32 threads), warpgroups (4 contiguous warps), threadblocks (i.e., cooperative thread arrays or CTAs), threadblock clusters (in Hopper), and grids.

These two hierarchies are closely interlinked. Threads in the same CTA are co-scheduled on the same SM, and CTAs in the same cluster are co-scheduled on the same GPC. SMEM is directly addressable by all threads within a CTA, whereas each thread has at most 256 registers (RMEM) private to itself.

Table 1: Thread-Memory hierarchy for the NVIDIA Hopper H100 SXM5 GPU.

|  Hardware Level | Parallel Agent | Data Locale | Capacity @ Bandwidth  |
| --- | --- | --- | --- |
|  Chip | Grid | GMEM | 80 GiB @ 3.35 TB/s  |
|  GPC | Threadblock Clusters | L2 | 50 MiB @ 12 TB/s  |
|  SM | Threadblock (CTA) | SMEM | 228 KiB per SM, 31TB/s per GPU  |
|  Thread | Thread | RMEM | 256 KiB per SM  |

Asynchrony and warp-specialization: GPUs are throughput processors that rely on concurrency and asynchrony to hide memory and execution latencies. For async memory copy between GMEM and SMEM, Hopper has the Tensor Memory Accelerator (TMA) as a dedicated hardware unit [35] §7.29]. Furthermore, unlike prior architectures such as Ampere, the Tensor Core of Hopper, exposed via the warpgroup-wide WGMMA instruction [40] §9.7.14], is also asynchronous and can source its inputs directly from shared memory.
</file>

Scoring unit: science-fa3-gradients. Joined pages use the same full output.

where softmax is applied row-wise and one typically sets \(\alpha = 1/\sqrt{d}\) as the scaling factor. In practice, we subtract rowmax(S) from S to prevent numerical instability with the exponential function. For multi-head attention (MHA), each head has its own set of query, key and value projections, and this computation parallelizes across multiple heads and batches to produce the full output tensor.

Now let \(\phi\) be a scalar loss function and let \(\mathbf{d}(-) = \partial \phi/\partial(-)\) be notation for the gradient. Given the output gradient \(\mathbf{dO} \in \mathbb{R}^{N \times d}\), we compute \(\mathbf{dQ}\), \(\mathbf{dK}\), and \(\mathbf{dV}\) according to the chain rule as follows:

\[
\begin{aligned}
\mathbf{dV} &= \mathbf{P}^{\mathsf{T}}\mathbf{dO} \in \mathbb{R}^{N \times d}\\
\mathbf{dP} &= \mathbf{dO}\mathbf{V}^{\mathsf{T}} \in \mathbb{R}^{N \times N}\\
\mathbf{dS} &= \operatorname{dsoftmax}(\mathbf{dP}) \in \mathbb{R}^{N \times N}\\
\mathbf{dQ} &= \alpha\mathbf{dS}\mathbf{K} \in \mathbb{R}^{N \times d}\\
\mathbf{dK} &= \alpha\mathbf{dS}^{\mathsf{T}}\mathbf{Q} \in \mathbb{R}^{N \times d},
\end{aligned}
\]

Here, we have that \(\mathbf{d}s = (\operatorname{diag}(p)-pp^{\mathsf{T}})\mathbf{d}p\) for \(p = \operatorname{softmax}(s)\) as a function of a vector \(s\), and we write \(\operatorname{dsoftmax}(\mathbf{dP})\) for this formula applied row-wise. Finally, this computation again parallelizes across the number of heads and batches for the backward pass of MHA.

## 2.2 GPU hardware characteristics and execution model

We describe the aspects of the GPU’s execution model relevant for FLASHATTENTION-3, with a focus on the NVIDIA Hopper architecture as a concrete instantiation of this model.

**Memory hierarchy:** The GPU’s memories are organized as a hierarchy of data locales, with capacity inversely related to bandwidth (Table 1⁴). Global memory (GMEM), also known as HBM, is the off-chip DRAM accessible to all streaming multiprocessors (SMs). Data from GMEM gets transparently cached into an on-chip L2 cache. Next, each SM contains a small on-chip, programmer-managed highly banked cache called shared memory (SMEM). Lastly, there is the register file within each SM.

**Thread hierarchy:** The GPU’s programming model is organized around logical groupings of execution units called threads. From the finest to coarsest level, the thread hierarchy is comprised of threads, warps (32 threads), warpgroups (4 contiguous warps), threadblocks (i.e., cooperative thread arrays or CTAs), threadblock clusters (in Hopper), and grids.

These two hierarchies are closely interlinked. Threads in the same CTA are co-scheduled on the same SM, and CTAs in the same cluster are co-scheduled on the same GPC. SMEM is directly addressable by all threads within a CTA, whereas each thread has at most 256 registers (RMEM) private to itself.

Table 1: Thread-Memory hierarchy for the NVIDIA Hopper H100 SXM5 GPU.

| Hardware Level | Parallel Agent | Data Locale | Capacity @ Bandwidth |
|---|---|---|---|
| Chip | Grid | GMEM | 80 GiB @ 3.35 TB/s |
| GPC | Threadblock Clusters | L2 | 50 MiB @ 12 TB/s |
| SM | Threadblock (CTA) | SMEM | 228 KiB per SM, 31 TB/s per GPU |
| Thread | Thread | RMEM | 256 KiB per SM |

**Asynchrony and warp-specialization:** GPUs are throughput processors that rely on concurrency and asynchrony to hide memory and execution latencies. For async memory copy between GMEM and SMEM, Hopper has the Tensor Memory Accelerator (TMA) as a dedicated hardware unit [38, §7.29]. Furthermore, unlike prior architectures such as Ampere, the Tensor Core of Hopper, exposed via the warpgroup-wide WGMMA instruction [40, §7.14], is also asynchronous and can source its inputs directly from shared memory.

⁴Luo et al. [34] reports shared memory bandwidth of 128 bytes per clock cycle per SM, and we multiply that by 132 SMs and the boost clock of 1830 MHz.

3

Scoring unit: science-fa3-gradients. Joined pages use the same full output.

![text](121,90,934,151) where softmax is applied row-wise and one typically sets $\alpha = 1/\sqrt{d}$ as the scaling factor. In practice, we subtract rowmax(S) from S to prevent numerical instability with the exponential function. For multi-head attention (MHA), each head has its own set of query, key and value projections, and this computation parallelizes across multiple heads and batches to produce the full output tensor.

![text](121,155,933,183) Now let $\phi$ be a scalar loss function and let $\mathbf{d}(-) = \partial\phi/\partial(-)$ be notation for the gradient. Given the output gradient $\mathbf{dO} \in \mathbb{R}^{N \times d}$ , we compute $\mathbf{dQ}$ , $\mathbf{dK}$ , and $\mathbf{dV}$ according to the chain rule as follows:

![formula](440,190,595,212) $$
\mathbf{dV} = \mathbf{P}^\top \mathbf{dO} \in \mathbb{R}^{N \times d}
$$

![formula](440,211,599,235) $$
\mathbf{dP} = \mathbf{dOV}^\top \in \mathbb{R}^{N \times N}
$$

![formula](449,233,650,259) $$
\mathbf{dS} = \text{dsoftmax}(\mathbf{dP}) \in \mathbb{R}^{N \times N}
$$

![formula](449,257,592,281) $$
\mathbf{dQ} = \alpha \mathbf{dSK} \in \mathbb{R}^{N \times d}
$$

![formula](445,280,598,301) $$
\mathbf{dK} = \alpha \mathbf{dS}^\top \mathbf{Q} \in \mathbb{R}^{N \times d},
$$

![text](121,311,933,353) Here, we have that $\mathbf{ds} = (\text{diag}(p) - pp^\top)\mathbf{dp}$ for $p = \text{softmax}(s)$ as a function of a vector $s$ , and we write $\text{dsoftmax}(\mathbf{dP})$ for this formula applied row-wise. Finally, this computation again parallelizes across the number of heads and batches for the backward pass of MHA.

![title](120,374,679,390) ## 2.2 GPU hardware characteristics and execution model

![text](121,400,933,428) We describe the aspects of the GPU’s execution model relevant for FLASHATTENTION-3, with a focus on the NVIDIA Hopper architecture as a concrete instantiation of this model.

![text](121,448,934,520) **Memory hierarchy:** The GPU’s memories are organized as a hierarchy of data locales, with capacity inversely related to bandwidth (Table 1) $^4$ Global memory (GMEM), also known as HBM, is the off-chip DRAM accessible to all streaming multiprocessors (SMs). Data from GMEM gets transparently cached into an on-chip L2 cache. Next, each SM contains a small on-chip, programmer-managed highly banked cache called shared memory (SMEM). Lastly, there is the register file within each SM.

![text](121,540,934,599) **Thread hierarchy:** The GPU’s programming model is organized around logical groupings of execution units called threads. From the finest to coarsest level, the thread hierarchy is comprised of threads, warps (32 threads), warpgroups (4 contiguous warps), threadblocks (i.e., cooperative thread arrays or CTAs), threadblock clusters (in Hopper), and grids.

![text](121,601,933,644) These two hierarchies are closely interlinked. Threads in the same CTA are co-scheduled on the same SM, and CTAs in the same cluster are co-scheduled on the same GPC. SMEM is directly addressable by all threads within a CTA, whereas each thread has at most 256 registers (RMEM) private to itself.

![caption](247,661,812,674) Table 1: Thread-Memory hierarchy for the NVIDIA Hopper H100 SXM5 GPU.

![table](200,687,858,759) <table>
  <tr>
    <th>Hardware Level</th>
    <th>Parallel Agent</th>
    <th>Data Locale</th>
    <th>Capacity @ Bandwidth</th>
  </tr>
  <tr>
    <td>Chip</td>
    <td>Grid</td>
    <td>GMEM</td>
    <td>80 GiB @ 3.35 TB/s</td>
  </tr>
  <tr>
    <td>GPC</td>
    <td>Threadblock Clusters</td>
    <td>L2</td>
    <td>50 MiB @ 12 TB/s</td>
  </tr>
  <tr>
    <td>SM</td>
    <td>Threadblock (CTA)</td>
    <td>SMEM</td>
    <td>228 KiB per SM, 31TB/s per GPU</td>
  </tr>
  <tr>
    <td>Thread</td>
    <td>Thread</td>
    <td>RMEM</td>
    <td>256 KiB per SM</td>
  </tr>
</table>

![text](121,794,933,866) **Asynchrony and warp-specialization:** GPUs are throughput processors that rely on concurrency and asynchrony to hide memory and execution latencies. For async memory copy between GMEM and SMEM, Hopper has the Tensor Memory Accelerator (TMA) as a dedicated hardware unit [38] §7.29]. Furthermore, unlike prior architectures such as Ampere, the Tensor Core of Hopper, exposed via the warpgroup-wide WGMMA instruction [40] §9.7.14], is also asynchronous and can source its inputs directly from shared memory.

![footnote](122,875,934,900) $^4$ Luo et al. [34] reports shared memory bandwidth of 128 bytes per clock cycle per SM, and we multiply that by 132 SMs and the boost clock of 1830 MHz.

![page_number](526,937,544,951) 3

Scoring unit: science-fa3-gradients. Joined pages use the same full output.

where softmax is applied row-wise and one typically sets $\alpha = 1/\sqrt{d}$ as the scaling factor. In practice, we subtract rowmax(**S**) from **S** to prevent numerical instability with the exponential function. For multi-head attention (MHA), each head has its own set of query, key and value projections, and this computation parallelizes across multiple heads and batches to produce the full output tensor.

Now let $\phi$ be a scalar loss function and let $d(-) = \partial \phi / \partial (-)$ be notation for the gradient. Given the output gradient $dO \in R^{N \times d}$, we compute $dQ$, $dK$, and $dV$ according to the chain rule as follows:

$$
\begin{aligned}
dV &= P^\top dO \in R^{N \times d} \\
dP &= dO V^\top \in R^{N \times N} \\
dS &= dsoftmax(dP) \in R^{N \times N} \\
dQ &= \alpha dS K \in R^{N \times d} \\
dK &= \alpha dS^\top Q \in R^{N \times d},
\end{aligned}
$$

Here, we have that $ds = (diag(p) - p p^\top) dp$ for $p = softmax(s)$ as a function of a vector $s$, and we write $dsoftmax(dP)$ for this formula applied row-wise. Finally, this computation again parallelizes across the number of heads and batches for the backward pass of MHA.

## 2.2 GPU hardware characteristics and execution model

We describe the aspects of the GPU's execution model relevant for FLASHATTENTION-3, with a focus on the NVIDIA Hopper architecture as a concrete instantiation of this model.

**Memory hierarchy:** The GPU's memories are organized as a hierarchy of data locales, with capacity inversely related to bandwidth (Table <mark>1</mark><sup>4</sup>). Global memory (GMEM), also known as HBM, is the off-chip DRAM accessible to all streaming multiprocessors (SMs). Data from GMEM gets transparently cached into an on-chip L2 cache. Next, each SM contains a small on-chip, programmer-managed highly banked cache called shared memory (SMEM). Lastly, there is the register file within each SM.

**Thread hierarchy:** The GPU's programming model is organized around logical groupings of execution units called threads. From the finest to coarsest level, the thread hierarchy is comprised of threads, warps (32 threads), warpgroups (4 contiguous warps), threadblocks (i.e., cooperative thread arrays or CTAs), threadblock clusters (in Hopper), and grids.

These two hierarchies are closely interlinked. Threads in the same CTA are co-scheduled on the same SM, and CTAs in the same cluster are co-scheduled on the same GPC. SMEM is directly addressable by all threads within a CTA, whereas each thread has at most 256 registers (RMEM) private to itself.

Table 1: Thread-Memory hierarchy for the NVIDIA Hopper H100 SXM5 GPU.

<table>
  <thead>
    <tr>
      <th>Hardware Level</th>
      <th>Parallel Agent</th>
      <th>Data Locale</th>
      <th>Capacity @ Bandwidth</th>
    </tr>
  </thead>
  <tbody>
    <tr>
      <td>Chip</td>
<td>Grid</td>
<td>GMEM</td>
<td>80 GiB @ 3.35 TB/s</td>
    </tr>
<tr>
      <td>GPC</td>
<td>Threadblock Clusters</td>
<td>L2</td>
<td>50 MiB @ 12 TB/s</td>
    </tr>
<tr>
      <td>SM</td>
<td>Threadblock (CTA)</td>
<td>SMEM</td>
<td>228 KiB per SM, 31TB/s per GPU</td>
    </tr>
<tr>
      <td>Thread</td>
<td>Thread</td>
<td>RMEM</td>
<td>256 KiB per SM</td>
    </tr>
  </tbody>
</table>

**Asynchrony and warp-specialization:** GPUs are throughput processors that rely on concurrency and asynchrony to hide memory and execution latencies. For async memory copy between GMEM and SMEM, Hopper has the Tensor Memory Accelerator (TMA) as a dedicated hardware unit [38] §7.29]. Furthermore, unlike prior architectures such as Ampere, the Tensor Core of Hopper, exposed via the warpgroup-wide WGMMA instruction [40] §9.7.14], is also asynchronous and can source its inputs directly from shared memory.

<sup>4</sup>Luo et al. [34] reports shared memory bandwidth of 128 bytes per clock cycle per SM, and we multiply that by 132 SMs and the boost clock of 1830 MHz.

3

Scoring unit: science-fa3-gradients. Joined pages use the same full output.

No output

Scoring unit: science-fa3-gradients. Joined pages use the same full output.

where softmax is applied row-wise and one typically sets α = 1/d as the scaling factor. In practice, we subtract
rowmax(S) from S to prevent numerical instability with the exponential function. For multi-head attention (MHA),
each head has its own set of query, key and value projections, and this computation parallelizes across multiple
heads and batches to produce the full output tensor.

Now let φ be a scalar loss function and let d(−) = ∂φ/∂(−) be notation for the gradient. Given the output
gradient do ∈RN×d, we compute dq, dk, and dV according to the chain rule as follows:

dV = PTdo ∈RN×d

dP = dOVT ∈RN×N

dS = dsoftmax(dp) ∈RN×N

dq = αdsK ∈RN×d

dk = αdsTQ ∈RN×d

Here, we have that ds = (diag(p) −ppT)dp for p = softmax(s) as a function of a vector s, and we write dsoftmax(dp)
for this formula applied row-wise. Finally, this computation again parallelizes across the number of heads and
batches for the backward pass of MHA.

## 2.2 GPU hardware characteristics and execution model

We describe the aspects of the GPU’s execution model relevant for FLASHATTENTION-3, with a focus on the
NVIDIA Hopper architecture as a concrete instantiation of this model.

**Memory hierarchy:** The GPU’s memories are organized as a hierarchy of data locales, with capacity inversely
related to bandwidth (Table 14). Global memory (GMEM), also known as HBM, is the off-chip DRAM accessible
to all streaming multiprocessors (SMs). Data from GMEM gets transparently cached into an on-chip L2 cache.
Next, each SM contains a small on-chip, programmer-managed highly banked cache called shared memory (SMEM).
Lastly, there is the register file within each SM.

**Thread hierarchy:** The GPU’s programming model is organized around logical groupings of execution units
called threads. From the finest to coarsest level, the thread hierarchy is comprised of threads, warps (32 threads),
warpggroups (4 contiguous warps), threadblocks (i.e., cooperative thread arrays or CTAs), threadblock clusters (in
Hopper), and grids.

These two hierarchies are closely interlinked. Threads in the same CTA are co-scheduled on the same SM, and
CTAs in the same cluster are co-scheduled on the same GPC. SMEM is directly addressable by all threads within a
CTA, whereas each thread has at most 256 registers (RMEM) private to itself.

Table 1: Thread-Memory hierarchy for the NVIDIA Hopper H100 SXM5 GPU.

<table><tr><th colspan="1" rowspan="1">Hardware Level</th><th colspan="1" rowspan="1">Parallel Agent</th><th colspan="1" rowspan="1">Data Locale</th><th colspan="1" rowspan="1">Capacity @ Bandwidth</th></tr><tr><td colspan="1" rowspan="1">Chip</td><td colspan="1" rowspan="1">Grid</td><td colspan="1" rowspan="1">GMEM</td><td colspan="1" rowspan="1">80 GiB @ 3.35 TB/s</td></tr><tr><td colspan="1" rowspan="1">GPC</td><td colspan="1" rowspan="1">Threadblock Clusters</td><td colspan="1" rowspan="1">L2</td><td colspan="1" rowspan="1">50 MiB @ 12 TB/s</td></tr><tr><td colspan="1" rowspan="1">SM</td><td colspan="1" rowspan="1">Threadblock (CTA)</td><td colspan="1" rowspan="1">SMEM</td><td colspan="1" rowspan="1">228 KiB per SM, 31TB/s per GPU</td></tr><tr><td colspan="1" rowspan="1">Thread</td><td colspan="1" rowspan="1">Thread</td><td colspan="1" rowspan="1">RMEM</td><td colspan="1" rowspan="1">256 KiB per SM</td></tr></table>

**Asynchrony and warp-specialization:** GPUs are throughput processors that rely on concurrency and asynchrony
to hide memory and execution latencies. For async memory copy between GMEM and SMEM, Hopper has the
Tensor Memory Accelerator (TMA) as a dedicated hardware unit [38 §7.29]. Furthermore, unlike prior architectures
such as Ampere, the Tensor Core of Hopper, exposed via the warpgroup-wide WGMMA instruction [40, §9.7.14], is
also asynchronous and can source its inputs directly from shared memory.

4Luo et al. [34 reports shared memory bandwidth of 128 bytes per clock cycle per SM, and we multiply that by 132 SMs and the
boost clock of 1830 MHz.

3

Scoring unit: science-fa3-gradients. Joined pages use the same full output.

where softmax is applied row-wise and one typically sets $\alpha = 1/\forall d as the scaling factor. In practice, we subtract rowmax(S) from S to prevent numerical instability with the exponential function. For multi-head attention (MHA), each head has its own set of query, key and value projections, and this computation parallelizes across multiple heads and batches to produce the full output tensor.

Now let $\phi be a scalar loss function and let d(-) = 0\phi/0(-) be notation for the gradient. Given the output gradient dO \in Rhhd, we compute dQ, dK, and dV according to the chain rule as follows:

## V = Pnd {E} d

## P = dV RN N S = d softmax (d) e R$ 3 Q = ad S K R d

## K = ad Q e Rand,

Here, we have that ds = (diag(p) - pp\top)dp for $p = softmax(s) as a function of a vector s, and we write dsoftmax(dP) for this formula applied row-wise. Finally, this computation again parallelizes across the number of heads and batches for the backward pass of MHA.

## 2.2 GPU hardware characteristics and execution model

We describe the aspects of the GPU’s execution model relevant for FLASHATTENTION-3, with a focus on the NVIDIA Hopper architecture as a concrete instantiation of this model.

Memory hierarchy: The GPU’s memories are organized as a hierarchy of data locales, with capacity inversely related to bandwidth (Table 1a] Global memory (GMEM), also known as HBM, is the off-chip DRAM accessible to all streaming multiprocessors (SMs). Data from GMEM gets transparently cached into an on-chip L2 cache. Next, each SM contains a small on-chip, programmer-managed highly banked cache called shared memory (SMEM). Lastly, there is the register file within each SM.

Thread hierarchy: The GPU’s programming model is organized around logical groupings of execution units called threads. From the finest to coarsest level, the thread hierarchy is comprised of threads, warps (32 threads), warpgroups (4 contiguous warps), threadblocks (i.e., cooperative thread arrays or CTAs), threadblock clusters (in Hopper), and grids.

These two hierarchies are closely interlinked. Threads in the same CTA are co-scheduled on the same SM, and CTAs in the same cluster are co-scheduled on the same GPC. SMEM is directly addressable by all threads within a CTA, whereas each thread has at most 256 registers (RMEM) private to itself.

Table 1: Thread-Memory hierarchy for the NVIDIA Hopper H100 SXM5 GPU.

<table>
<tr><th>Hardware Level</th><th>Parallel Agent</th><th>Data Locale</th><th>Capacity @ Bandwidth</th></tr>
<tr><td>Chip</td><td>Grid</td><td>GMEM</td><td>80 GiB @ 3.35 TB/s</td></tr>
<tr><td>GPC</td><td>Threadblock Clusters</td><td>L2</td><td>50 MiB @ 12 TB/s</td></tr>
<tr><td>SM</td><td>Threadblock (CTA)</td><td>SMEM</td><td>228 KiB per SM, 31TB/s per GPU</td></tr>
<tr><td>Thread</td><td>Thread</td><td>RMEM</td><td>256 KiB per SM</td></tr>
</table>

Asynchrony and warp-specialization: GPUs are throughput processors that rely on concurrency and asynchrony to hide memory and execution latencies. For async memory copy between GMEM and SMEM, Hopper has the Tensor Memory Accelerator (TMA) as a dedicated hardware unit [38] §7.29]. Furthermore, unlike prior architectures such as Ampere, the Tensor Core of Hopper, exposed via the warpgroup-wide WGMMA instruction [J0] $§9.7.14], is also asynchronous and can source its inputs directly from shared memory.

4Luo et al. [34] reports shared memory bandwidth of 128 bytes per clock cycle per SM, and we multiply that by 132 SMs and the boost clock of 1830 MHz.

Scoring unit: science-fa3-gradients. Joined pages use the same full output.

where softmax is applied row-wise and one typically sets α = 1/ √ d as the scaling factor. In practice, we subtract rowmax( S ) from S to prevent numerical instability with the exponential function. For multi-head attention (MHA), each head has its own set of query, key and value projections, and this computation parallelizes across multiple heads and batches to produce the full output tensor.

Now let φ be a scalar loss function and let d ( - ) = ∂φ/∂ ( - ) be notation for the gradient. Given the output gradient dO ∈$\_{R}$ N × $^{d}$, we compute dQ , dK , and dV according to the chain rule as follows:

$$d V & = P ^ { \top } d O \in \mathbb { R } ^ { N \times d } \\ d P & = d O V ^ { \top } \in \mathbb { R } ^ { N \times N } \\ d S & = d \text {softmax} ( d P ) \in \mathbb { R } ^ { N \times N } \\ d Q & = \alpha d S K \in \mathbb { R } ^ { N \times d } \\ d K & = \alpha d S ^ { \top } Q \in \mathbb { R } ^ { N \times d } ,$$

Here, we have that d s = (diag( p ) - pp $^{⊤}$) d p for p = softmax( s ) as a function of a vector s , and we write dsoftmax( dP ) for this formula applied row-wise. Finally, this computation again parallelizes across the number of heads and batches for the backward pass of MHA.

## 2.2 GPU hardware characteristics and execution model

We describe the aspects of the GPU's execution model relevant for FlashAttention-3, with a focus on the NVIDIA Hopper architecture as a concrete instantiation of this model.

Memory hierarchy: The GPU's memories are organized as a hierarchy of data locales, with capacity inversely related to bandwidth (Table 1). 4 , Global memory (GMEM), also known as HBM, is the off-chip DRAM accessible to all streaming multiprocessors (SMs). Data from GMEM gets transparently cached into an on-chip L2 cache. Next, each SM contains a small on-chip, programmer-managed highly banked cache called shared memory (SMEM). Lastly, there is the register file within each SM.

Thread hierarchy: The GPU's programming model is organized around logical groupings of execution units called threads. From the finest to coarsest level, the thread hierarchy is comprised of threads, warps (32 threads), warpgroups (4 contiguous warps), threadblocks (i.e., cooperative thread arrays or CTAs), threadblock clusters (in Hopper), and grids.

These two hierarchies are closely interlinked. Threads in the same CTA are co-scheduled on the same SM, and CTAs in the same cluster are co-scheduled on the same GPC. SMEM is directly addressable by all threads within a CTA, whereas each thread has at most 256 registers (RMEM) private to itself.

Table 1: Thread-Memory hierarchy for the NVIDIA Hopper H100 SXM5 GPU.

| Hardware Level   | Parallel Agent       | Data Locale   | Capacity @ Bandwidth           |
|------------------|----------------------|---------------|--------------------------------|
| Chip             | Grid                 | GMEM          | 80 GiB @ 3.35 TB/s             |
| GPC              | Threadblock Clusters | L2            | 50 MiB @  12 TB/s              |
| SM               | Threadblock (CTA)    | SMEM          | 228 KiB per SM, 31TB/s per GPU |
| Thread           | Thread               | RMEM          | 256 KiB per SM                 |

Asynchrony and warp-specialization: GPUs are throughput processors that rely on concurrency and asynchrony to hide memory and execution latencies. For async memory copy between GMEM and SMEM, Hopper has the Tensor Memory Accelerator (TMA) as a dedicated hardware unit [38, §7.29]. Furthermore, unlike prior architectures such as Ampere, the Tensor Core of Hopper, exposed via the wargroup-wide WGMMA instruction [40, §9.7.14], is also asynchronous and can source its inputs directly from shared memory.

$^{4}$Luo et al. [34] reports shared memory bandwidth of 128 bytes per clock cycle per SM, and we multiply that by 132 SMs and the boost clock of 1830 MHz.

Scoring unit: science-fa3-gradients. Joined pages use the same full output.

where softmax is applied row-wise and one typically sets \(\alpha = 1/\sqrt{d}\) as the scaling factor. In practice, we subtract rowmax(S) from S to prevent numerical instability with the exponential function. For multi-head attention (MHA), each head has its own set of query, key and value projections, and this computation parallelizes across multiple heads and batches to produce the full output tensor.
Now let \(\phi\) be a scalar loss function and let \(d(-) = \partial\phi/\partial(-)\) be notation for the gradient. Given the output gradient \(d\mathbf{O} \in \mathbb{R}^{N \times d}\), we compute \(d\mathbf{Q}\), \(d\mathbf{K}\), and \(d\mathbf{V}\) according to the chain rule as follows:
\[
d\mathbf{V} = \mathbf{P}^\top d\mathbf{O} \in \mathbb{R}^{N \times d}
\]
\( d\mathbf{P} = d\mathbf{O}\mathbf{V}^\top \in \mathbb{R}^{N \times N} \)
\( d\mathbf{S} = dsoftmax(d\mathbf{P}) \in \mathbb{R}^{N \times N} \)
\( d\mathbf{Q} = \alpha d\mathbf{SK} \in \mathbb{R}^{N \times d} \)
\( d\mathbf{K} = \alpha d\mathbf{S}^\top \mathbf{Q} \in \mathbb{R}^{N \times d}
\]
Here, we have that \(ds = (diag(p) - pp^\top)dp\) for \(p = softmax(s)\) as a function of a vector \(s\), and we write dsoftmax(dP) for this formula applied row-wise. Finally, this computation again parallelizes across the number of heads and batches for the backward pass of MHA.
2.2 GPU hardware characteristics and execution model
We describe the aspects of the GPU's execution model relevant for FLASHATTENTION-3, with a focus on the NVIDIA Hopper architecture as a concrete instantiation of this model.
Memory hierarchy: The GPU's memories are organized as a hierarchy of data locales, with capacity inversely related to bandwidth (Table 1.1). Global memory (GMEM), also known as HBM, is the off-chip DRAM accessible to all streaming multiprocessors (SMs). Data from GMEM gets transparently cached into an on-chip L2 cache. Next, each SM contains a small on-chip, programmer-managed highly banked cache called shared memory (SMEM). Lastly, there is the register file within each SM.
Thread hierarchy: The GPU's programming model is organized around logical groupings of execution units called threads. From the finest to coarsest level, the thread hierarchy is comprised of threads, warps (32 threads), warpgroups (4 contiguous warps), threadblocks (i.e., cooperative thread arrays or CTAs), threadblock clusters (in Hopper), and grids.
These two hierarchies are closely interlinked. Threads in the same CTA are co-scheduled on the same SM, and CTAs in the same cluster are co-scheduled on the same GPC. SMEM is directly addressable by all threads within a CTA, whereas each thread has at most 256 registers (RMEM) private to itself.
Table 1: Thread-Memory hierarchy for the NVIDIA Hopper H100 SXM5 GPU.
Hardware Level | Parallel Agent | Data Locale | Capacity @ Bandwidth
--- | --- | --- | ---
Chip | Grid | GMEM | 80 GiB @ 3.35 TB/s
GPC | Threadblock Clusters | L2 | 50 MiB @ 12 TB/s
SM | Threadblock (CTA) | SMEM | 228 KiB per SM, 31 TB/s per GPU
Thread | Thread | RMEM | 256 KiB per SM
Asynchrony and warp-specialization: GPUs are throughput processors that rely on concurrency and asynchrony to hide memory and execution latencies. For async memory copy between GMEM and SMEM, Hopper has the Tensor Memory Accelerator (TMA) as a dedicated hardware unit [38] [87.29]. Furthermore, unlike prior architectures such as Ampere, the Tensor Core of Hopper, exposed via the warpgroup-wide WGMMA instruction [40] [89.7.14], is also asynchronous and can source its inputs directly from shared memory.
\(^{4}\)Luo et al. [34] reports shared memory bandwidth of 128 bytes per clock cycle per SM, and we multiply that by 132 SMs and the boost clock of 1830 MHz.
3

Scoring unit: science-fa3-gradients. Joined pages use the same full output.

where softmax is applied row-wise and one typically sets a = 1/Vd as the scaling factor. In practice, we subtract
rowmax($) from § to prevent numerical instability with the exponential function. For multi-head attention (MHA),
each head has its own set of query, key and value projections, and this computation parallelizes across multiple
heads and batches to produce the full output tensor.
Now let ¢ be a scalar loss function and let d(-) = 0¢/0(—) be notation for the gradient. Given the output

gradient dO « RN*4, we compute dQ, dK, and dV according to the chain rule as follows:

dv =P™d0eR™4

dP =dOV' e RN*N

dS = dsoftmax(dP) « RN*

dQ = adSK € RN*¢

dK = adS'Qe RN”,
Here, we have that ds = (diag(p) - pp")dp for p = softmax(s) as a function of a vector s, and we write dsoftmax(dP)

for this formula applied row-wise. Finally, this computation again parallelizes across the number of heads and
batches for the backward pass of MHA.

2.2 GPU hardware characteristics and execution model

We describe the asp of the GPU’s execution model relevant for FLASHATTENTION-3, with a focus on the
NVIDIA Hopper architecture as a concrete instantiation of this model.

Memory hierarchy: The GPU’s memories are organized as a hierarchy of data locales, with capacity inversely
related to bandwidth (Table[I]}] Global memory (GMEM), also known as HBM, is the off-chip DRAM accessible
to all streaming multiprocessors (SMs). Data from GMEM gets transparently cached into an on-chip L2 cache.
Next, each SM contains
Lastly, there is the regi

small on-chip, programmer-managed highly banked cache called shared memory (SMEM).
cr file within cach SM.

Thread hierarchy: =The GPU’s programming model is organized around logical groupings of execution units
called threads. From the finest to coarsest level, the thread hierarchy is comprised of threads, warps (32 threa
warpgroups (4 contiguous warps), threadblocks (i.c., cooperative thread arrays or CTAS), threadblock clusters (in
Hopper), and grids.
These two hierarc

s are closely interlinked. Threads in the same CTA are co-scheduled on the same SM, and
CTAs in the same cluster are co-scheduled on the same GPC. SMEM is directly addressable by all threads within a
CTA, whereas each thread has at most 256 registers (RMEM) private to itself.

Table 1: Thread-Memory hierarchy for the NVIDIA Hopper H100 SXM5 GPU.

Hardware Level | Parallel Agent Data Locale | Capacity @ Bandwidth
Chip | Grid GMEM 80 GiB @ 3.35 TB/s
GPC | Threadblock Clusters | L2 50 MiB @ 12 TB/s
SM | Threadblock (CTA) SMEM 228 KiB per SM, 31TB/s per GPU
Thread | Thread RMEM 256 KiB per SM

Asynchrony and warp-specialization: GPUs are throughput processors that rely on concurrency and asynchrony
to hide memory and execution latencies. For async memory copy between GMEM and SMEM, Hopper has the
Tensor Memory Accelerator (TMA) as a dedicated hardware unit [38] §7.29]. Furthermore, unlike prior architectures
such as Ampere, the Tensor Core of Hopper, exposed via the warpgroup-wide WGMMA instruction [10) §9.7.14], is
also asynchronous and can source its inputs directly from shared memory.

4Luo et al. [34] reports shared memory bandwidth of 128 bytes per clock cycle per SM, and we multiply that by 132 SMs and the
boost clock of 1830 MHz.