diff --git a/zh/appendix/debugging_warp_specialized.md b/zh/appendix/debugging_warp_specialized.md index 5d9856db..ce57d261 100644 --- a/zh/appendix/debugging_warp_specialized.md +++ b/zh/appendix/debugging_warp_specialized.md @@ -1,10 +1,209 @@ ---- -orphan: true ---- - (chap_warp_spec_debug)= # 调试 Warp-Specialized Kernel -> 翻译状态：待翻译。对应英文章节：`appendix/debugging_warp_specialized.md`。 +{ref}`chap_gemm_advanced` 中的 GEMM Steps 7-9 会重叠 TMA load、`tcgen05` MMA 以及 TMEM/SMEM writeback。同样的调试方法也适用于 Flash Attention handoff：识别角色，识别每个角色拥有的 storage，然后用这个模型检查生成的 CUDA。 + +不要一开始就重写 kernel。先确认运行环境有效，再检查生成的 CUDA。排除环境和编译期问题后，这类 kernel 的运行时失败通常都会归结为某个 broken handoff：未初始化的 barrier、错误 arrival count、藏在 role guard 里的 collective、过期 barrier phase，或者 producer 还没让写入可见就复用了 storage。 + +## 调试 Kernel 之前 + +先排除运行时上下文问题： + +```bash +python -c "import tvm, tvm.tirx; print(tvm.__file__, tvm.__version__)" +python -c "import torch; print(torch.cuda.get_device_name(), torch.cuda.get_device_capability())" +``` + +这些 kernel 目标是 Blackwell（`sm_100a`）。如果 Python import 了过期 TVM checkout，或者 GPU 不是 Blackwell 级别，先修复这些问题再改 kernel。然后在看性能之前，先运行 kernel 最小的 correctness check，例如 `run_correctness()`。 + +## 调试流程 + +1. 用仍会失败的最小 shape 复现问题。如果 failure 是 illegal memory access，下一次运行前重启 Python。 +2. 如果 compilation fails，先检查安装的 API、target、`dispatch=` 和 buffer scope，再阅读 runtime synchronization 代码。 +3. 保存 `inspect_source("cuda")` 输出。在重新读 Python 之前，先搜索 role guard、`mbarrier_init`、`tcgen05`、`cp.async.bulk.tensor` 和 `cta_sync()`。 +4. 为失败的 kernel path 写出 roles / storage / handoff / lifetime 表。 +5. 用这张表检查生成的 CUDA：barrier init 是否位于 role branch 前；TMA producer、MMA issuer、writeback group 是否符合预期；warpgroup-only branch 内是否没有 CTA-wide collective。 +6. 把运行归类为 deadlock、crash、wrong result，或 correct-but-slow，然后使用下面对应小节。 +7. 一次只改一个 handoff：init count、arrive/wait phase、role guard、fence、TMA store drain、TMEM alloc/dealloc 或 tile-scheduler advance。 +8. 重新测性能前先重新跑 correctness。 + +## 需要记录什么 + +对任何异步 kernel，改代码前先做一个小 worksheet： + +| 项目 | 需要写下什么 | +|---|---| +| Roles | 发起每个 async operation 的精确线程、warp、warpgroup 或 CTA。 | +| Storage | 每个 tile 在每一步的 live 位置：GMEM、SMEM、TMEM 或 registers。 | +| Handoff | Producer、consumer、signal object、arrival count、phase，以及让数据可见的 fence 或 drain。 | +| Lifetime | 每个 storage slot 最早何时可以被复用、读回或释放。 | + +然后用 worksheet 检查生成的 CUDA： + +- Role guard 与 roles 表匹配。 +- Barrier init 出现在 guarded role branch 之前。 +- Collective operation 没有被 lane、warp 或 warpgroup guard 意外缩窄。 +- Arrive/wait phase 与 handoff 表匹配。 +- TMA store drain、TMEM dealloc 和 SMEM reuse 只在 lifetime 表允许之后发生。 + +同一张 worksheet 可以用于 TMA->MMA->writeback GEMM pipeline，也可以用于 Flash Attention 中的 score/softmax/value/correction handoff。 + +## 如果编译失败 + +先修复 compile-time failure，再调试 runtime synchronization： + +| 症状 | 可能区域 | 首先检查 | +|---|---|---| +| Unknown TIRx API 或 attribute error | 安装的 wheel 与教程代码不匹配 | 打印 `tvm.__file__` 和 `tvm.__version__`；用 {ref}`chap_language_reference` 对比 API 名称。 | +| Unsupported `dispatch=` | 选中的 target 或 primitive 不支持该路径 | 检查 `dispatch` 参数和 target capability；本教程中的 `tcgen05` 路径需要 Blackwell。 | +| Buffer scope mismatch | Buffer 正在通过错误硬件路径使用 | 检查 worksheet 中的 storage 行：TMEM 必须通过 `tcgen05` 访问，TMA operand 必须使用兼容的 GMEM/SMEM layout。 | +| 编译成功但生成 CUDA 缺少预期路径 | Dispatch 没有按预期 lower | 改算法前先检查生成 CUDA 中是否有 `tcgen05` 和 `cp.async.bulk.tensor`。 | + +## 检查生成代码 + +对任何已编译 kernel，保存 CUDA，方便搜索和 diff： + +```python +from pathlib import Path + +cuda_source = ex.mod.imports[0].inspect_source("cuda") +Path("artifacts").mkdir(exist_ok=True) +Path("artifacts/my_kernel.cu").write_text(cuda_source, encoding="utf-8") +print(cuda_source) +``` + +生成代码中 TIRx construct 到 CUDA 的映射如下： + +| TIRx | Generated CUDA | +|------|---------------| +| `wg_id == 0` | `(warp_id_in_cta >> 2) == 0` | +| `wg_id == 1` | `(warp_id_in_cta >> 2) == 1` | +| `warp_id == 0` | `(warp_id_in_cta & 3) == 0` | +| `warp_id == 3` | `(warp_id_in_cta & 3) == 3` | +| `lane_id == 0` | `(((int)threadIdx.x) % 32) == 0` | +| `.init()` internal guard | `((int)threadIdx.x) < 1`（仅 CTA thread 0） | +| `elect_sync()` | `tvm_builtin_elect_one_sync_op()` | + +读完整 kernel 前先扫描这些字符串： + +| Generated CUDA | 检查 | +|---|---| +| `if (threadIdx.x < 1)` | 单个 CTA-thread guard，常用于 barrier initialization | +| `mbarrier_init` | Barrier initialization 存在，并出现在 role branch 之前 | +| `tcgen05` | Tensor Core 路径已生成 | +| `cp.async.bulk.tensor` | Copy lowered 到 TMA | +| `cta_sync();` | CTA-wide barrier；它不能位于 `wg_id` branch 内 | + +## Step 7 参考骨架 + +正确编译的 Step 7 kernel 有如下 top-level 形状。下面的 guard 用 role name 写出以便阅读；在生成 CUDA 中，请搜索上表中的对应表达式。 + +```c +// (1) Barrier inits: top level, CTA thread 0 only +if (threadIdx.x < 1) { + mbarrier_init(tma2mma[0..1], 1); + mbarrier_init(mma2tma[0..1], 1); + mbarrier_init(mma2ld, 1); + mbarrier_init(ld2mma, 128); // arrived by all 128 WG0 threads +} + +// (2) TMEM alloc: WG0 warp 0, all lanes of the issuing warp +if (wg_id == 0 && warp_id == 0) tcgen05_alloc(..., 512); + +// (3) Fences + cta_sync, then phase init: producer=1, consumer=0 + +// (4) Warp-specialized loop +if (wg_id == 1 && warp_id == 3 && elect_sync) { /* TMA */ while(valid){ ... next_tile(); } } +if (wg_id == 1 && warp_id == 0 && elect_sync) { /* MMA */ while(valid){ ... next_tile(); } } +if (wg_id == 0) { /* WB */ while(valid){ ... next_tile(); } } + +// (5) Cleanup: issuing warp, no lane guard +cta_sync(); +if (warp_id == 0) { tcgen05_relinquish_alloc_permit(); tcgen05_dealloc(..., 512); } +``` + +改算法前先检查这些点： + +- Barrier init 位于 top level，而不是 `wg_id` guard 内。 +- `tcgen05_alloc` 和 `tcgen05_dealloc` 有 warp guard，但没有 lane guard；issuing warp 的所有 lane 都参与。 +- TMA 和 MMA loop 都迭代 `K_TILES` 次。 +- Phase init 是 producer=`1`，consumer=`0`。 + +## 症状映射 + +从症状开始，但把它当作线索，而不是最终诊断： + +| 线索 | 可能区域 | 首先检查 | +|---|---|---| +| Kernel hang，随后 runtime 报 unspecified launch failure | Deadlock | Barrier init 位置、arrival count、`cta_sync()` 位置和 `next_tile()` 参与情况 | +| Illegal memory access、XID，或后续无关 CUDA 调用也失败 | Crash / poisoned context | 重启 Python，然后检查 pointer range、storage lifetime 和 collective participation | +| 错误 row 以 128-row 或 tile-sized stripe 出现 | Sync race 或 tile-index mismatch | Producer/consumer phase、scheduler advance，以及哪个 warpgroup 拥有每个 row stripe | +| `NaN` 或明显 invalid values | Descriptor、operand setup 或未初始化 accumulation | SMEM/TMEM descriptor setup、swizzle/layout 和 accumulator initialization | +| 有限但带 pattern 的错误值 | Stale 或部分可见数据 | 缺少 fence、缺少 TMA store drain，或 storage 在 lifetime 表允许前被复用 | +| 输出正确但没有预期 speedup | Dispatch 或 resource 问题 | 生成 CUDA 路径、pipeline depth、occupancy 和 register spill | + +## 何时重启 Python + +CUDA error 不一定会自动清理状态。发生 illegal memory access、XID 或 “CUDA context poisoned” 错误后，后续无关调用如 `torch.randn` 可能继续失败。测试下一个修复前重启 Python process，否则你可能在调试上一次 crash，而不是当前代码。 + +## Deadlock + +按顺序检查这些点： + +- **Arrival count 与 init count 不匹配。** 常见情况：`MBarrier.init(128)`，但 `arrive` 被 `if warp_id == 0: if lane_id == 0:` guard 住，于是只有 1 个线程 arrive，wait 永远不返回。 + + | Barrier | init(count) | Who arrives | Arrivals | + |---|---|---|---| + | `TMABar` (tma->mma) | 1 | TMA engine 通过 `arrive(stage, bytes)` | 1 | + | `TCGen05Bar` (mma->tma, mma->ld) | 1 | MMA warp 通过 `tcgen05.commit` | 1 | + | `MBarrier` (ld->mma) | 128 | 所有 WG0 线程通过 `arrive` | 128 | + +- **Barrier init 嵌在 `wg_id` guard 内。** `.init()` 会 lower 成 `if threadIdx.x < 1:`，也就是 CTA thread 0。CTA thread 0 位于 WG0，所以 `if wg_id == 1:` 会阻止所有线程执行 init。Init 必须位于 top level；用 `inspect_source()` 中的 `grep mbarrier_init` 验证。 + +- **`cta_sync()` 位于 warpgroup branch 内。** `cta_sync` 是 `__syncthreads()`，要求所有 CTA 线程到达。放在 `if wg_id == 0:` 内时，WG1 永远不会到达。单个 warpgroup barrier 请使用 `T.cuda.warpgroup_sync(10)`。 + +- **`tile_scheduler.next_tile()` 被某些 consumer-warpgroup thread 跳过。** Scheduler 跟踪 per-thread state；跳过它的线程可能永远循环。 + +- **TMA 和 MMA 的 K-tile count 不一致。** 如果 MMA 做 `K_TILES - 1` 而不是 `K_TILES`，barrier phase 会 drift，第二个 outer tile 可能 deadlock。 + +- **`PipelineState` 初始 phase 错误。** Producer 从 `phase=1` 开始，使第一次 wait 通过；consumer 从 `phase=0` 开始，使第一次 wait 阻塞。如果两者从同一个 phase 开始，第一次 handoff 就可能立即 deadlock。 + +## Crash 和 Context Poisoning + +常见原因： + +- **`pool.commit()` 之后又 `pool.alloc`。** Barrier wrapper 内部会调用 `alloc`。正确顺序是：`tmem_addr -> barrier wrappers -> move_base_to(1024) -> Asmem / Bsmem / Dsmem -> commit()`。 +- **`tcgen05.alloc` 或 `tcgen05.dealloc` 带 lane guard。** Issuing warp 必须所有 lane 都参与。`if lane_id == 0:` 只运行一个线程，属于 undefined behavior。 +- **`tcgen05.dealloc` 前缺少 `cta_sync()`。** TMEM 在 writeback 仍在读取时被释放。 +- **GMEM 或 SMEM 越界访问。** 缩小到一个 tile，检查 scheduler 的 `m_idx` / `n_idx`，并检查当前 shape 是否是 kernel 的 tile 或 cluster tile 的倍数。 + +## 错误结果 + +猜测前先按 pattern 分类错误输出。整条 row stripe 通常指向 producer/consumer phase、tile-index 或 role-ownership mismatch。`NaN` 输出通常指向 descriptor setup、operand setup 或未初始化 accumulation。有限但带 pattern 的错误值通常意味着 consumer 读取了旧 tile、部分写入的 tile，或者 store 还没 drain 的数据。 + +- **`tcgen05.commit` 在 `elect_sync` 外。** 所有 32 个线程都会创建 commit group；其中 31 个空 group 会立即 signal mbarrier。TMA 可能在 MMA 读取 SMEM 前覆盖它。 +- **TMA store 前缺少 `fence.proxy_async("shared::cta")`。** TMA engine 可能看不到线程对 SMEM 的写入。 +- **TMA store 后缺少 `cp_async.bulk.commit_group()` 加 `wait_group(0)`。** 下一 tile 可能在 store drain 之前复用 Dsmem。 +- **Persistent kernel 在 1024x1024 等小尺寸上间歇失败。** 大尺寸可能用更长 K-loop 掩盖 race。重新检查 tile 之间的 phase reset 和 TMA-store commit/wait。 +- **`fence.after_thread_sync()` 通常不是修复。** MMA-completion mbarrier 已经携带 release-acquire 语义。Steps 8 和 9 只在 writeback edge 上保守添加它，也就是 `mma2ld.wait` 之后、第一个 `tcgen05.ld` 之前；不要在 TMA-to-MMA edge 上例行添加。 + +## 正确但很慢 + +如果输出正确但性能远低于预期，使用同样的 inspection loop： + +| 线索 | 可能区域 | 首先检查 | +|---|---|---| +| 生成 CUDA 没有 `cp.async.bulk.tensor` | Copy 没有 lower 到 TMA | 检查 `dispatch="tma"`、target capability 和 operand layout | +| 生成 CUDA 没有 `tcgen05` path | MMA 没有 lower 到 Blackwell Tensor Core 指令 | 检查 `dispatch="tcgen05"`、target capability 和 operand layout | +| TMA 和 MMA 没有重叠 | Pipeline 太浅或 phase 串行化了 producer/consumer | 检查生成 CUDA 中 wait/arrive/advance 的顺序 | +| 小 shape correctness 好，但大 shape 很慢 | Register spill、occupancy 或 staging-buffer pressure | 检查编译器 resource report；减小 tile size、chunk writeback 或降低 pipeline depth | + +## 提交高质量 Issue + +如果 failure 在上述检查后仍然存在，请先 reduce，再到 [Apache TVM GitHub 仓库](https://github.com/apache/tvm/issues)提交 issue。请包含： -本页用于放置调试 warp-specialized kernel 的中文翻译。 +- `tvm.__file__` / `tvm.__version__` 输出和 GPU capability； +- 能复现 failure 的最小 shape； +- failure 是 compile-time、deadlock、crash、wrong result，还是 correct-but-slow； +- 最小 kernel 或 notebook cell，以及对应 correctness check； +- 保存的 `inspect_source("cuda")` 输出，或能显示可疑 guard、barrier、dispatch path 的最小摘录。 diff --git a/zh/appendix/index.md b/zh/appendix/index.md index 3fd23d8f..a95227d3 100644 --- a/zh/appendix/index.md +++ b/zh/appendix/index.md @@ -1,10 +1,16 @@ ---- -orphan: true ---- - (chap_appendix)= # 参考资料 -> 翻译状态：待翻译。对应英文章节：`appendix/index.md`。 +主线内容贯穿第一部分到第四部分。参考资料收录的是阅读过程中需要随时查阅的材料： + +| 需求 | 位置 | +|------|-----| +| 查询 TIRx 语言特性 | **{ref}`chap_language_reference`** | +| 编译器内部机制（lowering pipeline） | **{ref}`chap_arch`** | +| 调试异步 GEMM/FA 的 hang、crash、错误结果或性能问题 | **{ref}`chap_warp_spec_debug`** | + +完整的 `tvm.tirx` Python API 请参阅 +[上游 TVM 文档](https://tvm.apache.org/docs/)。 -本页用于放置参考资料入口的中文内容。 +TIRx 原生层（{ref}`chap_tirx_primer`）和 tensor layout 模型 +（{ref}`chap_tirx_layout_api`）在第二部分介绍。 diff --git a/zh/chapter_async_barriers/index.md b/zh/chapter_async_barriers/index.md index 6d6660a2..1d22208b 100644 --- a/zh/chapter_async_barriers/index.md +++ b/zh/chapter_async_barriers/index.md @@ -1,52 +1,94 @@ (chap_async_barriers)= # 异步协作：mbarrier - - :::{admonition} 概览 :class: overview -- TODO：翻译本章 overview 第一条。 -- TODO：翻译本章 overview 第二条。 -- TODO：翻译本章 overview 第三条。 +- TMA 和 Tensor Core 都是异步的，所以“发起工作”和“工作完成”不是一回事；消费者需要显式的完成信号。 +- `mbarrier` 就是这个信号：producer arrive，consumer wait，并且它会跟踪 arrival count 以及 TMA 使用的 byte count。 +- 每个 barrier 都带有一个 *phase*，每完成一轮就翻转一次；consumer 等待正确的 phase，才能安全地越过这个同步点。 ::: -TODO：翻译导言部分。 +TMA（{ref}`chap_tma`）和 Tensor Core（{ref}`chap_tensor_cores`）操作都是异步的。当 kernel 发起 TMA load 或 `tcgen05` MMA 时，发起线程不会等待操作完成。指令只是被提交给硬件引擎；真正的数据搬运或矩阵运算会继续与程序的其他部分并行执行。 + +这很有用，因为它允许内存搬运和计算重叠。但这也意味着，单靠程序顺序无法证明数据已经准备好。后面的指令可能会在前面的异步操作完成之前运行。如果 TMA 仍在写 shared-memory tile 时 MMA 已经开始读取它，MMA 就会读到不完整的数据。如果 epilogue 在 Tensor Core 写完 accumulator 之前读取 TMEM，它会读到错误值。如果 kernel 等错了条件，它甚至可能永远无法继续前进。 + +因此，kernel 在每一个异步交接点都需要显式完成信号。`mbarrier` 就是这个信号。Producer 在工作完成时 arrive 到 barrier，consumer 在使用产物之前 wait 这个 barrier。同一个机制可以用于 TMA 到 MMA 的交接、MMA 到 epilogue 的交接，以及 pipeline stage 中的 buffer 复用。 + +Barrier 不只是一次性的 flag。它带有一个 phase bit，并且每当 barrier 完成一轮 arrival，这个 phase bit 都会改变。Phase 让同一个 barrier 可以在很多 loop iteration 中复用，而不会把一次 iteration 的完成误认为另一次 iteration 的完成。 ## mbarrier -TODO：翻译 “The mbarrier” 小节。 +`mbarrier` 是 memory barrier 的缩写，是存放在 shared memory 中的硬件同步对象。概念上，它包含两部分状态：arrival counter 和 phase bit。Counter 告诉 barrier 当前轮还缺多少次 arrival；phase bit 告诉 kernel 当前 barrier 处于哪一轮。 ```{raw} html +
+ style="width:100%; min-width:1320px; height:620px; border:1px solid var(--pst-color-border, #d0d0d0); border-radius:6px;"> +
``` +*交互图：`mbarrier` 的状态视图，展示 arrival counter、phase bit，以及 `init`、`arrive` 和 `wait` 操作；点击字段可以聚焦查看。* + +Barrier 从初始化开始。在 `init` 期间，kernel 设置这个 barrier 期望收到多少次 arrival。Barrier 从 phase 0 开始，并把 counter 设置为期望的 arrival count。从这之后，barrier 会等待所有需要的 producer 或资源使用者报告自己已经完成。 + +一次 arrival 会减少 barrier 仍在等待的工作量。Kernel 的不同部分可以用不同方式 arrive 到 barrier，这些差异很重要。 + +对于 TMA load，常见的 arrival 路径是 tx-count arrival。像 `mbarrier.arrive.expect_tx(bytes)` 这样的操作会做两件事。第一，它算作发起线程在 barrier 上的一次 arrival。第二，它记录 TMA 引擎预计要传输的 byte 数。Barrier 不会因为发起线程已经 arrive 就完成。它还会等待 TMA 引擎在传输完成时把 byte count 清零。只有两个条件都满足时 phase 才会翻转：普通 arrival count 到达 0，并且 pending tx byte count 也到达 0。 + +因此，不应该把 `expect_tx` 理解成“多一次普通 arrival”。它是在为异步 copy 设置一个 byte 预算。硬件随后通过 complete-tx 更新来记录真实 copy 的完成情况。只有 arrival 和 byte transfer 都完成之后，barrier 才完成。 + +对于 Tensor Core 工作，arrival 路径不同。`tcgen05` MMA 不会因为 MMA 被发起就自动推进一个 barrier。Kernel 必须显式把 barrier arrival 绑定到 commit 路径上，例如通过 `tcgen05.commit.mbarrier::arrive` 操作。当被 commit 的 group 完成时，Tensor Core 侧会执行 barrier arrival。如果 kernel 忘了这个 commit arrival，等待 barrier 的 consumer 就会永远等待下去。 + +普通线程也可以直接 arrive 到 barrier。当普通线程代码是 producer，或者一组线程要宣布它已经使用完某个资源时，就会用到这种方式。例如，consumer 读完一个 shared-memory buffer 后，可以 arrive 到一个 barrier，用来告诉 producer 这个 buffer 可以复用了。 -*点击图中组件查看细节：TODO：翻译 mbarrier mechanism 图注。* +等待是同一个协议的 consumer 侧。Consumer 会等待 barrier 完成当前 iteration 所期望的 phase。只有这样，读取数据或复用这个 barrier 保护的资源才是安全的。 + +关键点是：异步硬件不仅会跑在程序前面，它还会通过 barrier 把完成事件报告回来。TMA 可以发出 shared-memory tile 已经就绪的信号。Tensor Core 工作可以发出 TMEM 结果已经就绪的信号。普通线程可以发出某个 buffer 不再被使用的信号。Barrier 把这些情况统一成 producer-consumer 形状：producer arrive，consumer wait。 ## Phase Tracking -TODO：翻译 “Phase Tracking” 小节。 +Barrier 通常不会只分配给一次使用。一个 pipelined K-loop 可能会执行同一个交接数百次；如果每次 iteration 都分配一个新的 shared-memory barrier，并不现实。相反，kernel 会保留一小组固定 barrier，并在 loop 前进时复用它们。 + +Phase bit 让这种复用变得安全。 ```{raw} html +
+ style="width:100%; min-width:1320px; height:640px; border:1px solid var(--pst-color-border, #d0d0d0); border-radius:6px;"> +
``` +*交互图：多个 pipeline iteration 复用同一个 barrier，展示每一轮完成后 phase bit 如何翻转。* + +每当 barrier 完成当前轮的全部 arrival，它都会翻转 phase：phase 0 变成 phase 1，phase 1 变成 phase 0，如此循环。Wait 操作会检查 consumer 期望的 phase。这个期望 phase 由 kernel 保存在寄存器中。某个 stage 成功等待一轮之后，kernel 会在下一轮使用这个 barrier 前翻转自己的本地 phase 值。 + +这可以防止 kernel 把旧完成误认为新完成。假设一个 barrier 被用于某次 TMA load，并且已经完成。如果下一次 loop iteration 复用同一个 barrier 但没有跟踪 phase，那么 consumer 可能看到前一次完成，并错误地认为新的 load 已经就绪。Phase bit 把这两轮区分开。Iteration 0 等待一个 phase，iteration 1 等待相反的 phase，iteration 2 再等待第一个 phase，这个模式持续下去。 + +在真实 pipeline 中，记录通常按 stage 进行。Kernel 有固定数量的 shared-memory stage、匹配数量的 barrier，以及一小组保存在寄存器中的 phase 值。Loop 前进时，每个逻辑 iteration 映射到一个物理 stage，而 phase 值告诉 wait 操作它正在等待这个物理 barrier 的哪一轮。 + +这就是为什么后面的 GEMM 代码不需要每个 K tile 一个 barrier（{ref}`chap_gemm_async`）。它需要每个可复用 stage 一个 barrier，再加上 phase tracking。Stage index 选择 shared-memory buffer 和 barrier。Phase 值把当前对这个 stage 的使用与上一次使用区分开。 -*点击图中组件查看细节：TODO：翻译 phase tracking 图注。* +**Try with your agent**：给它一个 two-stage pipeline，并让它 trace 四次 iteration。对每次 iteration，列出 stage index、本地 phase 值、barrier 何时翻转，以及如果 stage 复用前没有翻转 phase 会出什么问题。 ## 同步规则 -TODO：翻译 “Synchronization Rules” 小节。 +理解 barrier 和 phase 机制之后，tensor-core kernel 中的同步模式就相当机械了。每当一条路径产生数据，或者释放另一条路径将要消费的资源时，这个交接都必须显式表达。 + +常见情况有三类。 + +第一类是线程代码为异步引擎生产数据。如果线程写 shared memory，而后续 TMA store 或 MMA 指令要读这个 shared memory，那么 kernel 必须先让线程写入对引擎可见。这需要合适的 thread-level synchronization 或 fence。具体指令取决于交接 scope，但原因始终相同：engine 不能在 producing threads 写完之前观察这个 shared-memory buffer。 + +第二类是 TMA 为 MMA 生产数据。TMA load 会异步填充一个 shared-memory tile。MMA 路径不能因为 TMA 指令已经发起就推断 tile 已经就绪。TMA 操作必须关联一个 `mbarrier`，MMA 路径必须在读取 tile 之前等待这个 barrier。 + +第三类是 MMA 为 epilogue 生产数据。`tcgen05` MMA 会异步地把结果写入 TMEM。Epilogue 不能在 Tensor Core 完成相关工作之前安全读取 accumulator。因此，MMA commit 路径会 arrive 到一个 completion barrier，而 epilogue 在读取 TMEM 之前等待这个 barrier。 ```{raw} html +
+ style="width:100%; min-width:1320px; height:700px; border:1px solid var(--pst-color-border, #d0d0d0); border-radius:6px;"> +
``` +*交互图：TMA load 通过 `mbarrier` 发出完成信号。MMA 路径在读取 shared-memory tile 之前等待 barrier。Tensor Core 到 epilogue 的交接也是同一个形状，只是 arrival 来自 Tensor Core commit 路径，而不是 TMA。* + +同一个思想也适用于资源复用。Barrier 不只是 data-ready 信号，也可以是“资源已经空闲”的信号。一个 shared-memory stage 不能在旧 tile 的所有 consumer 都完成之前被覆盖。一个 TMEM 区域不能在上一个使用者读写完成之前被复用。在这些情况下，arrival 表示“我已经用完这个资源”，wait 表示“现在可以安全地为下一个 stage 复用这个资源”。 -*点击图中组件查看细节：TODO：翻译 mbarrier TMA timeline 图注。* +这也是阅读 pipelined GEMM kernel 中同步逻辑的正确方式。Wait 和 arrive 并不是作为防御性编程到处散落。每一个都标记了一次具体的 ownership transfer：tile 变得可读、accumulator 变得可读，或者 buffer 变得可复用。一旦识别出这些交接，控制流就会容易跟踪得多。 diff --git a/zh/chapter_clc/index.md b/zh/chapter_clc/index.md index 739f05f1..65400086 100644 --- a/zh/chapter_clc/index.md +++ b/zh/chapter_clc/index.md @@ -1,31 +1,72 @@ (chap_clc)= # 进阶：Cluster Launch Control - - :::{admonition} 概览 :class: overview -- TODO：翻译本章 overview 第一条。 -- TODO：翻译本章 overview 第二条。 -- TODO：翻译本章 overview 第三条。 +- Persistent kernel 会让固定数量的 CTA 或 CTA cluster 驻留在 GPU 上（通常规模接近每个 SM 一个 active work owner，但不依赖严格的 1:1 映射），并让它们循环处理多个 output tile，而不是每个 tile 启动一个 CTA。 +- Cluster Launch Control 是 Blackwell 的硬件机制，它让一个驻留中的 cluster 可以在运行时请求另一个 tile。它是一条硬件 work-stealing 路径，围绕两条 PTX 指令构建：一条指令请求 work，另一条指令读回请求是否成功。 +- 主要收益是改善 tail 行为。当 tile 开销不均匀，或者 tile 数量不能被可用 SM 均匀整除时，提前完成的 CTA 可以继续拉取更多 work，而不是空等。 ::: -TODO：翻译导言部分。 +Persistent GEMM 不会把 CUDA grid 当作固定的 “one CTA per output tile” launch。相反，它启动较小数量的长生命周期 CTA 或 CTA cluster。每个 CTA/cluster 计算一个 tile，然后前进到另一个 tile，再继续计算，直到输出空间完成。这正是 {ref}`chap_gemm_advanced` 中逐步构建出的执行模式。 + +一旦 kernel 变成 persistent，主要调度问题就很简单：当一个 CTA 或 cluster 完成当前 tile 后，下一个 tile 从哪里来？ + +最简单的答案是静态公式。例如，kernel 可以根据 CTA id 计算 tile coordinate，然后按 grid stride 前进。这很容易实现，并且当所有 tile 开销大致相同、tile 数量在 GPU 上分布均匀时效果很好。但这个 schedule 在 work 真正运行之前就已经决定好了。如果少数 tile 更慢，或者最后几个 tile 分配不均匀，有些 SM 会提前完成自己的份额，而另一些 SM 仍在处理 tail。 + +Cluster Launch Control，简称 CLC，会改变这个调度模型。Persistent cluster 不再预先决定完整分配，而是可以向硬件 grid scheduler 请求另一个尚未 launch 的 cluster 的 work。如果请求成功，当前 cluster 会接管那个 cluster coordinate，并计算对应 tile。如果请求失败，就说明没有更多 work 可以 steal，loop 退出。 + +这和 thread block cluster 本身不是一回事。Thread block cluster（一起 launch 的 CTA，带有 cluster-level synchronization，并能访问 distributed shared memory）是在 Hopper 引入的（{ref}`chap_background`）。CLC 是 Blackwell 新增的机制，它让这些 cluster coordinate 上的调度变成动态的。Cluster 已经是 launch 的单位；CLC 让一个已经运行的 cluster 可以取消一个 pending launch，并继承它的 coordinate。 ## 两条指令 -TODO：翻译 “The Two Instructions” 小节。 +Cluster Launch Control 通过两条 PTX 指令暴露。第一条指令向 grid scheduler 发送一个异步请求。第二条指令读取响应。 + +请求指令是 `clusterlaunchcontrol.try_cancel.async`。 + +一次 `try_cancel` 会请求 scheduler 取消一个 pending cluster 的 launch，并把那个 cluster 的 coordinate 返回给调用方。响应会作为一个 16-byte record 写入 shared memory。由于请求是异步的，指令不会等待响应到达。相反，完成事件通过 `mbarrier` 报告，使用的还是 TMA 中同样的 barrier-and-phase 模型。 + +这个细节很重要，因为它意味着 CLC 不会引入新的等待模型。Kernel 发起请求，把它关联到一个 barrier，随后在读取响应之前等待这个 barrier。响应到达通过带 byte-count completion 的 barrier 发出信号，整体风格和其他异步硬件操作一致（见 {ref}`chap_async_barriers`）。 + +Barrier 触发后，kernel 使用 query 指令。 + +第一个 query 是 `clusterlaunchcontrol.query_cancel.is_canceled`。它返回一个 predicate，告诉 kernel cancellation 是否成功。True predicate 表示 scheduler 找到了一个 pending cluster launch，取消了它，并返回了它的 coordinate。False predicate 表示没有 pending work 可以拿了。 + +只有当 `is_canceled` 为 true 时，kernel 才应该读取 coordinate。读取使用 `clusterlaunchcontrol.query_cancel.get_first_ctaid`，它会提取被取消 cluster 的第一个 CTA id。这个 CTA id 是一个 coordinate vector，通常读作 `(x, y, z)`，kernel 会把它解码成下一步要计算的 output tile。 + +这个协议中没有数值形式的 sentinel tile id。Kernel 根据 predicate 分支。如果 predicate 为 true，coordinate 有效。如果 predicate 为 false，work-stealing loop 完成。 + +从底层看，这个形状直接来自 CLC 的实际行为。硬件不是从软件队列里分配一个抽象 task。它取消的是一个尚未发生的 cluster launch。因此，成功响应包含一个真实的 cluster coordinate。失败响应只是意味着 launch queue 已经耗尽。 ## Work-Stealing Loop -TODO：翻译 “The Work-Stealing Loop” 小节。 +有了这两条指令，persistent scheduler 就变成一个短循环。 + +在循环中的任意时刻，cluster 都负责计算一个 tile。在开始这个 tile 之前，它会为可能的下一个 tile 发送一个 `try_cancel` 请求。请求异步运行。当 scheduler 处理这个请求时，cluster 计算当前 tile。 + +当前 tile 完成后，cluster 等待与 `try_cancel` 响应关联的 `mbarrier`。响应就绪后，它调用 `query_cancel.is_canceled`。如果 predicate 为 true，它调用 `query_cancel.get_first_ctaid`，解码返回的 coordinate，并把它作为下一个 tile。如果 predicate 为 false，就说明没有 work 了，cluster 退出。 + +代码形状是： + +1. 为可能的下一个 tile 发起 `try_cancel`； +2. 在请求飞行期间计算当前 tile； +3. 等待响应 barrier； +4. 查询 cancellation 是否成功； +5. 要么用返回的 coordinate 继续，要么退出。 + +请求放置的位置让这个循环有价值。Cluster 不是等当前 tile 完成之后才请求更多 work。它先请求，再计算。这样 scheduler 请求就与有用计算重叠。当前 tile 完成时，下一个 tile 的答案往往已经可用。 + +这和 persistent kernel 在其他位置使用异步 copy 与 tensor-core barrier 的原因相同。Kernel 会避免把长延迟操作直接放在关键路径上。CLC 把同样的思想应用到 tile scheduling：提前请求下一份 work，计算当前 work，然后在需要时消费调度结果。 ## 与 Persistent GEMM 的关系 -TODO：翻译 “Relation to Persistent GEMM” 小节。 +{ref}`chap_gemm_advanced` 中的 persistent GEMM 在主要讲解里使用静态 scheduler。静态 scheduler 更容易解释，因为下一个 tile 可以直接从 loop state 计算出来。例如，`ClusterPersistentScheduler2D` 这样的 scheduler 可以用 output tile 空间上的 grid-stride pattern 分配 tile。 + +CLC 是这种静态分配的动态替代品。外层 loop 保持不变：每个驻留 cluster 反复计算一个 output tile，然后前进到另一个。变化的是下一个 tile 从哪里来。使用静态 scheduler 时，下一个 tile 由公式计算。使用 CLC 时，下一个 tile 由硬件 work stealing 返回。 + +这个差异在 launch 尾部最明显。在静态 schedule 中，剩余 work 可能分布不均。有些 SM 可能已经没有分配到的 tile，而另一些 SM 仍然还有几个 tile 要做。使用 CLC 时，提前完成的 cluster 会请求另一个 pending cluster coordinate。只要 launch queue 中还有 work，提前完成者就能继续拉取更多 tile。 + +当 tile 开销不均匀时，这也很重要。有些 GEMM tile 可能因为边界、masking、sparsity、grouped scheduling，或者主矩阵乘法周围的 fused work，而走不同路径。静态 schedule 假设 tile 分配在任何成本被观察到之前就足够好。CLC 不需要这个假设。它只在 cluster 变得可用之后才分配更多 work。 + +因此，在 TIRx 中，CLC 可以暴露成一个动态 tile scheduler。编程模型不需要改变单个 tile 的计算。Tile body 仍然是静态 scheduler 使用的同一个 persistent GEMM body。Scheduler 从“用公式计算我的下一个 tile coordinate”变成“向硬件请求下一个可用 cluster coordinate”。结果仍然是同一个 persistent loop，只是 work distribution 从固定的 launch-time schedule 变成硬件驱动。 diff --git a/zh/chapter_data_layout/index.md b/zh/chapter_data_layout/index.md index 19757ac2..cca80406 100644 --- a/zh/chapter_data_layout/index.md +++ b/zh/chapter_data_layout/index.md @@ -1,92 +1,168 @@ (chap_data_layout)= # 数据布局及其记号 - - :::{admonition} 概览 :class: overview -- TODO：翻译本章 overview 第一条。 -- TODO：翻译本章 overview 第二条。 -- TODO：翻译本章 overview 第三条。 +- *数据布局*把 tensor 的逻辑索引映射到物理位置，并决定 coalescing、bank conflict，以及某个硬件引擎是否能读这个 tile。 +- 本书用一套统一记号描述布局：`S[(shape) : (strides)]`，并配合命名轴（`@laneid`、`@TLane` 等）以及用于 broadcast 或复制数据的 replication term `R[...]`。 +- Swizzle 是一种基于 XOR 的地址重映射，用来消除 shared memory bank conflict。 ::: -TODO：翻译导言部分。 +同一组数值，如果以不同的物理排列写入内存，在同一块 GPU 上的运行速度可能相差一个数量级。 + +原因在于，tensor 的逻辑索引本身并不说明它的 byte 实际放在哪里。硬件对这个位置非常敏感：它决定了 32 个 lane 的 load 是合并成一次 transaction，还是分散成 32 次；决定了这些地址是落在不同 memory bank 上，还是撞到同一个 bank 后被串行化；甚至还决定了一个 tile 是否符合 Tensor Core 能够读取的 byte 排列。 + +机器学习程序通常用逻辑 shape 来描述 tensor。**数据布局**补上缺失的物理部分：它说明带有逻辑索引 `(i, j, …)` 的元素位于哪里，不论这个位置是在内存中、寄存器中，还是其他硬件存储中。 + +本章介绍现代 GPU 编程中会遇到的主要布局。为了让讨论可控，我们会发展一套紧凑的**记号**，用它描述机器学习系统中不同场景下的布局。最后，我们会讨论 **swizzling**：一种让同一个 tile 的 row-wise 和 column-wise 访问都高效的机制。 ## Shape-Stride 模型 -TODO：翻译 “The Shape-Stride Model” 小节。 +在进入 GPU-specific 的布局之前，值得先从最简单的布局开始，因为本章后面的内容都建立在它之上。布局的核心只有两件事：一个 **shape**，以及一组与之匹配的 **stride**。我们把这一对写作 `S[(shape) : (strides)]`。要找到某个逻辑索引的位置，只需要把这个索引与 strides 做点积。例如，一个 row-major 的 4×4 矩阵可以写成： + +```text +S[(4, 4) : (4, 1)] addr(i, j) = i·4 + j·1 +``` + +这只是经典 shape/stride 模型的紧凑写法（可以看作 CuTe 记号的 row-major 简化版），后面所有内容都从它扩展而来。 + +事实上，你几乎肯定已经用过这个模型。任何写过 PyTorch 或 NumPy 的人都用过，因为这些库中的 tensor 本质上就是一个 shape，加上 flat storage buffer 上的一组 stride： + +```python +import torch +t = torch.arange(12).reshape(3, 4) +t.shape # torch.Size([3, 4]) +t.stride() # (4, 1) ← exactly S[(3, 4) : (4, 1)] +``` + +一旦从这个角度看 tensor，很多 “reshape” 操作为什么完全不移动数据就很清楚了。它们只是重写 strides，并返回同一块 storage 上的一个 **view**。最清楚的例子是 transpose 或 permute： + +```python +tt = t.permute(1, 0) # or t.T +tt.shape # torch.Size([4, 3]) +tt.stride() # (1, 4) ← strides swapped, no data moved +tt.data_ptr() == t.data_ptr() # True, same bytes +``` + +这里 `t.permute(1, 0)` 是同一块内存上的 `S[(4, 3) : (1, 4)]`：transpose 只是 stride 改变，没有移动任何 byte。对 contiguous tensor 上的 `reshape` 或 `view` 也是同样道理：旧 storage 上的新 shape 和新 stride。（NumPy 的行为也相同；唯一差别是它的 `.strides` 以 byte 而不是 element 为单位。） + +GPU 上的布局也是这样工作的。本章剩下的内容本质上是一系列围绕同一思想的变体：一个 tile 的映射，不论是映射到内存，还是通过稍后介绍的命名轴映射到 lane 和寄存器，都是固定 buffer 上的 stride 规则。因此，重排 tile 通常是改变 *layout*，而不是 copy。不过，我们也要小心这个推理的边界。零拷贝的说法只对单个线性地址空间上的逻辑 view 清晰成立；在 GPU 上，它只在新的 view 与已有 byte 排列和 ownership 安排兼容时成立。一旦改变哪个线程或寄存器拥有某个元素，或者改变 SMEM swizzle，通常就需要真正的数据搬运：load、store、shuffle、`ldmatrix`、transpose。 ## Tile Layout -TODO：翻译 “Tile Layout” 小节。 +到目前为止，我们描述的是整个 tensor 的布局。但 GPU kernel 很少一次性操作整张矩阵；它们会处理较小的 tile，这些 tile 会被加载、变换，并由不同硬件部分参与计算。好消息是，tiling 不需要新概念。它仍然只是布局，只不过多了几个维度。把一个 8×8 矩阵切成 2×4 tile，我们会得到一个 4-D 布局，坐标是 `(tile_row, row_in_tile, tile_col, col_in_tile)`，stride 的选择让每个 tile 保持连续： + +```text +S[(4, 2, 2, 4) : (16, 4, 8, 1)] +``` + +逻辑 `(i, j)` 会先变成 `(i//2, j//4, i%2, j%4)`，再经过 strides 计算地址。值得注意的是，这套记号并没有引入任何特殊的 “tile” 概念，就表达了 tiling：它仍然是同一个 shape-stride 模型，只是把索引拆成了 outer 和 inner 坐标。 + +下面的交互图展示了一个逻辑矩阵索引如何分解成 tile 坐标，再映射到物理地址。 ```{raw} html + style="width:100%; min-width:1320px; height:640px; border:1px solid var(--pst-color-border, #d0d0d0); border-radius:6px;"> ``` - -*点击图中组件查看细节：TODO：翻译 tile layout 图注。* +*交互图：点击一个 cell，查看它的 tiled index 和 address。* ## 命名轴 -TODO：翻译 “Named Axes” 小节。 +到目前为止，`S[...]` 中的每个 stride 都表示线性内存中的 offset，我们也把 address 当作这个空间中的位置。但在 GPU 上，数据可以位于不止一个地方：除了内存，一个 tile 还可能分布在 warp lane 上、线程寄存器上，或者 TMEM lane 与 column 上。为了统一描述这些情况，我们给记号扩展出**命名轴**。核心想法是让每个 stride 系数携带一个 axis tag，说明它穿过的是哪个空间：`@m` 表示普通内存，`@laneid` 表示 warp lane，`@reg` 表示寄存器，`@warpid` 表示 warp，`@TLane` / `@TCol` 表示 TMEM 坐标。有了这些 tag，单个布局不仅能描述数据在内存中的位置，也能描述它如何分布在实际操作这些数据的硬件资源上。 + +把 memory tag 显式写出来之后，内存中的一个 row-major 8×16 tile 就是： + +```text +S[(8, 16) : (16@m, 1@m)] +``` + +当布局描述的数据是*分布在线程之间*，而不是放在线性内存中时，这些 tag 就开始发挥作用。例如 `S[(8, 4, 2) : (4@laneid, 1@laneid, 1@reg)]` 并不指向线性内存，而是把行列映射到 lane ID 和每个 lane 的寄存器上。这里 `laneid` 表示 warp 内的 lane index，也就是 `thread_index % warp_size`。这正是你会在 {ref}`chap_layout_generations` 中遇到的 tensor-core register fragment。 + +下面的交互图展示了一个布局如何把 tensor element 分布到 warp lane 和 per-lane register 上，而不是放进线性内存中。 ```{raw} html + style="width:100%; min-width:1320px; height:640px; border:1px solid var(--pst-color-border, #d0d0d0); border-radius:6px;"> ``` - -*点击图中组件查看细节：TODO：翻译 named axes 图注。* +*交互图：一个覆盖 `@laneid` 和 `@reg` 的布局；点击 cell 可以看到哪个 lane/register 持有它。* ## 分布式布局 -TODO：翻译 “Distributed Layout” 小节。 +命名轴最有用的地方在于，它们让我们可以统一描述系统中很多层级上的 placement，包括*跨设备*的 placement。刚才我们用它们描述单个 GPU 内部的 lane 和 register，但同样的思想可以继续向外延伸：`@gpuid_x` 和 `@gpuid_y` 这样的轴可以说明数据位于 GPU mesh 中的哪里。有了这些轴，记号就能捕捉分布式训练和推理中出现的 sharding pattern。不过，这些轴还没有描述*复制*，也就是被拷贝到多个位置的数据。因此我们加入 `R[n : stride]` 记号，其中 `R` 标记 replicated dimension。例如，`R[2 : 1@gpuid_x]` 表示沿 `@gpuid_x` 轴复制。把两者合在一起，单个表达式就能同时描述 tensor 在 2×2 GPU mesh 上的 sharding，以及沿一个轴的 replication： + +```text +S[(2, 4, 8) : (1@gpuid_y, 8@m, 1@m)] + R[2 : 1@gpuid_x] +``` + +下面的 demo 展示了一个小 GPU mesh 上 partition-and-replication 的组合模式。点击任意 cell 可以看到哪个 device 持有它，并观察 `@gpuid_x` replication 如何把同一个副本放到配对 device 上；按钮可以在 fully-sharded、shard + replica 和 shard + offset 布局之间切换。 ```{raw} html + style="width:100%; min-width:1320px; height:640px; border:1px solid var(--pst-color-border, #d0d0d0); border-radius:6px;"> ``` - -*点击图中组件查看细节：TODO：翻译 distributed layout 图注。* +*交互图：分布在 2×2 GPU mesh 上的布局；点击 cell 可以看到哪些 device 持有它。* ### Kernel 内复制模式：TMEM 中的 Scale Factor -TODO：翻译 “Intra-Kernel Replication Pattern: Scale Factors in TMEM” 小节。 +我们刚才为 GPU mesh 引入的 replication dimension `R[...]`，不只适用于多设备。事实证明，同一个结构也可以描述单个 kernel 内部发生的事情：硬件*跨 lane broadcast* 的数据。Blackwell 的 block-scaled MMA（{ref}`chap_layout_generations`）就是很好的例子。它的 scale factor 位于 TMEM 中，其中一个 128-row scale vector 只存储在 **32 个 TMEM lane** 中，逻辑行 `r` 映射到 TMEM lane `r % 32`，而 `r // 32` 沿 column 方向前进。随后这 32 个存储的 TMEM lane 会沿 TMEM `TLane` 轴**复制**，从 32 个 lane 扩展到 128 个 lane，这样读取 warpgroup 中的四个 warp 都能在自己的 32-lane TMEM window 中找到一份副本。这是一个 `warpx4` broadcast，我们用 replication dimension 写出它。读取本身由这些 warp 的线程执行： + +```text +S[(32, …) : (1@TLane, …)] + R[4 : 32@TLane] +``` + +这会在 32 个 TMEM lane 的 stride 上产生四份 replica：TMEM lane `l`、`l+32`、`l+64` 和 `l+96` 都保存同一个 scale。和前面一样，replication dimension 不携带新数据；它只是说明“同一个值位于四个 TMEM-lane 位置”，就像刚才 `@gpuid_x` 在 GPU mesh 上 broadcast 一行一样。 + +下面的交互 demo 同时展示两步：先紧凑打包到 32 个 TMEM lane 中，再通过 `warpx4` broadcast 扩展到 128 个读取 lane。 ```{raw} html + style="width:100%; min-width:1040px; height:560px; border:1px solid var(--pst-color-border, #d0d0d0); border-radius:6px;"> ``` +*交互图：点击一个 scale factor `SFA[m, sf]`；它会被打包到 TMEM lane `m mod 32`、column `(m // 32)·4 + sf`，随后沿 `TLane` 轴做 `warpx4` broadcast，产生四个 lane copy（`l`、`l+32`、`l+64`、`l+96`），每个 warp 的 32-lane window 各一份。* + +每个 column 内的 byte packing（`scale_vec` 的 1X/2X/4X 模式）以及 `cta_group::2` split 会在 {ref}`chap_layout_generations` 中介绍。 -*点击图中组件查看细节：TODO：翻译 scale factors in TMEM 图注。* +已经熟悉 CuTe 的读者可以把本章记号理解为 CuTe 的 row-major 变体，并额外扩展了显式硬件命名轴以及专用 replication 结构。 ## Swizzle Layout -TODO：翻译 “Swizzle Layout” 小节。 +本章最后一种布局是为了解决一个具体硬件问题。GPU 上的 shared memory 被组织成多个 memory bank。当不同 lane 落在不同 bank 上时，访问最快。反过来，如果多个 lane 访问同一个 bank 内的不同地址，硬件只能把它们串行化，我们就要付出 **bank conflict** 的代价。 + +在 tensor 程序中，这很难避免，因为内存访问并不总是纯线性顺序。处理矩阵时，我们经常需要从同一个 tile 中读取 row slice 和 column slice，这会产生真实的张力：对 row-wise 访问高效的布局，通常会让 column-wise 访问产生 bank conflict；而偏向 column 的布局又会伤害 row 访问。**Swizzling** 就是为了解开这个张力而设计的技术。 + +Swizzle 的想法是重排地址映射，通常是把 column index 与 row 做 XOR，使得 *row* 和 *column* 两类访问最终都分散到不同 bank 上。它提供的 conflict-free 保证是有条件的：只对匹配的 element width、swizzle mode 和 access pattern 成立，也就是对应 engine descriptor 所期望的模式；它并不保证任意 element width 或 alignment 都无冲突。 + +第一个交互 demo 让这个过程具体化。点击一个 column index，观察每个元素落到哪个 bank：左侧普通 row-major tile 中，一个 column 会把 8 个元素全部汇入同一个 bank，因此读取会串行化成 8 个 cycle；右侧 XOR-swizzled layout 中，同一个 column 被分散到 8 个不同 bank，可以在 1 个 cycle 中读完。 ```{raw} html + style="width:100%; min-width:1320px; height:640px; border:1px solid var(--pst-color-border, #d0d0d0); border-radius:6px;"> ``` +*交互图：一个 8×8 tile，普通 row-major 按列读取会产生 bank conflict，XOR swizzle 后变成 conflict-free。* + +这个小的 8×8 例子抓住了核心思想，但真实 GPU 内存的 bank 数远多于这个玩具图。为了让 swizzling 在完整尺度上工作，我们不会把整个 tile 当作一个单块对象处理。相反，我们会把内存切成小 segment，并在每个 segment 内应用 swizzle pattern。实践中最常见的是 `SWIZZLE_128B`，围绕 128-byte segment 组织，使同样的 row/column 重映射技巧自然适配 32-bank memory system。 -*点击图中组件查看细节：TODO：翻译 8x8 XOR swizzle 图注。* +下面的交互 demo 展示一个具体硬件 swizzle：`SWIZZLE_128B`。在我们推广到其他格式之前，可以先看到这个逐 segment 重复的 pattern。 ```{raw} html + style="width:100%; min-width:1320px; height:640px; border:1px solid var(--pst-color-border, #d0d0d0); border-radius:6px;"> ``` +*交互图：128-byte segment 内的 `SWIZZLE_128B` pattern；逐 cycle 查看读取过程，可以看到 `physical_sector = logical_sector XOR row` 如何把每个 column 分散到不同 bank。* + +同一个思想也可以扩展到 128-byte 之外。为了简化可视化，我们接下来会用单个彩色 block 表示一个 segment，而不是画出每个 bank。一般来说，硬件会定义一个小的重复 **atom**，在 atom 内应用 permutation；不同 swizzle mode 选择不同的 atom 大小。`SWIZZLE_128B` 使用 8 × 128 B atom，`SWIZZLE_64B` 使用 8 × 64 B atom，`SWIZZLE_32B` 使用 8 × 32 B atom；整个 tile 再由所选 atom 平铺而成。 -*点击图中组件查看细节：TODO：翻译 SWIZZLE_128B 图注。* +最后一个交互 demo 可以在这些格式之间切换（也包括一种 16 B interleaved mode），选择数据类型，并把鼠标悬停在任意 cell 上直接查看一个 atom 内部的 element 排列。对于判断某条 load/store 指令期望哪种 swizzle，这正是合适的细节层级。 ```{raw} html + style="width:100%; min-width:1320px; height:640px; border:1px solid var(--pst-color-border, #d0d0d0); border-radius:6px;"> ``` +*交互图：选择 swizzle format 和数据类型，查看它的 atom shape（8 × N B）；悬停在 cell 上可以看到其中的元素如何被 permutation。* + +应该选择哪种 mode？经验法则是优先选择 tile 能填满的*最大* atom。一个 N-byte atom 要求 tile 的 contiguous dimension 至少有 N byte，并且是 N 的倍数，因此 `SWIZZLE_128B` 只有在一行至少跨越 128 byte，也就是 64 个 `float16` 元素时才适用。只要能适用，它就是默认选择，因为它的 8 × 128 B atom 覆盖一整条 128-byte bank line，可以一次把一个 column 分散到全部 32 个 bank，在 fp16 中同时给 8 行和 8 列提供 conflict-free 访问。但如果问题 shape 迫使 contiguous dimension 变小，tile 无法填满 128 B atom，就要降到 `SWIZZLE_64B` 或 `SWIZZLE_32B`，选择 row 仍然可以覆盖的最大 atom。 + +你不需要手工推导这些 permutation address。这里值得精确说明 swizzle 与 `S[...]` 记号的关系：swizzle *不是* affine map 的一部分。它是组合在 affine map 之上的一个单独的非仿射层。`S[...]` 布局把元素放到线性内存（`@m`）地址上，swizzle 再重排这个地址。在 TIRx layout API 中，这写作 `ComposeLayout(swizzle, tile)`（{ref}`chap_tirx_layout_api`）。你的工作只是为所有会触碰这个 tile 的 op 选择一致的 mode，然后让组合布局处理剩下的事。 -*点击图中组件查看细节：TODO：翻译 swizzle atom 图注。* +硬件填充的也是同一个组合布局，这正是 swizzling 与 tiling 汇合的地方。TMA descriptor 是多维的，因此一个三维 box 可以同时描述 tile 的 atom tiling 以及每个 atom 内部的 swizzle；一次 TMA load 随后会 atom by atom 地布置 tile，并在写 shared memory 时应用 swizzle（{ref}`chap_tma`），不需要单独的 swizzling pass。每个 engine 需要*哪一种* swizzle 是 generation-specific 的，这就是下一章的主题。 diff --git a/zh/chapter_flash_attention/index.md b/zh/chapter_flash_attention/index.md index f00ed0b0..73ab1880 100644 --- a/zh/chapter_flash_attention/index.md +++ b/zh/chapter_flash_attention/index.md @@ -1,10 +1,607 @@ ---- -orphan: true ---- - (chap_flash_attention)= # Flash Attention 4 -> 翻译状态：待翻译。对应英文章节：`chapter_flash_attention/index.md`。 +:::{admonition} 概览 +:class: overview + +- Attention 会运行两个 MMA，中间夹着 softmax，因此它不能像 GEMM 那样只是重复一个 MMA。 +- Kernel 会把第一部分的硬件 primitive（TMA、`tcgen05`、TMEM、barrier）和第三部分的 GEMM 技术组合起来，再加入 warp role、online-softmax rescaling、causal masking 和 GQA。 +::: + +Attention 是决定 transformer 能否运行的 kernel，也是前面构建的所有东西最终必须协同工作的地方。我们为 GEMM 组装的每个部件都会带到这里：TMA tile movement、`tcgen05` MMA、TMEM、warpgroup register tile，以及显式 barrier。 + +挑战在于 attention 不是一个 MMA 的重复。它是两个 MMA，中间夹着真正的工作：online softmax、causal masking，以及让早期 block 和后续 block 保持同一尺度的 rescaling。 + +新的难点就在中间这一段。普通 matmul 只会加到 accumulator；attention 必须在新的 key 和 value stream in 时，重新访问并 rescale 已经计算过的结果。Softmax 本身也在两个 Tensor Core MMA 之间运行在 CUDA core 上，因此 exponential 和 row-wise reduction 直接位于关键路径上。 + +这就是为什么 attention 优化很大程度上就是 softmax 优化：重写 `exp`，并让 softmax 与 MMA 重叠，而不是在 softmax 上停住。 + +本章目标不是从零重新推导 Flash Attention。我们会保留足够算法背景，让 kernel 可读，然后把注意力放在真正的新内容上：这个算法如何变成 TIRx。 + +最清晰的入口是跟随一个 tile 在 kernel 中流动。`Q`、`K` 和 `V` 作为 input tile 进入，从 GMEM 加载到 SMEM。Score MMA 把 `Q` 和 `K` 相乘，得到 TMEM 中的 score tile `S`。Softmax 把 `S` 变成 numerator tile `P`，value MMA 再把 `P` 和 `V` 组合起来更新 output accumulator `O`。 + +到目前为止，这看起来像两个 matmul 粘在一起，但这里有一个 GEMM 从不需要处理的 twist：每当 running softmax maximum 改变，已经累加的 `O` 就突然处在错误尺度上。它必须先被 rescale，下一次 value MMA 才能安全地加到其中。下面各节会先追踪这条路径，然后展示 TIRx 如何把每个 stage 交给 warpgroup，并把 stage 串接起来。 + +## 算法形状 + +在把 tile 放进内存之前，我们需要知道这些 tile 服务的算法。对于一个 query block，Flash Attention 计算： + +$$O = \text{softmax}(QK^{\top} / \sqrt{d})V$$ + +字面看，公式说要构造完整 score matrix `S = QKᵀ`，对它做 softmax，再乘以 `V`。这正是我们不能使用的方法，因为完整 `S` 太大。seq=4096 时，每个 head 大约有 16M 个元素，fp32 下约 64 MB，比 SMEM 或单个 128×512 TMEM region 大几个数量级。片上根本没有地方放它。Flash Attention 的答案是永远不 materialize `S`。相反，它按 block stream `K/V`，并携带三个 per-row running state 来总结目前为止看到的全部内容： + +- `row_max`：目前见过的最大 score。 +- `row_sum`：softmax 的 running denominator。 +- `O`：running output accumulator。 + +Streaming update 会在新 block 到达时保持这些 state 正确。微妙之处在于，每处理一个 block，running max 都可能上升；一旦上升，在旧 max 下计算的一切都处在错误尺度上。因此在加入新贡献之前，我们必须先把旧状态拉回新尺度： + +```text +S = Q_block @ K_block.T +m_new = max(row_max, rowmax(S)) +scale = exp((row_max - m_new) / sqrt(d)) +P = exp((S - m_new) / sqrt(d)) +row_sum = row_sum * scale + rowsum(P) +O = O * scale + P @ V_block +row_max = m_new +``` + +单个 `scale` 因子在这里有双重作用：它同时 rescale running denominator 和 running output，使早期 block 和后续 block 的贡献最终位于同一尺度。 + +上面的 pseudocode 使用自然 `exp` 和显式 `/sqrt(d)`，因为这样最容易读，但 kernel 采用更便宜的路径。它把 `1/sqrt(d)` 和 `log2(e)` 合成一个常量 `scale_log2 = log2(e)/sqrt(d)`，并用硬件 `exp2` 对 raw score 计算所有 exponential，使用恒等式 `exp(x/sqrt(d)) = exp2(x · scale_log2)`。原因很简单：在这个硬件上，`exp2` 比自然 `exp` 更快。 + +继续之前，有一点要明确：这里的 `P` *不是*最终 normalized attention matrix。它只是当前 K/V block 的 softmax numerator。Normalization 被故意推迟，只有最后一个 block 之后，kernel 才写出 `O / row_sum`。 + +对于 TIRx，知道算法计算什么只是一半。另一半是 kernel 运行时*每个 tile 位于哪里*，因为这决定 layout 和 barrier 代码。`S`、`P` 和 `O` 都是 tile value，并且各自有家： + +- `S` 是 score tile。Score MMA 把它写入 TMEM。 +- `P` 是 softmax numerator tile。Softmax 从 TMEM 把 `S` 读入寄存器，计算 `P = exp((S - m_new) / sqrt(d))`，再把 `P` 写回 TMEM。 +- `O` 是 output accumulator tile。Value MMA 从 TMEM 读取 `P`，从 SMEM 读取 `V`，然后累加到 TMEM 中的 `O`。 + +前面提到的 rescale 也是一个 tile operation，不是 scalar bookkeeping：当 `row_max` 改变时，旧 `O` 会从 TMEM 读出，在寄存器中乘上 scale，再写回 TMEM，然后下一次 value MMA 才能累加进去。后面每节都会沿着同一结构展开：tile placement、hardware path，以及证明下一个 consumer 可以运行的 barrier。 + +## Tile-Primitive Graph + +有了 running state 和它们的位置，我们可以把算法展开成具体的 tile movement 序列。对于一个 K/V block，kernel 从上到下走这条 tile 路径： + +```text +Q, K, V in GMEM + -> Q, K, V in SMEM by TMA load + -> S in TMEM by score MMA: QK^T + -> P in TMEM by softmax numerator: TMEM -> RF -> TMEM + -> O in TMEM by value MMA: P V + -> O in GMEM by normalization, SMEM staging, and TMA store +``` + +与 GEMM 的差异归结为一行。GEMM 是重复的一条 MMA chain；FA4 有两个 MMA phase，中间夹着 softmax。后面几乎所有内容都是这个额外 stage 的后果。 + +如果把短路径展开成显式 producer-consumer edge，就得到完整 graph： + +| Stage | Tile movement or compute | TIRx primitive | Hardware path | +|-------|--------------------------|----------------|---------------| +| Load Q/K/V | GMEM tiles -> SMEM tiles | `Tx.copy_async(..., dispatch="tma")` | TMA load | +| Score MMA | SMEM 中的 Q 和 SMEM 中的 K -> TMEM 中的 score tile `S` | `Tx.warp.gemm_async(..., dispatch="tcgen05")` | `tcgen05.mma` | +| Softmax read | TMEM 中的 `S` -> warpgroup register tile | `Tx.wg.copy_async(reg, tmem)` | `tcgen05.ld` | +| Softmax write | registers 中的 numerator tile `P` -> fp16 TMEM view | `Tx.copy_async(tmem_as_f16, reg)` | TMEM store，随后 `tcgen05.wait.st()` | +| Value MMA | TMEM 中的 `P` 和 SMEM 中的 V -> TMEM 中的 output accumulator `O` | `Tx.warp.gemm_async(..., dispatch="tcgen05")` | 带 TMEM operand 的 `tcgen05.mma` | +| Correction | TMEM 中的 `O` -> registers -> TMEM 中的 `O` | TMEM readback、register multiply、TMEM store | `tcgen05.ld` / TMEM store | +| Epilogue | 最终 TMEM 中的 `O` -> registers -> SMEM -> GMEM | TMEM readback、`Tx.copy`、TMA store | `tcgen05.ld` + TMA store | + +新增行是 softmax 和 correction。二者都会加入 TMEM -> register -> TMEM 流量，也会在 score MMA 和 value MMA 之间创建额外交接。 + +**Try with your agent**：让它只 trace 上面的短路径。对每个箭头，命名 producer stage、consumer stage、source tile、destination tile 和 hardware path。然后问哪些箭头在 GEMM 章节中不存在。 + +## Warp 角色和 Scope + +数据路径确定后，自然的问题是谁实际运行每个 stage。这里每个 CTA 有 4 个 warpgroup，总共 512 个线程，它们不是按触碰的数据来切分，而是按 warpgroup 执行的*工作类型*切分： + +- WG3 驱动硬件引擎：TMA load、MMA 和 TMA store。 +- WG0、WG1 和 WG2 执行这些 engine 调用之间发生的 register-heavy math：softmax、correction 和 epilogue。 + +精确角色表是： + +| Owner | Role | What it does | +|-------|------|--------------| +| WG3, warp 1 | TMA load | 从 GMEM 加载 Q、K、V tile 到 SMEM | +| WG3, warp 0 | MMA | 发起 score MMA 和 value MMA | +| WG3, warp 2 | TMA store | 把最终 O tile 从 SMEM store 到 GMEM | +| WG0 | Q stage 0 的 softmax | 读取 TMEM 中的 S，计算 P，把 P 写回 TMEM | +| WG1 | Q stage 1 的 softmax | 为第二个 Q pipeline stage 执行同样工作 | +| WG2 | Correction 和 epilogue | Rescale TMEM 中的 O，normalize，并 stage 输出 | + +很容易把 “两个 Q stage” 误读成两个 attention head，但它们不是。它们只是 Q pipeline 中的两个 slot，由 WG0 拥有一个、WG1 拥有另一个，因此两个 Q tile 可以同时在飞行中。这就是 softmax 工作出现两次的原因：WG0 一次，WG1 一次。 + +代码用符号坐标选择这些角色： + +```python +wg_id = T.warpgroup_id([4]) +warp_id = T.warp_id_in_wg([4]) +``` + +阅读 kernel 时，先找到 role branch。它会告诉你嵌套在其中的每个 tile primitive 由哪一队线程拥有。 + +- WG3 warp 1 启动 TMA load 命令。一个 elected lane 发起 copy，TMA engine 移动 tile。 +- WG3 warp 0 发起 `tcgen05.mma` 指令。 +- WG0 和 WG1 以完整 warpgroup scope 运行 softmax。 +- WG2 以完整 warpgroup scope 运行 correction 和 epilogue。 + +有一个不对称性会塑造整个 barrier graph：*所有* MMA，无论 score 还是 value，都只由 WG3 warp 0 发起。WG0 和 WG1 完全不发起 MMA。它们只消费 score tile，运行 softmax，并把 `P` 写回 TMEM。 + +正因为这种分离，softmax 周围需要 barrier。`s_ready` 把 score tile 从 MMA warp 传给 softmax；`p_o_rescale` 传递 `P`，以及一个对 value MMA 安全的 `O` slot：要么已经 rescale，要么因为不需要 rescale 而被 release。本章后面会反复回到这两个名字。 + +## 读取代码片段 + +本章的代码片段摘自 [`flash_attention4.py`](https://github.com/mlc-ai/tirx-kernels/blob/main/tirx_kernels/attention/flash_attention4.py)，因此不可避免会引用一些没有在本章完整复现的 kernel 名称。自解释的名字（`wg_id`、`warp_id`、`BLK_M`/`BLK_N`、`HEAD_DIM`、`kv_stage`、`SMEM_PIPE_DEPTH_*` / `TMEM_PIPE_DEPTH` 深度、`should_accumulate`，以及这里为 1 的 `CTA_GROUP`）会在第一次相关时引入。其他名字在下表给出一行说明，这样当片段出现陌生名字时可以立即查到： + +| Name | Meaning | +|------|---------| +| `q_stage`, `i_q` | Q pipeline stage，0 或 1，即哪个 Q tile slot（`SMEM_PIPE_DEPTH_Q = 2`）。在 WG0/WG1 softmax 内部，warpgroup 自己的 `wg_id`（0 或 1）就是这个 stage index，因此 `S_region[q_stage]`、`P_region[wg_id]` 和 `O_region[i_q]` 都选择同一个 Q stage | +| `MMA_N` | score/output tile 在 TMEM column 上的宽度（128） | +| `MMA_K` | `P`/`V` column 上的 MMA inner-K step（16）；`K_SPLIT = 6 * MMA_K = 96` | +| `K_SPLIT` | value-MMA schedule 的 split point（见 *The Two MMA Phases*）；第一次 value MMA 覆盖 columns `0:K_SPLIT`（`6 * MMA_K = 96`） | +| `should_rescale` | WG2 per-row flag：旧 `O` 是否需要在下一次 value MMA 前 rescale（通过 `any_sync` 跨 warpgroup reduce） | +| `rescale_threshold` | 小 row-max 变化的 skip threshold；当前 kernel 使用 `8.0`，被跳过的 rescale 会把 `acc_scale` 精确设为 `1.0` | +| `scale_log2` | log2 单位下的 softmax scale，`log2(e)/√d`，因此 `P = exp2((S - m) · scale_log2)` | +| `acc_scale` | softmax 通过 SMEM mailbox 传给 WG2 的 per-row rescale factor | +| `chunk_start`/`chunk_end`, `p_start`/`p_end` | 正在读/写的 32-wide softmax chunk 的 column range | + +## 两个 MMA Phase + +对于每个 streamed K/V tile，Flash Attention 会运行两个 MMA phase，并由 softmax 连接： + +```text +Q, K -> score MMA -> S +S -> softmax -> P +P, V -> value MMA -> O +``` + +可以把它看成三个连续 producer 的 pipeline。第一个 MMA 产生 attention score `S`，softmax 把 `S` 变成 numerator `P`，第二个 MMA 消费 `P` 来更新 output accumulator `O`。对 `row_sum` 的 normalization 会推迟到 epilogue，等每个 K/V tile 都贡献之后再做。 + +下面每个 tile op 都会得到和 GEMM step 一样的 **scope / layout / dispatch** 卡片，并额外加入 **Handoff**，命名把 tile 交给下一角色的 barrier。 + +Compute 代码从不直接说 raw TMEM column number。Kernel 会把单个 TMEM allocation 切成 per-stage view（`S_region`、`P_region`、`O_region`），并用 pipeline stage 索引它们（`S_region[q_stage]`、`O_region[i_q]`、`P_region[i_q, 0:K_SPLIT]`）。这些 view 在 [TMEM Layout and Reuse](#tmem-layout-and-reuse) 一节中用 `T.TMEMStages` 定义；目前把每个 region 当作同一块物理 TMEM 的命名 slice 即可。 + +### Score MMA + +两个 phase 中的第一个是 score MMA，也就是每个 K/V iteration 开头的 matmul。它计算： + +$$S = Q_{\text{block}}K_{\text{block}}^{\top}$$ + +并把 `128 x 128` score tile 写入 TMEM： + +```python +Tx.warp.gemm_async( + S_region[q_stage], + Q_smem[q_stage, 0:BLK_M, 0:HEAD_DIM], + K_smem[kv_stage, 0:BLK_N, 0:HEAD_DIM], + dispatch="tcgen05", + cta_group=CTA_GROUP, +) +if T.ptx.elect_sync(): + s_ready.arrive(q_stage) +``` + +我们可以问 GEMM 章节对每个 tile op 都问过的四个问题：谁运行它、tile 位于哪里、如何 dispatch、如何 handoff： + +> **Tile-primitive readout：Score MMA** +> - Scope：WG3 warp 0 发起；一个 elected lane arrive `s_ready`。 +> - Layout：SMEM 中的 Q、K -> TMEM 中的 `S`（`S_region[q_stage]`）。 +> - Dispatch：`tcgen05`。 +> - Handoff：`s_ready`（-> softmax）。 + +单个 elected thread 在 `s_ready` 上 arrive，就是整个 handoff。它宣布这个 score tile 已经完成，softmax warpgroup 现在可以读取它。 + +### 两个 MMA 之间的 Softmax + +两个 MMA 中间是 softmax，它把 score tile `S` 变成 numerator tile `P`。它的 readout card 是： + +> **Tile-primitive readout：Softmax** +> - Scope：WG0（Q stage 0）/ WG1（Q stage 1），完整 warpgroup。 +> - Layout：TMEM 中的 `S` -> registers -> fp16 TMEM 中的 `P`（`P_region[wg_id]`）。 +> - Dispatch：用 `tcgen05.ld` 读取，用 TMEM store 写入；中间在 registers 中做 row-wise math。 +> - Handoff：等待 `s_ready`；arrive `p_o_rescale`（前 96 列）和 `p_ready_2`（最后 32 列）。 + +这个 stage 完全没有 GEMM 对应物。WG0/WG1 等待 score tile 在 `s_ready` 上到达，然后按 register-sized chunk 从 TMEM 读出： + +```python +Tx.copy_async( + s_chunk[:, chunk_start : chunk_end], + S_region[wg_id, chunk_start : chunk_end], +) +``` + +这是一个 warpgroup scope 下的 TMEM-to-register tile read。Scores 进入寄存器后，softmax warpgroup 按顺序做三件事： + +1. 计算 row max 和 row sum， +2. 计算 softmax numerator tile `P`， +3. 把 `P` 作为 fp16 写回 TMEM。 + +最后一步如下： + +```python +Tx.copy_async( + P_region[wg_id, p_start : p_end], + p_chunk[:, p_start : p_end], +) +``` + +既然刚刚在寄存器中算出了 `P`，为什么还要写回 TMEM？因为 value MMA 需要 `P` 作为一个 *tile operand*，而 MMA 不能把散落在每个线程中的 scalar register 当成矩阵读取。在这个 kernel 中，`P` 的 MMA-readable 形式是 `P_region`，也就是 fp16 TMEM alias `tmem_as_f16` 上的一个 view。因此这个 writeback 不是多余移动，而是把 `P` 放到下一次 MMA 真正能消费的唯一形状中。 + +### Value MMA + +第二个 phase，也是每个 K/V iteration 收尾的 phase，是 value MMA。它计算： + +$$O = O + P_{\text{block}}V_{\text{block}}$$ + +当这个 MMA 运行时，`O` 已经处在当前 K/V block 所需的正确状态：第一个 block 上被初始化，后续 block 上被 rescale。因此 MMA 只需要累加。它与 GEMM 的区别在于 operand 位于哪里：A operand 是 TMEM 中的 `P`，B operand 是 SMEM 中的 `V`，accumulator `O` 也在 TMEM 中： + +```python +# First sub-MMA: columns 0:K_SPLIT (the first 96 of P / rows of V). +Tx.warp.gemm_async( + O_region[i_q], + P_region[i_q, 0:K_SPLIT], + V_smem[kv_stage, 0:K_SPLIT, 0:HEAD_DIM], + transB=True, + accum=should_accumulate, + dispatch="tcgen05", + cta_group=CTA_GROUP, +) +# The second sub-MMA (same form, accum=True, gated on p_ready_2) covers the +# remaining columns K_SPLIT:BLK_N. +``` + +> **Tile-primitive readout：Value MMA** +> - Scope：WG3 warp 0。 +> - Layout：TMEM 中的 `P` + SMEM 中的 V -> TMEM 中的 `O`（`O_region[i_q]`）。 +> - Dispatch：带 TMEM operand 的 `tcgen05`。 +> - Handoff：等待 `p_o_rescale`、`p_ready_2`、`kv_load.full`；arrive `o_ready`（-> epilogue）。 + +这个 operand placement 是两个 MMA 之间的硬件差异： + +- Score MMA 从 SMEM 读取两个 operand：Q 和 K。 +- Value MMA 从 TMEM 读取一个 operand：`P`。 +- Value MMA 从 SMEM 读取另一个 operand：V。 +- 结果累加到 TMEM 中的 `O`。 + +`accum=should_accumulate` flag 实现了算法中的“初始化或累加”选择：它在 query block 的第一个 K/V tile 上为 false，在后续每个 tile 上为 true。 + +你还会注意到，value MMA 不是一次性运行，而是拆成 `96 + 32` schedule： + +1. Softmax 以四个 32-column chunk 写入 `P`。 +2. 前三个 chunk 就绪后，value MMA 立即开始处理 `P` 的前 96 列以及 V 中匹配的行。 +3. 最后 32 列等待 `p_ready_2`。 +4. 第二次 MMA 消费最后这个 chunk，并完成 tile。 + +拆分的原因是让 Tensor Core 保持忙碌。如果把 value MMA 作为一次单独指令运行，整个 phase 都会停到所有四个 32-column `P` chunk 都 exponentiate 并 store 完成。先对前三个 chunk firing，kernel 就能把最后一个 chunk 的 `exp` 和 TMEM write 与一个已经在飞行中的 96-wide MMA 重叠起来，把原本的空闲时间变成有用工作。 + +## TMEM Layout and Reuse + +`S`、`P` 和 `O` 都必须共享同一个 `128 x 512` TMEM allocation，而它们如何打包进这块空间，正是 barrier 和 layout 在这个 kernel 中不可分割的原因。 + +下图直接展示这种 packing：score slot、numerator slot 和 output slot 全都共享同一个 TMEM allocation，因此 barrier protocol 才能让复用合法。 + +![TMEM Layout](../../img/tmem_layout_v3.png) + +这张图可以读作一组 tile slot： + +- Score slot 保存 `S = QK^T`。 +- Numerator slot 保存 softmax exponentiation 之后的 `P` tile。 +- Output slot 保存 fp32 `O` accumulator。 + +这些不是独立 buffer。它们是*同一个* allocation 的区域，而这种共享不是风格选择，而是被迫的。Q-pipeline depth 为 2 时，两个 `S` slot（2 × MMA_N = 256 列）和两个 `O` slot（2 × MMA_N = 256 列）已经占满全部 512 个 fp32 列。没有剩余空间给 `P`，所以 `P` 只能通过更窄的 fp16 view alias 同一批 byte。这样安全的唯一原因是每个 region 都严格在前一个 consumer 完成后才复用，而这个时机正是 barrier 保证的。因此在 FA4 中，barrier 不只是 scheduling；它们首先让 layout 合法。 + +Aliasing 通过 `T.TMEMPool` 设置。Kernel 先取一个 fp32 view（`tmem`）用于 score 和 output accumulator，然后把 pool base rewind 回 0，再在*同一批*物理 byte 上取第二个 fp16 view（`tmem_as_f16`）： + +```python +tmem_pool = T.TMEMPool(pool, total_cols=N_COLS_TMEM, cta_group=CTA_GROUP, tmem_addr=tmem_addr) +tmem = tmem_pool.alloc((128, N_COLS_TMEM), "float32") +tmem_pool.move_base_to(0) +tmem_as_f16 = tmem_pool.alloc((128, N_COLS_TMEM * 2), "float16") +tmem_pool.commit() +``` + +因为 fp16 元素宽度只有一半，fp16 view 会在同样 byte 上暴露出两倍数量的可索引 column，这正是 `P` 所处的空间，也是 fp32 layout 没有空间容纳的区域。有了两个 view 后，kernel 用 `T.TMEMStages` 把 `S`、`P` 和 `O` slot 切成 staged region，使 compute 代码可以按 pipeline stage 索引，而不是按 raw column： + +```python +S_region = T.TMEMStages(tmem, col_start=0, width=MMA_N, stages=SMEM_PIPE_DEPTH_Q, stride=MMA_N) +O_region = T.TMEMStages(tmem, col_start=MMA_N * SMEM_PIPE_DEPTH_Q, width=MMA_N, stages=SMEM_PIPE_DEPTH_Q, stride=MMA_N) +P_region = T.TMEMStages(tmem_as_f16, col_start=MMA_N, width=BLK_N, stages=SMEM_PIPE_DEPTH_Q, stride=MMA_N * 2) +``` + +`P_region` stride 中的 `* 2` 是 aliasing 在代码中可见的地方。`S_region` 和 `O_region` 以 fp32 `tmem` column 计量，而 `P_region` 以 fp16 `tmem_as_f16` column 计量，后者宽度只有一半，因此 stage-to-stage movement 需要 doubled stride 才能落到同一批物理 byte 上。不过一旦 region 定义好，compute 代码就很干净：它写 `S_region[q_stage]`、读 `S_region[wg_id, ...]`、写 `P_region[wg_id, ...]`，并累加到 `O_region[i_q]`，完全不触碰 raw column index。 + +**Try with your agent**：让它解释这个 FA4 kernel 中的 fp32（`tmem`）和 fp16（`tmem_as_f16`）view。哪些物理 TMEM region 保存 `S`、`P` 和 `O`？为什么 `P_region` 的 stride 使用 `MMA_N * 2`？复用问题留到下一节：看完 barrier table 后，再检查每个 region 复用前必须等哪些 consumer 完成。 + +## Barrier 如何连接角色 + +这是 kernel 最难的部分，所以最好循序渐进。先从沿主 compute path 移动数据的几组 barrier 开始，把其他 barrier 当作稍后可查的 bookkeeping。Data-ready handoff 是： + +| Handoff | Meaning | +|---------|---------| +| TMA load -> score/value MMA | Q、K 或 V 已经到达 SMEM，可以喂给 MMA | +| score MMA -> softmax | `S` 已经在 TMEM 中 ready | +| softmax/correction -> value MMA | `P` 已经在 TMEM 中 ready，且 `O` 可以安全 accumulation | +| value MMA -> epilogue | 最终 `O` 已经在 TMEM 中 ready | +| epilogue -> TMA store | `O_smem` 已经 ready，可以 store | + +不在这个列表里的都是 pipeline bookkeeping：release 某个 SMEM、TMEM 或 staging buffer，让另一个角色可以复用它。有用之处在于，每个 barrier 无论传递数据还是只做 bookkeeping，都可以按同一种 tile handoff 来读。你只要问谁生产数据、谁消费数据，以及双方完成后哪个 buffer 变得空闲。 + +下一张图把这些 handoff 压缩成两个 MMA phase 的精确 readiness gate：score MMA 等待什么，value MMA 累加前必须等待什么。 + +![Flash Attention 4 MMA Input Gates](../../img/flash_attention_main_handoff.png) + +把这张图看成 correctness gate，而不是 schedule。它回答“这个 MMA firing 前必须满足什么”，不说明 timing。Score MMA 等待 Q 和 K 在 SMEM 中 ready，然后产生 `S`。Value MMA 同时等待三件事：SMEM 中的 V、softmax 产生的 `P` tile，以及 WG2 已经 release 或 rescale 过的 `O` slot。Softmax-to-value gate 被拆分，原因前面已经见过：当 `P` 的前 96 列就位时 value MMA 可以开始，`p_ready_2` 释放最后 32 列。 + +有一个 handoff 不符合 tile-readiness 模型：softmax-to-correction edge。Softmax 不是传递 tile，而是通过一个 one-slot SMEM mailbox 向 WG2 传递单个 scalar（K/V loop 中是 `acc_scale`，epilogue 中是最终 `row_sum`）。由于这个 slot 每次 iteration 都会复用，一对 `full`/`empty` barrier 必须保护它。 + +下图放大了 mailbox handshake，因此这对 barrier 应该读作 scalar producer-consumer channel，而不是 tile-ready gate。 + +![Flash Attention 4 Softmax Scale-Slot Handshake](../../img/flash_attention_softmax_correction.png) + +把 `softmax_corr.full` 和 `softmax_corr.empty` 读作一对 producer-consumer： + +1. Softmax 在复用 scale/sum slot 前等待 `softmax_corr.empty`。 +2. Softmax 把 `acc_scale` 或最终 `row_sum` 写入这个 slot。 +3. Softmax arrive `softmax_corr.full`。 +4. WG2 等待 `softmax_corr.full`，然后读取这个 slot。 +5. WG2 arrive `softmax_corr.empty`。 +6. Softmax warpgroup 可以在下一阶段复用这个 slot。 + +要小心 `softmax_corr.empty` 的含义和不含义。它只表示 WG2 已经消费 scale/sum slot。它不表示 `P` 已经 ready，更*不是*允许 value MMA 启动的 gate。那个 gate 是 `p_o_rescale`，它在 `P` 的前 96 列写入、且 `O` slot 可以安全累加时触发。混淆二者是经典 wrong-result bug 来源。 + +有了主路径之后，完整 barrier list 可以作为参考： + +| Barrier | Producer -> consumer | What becomes safe | +|---------|----------------------|-------------------| +| `q_load.full` | TMA load -> score MMA | Q SMEM tile 可以喂给 MMA | +| `q_load.empty` | 这个 Q stage 的所有 score MMA -> TMA load | Q SMEM stage 可以为下一个 task 复用 | +| `kv_load.full` | TMA load -> score/value MMA | K 或 V SMEM tile 可以喂给 MMA | +| `kv_load.empty` | score/value MMA -> TMA load | K/V SMEM stage 可以复用 | +| `s_ready` | score MMA -> softmax | S TMEM tile 可以读取 | +| `p_o_rescale` | softmax + WG2 -> value MMA | P 的前 96 列在 TMEM 中，且 O slot 可用于 value MMA | +| `p_ready_2` | softmax -> value MMA | P 的最后四分之一在 TMEM 中 | +| `o_ready` | value MMA -> epilogue | 最终 O accumulator ready | +| `softmax_corr.full` | softmax -> WG2 | `acc_scale` 或最终 `row_sum` 在 SMEM mailbox 中 ready | +| `softmax_corr.empty` | WG2 -> softmax | WG2 读取后，同一个 SMEM mailbox slot 可以复用 | +| `corr_epi.full` | epilogue -> TMA store | O_smem ready，可以 store | +| `corr_epi.empty` | TMA store -> epilogue | O_smem stage 可以复用 | + +和 GEMM 一样，可以从 signal producer 推断 barrier type： + +- TMA load 使用 `TMABar`，因为 TMA engine 会 byte-count 自己的完成。 +- MMA completion 使用 `TCGen05Bar`，因为 `tcgen05.commit` signal completion group。 +- 纯 thread-to-thread handoff 使用 `MBarrier`，参与线程显式 arrive。 + +拆分的 softmax-to-value handoff 值得再仔细看。它使用两个 gate： + +- `p_o_rescale` 在 `P` 的前 96 列写入且 `O` tile 可以安全累加时，允许 value MMA 开始。 +- `p_ready_2` 释放 `P` 的最后 32 列，匹配上一节的 `96 + 32` value-MMA schedule。 + +第一个 K/V block 是简单情况。WG2 会 pre-arrive `p_o_rescale`，因为还没有旧的 `O` tile 需要 rescale。 + +后续 block 必须更小心。WG2 只有在跳过不必要 rescale 或完成旧 `O` 的 rescale 后，才 arrive `p_o_rescale`。Skip test 故意保守：softmax 计算 log2-scaled delta `(m_old - m_new) * scale_log2`；如果这个值仍高于 `-rescale_threshold`，新 max 没有移动到值得 rescale 的程度，因此 kernel 保留旧 max，并把 `acc_scale` 精确设为 1.0。只有更大的 max jump 才会走 `exp2` 路径，并要求 WG2 rescale `O`。 + +WG2 随后用 `any_sync` 在 warpgroup 内 reduce `should_rescale`。如果没有 row 需要更新，它就让 `O` 保持原样。这个 skip 很重要，因为 rescale `O` 是覆盖整个 accumulator 的完整 TMEM -> RF -> TMEM read-modify-write；当 threshold logic 已经把 `acc_scale` 保持在 1.0 时，这就是纯浪费。 + +注意，所有新 barrier 都聚集在一个地方。`s_ready`、`p_o_rescale`、`p_ready_2` 和 softmax/correction pair 都是 softmax 周围的 barrier。它们存在的唯一原因是：score MMA 和 value MMA 不再相邻。Register math、TMEM rewrite 和 output rescaling 现在位于两者之间，每一步都需要自己的 handoff。 + +**Try with your agent**：让它 trace 一个 K/V block 穿过 `s_ready`、`p_o_rescale`、`p_ready_2` 和 `o_ready`。对每个 barrier，问谁 wait、谁 arrive、哪个 tile 变得可读，以及之后哪个 storage 可以复用。 + +## Pipelining Structure + +Barrier 告诉我们角色消费 tile 前必须有什么*ready*。但它们没有告诉我们实际有哪些东西*并发*运行，这正是现在要讨论的问题。二者确实不同：一个 correctness gate 可能在 producer 实际运行之前很久或之后很久才满足。 + +这里没有单一 pipeline depth，因为不同 tile stream 以不同速度移动。因此 kernel 为每类 stream 保持单独的 ring： + +- Q pipeline depth 2：一个 CTA 处理两个 Q stage。WG0 处理一个 stage，WG1 处理另一个。 +- KV pipeline depth 3：K 和 V block 在 inner loop 中 streaming，而同一批 Q stage 会复用。 +- TMEM pipeline depth 2：每个 Q stage 有自己的 S/P/O TMEM slot，这些 slot 会在匹配 barrier 触发后复用。 + +下图从 correctness gate 切换到 timeline view，展示这些独立 ring 进入飞行状态后，哪些角色可以大致同时 active。 + +![Flash Attention 4 Pipeline Structure](../../img/flash_attention_pipeline_v2.png) + +把它读作 timeline，而不是 barrier graph。它展示大致同一时间哪些角色 active；前面的 barrier-flow 图则用于检查精确 producer-consumer wait。两张图分别回答本节开头提出的两个不同问题。 + +每一行对应代码中的一个 role branch： + +- WG3 warp 1 发起 TMA load。 +- WG3 warp 0 发起 score MMA 和 value MMA。 +- WG0 和 WG1 为两个 Q stage 运行 softmax。 +- WG2 release 或 rescale `O`，稍后 normalize 最终输出。 +- WG3 warp 2 发起 TMA store。 + +从左到右沿图走，可以看到一个代表性的 pipeline wave。Load warp 从 `Q0`、`K[n-1]`、`Q1`、`V[n-1]` 开始，然后继续 stream 更低索引的 K/V block。MMA warp 发起第一批 score MMA 来产生 `S0` 和 `S1`，WG0/WG1 把它们变成 `P0` 和 `P1`。 + +重要的是，MMA warp 并不会先运行所有 score MMA，再运行所有 value MMA。一旦两个 Q stage 都 primed，它就会交错两类 MMA：当前 `V` block 的 value MMA，然后下一个 `K` block 的 score MMA，如此继续： + +```text +score Q0*K[n-1] +score Q1*K[n-1] +value P0*V[n-1] +score Q0*K[n-2] +value P1*V[n-1] +score Q1*K[n-2] +value P0*V[n-2] +... +``` + +这种 interleaving 正是图中 score、softmax、correction 和 value 行重叠，而不是整齐依次运行的原因。 + +WG2 行标注为 `release / rescale`，两个半部分对应我们已经见过的两种情况。第一个 K/V block 上还没有旧 `O`，因此 WG2 只参与让 value MMA 继续的 handoff；后续 block 上，它可能在 value MMA 累加前 rescale 旧 `O`。Normalization 和 TMA store 只在 attention task 的最后一个 K/V block 之后发生一次。 + +没有单个 GEMM-style pipeline 可以描述 FA4，因为 Q、K/V 和 TMEM slot 都按独立 schedule 前进。TIRx 把这些 schedule 显式保留下来，表现为单独的 tile buffer、`PipelineState` cursor 和 barrier phase，而不是把 kernel 藏在一个 monolithic primitive 后面。代价是 moving parts 更多；收益是复杂性保持可见、可检查。 + +## Rescaling 和 Writeback + +Rescale 是强制需要的，不是可以去掉的优化。Online softmax 可能随着每个新 score tile 抬高 per-row maximum；一旦抬高，早期 block 累加出的 `O` 是按*旧* maximum 缩放的。这会让早期每一项都比正确值大 `exp(m_new - m_old)` 倍。跳过 correction 后，这些 block 会被过度加权，最终输出就是错的。修复是一次 TMEM -> registers -> TMEM tile operation： + +$$O_{\text{old}} \leftarrow O_{\text{old}} \cdot e^{(m_{\text{old}} - m_{\text{new}}) / \sqrt{d}}$$ + +工作由两个角色拆分。Softmax 计算 per-row scale，并把它放入 SMEM mailbox；WG2 等待 `softmax_corr.full`，从 TMEM 读出当前 `O`，乘上这个 scale，再把 `O` 写回： + +```python +RESCALE_TILE = T.meta_var(16) +o_row = T.wg_reg_tile(RESCALE_TILE) +Tx.copy_async(o_row, O_region[i_q, d_start : d_start + RESCALE_TILE]) +Tx.mul(o_row, o_row, acc_scale) +Tx.copy_async(O_region[i_q, d_start : d_start + RESCALE_TILE], o_row) +T.ptx.tcgen05.wait.st() +``` + +需要强调，这是覆盖完整 `O` accumulator 的 TMEM -> registers -> TMEM tile operation，不是一点 scalar bookkeeping，它和其他 stage 一样有 readout card： + +> **Tile-primitive readout：Correction（rescale）** +> - Scope：WG2，完整 warpgroup。 +> - Layout：TMEM 中的 `O` -> registers -> TMEM 中的 `O`（`O_region[i_q]`）。 +> - Dispatch：用 `tcgen05.ld` 读取，用 TMEM store 写入；中间做 register multiply。 +> - Handoff：等待 `softmax_corr.full`；arrive `p_o_rescale`（-> value MMA）和 `softmax_corr.empty`（-> softmax）。 + +端到端同步如下： + +1. Softmax 把 scale value 写入 SMEM。 +2. WG2 等待 `softmax_corr.full`。 +3. WG2 在 TMEM 中 rescale `O`。 +4. WG2 arrive `p_o_rescale`。 +5. WG3 的 value MMA 现在可以消费 `P` 并累加到 rescaled `O` tile 中。 + +WG2 读取完 SMEM slot 后，`softmax_corr.empty` release 这个 slot，loop 闭合，softmax 可以在下一次 iteration 复用 mailbox。 + +K/V loop 结束后，WG2 从 correction 切换到 epilogue。它等待最终 `row_sum` 和 `o_ready`，从 TMEM 读取最终 `O`，乘以 `1 / row_sum`（最开始被我们推迟的 normalization），cast 到 fp16，并写入 `O_smem`。WG3 的 TMA store warp 随后把 `O_smem` 带回 GMEM。 + +对打算扩展这个 kernel 的人，有一个限制值得标明。它只计算 forward output，而 training forward pass 通常还会保存 backward pass 需要的 log-sum-exp（LSE）。加入 LSE 时有一个 scaling 细节要记住：这个 kernel 把 `row_max` 保持为*未缩放* raw `QK^T` score 的最大值，而 `row_sum` 累加的是 `exp((S - row_max) / sqrt(d))`。因此形成 natural-log LSE 时，必须把 `1/\sqrt{d}` 因子重新应用到 `row_max` 上： + +$$\mathrm{LSE}_i = \log(\mathrm{row\_sum}_i) + \mathrm{row\_max}_i / \sqrt{d}$$ + +这个实现只输出 forward result，不写 LSE。 + +## Causal Masking + +Causal attention 增加了一个约束：query 只能 attend 到自己位置及之前的 key。Kernel 用两个互补方式满足这个约束，一个便宜，一个精确。 + +便宜方式是完全跳过 work。很多 K/V block 对某个 Q block 来说完全位于 diagonal 上方，没有任何贡献，因此 `get_n_block_max(...)` 会计算这个 block 可能需要的最后一个 block，loop 直接不加载也不计算剩下部分。 + +精确方式处理跨越 diagonal 的 block，其中有些 column 有效，有些无效。这些 block 仍然运行 score MMA，但 softmax 会在 exponentiation 前 mask 掉无效 column。对于每一行，它从该行 query 位置和 block offset 推导 column limit，保留不超过这个 limit 的 column，并把超过它的每个 column 在寄存器中设为 `-inf`，因此这些 column 不会贡献 row max，也不会贡献 `exp2` numerator。 + +实现不会逐元素 branch，而是用 `mask_r2p(...)` 应用这个 limit，把它变成覆盖整个 32-wide score chunk 的 bit mask，并一次性 mask 这个 chunk。完全位于 diagonal 下方的 block 保留所有 column，不需要 mask。 + +从 tile-primitive 视角看，causal mode 完全不重写数据路径。它只裁剪 K/V trip count，并在 register-resident softmax 中、score MMA 和 `P` writeback 之间插入一个 masking step。 + +## GQA 支持 + +Grouped Query Attention 让多个 query head 共享一个 K/V head。这节省 memory bandwidth，但提出一个 packing 问题：如何保持只有一个 K/V tile，同时仍然把许多 query head 喂给它？Kernel 的答案是一次处理一个 query-head group，让它们共同使用一个 scheduled `kv_head_idx`： + +```python +GQA_RATIO = num_qo_heads // num_kv_heads +SEQ_Q_PER_TILE = BLK_M // GQA_RATIO +``` + +技巧是重新解释 128 个 Q-tile row。对于 `GQA_RATIO=4`，它们不再表示 128 个 sequence position；而是表示 32 个 sequence position 乘以 4 个 query head，并打包在一起，使四个 head 搭乘同一个 K/V tile。Row 解码为： + +```text +seq_pos = row // GQA_RATIO +q_head = row % GQA_RATIO +``` + +Q load 用一个 3D view 表达这种 packing。Source 是自然的 `Q[batch, seq, qo_head, dim]` 布局，destination 是稍后 score MMA 会作为 flat `128 x HEAD_DIM` operand 读取的同一个 SMEM tile。View 负责调和两者，而且不需要任何 copy： + +```python +Q_smem_3d = Q_smem.view(SMEM_PIPE_DEPTH_Q, SEQ_Q_PER_TILE, GQA_RATIO, HEAD_DIM) +Tx.copy_async( + Q_smem_3d[i_q, :, :, :], + Q[batch_idx, + m_start : m_start + SEQ_Q_PER_TILE, + kv_head_idx * GQA_RATIO : (kv_head_idx + 1) * GQA_RATIO, + :], + **tma_copy_q, +) +``` + +K 和 V 从不在内存中扩展，这正是 GQA 的意义：`kv_head_idx` 对应的单个 K/V tile 被打包到 Q row 中的所有 `GQA_RATIO` query head 复用。输出侧镜像输入：epilogue 后用匹配的 3D view 把 packed row store 回 `O[batch, seq, qo_head, dim]`。 + +结果是 GQA 完全存在于 Q-load 和 O-store 边界。Compute path 内部，score MMA 仍然看到普通的 `128 x HEAD_DIM` Q tile，tile-primitive graph 的其余部分不变。 + +## Tile Scheduling + +Scheduler 的工作是把每个 CTA 映射到一个 `(batch, kv_head, m_block)` attention task，正确策略取决于 masking 是否让这些 task 成本相同： + +- Non-causal mode 使用 `FlashAttentionLinearScheduler`。每个 task 工作量相同，因此固定 CTA pool 按 `num_ctas` 前进就足以均匀分配。 +- Causal mode 使用 `FlashAttentionLPTScheduler`，因为 causal masking 会让 work 极不均匀：靠近开头的 Q block 大约只 attend 一个 K/V block，而靠近结尾的 block 会 attend 全部 block。朴素切分会让一些 CTA 比其他 CTA 晚很多完成，所以 longest-processing-time scheduler 会优先安排重 block，使结束时间更均衡，同时仍然把邻近 batch/head task 放在一起以改善 L2 locality。 + +二者虽然不同，但暴露同样的 loop interface： + +```python +while scheduler.valid(): + m_block_idx = scheduler.m_block_idx + batch_idx = scheduler.batch_idx + kv_head_idx = scheduler.head_idx + # process one Q block against its K/V block range + scheduler.next_tile() +``` + +唯一行为差异在 `next_tile()` 做什么：non-causal mode 中它把 CTA 前进到另一个 task；causal mode 中它在当前 task 后结束 loop。无论如何，这只是 scheduling decision：它选择 CTA 拥有*哪个* attention tile，从不改变 tile 如何计算。Loop 内部运行同样的 local primitive：TMA load、score MMA、softmax、value MMA、correction、TMA store。 + +## 编译和验证 + +上面都是摘录，所以要把它们合起来实际运行 kernel，我们从 `tirx-kernels` import 真实实现，编译它，并用 torch reference 检查。完整 kernel，即本章讲过的所有部分组合成一个文件，位于 `tirx-kernels` 仓库中的 [`flash_attention4.py`](https://github.com/mlc-ai/tirx-kernels/blob/main/tirx_kernels/attention/flash_attention4.py)。与 GEMM verify cell 有两个不同点：Flash Attention 有更丰富的入口（`get_flash_attention4_kernel`），并且额外接受 `profiler_buf` 参数用于内置 profiler。整个章节只需要运行这一段： + +```python +import torch +import torch.nn.functional as F +import tvm +from tirx_kernels.attention.flash_attention4 import ( + get_flash_attention4_kernel, PROFILER_BUFFER_SIZE) + +B, S, Hq, Hkv, D = 1, 1024, 32, 8, 128 # GQA: 32 query heads share 8 KV heads +Q = torch.randn(B, S, Hq, D, dtype=torch.float16, device="cuda") +K = torch.randn(B, S, Hkv, D, dtype=torch.float16, device="cuda") +V = torch.randn(B, S, Hkv, D, dtype=torch.float16, device="cuda") +O = torch.empty(B, S, Hq, D, dtype=torch.float16, device="cuda") +prof = torch.zeros(PROFILER_BUFFER_SIZE, dtype=torch.uint64, device="cuda") + +kernel = get_flash_attention4_kernel(B, S, S, Hq, Hkv, D, is_causal=False) +target = tvm.target.Target("cuda") +with target: + ex = tvm.compile(tvm.IRModule({"main": kernel}), target=target, tir_pipeline="tirx") +ex.mod(Q, K, V, O, prof) # ex.mod takes torch tensors directly, like every other chapter +torch.cuda.synchronize() + +# torch reference; enable_gqa lets the 32 query heads share the 8 KV heads +qt, kt, vt = (x.transpose(1, 2).float() for x in (Q, K, V)) +ref = F.scaled_dot_product_attention(qt, kt, vt, enable_gqa=True).transpose(1, 2).half() +torch.testing.assert_close(O, ref, rtol=1e-2, atol=1e-2) +print(f"FA4: B={B} S={S} Hq={Hq} Hkv={Hkv} D={D}, non-causal -> PASS") +``` + +**期望输出**：`... -> PASS`。Kernel 以 fp32 累加 online softmax，但它的结果与高精度 reference 之间仍然存在多个近似来源。包括 input 和 operand 的 fp16 存储与 rounding；基于 `exp2` 的 softmax 重写（把每个 exponential 重写成 `scale_log2 = log2(e)/√d` 形式）；online-softmax 的重排和 per-row rescaling，它在 running scale 中分块求和，而不是一次性求和；最后还有 writeback 时对 `O` 的 fp16 cast。这里选择的 `rtol`/`atol` 与源码 kernel 自己测试使用的 tolerance 相同，是为了同时覆盖这些因素相对 torch reference 的误差，而不是只覆盖 fp16 rounding。所以如果这里出现真正 failure，而不是边缘 near-miss，请把它理解成指向 softmax path 的路标：可能漏了 `s_ready` / `p_o_rescale` / `p_ready_2` wait，或者 `row_max` / `row_sum` 更新没有正确应用到 rescale step。这些正是本章花大量 barrier 讨论的 handoff。 + +## 与 GEMM 的差异 + +下表沿变化的轴比较 FA4 和 GEMM： + +| Aspect | GEMM | Flash Attention 4 | +|--------|------|-------------------| +| MMA phases | 一个重复的 MMA | score MMA 和 value MMA | +| Work between MMAs | 除 pipeline handoff 外没有 | online softmax、masking 和 O rescaling | +| Running state | 只有 accumulator | row max、row sum、O accumulator | +| Main intermediate | accumulator TMEM tile | S、P 和 O TMEM tile region | +| Warp roles | TMA producer、MMA consumer、writeback | TMA load、MMA、softmax、correction、TMA store | +| Barriers | 主要是 load/compute/writeback handoff | 额外的 score/softmax/value/correction handoff | +| Scheduling unit | output matrix tile | attention task：`(batch, kv_head, m_block)` | + +这些差异都可以追溯到本章开头说的结构变化：第二个 MMA，以及夹在两个 MMA 之间的 softmax。底层 TIRx contract 完全没有改变： + +- tile primitive 说明哪个 tile 被移动或计算， +- 周围 scope 说明哪些线程协作， +- layout 说明 tile 位于哪里， +- barrier 说明下一个角色何时可以消费它。 + +因此，FA4 比 GEMM 难，不是因为它依赖不同硬件，而是因为 tile value 更多、它们之间的 handoff 更多。 + +## 练习 + +1. 与 GEMM 相比，FA4 的两个 MMA phase 之间出现了什么新的 tile handoff？命名 producer、TMEM tile 和 consumer。 +2. 为什么 softmax 要把 numerator tile `P` 写回 TMEM，而不是只把它保存在寄存器里供 value MMA 使用？ +3. 选择 `p_o_rescale` 或 `p_ready_2`。这个 barrier 精确证明了什么？如果 value MMA 跳过这个 wait，会出什么问题？ -本页用于放置 Flash Attention 4 kernel 构建过程的中文翻译。 +**Try with your agent**：选择一个没有标注的 tile primitive，例如 epilogue 中的 `Tx.copy_async`、fp32 -> fp16 的 `Tx.cast`，或第二个 `gemm_pv` sub-MMA。让它给出 scope / layout / dispatch / handoff 卡片，然后对照源码中的 guard、allocation 和 wait 检查答案。 diff --git a/zh/chapter_gemm_advanced/index.md b/zh/chapter_gemm_advanced/index.md index 21ef23c0..61833874 100644 --- a/zh/chapter_gemm_advanced/index.md +++ b/zh/chapter_gemm_advanced/index.md @@ -1,13 +1,910 @@ ---- -orphan: true ---- - (chap_gemm_advanced)= +# 用 Warp Specialization 和 Cluster 扩展 GEMM + +:::{admonition} 概览 +:class: overview + +- Pipelined GEMM 仍然让一个 warpgroup 按顺序执行 load、MMA 和 writeback，本章要移除的瓶颈正是这一点。 +- Step 7 把 warp specialization 成不同角色，Step 8 加入 2-CTA cluster，Step 9 加入多个 consumer。 +- 每一步都移除一个串行瓶颈，最终达到接近 state-of-the-art 的吞吐。 +::: + +上一章（{ref}`chap_gemm_async`）中的 pipelined GEMM 已经很快，但它仍然要求一个 warpgroup 做所有事：发起 load、运行 MMA、再把结果写回。即使有 software pipeline，这一组线程仍然是三个硬件引擎汇合的地方。 + +症状很容易看出来。Tensor Core 运行时 TMA unit 变安静，结果写回 memory 时 Tensor Core 变安静，每个引擎都通过同一组线程等待其他引擎。突破这个问题的方式，是不再让一组线程做所有事。 + +我们通过三个逐步扩大协作范围的 step 来贯彻这个想法。Step 7（{ref}`chap_warp_specialization`）把 warp specialization 成 producer、consumer 和 writeback 角色。Step 8（{ref}`chap_cta_cluster`）把两个 CTA 组成一个 cluster，并让它们跨 shared memory 共享 operand。Step 9（{ref}`chap_multi_consumer`）加入第二个 MMA consumer，让一个 staged tile 喂给两份数学计算。 + +把这三个 step 看作同一个 pattern 在不同尺度上的应用会很有帮助。Step 7 把完整 pipeline 保持在一个 CTA 内部：TMA 和 MMA 共享一个 warpgroup，而 writeback 在另一个 warpgroup 中运行。Step 8 把协作扩大到 CTA 之间，产生跨越两个 CTA 的 256×256 tile。Step 9 进一步提高 compute density：cluster output 增长到 512×256，每个 staged B tile 被两个 consumer 复用，我们到达本教程中最密集的变体。 + +贯穿这些变化的有一件事保持不变。SMEM、TMEM 和 register layout 仍然遵守前两章建立的 contract；变化的是*谁协作*，而不是数据如何布局。Step 8 是合作 scope 第一次越过单个 CTA，因此 operand tile 会切分到两个 CTA 的 shared memory 中，并且一个布局会沿 `cbx` cluster 轴跨越两个 CTA。 + + (chap_warp_specialization)= +## Step 7：Warp Specialization + Pipeline + +Single-warpgroup kernel 留下性能的原因很简单：每个线程都走同一条路径，先 load，再 compute，再 write。所以在它 load 时 Tensor Core 没事做，在它 compute 时 TMA engine 没事做。修复方式是 *warp specialization*。我们不再要求一组线程依次做每项工作，而是把每项工作交给专门的 warp，并让这些 warp 同时运行，中间由 software pipeline 串接。这是 GEMM 路径中最大的架构变化，本章剩余部分都建立在它之上。本节 benchmark 使用 M=N=K=4096。 + +> **本 step 改变的内容：Scope** +> - Scope：一个 warpgroup 顺序执行 load -> MMA -> writeback，变成三个并发角色（TMA producer、MMA consumer、writeback），并由 full/empty barrier 连接。 +> - Layout：不变，仍然是 Step 6 中的 SMEM stage 和 TMEM accumulator。 +> - Dispatch：不变，TMA load 和 `tcgen05` MMA。 + +**主题。** + +- Warp specialization：把不同 warp/warpgroup 专门分配给不同任务 + +- 高层 barrier abstraction：`TMABar`、`TCGen05Bar`、`MBarrier` + +- 用 `PipelineState` 自动管理 stage/phase + +- 用于 per-warpgroup synchronization 的 `warpgroup_sync` barrier ID + +（多 stage SMEM pipeline 和 persistent `ClusterPersistentScheduler2D` 会从 Steps 5-6 原样复用；这里只新增 scope split。） + +### 从顺序到并发 + +在介绍角色和 barrier 之前，先隔离 warp specialization 要移除的 scheduling bottleneck 会更清楚。下图用 Step-4-style sequential timeline 紧凑表示 Steps 4-6 中 specialization 前的 kernel，然后把它放在 Step 7 warp-specialized schedule 上方，让 engine utilization 的差异一眼可见。 + +![Warp Specialization Timeline](../../img/warp_specialization_timeline.png) + +上方是 specialization 前的 single-warpgroup pattern：同一组未 specialization 的线程同时拥有 load path 和 MMA path，所以一个 engine 在另一个 active 时很容易 idle。Steps 5 和 6 通过 double buffering 和 persistent scheduling 改善这个 baseline，但它们还没有把 loading 和 compute 拆成独立 producer/consumer 角色。下方的 specialization 打破了这种轮流执行。TMA producer 在 MMA consumer 忙于计算时 prefetch 下一个 tile，writeback 独立前进。Producer warp 3 在 consumer warp 0 仍在处理当前 MMA 时发起下一次 load，因此两个 engine 都不必等待对方。Load/MMA 交接使用两个 barrier： + +- **`tma2mma`**（TMA -> MMA）：表示已加载的 SMEM 数据已经准备好给 MMA 消费。 +- **`mma2tma`**（MMA -> TMA）：表示 MMA 已经读完一个 buffer，TMA 可以把它复用于下一次 load。 + +图中有一个细节第一眼看起来像错误：`mma2tma` 箭头会跳过一个 stage。原因是 ring buffer。`PIPE_DEPTH=2` 时有两个 SMEM buffer，stage 0 和 stage 1；TMA Load k=0 填充 buffer 0，TMA Load k=1 填充 buffer 1。当 MMA Compute k=0 读完 buffer 0 后，它 signal `mma2tma` 表示这个 buffer 空闲，但真正想要拿回 buffer 0 的 load 是 TMA Load k=2，而不是 k=1（它正在使用 buffer 1）。因此从 MMA Compute k=0 发出的 `mma2tma` 箭头会一直指向 TMA Load k=2。Release 跳过一个 stage，只是因为这个 ring 有两个 slot。 + +### Warp 角色 + +Timeline 展示了*为什么*要拆分 work；下一个问题是*谁*做每个部分。Specialization 把 load、compute、writeback 三项工作分配给具体 warp，使它们可以同时运行。使用 `WG_NUMBER=2` 时，kernel 使用两个 warpgroup（角色表中缩写为 WG）： + +| Actor | Location | Job | +|-------|----------|-----| +| **TMA Producer** | Warpgroup 1, warp 3 | 持续通过 TMA 加载 A 和 B tile | +| **MMA Consumer** | Warpgroup 1, warp 0 | 数据就绪后立即运行 MMA | +| **Writeback** | Warpgroup 0（所有 warp） | 读取 TMEM 结果并写入 GMEM | + +### 4 个 Barrier + +三个并发 actor 需要四个 barrier，这四个 barrier 可以整齐地分成两个方向。前向路径（TMA -> MMA -> Writeback）发出数据*就绪*信号；它传递的信息是“你等待的 tile 已经到了”。反向路径（Writeback -> MMA -> TMA）发出 buffer *release* 信号：“你想要的 slot 又空出来了”。理解命名约定后，名字本身就能读懂：每个名字都是 `source2destination`，所以 `tma2mma` 就是 TMA 向 MMA 发信号的 barrier。 + +| Barrier | Type | Direction | Meaning | +|---------|------|-----------|---------| +| **tma2mma** | `TMABar` | TMA -> MMA | “SMEM data is ready” | +| **mma2tma** | `TCGen05Bar` | MMA -> TMA | “SMEM buffer can be reused” | +| **mma2ld** | `TCGen05Bar` | MMA -> Writeback | “TMEM results are ready” | +| **ld2mma** | `MBarrier` | Writeback -> MMA | “TMEM is free for next tile” | + +为什么每个 barrier 有对应的 *type*？Type 来自 producer 宣布完成的方式。**TMA load** 使用 `TMABar`，也就是带 byte counting 的 mbarrier：TMA hardware 自己会在 transfer byte 落地后 arrive 到 barrier，consumer 因此无需 thread polling 就能知道数据就绪。**TMA store** 不能使用这个方式（store 没有要通知的对象），所以它回到 `cp_async.bulk.commit_group()` + `wait_group(0)`，由 issuing thread 等待自己的写入 drain。**MMA operation** 使用 `TCGen05Bar`，其中 `tcgen05.commit()` 指令会在 MMA 完成时 signal barrier。 + +这里有个小细节会在 Step 8 派上用场。`arrive` 调用会传入 `cta_mask=0`，因为在 single-CTA kernel 中没有其他 CTA 需要 signal。当 Step 8 形成 cluster 时，正是这个参数变成非零值，并成为唤醒协作 CTA 的机制。 + +### PipelineState + +四个 barrier 告诉角色 buffer *何时*就绪；但还需要有东西跟踪 pipeline 循环时每个角色位于*哪个* buffer 上。这正是 `PipelineState` 管理的 bookkeeping。Ring buffer 同时携带两类 bookkeeping：当前在哪个 slot，以及正在等待这个 slot 的 barrier 的哪个 “phase”。在 pipelined loop 中手工跟踪这两者，正是容易产生 off-by-one error 的地方，而这里的 off-by-one 会让整个 kernel deadlock。`PipelineState` 把两者绑在一起，避免你手写： + +```python +tma_ps = PipelineState(PIPE_DEPTH, phase=1) # Producer starts ready (phase=1) +# tma_ps.stage = current stage index +# tma_ps.phase = current phase (0 or 1) +tma_ps.advance() # Advance to next stage +``` + +初始 `phase` 决定某个角色第一次 `wait` 是直接通过还是阻塞，而 pipe 两端正确答案正好相反，这也是容易出错的部分： + +- `phase=1`（producer）-> 第一次 `wait(phase=1)` 看到 barrier 仍处于 phase 0，因为 0 != 1，所以会**立即通过**。这正是我们想要的，因为 buffer 初始为空，producer 应该可以立刻开始填充。 + +- `phase=0`（consumer）-> 第一次 `wait(phase=0)` 看到 barrier 处于 phase 0，因为 0 == 0，所以会**阻塞**。这也是我们想要的，因为还没有数据，consumer 在 producer arrive 前没有东西可读。 + +如果给两端相同 starting phase，就会得到 deadlock，或者更糟的静默 corruption，因此这个选择值得认真对待。 + +### `warpgroup_sync` Barrier ID + +Specialization 会引入一个很容易踩到的 synchronization hazard。一旦每个 warpgroup 走不同 code path，熟悉的 `cta_sync()` 就会 deadlock：它使用硬件 barrier #0，并要求*每个* CTA thread 都 arrive，但在 warpgroup branch 内只有部分线程存在。我们需要的是作用域限定在单个 warpgroup 内的 barrier。GPU 提供 16 个 named barrier（ID 0-15），所以 kernel 使用 `warpgroup_sync(10)`，它只同步一个 warpgroup 内的线程。当多个 warpgroup 各自需要同步时，例如 multi-consumer Step 9，会通过 `warpgroup_sync(wg_id + 10)` 使用不同 ID，避免撞到同一个硬件 barrier。 + +**实现。** + +这里使用 `PIPE_DEPTH=2`，这是仍能让 load 和 compute 重叠的最小深度。更深 pipeline 可以隐藏更多 memory latency，直到 SMEM 预算达到上限；下面的 *When Step 7 misbehaves* 会详细讨论这个取舍。现在所有组件都齐了：角色、四个 barrier、`PipelineState` 和 warpgroup-scoped sync，可以组合出完整 kernel： + +```python +import tvm +from tvm.script import tirx as T +from tvm.script.tirx import tile as Tx +from tvm.tirx.layout import TileLayout, S, TLane, TCol, tid_in_wg +from tvm.tirx.cuda.operator.tile_primitive.tma_utils import tma_shared_layout, SwizzleMode +from tvm.tirx.lang.pipeline import TMABar, TCGen05Bar, MBarrier, PipelineState +from tvm.tirx.lang.tile_scheduler import ClusterPersistentScheduler2D + +SM_COUNT = 148 # Number of SMs on NVIDIA B200 GPU +F16_SIZE = 2 + +def hgemm_v7(M, N, K): + a_type = tvm.DataType("float16") + b_type = tvm.DataType("float16") + d_type = tvm.DataType("float16") + acc_type = tvm.DataType("float32") + + BLK_M, BLK_N, BLK_K = 128, 128, 64 + K_TILES = K // BLK_K + PIPE_DEPTH = 2 + WG_NUMBER = 2 + + A_layout = tma_shared_layout(a_type, SwizzleMode.SWIZZLE_128B_ATOM, (PIPE_DEPTH, BLK_M, BLK_K)) + B_layout = tma_shared_layout(b_type, SwizzleMode.SWIZZLE_128B_ATOM, (PIPE_DEPTH, BLK_N, BLK_K)) + D_layout = tma_shared_layout(d_type, SwizzleMode.SWIZZLE_128B_ATOM, (BLK_M, BLK_N)) + + @T.prim_func + def kernel( + A: T.Buffer((M, K), a_type), + B: T.Buffer((N, K), b_type), + D: T.Buffer((M, N), d_type), + ): + T.device_entry() + bx = T.cta_id([SM_COUNT]) + wg_id = T.warpgroup_id([WG_NUMBER]) + warp_id = T.warp_id_in_wg([4]) + lane_id = T.lane_id([32]) + + # --- Allocation --- + pool = T.SMEMPool() + tmem_addr = pool.alloc((1,), "uint32") + tma2mma = TMABar(pool, PIPE_DEPTH) + mma2tma = TCGen05Bar(pool, PIPE_DEPTH) + mma2ld = TCGen05Bar(pool, 1) + ld2mma = MBarrier(pool, 1) + pool.move_base_to(1024) + Asmem = pool.alloc((PIPE_DEPTH, BLK_M, BLK_K), a_type, layout=A_layout) + Bsmem = pool.alloc((PIPE_DEPTH, BLK_N, BLK_K), b_type, layout=B_layout) + Dsmem = pool.alloc((BLK_M, BLK_N), d_type, layout=D_layout) + + # --- Barrier init --- + tma2mma.init(1) + mma2tma.init(1) + mma2ld.init(1) + ld2mma.init(128) # all 128 Warpgroup 0 threads arrive + pool.commit() + + # --- TMEM alloc + fence --- + if wg_id == 0: + if warp_id == 0: + T.ptx.tcgen05.alloc(T.address_of(tmem_addr), n_cols=512, cta_group=1) + T.ptx.fence.proxy_async("shared::cta") + T.ptx.fence.mbarrier_init() + T.cuda.cta_sync() + + tmem = T.decl_buffer( + (128, 512), acc_type, scope="tmem", allocated_addr=tmem_addr[0], + layout=TileLayout(S[(128, 512) : (1@TLane, 1@TCol)])) + + # --- Tile scheduler --- + tile_scheduler = ClusterPersistentScheduler2D( + "ts", num_m_tiles=M // BLK_M, num_n_tiles=N // BLK_N, + l2_group_size=8, num_clusters=SM_COUNT) + tile_scheduler.init(bx) + m_st = T.meta_var(tile_scheduler.m_idx * BLK_M) + n_st = T.meta_var(tile_scheduler.n_idx * BLK_N) + + # ============================================= + # Warpgroup 1: TMA Producer (warp 3) + MMA Consumer (warp 0) + # ============================================= + if wg_id == 1: + if warp_id == 3: + # === TMA Producer === + tma_ps = PipelineState(PIPE_DEPTH, phase=1) + + @T.inline + def tma_load(k_offset): + Tx.copy_async(Asmem[tma_ps.stage, :, :], + A[m_st:m_st+BLK_M, k_offset:k_offset+BLK_K], + dispatch="tma", cta_group=1, + mbar=tma2mma.ptr_to([tma_ps.stage])) + Tx.copy_async(Bsmem[tma_ps.stage, :, :], + B[n_st:n_st+BLK_N, k_offset:k_offset+BLK_K], + dispatch="tma", cta_group=1, + mbar=tma2mma.ptr_to([tma_ps.stage])) + + if T.filter(lane_id, T.ptx.elect_sync()): + while tile_scheduler.valid(): + for k in range(K_TILES): + mma2tma.wait(tma_ps.stage, tma_ps.phase) + tma_load(k * BLK_K) + tma2mma.arrive(tma_ps.stage, + (BLK_M * BLK_K + BLK_N * BLK_K) * F16_SIZE) + tma_ps.advance() + tile_scheduler.next_tile() + + elif warp_id == 0: + # === MMA Consumer === + mma_ps = PipelineState(PIPE_DEPTH, phase=0) + ld_ps = PipelineState(1, phase=1) + + if T.filter(lane_id, T.ptx.elect_sync()): + while tile_scheduler.valid(): + # Wait for TMEM to be free from previous tile's writeback + ld2mma.wait(ld_ps.stage, ld_ps.phase) + ld_ps.advance() + + for k in range(K_TILES): + tma2mma.wait(mma_ps.stage, mma_ps.phase) + Tx.gemm_async( + tmem[:, :BLK_N], + Asmem[mma_ps.stage, :, :], + Bsmem[mma_ps.stage, :, :], + accum=(k != 0), dispatch="tcgen05", cta_group=1) + mma2tma.arrive(mma_ps.stage, cta_group=1, cta_mask=0) + mma_ps.advance() + + # Signal results ready for writeback + mma2ld.arrive(0, cta_group=1, cta_mask=0) + tile_scheduler.next_tile() + + # ============================================= + # Warpgroup 0: Writeback + # ============================================= + elif wg_id == 0: + wb_ps = PipelineState(1, phase=0) + reg_f16 = T.alloc_local((BLK_N,), d_type) + + while tile_scheduler.valid(): + # Wait for MMA results + mma2ld.wait(wb_ps.stage, wb_ps.phase) + wb_ps.advance() + + # Read TMEM -> registers (warpgroup scope) + reg = T.alloc_local((BLK_N,), acc_type) + reg_wg = reg.view(128, BLK_N, + layout=TileLayout(S[(128, BLK_N) : (1@tid_in_wg, 1)])) + Tx.wg.copy_async(reg_wg[:], tmem[:, :BLK_N]) + T.ptx.tcgen05.wait.ld() + + # Signal TMEM free (all 128 threads arrive) + ld2mma.arrive(0, cta_id=0, pred=True) + + # Cast fp32 -> fp16 + Tx.cast(reg_f16[:], reg[:]) + + # Write to Dsmem + TMA store + Tx.copy(Dsmem[warp_id * 32 + lane_id, :], reg_f16[:]) + T.ptx.fence.proxy_async("shared::cta") + T.cuda.warpgroup_sync(10) + if warp_id == 0: + if lane_id == 0: + Tx.copy_async(D[m_st:m_st+BLK_M, n_st:n_st+BLK_N], + Dsmem[:, :], dispatch="tma") + T.ptx.cp_async.bulk.commit_group() + T.ptx.cp_async.bulk.wait_group(0) + T.cuda.warpgroup_sync(10) + + tile_scheduler.next_tile() + + # --- Cleanup --- + T.cuda.cta_sync() + if warp_id == 0: + T.ptx.tcgen05.relinquish_alloc_permit(cta_group=1) + T.ptx.tcgen05.dealloc(tmem_addr[0], n_cols=512, cta_group=1) + + return kernel +``` + +要运行这些 kernel，复用 Step 1 中展示过的 compile / run / check harness（{ref}`chap_gemm_basics`）：把 `hgemm_v1` 换成 `hgemm_v7`、`hgemm_v8` 或 `hgemm_v9`，并选择类似 `M=N=K=4096` 的问题规模。注意 clustered step 要求 `M` 和 `N` 是 cluster tile 的倍数（Step 8 是 `256×256`，Step 9 是 `512×256`），所以很小的 `128×128` 尺寸不会产生任何 tile。每个新的 Python session 只编译一个 step；切换 step 前重启 kernel，因为这些 kernel 会复用内部名字，而编译器持有 per-session state。各 step 的 timing 汇总在下面的 *End-to-End Result* 中。 + +### Epilogue（Writeback）细节 + +Step 7 可以使用一个很简单的 epilogue。由于只有 `BLK_N=128` 列，writeback warpgroup 可以一次把整个 TMEM tile 读入寄存器，然后发起一次 TMA store。Steps 8 和 9 没有这个便利，这也是它们后面要引入 chunking 的原因；但目前序列是： + +1. 等待 MMA：`mma2ld.wait(phase)`。本教程中的 Steps 8 和 9 会在这里加一个 `fence.after_thread_sync()` 作为保守额外操作；MMA-completion mbarrier 已经覆盖 ordering，大多数 kernel（包括 CUTLASS）会省略它，所以 Step 7 也省略。 +2. 读取 TMEM -> registers（每个线程 128 个 fp32，warpgroup scope，通过 `Tx.copy_async(reg_wg, tmem[:, :BLK_N])` 后接 `T.ptx.tcgen05.wait.ld()`）。 +3. Signal MMA：`ld2mma.arrive(0, cta_id=0, pred=True)`（全部 128 个线程 arrive）；TMEM 现在可以给下一个 tile 复用。两个 `arrive` kwarg 会在 clustered step 中再次出现：`cta_id` 指定 signal *哪个 CTA 的* barrier copy（`0` = 当前 CTA，也就是 local barrier；Step 8 中 cooperative arrive 会改为通过 `cta_mask` 指向 CTA-0），`pred` 是 per-thread predicate，决定这个线程是否真的 arrive（这里是 `True`，所以每个 writeback thread 都计入 arrival total）。 +4. 在寄存器中 cast fp32 -> fp16。 +5. 写 registers -> Dsmem，然后用 `fence.proxy_async("shared::cta") + warpgroup_sync(10)` flush。 +6. 通过 `cp_async.bulk.commit_group() + wait_group(0)` 执行 TMA store Dsmem -> GMEM。 + +Step 8（`BLK_N=256`）和 Step 9（每个 consumer 的 `MMA_N=256`）无法保持这种 one-pass 形式，原因是寄存器压力。每个线程读取 256 个 fp32 值，意味着每个线程的寄存器中同时要保存 256 × 4 = 1024 byte，这可能 spill 到 local memory，并且还会迫使 Dsmem buffer 变大。因此这些 step 会把 writeback 拆成 `EPI_N` 列的 chunk（`EPI_N=64`）：每次 iteration 只保持 `EPI_N` 个 fp32 register live，并发起对应更小的 TMA store，用更多 store 指令换取舒适的寄存器预算。 + +**实现备注。** + +- **Persistent kernel**：`bx = T.cta_id([SM_COUNT])`，每个 SM 一个 CTA，循环处理多个 tile + +- **L2-friendly scheduling**：`ClusterPersistentScheduler2D` 按 cache locality 排列 tile + +- 这种 pattern，即 warp specialization 加 software pipelining，在高性能 GEMM kernel 中很常见，包括 CUTLASS 风格设计。 + +### 当 Step 7 出问题时 + +Step 7 是第一个让 TMA load、`tcgen05` MMA 和 writeback 同时在飞行中的 GEMM kernel。同样的 failure pattern 会在 Steps 8 和 9 中反复出现：barrier count 不匹配、role guard 放错位置、缺少 fence，或者 staging buffer 在 TMA store drain 之前被复用。用于这些情况的调试 checklist 收集在 {ref}`chap_warp_spec_debug` 中。 + +**Pipeline depth tuning。** Step 7 kernel 使用最小的 `PIPE_DEPTH=2`。把它推到 4 或 6 可以让 TMA producer 更早跑在 MMA consumer 前面，隐藏更多 memory latency，但代价是消耗更多 SMEM，而 SMEM 有限。B200 每个 SM 提供 228 KB（见 {ref}`chap_background` 中的 *Numbers to Keep in Mind*）。在 `BLK_M=BLK_N=128, BLK_K=64, fp16` 下，每个 pipeline stage 的 A 和 B 合起来消耗 `(128*64 + 128*64) * 2 = 32 KB`，`Dsmem` writeback staging buffer 额外增加 32 KB。因此 `PIPE_DEPTH=4` 大约是 160 KB，`PIPE_DEPTH=6` 大约是 224 KB，已经贴近预算。想更深，就必须重新考虑 writeback staging 策略。 + +--- + +Warp specialization 让一个 CTA 内的线程开始协作。下一步会把协作范围扩大到 CTA 边界之外，让两个 CTA 为一个更大的 tile 工作。 + + (chap_cta_cluster)= +## Step 8：2-CTA Cluster + +Step 7 让各个 engine 重叠起来，但每个 CTA 仍然独立计算自己的 128×128 tile，并重新加载邻居无法借用的 operand。Step 8 打破这种隔离。两个 CTA 组成一个 cluster，并获得访问彼此 shared memory 的能力，因此单个 cooperative `tcgen05` MMA 可以产生跨越两个 CTA 的 256×256 tile，而一次 B load 现在可以喂给两倍的 MMA work。和之前一样，M=N=K=4096。 + +> **本 step 改变的内容：Scope + Layout + Dispatch** +> - Scope：协作范围现在跨越 cluster 中的两个 CTA，而不是一个 CTA。 +> - Layout：operand tile 切分到两个 CTA 的 SMEM 中；CTA 0 拥有共享 completion barrier（`remote_view`）。 +> - Dispatch：MMA 获得 `cta_group` / `cta_mask`，使 `tcgen05` 作为 2-CTA cooperative op 运行。 + +**主题。** + +- CTA cluster：多个 CTA 协作处理更大的 tile + +- 通过 `map_shared_rank` 做 cross-CTA SMEM access + +- 用 `cta_group=2` 在 256x256 cluster tile 上执行 cooperative MMA + +- 用 `cta_mask` 做 cross-CTA barrier signaling + + +### Cluster Tile Shape + +整个优化建立在一个硬件能力上：使用 `cta_group=2` 时，MMA 可以读取*两个* CTA staged 的 operand tile，而不仅仅是自己所在 CTA 的 tile。每个 CTA 加载 stored B 的一个 128-row slice，经过 transpose 后，它会变成 128 个逻辑 output column；cooperative MMA 会把两个 slice 重新拼成一个 operand。下图展示两个 CTA 的 A 和 B slice 如何组合成一个 256×256 cluster tile： + +```{raw} html +
+ +
+``` +*交互图：每个 CTA 拥有一个 A row slice 和一个 stored-B row slice，然后通过 cluster（DSMEM）读取另一个 CTA 的 stored-B slice。经过 `B.T` 后，两个 stored-B slice 覆盖完整 output-column span，因此这一对 CTA 产生一个 256×256 output tile。* + +**为什么 A 和 B 要在 cluster 中切分**：要理解 256×256 tile 如何 partition，回想本教程把 GEMM 写成 `D = A @ B.T`，其中 stored B 的 shape 是 `N x K`。两个 CTA 组成 cluster 后，切分方式很自然： + +- **A 竖直切分**：CTA-0 持有 A0（rows 0-127），CTA-1 持有 A1（rows 128-255）。叠起来是 `[A0; A1]`（256 行）。 +- **Stored B 按 row 切分**：CTA-0 加载 B rows 0-127，CTA-1 加载 B rows 128-255。因为数学使用 `B.T`，这两个 stored row slice 会变成逻辑右操作数的两个 128-column slice。 +- 使用 `cta_group=2` 时，MMA 硬件通过 cross-CTA shared memory access 从**两个** CTA 的 SMEM 中读取 B，因此它看到完整的逻辑 output-column span。 +- 结果：两个 CTA 协作处理一个 256x256 output tile。每个 CTA 写这个 tile 的一个 128x256 row stripe。 + +这里值得停一下，看看为什么这是真正收益而不只是 work 重新排列。每个 CTA 仍然只加载 128×K 的 A 和 128×K 的 B，因此整个 cluster staged 的 operand 大约是单个 CTA 的 2×，但它产生的是 256×256 tile，拥有 128×128 tile 大约 4× 的 output FLOP。MMA 因此对每个 staged-operand byte 做大约两倍工作，因为每个 CTA 的 B slice 会通过 cooperative MMA 与另一个 CTA 的 A slice 复用。换句话说，算术强度大约翻倍，而这正是仍偏 memory-leaning 的 kernel 需要的杠杆：End-to-End 表中的约 2.2× speedup 来自把同样 byte 喂给更多数学计算。 + +### Tile Address Calculation + +现在 cluster 是 work 的单位，tile scheduler 也必须按 cluster tile 计数。它返回的每个 `(m_idx, n_idx)` 都命名一个完整 256×256 region，而 cluster 内的两个 CTA 会切分这个 region。把 cluster coordinate 转换成每个 CTA 实际加载的 per-CTA slice，如下： + +```python +m_st = (m_idx * CTA_GROUP + cbx) * BLK_M +n_st = (n_idx * CTA_GROUP + cbx) * BLK_N +``` + +两个 CTA 处理的是*同一个* 256×256 cluster tile，单个 coordinate `cbx`（CTA 在 cluster 内的位置，0 或 1）选择这个 CTA 沿两个轴贡献的部分。`m_st` 选择该 CTA 拥有的 output row stripe，`n_st` 选择它喂给 cooperative MMA 的 stored-B slice，writeback 稍后会写出 256-column output span 的两个 128-column half。还要注意，`num_m_tiles = M // 256` 和 `num_n_tiles = N // 256` 计数的是 cluster tile，不是单个 CTA tile。 + +乍一看，`cbx` 同时出现在 `m_st` 和 `n_st` 中，像是 row offset 泄漏到了 column 里，但两处都是正确的，值得拆开理解。在 writeback 路径上，`cbx` 只属于 M 轴：每个 CTA 拥有不同的 128-row stripe（`m_st = (m_idx * CTA_GROUP + cbx) * BLK_M`，所以 CTA-0 写 rows `m_idx*256 .. +128`，CTA-1 写接下来的 128），但两个 CTA 都写 cluster tile 的*完整* 256 个 output column。因此 store 的 column 来自 cluster 的 `n_idx`（`n_st_epi = n_idx * 256 + no * 128`，完全没有 `cbx`），而不是来自 per-CTA 的 `n_st`。`n_st` 带有 `cbx` 的原因是每个 CTA 要把不同 stored-B row slice 加载进 MMA：在那里，`cbx` 是一个 *load* offset，不是该 CTA 的 output-column offset。 + +### 与 Step 7 的代码差异 + +相对 Step 7 的 diff 有六处，每一处都编码了刚才描述的 cluster contract 的一个部分： + +```python +# 1. Cluster launch +cbx, cby = T.cta_id_in_cluster([CTA_GROUP, 1]) # cbx = CTA index within cluster (0 or 1) + +# 2. Cooperative MMA (was cta_group=1) +Tx.gemm_async(..., cta_group=2) + +# 3. Cross-CTA shared memory access +B_remote = T.ptx.map_shared_rank(Bsmem, cta_id=1) + +# 4. Cross-CTA barrier +tma2mma_cta0 = T.decl_buffer( + [CTA_GROUP], "uint64", + data=T.ptx.map_shared_rank(tma2mma.ptr_to([0]), 0), + scope="shared" +) + +# 5. mma2tma / mma2ld arrives go from cta_mask=0 (single CTA, Step 7) +# to cta_mask=3 (signal both CTAs in the cluster) +mma2tma.arrive(mma_ps.stage, cta_group=CTA_GROUP, cta_mask=3) +mma2ld.arrive(0, cta_group=CTA_GROUP, cta_mask=3) + +# 6. Cluster sync replaces cta_sync at the end +T.cuda.cluster_sync() +``` + + +### Cluster-Scope 变化 + +这六处修改都来自同一个转变：协作 scope 现在是 cluster，而不是单个 CTA。下面逐项说明这种扩大在实践中意味着什么：每个 CTA 如何找到自己的位置，cluster 通过谁的 barrier 协调，以及哪个 CTA 实际发起 cooperative MMA。 + +- **Cluster CTA ID**：`cbx` 告诉每个 CTA 自己在 cluster 中的位置（0 或 1）。CTA-0 处理 A rows 0-127，CTA-1 处理 rows 128-255。 + +- **Remote barrier view**：在一个 cluster 中，每个 CTA 都有自己的 SMEM 和自己的 barrier，这带来一个明显问题：如果 CTA-1 需要等待 CTA-0 产生的东西，它实际触碰谁的 barrier？答案是指定 CTA-0 的 barrier 作为唯一 coordination point，并允许 cluster 中任意 CTA 访问它。`map_shared_rank(tma2mma.ptr_to([0]), 0)` 返回一个指向 CTA-0 barrier 的 cluster-wide pointer；TIRx wrapper `tma2mma.remote_view(0)` 提供这个能力，从此之后每个 arrive 和 wait 都指向 CTA-0 的 copy。 + +- **MMA dispatch 只来自 CTA-0**：容易把 `cta_group=2` 理解成两个 engine 并行 firing，但实际不是这样。CTA-0 发起恰好一次 `tcgen05.mma`，硬件随后驱动一个跨两个 CTA 的*单个 cooperative* MMA，从两个 SM 的 SMEM 读取 operand，并把 accumulator 写到两个 SM 的 TMEM 中。CTA-1 不发起 MMA。（每个 SM 只有一个 `tcgen05` engine，因此 `cta_group=2` 是一个 cross-SM MMA，而不是两个 engine side by side 运行。）这就是为什么代码用 `if cbx == 0:` guard MMA。 + +- **Multicast arrive**：`tcgen05.commit(..., cta_group=2, cta_mask=3)` 只由 CTA-0 发起，但会 signal 两个 CTA 的 barrier。`cta_mask=3`（binary `11`）表示目标是 CTA-0 和 CTA-1。 + +- **ld2mma init count**：`init(128 * CTA_GROUP)`，两个 CTA 的 writeback warpgroup（每个 128 线程）都会 arrive。 + + +**实现。** + +```python +def hgemm_v8(M, N, K): + a_type = tvm.DataType("float16") + b_type = tvm.DataType("float16") + d_type = tvm.DataType("float16") + acc_type = tvm.DataType("float32") + + CTA_GROUP = 2 + BLK_M, BLK_N, BLK_K = 128, 128, 64 + MMA_M, MMA_N = 256, 256 + K_TILES = K // BLK_K + PIPE_DEPTH = 4 + WG_NUMBER = 2 + F16_SIZE = 2 # fp16 + + A_layout = tma_shared_layout(a_type, SwizzleMode.SWIZZLE_128B_ATOM, (PIPE_DEPTH, BLK_M, BLK_K)) + B_layout = tma_shared_layout(b_type, SwizzleMode.SWIZZLE_128B_ATOM, (PIPE_DEPTH, BLK_N, BLK_K)) + D_layout = tma_shared_layout(d_type, SwizzleMode.SWIZZLE_128B_ATOM, (BLK_M, 128)) + + @T.prim_func + def kernel( + A: T.Buffer((M, K), a_type), + B: T.Buffer((N, K), b_type), + D: T.Buffer((M, N), d_type), + ): + T.device_entry() + bx = T.cta_id([SM_COUNT]) + cbx, cby = T.cta_id_in_cluster([CTA_GROUP, 1]) + wg_id = T.warpgroup_id([WG_NUMBER]) + warp_id = T.warp_id_in_wg([4]) + lane_id = T.lane_id([32]) + + # --- Allocation --- + pool = T.SMEMPool() + tmem_addr = pool.alloc((1,), "uint32") + tma2mma = TMABar(pool, PIPE_DEPTH) + mma2tma = TCGen05Bar(pool, PIPE_DEPTH) + mma2ld = TCGen05Bar(pool, 1) + ld2mma = MBarrier(pool, 1) + pool.move_base_to(1024) + Asmem = pool.alloc((PIPE_DEPTH, BLK_M, BLK_K), a_type, layout=A_layout) + Bsmem = pool.alloc((PIPE_DEPTH, BLK_N, BLK_K), b_type, layout=B_layout) + Dsmem = pool.alloc((BLK_M, 128), d_type, layout=D_layout) + + # --- Barrier init --- + tma2mma.init(1) + mma2tma.init(1) + mma2ld.init(1) + ld2mma.init(128 * CTA_GROUP) # both CTAs' writeback threads + pool.commit() + + # --- TMEM alloc (cooperative) --- + if wg_id == 0: + if warp_id == 0: + T.ptx.tcgen05.alloc(T.address_of(tmem_addr), n_cols=512, cta_group=CTA_GROUP) + T.ptx.fence.proxy_async("shared::cta") + T.ptx.fence.mbarrier_init() + T.cuda.cta_sync() + + tmem = T.decl_buffer( + (128, 512), acc_type, scope="tmem", allocated_addr=tmem_addr[0], + layout=TileLayout(S[(128, 512) : (1@TLane, 1@TCol)])) + + # --- Tile scheduler (cluster tiles) --- + tile_scheduler = ClusterPersistentScheduler2D( + "ts", num_m_tiles=M // 256, num_n_tiles=N // 256, + l2_group_size=8, num_clusters=SM_COUNT // CTA_GROUP) + tile_scheduler.init(bx // CTA_GROUP) + m_idx = T.meta_var(tile_scheduler.m_idx) + n_idx = T.meta_var(tile_scheduler.n_idx) + m_st = T.meta_var((m_idx * CTA_GROUP + cbx) * BLK_M) + n_st = T.meta_var((n_idx * CTA_GROUP + cbx) * BLK_N) + + # --- Cross-CTA barrier view --- + tma2mma_cta0 = tma2mma.remote_view(0) + + # ============================================= + # Warpgroup 1: TMA Producer (warp 3) + MMA Consumer (warp 0) + # ============================================= + if wg_id == 1: + if warp_id == 3: + tma_ps = PipelineState(PIPE_DEPTH, phase=1) + + @T.inline + def tma_load(k_offset): + Tx.copy_async(Asmem[tma_ps.stage, :, :], + A[m_st:m_st+BLK_M, k_offset:k_offset+BLK_K], + dispatch="tma", cta_group=CTA_GROUP, + mbar=tma2mma_cta0.ptr_to([tma_ps.stage])) + Tx.copy_async(Bsmem[tma_ps.stage, :, :], + B[n_st:n_st+BLK_N, k_offset:k_offset+BLK_K], + dispatch="tma", cta_group=CTA_GROUP, + mbar=tma2mma_cta0.ptr_to([tma_ps.stage])) + + if T.filter(lane_id, T.ptx.elect_sync()): + while tile_scheduler.valid(): + for k in range(K_TILES): + mma2tma.wait(tma_ps.stage, tma_ps.phase) + tma_load(k * BLK_K) + if cbx == 0: + tma2mma_cta0.arrive(tma_ps.stage, + CTA_GROUP * (BLK_M * BLK_K + BLK_N * BLK_K) * F16_SIZE) + tma_ps.advance() + tile_scheduler.next_tile() + + elif warp_id == 0: + mma_ps = PipelineState(PIPE_DEPTH, phase=0) + ld_ps = PipelineState(1, phase=1) + + if cbx == 0: + if T.filter(lane_id, T.ptx.elect_sync()): + while tile_scheduler.valid(): + ld2mma.wait(ld_ps.stage, ld_ps.phase) + ld_ps.advance() + + for k in range(K_TILES): + tma2mma.wait(mma_ps.stage, mma_ps.phase) + Tx.gemm_async( + tmem[:, :MMA_N], + Asmem[mma_ps.stage, :, :], + Bsmem[mma_ps.stage, :, :], + accum=(k != 0), dispatch="tcgen05", cta_group=CTA_GROUP) + mma2tma.arrive(mma_ps.stage, cta_group=CTA_GROUP, cta_mask=3) + mma_ps.advance() + + mma2ld.arrive(0, cta_group=CTA_GROUP, cta_mask=3) + tile_scheduler.next_tile() + + # ============================================= + # Warpgroup 0: Writeback (256 columns in 2 x 128-column chunks) + # ============================================= + elif wg_id == 0: + wb_ps = PipelineState(1, phase=0) + reg_f16 = T.alloc_local((128,), d_type) + + while tile_scheduler.valid(): + mma2ld.wait(wb_ps.stage, wb_ps.phase) + wb_ps.advance() + T.ptx.tcgen05.fence.after_thread_sync() + + for no in T.unroll(2): # 2 chunks of 128 columns = 256 total + reg = T.alloc_local((128,), acc_type) + reg_wg = reg.view(128, 128, + layout=TileLayout(S[(128, 128) : (1@tid_in_wg, 1)])) + Tx.wg.copy_async(reg_wg[:], tmem[:, no * 128:(no + 1) * 128]) + T.ptx.tcgen05.wait.ld() + Tx.cast(reg_f16[:], reg[:]) + Tx.copy(Dsmem[warp_id * 32 + lane_id, :], reg_f16[:]) + T.ptx.fence.proxy_async("shared::cta") + T.cuda.warpgroup_sync(10) + if warp_id == 0: + if lane_id == 0: + n_st_epi = T.meta_var(n_idx * 256 + no * 128) + Tx.copy_async(D[m_st:m_st+BLK_M, n_st_epi:n_st_epi+128], + Dsmem[:, :], dispatch="tma") + T.ptx.cp_async.bulk.commit_group() + T.ptx.cp_async.bulk.wait_group(0) + T.cuda.warpgroup_sync(10) + + ld2mma.arrive(0, cta_id=0, pred=True) + tile_scheduler.next_tile() + + # --- Cleanup --- + T.cuda.cluster_sync() + if warp_id == 0: + T.ptx.tcgen05.relinquish_alloc_permit(cta_group=CTA_GROUP) + T.ptx.tcgen05.dealloc(tmem_addr[0], n_cols=512, cta_group=CTA_GROUP) + + return kernel +``` + +**2 个 CTA 带来的变化。** + +- `CTA_GROUP = 2`，`MMA_N = BLK_N * CTA_GROUP = 256` + +- `ld2mma.init(128 * CTA_GROUP)`，两个 CTA 的 writeback WG 都会 arrive + +- TMA arrive byte count 包含两个 CTA：`CTA_GROUP * (BLK_M * BLK_K + BLK_N * BLK_K) * F16_SIZE` + +- `tcgen05.alloc` 和 `tcgen05.dealloc` 必须使用 `cta_group=2` + +- Writeback 把 256 个 output column 拆成两个 128-column chunk；一次性读取全部 256 个 TMEM column 会超过寄存器容量。Step 9 会进一步把 chunk 缩小到 `EPI_N=64` + +- 末尾用 `cluster_sync()` 替代 `cta_sync()`，确保所有 CTA 在 TMEM dealloc 前都完成 + +这些额外算术强度会直接体现在 wall clock 上：Step 8 在 4096³ 上达到 **0.104 ms**，相对同规模下 70 ms 的 Step-1 算法约 676×（见 End-to-End 表）。Kernel 现在开始偏向 compute-bound，这正好为 Step 9 做铺垫：我们会加入第二个 MMA consumer，让更多 Tensor Core work 保持在飞行中。 + +如果 Step 8 反而比 Step 7 *更慢*，问题几乎总是某个新的 cluster contract 写错。优先检查三件事：TMA arrive byte count 是否是 `CTA_GROUP * (BLK_M*BLK_K + BLK_N*BLK_K) * F16_SIZE`；scheduler 维度是否是 256×256 cluster tile 对应的 `num_m_tiles=M//256, num_n_tiles=N//256`；writeback 是否发起两次 TMA store，每个 128-column chunk 一次，并且每次都在 Dsmem 复用前 drain。 + +--- + +Cluster 提高了 CTA *之间*的复用。最后一步转向内部，在每个 CTA 内通过给 producer 增加第二个 MMA consumer 来提高 compute density。 + + (chap_multi_consumer)= -# 用 Warp Specialization 和 Cluster 扩展 GEMM +## Step 9：Multi-Consumer Warp Specialization + +到 Step 8 时，MMA 已经真正忙起来了，但一个 consumer warp 消费 staged B tile 的速度有限，而这个 B tile 在 SMEM 中一直可用，任何愿意读取它的人都能读。最后一个优化就利用这一点：加入第二个 MMA consumer，让它把*不同* A block 与*同一个* B tile 相乘。每个 CTA 的 compute density 翻倍，cluster output 从 256×256 增长到 512×256。和之前一样，M=N=K=4096。 + +> **本 step 改变的内容：Scope + Layout** +> - Scope：一个 MMA consumer 变成两个，由 `warp_id` 选择。 +> - Layout：一个 staged B tile 被两个 consumer 复用；A 增加 consumer axis。 +> - Dispatch：不变。 + +**主题。** + +- 多个 MMA warp（consumer）提高吞吐 + +- 多个 writeback warpgroup 使用独立 barrier slot + +- 本教程中最优化 GEMM 变体使用的结构 + + +### Multi-Consumer Structure + +加入第二个 consumer 后，kernel 现在需要安排更多不同角色：两个 MMA warp，而不是一个；并配套第二个 writeback warpgroup 来 drain 额外 accumulator。使用 `NUM_CONSUMER=2` 和 `WG_NUMBER=3` 时，kernel 跨三个 warpgroup（角色表中缩写为 WG）： + +| Warpgroup | Warp | Role | +|-----------|------|------| +| **WG 2** | warp 0 | MMA consumer 0：`Asmem[..., 0] x B` -> TMEM cols `[0:256]` | +| **WG 2** | warp 1 | MMA consumer 1：`Asmem[..., 1] x B` -> TMEM cols `[256:512]` | +| **WG 2** | warp 3 | TMA producer：每个 stage 加载 2x A block + 1x B block | +| **WG 0** | all | consumer 0 的 writeback：读取 TMEM `[0:256]` | +| **WG 1** | all | consumer 1 的 writeback：读取 TMEM `[256:512]` | + +整个安排依赖一个不对称性。每个 consumer 都把自己的 A block 与*同一个* staged B tile 相乘，因此单次 B load 现在喂给 2× MMA work，B 的每有效 FLOP load cost 实际减半。我们共享 B 而不是 A 的原因是，两个 consumer 覆盖不同 M-row stripe：它们的 A block 真实不同，而 B 对二者相同。练习 3 会要求你说服自己，这是唯一可行的共享。 + +### 相对 Step 8 的变化 + +具体来说，支持第二个 consumer 会触碰 kernel 中少数几处，而每个变化都可以追溯到同一个事实：现在每个 stage 有两个 A block 和两个 TMEM range 需要喂给并 drain，而 B 保持共享。下面的修改会 stage 额外 A block，给每个 consumer 自己的 barrier slot，并为更高的 512×256 cluster tile 调整 tile addressing。 + +- `Asmem = pool.alloc((PIPE_DEPTH, NUM_CONSUMER, BLK_M, BLK_K), ...)`：每个 stage 两个 A block，每个 consumer 一个 + +- TMA 同时加载 `Asmem[stage, 0]` 和 `Asmem[stage, 1]`，TMA arrive byte 现在是 `CTA_GROUP * (NUM_CONSUMER * BLK_M * BLK_K + BLK_N * BLK_K) * F16_SIZE`（额外 A block） + +- MMA warp `warp_id` 选择哪个 A block 和 TMEM range + +- `mma2tma.init(NUM_CONSUMER)`：每个 stage 两个 consumer 都要向 TMA signal + +- `mma2ld` 和 `ld2mma` 的 `depth=NUM_CONSUMER`：每个 consumer 使用自己的 barrier slot（MMA 侧用 `warp_id`，writeback 侧用 `wg_id`） + +- Tile address：`m_st = (m_idx * NUM_CONSUMER * CTA_GROUP + cbx) * BLK_M`。M 方向有额外 `NUM_CONSUMER` 因子，因为每个 cluster tile 现在沿 M 跨越 `NUM_CONSUMER` 个 consumer。Tile scheduler 使用 `num_m_tiles = M // 256 // NUM_CONSUMER`（cluster tile 是 512x256） + +- Writeback 使用 chunked `EPI_N`，让每次 iteration 中 live 的 TMEM-readback register 更少 + + +**实现。** + +```python +def hgemm_v9(M, N, K): + a_type = tvm.DataType("float16") + b_type = tvm.DataType("float16") + d_type = tvm.DataType("float16") + acc_type = tvm.DataType("float32") + + CTA_GROUP = 2 + NUM_CONSUMER = 2 + BLK_M, BLK_N, BLK_K = 128, 128, 64 + MMA_N = BLK_N * CTA_GROUP # 256 + K_TILES = K // BLK_K + PIPE_DEPTH = 4 + EPI_N = 64 + WG_NUMBER = 3 + F16_SIZE = 2 # fp16 + + A_layout = tma_shared_layout(a_type, SwizzleMode.SWIZZLE_128B_ATOM, + (PIPE_DEPTH, NUM_CONSUMER, BLK_M, BLK_K)) + B_layout = tma_shared_layout(b_type, SwizzleMode.SWIZZLE_128B_ATOM, + (PIPE_DEPTH, BLK_N, BLK_K)) + D_layout = tma_shared_layout(d_type, SwizzleMode.SWIZZLE_128B_ATOM, + (NUM_CONSUMER, BLK_M, EPI_N)) + + @T.prim_func + def kernel( + A: T.Buffer((M, K), a_type), + B: T.Buffer((N, K), b_type), + D: T.Buffer((M, N), d_type), + ): + T.device_entry() + bx = T.cta_id([SM_COUNT]) + cbx, cby = T.cta_id_in_cluster([CTA_GROUP, 1]) + wg_id = T.warpgroup_id([WG_NUMBER]) + warp_id = T.warp_id_in_wg([4]) + lane_id = T.lane_id([32]) + + # --- Allocation --- + pool = T.SMEMPool() + tmem_addr = pool.alloc((1,), "uint32") + tma2mma = TMABar(pool, PIPE_DEPTH) + mma2tma = TCGen05Bar(pool, PIPE_DEPTH) + mma2ld = TCGen05Bar(pool, NUM_CONSUMER) # depth=2, one slot per consumer + ld2mma = MBarrier(pool, NUM_CONSUMER) # depth=2, one slot per consumer + pool.move_base_to(1024) + Asmem = pool.alloc((PIPE_DEPTH, NUM_CONSUMER, BLK_M, BLK_K), a_type, layout=A_layout) + Bsmem = pool.alloc((PIPE_DEPTH, BLK_N, BLK_K), b_type, layout=B_layout) + Dsmem = pool.alloc((NUM_CONSUMER, BLK_M, EPI_N), d_type, layout=D_layout) + + # --- Barrier init --- + tma2mma.init(1) + mma2tma.init(NUM_CONSUMER) # each stage expects 2 arrivals + mma2ld.init(1) # each slot gets 1 arrival + ld2mma.init(128 * CTA_GROUP) # both CTAs' writeback threads + pool.commit() + + # --- TMEM alloc (cooperative) --- + if wg_id == 0: + if warp_id == 0: + T.ptx.tcgen05.alloc(T.address_of(tmem_addr), n_cols=512, cta_group=CTA_GROUP) + T.ptx.fence.proxy_async("shared::cta") + T.ptx.fence.mbarrier_init() + T.cuda.cta_sync() + + tmem = T.decl_buffer( + (128, 512), acc_type, scope="tmem", allocated_addr=tmem_addr[0], + layout=TileLayout(S[(128, 512) : (1@TLane, 1@TCol)])) + + # --- Tile scheduler (512x256 cluster tiles) --- + tile_scheduler = ClusterPersistentScheduler2D( + "ts", num_m_tiles=M // 256 // NUM_CONSUMER, num_n_tiles=N // 256, + l2_group_size=8, num_clusters=SM_COUNT // CTA_GROUP) + tile_scheduler.init(bx // CTA_GROUP) + m_idx = T.meta_var(tile_scheduler.m_idx) + n_idx = T.meta_var(tile_scheduler.n_idx) + m_st = T.meta_var((m_idx * NUM_CONSUMER * CTA_GROUP + cbx) * BLK_M) + n_st = T.meta_var((n_idx * CTA_GROUP + cbx) * BLK_N) + + tma2mma_cta0 = tma2mma.remote_view(0) + + # ============================================= + # Warpgroup 2: TMA Producer (warp 3) + 2 MMA Consumers (warp 0, 1) + # ============================================= + if wg_id == 2: + if warp_id == 3: + # === TMA Producer: loads 2 A blocks + 1 B block per stage === + tma_ps = PipelineState(PIPE_DEPTH, phase=1) + + @T.inline + def tma_load(k_offset): + m_st_c1 = T.meta_var(m_st + CTA_GROUP * BLK_M) + Tx.copy_async(Asmem[tma_ps.stage, 0, :, :], + A[m_st:m_st+BLK_M, k_offset:k_offset+BLK_K], + dispatch="tma", cta_group=CTA_GROUP, + mbar=tma2mma_cta0.ptr_to([tma_ps.stage])) + Tx.copy_async(Asmem[tma_ps.stage, 1, :, :], + A[m_st_c1:m_st_c1+BLK_M, k_offset:k_offset+BLK_K], + dispatch="tma", cta_group=CTA_GROUP, + mbar=tma2mma_cta0.ptr_to([tma_ps.stage])) + Tx.copy_async(Bsmem[tma_ps.stage, :, :], + B[n_st:n_st+BLK_N, k_offset:k_offset+BLK_K], + dispatch="tma", cta_group=CTA_GROUP, + mbar=tma2mma_cta0.ptr_to([tma_ps.stage])) + + if T.filter(lane_id, T.ptx.elect_sync()): + while tile_scheduler.valid(): + for k in range(K_TILES): + mma2tma.wait(tma_ps.stage, tma_ps.phase) + tma_load(k * BLK_K) + if cbx == 0: + tma2mma_cta0.arrive(tma_ps.stage, + CTA_GROUP * (NUM_CONSUMER * BLK_M * BLK_K + BLK_N * BLK_K) * F16_SIZE) + tma_ps.advance() + tile_scheduler.next_tile() + + elif warp_id < NUM_CONSUMER: + # === MMA Consumer: warp_id selects A block and TMEM range === + mma_ps = PipelineState(PIPE_DEPTH, phase=0) + ld_ps = PipelineState(1, phase=1) + + if cbx == 0: + if T.filter(lane_id, T.ptx.elect_sync()): + while tile_scheduler.valid(): + ld2mma.wait(warp_id, ld_ps.phase) + ld_ps.advance() + + for k in range(K_TILES): + tma2mma.wait(mma_ps.stage, mma_ps.phase) + Tx.gemm_async( + tmem[:, warp_id * MMA_N:warp_id * MMA_N + MMA_N], + Asmem[mma_ps.stage, warp_id, :, :], + Bsmem[mma_ps.stage, :, :], + accum=(k != 0), dispatch="tcgen05", cta_group=CTA_GROUP) + mma2tma.arrive(mma_ps.stage, cta_group=CTA_GROUP, cta_mask=3) + mma_ps.advance() + + mma2ld.arrive(warp_id, cta_group=CTA_GROUP, cta_mask=3) + tile_scheduler.next_tile() + + # ============================================= + # Warpgroup 0/1: Writeback (each reads its consumer's TMEM range) + # ============================================= + elif wg_id < NUM_CONSUMER: + wb_ps = PipelineState(1, phase=0) + reg_f16 = T.alloc_local((EPI_N,), d_type) + + while tile_scheduler.valid(): + mma2ld.wait(wg_id, wb_ps.phase) # wait for THIS consumer + wb_ps.advance() + T.ptx.tcgen05.fence.after_thread_sync() + + # Read TMEM in EPI_N=64 column chunks (4 iterations for 256 cols) + for i in T.unroll(MMA_N // EPI_N): + reg = T.alloc_local((EPI_N,), acc_type) + reg_wg = reg.view(128, EPI_N, + layout=TileLayout(S[(128, EPI_N) : (1@tid_in_wg, 1)])) + col_st = T.meta_var(wg_id * MMA_N + i * EPI_N) + col_end = T.meta_var(wg_id * MMA_N + i * EPI_N + EPI_N) + Tx.wg.copy_async(reg_wg[:], tmem[:, col_st:col_end]) + T.ptx.tcgen05.wait.ld() + Tx.cast(reg_f16[:], reg[:]) + Tx.copy(Dsmem[wg_id, warp_id * 32 + lane_id, :], reg_f16[:]) + T.ptx.fence.proxy_async("shared::cta") + T.cuda.warpgroup_sync(wg_id + 10) + if warp_id == 0: + if lane_id == 0: + m_st_epi = T.meta_var( + (m_idx * NUM_CONSUMER * CTA_GROUP + wg_id * CTA_GROUP + cbx) * BLK_M) + n_st_epi = T.meta_var(n_idx * MMA_N + i * EPI_N) + Tx.copy_async( + D[m_st_epi:m_st_epi+BLK_M, n_st_epi:n_st_epi+EPI_N], + Dsmem[wg_id, :, :], dispatch="tma") + T.ptx.cp_async.bulk.commit_group() + T.ptx.cp_async.bulk.wait_group(0) + T.cuda.warpgroup_sync(wg_id + 10) + + ld2mma.arrive(wg_id, cta_id=0, pred=True) + tile_scheduler.next_tile() + + # --- Cleanup --- + T.cuda.cluster_sync() + if warp_id == 0: + T.ptx.tcgen05.relinquish_alloc_permit(cta_group=CTA_GROUP) + T.ptx.tcgen05.dealloc(tmem_addr[0], n_cols=512, cta_group=CTA_GROUP) + + return kernel +``` + +**实现备注。** + +- 在 Step 9 设计中，`mma2ld` 和 `ld2mma` 都是一个 `depth=NUM_CONSUMER` 的共享对象，而不是每个 consumer 一个对象。Slot 0 连接 MMA warp 0 和 Warpgroup 0，slot 1 连接 MMA warp 1 和 Warpgroup 1；MMA 侧用 `warp_id` 索引，writeback 侧用 `wg_id` 索引。 + +## End-to-End Result + +下表报告了从 naive baseline 到 warp-specialized cluster kernel 的实测里程碑，并列出 cuBLAS reference。参考数字来自 NVIDIA B200，M=N=K=4096，fp16，locked clocks，1000-iteration timed benchmark： + +| Step | Technique | Time | Speedup | +|------|-----------|------|---------| +| 1 | Sync load + MMA | 70 ms | 1× | +| 2 | K-loop accumulation | --- | Handle K larger than one tile | +| 3 | Spatial tiling | 53.6 ms | ~1.3× | +| 4 | TMA async load | 0.49 ms | ~142× | +| 5 | Software pipeline | --- | Overlap load + compute | +| 6 | Persistent kernel | --- | L2 cache locality | +| 7 | Warp specialization | 0.23 ms | ~309× | +| 8 | 2-CTA cluster | 0.104 ms | ~676× | +| 9 | Multi-consumer | 0.094 ms | ~744× | +| --- | cuBLAS (reference) | 0.094 ms | ~744× | + +这张表中的所有时间，包括 70 ms 的 Step 1 baseline，都在同一个 M=N=K=4096 规模下测得，因此 speedup chain 可以端到端比较。这里需要精确说明 70 ms 到底是什么，因为很容易误读。它*不是* {ref}`chap_gemm_basics` 中的 single-tile Step-1 kernel 在 4096³ 上运行；那个 kernel 只会计算一个 128×128 tile，并且只在小规模上运行。70 ms 指的是一个 naive full-size baseline，它采用同样的 sequential、single-tile 思路，并扩展到完整 4096³ 问题。Steps 1-3 在 {ref}`chap_gemm_basics` 中以小规模（128×128 和 256³）介绍，是为了让最初的 walkthrough 简单；这里 Step 1 和 Step 3 行是它们的 full-size benchmark counterpart。剩下的 dash（Steps 2、5、6）表示这些 step 用于展示结构，但没有单独计时。 + +请把这些数字看作受控条件下的一次 B200 reference run，而不是 leaderboard。每个 step 中嵌入的 `{.python .input}` benchmark cell 是 smoke benchmark：适合观察趋势，不适合声明峰值性能。 + +几乎所有收益都来自四项技术： + +1. **TMA Async Data Movement**：硬件 copy engine 替代 software copy（Step 1 -> Step 4 约 142×）。要正确理解这个 142×：它反映的是从单个 128×128-tile kernel（grid 1×1）一路变成带 K-loop、spatial tiling、多 CTA 和 TMA 的 full tiled-and-parallel kernel；这不是 TMA 单独贡献。要隔离 TMA，需要比较两个只在 copy mechanism 上不同的 full-size kernel。 +2. **Software Pipelining + Warp Specialization**：通过给 load 和 compute 各自专门角色来重叠二者（Step 4 -> Step 7 约 2.2×）。 +3. **CTA Clusters**：2-SM cooperative MMA 提高 CTA 间的 B-tile 复用（本 benchmark 中 Step 7 -> Step 8 约 2.2×）。 +4. **Multi-Consumer**：两个 MMA warp 提高 compute density（Step 8 -> Step 9 约 10%）。 + +把这些 measured milestone 画出来，同样四项贡献会描出从同步 tiled kernel 走向 cuBLAS reference 的下降曲线。下图展示选中的 measured point： + +![GEMM Optimization Journey](../../img/gemm_perf.png) + +注意，越往后 gain 越小，这背后有结构性原因，而不是优化力度下降。早期 step 攻击的是*内存*瓶颈（TMA 替代 software copy，cluster 提高 arithmetic intensity），而 70 ms 中大部分时间确实花在这里，所以这些 step 收益最大。到 Step 8，kernel 已经在 cuBLAS 的约 10% 以内（0.104 vs 0.094 ms），并接近 *compute-bound*，这意味着几乎没有多少 memory stall 可隐藏；Step 9 的 multi-consumer overlap 收回了剩下的大部分空间。接近 compute ceiling 时，最后约 10% gain 正是预期中的 diminishing return：问题几乎已经解决，而不是优化本身薄弱。 + +本章构建的所有内容（TMA load、`tcgen05` MMA、TMEM readback 和 warp-specialized barrier）都会直接带到下一章。Flash Attention 会复用所有这些能力，然后通过在两个 MMA phase 之间插入 online-softmax step，而不是简单重复同一个 MMA，进一步提高难度。 + + +## 练习 -> 翻译状态：待翻译。对应英文章节：`chapter_gemm_advanced/index.md`。 +1. 如果在 Step 7 中把 TMA 和 MMA `PipelineState` 的初始 `phase` 都设为 `0`，会发生什么？画出 deadlock 场景。 +2. Step 8 使用 `cta_group=2` 时，TMA arrive byte count 是 `CTA_GROUP * (BLK_M*BLK_K + BLK_N*BLK_K) * F16_SIZE`。既然每个 CTA 加载自己的数据，为什么还要乘以 `CTA_GROUP`？ +3. Step 9 中，每个 consumer 处理不同 M row，但使用同一个 B tile。为什么共享 B 而不是 A 是正确选择？ -本页用于放置 warp specialization、persistent scheduling 和 cluster GEMM 相关内容的中文翻译。 +**Try with your agent**：粘贴 Step 7 kernel，让它 trace 一个 K-tile 穿过四个 barrier（`tma2mma`、`mma2tma`、`mma2ld`、`ld2mma`）。对每个 barrier，问清楚谁 wait、谁 arrive、哪个 tile 变得可读，以及哪个 buffer 随后变得可复用。 diff --git a/zh/chapter_gemm_async/index.md b/zh/chapter_gemm_async/index.md index 43dc3fe2..d7aca284 100644 --- a/zh/chapter_gemm_async/index.md +++ b/zh/chapter_gemm_async/index.md @@ -1,13 +1,686 @@ ---- -orphan: true ---- - (chap_gemm_async)= +# 用 TMA 为 GEMM 建立 Pipeline + +:::{admonition} 概览 +:class: overview + +- 基础 GEMM 在两个本可同时运行的阶段之间轮流执行：copy 一个 tile、compute、再 copy 下一个 tile，因此浪费了大量时间。 +- Step 4 切换到 TMA async load，Step 5 对 SMEM 做 double buffering 并 prefetch（`PIPE_DEPTH=2`）；完整 load/compute overlap 要等到 Step 7 的 warp specialization，Step 6 则通过 tile scheduler 把 kernel 变成 persistent kernel。 +- 目标是在 Tensor Core 计算当前 tile 的同时，加载下一个 tile。 +::: + +Tensor Core 是芯片上最昂贵的单元，而上一章正确的 tiled GEMM 让它在大部分时钟周期里空闲。Kernel 轮流执行：线程把一个 tile copy 到 shared memory，Tensor Core 消费它，线程再 copy 下一个 tile，Tensor Core 等待。每个阶段都卡在前一个阶段之后，尽管加载下一个 tile 和计算当前 tile 使用的是完全不同的硬件，本可以同时运行。要缩小这个差距，不需要新数据路径；tile、layout 和数学都已经正确。需要改变的是 work *何时*发生，以及由*谁*调度。本章保持 tile 数据路径完全不变，直接攻击空闲时间。 + +我们通过三个递进 step 到达那里。先知道终点会有帮助。Step 4 把 bulk GMEM <-> SMEM transfer 交给 TMA，让专用 copy 硬件移动 tile，而不是由线程移动。Step 5 加入两级 software pipeline，让下一个 K tile 在当前 tile 仍被乘法消费时有地方落下。Step 6 把 launch 重塑为由 tile scheduler 驱动的 persistent kernel，摊销 per-tile setup，并允许我们选择能让 operand 保持 hot 的 tile order。整个过程中，SMEM、TMEM 和 register layout 都保持上一章留下的样子。真正的新思想只有硬件单元之间的异步交接：让一个引擎跑在另一个前面，而不是让它们 lockstep 前进。 + (chap_tma_async)= +## Step 4：TMA Async Load + +第一步是把 copy 本身移出关键路径。回想 Steps 1-3 中 CTA 在做什么：每个线程都计算地址并发出 load 指令，目的只是把 tile 搬进 SMEM。这些指令带宽花在了 plumbing 上，而不是数学计算上。Step 4 用 TMA 替换同步 `Tx.copy`：一个线程发出一条命令，TMA engine 自己完成整个 tile transfer。从这里开始，示例使用完整 M=N=K=4096 规模，而不再使用 Steps 1-3 中的小规模；它们的 end-to-end timing 会出现在 {ref}`chap_gemm_advanced` 末尾的 *End-to-End Result* 表中。 + +> **本 step 改变的内容：Dispatch** +> - Scope：不变，一个 warpgroup。 +> - Layout：不变，同样的 SMEM/TMEM/register tile。 +> - Dispatch：GMEM -> SMEM load 从同步 `Tx.copy` 切换到 TMA engine。 + +### TMA 发起模式 + +Step 4 的唯一变化是把同步 tile copy 换成 TMA load，因此值得仔细看这个 load 如何发起。源码修改只有几行，但这些行背后的执行模型完全不同。同步 `Tx.copy` 是 CTA 线程自己用自己的指令完成的 work；TMA copy 是一个线程发出的命令，之后所有移动都由 TMA 硬件完成。把两者并排看最清楚。 + +**之前（Step 3）**：全部 128 个线程参与 copy，然后 `cta_sync` 让 shared-memory write 可见： + +```python +Tx.cta.copy(Asmem[:, :], A[m_st:m_st+BLK_M, i*BLK_K:(i+1)*BLK_K]) # all 128 threads +Tx.cta.copy(Bsmem[:, :], B[n_st:n_st+BLK_N, i*BLK_K:(i+1)*BLK_K]) +T.cuda.cta_sync() +``` + +**之后（Step 4）**：一个线程发起 TMA load，mbarrier 跟踪硬件 transfer 何时完成： + +```python +tid = warp_id * 32 + lane_id # 0..127 within the warpgroup +if tid == 0: # exactly one thread starts TMA + Tx.copy_async(Asmem, A[...], dispatch="tma") + Tx.copy_async(Bsmem, B[...], dispatch="tma") + T.ptx.mbarrier.arrive.expect_tx(tma_bar, byte_count) # bytes expected from TMA +T.ptx.mbarrier.try_wait(tma_bar, phase) # wait before MMA reads SMEM +``` + +注意，load 用 `tid == 0` gate，而不是用 `elect_sync()`；这个差异比看起来更重要。`elect.sync` 会在*每个 warp* 中选择一个 active lane，而一个 warpgroup 有四个 warp，因此 `elect_sync()` 实际会让四个线程进入 load protocol。问题在于，这个 protocol 会向 mbarrier 宣告 expected byte count，而且必须只宣告一次；四次宣告会破坏计数，让 wait 无法正确 release。用 warpgroup-wide id 精确选择一个线程，是避免这个问题的干净方法。 + +也要诚实说明 speedup 来自哪里。Step 4 仍然在每次 TMA load 后等待，所以还没有让 load 与 compute 重叠；那是 Step 5 的工作。这里的收益纯粹来自 data-movement path 的改变： + +- `Tx.copy` 使用 CTA 线程计算地址，并发出 load/store 指令。 +- TMA 使用一条发起命令启动硬件 tile transfer。地址生成、coalescing 和 swizzling 由 TMA descriptor 描述，并由 TMA engine 执行。 + +因此，即使 Step 4 仍然阻塞在每次 load 上，它也会更快。TMA 吸收了 bulk transfer，让 CTA 线程不必花指令带宽搬 tile；仅这一点就足以改善性能。 + +### TMA Load 和 Store 同步 + +我们已经看过 TMA copy 如何发起；故事的另一半是如何知道它已经完成。切换到 TMA 会同时改变两件事：谁开始 copy，以及代码如何知道它完成。第一件事在代码中很明显；第二件事很容易忽略，而一旦搞错，得到的是静默 correctness bug，不是 crash。使用 `Tx.cta.copy` 时，CTA 线程共同完成 copy，后面的 `cta_sync()` 足以说明它完成。使用 TMA 时，一个被选中的线程发起 `Tx.copy_async(..., dispatch="tma")`，engine 按自己的 schedule 执行 transfer，并通过 mbarrier 发出完成信号。 + +这正是为什么 `cta_sync()` 不再足够。`cta_sync()` 只等待 CTA 自己的线程，并且只排序这些线程的 shared-memory write；它对正在飞行的 TMA transfer 一无所知，因此可能在 tile 仍在到达时就返回。修复方法是显式表达完成：对于 TMA load，被选中的线程先告诉 mbarrier 要期待多少 byte，然后 CTA 在任何 MMA 触碰 SMEM tile 之前等待*这个* mbarrier。下图完整追踪了这个 handshake。 + +![TMA Async Load: Synchronization Flow](../../img/tma_sync_flow.png) + +上图隔离了 load 侧的 handshake：一个选中的线程 launch TMA，mbarrier 计数 expected byte，MMA 在读取 SMEM 之前等待 release。图中 “Elected Thread” 指的是发起 TMA 的 selected thread，在我们的代码里就是 `tid == 0` 线程，而不是 `elect_sync()` 选出的 lane。 + +把 load 路径合起来看：selected thread 发起两个 `copy_async` 调用，然后执行 `arrive.expect_tx(total_bytes)`，其中 byte count 精确说明 mbarrier 应该等待多少数据。Engine 移动完这些 byte 后，匹配的 `mbarrier.try_wait(phase)` 才 release，只有这时 SMEM tile 才能安全喂给 MMA。 + +Store 侧经过同一套硬件，但等待方式不同，所以最好在脑中清楚区分两个 protocol：load 用 mbarrier 和 byte count 跟踪完成，而 store 用 commit group 和 wait group 跟踪完成。线程把 fp16 结果写入 `Dsmem` 并同步后，一个 selected thread 发起 `Tx.copy_async(D[...], Dsmem, dispatch="tma")`，然后 `cp_async.bulk.commit_group()` 加 `cp_async.bulk.wait_group(0)` 会阻塞直到 store drain。这个 wait 不是可选的：上一轮 store 完成前，`Dsmem` 不能被下一个 tile 复用。 + +**Try with your agent**：Trace Step 4 中一个 K tile 的 load 和 store 同步。指出哪个线程启动每条 TMA 命令，哪个 mbarrier 或 commit group 跟踪完成，哪个 wait 保护 MMA 对 `Asmem` 和 `Bsmem` 的读取，哪个 wait 保护 `Dsmem` 的复用。为什么这里用 `elect_sync()` 选择 TMA load 线程是错误的？ + +### 完整 Kernel + +完整 kernel 把 TMA load 和 store 合入 Step 3 的结构，其余结构保持不变。Import 与之前相同： + +```python + +import tvm +from tvm.script import tirx as T +from tvm.script.tirx import tile as Tx +from tvm.tirx.layout import TileLayout, S, TLane, TCol, tid_in_wg +from tvm.tirx.cuda.operator.tile_primitive.tma_utils import tma_shared_layout, SwizzleMode +``` + +它包在 `hgemm_v4(M, N, K)` 中，这是本书一直采用的模式：wrapper 把 shape-dependent constant 和 layout 放在使用它们的 kernel 旁边。 + +```python +def hgemm_v4(M, N, K): + a_type = tvm.DataType("float16") + b_type = tvm.DataType("float16") + d_type = tvm.DataType("float16") + acc_type = tvm.DataType("float32") + + BLK_M, BLK_N, BLK_K = 128, 128, 64 + K_TILES = K // BLK_K + F16_SIZE = 2 + + A_layout = tma_shared_layout(a_type, SwizzleMode.SWIZZLE_128B_ATOM, (BLK_M, BLK_K)) + B_layout = tma_shared_layout(b_type, SwizzleMode.SWIZZLE_128B_ATOM, (BLK_N, BLK_K)) + D_layout = tma_shared_layout(d_type, SwizzleMode.SWIZZLE_128B_ATOM, (BLK_M, BLK_N)) + + @T.prim_func + def kernel( + A: T.Buffer((M, K), a_type), + B: T.Buffer((N, K), b_type), + D: T.Buffer((M, N), d_type), + ): + T.device_entry() + bx, by = T.cta_id([M // BLK_M, N // BLK_N]) + wg_id = T.warpgroup_id([1]) + warp_id = T.warp_id_in_wg([4]) + lane_id = T.lane_id([32]) + + # --- SMEM allocation (now includes Dsmem for TMA store) --- + pool = T.SMEMPool() + tmem_addr = pool.alloc((1,), "uint32") + tma_bar = pool.alloc((1,), "uint64", align=8) + mma_bar = pool.alloc((1,), "uint64", align=8) + pool.move_base_to(1024) + Asmem = pool.alloc((BLK_M, BLK_K), a_type, layout=A_layout) + Bsmem = pool.alloc((BLK_N, BLK_K), b_type, layout=B_layout) + Dsmem = pool.alloc((BLK_M, BLK_N), d_type, layout=D_layout) + pool.commit() + + # --- Barrier + TMEM init --- + if warp_id == 0 and lane_id == 0: + T.ptx.mbarrier.init(mma_bar.ptr_to([0]), 1) + T.ptx.mbarrier.init(tma_bar.ptr_to([0]), 1) + if warp_id == 0: + T.ptx.tcgen05.alloc(T.address_of(tmem_addr), n_cols=512, cta_group=1) + + T.ptx.fence.proxy_async("shared::cta") + T.ptx.fence.mbarrier_init() + T.cuda.cta_sync() + + tmem = T.decl_buffer( + (128, 512), "float32", scope="tmem", allocated_addr=tmem_addr[0], + layout=TileLayout(S[(128, 512) : (1@TLane, 1@TCol)]) + ) + + m_st = T.meta_var(bx * BLK_M) + n_st = T.meta_var(by * BLK_N) + phase_tma: T.int32 = 0 + phase_mma: T.int32 = 0 + + # --- Inline helpers --- + @T.inline + def tma_load(k_st): + tma_config = T.meta_var({ + "dispatch": "tma", "cta_group": 1, + "mbar": tma_bar.ptr_to([0]) + }) + Tx.copy_async(Asmem[:, :], + A[m_st : m_st + BLK_M, k_st : k_st + BLK_K], + **tma_config) + Tx.copy_async(Bsmem[:, :], + B[n_st : n_st + BLK_N, k_st : k_st + BLK_K], + **tma_config) + T.ptx.mbarrier.arrive.expect_tx( + tma_bar.ptr_to([0]), + (BLK_M * BLK_K + BLK_N * BLK_K) * F16_SIZE + ) + + @T.inline + def mma(accum): + Tx.gemm_async( + tmem[:, :BLK_N], Asmem[:, :], Bsmem[:, :], + accum=accum, dispatch="tcgen05", cta_group=1 + ) + T.ptx.tcgen05.commit(mma_bar.ptr_to([0]), cta_group=1) + + # --- K-loop with TMA async --- + tid = T.meta_var(warp_id * 32 + lane_id) + for k in range(K_TILES): + k_st = T.meta_var(k * BLK_K) + + # Single thread issues TMA load + if tid == 0: + tma_load(k_st) + + # Wait for TMA to finish; the mbarrier release carries SMEM + # visibility to the subsequent MMA, so no extra fence is needed. + T.ptx.mbarrier.try_wait(tma_bar.ptr_to([0]), phase_tma) + + # Single thread issues MMA + if tid == 0: + mma(accum=k != 0) + + # Wait for MMA to finish + T.ptx.mbarrier.try_wait(mma_bar.ptr_to([0]), phase_mma) + phase_tma ^= 1 + phase_mma ^= 1 + + # --- TMA Store Writeback --- + Dreg = T.alloc_local((BLK_N,), acc_type) + Dreg_f16 = T.alloc_local((BLK_N,), d_type) + Dreg_wg = Dreg.view(128, BLK_N, + layout=TileLayout(S[(128, BLK_N) : (1@tid_in_wg, 1)])) + + # Read TMEM -> registers (async; wait.ld then cta_sync to ensure read completes) + Tx.wg.copy_async(Dreg_wg[:, :], tmem[:, :BLK_N]) + T.ptx.tcgen05.wait.ld() + T.cuda.cta_sync() + # Cast fp32 -> fp16 + Tx.cast(Dreg_f16[:], Dreg[:]) + # Write registers -> Dsmem, flush, then sync + Tx.copy(Dsmem[warp_id * 32 + lane_id, 0:BLK_N], Dreg_f16[:]) + T.ptx.fence.proxy_async("shared::cta") + T.cuda.warpgroup_sync(10) + # TMA store: Dsmem -> GMEM. One selected thread starts the store and drains the + # store group before Dsmem is reused. + if tid == 0: + Tx.copy_async(D[m_st : m_st + BLK_M, n_st : n_st + BLK_N], + Dsmem[:, :], dispatch="tma") + T.ptx.cp_async.bulk.commit_group() + T.ptx.cp_async.bulk.wait_group(0) + T.cuda.warpgroup_sync(10) + + # --- Deallocate TMEM --- + T.cuda.cta_sync() + if warp_id == 0: + T.ptx.tcgen05.relinquish_alloc_permit(cta_group=1) + T.ptx.tcgen05.dealloc(tmem_addr[0], n_cols=512, cta_group=1) + + return kernel +``` + +### Kernel 中的 TMA 配置 + +这个 kernel 里的大部分内容都继承自 Step 3。真正携带 TMA 语义的配置点只有五个，值得逐一认识： + +- **TMA config**：`{"dispatch": "tma", "cta_group": 1, "mbar": tma_bar.ptr_to([0])}` 告诉 `Tx.copy_async` 使用 TMA，并通过 `tma_bar` 报告 load completion。 + +- **Byte count**：`(BLK_M * BLK_K + BLK_N * BLK_K) * 2` 是两个 fp16 operand tile 加载的 byte 数。`arrive.expect_tx(...)` 把这个 count 交给 mbarrier。 + +- **mbarrier initialization**：`init(tma_bar.ptr_to([0]), 1)` 创建 TMA load 使用的 completion barrier。 + +- **`@T.inline`**：`tma_load(...)` 和 `mma(...)` 是 helper function。它们在编译时展开到 kernel body 中，并且可以使用外围 kernel 的变量。 + +- **TMA store synchronization**：Epilogue 先把 fp16 row 写入 `Dsmem`。`fence.proxy_async` 和 `warpgroup_sync` 让这些由线程写入的 SMEM 值准备好供 TMA store path 使用。Store 随后使用 `commit_group()` 和 `wait_group(0)` 等待 SMEM-to-GMEM transfer 完成。 + +现在我们有了正确的组件，但节奏仍然不对。Step 4 仍然会在开始匹配 MMA 前完成每次 load，所以 load 和 multiply 从未真正同时运行；我们努力分开的两个引擎仍然在轮流工作。下一步保持 TMA load 和 store 路径完全不变，只重新安排 schedule，让一个 K tile 的加载可以在另一个 tile 上的 compute 运行时进行。 + (chap_software_pipeline)= +## Step 5：Software Pipeline（PIPE_DEPTH=2） + +既然 load 和 compute 明显是独立引擎，为什么 Step 4 仍然无法重叠它们？障碍其实是 storage。只有一对 SMEM tile 时，下一个 load 没地方落：它必须等当前 MMA 读完这对 tile 后才能开始，否则会覆盖仍在使用的数据。Step 5 通过 double-buffering shared memory 消除这个 storage conflict。Single-warpgroup loop 仍然会在 launch 下一次 TMA load 之前等待每次 MMA，但现在有不同 stage 可以 prefetch 和复用。我们仍然使用完整 M=N=K=4096 规模。 + +> **本 step 改变的内容：Layout** +> - Scope：不变，一个 warpgroup。 +> - Layout：单个 SMEM tile pair 变成 `PIPE_DEPTH` stage 的 ring buffer。 +> - Dispatch：不变，仍然是 TMA load 和 `tcgen05` MMA；这个 step 加入 prefetch 和 stage reuse，完整 load/compute overlap 会在 Step 7 到来。 + +### Pipeline Walkthrough + +当 `PIPE_DEPTH=2` 时，kernel 会分配两个 SMEM stage，让 load 路径和 MMA 路径有不同 slot 可以工作。 + +请把下图理解成两级 buffer 想要支持的 pipeline structure，而不是这个 single-warpgroup kernel 的精确执行 trace。Step 5 构建 ring buffer 并 prefetch 后续 stage，但主循环仍然在发起下一次 TMA load 前等待当前 MMA。完整 load/compute overlap 会在 Step 7 到来，那时 warp specialization 会把 TMA 和 MMA 分成不同角色。 + +![*Pipeline PIPE_DEPTH=2, the target schedule; this single-warpgroup step only prefetches, full overlap arrives with warp specialization in Step 7*](../../img/pipe_depth2.png) + +Primed 之后，loop 会在两个 stage 间交替。最开始两次 TMA load 会填充两个 stage；之后，loop 等待当前 stage、在其上运行 MMA、等待该 MMA 读完这个 stage，然后把 `k + PIPE_DEPTH` 的 load 发到刚刚变得可复用的 stage 中。这还不是并发 TMA/MMA schedule，但它建立了 Step 7 会拆分到 producer 和 consumer role 中的 ring-buffer 结构。 + +具体来说，代码与 Step 4 有四处不同： + +1. `Asmem` 和 `Bsmem` 增加一个前导 `PIPE_DEPTH` 维度，因此每个 stage 都有自己的 SMEM storage。 +2. `tma_bar` 变成一个 array，每个 stage 一个 mbarrier。 +3. 主 K loop 之前，kernel prefetch 前两个 stage。 +4. K loop 使用 `stage = k % PIPE_DEPTH`：等待当前 stage，在其上运行 MMA，然后把这个 stage 复用于 `k + PIPE_DEPTH`。 + +### Pipeline 机制 + +**1. Prefetch**：主循环开始前，我们加载前 `PIPE_DEPTH` 个 stage，因此 loop 第一次 iteration 就能发现数据已经在等它： + +```python +for s in range(min(PIPE_DEPTH, K_TILES)): + tma_load(s, s * BLK_K) +``` + +**2. Main loop**：对每个 K tile，等待对应 stage ready，在其上运行 MMA，然后立刻把这个刚释放的 stage 重新投入工作，发起领先 `PIPE_DEPTH` 的 tile load： + +```python +stage = k % PIPE_DEPTH +wait(tma_bar[stage], phase_tma) +mma(stage, accum) +wait(mma_bar[0], phase_mma) +phase_mma ^= 1 +tma_load(stage, next_k * BLK_K) +``` + +**3. Phase management**：这是容易绊倒人的地方，但规则比第一眼看上去简单。每个 barrier 的 phase-flip 规则直接来自这个 barrier 有多少个 slot，因此两个 barrier 的翻转节奏不同。MMA accumulator 位于一个 TMEM slot 中，所以 `mma_bar` 是单个 barrier（`mma_bar.ptr_to([0])`），每次 iteration 都会再次访问；每次 iteration 都会访问的 barrier，每次 iteration 都必须 flip phase。TMA barrier 的故事不同：它们形成一个长度为 `PIPE_DEPTH` 的 array，每个 stage 一个 barrier；某个具体 stage 的 barrier 只有在 ring 转一圈后才会再次被访问。因此 `phase_tma` 只在 stage index wrap 回 0 前后，也就是当前 stage 是最后一个 stage 时翻转： + +```python +if stage == PIPE_DEPTH - 1: + phase_tma ^= 1 +``` + +**Try with your agent**：设 `PIPE_DEPTH=2` 且 `K_TILES=5`，让它 trace 主循环。对每个 `k`，列出 `stage`、传给 wait 的 `phase_tma` 和 `phase_mma` 值，以及是否发起新的 prefetch。`phase_tma` 精确在哪里 flip？为什么最后两个 iteration 没有 prefetch？ + +### 完整 Kernel + +完整 kernel 保持 Step 4 的 TMA load 和 store 路径原样，然后把它包进刚才描述的 staged buffer 和 phase 逻辑中。Import 不变： + +```python + +import tvm +from tvm.script import tirx as T +from tvm.script.tirx import tile as Tx +from tvm.tirx.layout import TileLayout, S, TLane, TCol, tid_in_wg +from tvm.tirx.cuda.operator.tile_primitive.tma_utils import tma_shared_layout, SwizzleMode +``` + +它包在 `hgemm_v5(M, N, K)` 中。常量 `PIPE_DEPTH=2` 设置 pipeline stage 数量（这里是两个，正好是 double buffering）： + +```python +PIPE_DEPTH = 2 + +def hgemm_v5(M, N, K): + a_type = tvm.DataType("float16") + b_type = tvm.DataType("float16") + d_type = tvm.DataType("float16") + acc_type = tvm.DataType("float32") + F16_SIZE = 2 + BLK_M, BLK_N, BLK_K = 128, 128, 64 + K_TILES = K // BLK_K + + # Double-buffered layouts: first dimension is pipeline stage + A_layout = tma_shared_layout(a_type, SwizzleMode.SWIZZLE_128B_ATOM, + (PIPE_DEPTH, BLK_M, BLK_K)) + B_layout = tma_shared_layout(b_type, SwizzleMode.SWIZZLE_128B_ATOM, + (PIPE_DEPTH, BLK_N, BLK_K)) + D_layout = tma_shared_layout(d_type, SwizzleMode.SWIZZLE_128B_ATOM, + (BLK_M, BLK_N)) + + @T.prim_func + def kernel( + A: T.Buffer((M, K), a_type), + B: T.Buffer((N, K), b_type), + D: T.Buffer((M, N), d_type), + ): + T.device_entry() + bx, by = T.cta_id([M // BLK_M, N // BLK_N]) + wg_id = T.warpgroup_id([1]) + warp_id = T.warp_id_in_wg([4]) + lane_id = T.lane_id([32]) + + # --- SMEM allocation --- + pool = T.SMEMPool() + tmem_addr = pool.alloc((1,), "uint32") + # Double-buffered TMA barriers (one per stage), single MMA barrier + tma_bar = pool.alloc((PIPE_DEPTH,), "uint64", align=8) + mma_bar = pool.alloc((1,), "uint64", align=8) + pool.move_base_to(1024) + Asmem = pool.alloc((PIPE_DEPTH, BLK_M, BLK_K), a_type, layout=A_layout) + Bsmem = pool.alloc((PIPE_DEPTH, BLK_N, BLK_K), b_type, layout=B_layout) + Dsmem = pool.alloc((BLK_M, BLK_N), d_type, layout=D_layout) + pool.commit() + + # Initialize barriers: PIPE_DEPTH for TMA, 1 for MMA + if warp_id == 0: + if lane_id == 0: + T.ptx.mbarrier.init(mma_bar.ptr_to([0]), 1) + for s in range(PIPE_DEPTH): + T.ptx.mbarrier.init(tma_bar.ptr_to([s]), 1) + if warp_id == 0: + T.ptx.tcgen05.alloc(T.address_of(tmem_addr), n_cols=512, cta_group=1) + + T.ptx.fence.proxy_async("shared::cta") + T.ptx.fence.mbarrier_init() + T.cuda.cta_sync() + + tmem = T.decl_buffer( + (128, 512), acc_type, scope="tmem", allocated_addr=tmem_addr[0], + layout=TileLayout(S[(128, 512) : (1@TLane, 1@TCol)]) + ) + + m_st = T.meta_var(bx * BLK_M) + n_st = T.meta_var(by * BLK_N) + phase_tma: T.int32 = 0 + phase_mma: T.int32 = 0 + + @T.inline + def tma_load(stage, k_offset): + tma_config = T.meta_var({ + "dispatch": "tma", "cta_group": 1, + "mbar": tma_bar.ptr_to([stage]) + }) + Tx.copy_async(Asmem[stage, :, :], + A[m_st:m_st+BLK_M, k_offset:k_offset+BLK_K], + **tma_config) + Tx.copy_async(Bsmem[stage, :, :], + B[n_st:n_st+BLK_N, k_offset:k_offset+BLK_K], + **tma_config) + T.ptx.mbarrier.arrive.expect_tx( + tma_bar.ptr_to([stage]), + (BLK_M * BLK_K + BLK_N * BLK_K) * F16_SIZE) + + @T.inline + def mma(stage, accum): + Tx.gemm_async(tmem[:, :BLK_N], Asmem[stage, :, :], Bsmem[stage, :, :], + accum=accum, dispatch="tcgen05", cta_group=1) + T.ptx.tcgen05.commit(mma_bar.ptr_to([0]), cta_group=1) + + tid = T.meta_var(warp_id * 32 + lane_id) + + # === Prefetch: load first PIPE_DEPTH stages === + if tid == 0: + for s in range(min(PIPE_DEPTH, K_TILES)): + tma_load(s, s * BLK_K) + + # === Main loop === + for k in range(K_TILES): + stage = k % PIPE_DEPTH + + # Wait for TMA to finish loading this stage + T.ptx.mbarrier.try_wait(tma_bar.ptr_to([stage]), phase_tma) + + # MMA on this stage's data + if tid == 0: + mma(stage, accum=(k != 0)) + + T.ptx.mbarrier.try_wait(mma_bar.ptr_to([0]), phase_mma) + phase_mma ^= 1 + + # Issue next prefetch load (k + PIPE_DEPTH) + next_k = k + PIPE_DEPTH + if next_k < K_TILES: + if tid == 0: + tma_load(stage, next_k * BLK_K) + + # TMA phase flips when stage wraps around + if stage == PIPE_DEPTH - 1: + phase_tma ^= 1 + + # === TMA Store Writeback: TMEM -> RF -> Dsmem -> TMA -> GMEM === + Dreg = T.alloc_local((BLK_N,), acc_type) + Dreg_f16 = T.alloc_local((BLK_N,), d_type) + Dreg_wg = Dreg.view(128, BLK_N, + layout=TileLayout(S[(128, BLK_N) : (1@tid_in_wg, 1)])) + Tx.wg.copy_async(Dreg_wg[:, :], tmem[:, :BLK_N]) + T.ptx.tcgen05.wait.ld() + T.cuda.cta_sync() + Tx.cast(Dreg_f16[:], Dreg[:]) + Tx.copy(Dsmem[warp_id * 32 + lane_id, 0:BLK_N], Dreg_f16[:]) + T.ptx.fence.proxy_async("shared::cta") + T.cuda.warpgroup_sync(10) + if tid == 0: + Tx.copy_async(D[m_st : m_st + BLK_M, n_st : n_st + BLK_N], + Dsmem[:, :], dispatch="tma") + T.ptx.cp_async.bulk.commit_group() + T.ptx.cp_async.bulk.wait_group(0) + T.cuda.warpgroup_sync(10) + + # Deallocate TMEM + T.cuda.cta_sync() + if warp_id == 0: + T.ptx.tcgen05.relinquish_alloc_permit(cta_group=1) + T.ptx.tcgen05.dealloc(tmem_addr[0], n_cols=512, cta_group=1) + + return kernel +``` + (chap_persistent_kernel)= -# 用 TMA 为 GEMM 建立 Pipeline +## Step 6：Persistent Kernel + Tile Scheduler + +到目前为止，我们优化的都是单个 tile 内部的工作。Step 6 改变问题的尺度，开始跨 tile 优化。 + +Step 5 为每个 128 x 128 output tile launch 一个 CTA。对于 4096 x 4096 输出，这意味着 1024 个独立 CTA，每个都支付自己的 setup cost，并在完成 tile 后立刻消失。 + +Step 6 改为 launch 固定数量的 CTA pool，然后让每个 CTA 依次处理多个 tile。这带来两个收益：setup work 被摊销到多个 tile 上；tile assignment 进入 kernel 内部，scheduler 可以选择复用 operand 的顺序。我们仍然保持完整 M=N=K=4096 规模。 + +> **本 step 改变的内容：Scope** +> - Scope：固定数量的 persistent CTA，每个 CTA 通过 scheduler loop 处理多个 output tile。 +> - Layout：不变，同样的 per-tile SMEM/TMEM/register 路径。 +> - Dispatch：不变。 + +### Persistent Scheduling + +Persistent kernel 的定义性思想是按硬件规模而不是问题规模设置 grid。它 launch `SM_COUNT` 个 CTA，大致每个 SM 一个，不论 output tile 实际有多少，目标是让每个 SM 持续有工作。这里故意说“大致”：精确 1:1 residency 并不保证，因为它取决于 occupancy，也取决于硬件如何调度 CTA。 + +在这里目标的 B200 上，`SM_COUNT=148`。这 148 个 CTA 中的每一个，都会循环处理 `ClusterPersistentScheduler2D` 分配给它的 tile。 + +第一个收益是摊销。TMEM allocation、barrier initialization 和 scheduler state 现在每个 CTA 只发生一次，并在这个 CTA 处理的大约 7 个 tile 之间复用，而不是在 1024 个一次性 CTA 上反复执行。 + +第二个收益来自 scheduler 选择的顺序。设置 `l2_group_size=8` 会把邻近 tile 分组在一起，因此共享同一 row band 的 tile 会复用同样的 A row-tile，共享同一 column band 的 tile 会复用同样的 B tile。连续运行这些 tile 可以让 operand 保持在 L2 中，而不是从 HBM 重新 fetch。这正是 Step 3 没有利用的复用。 + +```python +bx = T.cta_id([SM_COUNT]) # 1D grid, one CTA per SM + +tile_scheduler = ClusterPersistentScheduler2D( + "ts", + num_m_tiles=M // BLK_M, + num_n_tiles=N // BLK_N, + l2_group_size=8, # Group 8 nearby tiles together + num_clusters=SM_COUNT +) +tile_scheduler.init(bx) +``` + +循环处理多个 tile 带来一个容易忽略的 correctness consequence。每个 tile 都运行自己全新的 K-loop，这意味着它的 barrier phase 必须从已知状态开始。Step 5 中一个 CTA 只处理一个 tile，因此只初始化一次 `phase_tma` 和 `phase_mma` 完全没问题。Step 6 中，这些 initializer 必须移动到 `while tile_scheduler.valid()` loop *内部*，让每个 tile 都用与自己 TMA 和 MMA work 匹配的 phase state 开始，而不是继承上一个 tile 恰好留下的状态： + +```python +while tile_scheduler.valid(): + phase_tma: T.int32 = 0 + phase_mma: T.int32 = 0 + ... +``` + +### 完整 Kernel + +结构上，这个 kernel 只是把 Step 5 的 pipeline 包在 tile-level outer loop 中。唯一新的依赖是 scheduler 本身，我们和其他内容一起 import： + +```python + +import tvm +from tvm.script import tirx as T +from tvm.script.tirx import tile as Tx +from tvm.tirx.layout import TileLayout, S, TLane, TCol, tid_in_wg +from tvm.tirx.cuda.operator.tile_primitive.tma_utils import tma_shared_layout, SwizzleMode +from tvm.tirx.lang.tile_scheduler import ClusterPersistentScheduler2D +``` + +Grid dimension 现在只是 `SM_COUNT`，而不再是 `(M//BLK_M, N//BLK_N)`；`ClusterPersistentScheduler2D` 接管了给每个 CTA 分配 tile 的工作： + +```python +SM_COUNT = 148 # Number of SMs on NVIDIA B200 GPU +PIPE_DEPTH = 2 + +def hgemm_v6(M, N, K): + a_type = tvm.DataType("float16") + b_type = tvm.DataType("float16") + d_type = tvm.DataType("float16") + acc_type = tvm.DataType("float32") + F16_SIZE = 2 + BLK_M, BLK_N, BLK_K = 128, 128, 64 + K_TILES = K // BLK_K + + A_layout = tma_shared_layout(a_type, SwizzleMode.SWIZZLE_128B_ATOM, + (PIPE_DEPTH, BLK_M, BLK_K)) + B_layout = tma_shared_layout(b_type, SwizzleMode.SWIZZLE_128B_ATOM, + (PIPE_DEPTH, BLK_N, BLK_K)) + D_layout = tma_shared_layout(d_type, SwizzleMode.SWIZZLE_128B_ATOM, + (BLK_M, BLK_N)) + + @T.prim_func + def kernel( + A: T.Buffer((M, K), a_type), + B: T.Buffer((N, K), b_type), + D: T.Buffer((M, N), d_type), + ): + T.device_entry() + # 1D grid: one CTA per SM (not a 2D grid anymore!) + bx = T.cta_id([SM_COUNT]) + wg_id = T.warpgroup_id([1]) + warp_id = T.warp_id_in_wg([4]) + lane_id = T.lane_id([32]) + + # --- SMEM allocation (same as Step 5) --- + pool = T.SMEMPool() + tmem_addr = pool.alloc((1,), "uint32") + tma_bar = pool.alloc((PIPE_DEPTH,), "uint64", align=8) + mma_bar = pool.alloc((1,), "uint64", align=8) + pool.move_base_to(1024) + Asmem = pool.alloc((PIPE_DEPTH, BLK_M, BLK_K), a_type, layout=A_layout) + Bsmem = pool.alloc((PIPE_DEPTH, BLK_N, BLK_K), b_type, layout=B_layout) + Dsmem = pool.alloc((BLK_M, BLK_N), d_type, layout=D_layout) + pool.commit() + + # --- Barrier + TMEM init (same as Step 5) --- + if warp_id == 0 and lane_id == 0: + T.ptx.mbarrier.init(mma_bar.ptr_to([0]), 1) + for s in range(PIPE_DEPTH): + T.ptx.mbarrier.init(tma_bar.ptr_to([s]), 1) + if warp_id == 0: + T.ptx.tcgen05.alloc(T.address_of(tmem_addr), n_cols=512, cta_group=1) + T.ptx.fence.proxy_async("shared::cta") + T.ptx.fence.mbarrier_init() + T.cuda.cta_sync() + + tmem = T.decl_buffer( + (128, 512), acc_type, scope="tmem", allocated_addr=tmem_addr[0], + layout=TileLayout(S[(128, 512) : (1@TLane, 1@TCol)]) + ) + + # Tile scheduler: assigns tiles to CTAs in L2-friendly order + tile_scheduler = ClusterPersistentScheduler2D( + "ts", + num_m_tiles=M // BLK_M, + num_n_tiles=N // BLK_N, + l2_group_size=8, + num_clusters=SM_COUNT + ) + tile_scheduler.init(bx) + + tid = T.meta_var(warp_id * 32 + lane_id) + + @T.inline + def tma_load(stage, k_offset, m_st, n_st): + tma_config = T.meta_var({ + "dispatch": "tma", "cta_group": 1, + "mbar": tma_bar.ptr_to([stage]) + }) + Tx.copy_async(Asmem[stage, :, :], + A[m_st:m_st+BLK_M, k_offset:k_offset+BLK_K], + **tma_config) + Tx.copy_async(Bsmem[stage, :, :], + B[n_st:n_st+BLK_N, k_offset:k_offset+BLK_K], + **tma_config) + T.ptx.mbarrier.arrive.expect_tx( + tma_bar.ptr_to([stage]), + (BLK_M * BLK_K + BLK_N * BLK_K) * F16_SIZE) + + @T.inline + def mma(stage, accum): + Tx.gemm_async(tmem[:, :BLK_N], Asmem[stage, :, :], Bsmem[stage, :, :], + accum=accum, dispatch="tcgen05", cta_group=1) + T.ptx.tcgen05.commit(mma_bar.ptr_to([0]), cta_group=1) + + # === Outer loop: iterate over tiles === + while tile_scheduler.valid(): + # Get current tile position from scheduler + m_st = T.meta_var(tile_scheduler.m_idx * BLK_M) + n_st = T.meta_var(tile_scheduler.n_idx * BLK_N) + + # === Inner loop: same pipeline as Step 5 === + phase_tma: T.int32 = 0 + phase_mma: T.int32 = 0 + + # Prefetch first PIPE_DEPTH stages + if tid == 0: + for s in range(min(PIPE_DEPTH, K_TILES)): + tma_load(s, s * BLK_K, m_st, n_st) + + # Main K-loop + for k in range(K_TILES): + stage = k % PIPE_DEPTH + T.ptx.mbarrier.try_wait(tma_bar.ptr_to([stage]), phase_tma) + if tid == 0: + mma(stage, accum=(k != 0)) + T.ptx.mbarrier.try_wait(mma_bar.ptr_to([0]), phase_mma) + phase_mma ^= 1 + next_k = k + PIPE_DEPTH + if next_k < K_TILES: + if tid == 0: + tma_load(stage, next_k * BLK_K, m_st, n_st) + if stage == PIPE_DEPTH - 1: + phase_tma ^= 1 + + # === TMA Store Writeback: TMEM -> RF -> Dsmem -> TMA -> GMEM === + Dreg = T.alloc_local((BLK_N,), acc_type) + Dreg_f16 = T.alloc_local((BLK_N,), d_type) + Dreg_wg = Dreg.view(128, BLK_N, + layout=TileLayout(S[(128, BLK_N) : (1@tid_in_wg, 1)])) + Tx.wg.copy_async(Dreg_wg[:, :], tmem[:, :BLK_N]) + T.ptx.tcgen05.wait.ld() + T.cuda.cta_sync() + Tx.cast(Dreg_f16[:], Dreg[:]) + Tx.copy(Dsmem[warp_id * 32 + lane_id, 0:BLK_N], Dreg_f16[:]) + T.ptx.fence.proxy_async("shared::cta") + T.cuda.warpgroup_sync(10) + if tid == 0: + Tx.copy_async(D[m_st : m_st + BLK_M, n_st : n_st + BLK_N], + Dsmem[:, :], dispatch="tma") + T.ptx.cp_async.bulk.commit_group() + T.ptx.cp_async.bulk.wait_group(0) + T.cuda.warpgroup_sync(10) + + T.cuda.cta_sync() + tile_scheduler.next_tile() # Move to next tile + + # Deallocate TMEM + T.cuda.cta_sync() + if warp_id == 0: + T.ptx.tcgen05.relinquish_alloc_permit(cta_group=1) + T.ptx.tcgen05.dealloc(tmem_addr[0], n_cols=512, cta_group=1) + + return kernel +``` -> 翻译状态：待翻译。对应英文章节：`chapter_gemm_async/index.md`。 +## 练习 -本页用于放置 GEMM TMA pipeline 相关内容的中文翻译。 +1. 在 Step 4 中，`arrive.expect_tx` 使用 `(BLK_M * BLK_K + BLK_N * BLK_K) * 2` byte。如果这个 byte count 太小或太大，mbarrier 会等待什么？ +2. 在 Step 5 中，为什么每个 SMEM stage 都需要自己的 TMA barrier，而不是两个 stage 共用同一个 `tma_bar`？ +3. 在 Step 6 中，一个 4096 x 4096 output 且 `BLK_M=BLK_N=128` 时有多少 output tile？如果 `SM_COUNT=148`，每个 persistent CTA 平均处理多少 tile？ diff --git a/zh/chapter_gemm_basics/index.md b/zh/chapter_gemm_basics/index.md index 0fa614b3..5892b5d8 100644 --- a/zh/chapter_gemm_basics/index.md +++ b/zh/chapter_gemm_basics/index.md @@ -1,13 +1,602 @@ ---- -orphan: true +(chap_gemm_basics)= +# 构建 Tiled GEMM + +:::{admonition} 概览 +:class: overview + +- 从 TIRx tile primitive 出发，构建一个正确的 tiled GEMM，起点是单个 output tile。 +- Step 1 是 single-tile GEMM，Step 2 加入 K-loop accumulation，Step 3 在 CTA 之间做 spatial tiling 以覆盖完整矩阵。 +- 先保证正确性；性能优化交给接下来的两章。 +::: + +GEMM 是本书围绕构建的核心 workload。它位于 linear layer、attention projection 和 convolution 的底层，而这些操作占据了 GPU 的大部分时间。因此，一个只“正确”的 GEMM 和一个“快速”的 GEMM 之间的差别，往往就是芯片大部分算力空转和接近饱和之间的差别。 + +这个差距太大，无法一步跨过去。一个接近饱和的 kernel 会让你同时调试内存搬运、accumulation、tiling 和 Tensor Core scheduling，而且一开始还没有可信的 baseline 可以对照。更稳妥的路径是从能产生正确答案的最小 kernel 开始，然后一次只增加一个设计决策。 + +本章会写出第一个正确的 tiled GEMM。前面章节以抽象形式介绍了 TIRx 的 scope / layout / dispatch 模型；这里我们把它应用到真实 kernel 上。我们从一个 128 x 128 output tile 开始，然后把它扩展成可以处理完整矩阵的 kernel：先加入 K 维度 accumulation，再加入跨多个 CTA 的 spatial tiling。 + +这是三章 GEMM 优化路径中的第一章。这条路径会从头到尾走完一个 GEMM kernel 的演化。本章只构建正确的 tiled kernel 并到此为止。下一章（{ref}`chap_gemm_async`）会把 thread copy 换成 TMA，并通过 pipelining 让数据搬运和计算重叠；{ref}`chap_gemm_advanced` 会进一步加入 warp specialization 和 CTA cluster。每章都建立在前一章之上，因此 kernel 会逐步积累功能，而不是每章重新开始。 + +阅读每个 step 时，可以把它看作对同一份三项 contract 的一次编辑：哪个 **scope** 执行操作，operand tile 使用哪个 **layout**，以及通过哪条 **dispatch** 路径执行。多数 step 都只有一个主要变化，因此我们会先用一个小卡片指出变化是什么，并标出让复用安全所需的同步细节。Step 1 会建立后续所有修改的 baseline。 + +## GEMM + +GEMM 是 dense matrix multiply，位于 linear layer、attention projection 和许多 convolution 实现的底层，所以快速 GEMM kernel 几乎在任何地方都有收益。本教程中的例子使用 $D = A B^{\top}$： + +- $A$ 的 shape 是 $M \times K$。 +- $B$ 的 shape 是 $N \times K$。 +- $D$ 的 shape 是 $M \times N$。 +- $D[m,n] = \sum_k A[m,k] \cdot B[n,k]$。 + +这里的 transpose 不是我们额外选择执行的操作；它来自数据的存储方式。示例保持 $B$ 为 $N$ 行、每行长度 $K$，这也是 linear-layer weight 通常使用的布局。因此沿 $K$ contraction 时，自然就是在读取 $B^{\top}$，不需要实际重排。 + +整个教程中，我们用 TFLOPS 衡量 kernel 的吞吐，把每次 multiply-add 计为两个 floating-point operation，并除以 wall-clock time： + +$$\text{TFLOPS} = \frac{2 \times M \times N \times K}{t_{\text{seconds}} \times 10^{12}}$$ + +### GEMM 数据路径 + +本教程中的每个优化最终都归结为数据位于哪里、如何移动。因此在写代码前，先把这条路径画出来是值得的。一个 Blackwell GEMM kernel 的核心只有两类活动：在不同 memory 之间搬运 tile，以及在 tile 上计算。下图追踪一个 tile 从输入到输出会触碰的每一种 memory： + +![*Memory Data Flow*](../../img/memory_dataflow.png) + +上图展示了 baseline 路径。之后每个优化都会编辑这条路径，但不会替换它。从左到右读：operand tile 先从 GMEM 移到 SMEM；随后 `tcgen05.mma` 消费 SMEM operand，并把 accumulator 写入 TMEM；最后 epilogue 把 TMEM 读回寄存器，再把结果 store 到 GMEM。请记住这条链路，因为下面每一步都会改变其中某一跳*如何*发生，但不会改变这些跳本身。 + +## 优化路径 + +上面的朴素数据路径已经足以得到正确答案，但会让大部分硬件空闲。教程剩余部分会一次加入一个 Blackwell 特性来缩小这个差距，每个特性都通过 TIRx tile primitive 表达。我们将依次经过这些特性： + +- **TMA async movement** 通过 Blackwell 的硬件 copy 路径移动 GMEM <-> SMEM tile，并用 barrier 跟踪完成。 +- **Software pipelining** 使用多个 SMEM stage，让下一块 K tile 的数据搬运可以与当前 tile 上的 Tensor Core compute 重叠。 +- **Persistent scheduling** 保持一组固定 CTA，让每个 CTA 通过 tile scheduler 处理多个 output tile，而不是每个 tile launch 一个 CTA。 +- **Warp specialization** 把 producer、MMA consumer 和 writeback 角色拆分到不同 warpgroup 上。 +- **CTA clusters** 让两个 CTA 协作处理一个更大的 Blackwell MMA tile。 +- **Multi-consumer execution** 使用多个 consumer warpgroup 同时计算 tile 的不同部分，提高 compute density。 + --- -(chap_gemm_basics)= (chap_single_tile)= +## Step 1：顺序 Single-Tile GEMM + +仍然能覆盖完整硬件路径的最简单 GEMM，是计算单个 output tile 的 GEMM。因此我们从这里开始。Step 1 计算一个 128 x 128 output tile，K = 64；这个规模小到不需要 loop，并且数据路径中的每个部分都只出现一次。没有重复结构时，我们可以先单独看清每一跳，然后再开始推理循环。 + +> **本 step 建立的内容：baseline** +> - Scope：一个 128 线程的 single warpgroup 按顺序走完整条路径，一阶段接一阶段。 +> - Layout：A 和 B tile 位于 SMEM，accumulator 位于 TMEM，结果通过寄存器 staged out。 +> - Dispatch：同步 `Tx.copy` 执行 load，`tcgen05` 执行 MMA。 + +### Single-Tile Dataflow + +Baseline contract 固定后，下一件事是确定一个 tile 按什么顺序穿过它。第一个 kernel 会完整走一次核心 GEMM 数据路径，也就是 data-flow 图里的同一条 GMEM -> SMEM -> TMEM -> registers -> GMEM 链路，外面没有包任何 loop。它分配工作内存、加载 operand、计算乘积、写回结果，并清理自己使用的资源： + +1. **Allocate**：SMEM（pool allocator）、TMEM（`tcgen05.alloc`）、mbarrier +2. **Load**：全部 128 个线程协作把 A 和 B tile 从 GMEM copy 到 SMEM（sync `Tx.copy`） +3. **Compute**：一个 elected thread 发起 `Tx.gemm_async` + `tcgen05.commit`；所有线程在 mbarrier 上等待 +4. **Writeback**：Warpgroup 读取 TMEM -> registers；每个线程把 fp32 cast 到 fp16 并写入 GMEM +5. **Deallocate**：释放 TMEM + +### 第一个 Kernel 的四个部分 + +完整 kernel 只有几十行，但分段读更容易消化。我们会按四部分阅读它：memory allocation、同步 load、MMA dispatch 和 writeback；之后再把它们拼成一个 kernel。沿途出现的 API 名称，是第二部分介绍过的 TIRx tile-primitive 词汇（{ref}`chap_tirx_primer`、{ref}`chap_tirx_layout_api`）。 + +**Memory allocation。** Kernel 首先从 shared memory 中切出 operand 所需空间，以及 TMEM address 和 mbarrier 的位置： + +```python +pool = T.SMEMPool() +tmem_addr = pool.alloc((1,), "uint32") # TMEM address (4 bytes) +mma_bar = pool.alloc((1,), "uint64", align=8) # mbarrier (8 bytes) +pool.move_base_to(1024) # Skip to offset 1024 +Asmem = pool.alloc((BLK_M, BLK_K), a_type, layout=A_layout) # 128×64 fp16 +Bsmem = pool.alloc((BLK_N, BLK_K), b_type, layout=B_layout) # 128×64 fp16 +pool.commit() +``` + +这里有两个细节值得停一下。`pool.move_base_to(1024)` 会把 Asmem 和 Bsmem 推到 offset 1024，给前面较小的 metadata 预留低地址区域，使 bulky operand tile 落在干净边界上。`layout=A_layout` 会让 `tma_shared_layout` 提供一个 swizzled SMEM placement，这个 placement 能被 TMA 和 `tcgen05.mma` 直接读取，正是第二部分所说的 layout-as-contract。 + +**Synchronous load。** Buffer 到位后，operand 还需要到达 SMEM。在第一个版本中，我们让 CTA 自己的线程执行 copy： + +```python +Tx.cta.copy(Asmem[:, :], A[:, :]) +Tx.cta.copy(Bsmem[:, :], B[:, :]) +T.cuda.cta_sync() +``` + +因为这里总共只有一个 tile（M=N=128, K=64），copy 完整 A 和 B 就是整个 load。`Tx.cta.copy(...)` 让 CTA 在这次 copy 上协作，每个线程负责自己那一片数据。后面的 `T.cuda.cta_sync()` 有双重作用：它等待每个线程完成，并发布这些线程对 shared memory 的写入，因此后续 MMA 读取 `Asmem` 和 `Bsmem` 时看到的是完整 tile，而不是半填充 buffer。这个 thread-driven copy 也是我们首先会替换的东西；下一章（{ref}`chap_gemm_async`）会把它换成 TMA。 + +**MMA dispatch。** Operand 已经位于 SMEM 中，现在可以发起 MMA，并且由一个 elected thread 来做： + +```python +if warp_id == 0: + if T.ptx.elect_sync(): + Tx.gemm_async(tmem[:, :BLK_N], Asmem[:, :], Bsmem[:, :], + accum=False, dispatch="tcgen05", cta_group=1) + T.ptx.tcgen05.commit(mma_bar.ptr_to([0]), cta_group=1) +``` + +两层 guard 会分两步把 issuer 缩小到一个线程。外层 `if warp_id == 0` 只保留 warpgroup 中的 warp 0，内层 `if T.ptx.elect_sync():` 再从这个 warp 的 active lane 中选出一个。合起来，只剩一个线程执行 `Tx.gemm_async` 和 `tcgen05.commit`。 + +这里需要明确说明单个线程意味着什么、不意味着什么，因为直觉读法很容易误导。单个 issuing thread 并不意味着单线程矩阵乘法。计算仍然是完整的 tile-level MMA：硬件会根据 SMEM operand layout 和 TMEM accumulator layout 描述的 tile 执行协作矩阵乘法。关键在于 `Tx.gemm_async` 是一个 *tile operation*，不是一条硬件指令。K = 64 tile 比硬件 MMA K-atom（`MMA_K = 16`）更宽，所以这个 tile op 会 lower 成沿 K 前进的一小段 raw `tcgen05.mma` 指令序列，而 warpgroup 会协作驱动每一条。之所以只有一个线程发起 tile op，是因为底层每条 `tcgen05.mma` 本身就是一条 cooperative op：一次 launch 驱动这个 K-atom 的 tile MMA。如果 128 个线程都发起同一段序列，同样的工作就会被 launch 128 次。最后，`accum=False` 告诉 MMA 覆盖 TMEM destination，而不是加到已有值上；这里没有先前 partial sum，所以这正是我们想要的。 + +**Writeback。** 乘积现在位于 TMEM 中，但调用方希望在 GMEM 中得到 fp16 结果。因此 epilogue 必须先把结果通过寄存器带下来，并在途中 cast： + +```python +Dreg = T.alloc_local((BLK_N,), acc_type) # per-thread fp32 register row +Dreg_f16 = T.alloc_local((BLK_N,), d_type) # same row, cast to fp16 +Dreg_wg = Dreg.view(128, BLK_N, layout=TileLayout(S[(128, BLK_N) : (1@tid_in_wg, 1)])) +Tx.wg.copy_async(Dreg_wg[:, :], tmem[:, :BLK_N]) +T.ptx.tcgen05.wait.ld() +Tx.cast(Dreg_f16[:], Dreg[:]) +m_thr = T.meta_var(m_st + warp_id * 32 + lane_id) +Tx.copy(D[m_thr, n_st : n_st + BLK_N], Dreg_f16[:]) +``` + +MMA 会在 TMEM 中留下一个 128 x 128 fp32 accumulator tile。使用 fp32 是有意的：GEMM 会沿 K 累加很多乘积，用更高精度保存 running sum 可以降低累积的 rounding error。但 `D` 是 fp16，所以这些值不能直接写出。它们先进入寄存器，在那里 narrow 到 fp16，然后才到达 GMEM。 + +两个 register buffer 作用不同。`Dreg` 是每个线程自己的 `BLK_N` 元素 buffer，而 `Dreg_wg` 是同一组寄存器在所选 layout 下的 warpgroup-wide *view*： + +```python +TileLayout(S[(128, BLK_N) : (1@tid_in_wg, 1)]) +``` + +这个 layout 把 tile 的第一维映射到 warpgroup 的线程上：thread 0 拥有 row 0，thread 1 拥有 row 1，一直到 row 127。第二维留在每个线程自己的 register buffer 中，因此单个线程持有自己那一行的所有列。Warpgroup 有 128 个线程，tile 有 128 行，所以 128 x 128 输出刚好分成每个线程一行。 + +在这个 view 下读取 accumulator，正是 `Tx.wg.copy_async(Dreg_wg, tmem)` 所做的事，它会 lower 到 Blackwell TMEM load 路径 `tcgen05.ld`。由于这个 load 是异步的，任何线程触碰 `Dreg` 之前都必须先完成 `T.ptx.tcgen05.wait.ld()`；否则线程可能读取尚未被 load 填好的寄存器。 + +Wait 返回后，每个线程私有的 `Dreg[:]` 保存自己那一条逻辑输出行的 fp32 值。线程把它们 narrow 到 `Dreg_f16` 中，计算自己负责的 global row： + +```python +m_thr = T.meta_var(m_st + warp_id * 32 + lane_id) +``` + +然后写入 `D[m_thr, n_st:n_st + BLK_N]`。这些 row 在四个 warp 之间整齐切分：warp 0 写 rows 0-31，warp 1 写 rows 32-63，warp 2 写 rows 64-95，warp 3 写 rows 96-127。 + +### 完整 Kernel + +现在把四个部分拼回一个可运行 kernel（M=N=128, K=64）。Import 先出现： + +```python + +import tvm +from tvm.script import tirx as T +from tvm.script.tirx import tile as Tx +from tvm.tirx.cuda.operator.tile_primitive.tma_utils import tma_shared_layout, SwizzleMode +from tvm.tirx.layout import TileLayout, S, TLane, TCol, tid_in_wg +``` + +Kernel 包在后续 step 都会使用的 `hgemm_vX(M, N, K)` 风格中。Step 1 使用 `M=N=128, K=64`，因此 launch 中刚好有一个 output tile： + +```python +def hgemm_v1(M, N, K): + a_type = tvm.DataType("float16") + b_type = tvm.DataType("float16") + d_type = tvm.DataType("float16") + acc_type = tvm.DataType("float32") + + BLK_M, BLK_N, BLK_K = 128, 128, 64 + # MMA_M/MMA_N/MMA_K document the underlying hardware MMA tile; they are not + # passed to gemm_async (which derives the MMA shape from the operand and + # accumulator tiles), so the later steps omit them. + MMA_M, MMA_N, MMA_K = 128, 128, 16 + + A_layout = tma_shared_layout(a_type, SwizzleMode.SWIZZLE_128B_ATOM, (BLK_M, BLK_K)) + B_layout = tma_shared_layout(b_type, SwizzleMode.SWIZZLE_128B_ATOM, (BLK_N, BLK_K)) + + @T.prim_func + def kernel( + A: T.Buffer((M, K), a_type), + B: T.Buffer((N, K), b_type), + D: T.Buffer((M, N), d_type), + ): + T.device_entry() + # Step 1 is a single-tile kernel: M = BLK_M and N = BLK_N, so the grid + # is 1x1. Starting with a 1x1 grid keeps the per-CTA tile offsets + # (m_st, n_st) trivially zero; Steps 3+ generalise this to larger M / N. + bx, by = T.cta_id([M // BLK_M, N // BLK_N]) + wg_id = T.warpgroup_id([1]) # single warpgroup, so wg_id is always 0 (unused below) + warp_id = T.warp_id_in_wg([4]) + lane_id = T.lane_id([32]) + + # --- SMEM allocation --- + pool = T.SMEMPool() + tmem_addr = pool.alloc((1,), "uint32") + mma_bar = pool.alloc((1,), "uint64", align=8) + pool.move_base_to(1024) + Asmem = pool.alloc((BLK_M, BLK_K), a_type, layout=A_layout) + Bsmem = pool.alloc((BLK_N, BLK_K), b_type, layout=B_layout) + pool.commit() + + # --- Barrier + TMEM init (warp 0 only) --- + if warp_id == 0: + if lane_id == 0: + T.ptx.mbarrier.init(mma_bar.ptr_to([0]), 1) + T.ptx.tcgen05.alloc(T.address_of(tmem_addr), n_cols=512, cta_group=1) + + T.ptx.fence.proxy_async("shared::cta") + T.ptx.fence.mbarrier_init() + T.cuda.cta_sync() + + tmem = T.decl_buffer( + (128, 512), "float32", scope="tmem", allocated_addr=tmem_addr[0], + layout=TileLayout(S[(128, 512) : (1@TLane, 1@TCol)]) + ) + + m_st = T.meta_var(bx * BLK_M) + n_st = T.meta_var(by * BLK_N) + phase_mma: T.int32 = 0 + + # --- Load: all threads copy global -> shared (synchronous). + # With M=BLK_M and N=BLK_N the slices below cover the full matrices; + # the slice form is kept so the diff to Step 3 (multi-tile) is minimal. + Tx.cta.copy(Asmem[:, :], A[m_st:m_st + BLK_M, :]) + Tx.cta.copy(Bsmem[:, :], B[n_st:n_st + BLK_N, :]) + T.cuda.cta_sync() + + # --- Compute: single elected thread issues MMA --- + if warp_id == 0: + if T.ptx.elect_sync(): + Tx.gemm_async( + tmem[:, :BLK_N], Asmem[:, :], Bsmem[:, :], + accum=False, dispatch="tcgen05", cta_group=1 + ) + T.ptx.tcgen05.commit(mma_bar.ptr_to([0]), cta_group=1) + + T.ptx.mbarrier.try_wait(mma_bar.ptr_to([0]), phase_mma) + + # --- Writeback: TMEM -> RF -> GMEM --- + Dreg = T.alloc_local((BLK_N,), acc_type) + Dreg_f16 = T.alloc_local((BLK_N,), d_type) + Dreg_wg = Dreg.view(128, BLK_N, + layout=TileLayout(S[(128, BLK_N) : (1@tid_in_wg, 1)])) + Tx.wg.copy_async(Dreg_wg[:, :], tmem[:, :BLK_N]) + T.ptx.tcgen05.wait.ld() + Tx.cast(Dreg_f16[:], Dreg[:]) + m_thr = T.meta_var(m_st + warp_id * 32 + lane_id) + Tx.copy(D[m_thr, n_st : n_st + BLK_N], Dreg_f16[:]) + + # --- Deallocate TMEM --- + T.cuda.cta_sync() + if warp_id == 0: + T.ptx.tcgen05.relinquish_alloc_permit(cta_group=1) + T.ptx.tcgen05.dealloc(tmem_addr[0], n_cols=512, cta_group=1) + + return kernel +``` + +后面的每个 GEMM step 都会用同样方式编译、运行并检查自己，因此这个 scaffolding 只完整写一次。从此之后，我们只展示 kernel。要运行后续 step，把下面的 `hgemm_vX` 和匹配的问题规模换成对应版本即可。有一个注意事项：每次新的 Python session 只编译一个 step，尝试另一个 step 前先重启，因为这些示例会复用内部名字，而编译器持有 per-session state。 + +```python +import torch + +target = tvm.target.Target("cuda") +device = torch.device('cuda') # gpu(0) + +M, N, K = 128, 128, 64 +kernel = hgemm_v1(M, N, K) +with target: + ex = tvm.compile(tvm.IRModule({"main": kernel}), target=target, tir_pipeline="tirx") + +torch.cuda.empty_cache() +torch.cuda.synchronize() +A_tensor = torch.randn(M, K, dtype=torch.float16, device=device) +B_tensor = torch.randn(N, K, dtype=torch.float16, device=device) +D_tensor = torch.zeros(M, N, dtype=torch.float16, device=device) + +# ex.mod(...) takes torch tensors directly, the same call form used in every chapter. +ex.mod(A_tensor, B_tensor, D_tensor) + +D_ref = (A_tensor.float() @ B_tensor.float().T).half() +max_err = float((D_tensor - D_ref).abs().max()) +print(f"Max error vs torch reference: {max_err:.6f}") +# Relative tolerance, like the warp-specialization and Flash Attention cells: +# output magnitude grows with K, so a fixed absolute bound would fail at larger K. +torch.testing.assert_close(D_tensor, D_ref, rtol=2e-2, atol=1e-2) +print("PASS") + +# Optional timing for larger kernels. +ITERS = 10 +for _ in range(3): + ex.mod(A_tensor, B_tensor, D_tensor) +torch.cuda.synchronize() +start = torch.cuda.Event(enable_timing=True) +end = torch.cuda.Event(enable_timing=True) +start.record() +for _ in range(ITERS): + ex.mod(A_tensor, B_tensor, D_tensor) +end.record() +torch.cuda.synchronize() +ms = start.elapsed_time(end) / ITERS +tflops = 2 * M * N * K / ms / 1e9 +print(f"Performance: {ms:.3f} ms, {tflops:.1f} TFLOPS") +``` + +Steps 1 到 3 会刻意使用较小规模（这里是 128×128，Step 3 是 256³），使最初几个 walkthrough 容易跟上。{ref}`chap_gemm_advanced` 末尾的跨 step *End-to-End Result* 表则采用相反策略：它把每个 step，包括这个 Step 1 算法，都放到统一的 M=N=K=4096 规模下测量，因此 speedup ratio 可以直接比较。 + +### Single-Tile Kernel 的限制 + +这个 kernel 是正确的，这正是 Step 1 的目标，但它只在非常窄的设定下正确。这里有四个限制是有意留下的，后续优化路径会逐个解除它们： + +- 它只处理单个 K tile，因此不能对很大的 K 做 contraction。 +- 它只处理单个 output tile，因此 M 和 N 被固定在 128。 +- 它使用同步 GMEM -> SMEM copy，而不是 TMA。 +- 它没有重叠数据搬运和计算，所以两者不会同时运行。 + +--- + (chap_k_loop)= +## Step 2：K-Loop Accumulation + +第一个要移除的是最小的限制。Step 1 只处理单个宽度为 64 的 K tile，但真实矩阵会沿远大于 64 的 K 做 contraction。在 Step 2 中，我们仍然只计算单个 output tile，但允许 K 跨越多个 64-wide chunk。 + +想法很直接：对每个 chunk 重复 load -> MMA -> wait 序列，并让每次 MMA 累加到同一个 TMEM slot 中。真正需要小心的地方是同步。跨 iteration 复用同一个 mbarrier，会引入本章第一个真实 correctness hazard。如果代码跟踪了错误 phase，某次 wait 可能在对应 MMA 真正完成*之前*返回，静默破坏结果。下面的机制会精确说明这个错误如何发生，以及如何避免。 + +> **本 step 改变的内容：Layout reuse** +> - Scope：不变，仍然是 single warpgroup。 +> - Layout/reuse：同一对 SMEM tile 和同一个 TMEM accumulator slot 会在 K-loop 中复用。不分配新 storage；operand tile 流经固定的一对 buffer，accumulator state 保持在一个 TMEM slot 中。 +> - Synchronization：复用的 MMA barrier 必须在每个 K chunk 上推进到正确 phase，否则后续 wait 可能观察到更早的 completion。 +> - Dispatch：不变。 + +### K-Loop 机制 + +Step 1 只 contraction 了单个 64-wide K tile；这里我们保留它的 single output tile，但让 K 按矩阵需要的长度前进。为了覆盖大于 64 的 K，我们以 `BLK_K=64` 为 chunk 沿 K 迭代。每次 iteration 加载下一段 A 和 B 的 K-slice 到 SMEM，并发起 `Tx.gemm_async`。`accum` flag 把这些 chunk 拼成同一个 dot product：第一个 chunk 上 `accum=False` 初始化 TMEM accumulator，之后每个 chunk 上 `accum=True` 把该 chunk 的乘积加到 TMEM 中已经存在的 running sum 上。 + +同步是需要谨慎的地方。我们为每次 MMA completion 复用同一个 mbarrier，而安全复用的关键是跟踪正在等待哪个 barrier phase。一个 mbarrier 带有 1-bit phase，可以是 0 或 1；每当期望的 arrival 到达，它就翻转到另一个值。微妙之处在于 wait 条件本身：`try_wait(bar, phase)` 会阻塞，直到 barrier 的内部 phase *不同于* `phase` 参数。因此我们传入的参数必须命名我们期望离开的 phase，而不是等待抵达的 phase： + +| K iteration | Wait 前本地 `phase_mma` | `try_wait` 等待什么 | Wait 后本地更新 | +|---|---:|---|---:| +| 0 | 0 | barrier flips to 1 | `phase_mma = 1` | +| 1 | 1 | barrier flips to 0 | `phase_mma = 0` | +| 2 | 0 | barrier flips to 1 | `phase_mma = 1` | + +`phase_mma ^= 1` 这一行正是保持这张表正确的原因。去掉它后，第二次 iteration 仍然调用 `try_wait(bar, 0)`，但 barrier 在第一次 MMA 后已经翻到了 phase 1，因此 wait 看到 mismatch 就立即返回，而此时第二次 MMA 还没完成。Kernel 随后会读取半计算的 accumulator，并在没有任何 error 的情况下给出错误答案。这个 bug 可以完美编译和运行，这就是 phase flip 值得如此强调的原因。 + +### 完整 Kernel + +下面的完整 kernel 只是 Step 1 加入 K-loop 和 phase flip。Import 和之前相同： + +```python + +import tvm +from tvm.script import tirx as T +from tvm.script.tirx import tile as Tx +from tvm.tirx.cuda.operator.tile_primitive.tma_utils import tma_shared_layout, SwizzleMode +from tvm.tirx.layout import TileLayout, S, TLane, TCol, tid_in_wg +``` + +它包在 `hgemm_v2(M, N, K)` 中。Grid 仍然是 `[1, 1]`，因为我们仍然只计算单个 output tile；增长的只是 K extent： + +```python +def hgemm_v2(M, N, K): + a_type = tvm.DataType("float16") + b_type = tvm.DataType("float16") + d_type = tvm.DataType("float16") + acc_type = tvm.DataType("float32") + + BLK_M, BLK_N, BLK_K = 128, 128, 64 + K_TILES = K // BLK_K + + A_layout = tma_shared_layout(a_type, SwizzleMode.SWIZZLE_128B_ATOM, (BLK_M, BLK_K)) + B_layout = tma_shared_layout(b_type, SwizzleMode.SWIZZLE_128B_ATOM, (BLK_N, BLK_K)) + + @T.prim_func + def kernel( + A: T.Buffer((M, K), a_type), + B: T.Buffer((N, K), b_type), + D: T.Buffer((M, N), d_type), + ): + T.device_entry() + bx, by = T.cta_id([M // BLK_M, N // BLK_N]) # still one output tile (M=N=128) + wg_id = T.warpgroup_id([1]) + warp_id = T.warp_id_in_wg([4]) + lane_id = T.lane_id([32]) + + pool = T.SMEMPool() + tmem_addr = pool.alloc((1,), "uint32") + mma_bar = pool.alloc((1,), "uint64", align=8) + pool.move_base_to(1024) + Asmem = pool.alloc((BLK_M, BLK_K), a_type, layout=A_layout) + Bsmem = pool.alloc((BLK_N, BLK_K), b_type, layout=B_layout) + pool.commit() + + if warp_id == 0: + if lane_id == 0: + T.ptx.mbarrier.init(mma_bar.ptr_to([0]), 1) + T.ptx.tcgen05.alloc(T.address_of(tmem_addr), n_cols=512, cta_group=1) + + T.ptx.fence.proxy_async("shared::cta") + T.ptx.fence.mbarrier_init() + T.cuda.cta_sync() + + tmem = T.decl_buffer( + (128, 512), "float32", scope="tmem", allocated_addr=tmem_addr[0], + layout=TileLayout(S[(128, 512) : (1@TLane, 1@TCol)])) + + phase_mma: T.int32 = 0 + m_st = T.meta_var(bx * BLK_M) + n_st = T.meta_var(by * BLK_N) + + # === K-loop: iterate over K in chunks of BLK_K === + for i in T.serial(K_TILES): # serial device loop (keeps the full-K A/B parameters correctly shaped) + # Load the i-th K chunk + Tx.cta.copy(Asmem[:, :], A[:, i*BLK_K:(i+1)*BLK_K]) + Tx.cta.copy(Bsmem[:, :], B[:, i*BLK_K:(i+1)*BLK_K]) + + T.cuda.cta_sync() + + # MMA: accum=False for first tile, True for rest + if warp_id == 0: + if T.ptx.elect_sync(): + Tx.gemm_async(tmem[:, :BLK_N], Asmem[:, :], Bsmem[:, :], + accum=(i != 0), dispatch="tcgen05", cta_group=1) + T.ptx.tcgen05.commit(mma_bar.ptr_to([0]), cta_group=1) + + # Wait for MMA, then flip phase + T.ptx.mbarrier.try_wait(mma_bar.ptr_to([0]), phase_mma) + phase_mma ^= 1 + + # === Writeback (same as Step 1) === + Dreg = T.alloc_local((BLK_N,), acc_type) + Dreg_f16 = T.alloc_local((BLK_N,), d_type) + Dreg_wg = Dreg.view(128, BLK_N, + layout=TileLayout(S[(128, BLK_N) : (1@tid_in_wg, 1)])) + + Tx.wg.copy_async(Dreg_wg[:, :], tmem[:, :BLK_N]) + T.ptx.tcgen05.wait.ld() + + Tx.cast(Dreg_f16[:], Dreg[:]) + m_thr = T.meta_var(m_st + warp_id * 32 + lane_id) + Tx.copy(D[m_thr, n_st : n_st + BLK_N], Dreg_f16[:]) + + T.cuda.cta_sync() + if warp_id == 0: + T.ptx.tcgen05.relinquish_alloc_permit(cta_group=1) + T.ptx.tcgen05.dealloc(tmem_addr[0], n_cols=512, cta_group=1) + + return kernel +``` + +--- + (chap_spatial_tiling)= -# 构建 Tiled GEMM +## Step 3：Spatial Tiling（Multi-CTA） + +K-loop 处理了 contraction dimension，但 M 和 N 仍然固定在单个 128 x 128 tile 上。真实输出远大于一个 tile，所以基础 kernel 的最后一块是用多个 tile 同时覆盖 M 和 N。Step 3 会 launch 一个 2D CTA grid，每个 output tile 一个 CTA，让 GPU 并行计算所有 tile。示例使用 M=N=K=256，对应 2x2 tile grid，刚好能让 indexing 不再平凡，又不至于把重点埋掉。 + +> **本 step 改变的内容：Scope** +> - Scope：一个 2D CTA grid，每个 CTA 拥有一个 128 x 128 output tile。 +> - Layout：不变；在每个 CTA 内部，这仍然是 Step 2 的同一条 SMEM/TMEM/register 路径。 +> - Dispatch：不变。 + +### Grid Mapping + +Grid shape 直接来自 tiling：每个 128 x 128 output tile 一个 CTA，所以总共需要 `[M // BLK_M, N // BLK_N]` 个 CTA。相对 Step 2，唯一真正新增的工作，是让每个 CTA 知道矩阵中哪一片是*自己*要计算的 slice。 + +CTA `(bx, by)` 拥有这个 output region： + +```text +D[bx * BLK_M : (bx + 1) * BLK_M, + by * BLK_N : (by + 1) * BLK_N] +``` + +为了产生它，该 CTA 的 K-loop 会反复加载自己 A row band 和 B column band 对应的 K-slice： + +```text +A[bx * BLK_M : (bx + 1) * BLK_M, k : k + BLK_K] +B[by * BLK_N : (by + 1) * BLK_N, k : k + BLK_K] +``` + +Indexing 直接来自 `D = A @ B.T` 约定：`bx` 选择 A 和 D 的行，而 `by` 选择 B 的行；transpose 应用之后，这些 B 行会变成 D 的列。 + +每个 CTA 一个 tile 是最简单可行的映射，但它也浪费。Row 中的每个 CTA 都会从 GMEM 重新加载同样的 A tile，column 中的每个 CTA 都会重新加载同样的 B tile，因此没有复用邻近 CTA 已经拉进来的数据。我们暂时保留这个浪费；persistent scheduling（{ref}`chap_gemm_async` 中的 Step 6）会回到这个问题，并让这些共享 operand 在 L2 中保持 hot。 + +**Try with your agent**：设 `M=N=K=256`、`BLK_M=BLK_N=128`、`BLK_K=64`，让它 trace CTA `(1, 0)` 和 CTA `(0, 1)`。对每个 CTA，列出 `m_st`、`n_st`、每次 K iteration 加载的 A/B slice，以及写入的 D region。哪些 B row 因为 kernel 计算 `D = A @ B.T` 而变成 D column？ + +### 完整 Kernel + +这个 kernel 再次从 Step 2 发展而来，这次只有两个变化：grid shape 和 per-CTA offset。内部 K-loop 和 writeback 不变。Import 相同： + +```python + +import tvm +from tvm.script import tirx as T +from tvm.script.tirx import tile as Tx +from tvm.tirx.cuda.operator.tile_primitive.tma_utils import tma_shared_layout, SwizzleMode +from tvm.tirx.layout import TileLayout, S, TLane, TCol, tid_in_wg +``` + +Grid 从 `[1, 1]` 变成 `[M // BLK_M, N // BLK_N]`，load 和 store 现在都会加上 CTA 自己的 `m_st` 和 `n_st` offset： + +```python +def hgemm_v3(M, N, K): + a_type = tvm.DataType("float16") + b_type = tvm.DataType("float16") + d_type = tvm.DataType("float16") + acc_type = tvm.DataType("float32") + + BLK_M, BLK_N, BLK_K = 128, 128, 64 + K_TILES = K // BLK_K + + A_layout = tma_shared_layout(a_type, SwizzleMode.SWIZZLE_128B_ATOM, (BLK_M, BLK_K)) + B_layout = tma_shared_layout(b_type, SwizzleMode.SWIZZLE_128B_ATOM, (BLK_N, BLK_K)) + + @T.prim_func + def kernel( + A: T.Buffer((M, K), a_type), + B: T.Buffer((N, K), b_type), + D: T.Buffer((M, N), d_type), + ): + T.device_entry() + # 2D grid: one CTA per 128x128 output tile + bx, by = T.cta_id([M // BLK_M, N // BLK_N]) + wg_id = T.warpgroup_id([1]) + warp_id = T.warp_id_in_wg([4]) + lane_id = T.lane_id([32]) + + pool = T.SMEMPool() + tmem_addr = pool.alloc((1,), "uint32") + mma_bar = pool.alloc((1,), "uint64", align=8) + pool.move_base_to(1024) + Asmem = pool.alloc((BLK_M, BLK_K), a_type, layout=A_layout) + Bsmem = pool.alloc((BLK_N, BLK_K), b_type, layout=B_layout) + pool.commit() + + if warp_id == 0: + if lane_id == 0: + T.ptx.mbarrier.init(mma_bar.ptr_to([0]), 1) + T.ptx.tcgen05.alloc(T.address_of(tmem_addr), n_cols=512, cta_group=1) + + T.ptx.fence.proxy_async("shared::cta") + T.ptx.fence.mbarrier_init() + T.cuda.cta_sync() + + tmem = T.decl_buffer( + (128, 512), "float32", scope="tmem", allocated_addr=tmem_addr[0], + layout=TileLayout(S[(128, 512) : (1@TLane, 1@TCol)])) + + phase_mma: T.int32 = 0 + + # Per-CTA tile offsets + m_st = T.meta_var(bx * BLK_M) + n_st = T.meta_var(by * BLK_N) + + # K-loop with offset A and B slices + for i in T.serial(K_TILES): # serial device loop (keeps the full-K A/B parameters correctly shaped) + Tx.cta.copy(Asmem[:, :], A[m_st:m_st+BLK_M, i*BLK_K:(i+1)*BLK_K]) + Tx.cta.copy(Bsmem[:, :], B[n_st:n_st+BLK_N, i*BLK_K:(i+1)*BLK_K]) + + T.cuda.cta_sync() + + if warp_id == 0: + if T.ptx.elect_sync(): + Tx.gemm_async(tmem[:, :BLK_N], Asmem[:, :], Bsmem[:, :], + accum=(i != 0), dispatch="tcgen05", cta_group=1) + T.ptx.tcgen05.commit(mma_bar.ptr_to([0]), cta_group=1) + + T.ptx.mbarrier.try_wait(mma_bar.ptr_to([0]), phase_mma) + phase_mma ^= 1 + + # Writeback to the correct output tile + Dreg = T.alloc_local((BLK_N,), acc_type) + Dreg_f16 = T.alloc_local((BLK_N,), d_type) + Dreg_wg = Dreg.view(128, BLK_N, + layout=TileLayout(S[(128, BLK_N) : (1@tid_in_wg, 1)])) + + Tx.wg.copy_async(Dreg_wg[:, :], tmem[:, :BLK_N]) + T.ptx.tcgen05.wait.ld() + + Tx.cast(Dreg_f16[:], Dreg[:]) + m_thr = T.meta_var(m_st + warp_id * 32 + lane_id) + Tx.copy(D[m_thr, n_st:n_st+BLK_N], Dreg_f16[:]) + + T.cuda.cta_sync() + if warp_id == 0: + T.ptx.tcgen05.relinquish_alloc_permit(cta_group=1) + T.ptx.tcgen05.dealloc(tmem_addr[0], n_cols=512, cta_group=1) + + return kernel +``` -> 翻译状态：待翻译。对应英文章节：`chapter_gemm_basics/index.md`。 +## 练习 -本页用于放置基础 tiled GEMM 构建过程的中文翻译。 +1. 在 Steps 1-3 中，`Tx.copy` 会在 MMA 之前把 A 和 B tile 移入 SMEM。为什么 kernel 需要在 `Tx.gemm_async` 读取这些 SMEM tile 前执行 `T.cuda.cta_sync()`？ +2. 在 Step 2 中，如果从 K-loop 中移除 `phase_mma ^= 1` 会发生什么？Kernel 是否会等待每一次 MMA，还是后续 wait 可能过早通过？ +3. 对于 M=N=4096 且 BLK_M=BLK_N=128 的情况，Step 3 会 launch 多少个 CTA？哪些 operand tile 在逻辑上会被相邻 CTA 复用？Step 3 是否利用了这种复用？ diff --git a/zh/chapter_intro_tirx/index.md b/zh/chapter_intro_tirx/index.md index a29ed160..4c6e3f96 100644 --- a/zh/chapter_intro_tirx/index.md +++ b/zh/chapter_intro_tirx/index.md @@ -1,10 +1,222 @@ ---- -orphan: true ---- - (chap_tirx_primer)= # TIRx 入门 -> 翻译状态：待翻译。对应英文章节：`chapter_intro_tirx/index.md`。 +:::{admonition} 概览 +:class: overview + +- TIRx 是一个用于在 IR 层编写 GPU kernel 的 Python DSL：你会直接命名硬件，但通过结构化 IR 来表达。 +- 每个 tile 操作都由三个设计要素控制：*scope*（哪些线程执行）、*layout*（tile 位于哪里）和 *dispatch*（走哪条硬件路径）。 +- 一个可运行的 single-MMA GEMM 会同时展示这三者；本书后续内容就是把这些设计要素放大到真实规模。 +::: + +:::{admonition} 运行示例 +:class: note + +这些示例需要 Blackwell GPU（`sm_100a`，例如 B200）。TIRx 编译器作为 Apache TVM wheel 中的 `tvm.tirx` 模块发布；请和 CUDA 版本的 PyTorch 一起安装： + +```bash +pip install apache-tvm +``` + +用 `python -c "import tvm, tvm.tirx; print(tvm.__version__)"` 确认它可以 import。同样的环境可以运行本书所有可运行示例。 +::: + +第一部分解释了硬件是什么。要让硬件真正计算，我们还需要一种编程方式。 + +我们可以直接写 CUDA 或 PTX，很多快速 kernel 也确实是这样写的。问题在于，真正决定 kernel 行为的决策在那里很难看清：哪些线程执行某个操作、每个数据 tile 位于哪里、以及由哪条硬件路径执行。这些选择会埋在 intrinsic 参数、地址计算和约定之中。 + +TIRx（Tensor IR neXt）是一个 Python DSL，它把这三个决策明确提升出来：**scope**（哪些线程执行操作）、**layout**（operand tile 位于哪里）和 **dispatch**（使用哪条硬件路径执行）。它仍然直接命名硬件概念，包括线程、shared memory、tensor memory、barrier 和 `tcgen05` MMA。区别是，这些选择现在变成结构化 IR，编译器可以 lower、检查和调度。 + +我们不会先抽象地介绍这些概念，而是从一个完整 kernel 开始：最小 single-MMA GEMM。我们先让它跑起来，然后逐行读回去，看 scope、layout 和 dispatch 分别如何塑造它，以及 kernel 如何被编译。Kernel 依赖的 tensor layout 模型会在 {ref}`chap_tirx_layout_api` 中单独展开，完整语言特性集合在 {ref}`chap_language_reference` 中介绍；这里我们聚焦这一个 kernel 和三个设计要素。 + +## 第一个 Kernel：Single-MMA GEMM + +我们承诺的 kernel 是一个最小 GEMM，删减到仍然能使用 Tensor Core 的最小版本。它计算 `D = A B^T` 的单个 128 x 128 output tile，K = 64。整个计算从头到尾被表达成一次 `Tx.gemm_async` tile operation。（这一个 tile operation 并不映射到单条硬件指令：因为硬件 MMA 的 K atom 是 16，K=64 的 tile 会 lower 成沿 K 前进的一小段 `tcgen05.mma` 指令序列。DSL 的重点正是我们写 tile，而不是手写这段序列。）在这个操作周围，kernel 做常规工作：分配 shared memory（SMEM）和 tensor memory（TMEM），把 A 和 B 从 global copy 到 shared memory，发起 tile MMA 并把结果写入 TMEM accumulator，再把 accumulator 通过寄存器读回并 store 结果。虽然它很小，但这个 kernel 就是 {ref}`chap_gemm_basics` 中 GEMM 阶梯的 Step 1，那里会完整讲解它。 + +每个 TIRx kernel 都从同一组 import 开始，所以值得先看一次： + +```python + +import tvm +from tvm.script import tirx as T +from tvm.script.tirx import tile as Tx +from tvm.tirx.cuda.operator.tile_primitive.tma_utils import tma_shared_layout, SwizzleMode +from tvm.tirx.layout import TileLayout, S, TLane, TCol, tid_in_wg +``` + +我们把 kernel 包在一个小 builder `hgemm_v1(M, N, K)` 中，它接受问题 shape 并返回一个 `PrimFunc`。对于我们选择的 shape，`M=N=128, K=64`，launch 中刚好只有一个 output tile，这让第一个版本足够简单，可以一次读完： + +```python +def hgemm_v1(M, N, K): + a_type = tvm.DataType("float16") + b_type = tvm.DataType("float16") + d_type = tvm.DataType("float16") + acc_type = tvm.DataType("float32") + + BLK_M, BLK_N, BLK_K = 128, 128, 64 + # MMA_M/MMA_N/MMA_K document the underlying hardware MMA tile; they are not + # passed to gemm_async (which derives the MMA shape from the operand and + # accumulator tiles), so the later steps omit them. + MMA_M, MMA_N, MMA_K = 128, 128, 16 + + A_layout = tma_shared_layout(a_type, SwizzleMode.SWIZZLE_128B_ATOM, (BLK_M, BLK_K)) + B_layout = tma_shared_layout(b_type, SwizzleMode.SWIZZLE_128B_ATOM, (BLK_N, BLK_K)) + + @T.prim_func + def kernel( + A: T.Buffer((M, K), a_type), + B: T.Buffer((N, K), b_type), + D: T.Buffer((M, N), d_type), + ): + T.device_entry() + # Step 1 is a single-tile kernel: M = BLK_M and N = BLK_N, so the grid + # is 1x1. Starting with a 1x1 grid keeps the per-CTA tile offsets + # (m_st, n_st) trivially zero; Steps 3+ generalise this to larger M / N. + bx, by = T.cta_id([M // BLK_M, N // BLK_N]) + wg_id = T.warpgroup_id([1]) # single warpgroup, so wg_id is always 0 (unused below) + warp_id = T.warp_id_in_wg([4]) + lane_id = T.lane_id([32]) + + # --- SMEM allocation --- + pool = T.SMEMPool() + tmem_addr = pool.alloc((1,), "uint32") + mma_bar = pool.alloc((1,), "uint64", align=8) + pool.move_base_to(1024) + Asmem = pool.alloc((BLK_M, BLK_K), a_type, layout=A_layout) + Bsmem = pool.alloc((BLK_N, BLK_K), b_type, layout=B_layout) + pool.commit() + + # --- Barrier + TMEM init (warp 0 only) --- + if warp_id == 0: + if lane_id == 0: + T.ptx.mbarrier.init(mma_bar.ptr_to([0]), 1) + T.ptx.tcgen05.alloc(T.address_of(tmem_addr), n_cols=512, cta_group=1) + + T.ptx.fence.proxy_async("shared::cta") + T.ptx.fence.mbarrier_init() + T.cuda.cta_sync() + + tmem = T.decl_buffer( + (128, 512), "float32", scope="tmem", allocated_addr=tmem_addr[0], + layout=TileLayout(S[(128, 512) : (1@TLane, 1@TCol)]) + ) + + m_st = T.meta_var(bx * BLK_M) + n_st = T.meta_var(by * BLK_N) + phase_mma: T.int32 = 0 + + # --- Load: all threads copy global -> shared (synchronous). + # With M=BLK_M and N=BLK_N the slices below cover the full matrices; + # the slice form is kept so the diff to Step 3 (multi-tile) is minimal. + Tx.cta.copy(Asmem[:, :], A[m_st:m_st + BLK_M, :]) + Tx.cta.copy(Bsmem[:, :], B[n_st:n_st + BLK_N, :]) + T.cuda.cta_sync() + + # --- Compute: single elected thread issues MMA --- + if warp_id == 0: + if T.ptx.elect_sync(): + Tx.gemm_async( + tmem[:, :BLK_N], Asmem[:, :], Bsmem[:, :], + accum=False, dispatch="tcgen05", cta_group=1 + ) + T.ptx.tcgen05.commit(mma_bar.ptr_to([0]), cta_group=1) + + T.ptx.mbarrier.try_wait(mma_bar.ptr_to([0]), phase_mma) + + # --- Writeback: TMEM -> RF -> GMEM --- + Dreg = T.alloc_local((BLK_N,), acc_type) + Dreg_f16 = T.alloc_local((BLK_N,), d_type) + Dreg_wg = Dreg.view(128, BLK_N, + layout=TileLayout(S[(128, BLK_N) : (1@tid_in_wg, 1)])) + Tx.wg.copy_async(Dreg_wg[:, :], tmem[:, :BLK_N]) + T.ptx.tcgen05.wait.ld() + Tx.cast(Dreg_f16[:], Dreg[:]) + m_thr = T.meta_var(m_st + warp_id * 32 + lane_id) + Tx.copy(D[m_thr, n_st : n_st + BLK_N], Dreg_f16[:]) + + # --- Deallocate TMEM --- + T.cuda.cta_sync() + if warp_id == 0: + T.ptx.tcgen05.relinquish_alloc_permit(cta_group=1) + T.ptx.tcgen05.dealloc(tmem_addr[0], n_cols=512, cta_group=1) + + return kernel +``` + +在读这个 kernel 之前，我们先确认它能工作。我们编译它，并用 torch reference 检查输出。这里不需要显式写出具体架构：arch（例如 `sm_100a`）会从设备自动检测，所以 target `"cuda"` 就足够了，`tir_pipeline="tirx"` 选择 TIRx lowering pipeline。编译完成后，`ex.mod(...)` 可以直接接受 torch tensor，中间不需要手动转换。 + +```python +import torch + +target = tvm.target.Target("cuda") +device = torch.device('cuda') # gpu(0) + +M, N, K = 128, 128, 64 +kernel = hgemm_v1(M, N, K) +with target: + ex = tvm.compile(tvm.IRModule({"main": kernel}), target=target, tir_pipeline="tirx") + +torch.cuda.empty_cache() +torch.cuda.synchronize() +A_tensor = torch.randn(M, K, dtype=torch.float16, device=device) +B_tensor = torch.randn(N, K, dtype=torch.float16, device=device) +D_tensor = torch.zeros(M, N, dtype=torch.float16, device=device) + +# ex.mod(...) takes torch tensors directly, the same call form used in every chapter. +ex.mod(A_tensor, B_tensor, D_tensor) + +D_ref = (A_tensor.float() @ B_tensor.float().T).half() +max_err = float((D_tensor - D_ref).abs().max()) +print(f"Max error vs torch reference: {max_err:.6f}") +torch.testing.assert_close(D_tensor, D_ref, rtol=2e-2, atol=1e-2) +print("PASS") +``` + +## Scope、Layout、Dispatch + +现在 kernel 已经能跑，我们可以回过头来读它，问每一行到底决定了什么。从这个角度看，整个 kernel 是围绕三个设计要素的一组选择。里面的每个操作都回答同样三个问题：*谁*执行它、它的数据*在哪里*、它*如何*执行；这三个答案就是 scope、layout 和 dispatch。本节会依次讨论这些设计要素；下面的交互 demo 可以看到每个设计要素控制了 kernel 的哪些行。 + +```{raw} html + +``` +*交互图：点击 Scope / Layout / Dispatch，高亮 kernel 中由每个设计要素控制的行。* + +使用 demo 时，请关注三个问题： + +- **Scope：谁执行这个操作？** `Tx.cta.copy(...)` 是 CTA-scoped，因此全部 128 个线程都会帮助完成 GMEM -> SMEM copy。`Tx.gemm_async(...)` 由一个 elected thread 发起，因为 lowering 后的每条 `tcgen05.mma` 指令本身已经是一次 cooperative MMA launch。`Tx.wg.copy_async(...)` 是 warpgroup-scoped，因此 warpgroup 的 128 个线程会按行拆分 TMEM readback。 +- **Layout：每个 tile 位于哪里？** A 和 B 使用 `tcgen05.mma` 期望的 swizzled SMEM layout。Accumulator 位于 TMEM 中，使用 `TLane`/`TCol` 布局。Register readback view 把 row 映射到 `tid_in_wg`，因此每个 warpgroup thread 拥有一个 row fragment。 +- **Dispatch：哪条硬件路径执行它？** `Tx.gemm_async(..., dispatch="tcgen05", ...)` 选择 Blackwell Tensor Core 路径。Copy 操作也有 dispatch 选择：第一个 kernel 使用普通 thread copy，后续 GEMM step 会把这些 copy 换成 TMA，而不改变周围的 scope 或 layout。 + +**Try with your agent**：从第一个 kernel 里挑三行：一个 copy、一个 MMA、一个 TMEM readback。让它用 scope、layout 和 dispatch 标注每一行，然后检查答案是否与 guard、buffer layout 和 `dispatch=` 参数一致。 + +## 编译如何工作 + +我们已经在上面编译过 kernel 来测试它；现在稍微近距离看一下这个步骤做了什么。流程很短：把 `PrimFunc` 包进一个 `IRModule`，再交给 `tvm.compile(mod, target=..., tir_pipeline="tirx")`。这会运行 TIRx lowering pipeline，并返回一个可以直接调用的 `Executable`。 + +```python +target = tvm.target.Target("cuda") +ex = tvm.compile(tvm.IRModule({"main": kernel}), target=target, tir_pipeline="tirx") +``` + +至少大致知道 `tir_pipeline="tirx"` 会触发什么，是有帮助的。Pipeline 的核心 pass 是 `LowerTIRx`，它会根据每个 tile primitive 的 scope / layout / dispatch contract 进行解析：这里正是我们刚才讨论的三个设计要素真正兑换成指令的地方。之后，常规 host/device split 和 finalize step 会产生可 launch 的 module。如果愿意，也可以在 `with target:` block 内编译，这样 kernel 可以获取外层 target context。 + +这个流程的一个好处是没有东西对你隐藏：结果可以在两个层级上检查。你可以用 `.show()` 或 `.script()` 读取 IR 本身，也可以从编译后的 module 中直接查看编译器最终生成的 CUDA C。 + +```python +kernel.show() # pretty-print the TIRx (TVMScript) +print(kernel.script()) # ... the same, as a string + +# the generated CUDA C source, from the compiled Executable: +print(ex.mod.imports[0].inspect_source()) +``` + +这里只是一个概览。完整 lowering 过程，包括所有 pass、tile-primitive dispatch 如何解析，以及 host/device split 如何完成，请见 {ref}`chap_arch`。 + +## 下一步 + +一个 kernel 已经足够让我们认识 scope、layout 和 dispatch，并看到它们如何被编译和运行。这三个设计要素以及 kernel 本身，分别通向后续章节： -本页用于放置 TIRx 入门 walkthrough 的中文翻译。 +- {ref}`chap_tirx_layout_api`：tensor layout 模型（`TileLayout`、命名轴、swizzle），上面 operand 和 accumulator 的 placement 都建立在它之上。如果 layout 这个设计要素最让你困惑，可以从这里继续。 +- {ref}`chap_language_reference`：完整语言特性集合，包括 parser utility、data type、buffer 和 memory、control flow，以及 thread synchronization；当你需要完整词汇表而不是导览时可以查这里。 +- {ref}`chap_gemm_basics`：这个 kernel 作为 GEMM 优化路径的 Step 1，并进一步加入 K-loop accumulation、spatial tiling、TMA 和 warp specialization。如果你想看同样三个设计要素如何扩展到真实 kernel，这是自然的下一站。 diff --git a/zh/chapter_layout_generations/index.md b/zh/chapter_layout_generations/index.md index 359d216e..5367939a 100644 --- a/zh/chapter_layout_generations/index.md +++ b/zh/chapter_layout_generations/index.md @@ -1,79 +1,288 @@ (chap_layout_generations)= # 跨 GPU 世代的 Tensor Core 操作数布局 - - :::{admonition} 概览 :class: overview -- TODO：翻译本章 overview 第一条。 -- TODO：翻译本章 overview 第二条。 -- TODO：翻译本章 overview 第三条。 +- 从 Ampere 到 Hopper 再到 Blackwell，Tensor Core 执行的高层操作仍然相同：`D = A B + C`。 +- 每一代发生变化的是：操作数如何到达 Tensor Core，支持哪些 tile shape 与 dtype，以及 accumulator 放在哪里。 +- Ampere 使用 warp-level register fragment。Shared memory tile 通过 `ldmatrix` 加载到 fragment 中，accumulator 保持在寄存器里。 +- Hopper 让 `wgmma` 通过 matrix descriptor 直接从 shared memory 读取操作数。Descriptor 会命名 Tensor Core 期望的 shared-memory swizzle format。 +- Blackwell 保留 shared-memory operand path，但把 accumulator 移到 TMEM 中。Block-scaled MMA 还会通过 TMEM staging scale factor。 +- 两个内存约束跨所有世代都始终存在：global memory coalescing 和 shared memory bank conflict。 ::: -TODO：翻译导言部分。 +从远处看，Tensor Core 操作似乎一直很稳定。它把 A 和 B 的 tile 相乘，加上 accumulator C，并产生 D。这个形式从 Volta 起就是这样。 + +但围绕这个操作的细节并没有保持不变。在某一代上很快的 kernel，到了下一代可能会变慢。使用错误布局的 kernel 也可能算出错误答案，即使逻辑数学仍然写作 `D = A B + C`。原因是 Tensor Core 消费的不是抽象矩阵，而是非常具体的硬件布局中的操作数。 + +本章会沿着三代硬件追踪这份布局契约。Ampere 通过 warp-level register fragment 暴露 Tensor Core。Hopper 把输入操作数移到 shared memory descriptor 上。Blackwell 保留 shared memory 操作数，但把 accumulator 移到 TMEM 中。操作仍然是 matrix-multiply-accumulate，但进入 Tensor Core 和离开 Tensor Core 的路径每一代都在改变。 + +{ref}`数据布局 ` 一章中的布局记号，是我们描述这些契约的语言。Blackwell TMEM 的细节会在 {ref}`chap_tmem` 中单独介绍。 ## 两个始终存在的约束 -TODO：翻译 “Two Constraints That Never Went Away” 小节。 +在 Tensor Core 参与之前，两个普通内存约束就已经在塑造 GPU kernel 的布局。 + +第一个是 global memory coalescing。当一个 warp 的 32 个 lane 发起 global memory load 时，内存系统希望这些地址落在少量连续且对齐的 memory segment 中。如果地址分散，warp load 就会变成多次 memory transaction。同样的逻辑数据搬运会消耗更多带宽和更多时间。 + +第二个是 shared memory bank conflict。Shared memory 被分成 32 个 bank。如果 warp 中多个 lane 访问映射到同一个 bank 的不同地址，这些访问无法同时服务，硬件会把它们串行化。因此，一个看起来只是扁平 shared memory array 的布局，可能因为 bank pattern 而变慢。 + +Swizzling 是修复 shared memory 侧问题的常见方法。逻辑 tile 保持不变，但物理地址映射被重排，让访问 pattern 分散到不同 bank 上，而不是堆到同一个 bank。 + +这两个约束即使在完全不用 Tensor Core 的 kernel 中也存在。Tensor Core kernel 会再加上第三个约束：操作数必须按照 Tensor Core 指令自身期望的布局排列。本章剩下的内容，就是看这个第三个约束如何在 Ampere、Hopper 和 Blackwell 之间变化。 ## Ampere：Warp Lane 上的寄存器 Fragment -TODO：翻译 “Ampere: Register Fragments over Warp Lanes” 小节。 +在 Ampere 级别的 GPU 上，主要 Tensor Core 指令是 warp-level 的 `mma.sync.aligned.m16n8k*` 系列。最重要的事实是这条指令从哪里读写数据：寄存器。 + +A、B，以及 C 或 D accumulator，都是分布在 warp 的 32 个 lane 上的 per-thread register fragment。Shared memory 只是 staging area。在 MMA 运行之前，operand tile 必须从 shared memory 移到这条指令期望的精确 register fragment 布局中。 + +数据路径如下： + +```text +SMEM to registers with ldmatrix +registers to registers with mma.sync +registers back to SMEM with ordinary stores +``` + +Ampere 的大部分布局故事都来自这条路径。Kernel 必须先把 tile 以一种能高效加载的形式存入 shared memory，然后用 `ldmatrix` 产生 `mma.sync` 需要的 register fragment。 ## Ampere Tensor Core 期望的输入 -TODO：翻译 “What the Ampere Tensor Core Expects” 小节。 +Ampere Tensor Core 读取由 8 by 8 subtile unit 构成的 register fragment。这些 unit 正是 `ldmatrix` 加载的单位，也是 MMA 消费的单位。 + +以带 fp16 或 bf16 输入、fp32 accumulation 的 `mma.m16n8k16` 为具体例子。Accumulator tile 的形状是 `16 by 8`，并以固定 pattern 分布到 32 个 lane 上。 + +对于 C 或 D accumulator，lane `l` 持有的行是： + +```text +l / 4 +l / 4 + 8 +``` + +列是： + +```text +2 * (l % 4) +2 * (l % 4) + 1 +``` + +因此每个 lane 拥有四个 fp32 accumulator 值：来自两个 8-row half 的两行，与两个相邻列交叉组合。连续四个 lane 覆盖某一行的八个列。 + +A operand 使用同样的 M-side row carve。K 维度分散在 `l % 4` 以及 lane 持有的寄存器中。对于 fp16 或 bf16，每个 32-bit 寄存器会打包两个 K 值。 + +B operand 使用匹配的 K placement，并把 N 侧分散到 lane group 和寄存器中。 -## `ldmatrix`：从共享内存到寄存器 Fragment +具体细节会随 instruction shape 和 dtype 而变，但原则固定：Tensor Core 期望一个特定的 per-lane register fragment。如果值没有按这个 pattern 放到这些寄存器里，指令就会把错误元素相乘。 -TODO：翻译 “ldmatrix: Shared Memory to Register Fragment” 小节。 +在布局记号中，m8n8 fragment 就是用命名 lane 轴写出的那类 pattern，例如： -![TODO：翻译 ldmatrix 图注](../../img/ldstmatrix.svg) +```text +S[(8, 4, 2) : (4@laneid, 1@laneid, 1@m)] +``` + +两个 `laneid` 迭代器共同描述 row 和 column piece 如何分散到 lane 上，而最后的 `m` component 描述 per-lane register slot。 + +## `ldmatrix`：从 Shared Memory 到寄存器 Fragment + +`ldmatrix` 是 Ampere 上连接 shared memory 与 Tensor Core register fragment 的指令。它是 warp-collective load。一条指令会把一个或多个 8 by 8 的 16-bit 矩阵从 shared memory 移到 `mma.sync` 期望的分布式 register layout 中。 + +指令形式是： + +```text +ldmatrix.sync.aligned.m8n8.x1.shared.b16 +ldmatrix.sync.aligned.m8n8.x2.shared.b16 +ldmatrix.sync.aligned.m8n8.x4.shared.b16 +``` + +并且可以带一个可选的 `.trans` qualifier。 + +`.x1`、`.x2` 和 `.x4` 形式分别加载一个、两个或四个 8 by 8 矩阵。Row base address 由 lane 提供。对于矩阵 `m` 和行 `r`，base address 来自 lane `m * 8 + r`。这意味着 `.x1` 使用 lane 0 到 7 作为 row address，`.x2` 使用 lane 0 到 15，`.x4` 使用 lane 0 到 31。 + +结果会直接落入 MMA fragment。对于基本的 8 by 8 情况，lane `l` 会收到 Tensor Core 期望的 row/column pair。普通 per-lane `ld.shared` 指令循环必须手动复现这种 scatter。`ldmatrix` 则把 shared-memory 到 fragment 的重排作为一条 warp-collective 指令完成。 + +`.trans` 形式会在加载每个 8 by 8 矩阵时做 transpose。当 operand 的存储方向与 MMA 指令期望方向相反时，会使用这个形式。 + +![ldmatrix loads an 8x8 shared memory tile into the warp register fragment; the reverse direction on Ampere uses ordinary stores, and a dedicated stmatrix instruction appears later on Hopper](../../img/ldstmatrix.svg) ## 写回 Ampere Fragment -TODO：翻译 “Writing the Ampere Fragment Back” 小节。 +`mma.sync` 完成之后，accumulator 仍然是 register fragment。Epilogue 必须把这个 fragment 移出去。 + +Ampere 上没有专门的 `ldmatrix` 反向指令。Kernel 使用普通 per-thread store，有时在 store 之前配合 warp shuffle 或局部重排，把 accumulator 写入 shared memory 或 global memory 中的有用布局。 + +这让 Ampere 模型保持简单，但也把很多布局工作暴露给 kernel。输入侧用 `ldmatrix` 创建 fragment。Compute 指令读写 register fragment。输出侧由这些 fragment 上的普通 store 处理。 ## Ampere 上的 Swizzle -TODO：翻译 “Swizzle on Ampere” 小节。 +Ampere kernel 已经需要 shared memory swizzle。原因是 shared memory tile 通常以一种访问 pattern 写入，又以另一种访问 pattern 读取。 + +假设一个 tile 沿 row 从 global memory 填充。Row-major layout 让这种写入 coalesced 且 bank friendly。但 `ldmatrix` 稍后可能用一种等价于沿 column 或跨 8 by 8 subtile 行走的 pattern 读取这个 tile。使用普通 row-major layout 时，这些读取可能堆到同一个 shared memory bank 上。 + +对于一个简单的 `(8, 64)` float16 tile，一行是： + +```text +64 * 2 bytes = 128 bytes +``` + +这刚好是一整条 shared memory bank line。沿固定 column 向下走时，每行前进 128 byte，因此 bank index 会重复。八行可能全部落到同一个 bank 上，产生 8-way conflict。 + +改成普通 column-major layout 并不能完整解决问题。它通常只是把 conflict 移到另一次访问上。Row write 变差，而 column-style read 变好。 + +XOR swizzle 通过让物理 column 依赖 row 来解决这个问题。一个简单版本是： + +```text +physical_col = logical_col xor row +``` + +逻辑 tile 不变。Shared memory 中的物理 placement 被重排，使 row-style write 和 Tensor Core read pattern 都能避免 bank conflict。 + +在 Ampere 上，这种 swizzle 通常通过手写 shared memory index math 表达。后续世代会把它变成硬件 engine 使用的 descriptor format 的一部分。 -![TODO：翻译 swizzle conflict 图注](../../img/swizzle_conflict.svg) +![On a plain row-major tile a row write spreads across banks while a column read collides on one bank; the XOR swizzle scatters the column read across banks without giving up the coalesced row write](../../img/swizzle_conflict.svg) -## Hopper：`wgmma`、共享内存 Descriptor 和 Swizzle Format +## Hopper：`wgmma`、Shared Memory Descriptor 和 Swizzle Format -TODO：翻译 “Hopper: wgmma, Shared Memory Descriptors, and Swizzle Formats” 小节。 +Hopper 改变了 Tensor Core 路径的输入侧。Hopper `wgmma` 不再要求每个 operand 都通过 `ldmatrix` 加载到寄存器，而是可以直接从 shared memory 读取 operand。 + +B operand 从 shared memory matrix descriptor 中读取。A operand 可以从 shared memory descriptor 读取，也可以从寄存器读取，对应 `.ss` 和 `.rs` 形式。 + +这移除了 SMEM-sourced operand 上显式的 `ldmatrix` 步骤。但它没有移除布局要求。Tensor Core 仍然期望 operand 以精确的 shared memory format 存放。区别在于，现在这个 format 通过 matrix descriptor 告诉硬件。 ## Hopper Tensor Core 期望的输入 -TODO：翻译 “What the Hopper Tensor Core Expects” 小节。 +Hopper shared memory matrix descriptor 是 shared memory 中矩阵 tile 的紧凑描述。它告诉 `wgmma` 如何把逻辑 operand coordinate 转换成 shared memory address。 + +Descriptor 包含如下字段： + +```text +start address +leading dimension offset +stride dimension offset +swizzle mode +base offset +``` + +具体解释取决于 operand major mode。对于 K-major tile，一个 stride 沿 K 前进，另一个沿 M 前进。对于 MN-major tile，这些角色会交换。 + +Swizzle mode 是 shared memory descriptor format 之一，例如： -![TODO：翻译 Hopper shared memory descriptor 图注](../../img/smem_descriptor.svg) +```text +SWIZZLE_NONE +SWIZZLE_32B +SWIZZLE_64B +SWIZZLE_128B +``` + +Swizzle mode 决定两件事。它决定 descriptor 使用的 atom shape，也决定在 atom 内应用的 XOR permutation。例如，128-byte swizzle mode 会把 operand 看作 8-row by 128-byte atom 的网格，并在每个 atom 内应用 swizzle。 + +Kernel 仍然必须正确放置 byte。TMA 通常负责填充 shared memory tile，而 TMA descriptor 必须使用与后续 `wgmma` descriptor 命名的 swizzle format 相同的格式。如果 TMA 写入一个 128-byte swizzled tile，`wgmma` descriptor 就必须把它作为 128-byte swizzled tile 读取。如果 descriptor 和数据不一致，Tensor Core 就会读取被打乱的 operand。 + +这是相对 Ampere 的主要变化。Swizzle 不再只是隐藏在手写 shared memory indexing 里。Hopper 把它变成一等 descriptor format。写 tile 的 TMA load 和读 tile 的 `wgmma` 指令都可以命名同一种 format。 + +![A Hopper shared memory matrix descriptor maps operand coordinates into swizzled shared memory atoms: the descriptor strides choose the atom, and the swizzle chooses the byte position inside the atom](../../img/smem_descriptor.svg) ## Hopper 输出仍然使用寄存器 -TODO：翻译 “Hopper Output Still Uses Registers” 小节。 +Hopper 改变了输入路径，但 accumulator 仍然位于寄存器中。 + +一条 `wgmma` 指令会把 accumulator 写成 per-thread register fragment。具体 fragment 大小和寄存器数量取决于 instruction shape，例如 `m64nNk16` 中 N 会改变 accumulator register 的数量。但基本思想和 Ampere 相同：epilogue 消费的是 register fragment。 + +因此，Hopper 具有混合布局模型。输入 operand 可以直接来自 shared memory descriptor，并由硬件描述 swizzle。输出 accumulator 仍然是一个 register layout 问题。 + +Blackwell 改变了输出侧。 ## Blackwell：`tcgen05` 和 TMEM -TODO：翻译 “Blackwell: tcgen05 and TMEM” 小节。 +Blackwell 保留了数据 operand 上的 shared memory descriptor 思想。A 和 B 仍然以 Tensor Core 期望的布局准备在 shared memory 中。有些 mode 也可以从 TMEM 读取 A operand。 + +主要变化是 accumulator。`tcgen05.mma` 会把 accumulator 写入 Tensor Memory，也就是 TMEM，而不是把它作为长生命周期 register fragment 保持住。在 compute phase 中，accumulator 位于 TMEM。Epilogue 随后用 `tcgen05.ld` 把它加载回寄存器。 + +这把输出布局问题从寄存器移动到了 TMEM。Kernel 必须分配 TMEM，选择正确的 TMEM 布局，等待 MMA 完成，然后使用匹配的 `tcgen05.ld` 路径把 accumulator fragment 恢复出来供 epilogue 使用。 + +`cta_group::1` 和 `cta_group::2` 如何把 accumulator 切分到一个或两个 CTA 上，会在 {ref}`chap_tensor_cores` 中介绍。与早期世代差异最大的布局，是 block-scaled scale-factor layout。 ## TMEM 中的 Scale Factor Layout -TODO：翻译 “Scale Factor Layout in TMEM” 小节。 +Block-scaled MMA mode，例如 `mxfp8` 和 `nvfp4`，会添加 scale-factor operand。除了 A 和 B，MMA 还会读取： + +```text +SFA(M, SFK) +SFB(N, SFK) +``` + +其中 `SFK` 是 K scale block 的数量。 + +数据 operand A 和 B 位于 shared memory。Scale factor 位于 TMEM。因此它们有不同的数据移动路径。 + +TMA 从 global memory 加载到 shared memory。它不会直接加载到 TMEM。因此 scale factor 通常分两步移动： + +```text +global memory to shared memory with TMA +shared memory to TMEM with tcgen05.cp +``` -![TODO：翻译 scale_vec byte packing 图注](../../img/sf_scale_vec.svg) +只有完成这次 copy 后，scale factor 才会位于 `tcgen05.mma` 期望读取的内存空间中。 + +TMEM scale-factor layout 使用 TMEM 的硬件坐标 Lane 和 Col。在 TIRx 布局记号中，这些轴写作 `TLane` 和 `TCol`。 + +一个 128-row scale vector 会紧凑压缩到一个 32-lane group 中，然后复制到 TMEM 的四个 32-lane window 上。在布局记号中，核心 pattern 是： + +```text +S[(32, sf_per_mma) : (1@TLane, 1@TCol)] + R[4 : 32@TLane] +``` + +Shard 放置 base 32-row group： + +```text +TLane = r +TCol = s +``` + +Replica term 在 lane offset 0、32、64 和 96 处添加副本： + +```text +TLane = r + 32 * q, where q in {0, 1, 2, 3} +TCol = s +``` + +这就是 `warpx4` broadcast pattern。同一组紧凑 scale factor 会在完整 128-lane TMEM 空间中变得可见。 + +32-bit `TCol` cell 内部还有 byte packing。Packing 取决于 `scale_vec` mode： + +```text +1X: one scale value is broadcast across the 32-bit cell +2X: two scale values are packed, each duplicated +4X: four K-block scale values are packed +``` + +![scale_vec byte packing: 1X broadcasts one scale across the 4-byte cell; 2X packs two scales, each duplicated; 4X packs four K-block scales](../../img/sf_scale_vec.svg) + +这个 packing 在 Ampere 或 Hopper 中没有直接对应物，因为那些世代没有用于 `tcgen05` block-scaled MMA 的 TMEM scale-factor operand。 + +在 `cta_group::2` 中，scale factor 会跟随它们缩放的数据。SFA 缩放 A，因此它会按 M 在两个 CTA 之间切分，匹配每个 CTA 拥有的 A 行。SFB 缩放 B，而 B 被计算的两个 CTA half 共享，因此 SFB 会 multicast 给两个 CTA（{ref}`chap_tensor_cores`）。 ## 一个反复出现的 Fragment -TODO：翻译 “A Recurring Fragment” 小节。 +虽然周围的内存路径不断变化，但有一个结构反复出现：m8n8-style register fragment。 + +在 Ampere 上，`ldmatrix` 构造这个 fragment，供 `mma.sync` 读取。 + +在 Hopper 上，`wgmma` 把 accumulator 写成 register fragment，供 epilogue 使用。 + +在 Blackwell 上，accumulator 在 compute 期间位于 TMEM，但 `tcgen05.ld` 会在 epilogue 处理和存储之前，把它重新加载到 register fragment 中（{ref}`chap_tmem`）。 + +所以 fragment 并没有消失，它的角色变了。早期世代会让 accumulator 在整个 compute phase 中都留在这里。Blackwell 主要在 TMEM 与 epilogue 之间的边界上使用它。 ## 主线 -TODO：翻译 “The Throughline” 小节。 +在 Ampere 上，kernel 显式构建 Tensor Core register fragment。Shared memory swizzle 大多由 kernel 通过 index math 负责。 + +在 Hopper 上，Tensor Core 可以通过 matrix descriptor 直接从 shared memory 读取 operand。Swizzle 变成 TMA 和 `wgmma` 共享的命名 descriptor format。 + +在 Blackwell 上，输入侧仍然使用 shared memory operand，但 accumulator 移到 TMEM 中。Block-scaled MMA 还添加了必须 staged 到 TMEM 的 scale-factor operand。 + +Descriptor 并不会消除布局工作。它们只是把契约显式化。Kernel 仍然必须确保数据搬运路径、内存布局和 Tensor Core 指令全都一致。写入 swizzled SMEM tile 的 TMA descriptor、读取这个 tile 的 MMA descriptor，以及附着在 buffer 上的布局，都必须描述同一个物理排列。 + +如果其中任意一项不一致，硬件仍然会运行。它只会读到错误的 byte，或者以很慢的方式读取它们。这就是为什么布局不是 Tensor Core kernel 周围的装饰；它是指令接口的一部分。 diff --git a/zh/chapter_tensor_cores/index.md b/zh/chapter_tensor_cores/index.md index fc752c80..35949513 100644 --- a/zh/chapter_tensor_cores/index.md +++ b/zh/chapter_tensor_cores/index.md @@ -1,84 +1,187 @@ (chap_tensor_cores)= # Tensor Core：`tcgen05` - - :::{admonition} 概览 :class: overview -- TODO：翻译本章 overview 第一条。 -- TODO：翻译本章 overview 第二条。 -- TODO：翻译本章 overview 第三条。 +- `tcgen05` 是 Blackwell 的 Tensor Core 指令族。它的 MMA 指令会协作执行 tile matrix-multiply-accumulate 工作，并由一个被选中的线程 commit 指令。 +- Accumulator 位于 TMEM，而不是寄存器。Epilogue 随后用 `tcgen05.ld` 把它带回寄存器。 +- `cta_group::1` 和 `cta_group::2` 控制一个 CTA 还是两个 CTA 协作执行 MMA。这个选择也会改变 M 维度映射到 TMEM 的方式。 +- Block-scaled MMA mode，例如 `mxfp8` 和 `nvfp4`，会添加 scale-factor operand。数据 operand 位于 SMEM，而 scale factor 会通过 TMEM staged。 ::: -TODO：翻译导言部分。 +Dense linear algebra 是现代 GPU 花费最多有效工作的地方。普通 CUDA core 矩阵乘法无法接近芯片标称峰值（{ref}`chap_background`）。快速 GEMM 和 attention kernel 通过给 Tensor Core 喂入正确的 tile shape、layout 和 synchronization，才能达到这个峰值。 + +基本操作从 Volta 以来在精神上没有改变。Tensor Core 消费矩阵 tile，相乘，并把结果累加起来。每一代变化的是操作如何发起、操作数如何布局，以及 accumulator 位于哪里。 + +Blackwell 对最后一部分做了很大改变。`tcgen05` 的 accumulator 不再作为长生命周期 register fragment 保存。它会写入 Tensor Memory，也就是 TMEM（{ref}`chap_tmem`）。这一个变化会影响整个 kernel。MMA 写入 TMEM。完成状态被异步跟踪。Epilogue 随后从 TMEM 中加载 accumulator，并把它重新变成用于转换和 store 的 register fragment。 + +本章聚焦 compute instruction 本身。TMA（{ref}`chap_tma`）负责把 operand 移入 SMEM。TMEM 负责保存 accumulator 和某些 scale-factor operand。`tcgen05.mma` 是位于这两类数据移动之间的 Tensor Core 操作。 ```{raw} html +
+ style="width:100%; min-width:1320px; height:640px; border:1px solid var(--pst-color-border, #d0d0d0); border-radius:6px;"> +
``` - -*点击图中组件查看细节：TODO：翻译 tcgen05 and Tensor Memory 图注。* +*交互图：`tcgen05` accumulator 行为。切换 A 或 B 的 transpose，选择输出宽度 `N`，并逐步执行 `K` iteration，观察 partial sum 如何在 TMEM 中累加。* ## `tcgen05` MMA -TODO：翻译 “The tcgen05 MMA” 小节。 +`tcgen05` MMA 是 Blackwell Tensor Core 的 matrix-multiply-accumulate 指令。它是一条协作指令。工作由一个 warpgroup 执行，并且在某些 mode 中可以涉及同一 cluster 中的两个 CTA。指令不是由每个线程独立发起的。一个被选中的线程会代表参与的 group commit 这个操作。 + +把 MMA 拆成三个问题来看会更清楚。 + +第一个问题是谁参与协作。普通 mode 使用一个 CTA，写作 `cta_group::1`。更大的 mode 使用 cluster 中的两个 CTA，写作 `cta_group::2`。在两种情况下，这条指令都表示一个 tile 上的一次 Tensor Core 操作，而不是某个线程执行的 scalar operation。 + +第二个问题是 operand 和 result 位于哪里。数据 operand 通常位于 SMEM。某些变体也可以从 TMEM 读取 A operand。Accumulator 写入 TMEM。Operand layout 必须匹配 Tensor Core 的期望，包括数据 operand 使用的 swizzled shared-memory layout（{ref}`chap_data_layout`）。 + +第三个问题是如何观察完成。`tcgen05.mma` 是异步的。发起 MMA 并不表示 multiply-accumulate 已完成。指令在操作 commit 后返回，而 Tensor Core 会继续运行。Kernel 使用 commit group 和 `mbarrier` 来知道结果何时就绪（{ref}`chap_async_barriers`）。 + +这种异步行为让 overlap 成为可能。快速 kernel 不会发起 MMA 后立即停下来等待它完成。它可以发起 MMA，开始准备后续 tile，并且只在真正需要结果时等待。代价是每个交接都必须显式表达。如果 epilogue 在 MMA completion barrier 触发前读取 TMEM，那就是读得太早。 ## Accumulator 位于 TMEM -TODO：翻译 “The Accumulator Lives in TMEM” 小节。 +在 Ampere 和 Hopper 上，accumulator 以寄存器形式暴露给程序。MMA 产生一个 per-lane register fragment，epilogue 直接消费这个 fragment。这很简单，但它把 accumulator 大小绑定到了每个线程的寄存器预算上。 + +Blackwell 切断了这条绑定。`tcgen05.mma` 会把 accumulator 写入 TMEM，这是 Blackwell 上作用域属于 CTA 的内存空间。Accumulator 可以在 compute phase 中保存在 TMEM 中，epilogue 随后用 `tcgen05.ld` 把它加载回寄存器。 + +这改变了 kernel 的形状。Register fragment 在边界上仍然重要。Epilogue 仍然需要寄存器，以便转换、执行 elementwise 工作并 store 结果。但长生命周期 accumulator state 不再是寄存器分配问题，而是 TMEM 分配和布局问题（{ref}`chap_tmem`）。 + +这就是为什么必须把 `tcgen05` 和 TMEM 放在一起理解。MMA 指令决定计算哪个 tile。TMEM 决定 accumulator 落在哪里。Epilogue 必须使用匹配的 load path，把 accumulator 恢复到它期望的 register layout。 ## `cta_group::1` 和 `cta_group::2` -TODO：翻译 “cta_group::1 and cta_group::2” 小节。 +`tcgen05` MMA 可以在 `cta_group::1` 或 `cta_group::2` mode 中运行。 + +在 `cta_group::1` 中，一个 CTA 拥有这次 MMA。它的 operand 位于该 CTA 的 SMEM 中，accumulator 写入该 CTA 的 TMEM。 + +在 `cta_group::2` 中，cluster 中的两个 CTA 协作执行一个 MMA tile。每个 CTA 都有自己的 SMEM 和自己的 TMEM。Accumulator 并不是存储在跨两个 CTA 的单个物理 TMEM 区域中。它会在两个 CTA 之间切分，每个 CTA 保存自己的部分。偶数 CTA 负责发起指令，并为这对 CTA commit completion barrier。 + +这个选择很重要，因为它会改变逻辑 accumulator tile `C(M, N)` 到 TMEM 的映射。TMEM 有 128 个硬件 Lane 行和最多 512 个硬件 Col 列。在 TIRx 布局记号中，这些轴写作 `TLane` 和 `TCol`。MMA mode 决定 C 的行和列如何放到这些 TMEM 轴上。 + +有四种有用情况需要记住。 + +下面的图沿用 demo 中的颜色约定：紫色表示 SMEM operand，橙色表示 TMEM accumulator state，绿色表示 Tensor Core MMA 路径。CTA identity 通过标签和位置表示，而不是改变这些硬件颜色。 ### `cta_group::1`，`M = 128` -TODO：翻译本小节。 +这是最简单的情况。一个 CTA 计算一个 128-row tile。TMEM 也有 128 个 Lane 行。因此映射是直接的：accumulator 的 row `m` 映射到 Lane `m`，N 维度映射到 TMEM column。 + +结果填满 128 个 Lane 行乘以 N 个 Col 列。这是基线图。CTA 在 SMEM 中拥有 A 和 B，并在自己的 TMEM 中拥有完整 accumulator tile。 -![TODO：翻译 cta_group::1 M=128 图注](../../img/mma_cg1_m128.svg) +![cta_group::1, M=128: row m maps directly to TMEM Lane m](../../img/mma_cg1_m128.svg) ### `cta_group::1`，`M = 64` -TODO：翻译本小节。 +当 `M = 64` 时，accumulator 只有 64 行，但 TMEM 仍然有 128 个 Lane 行。硬件并不是简单地把 row 0 到 63 打包到 lane 0 到 63。相反，它会把这些行分散到四段 16-row run 中。 + +Rows 0 到 15 放到 lanes 0 到 15。Rows 16 到 31 放到 lanes 32 到 47。Rows 32 到 47 放到 lanes 64 到 79。Rows 48 到 63 放到 lanes 96 到 111。 -![TODO：翻译 cta_group::1 M=64 图注](../../img/mma_cg1_m64.svg) +这会在 lanes 16 到 31、48 到 63、80 到 95、112 到 127 留出空隙。这些空隙是有意的。使用不同 lane alignment 时，另一个独立的 `M = 64` MMA 可以占用互补 lane。这样两个较小的 M tile 可以共享 128-lane TMEM 结构，而不会互相覆盖。 + +N 维度仍然映射到 TMEM column。不寻常的地方只有 M row 在 Lane 上的 placement。 + +![cta_group::1, M=64: four 16-row runs at a Lane stride of 32, leaving space for another aligned M=64 tile](../../img/mma_cg1_m64.svg) ### `cta_group::2`，`M = 256` -TODO：翻译本小节。 +当 M 维度大到单个 CTA 自然容纳不下时，MMA 可以使用 `cta_group::2`。对于 `M = 256`，切分是直接的。CTA 0 持有 rows 0 到 127。CTA 1 持有 rows 128 到 255。 + +每个 CTA 使用自己的 TMEM Lane rows 0 到 127，以及完整 N columns。物理上，这是两个独立的 128-row TMEM 区域，每个 CTA 一个。逻辑上，它们组成一个 256 by N 的 accumulator tile。 + +每个 CTA 也提供 A 中对应自己 M row 的部分。B 会按照 mode 的要求对两个 CTA 可用。偶数 CTA 负责发起 MMA，并为这对 CTA commit completion barrier。 -![TODO：翻译 cta_group::2 M=256 图注](../../img/mma_cg2_m256.svg) +这是 {ref}`chap_gemm_advanced` 中 two-CTA cluster GEMM 使用的 mode。 + +![cta_group::2, M=256: M split contiguously across two CTAs, 128 rows per CTA](../../img/mma_cg2_m256.svg) ### `cta_group::2`，`M = 128` -TODO：翻译本小节。 +`cta_group::2`, `M = 128` mode 仍然使用两个 CTA，但 M 维度更短。因为总共只有 128 行，每个 CTA 得到 64 个 M row。 + +剩余 lane capacity 被用来打包 N 维度。在每个 CTA 内部，N 的一半占用 lanes 0 到 63，另一半占用 lanes 64 到 127。这样即使每个 CTA 只拥有 64 个 M row，也能使用全部 128 个 Lane 行。 + +因此这个 split 有两部分。M 在 CTA pair 之间切分，每个 CTA 64 行。N 随后在每个 CTA 内部切分到 TMEM Lane 行的 lower half 和 upper half。 -![TODO：翻译 cta_group::2 M=128 图注](../../img/mma_cg2_m128.svg) +![cta_group::2, M=128: 64 M rows per CTA, with the two halves of N stacked across the lower and upper Lane halves](../../img/mma_cg2_m128.svg) + +在这些 mode 中，原则相同。`tcgen05.mma` 计算一个逻辑 accumulator tile，但这个 tile 必须放入物理的 128 Lane by up to 512 Col TMEM 空间中。Mode 和 M shape 决定这种 placement。Kernel 的其他部分稍后读取 accumulator 时，必须使用同一映射。 + +对于这里的 kernel，accumulator 通常在 TMEM 中使用 f32。这是常见的高精度路径。它不是唯一可能的 accumulator type。`.kind::f16` 路径可以用 f16 accumulation。 ## Operand Placement -TODO：翻译 “Operand Placement” 小节。 +对于 dense MMA mode，A 和 B 在 MMA 运行前准备在 SMEM 中。TMA 负责把 global memory tile 移入 SMEM。Kernel 把这些 SMEM tile 安排成 Tensor Core 期望的布局，包括任何必要的 swizzle。 + +Accumulator C 写入 TMEM。这是相对早期世代的主要区别。Epilogue 不会直接收到 MMA 指令输出的 accumulator。它必须用 `tcgen05.ld` 显式从 TMEM 加载。 + +在 `cta_group::1` 中，一个 CTA 提供 operand 并拥有 accumulator。在 `cta_group::2` 中，每个 CTA 从自己的 SMEM 中提供自己那一侧的 operand，并拥有 accumulator 中自己的 TMEM 部分。当 A 按 M 切分时，每个 CTA 保留自己 M slice 对应的 A row。B 根据 mode 被共享，因为两个 M slice 都要乘同一个 N by K tile。 + +阅读 kernel 时，这种分离很重要。SMEM placement 回答 Tensor Core 如何读取 A 和 B。TMEM placement 回答 accumulator 去哪里。这两个 layout 由 MMA mode 关联起来，但它们不是同一个内存空间，不能互换对待。 ## Block-Scaled MMA -TODO：翻译 “Block-Scaled MMA” 小节。 +Dense mode 直接从 SMEM 读取数据 operand，并累加到 TMEM。Block-scaled MMA 添加了两个额外 operand：A 和 B 的 scale-factor tensor。 + +这用于 `mxfp8` 和 `nvfp4` 等非常低精度格式。低精度格式效率高，但动态范围小。单个 global scale 通常太粗。如果 scale 按最大值选择，小值会丢精度。如果 scale 按小值选择，大值可能溢出或被截断。 + +Block scaling 通过给较小的 K block 分配 scale factor 来解决这个问题。一组连续 K element 共享一个 scale。MMA 在概念上先用 scale 对每个 block dequantize，再把乘积累加到 accumulator type 中。 + +对于 A 和 B，这引入两个 scale-factor tensor： + +```text +SFA(M, SFK) +SFB(N, SFK) +``` + +其中 `SFK = K / B`，`B` 是沿 K 的 block size。 + +具体 block size 取决于格式。重要的是，scale axis 以更粗粒度跟随 K。每个 scale factor 描述一块 K value，而不是单个元素，也不是整张矩阵。 + +数学形状是： + +```text +acc += (Aq * scale_a) * (Bq * scale_b) +``` + +其中 `Aq` 和 `Bq` 是量化的低精度值，scale 在 accumulation 前恢复它们的近似 magnitude。 + +Scale dtype 也很重要。使用 `e8m0` scale 时，每个 scale 实际上是 2 的幂。使用 `e4m3` scale 时，例如 `nvfp4` 中的 scale，它是一个小浮点值，可以表示两个 2 的幂之间的数值。 ## Scale Factor 放在哪里 -TODO：翻译 “Where the Scale Factors Live” 小节。 +Block-scaled `tcgen05.mma` 与 dense MMA 有一个重要 placement 规则不同：scale factor 从 TMEM 读取。 + +数据 operand A 和 B 仍然 staged 在 SMEM 中。Scale factor SFA 和 SFB 通过 TMEM staged。由于 TMA 加载到 SMEM，scale factor 通常需要额外一步。Kernel 先把它们加载到 SMEM，再用 `tcgen05.cp` 从 SMEM copy 到 TMEM。只有 scale factor 位于 TMEM 后，block-scaled MMA 才能读取它们。 + +这给 scale factor 一条不同于数据 operand 的移动路径： + +```text +A, B: global memory to SMEM, then MMA reads SMEM +SFA, SFB: global memory to SMEM, then tcgen05.cp copies SMEM to TMEM, then MMA reads TMEM +``` + +Scale factor 的 TMEM layout 是紧凑的。一个 128-row scale vector 可以打包到 32 个 Lane row 中，lane 位置基于 `r % 32` 映射，`r / 32` 沿 column 方向前进。随后数据可以 broadcast 到读取完整 128 Lane 空间的四个 warp 上（{ref}`chap_layout_generations`）。 + +这很好地说明为什么 TMEM layout 必须显式。Accumulator layout 和 scale-factor layout 都在 TMEM 中，但它们不是同一个 layout。Accumulator 使用 MMA output mapping。Scale factor 使用 block-scaled MMA 期望的紧凑布局。 ## `cta_group::2` 中的 Scale Factor -TODO：翻译 “Scale Factors in cta_group::2” 小节。 +在 two-CTA 情况下，scale factor 跟随它们缩放的数据。 + +SFA 缩放 A。由于 A 在 CTA pair 之间按 M 切分，SFA 也按 M 切分。每个 CTA 持有与自己 A row 对应的 SFA row。 + +SFB 缩放 B。由于两个 CTA 都要乘同一个 B tile，SFB 必须对两个 CTA 可见。实践中，这意味着 SFB 会 multicast 到 CTA pair。 -![TODO：翻译 block-scaled MMA placement 图注](../../img/mma_block_scaled.svg) +这是 block-scaled cluster GEMM 中常见 loading pattern 的来源。SFA 按 CTA 加载，使用该 CTA 自己 M slice 的 mask。SFB broadcast 给这对 CTA，因为两个 CTA 都需要同一组 N-side scale factor。 + +![Block-scaled MMA placement: A and B packed in SMEM; SFA, SFB, and C in TMEM, with SFA split by M across CTAs and SFB multicast across the CTA pair](../../img/mma_block_scaled.svg) ## 保持 MMA Contract 匹配 -TODO：翻译 “Keeping the MMA Contracts Matched” 小节。 +一个 Blackwell GEMM tile 会经过多条专用路径。 + +TMA 把 A 和 B 从 global memory 带入 SMEM。对于 block-scaled mode，它也会把 scale factor 带入 SMEM。需要时，`tcgen05.cp` 把这些 scale factor 移到 TMEM。`tcgen05.mma` 读取 operand，在 Tensor Core 上异步运行，并累加到 TMEM。Completion barrier 告诉 kernel accumulator 何时就绪。Epilogue 随后用 `tcgen05.ld` 从 TMEM 把 accumulator 加载回寄存器，并 store 最终输出。 + +在这些路径之间，kernel 必须保持三份 contract 匹配：SMEM operand layout、TMEM accumulator 或 scale-factor layout，以及让下一个 consumer 安全运行的异步完成信号。 diff --git a/zh/chapter_tirx_layout_api/index.md b/zh/chapter_tirx_layout_api/index.md index 2c7388cc..e559987e 100644 --- a/zh/chapter_tirx_layout_api/index.md +++ b/zh/chapter_tirx_layout_api/index.md @@ -1,10 +1,711 @@ ---- -orphan: true ---- - (chap_tirx_layout_api)= # TIRx Layout API -> 翻译状态：待翻译。对应英文章节：`chapter_tirx_layout_api/index.md`。 +:::{admonition} 概览 +:class: overview + +- TIRx layout API 会把 {ref}`chap_data_layout` 中的布局记号变成编译器对象。主要对象是 `TileLayout`、`SwizzleLayout` 和 `ComposeLayout`。 +- `TileLayout` 描述命名硬件轴上的 affine placement。它由 shard spec `S[...]`、replica spec `R[...]` 和可选 offset 构成。 +- 一个布局会把一个逻辑坐标映射到一个或多个物理坐标。`layout.apply()` 会计算这个映射。 +- `SwizzleLayout` 描述用于避免 bank conflict 的 XOR-based shared-memory swizzle。`ComposeLayout` 会把一个 swizzle 叠加到 tile layout 上。 +- `tmem_datapath_layout`、`tcgen05_atom_layout` 和 `wg_local_layout` 这类现成 constructor 覆盖了 kernel 中反复出现的硬件布局。 +::: + +{ref}`chap_data_layout` 介绍了本书一直使用的记号：tile shape、命名轴上的一组 stride，以及可选的 replication term，用于表示被复制而不是被切分的值。本章会把这套记号变成编译器使用的 API。 + +目标是让页面上的记号和 kernel 里的代码几乎长得一样。当你写出这样的布局时： + +```python +S[(128, 256) : (1@TLane, 1@TCol)] +``` + +你写的不只是解释文字。你正在构造一个可以挂到 buffer 上的 `TileLayout` 对象。之后，每个触碰这个 buffer 的 tile operation 都能从 layout 中读取它的 placement。Placement 只写一次、检查一次，并由编译器复用。 + +布局会在从 pool 分配时，或声明 buffer 时附着上去： + +```python +pool.alloc(shape, dtype, layout=layout) + +T.decl_buffer(shape, dtype, scope=scope, layout=layout) +``` + +从那时起，buffer 就携带了自己的物理 placement。Tile operation 不需要反复说明每个元素位于哪里。 + +这些 layout object 位于同一个模块中： + +```python +from tvm.tirx.layout import ( + TileLayout, + SwizzleLayout, + ComposeLayout, + S, + R, + laneid, + warpid, + tid_in_wg, + TLane, + TCol, + m, + tcgen05_atom_layout, + tmem_datapath_layout, +) +``` + +这个 API 背后有一个核心思想：布局不一定把逻辑索引映射到单个物理地址。它会把逻辑索引映射到命名轴上的一组物理坐标。通常这组坐标只有一个元素。当存在 replication 时，同一个逻辑元素会有多个物理 placement。 + +这就是为什么 layout model 有三部分：shard、replica 和 offset。Shard 放置元素。Replica 把它复制到额外坐标。Offset 移动整个 placement。 + +## 通过例子看布局 + +下面的例子展示 API 的基本形状。 + +TMEM 中的 accumulator 可以写成覆盖 TMEM 轴的直接 placement： + +```python +acc = TileLayout(S[(128, 256) : (1@TLane, 1@TCol)]) +``` + +这里逻辑行映射到 `TLane`，逻辑列映射到 `TCol`。在 {ref}`chap_tmem` 中，硬件坐标称为 Lane 和 Col。在 TIRx 布局记号中，这些硬件轴写作 `TLane` 和 `TCol`。 + +Block-scaled MMA 的 scale-factor layout 会使用 replication： + +```python +scale_factor_layout = TileLayout( + S[(32, sf_per_mma) : (1@TLane, 1@TCol)] + R[4 : 32@TLane] +) +``` + +Shard 会把一个 32-row group 放入 TMEM。Replica 以 32 个 lane 为 stride，把这个 group 重复四次，使这个 32-row group 在完整 128-lane TMEM 空间中都可见。 + +Tensor-core register fragment 可以分布在 lane 和 warp 上： + +```python +frag = TileLayout( + S[(8, 2, 4, 2) : (4@laneid, 1@warpid, 1@laneid, 1)] +) +``` + +同一个物理轴可以出现多次。在这个例子中，两个不同 iter 都会贡献到 `laneid`。没有显式 axis 的 stride 使用默认 memory axis `m`。 + +在真实 kernel 中，常见硬件布局通常来自 constructor： + +```python +acc = tmem_datapath_layout("D", 128, 256) + +ld = tcgen05_atom_layout("32x32b", (128, 64), "float32") +``` + +这些 constructor 返回普通 `TileLayout` 对象。它们只是方便写法，不是单独机制。你可以检查返回的 layout，把它和其他 layout 组合，或者在 shape 不寻常时手写底层 `S[...]` 和 `R[...]` 形式。 + +## 交互 Demo + +在进入机制之前，先有一个可以操作的具体对象会更容易理解。下面的 demo 允许你选择 preset layout、编辑 logical shape 和 `S` 或 `R` term、选择 dtype 和 swizzle mode，并点击某个 element 查看哪个物理坐标或哪些物理坐标拥有它。 + +```{raw} html +

+ ▶ 全屏打开 demo ↗ +

+ + +``` + +这个 demo 很有用，因为 API 的大部分内容只是把 demo 中展示的过程精确定义出来。一个 logical element 进入 layout。Layout 会把它 flatten，按 iter 拆分，在命名轴上累加坐标，然后在需要时应用 replication。 + +## TileLayout + +`TileLayout` 是主要的 affine layout object。它通常用正文中同样的记号写出： + +```python +TileLayout(S[shape : strides]) +``` + +`S` term 是 shard spec。可以这样读：取一个具有这个 shape 的逻辑 tile，并用这些命名轴 stride 放置它。 + +当某个值需要出现在多个地方时，shard spec 会扩展出 replica spec： + +```python +TileLayout(S[shape : strides] + R[replica_shape : replica_stride]) +``` + +还可以加入可选 offset： + +```python +TileLayout(S[shape : strides] + R[replica_shape : replica_stride] + offset) +``` + +在表面之下，这些部分由 iter 表示。一个 iter 是三元组： + +```text +(extent, stride, axis) +``` + +它描述命名轴上的一次 strided walk。Extent 表示这个 iter 有多少个位置。Stride 表示每一步移动多远。Axis 表示被改变的是哪个硬件坐标。 + +一个 layout 有三部分。 + +### Shard + +Shard，也就是 `D`，是由 `S[...]` 构建的部分。它把逻辑索引切分到一个或多个 iter 上，并产生 base physical coordinate。 + +例如： + +```python +S[(8, 2, 4, 2) : (4@laneid, 1@warpid, 1@laneid, 1)] +``` + +有四个 shard iter。它们的 extent 是 `8`、`2`、`4` 和 `2`。它们的 stride 会把数据放到 `laneid`、`warpid`、再次 `laneid`，以及默认 memory axis `m` 上。 + +这推广了普通 shape-and-stride 规则。差别在于，这里的 stride 附着在命名硬件轴上，而不是单个 flat address 上。 + +### Replica + +Replica，也就是 `R`，描述同一个逻辑元素的额外物理副本。Replica iter 与逻辑索引无关。它们枚举硬件空间中的额外 offset。 + +例如： + +```python +R[2 : 4@warpid] +``` + +会在 `warpid` 轴上创建两个相距四个 warp 的副本。 + +Replication 不是为了方便写法的技巧。它描述真实硬件行为。有些数据会 broadcast 到多个 warp、lane 或 memory region。Logical-to-physical mapping 很自然地支持这一点，因为一个逻辑元素可以映射到一组物理坐标。 + +### Offset + +Offset，也就是 `O`，是加到每个结果上的固定坐标。 + +例如： + +```python +5@warpid +``` + +会把整个 placement 在 `warpid` 轴上移动 5。 + +Offset 用于把 tile 放到选定 base coordinate、为独占使用预留区域，或者描述在同一资源中位于另一个 tile 之后开始的 tile。 + +### 把三部分合在一起 + +Layout 会按顺序应用这三部分。 + +首先，shard 计算 base coordinate。然后，replica 把这个 coordinate fan out 成零个或多个额外副本。最后，offset 移动每个 coordinate。 + +对于逻辑坐标 `x`，结果是： + +```text +L(x) = { D(x) + r + O | r in R } +``` + +如果没有 replica，`R` 只包含 zero offset，因此结果是 singleton set。如果有 replica，结果中会有每个 replica position 对应的一个 coordinate。 + +在 TIRx 语法中，一个完整 layout 可以写作： + +```python +layout = TileLayout( + S[(8, 2, 4, 2) : (4@laneid, 1@warpid, 1@laneid, 1)] + + R[2 : 4@warpid] + + 5@warpid +) +``` + +从左到右读，shard 放置逻辑 tile，replica 在四个 warp ID 之外创建第二份副本，offset 把整个 placement 移动到从 `warpid = 5` 开始。 + +如果 iter 已经被构造成对象，同一个 layout 也可以直接构造： + +```python +TileLayout.from_iters(shard, replica, offset) +``` + +大多数用户代码使用 `S[...]` 和 `R[...]` 记号，因为它更接近数学形式。 + +## 命名轴 + +Layout 中的轴不是匿名维度。每个轴都命名一个真实硬件坐标，或者一个编译器级 placement 坐标。 + +例子包括： + +```text +bx, by, bz +cbx, cby, cbz +tx +warpid +laneid +wgid +tid_in_wg +wid_in_wg +m +P, F +Bank +TLane, TCol +``` + +`bx`、`by`、`bz` 这样的 grid axis 把 work 放到 CTA 上。`cbx`、`cby`、`cbz` 这样的 cluster axis 把 work 放到 CTA cluster 内部。`tx`、`warpid`、`laneid`、`tid_in_wg` 和 `wid_in_wg` 这样的 thread axis 描述 CTA 或 warpgroup 内的 ownership。`m` 是默认 linear memory axis。`P` 和 `F` 用于二维 scratchpad-style placement。`Bank` 命名 shared memory bank。`TLane` 和 `TCol` 是 TIRx 布局中对 TMEM Lane 和 Col 坐标的命名。 + +Axis name 是 layout 的一部分。这一点很重要，因为两个整数值相同的坐标可能表示不同硬件事物。`1@tx` 不是 `1@tid_in_wg`。`1@laneid` 不是 `1@TLane`。Layout 会让这些含义保持显式。 + +## Forward Mapping + +求值一个 layout，意味着取一个逻辑坐标并计算它物理上落在哪里。API 方法是： + +```python +layout.apply(*coord) +``` + +对于没有 replication 的 layout，结果是一个 coordinate dictionary。带有 replication 时，结果是一组 coordinate dictionary。Coordinate dictionary 把 axis name 映射到整数位置，例如： + +```python +{"laneid": 7, "warpid": 2, "m": 1} +``` + +求值规则有四步。 + +第一步，按 row-major 顺序 flatten 逻辑坐标。对于逻辑 shape： + +```text +(S0, S1, ..., Sr-1) +``` + +中的逻辑坐标： + +```text +x = (x0, x1, ..., xr-1) +``` + +flat index 是： + +```text +flat = x0 * S1 * S2 * ... * Sr-1 + + x1 * S2 * ... * Sr-1 + + ... + + xr-2 * Sr-1 + + xr-1 +``` + +第二步，把这个 flat index 按 shard extents 拆分。如果 shard extents 是： + +```text +(e0, e1, ..., en-1) +``` + +那么拆分会产生 components： + +```text +c0, c1, ..., cn-1 +``` + +使用的是同样的 row-major 顺序，只不过作用在 shard extents 上。 + +第三步，用每个 component 的 stride 把它累加到对应 axis 上。如果 shard iter `k` 的 extent 是 `ek`、stride 是 `sk`、axis 是 `ak`，那么 component `ck` 的贡献是： + +```text +ck * sk @ ak +``` + +对同一个 axis 的所有贡献会相加。然后加入 offset。 + +第四步，应用 replica iter。每个 replica iter 都贡献一个与逻辑坐标无关的额外 offset。如果有多个 replica iter，layout 会枚举所有组合。 + +这个规则有一个有用结果：layout 不需要 hard-code 输入 shape。它需要的是逻辑 tile 的元素总数等于 shard extents 的乘积。只要这个条件成立，flattening 和 splitting 就定义了映射。 + +## Case Study：Tensor Core Register Tile + +考虑一个逻辑 `(8, 16)` tile，它分布在两个 warp 上，每个 warp 有 32 个 lane。每个 lane 拥有一个小 register fragment。Register slot 用默认 memory axis `m` 表示。 + +```python +layout = TileLayout( + S[(8, 2, 4, 2) : (4@laneid, 1@warpid, 1@laneid, 1)] + + R[2 : 4@warpid] + + 5@warpid +) +``` + +取 `(8, 16)` tile 中的一个逻辑元素 `(i, j)`。 + +Row-major flat index 是： + +```text +flat = 16 * i + j +``` + +按 shard extents `(8, 2, 4, 2)` 拆分得到： + +```text +c0 = i +c1 = floor(j / 8) +c2 = floor(j / 2) mod 4 +c3 = j mod 2 +``` + +Shard 贡献是： + +```text +laneid = 4 * c0 + c2 +warpid = c1 +m = c3 +``` + +加上 offset `5@warpid` 后变成： + +```text +laneid = 4 * i + floor(j / 2) mod 4 +warpid = floor(j / 8) + 5 +m = j mod 2 +``` + +Replica term： + +```python +R[2 : 4@warpid] +``` + +会向 `warpid` 添加 `0` 或 `4`。因此完整映射是： + +```text +laneid = 4 * i + floor(j / 2) mod 4 +warpid = floor(j / 8) + 5 + 4 * r, where r in {0, 1} +m = j mod 2 +``` + +Shard 把 tile 放到 warps 5 和 6 上。Replica 随后把它复制到 warps 9 和 10。于是同一个逻辑元素会出现在两个 warp 位置。 + +这个例子说明了为什么模型使用一组物理坐标。Replication 很难自然表示成从物理坐标到逻辑坐标的函数。它更自然地表示成从一个逻辑坐标到多个物理坐标的函数。 + +## Case Study：Blackwell Tensor Memory + +同一个 layout model 也适用于内存 placement。Axis 不一定是 thread axis，也可以是 memory axis。 + +TMEM 通过硬件 Lane 和 Col 坐标寻址。在 TIRx 布局记号中，这些轴写作 `TLane` 和 `TCol`。 + +考虑这个 layout： + +```python +layout = TileLayout( + S[(2, 128, 112) : (112@TCol, 1@TLane, 1@TCol)] +) +``` + +如果逻辑 tile shape 是 `(2, 128, 112)`，split component 就是逻辑坐标本身。对于元素 `(a, l, c)`，映射是： + +```text +TLane = l +TCol = 112 * a + c +``` + +Extent-128 iter 以 stride `1@TLane` 填充 128 个 TMEM Lane row。Extent-2 iter 以 stride `112@TCol`，extent-112 iter 以 stride `1@TCol`，两者共同覆盖 224 列： + +```text +TCol in [0, 224) +``` + +224-column span 是有意选择的。TMEM layout 不一定是 2 的幂。Block-scaled FP8 GEMM 可能选择 224-column accumulator，因为完整 256-column tile 可能无法为两个 accumulator stage 加 scale factor 留出足够 TMEM 容量。Layout API 可以直接表达这种 shape。 + +## Scale Factor Layout + +上面的 accumulator layout 是纯 placement。每个逻辑 accumulator element 映射到一个 TMEM coordinate。Block-scaled MMA 的 scale factor 不同，因为同一个物理 group 可能需要在多个 warp window 中可见。这正是 replication 有用的地方。 + +一个紧凑 scale-factor layout 可以写作： + +```python +scale = TileLayout( + S[(32, sf_per_mma) : (1@TLane, 1@TCol)] + + R[4 : 32@TLane] +) +``` + +Shard 把一个 32-row scale-factor group 放入 TMEM： + +```text +TLane = r +TCol = s +``` + +对于逻辑 scale coordinate `(r, s)`。 + +Replica term 会创建四份副本，间隔 32 个 lane： + +```text +TLane = r + 32 * q, where q in {0, 1, 2, 3} +TCol = s +``` + +所以 32-row group 会在 TMEM lanes 0 到 31、32 到 63、64 到 95 和 96 到 127 处可见。这就是 `warpx4` broadcast pattern（{ref}`chap_layout_generations`）。四个 warp-sized TMEM lane window 中的每一个，都能看到同一个 scale-factor group。 + +在完整 block-scaled MMA layout 中，这个 atom 会与 M row 和 K scale-factor group 上的 outer iter 组合在一起。多个 scale factor 也可能被打包到同一个 32-bit `TCol` cell 中，具体取决于 scale-factor dtype。例如，fp8 scale factor 可以把四个值打包到一个 32-bit column cell 中。可选的 stride-zero reuse 和 pipeline-depth iter 可以进一步描述多个 MMA 之间的 scale reuse 以及 double buffering。 + +重要的是，同一个 `TileLayout` model 描述了两种情况。Accumulator 是 TMEM 中的单一 placement。Scale factor 是同一 TMEM 地址空间中的 replicated placement。 + +## 现成 Layout + +大多数 kernel 不会手写每一个硬件布局。TIRx 为常见布局提供了 constructor。 + +```python +tmem_datapath_layout(datapath, rows, cols) +``` + +返回 `tcgen05.mma` 写出的 TMEM accumulator layout。`datapath` 参数选择 row placement pattern。例如，`"D"` 对应 `M = 128` identity-style placement，而 `"F"` 对应 `M = 64` scattered placement。 + +```python +tcgen05_atom_layout(instr_shape, tensor_shape, dtype) +``` + +返回由 `tcgen05.ld` 或 `tcgen05.st` atom 移动的 register tile layout。Instruction shape 的例子包括 `.32x32b`、`.16x64b`、`.16x128b` 以及相关形式。在 DSL 层面，这是一个 warpgroup-distributed tile。Lowering 时，它会变成四条 warp-collective `tcgen05.ld` 或 `tcgen05.st` 指令，每个 warp 一条，并且每个 warp 处理自己的 32 个 TMEM lane。 + +```python +wg_local_layout(cols, rows=128) +``` + +返回一个 warpgroup-local register tile，通常在 `tid_in_wg` 上每个线程对应一行。 + +这些 helper 的作用是避免手写常见硬件映射。它们不会隐藏模型。每个 helper 都返回一个普通 `TileLayout`，由上面描述的同一组 `S` 和 `R` 片段构成。 + +## SwizzleLayout 和 ComposeLayout + +`TileLayout` 是 affine 的。它可以在命名轴上表达 stride、replication 和 offset。这足以描述很多 placement，包括 thread fragment、TMEM tile 和紧凑 scale-factor layout。 + +Shared memory swizzle 需要另一种东西。用于避免 bank conflict 的 swizzle 不是 affine stride pattern。它是对线性 shared-memory address 的 XOR-based permutation。 + +因此，TIRx 把 swizzling 保持为一个单独的 layout object： + +```python +SwizzleLayout(...) +``` + +并把它与 tile layout 组合： + +```python +ComposeLayout(swizzle, tile) +``` + +Tile layout 先产生线性 memory address。Swizzle 随后重排这个 address。把两层分开，比强行把 XOR permutation 塞进 affine layout model 更干净。 + +## 为什么需要 Swizzle + +Shared memory 分成 32 个 bank，每个 bank word 保存 4 byte。当一次访问中的多个 lane 触碰同一 bank 中的不同地址时，访问会因为 bank conflict 被串行化。 + +普通 row-major tile 会结构性地产生这种 conflict。考虑一个 row-major layout 的 `(8, 64)` float16 tile： + +```python +TileLayout(S[(8, 64) : (64@m, 1@m)]) +``` + +逻辑元素 `(i, j)` 的线性 element address 是： + +```text +m = 64 * i + j +``` + +每行有 64 个 float16 值，也就是 128 byte。这刚好是一整条 shared memory bank line。如果一个 warp 沿固定 `j` 读取 column，每向下一行都会前进完整 128-byte line。Bank index 重复，因此 column read 会跨行塌缩到同一个 bank 上。 + +Swizzle 通过让低地址位依赖更高的 row bit 来改变这一点。原本会反复落到同一个 bank 的 column 会被分散到不同 bank 上。 + +## Swizzle Transform + +`SwizzleLayout` 由三个整数参数控制： + +```text +per_element = M +swizzle_len = B +atom_len = S +``` + +输入是一个线性 element address `m`。 + +`m` 的低 `M` 位保持不变。这会保留一个小的连续 element group。更高的 bit 会被右移到一个临时值中： + +```text +x = m >> M +``` + +然后，`x` 中位于 `[S, S + B)` 的 bit group 会 XOR 到 `x` 中位于 `[0, B)` 的 bit group 上。最后把保持不变的低 `M` bit 放回去，形成 swizzled address。 + +等价地： + +```text +mask = (1 << B) - 1 + +low = m & ((1 << M) - 1) +x = m >> M +x2 = x ^ ((x >> S) & mask) + +addr = (x2 << M) | low +``` + +要让 layout well formed，`S` 必须至少为 `B`。 + +这个 transform 的目的不是改变 tile 中有哪些逻辑元素。它改变的是这些元素落在 shared memory 中的位置。MMA 仍然读取同一个逻辑 tile。Swizzle 让物理 bank pattern 更好。 + +## 选择 Swizzle 参数 + +正常使用时，swizzle 参数由 dtype 和 shared-memory swizzle mode 共同决定。常见 mode 是 32-byte、64-byte 和 128-byte swizzle。 + +`per_element` 参数的选择要保证一个小的 vector-sized group 保持连续。对于 float16，一个 16-byte vector 包含 8 个元素，因此： + +```text +M = log2(8) = 3 +``` + +使用 128-byte swizzle 时，layout 使用： + +```python +SwizzleLayout(per_element=3, swizzle_len=3, atom_len=3) +``` + +这会保持 16-byte vector group 完整，同时仍然充分重排更大的 shared-memory address pattern，以打破 column bank conflict。 + +大多数代码不应该手工推导这些参数。Dtype 和 descriptor mode 通常会决定它们。对程序员来说，重要的是 TIRx layout 中的 swizzle、TMA descriptor 和 MMA 期望三者保持匹配。 + +一个 swizzled shared memory allocation 因此会写成： + +```python +tile = TileLayout(S[(8, 64) : (64@m, 1@m)]) +swizzle = SwizzleLayout(per_element=3, swizzle_len=3, atom_len=3) + +layout = ComposeLayout(swizzle, tile) +``` + +组合后的 layout 会附着到 shared memory buffer 上。 + +## 元素的 Bank 和 Line + +要判断 swizzle 是否有帮助，可以把 swizzled element address 转回 shared memory bank。 + +设 `addr` 为 swizzled element address，`b` 为 element size in bytes。Byte address 是： + +```text +byte = addr * b +``` + +Bank 是： + +```text +bank = floor(byte / 4) mod 32 +``` + +128-byte bank line 是： + +```text +line = floor(byte / 128) +``` + +对于 float16，`b = 2`，因此 bank 公式变成： + +```text +bank = floor(addr / 2) mod 32 +``` + +这就是下面 worked example 使用的公式。 + +## Worked Example：`(8, 64)` float16 Tile 上的 128B Swizzle + +回到 row-major float16 tile： + +```text +m = 64 * i + j +``` + +使用： + +```python +SwizzleLayout(per_element=3, swizzle_len=3, atom_len=3) +``` + +Transform 变成： + +```text +x = m >> 3 +addr = ((x ^ ((x >> 3) & 7)) << 3) | (m & 7) +``` + +因为： + +```text +m = 64 * i + j +``` + +我们可以写： + +```text +q = floor(j / 8) +r = j mod 8 +``` + +swizzled address 是： + +```text +addr = 64 * i + 8 * (q xor i) + r +``` + +现在看 column `j = 0`。此时 `q = 0` 且 `r = 0`，因此： + +```text +addr = 72 * i +``` + +对于 float16，bank 是： + +```text +bank = floor(addr / 2) mod 32 +``` + +所以八行映射到： + +```text +i = 0: bank 0 +i = 1: bank 4 +i = 2: bank 8 +i = 3: bank 12 +i = 4: bank 16 +i = 5: bank 20 +i = 6: bank 24 +i = 7: bank 28 +``` + +这个 column 现在触碰八个不同 bank。Conflict 消失了。 + +如果没有 swizzling，同一个 column 的地址是： + +```text +m = 64 * i +``` + +因此： + +```text +bank = floor(64 * i / 2) mod 32 = 0 +``` + +每一行都落在 bank 0 上，所以访问会被串行化。Swizzle 只改变物理 placement，但这已经足以把 column access 变成 conflict-free。 + +这个保证依赖于按设计方式使用 swizzle。Dtype、swizzle width 和 access shape 必须匹配 TMA 和 MMA descriptor mode。128-byte float16 swizzle 围绕相关的 16-byte row chunk 和 Tensor Core access pattern 设计。它并不承诺任意 shared memory access 都会变成 conflict free。本章开头的 demo 可以看到这一点：选择 dtype 和 swizzle mode，观察无 swizzle 时一个 column 如何塌缩到一个 bank 上，再观察匹配 swizzle 应用后它如何分散到 bank view 中。 + +## 设计理由 + +Layout API 遵循三个设计选择。 + +第一，它支持通用 shape。硬件 tile 并不总是 2 的幂。Global tensor、shared memory stage、TMEM accumulator 和 scale-factor buffer 的 shape 常常来自容量限制或算法选择。Layout model 把这些 shape 当作正常情况处理。 + +第二，映射方向是从逻辑坐标到物理坐标。这个方向很重要，因为 replication 很常见。一个逻辑元素可能位于多个物理位置。Logical-to-physical map 可以直接把它表示成一组坐标。 + +第三，硬件轴是显式的。Layout 不使用匿名维度，再依赖上下文在事后解释它们。`tx`、`tid_in_wg`、`laneid`、`warpid`、`TLane` 和 `TCol` 之间的差异直接写在 layout 里。 -本页用于放置 TIRx layout API 相关内容的中文翻译。 +Legality 和 feasibility 检查不只由 layout object 自己负责。Layout 可以说明数据放在哪里。更高层的 tile primitive 会决定某个操作能否合法且高效地使用这个 placement。这种分离让 layout API 保持小而清晰，同时仍然给编译器足够信息去 dispatch 真实硬件操作。 diff --git a/zh/chapter_tma/index.md b/zh/chapter_tma/index.md index f45a98ab..d45eca45 100644 --- a/zh/chapter_tma/index.md +++ b/zh/chapter_tma/index.md @@ -1,66 +1,147 @@ (chap_tma)= # 异步数据搬运：TMA - - :::{admonition} 概览 :class: overview -- TODO：翻译本章 overview 第一条。 -- TODO：翻译本章 overview 第二条。 -- TODO：翻译本章 overview 第三条。 +- TMA 是在 global memory 和 shared memory 之间异步拷贝 tile 的硬件引擎。一个线程发起 copy，硬件引擎移动 byte。 +- 一次 TMA copy 由 tensor-map descriptor 描述。Descriptor 告诉引擎 global tensor shape、stride、tile coordinate，以及 shared-memory swizzle mode。 +- 在 load 路径上，TMA 可以在写 shared memory 时对 tile 应用 swizzle，使 tile 直接落到 Tensor Core 期望的布局中。 +- TMA load 通过带 byte-count tracking 的 `mbarrier` 完成。TMA store 使用 commit group 和 wait group。 ::: -TODO：翻译导言部分。 +只有当数据已经准备好供 Tensor Core 消费时，Tensor Core 才有用。在 GEMM 或 attention kernel 中，一旦 pipeline 被填满，数学部分可能是 compute-bound（{ref}`chap_performance`），但只有下一块 operand tile 及时到达，pipeline 才能一直保持填满。 + +移动 tile 的旧方式是让线程自己 copy。每个线程计算地址，从 global memory 发起 load，再把值 store 到 shared memory。这样当然可行，但它把 warp 指令花在地址计算和 copy bookkeeping 上，而不是计算上。它也让 copy 路径暴露在同一批本应喂给 Tensor Core 的 warp 的指令流中。 + +Tensor Memory Accelerator，简称 TMA，会把这项工作移动到硬件 copy engine 中。一个线程发起一次 tile copy，copy engine 随后在 global memory 和 shared memory 之间异步移动一个矩形 tile。当引擎移动 byte 时，CTA 的其他部分可以继续执行别的工作。 + +TMA 还会处理一部分布局问题。Tensor Core 不只是需要 shared memory 中有正确的值。它还需要这些值位于正确的 shared-memory layout 中。在 load 路径上，TMA 可以在写 tile 时应用 shared-memory swizzle。这样 tile 会直接落到后续 MMA 期望的布局中。 ```{raw} html +
+ style="width:100%; min-width:1320px; height:640px; border:1px solid var(--pst-color-border, #d0d0d0); border-radius:6px;"> +
``` - -*点击图中组件查看细节：TODO：翻译 TMA intro 图注。* +*交互图：TMA 把 tile 从 global memory 拷贝到 shared memory。切换 swizzle mode，并悬停在 source cell 上查看它在 shared memory 中落到哪里。* ## 一个线程发起，硬件搬运 Tile -TODO：翻译 “One Thread Issues, Hardware Moves the Tile” 小节。 +一次 TMA copy 从一个 issuing thread 开始。这个线程不会遍历 tile 中的所有元素。它会把 copy 的描述交给硬件，然后由 TMA engine 执行整个 transfer。 + +主要输入是 tensor-map descriptor。Descriptor 描述 global tensor，以及应该如何从中读取一个 tile。它会记录 tensor shape、stride、element size、tile shape 和 swizzle mode 等信息。Issuing thread 还会提供 tile 应该落到的 shared-memory address。 + +指令发出之后，copy 会异步运行。Issuing thread 可以继续执行。CTA 中的其他线程也可以继续执行。Transfer 现在由 TMA engine 负责，而不是由普通 load/store 指令组成的 loop 负责。 + +这让 kernel 有两种不同方式表达同一个逻辑操作：“copy 这个 tile”。 + +一条路径是 thread copy。线程协作从 global memory load，并 store 到 shared memory。这样 kernel 可以直接控制每一次访问，但会消耗线程指令和寄存器来做地址计算。 + +另一条路径是 TMA copy。一个线程发起 transfer，硬件 copy engine 执行矩形 copy。对于大的规则 tile，尤其是 Tensor Core kernel 使用的 operand tile，这是自然路径。 + +这两条路径有不同的同步规则和不同的性能行为。选择其中哪一条是一个 dispatch decision。Layout 告诉 kernel 想要什么内存排列。Scope 告诉它哪些线程或 CTA 参与。Dispatch 决定这次 copy 是由普通 thread copy 实现，还是由 TMA 实现。 ## Swizzled Layout -TODO：翻译 “Swizzled Layouts” 小节。 +移动 tile 本身还不够。Tile 还必须以 Tensor Core 能高效读取的布局放入 shared memory。 + +这就是 TMA swizzling 的用途。当 TMA 把 tile 写入 shared memory 时，它可以重排 shared-memory address pattern。Global memory tile 仍然是一个逻辑矩形，但 shared memory 中的 destination layout 可以是 swizzled 的。 + +Swizzle mode 是 TMA descriptor 的一部分。Descriptor 设置好之后，issuing thread 不需要手工应用 swizzle。Engine 会在 byte 落入 shared memory 时应用它。 + +重要要求是一致性。TMA descriptor、shared-memory tile layout，以及后续 MMA 指令，都必须描述同一个布局（{ref}`chap_data_layout`）。如果 TMA 用一种 swizzle 写入 tile，而 MMA 以为它是另一种 swizzle，硬件仍然会忠实执行收到的指令。只是这些 byte 对计算来说会排列错误。 + +这正是布局记号不再只是 bookkeeping 的地方。DSL 使用的布局必须匹配 TMA descriptor 和 Tensor Core 指令使用的硬件布局。例如，如果 kernel 说某个 operand tile 存储在 128-byte swizzled layout 中，TMA descriptor 就必须使用匹配的 swizzle mode，MMA dispatch 也必须期望同样的 shared-memory arrangement。上面的 demo 可以在 no swizzle 和 128-byte swizzle 之间切换；悬停在 source element 上可以看到 swizzle 应用后它落在哪里。 + +理解 swizzle 的一个有用方式是：TMA 并没有改变逻辑 tile。它改变的是逻辑元素在 shared memory 中的物理落点。后续 MMA 仍然消费同一个逻辑 A 或 B tile。Swizzle 只决定这个 tile 如何排列在 shared memory bank 上。 ## 用 3D TMA 表达 Tiling 和 Swizzling -TODO：翻译 “3D TMA for Tiling and Swizzling” 小节。 +普通 TMA copy 会移动一个扁平 2D tile，但 Tensor Core 期望的 shared-memory layout 通常会被 *tiled* 成 swizzle atom（来自 {ref}`chap_data_layout` 的 8 x 128-byte atom）。TMA 通过额外的 descriptor dimension 处理这一点。**3D TMA** 把 shared-memory box 描述成 `(group, row, col)`，其中 group 维度沿 atom 前进，内部两个维度在一个 atom 内寻址。一次 3D copy 随后既会 atom by atom 地布置 tile（tiling），又会在每个 atom 内应用 swizzle，因此数据到达时就已经在 MMA 期望的布局中，不需要单独的 tiling 或 swizzling pass。 ```{raw} html +
+ style="width:100%; min-width:1320px; height:640px; border:1px solid var(--pst-color-border, #d0d0d0); border-radius:6px;"> +
``` +*交互图：一次 3D TMA copy，以 (group, row, col) 寻址，并 tiled 到 swizzled shared memory 中。* -*点击图中组件查看细节：TODO：翻译 3D TMA 图注。* +选择 swizzle *format* 与这种 tiling 绑定在一起。更宽的 swizzle 会把一个 column 分散到更多 bank，所以只要能适用，128-byte swizzle 就是默认选择。但一个 N-byte atom 要求 tile 的 contiguous dimension 能填满它。因此，一个由于 shape constraint 而变小的 tile 不能使用 128-byte swizzle，必须降到 64-byte 或 32-byte：经验法则是选择 tile 能填满的最大 swizzle（{ref}`chap_data_layout`）。下面的 demo 直接展示这个约束：16 x 16 tile 上的 128-byte swizzle 只有在 tile 被切成匹配 atom 的 16 x 8 group 之后，才会变成 conflict-free。 ```{raw} html +
+ style="width:100%; min-width:1320px; height:640px; border:1px solid var(--pst-color-border, #d0d0d0); border-radius:6px;"> +
+ ``` - -*点击图中组件查看细节：TODO：翻译 tiling constraint 图注。* +*交互图：16 x 16 tile 上的 128-byte swizzle；切成 16 x 8 group 后变成 conflict-free。* ## 完成通知：Load -TODO：翻译 “Completion: Loads” 小节。 +Copy 是异步的，所以发起它还不够。Consumer 不能因为 TMA 指令已经发出就读取 shared-memory tile。只有在引擎写完 byte 之后，读取 tile 才是安全的。 + +对于 TMA load，完成信号是 `mbarrier`（{ref}`chap_async_barriers`）。 -![TODO：翻译 TMA load synchronization flow 图注](../../img/tma_sync_flow.png) +常见顺序是： + +1. 为 pipeline stage 初始化或复用一个 `mbarrier`； +2. 告诉 barrier 这次 TMA transfer 预计会写入多少 byte； +3. 发起 TMA load； +4. 让 TMA engine 在 byte 到达时更新 barrier； +5. consumer 在读取 shared-memory tile 之前等待这个 barrier phase。 + +Byte count 通过如下操作设置： + +```text +mbarrier.arrive.expect_tx(bytes) +``` + +它做两件事。第一，记录期望 transfer size。第二，它也执行 issuing thread 在 barrier 上的 arrival。Barrier 不会因为这个调用发生了就完成。它仍然等待 TMA engine 报告期望 byte 已经到达。 + +Transfer 进行时，引擎会对 barrier 执行 complete-tx 更新。只有两个条件都满足时，barrier phase 才会翻转：arrival count 已经满足，并且 pending byte count 到达 0。 + +Consumer 随后等待这个 barrier。对期望 phase 的 wait 完成后，shared-memory tile 就准备好了。此时 MMA 路径可以安全读取它。 + +![TMA load synchronization flow](../../img/tma_sync_flow.png) + +这与其他异步 producer-consumer 交接使用的是同一个 barrier 模型。Producer 是 TMA engine。Consumer 是 MMA 路径，或者任何其他读取 shared-memory tile 的代码。Barrier 是它们之间的显式交接。 ## 完成通知：Store -TODO：翻译 “Completion: Stores” 小节。 +TMA store 沿相反方向移动数据，从 shared memory 到 global memory。它们也是异步的，但完成机制不同。 + +TMA load 通常喂给同一个 kernel 内部的 consumer。MMA 路径需要知道 shared-memory tile 何时就绪。因此 load 路径使用 `mbarrier`。 + +TMA store 通常把最终数据写出到 global memory。通常没有立即在 kernel 内等待存储结果的 consumer。Kernel 主要需要知道的是，什么时候可以复用 shared-memory buffer，或者结束 store 序列。 + +为此，TMA store 使用 commit group 和 wait group。Kernel 发起一次或多次 store，commit 这个 group，随后等待这个 group drain。Wait 完成后，从 kernel 视角看，这个 group 中的 store 已经完成，store 使用的 shared-memory 区域可以安全复用。 + +规则很简单： + +```text +TMA load: wait through an mbarrier with byte-count tracking +TMA store: wait through a commit group and wait group +``` + +这两个机制在不同交接点服务于同一个目的。Load 需要让 shared-memory tile 对后续 consumer 可见。Store 需要确保 outgoing transfer 在 kernel 复用 source storage 或依赖 store 已 drain 之前完成。 ## 为什么 TMA 对 Pipelining 很重要 -TODO：翻译 “Why TMA Matters for Pipelining” 小节。 +TMA 最有用的场景是作为 pipeline 的一部分。Kernel 可以在 Tensor Core 计算当前 tile 时，发起未来 tile 的 load。Load 在后台运行。Compute 在前台运行。当未来 tile 变成当前 tile 时，barrier 把两者连接起来。 + +典型 GEMM loop 会反复使用这个结构。Shared memory 的一个 stage 保存当前被 MMA 消费的 tile。另一个 stage 正在被 TMA 填充。Loop 前进时，这些角色轮换。在 MMA 读取某个 stage 之前，它会等待该 stage 的 load barrier。在 TMA 覆盖某个 stage 之前，kernel 会确保上一个 consumer 已经使用完它。 + +这就是为什么 TMA 和 `mbarrier` 通常一起出现在 Blackwell 和 Hopper 风格的 kernel 中。TMA 给 kernel 一个异步 copy engine。Barrier 给 kernel 一个精确方式来知道 copy 出来的 byte 何时准备好。 diff --git a/zh/chapter_tmem/index.md b/zh/chapter_tmem/index.md index 99746ba0..a66aea83 100644 --- a/zh/chapter_tmem/index.md +++ b/zh/chapter_tmem/index.md @@ -1,35 +1,76 @@ (chap_tmem)= # 特殊内存：TMEM - - :::{admonition} 概览 :class: overview -- TODO：翻译本章 overview 第一条。 -- TODO：翻译本章 overview 第二条。 -- TODO：翻译本章 overview 第三条。 +- TMEM 是 Blackwell 上 `tcgen05` 使用的专用内存空间。它是每个 SM 上的二维暂存区，包含 128 个 Lane 行和最多 512 个 Col 列。 +- `tcgen05.mma` 会把 accumulator 写入 TMEM。Block-scaled MMA 也会用 TMEM 保存 scale factor。 +- TMEM 通过 Lane 和 Col 寻址。在 TIRx 的布局记号里，这两个硬件轴写作 `TLane` 和 `TCol`。 +- TMEM 不像寄存器那样自动分配。Kernel 必须以 32 列为单位显式分配和释放 TMEM。 +- 普通 shared-memory load/store 不能访问 TMEM。TMEM、寄存器和 shared memory 之间的数据搬运需要通过专用的异步 `tcgen05` 指令完成。 ::: -TODO：翻译导言部分。 +在 Hopper 以及更早的 GPU 上，Tensor Core（{ref}`chap_tensor_cores`）的 accumulator 位于寄存器中。这个模型很容易理解：MMA 指令产生一个寄存器 fragment，kernel 在计算阶段让这个 fragment 保持存活，epilogue 随后读取它、做类型转换，并把结果存出去。 + +问题在于寄存器压力。寄存器是固定的 per-thread 资源。随着 MMA tile 变大，accumulator fragment 也会变大。到一定程度之后，accumulator 会挤占线程需要保存的其他值。更大的 tile 有利于 Tensor Core 吞吐，但把整个 accumulator 都放在寄存器里，会让这些大 tile 更难使用。 + +Blackwell 改变了这段数据路径。`tcgen05` 的 accumulator 不必在整个计算阶段都留在寄存器里。相反，`tcgen05.mma` 会把 accumulator 写入 Tensor Memory，也就是 TMEM。TMEM 是早期 NVIDIA GPU 没有的内存空间。它是 SM 上的一个二维暂存区，形状是 128 个 Lane 行乘以最多 512 个 Col 列，并且作用域属于使用它的 CTA。 + +这个额外的内存空间让 Blackwell 可以支持更大的 Tensor Core tile，而不必把完整 accumulator 压到每个线程的寄存器里。但 TMEM 并不像寄存器那样自动存在。编译器不会把它当作普通寄存器存储直接分配给程序。Kernel 必须分配 TMEM，用正确的布局寻址，通过正确的指令搬入搬出，并在 CTA 完成后释放它。 ## 二维地址空间 -TODO：翻译 “A 2D Address Space” 小节。 +TMEM 不是一个扁平的 byte array。它是二维地址空间。硬件把两个坐标称为 Lane 和 Col。TMEM 有 128 个 Lane 行，最多 512 个 Col 列。每个 Col 是一个 32-bit 列。 + +这个形状很重要，因为 `tcgen05.mma` 会按照这个二维结构把 accumulator 写入 TMEM。一个 TMEM 位置由 Lane 坐标和 Col 坐标描述，而不是由一个类似 shared memory byte offset 的单一地址描述。 -![TODO：翻译 TMEM 2D grid 图注](../../img/tmem_grid.png) +当 kernel 在 TIRx 中声明 TMEM buffer 时，会给这个 buffer 一个覆盖这两个硬件坐标的布局。在布局记号（{ref}`chap_data_layout`）里，我们把 TMEM Lane 轴写作 `TLane`，把 TMEM Col 轴写作 `TCol`。这些名字不是要替代官方硬件术语，而是 DSL 中的布局轴名，用来显式表达 TMEM 的两个维度。 + +例如，一个 accumulator tile 可以写作： + +```text +S[(128, N) : (1@TLane, 1@TCol)] +``` + +这表示 tile 沿硬件 Lane 维度有 128 行，沿硬件 Col 维度有 `N` 列。在布局记号中，这两个维度分别是 `TLane` 和 `TCol`。这个布局是直接映射：相邻行沿 `TLane` 移动，相邻列沿 `TCol` 移动。下图展示了这个网格，其中硬件 Lane 沿 128 行向下，硬件 Col 沿列方向展开。 + +![TMEM as a 2D grid: TLane rows × TCol columns](../../img/tmem_grid.png) + +核心要点是：TMEM 是 tile 布局问题的一部分。它不是 Tensor Core 背后的隐藏存储。Kernel 必须命名这块内存，从中分配列，并使用与 `tcgen05` 指令读写方式匹配的布局。 ## 分配 -TODO：翻译 “Allocation” 小节。 +Kernel 使用 TMEM 之前，必须先预留空间。这一点不同于寄存器。寄存器由编译器分配，而 TMEM 由 kernel 显式分配。 + +分配按 CTA 进行。CTA 中的一个 warp 会请求一段 TMEM 列。请求以 32 列为单位，列数会按照硬件分配规则向上取整。分配完成后，CTA 会得到一个 base TMEM address。后续 `tcgen05` 指令用这个 base address 访问预留区域。 + +把 TMEM 理解成一种有预算的 CTA 资源很有用，它类似 shared memory。CTA 拥有它分配到的 TMEM 列。Kernel 决定 accumulator、scale factor 或临时 staging 需要多少列。CTA 完成后，必须释放这段分配。 + +因此，TMEM 也是 kernel 资源规划的一部分。更大的 accumulator tile 可能提高 Tensor Core 吞吐，但会消耗更多 TMEM 列。Block-scaled MMA 可能还需要额外 TMEM 空间来保存 scale factor。Kernel 必须让这些用途都落在可用的 TMEM 预算内，就像必须让 shared-memory buffer 落在 SMEM 预算内一样。 ## 读写 TMEM -TODO：翻译 “Reading and Writing TMEM” 小节。 +普通 `ld.shared` 和 `st.shared` 指令不能访问 TMEM。TMEM 是单独的地址空间，因此数据必须通过专用 `tcgen05` 指令移动。 + +主要有三条路径。 + +第一条路径是 `tcgen05.ld`，它把数据从 TMEM 加载到寄存器中。这是 MMA 阶段之后 epilogue 使用的路径。Accumulator 已经在 TMEM 中产生，但 epilogue 通常需要寄存器 fragment，才能做类型转换、执行 elementwise 操作并写出最终结果。 + +在 DSL 层面，TMEM load 分布在一个 warpgroup 上。它会 lowering 成四个 warp-level 的 `tcgen05.ld` 操作，每个 warp 一个。每个 warp 处理 128 个 TMEM Lane 行中的 32 行，因此四个 warp 合起来覆盖完整的 Lane 维度。在布局记号里，这个完整维度就是 `TLane` 轴。 + +指令本身来自一组 load shape，例如 `.16x64b`、`.16x128b`、`.16x256b`、`.32x32b` 和 `.16x32bx2`，并带有从 `.x1` 到 `.x128` 的 repeat factor。选择的 shape 决定读取多少个 TMEM 列，以及每个线程会收到多少个寄存器。 + +重要结果是寄存器 fragment 的布局。对于常见的 epilogue 路径，lane `l` 会收到来自 TMEM 行 `l / 4` 以及两个列位置的值。这会产生与早期世代从 MMA 直接暴露出的 per-lane accumulator fragment 同类的布局（{ref}`chap_layout_generations`）。这种连续性很重要：虽然 Blackwell 的 accumulator 在计算阶段位于 TMEM 中，但 epilogue 仍然可以复用 Ampere `mma` 或 Hopper `wgmma` 中已经使用过的寄存器级 cast 和 store 结构。 + +![tcgen05.ld / st move the TMEM accumulator to and from registers in the m8n8 fragment (lane l → row l/4, two columns)](../../img/tcgen05_ldst.svg) + +第二条路径是 `tcgen05.st`，它把数据从寄存器写回 TMEM。这是 `tcgen05.ld` 的反方向。当线程已经持有一个寄存器 fragment，并且需要把它放入 TMEM 时会使用这条路径。例如，某些操作数或中间值可能会先经过寄存器 staging，再写入 TMEM 供后续 `tcgen05` 操作使用。 + +第三条路径是 `tcgen05.cp`，它把数据从 shared memory 拷贝到 TMEM。这是一条 bulk copy 路径，常用于 block-scaled MMA 的 scale factor。在这种情况下，TMA 或普通线程代码会先把 scale 数据准备到 shared memory 中，然后 `tcgen05.cp` 把它移动到 Tensor Core 期望的 TMEM 布局里。 + +这三条路径都是异步的。`tcgen05.ld`、`tcgen05.st` 或 `tcgen05.cp` 指令可能在数据搬运完成之前就返回。因此，kernel 在消费结果或复用存储之前，必须使用正确的完成机制（{ref}`chap_async_barriers`）。 + +等待路径取决于具体指令。`tcgen05.ld` 通过 `tcgen05.wait::ld` 完成。`tcgen05.st` 通过 `tcgen05.wait::st` 完成。`tcgen05.cp` 和 `tcgen05.mma` 一样，通过 commit group 和 `mbarrier` 完成。如果数据需要从一组线程交给另一组线程，kernel 还可能需要 fence，确保接收方线程按预期顺序看到已完成的写入。 -![TODO：翻译 tcgen05.ld/st 图注](../../img/tcgen05_ldst.svg) +TMEM 位于 Blackwell Tensor Core 数据路径的中间。TMA 把操作数 staged 到 shared memory。`tcgen05.mma` 读取操作数并累加到 TMEM。对于 block-scaled MMA，scale factor 也可以 staged 到 TMEM。计算阶段结束后，`tcgen05.ld` 把 accumulator 带回寄存器，epilogue 再转换并存出最终输出。 diff --git a/zh/conf.py b/zh/conf.py index 06ccfb07..70f80e5f 100644 --- a/zh/conf.py +++ b/zh/conf.py @@ -28,23 +28,6 @@ "**/README.md", "_*.md", "**/_*.md", - # Release the Chinese edition chapter by chapter. Keep draft sources in - # zh/, but exclude unreleased pages so they are not published or searchable. - "appendix/**", - "chapter_async_barriers/**", - "chapter_clc/**", - "chapter_data_layout/**", - "chapter_flash_attention/**", - "chapter_gemm_advanced/**", - "chapter_gemm_async/**", - "chapter_gemm_basics/**", - "chapter_intro_tirx/**", - "chapter_layout_generations/**", - "chapter_tensor_cores/**", - "chapter_tirx_layout_api/**", - "chapter_tma/**", - "chapter_tmem/**", - "tirx_guide/**", ] html_theme = "sphinx_book_theme" diff --git a/zh/index.md b/zh/index.md index 402b4224..9cd9b6ca 100644 --- a/zh/index.md +++ b/zh/index.md @@ -19,18 +19,51 @@ - **第四部分：Flash Attention 4。** 这一部分基于第三部分的技术构建完整的 attention kernel：两个 MMA，中间插入 softmax，并包含 online-softmax rescaling、causal masking 和 GQA。 - **附录。** TIRx API 和编译器内部机制说明。 -## 已发布章节 - ```{toctree} :caption: 第一部分：理解 GPU :maxdepth: 1 chapter_background/index chapter_performance/index +chapter_data_layout/index +chapter_layout_generations/index +chapter_tma/index +chapter_tensor_cores/index +chapter_tmem/index +chapter_async_barriers/index +chapter_clc/index +``` + +```{toctree} +:caption: 第二部分：TIRx 概览 +:maxdepth: 1 + +chapter_intro_tirx/index +chapter_tirx_layout_api/index +``` + +```{toctree} +:caption: 第三部分：GEMM：从 Tiled 到 SOTA +:maxdepth: 2 + +chapter_gemm_basics/index +chapter_gemm_async/index +chapter_gemm_advanced/index +``` + +```{toctree} +:caption: 第四部分：Flash Attention 4 +:maxdepth: 2 + +chapter_flash_attention/index ``` - +```{toctree} +:caption: 参考资料 +:maxdepth: 1 + +appendix/index +appendix/debugging_warp_specialized +tirx_guide/arch/index +tirx_guide/language_reference/index +``` diff --git a/zh/tirx_guide/arch/index.md b/zh/tirx_guide/arch/index.md index 9dcd21f3..faeee895 100644 --- a/zh/tirx_guide/arch/index.md +++ b/zh/tirx_guide/arch/index.md @@ -1,10 +1,10 @@ ---- -orphan: true ---- - (chap_arch)= # 编译器内部机制 -> 翻译状态：待翻译。对应英文目录：`tirx_guide/arch/`。 +面向贡献者的 TIRx 编译器内部机制说明。 + +```{toctree} +:maxdepth: 1 -本页用于放置 TIRx 编译器内部机制的中文参考内容。 +lowering_pipeline +``` diff --git a/zh/tirx_guide/arch/lowering_pipeline.rst b/zh/tirx_guide/arch/lowering_pipeline.rst new file mode 100644 index 00000000..c0204673 --- /dev/null +++ b/zh/tirx_guide/arch/lowering_pipeline.rst @@ -0,0 +1,158 @@ +TIRx lowering pipeline +====================== + +``tvm.compile(mod, target, tir_pipeline="tirx")`` 会把作者写出的 TIRx module 送入 **tirx pipeline**。这是一组有序的 TIR pass，会把你写下的高层构造（tile primitive、带 ``TileLayout`` 的 buffer、execution-scope id）变成拆分后的 **host** + **device** function，然后由 CUDA backend 渲染成源码。Pipeline 定义在 ``python/tvm/tirx/compilation_pipeline.py``（``tirx_pipeline``）中；本页按顺序介绍这些 pass。 + +它位于哪里 +---------- + +``tvm.compile`` 会先绑定 target，运行 **tirx pipeline**（下面这些 module-level pass），然后分别对 host 和 device function 应用 **finalization** pass，最后把每个 device function 交给 CUDA code generator： + +.. code-block:: text + + authored TIRx ──BindTarget──▶ tirx_pipeline ──▶ host func ──host finalize──▶ C/LLVM + │ + └──────────▶ device func ──device finalize──▶ CUDA + +这些 Pass +--------- + +``tirx_pipeline`` module pass 会应用下面这个精确序列（少数 pass 受 ``PassContext`` config 控制）： + +.. list-table:: + :header-rows: 1 + :widths: 6 32 62 + + * - # + - Pass + - 作用 + * - 1 + - ``LowerTIRx`` + - 核心 lowering，见下方 `Inside LowerTIRx`_ + * - 2 + - ``UnifyThreadBinding`` + - 合并等价的 thread-axis binding，使每个 ``threadIdx`` / ``blockIdx`` 轴只声明一次 + * - 3 + - ``StmtSimplify`` + - statement-level 算术简化（arith analyzer） + * - 4 + - ``LowerTIRxOpaque`` + - 把剩余 opaque TIRx 构造 lower 成普通 TIR + * - 5 + - ``FlattenBuffer`` + - 把多维 ``BufferLoad`` / ``BufferStore`` 展平成 1-D + * - 6 + - ``BF16ComputeLegalize`` + - 把 ``bfloat16`` compute 重写成合法形式（f32 up-cast） + * - 7 + - ``NarrowDataType(32)`` + - 在可证明安全时把 index/loop ``PrimExpr`` dtype 缩窄到 32-bit + * - 8 + - ``VectorizeLoop`` + - 把 ``T.vectorized`` loop 变成 vector op（如果 ``tir.disable_vectorize`` 则跳过） + * - 9 + - ``UnrollLoop`` + - 展开标记为 ``T.unroll`` 的 loop（以及小的 constant loop） + * - 10 + - ``StmtSimplify`` + - 在 vectorize/unroll 暴露常量后再次简化 + * - 11 + - ``CommonSubexprElim`` + - 把重复 subexpression hoist 到临时变量（如果 ``tir.disable_cse_tir`` 则跳过） + * - 12 + - ``FP8ComputeLegalize`` + - 把 ``float8`` compute 重写成合法形式 + * - 13 + - ``VerifyMemory`` + - 检查 host-side code 不会直接 dereference device memory（安全门） + * - 14 + - ``AnnotateEntryFunc`` + - 把单个 PrimFunc 标记为 module entry point + * - 15 + - ``SplitHostDevice`` + - 在 ``launch_thread`` 边界把每个 kernel 拆成 **host** function 和 **device** function + * - 16 + - ``MakePackedAPI`` + - 把 host function 重写成 packed-func ABI（TVM 调用的 launcher） + * - 17 + - ``FP8StorageLegalize`` + - legalize ``float8`` storage（打包进支持的 container type） + * - 18 + - ``BF16StorageLegalize`` + - legalize ``bfloat16`` storage + +**Finalization** 随后按 function kind 分别运行： + +- **host**：``LowerTVMBuiltin`` （lower ``tvm_*`` builtin）、``LowerIntrin`` （target-specific intrinsic） +- **device**：``LowerWarpMemory`` （warp-scoped buffer -> shuffle）、``StmtSimplify``、``LowerIntrin`` + +Inside LowerTIRx +---------------- + +``LowerTIRx`` 本身是一个小序列（``src/tirx/transform/lower_tirx.cc``）： + +.. code-block:: text + + LowerTIRx = Sequential([ TilePrimitiveDispatch, LowerTIRxCleanup ]) + +- **``TilePrimitiveDispatch``** 会把每个 ``TilePrimitiveCall``（``copy``、``gemm``、``reduction`` 等）替换成其选定 backend dispatch 生成的 body，也就是 variant selection 和 codegen。 +- **``LowerTIRxCleanup``** 会运行 ``LayoutApplier``：把每个带 ``TileLayout`` 的 buffer access 解析成具体物理地址算术（``addr = data + elem_offset + layout.apply(coord)``），展平 buffer，并把 execution-scope id（``T.cta_id`` / ``T.thread_id`` 等）lower 成 ``launch_thread`` 上的 ``blockIdx`` / ``threadIdx``。 + +因此，``LowerTIRx`` 之后 module 已经是普通 TIR：没有 tile primitive，没有 ``TileLayout`` 间接层，scope id 已解析成 thread axis。 + +一个 worked example +------------------- + +看一个一行 scale kernel： + +.. code-block:: python + + @T.prim_func + def scale(A_ptr: T.handle, B_ptr: T.handle): + A = T.match_buffer(A_ptr, (256,), "float32") + B = T.match_buffer(B_ptr, (256,), "float32") + T.device_entry(); bx = T.cta_id([1]); tx = T.thread_id([256]) + B[tx] = A[tx] * T.float32(2.0) + +**``LowerTIRx`` 之后**，scope id 变成真实 thread axis，layout 已经应用（``A_1`` / ``B_1`` 是展平后的 1-D view）： + +.. code-block:: python + + with T.launch_thread("blockIdx.x", 1) as blockIdx_x: + threadIdx_x = T.launch_thread("threadIdx.x", 256) + bx: T.let = blockIdx_x + tx: T.let = threadIdx_x + B_1[threadIdx_x] = A_1[threadIdx_x] * T.float32(2.0) + +**``SplitHostDevice`` + ``MakePackedAPI`` 之后**，一个 function 变成两个：host launcher 和 device kernel： + +.. code-block:: python + + @I.ir_module + class Module: + def main(...): # host: packed-API launcher (computes the grid/block, launches) + ... + def scale_kernel(...): # device: the __global__ body, run on the GPU + +CUDA backend 随后把 ``scale_kernel`` 渲染成 ``__global__`` function（``B_ptr[threadIdx.x] = A_ptr[threadIdx.x] * 2.0f``）。 + +自己复现 +-------- + +你可以手动运行 pipeline 的任意前缀来检查某个阶段。这些文档中的 IR snippet 就是这样生成的： + +.. code-block:: python + + from tvm.tirx import transform as TT + + target = tvm.target.Target("cuda") + mod = TT.BindTarget(target.with_host("llvm"))(tvm.IRModule({"main": scale})) + mod = TT.LowerTIRx()(mod) # tile primitives dispatched, layouts applied + print(mod.script()) # inspect the lowered TIRx IR + +也可以编译整个 module 并读取生成的 CUDA： + +.. code-block:: python + + exe = tvm.compile(tvm.IRModule({"main": scale}), target=target, tir_pipeline="tirx") + print(exe.mod.imports[0].inspect_source()) diff --git a/zh/tirx_guide/language_reference/cuda/buffers.rst b/zh/tirx_guide/language_reference/cuda/buffers.rst new file mode 100644 index 00000000..27377b92 --- /dev/null +++ b/zh/tirx_guide/language_reference/cuda/buffers.rst @@ -0,0 +1,366 @@ +Buffer 与 memory +================ + +参数 buffer 使用 ``T.match_buffer`` 绑定；临时 scratch buffer 在函数体中使用下面两类声明 API 创建。用 ``A[i, j]`` 索引 buffer，用 ``A[m0:m0+BM, 0:BK]`` 切片（得到 ``BufferRegion``），并用 ``A.ptr_to([i, j])`` 取得 pointer，或用 ``A.data`` 取得原始 data pointer。 + +声明 buffer +----------- + +有两个基础 API 会创建 buffer： + +- ``T.alloc_buffer(shape, dtype, scope=..., ...)`` — 分配新的 storage（生成 ``AllocBuffer`` 节点），并返回 ``Buffer``。``T.alloc_shared`` / ``T.alloc_local`` 只是带有 ``scope="shared"`` / ``scope="local"`` 的 ``alloc_buffer``。 +- ``T.decl_buffer(shape, dtype, data=..., ...)`` — 在已有 pointer ``data`` 上声明一个 view（不分配）；用于给 storage 起别名或重新解释 storage，例如 pool 的某个子区域，或一个 tensor-memory address。当 ``data=None`` 时，它会像 ``alloc_buffer`` 一样分配。 + +Buffer 的 ``data`` pointer 是一个 immutable ``Var``（``alloc_buffer`` 会定义它；``decl_buffer`` 会接收它）。如果要用一个 pointer *expression* 作为 buffer backing，需要先把它绑定起来，参见 :doc:`data_types`。 + +二者共用同一类 descriptor；最重要的参数如下： + +.. list-table:: + :header-rows: 1 + :widths: 28 72 + + * - 参数 + - 含义 + * - ``dtype`` + - element type，例如 ``"float32"``、``"float16"``、``"float4_e2m1fn"`` 等 + * - ``shape`` + - 逻辑 shape（一组 extent） + * - ``layout`` + - 物理映射（:ref:`TileLayout `）；``"default"`` 表示 dense row-major + * - ``elem_offset`` / ``allocated_addr`` + - ``elem_offset``（或 ``byte_offset``）把一个 *view* 放到 ``data`` 中的某个 offset；``allocated_addr`` 携带预先分配的 address（tensor memory） + * - ``align`` + - data pointer 的对齐，以 byte 为单位 + +``scope`` 参数选择 memory space： + +.. list-table:: + :header-rows: 1 + :widths: 26 22 52 + + * - Scope + - 简写 + - Memory + * - ``"global"`` + - （默认） + - device global memory + * - ``"shared"`` + - ``T.alloc_shared`` + - static shared memory（``__shared__``） + * - ``"shared.dyn"`` + - （pool） + - dynamic shared memory（pooled，见下文） + * - ``"local"`` + - ``T.alloc_local`` + - per-thread register + * - ``"tmem"`` + - （TMEM pool） + - Blackwell tensor memory（见下文） + +.. code-block:: python + + A = T.match_buffer(A_ptr, (M, K), "float16", align=16) # parameter buffer + As = T.alloc_shared((BM, BK), "float16") # new shared tile + acc = T.alloc_local((4,), "float32") # register accumulator + view = T.decl_buffer((BM, BK), "float16", data=As.data) # a view over As + +**基于 pointer 的 buffer 只是 pointer 之上的 metadata。** 对任何非 tmem buffer 而言，声明就是一个 pointer 加一个 layout，索引会解析成 address:: + + addr(buffer[coord]) = buffer.data + elem_offset + layout.apply(coord, shape=shape)["m"] + +（``layout.apply`` 返回逐轴映射；其中 ``"m"`` 分量是 element offset。）因此，同一个逻辑访问会完全根据 buffer metadata 编译成不同的 address arithmetic。在 4×8 区域上写 ``B[i, j] = A[i, j] + 1``，如果用四种方式声明 ``B``： + +.. code-block:: python + + from tvm.tirx.layout import TileLayout, S + + B = T.match_buffer(p, (4, 8), "float32") # row-major + B = T.match_buffer(p, (4, 8), "float32", layout=TileLayout(S[(4, 8):(1, 4)])) # column-major + B = T.match_buffer(p, (4, 8), "float32", elem_offset=64) # shifted view + B = T.match_buffer(p, (4, 8), "float32", layout=TileLayout(S[(4, 8):(16, 1)])) # row stride 16 + +那么每一种都会让 ``B[i, j]`` lower 成生成 CUDA 中不同的 index（``A[i, j]`` load 仍然是 ``i*8 + j``，只有 ``B`` 的 metadata 发生了变化）： + +.. code-block:: c++ + + B_ptr[((i * 8) + j)] = ...; // row-major: i*8 + j + B_ptr[((j * 4) + i)] = ...; // column-major: j*4 + i + B_ptr[(((i * 8) + j) + 64)] = ...; // elem_offset=64: i*8 + j + 64 + B_ptr[((i * 16) + j)] = ...; // row stride 16: i*16 + j + +Shared memory +------------- + +Shared memory 有两种形式：static（编译时固定）和 dynamic（launch 时指定大小），此外还有一个 pool helper 用于管理 dynamic 情况。 + +Static +~~~~~~ + +最简单的 shared buffer 是 **static** buffer，也就是 ``T.alloc_shared``（即 ``scope="shared"``），其大小在编译时确定。把数据 stage 到其中，执行 ``cta_sync`` 让整个 block 看到写入，然后再读回： + +.. code-block:: python + + @T.prim_func + def smem_demo(A_ptr: T.handle, B_ptr: T.handle): + A = T.match_buffer(A_ptr, (128,), "float32") + B = T.match_buffer(B_ptr, (128,), "float32") + T.device_entry() + bx = T.cta_id([1]) + tx = T.thread_id([128]) + sm = T.alloc_shared((128,), "float32") # static shared memory + sm[tx] = A[tx] + T.cuda.cta_sync() + B[tx] = sm[tx] * T.float32(2.0) + +它会 lower 成普通的 ``__shared__`` array（生成 CUDA，省略 boilerplate）： + +.. code-block:: c++ + + extern "C" __global__ void __launch_bounds__(128) + smem_demo_kernel(float* __restrict__ A_ptr, float* __restrict__ B_ptr) { + int tx = ((int)threadIdx.x); + __shared__ alignas(64) float sm_ptr[128]; // T.alloc_shared + sm_ptr[tx] = A_ptr[tx]; + __syncthreads(); // T.cuda.cta_sync() + B_ptr[tx] = sm_ptr[tx] * 2.0f; + } + +Dynamic +~~~~~~~ + +**Dynamic** shared memory（``scope="shared.dyn"``）的大小按 launch 设置（``sharedMemBytes`` launch parameter），而不是在编译时固定。一个 kernel 只能有 **一个** dynamic-shared allocation，也就是 *arena*。因此需要分配一次 arena，再用 ``T.decl_buffer`` 和 arena pointer 加 ``elem_offset`` 把每个 buffer 声明成其中的一个 view： + +.. code-block:: python + + arena = T.alloc_buffer((128,), "float32", scope="shared.dyn") # the one arena + As = T.decl_buffer((64,), "float32", data=arena.data, scope="shared.dyn") # offset 0 + Bs = T.decl_buffer((64,), "float32", data=arena.data, elem_offset=64, scope="shared.dyn") # offset 64 + As[tx] = A[tx] + Bs[tx] = B[tx] + T.cuda.cta_sync() + C[tx] = As[tx] + Bs[tx] + +两个 view 共享同一个 ``extern __shared__`` arena（生成 CUDA，省略 boilerplate；这里为清晰起见把 arena 命名为 ``smem``）： + +.. code-block:: c++ + + extern __shared__ __align__(64) float smem[]; // the one dynamic-shared arena + smem[tx] = A_ptr[tx]; // As — view at offset 0 + smem[tx + 64] = B_ptr[tx]; // Bs — view at offset 64 + __syncthreads(); + C_ptr[tx] = smem[tx] + smem[tx + 64]; + +（两个独立的 ``alloc_buffer(scope="shared.dyn")`` 调用是错误的，**只允许一次 dynamic shared memory allocation**。）所以，static shared memory 的大小在编译时确定（``__shared__ T x[N];``）；dynamic shared memory 则是这个按 launch 指定大小的唯一 arena，其他 buffer 是在其中以 offset 声明出来的 view。 + +.. note:: + + **TVM 如何标注 dynamic-shared 大小。** Arena 的大小在编译时已知（这里是 ``128`` 个 float，即 ``512`` byte）。lowering 时，TVM 会向 device kernel 的 ``tirx.kernel_launch_params`` 追加一个 ``"tirx.use_dyn_shared_memory"`` tag，host launcher 会计算总 byte 数，并作为最后一个 launch argument 传入： + + .. code-block:: python + + # device kernel attribute: + "tirx.kernel_launch_params": ["blockIdx.x", "threadIdx.x", "tirx.use_dyn_shared_memory"] + + # host-side launch call (..., gridDim.x, blockDim.x, dyn_shared_bytes): + T.call_packed("dyn_kernel", A.data, B.data, C.data, 1, 64, 512) + + 运行时这个 ``512`` 会变成 ``cuLaunchKernelEx`` 调用中的 ``config.sharedMemBytes``。你不需要手动设置它；它由 ``shared.dyn`` allocation 的大小推导出来。 + +Pool sugar +~~~~~~~~~~ + +``T.SMEMPool`` 会自动处理 arena bookkeeping：它以 bump allocation 的方式分配 offset，因此不需要手写 ``decl`` view。除了 ``alloc`` / ``commit`` 之外，它还提供每个 buffer 的 ``align=``、``alloc_mma`` helper（自动构造 MMA 兼容的 swizzle layout），以及 ``move_base_to``，用于回退 cursor 并复用空间： + +.. code-block:: python + + pool = T.SMEMPool() # bump allocator over shared.dyn + As = pool.alloc((BM, BK), "float16", align=128) # carve a tile + Bs = pool.alloc((BK, BN), "float16", align=128) + Cs = pool.alloc_mma((BM, BN), "float16") # MMA-compatible, swizzle inferred + pool.commit() # finalize the pool's size + # pool.move_base_to(offset) rewinds the cursor to reuse space + +TMEM pool（见下文 `Tensor memory`_）建立在 ``SMEMPool`` 之上。 + +Registers +--------- + +Per-thread scratch 位于 register 中。用 ``T.alloc_local(shape, dtype)``（即 ``scope="local"``）分配它：它对每个 thread 私有，并 lower 成保存在 register 中的 local array。 + +.. code-block:: python + + r = T.alloc_local((4,), "float32") # per-thread register array + for k in T.unroll(4): + r[k] = A[tx, k] + # ... compute on r[0..3] ... + +.. code-block:: c++ + + alignas(64) float r_ptr[4]; // per-thread, register-resident + r_ptr[0] = A_ptr[tx * 4 + 0]; + r_ptr[1] = A_ptr[tx * 4 + 1]; + // ... + +.. note:: + + ``alignas(64)`` 是 *默认* buffer alignment：buffer 的 ``data_alignment`` 默认是 ``runtime::kAllocAlignment``（64 byte），CUDA codegen 会把它标到每个 allocation 上，包括 per-thread ``local`` array，即使这里没有实际意义。对这些 register-resident array 来说，它 **没有性能影响**：带有静态可解析 index 的 thread-local array 会被 nvcc/ptxas 提升到 register 中（scalar replacement of aggregates, SROA），因此它永远不会存在于可寻址的 local memory 中，alignment 也就是 no-op。（如果动态索引 array spill 到 local memory，它确实会带上这个过度对齐，但这不是常见情况。）register local 的这种过度对齐是一个已知粗糙点，我们计划修复（对 ``local`` scope 使用 dtype 的自然对齐）。 + +Scalar +~~~~~~ + +Scalar 只是只有 **一个 element** 的 register array；严格来说，不需要单独的概念。你可以分配一个大小为 1 的 ``local`` buffer 并索引 ``[0]``： + +.. code-block:: python + + phase = T.alloc_local((1,), "int32") # 1-element register array + phase[0] = 0 + while phase[0] < 4: + acc = acc + A[tx, phase[0]] + phase[0] += 1 + +但到处写 ``phase[0]`` 很笨重，所以 **scalar** 正是这件事的语法糖：一个单元素 register buffer，可以 **按名字** 读写： + +.. code-block:: python + + phase: T.int32 = 0 # mutable scalar (sugar for the above) + while phase < 4: + acc = acc + A[tx, phase] + phase += 1 + + s = T.local_scalar("int32") # explicit form; assign by name (s = ..., not s[0]) + acc: T.float32 = 0.0 # a type-annotated assignment also makes one + +二者不只是相似，而是会 parse 成 **结构完全相同的 TIRx**。这个语法糖完全在 parser 中消解：``phase: T.int32`` *就是* 那个单元素 ``local`` buffer，``phase`` / ``phase += 1`` *就是* ``phase[0]`` / ``phase[0] += 1``。对两个 kernel 调用 ``tvm.ir.assert_structural_equal`` 会通过，printer 甚至会把显式的 ``alloc_local`` + ``[0]`` 形式 **重新打印** 成 scalar 形式。因此，一旦 parsing 完成，二者完全没有区别。二者都会 lower 成同一个 ``alignas(64) int phase_ptr[1];``；scalar 只是让你省掉 ``[0]``。（``T.local_scalar`` / ``T.shared_scalar`` / ``T.alloc_scalar`` 会显式选择 scope。） + +.. note:: + + 为什么不用 ``Var`` ？TIRx ``Var`` 是 *immutable* 的，也就是一次性的静态绑定（正是下面 ``T.let`` 产生的东西）。Scalar 需要是 *mutable* 的：你会在 loop 和 accumulator 中重新赋值。因此它必须由可以反复 store 的单元素 buffer backing，而不是由 ``Var`` backing。 + +``let`` +~~~~~~~ + +``T.let`` binding 是 **immutable** 的，也就是一个 ``LetStmt`` （一个具名值，不是 buffer）。用它表示派生常量： + +.. code-block:: python + + n: T.let = M * K # immutable binding (LetStmt) + half: T.let[T.int32] = N // 2 # ... with an explicit type + +它会 lower 成 **普通的 scalar C variable**，而不是 buffer（没有 array，也没有 ``[0]``）。例如 ``half: T.let = m * 2``（其中 ``m`` 是 runtime 值）： + +.. code-block:: c++ + + int half = m * 2; // the `let` -> a const-like local + +由于值是 immutable 的，simplifier 可以自由传播它并对它做 CSE，所以在使用处你经常会看到 ``m * 2`` 被直接替换进去（或通过 common-subexpression 临时变量共享），而不是看到对 ``half`` 的引用。 + +.. note:: + + **为什么需要 immutable binding？** 因为值不能改变，arithmetic analyzer 会把 var 绑定到该值（它简化 ``LetStmt`` 时会执行 ``analyzer.Bind(var, value)``），所以关于该值证明出来的事实，包括常量边界、modular set（divisibility / alignment）和 range，都会 **传播到每一次使用**。这会帮助 index simplification、bounds-check elimination，以及 alignment/vectorization 决策。*Mutable* scalar 是一次 memory load（``buf[0]``）：analyzer 不能假设它保持不变，因此这些属性都无法传递。``let`` 也是一个纯 value：没有 allocation，并且可以自由 inline、substitute 或 CSE；而 scalar 是带有 load/store 语义的单元素 buffer。 + +Tensor memory +------------- + +Blackwell *tensor memory* 不是普通的 scratch scope：它必须通过 warp-uniform 的 ``T.ptx.tcgen05.alloc`` / ``tcgen05.dealloc`` intrinsic 显式 reserve 和 free；每个 tensor 都是其中的一个 view，通过 ``T.decl_buffer(..., scope="tmem", allocated_addr=, layout=)`` 声明。``allocated_addr``（column offset）是必需的，tensor-core dispatch 会断言它存在，因此 ``T.alloc_buffer(scope="tmem")``（它 **不会** 设置该字段）不能工作。与 shared memory 不同，tensor memory 不能被直接寻址：只能通过 ``tcgen05`` ``mma`` / ``ld`` / ``st`` / ``cp`` 读写。 + +手写时，一个 warp 把 allocation 发到 shared slot 中，你用 column offset 把每个 tensor ``decl`` 成一个 view，最后由一个 warp free 它： + +.. code-block:: python + + addr = T.alloc_shared((1,), "uint32") # slot for the allocated base + if warp_id == alloc_warp: # tcgen05.alloc is warp-uniform + T.ptx.tcgen05.alloc(T.address_of(addr), n_cols=512, cta_group=cta_group) + acc = T.decl_buffer((CTA_M, 512), "float32", scope="tmem", + allocated_addr=0, layout=tmem_layout) # view at column 0 + # ... use acc as a gemm_async / copy_async operand ... + if warp_id == alloc_warp: + T.ptx.tcgen05.relinquish_alloc_permit(cta_group=cta_group) + T.ptx.tcgen05.dealloc(addr, n_cols=512, cta_group=cta_group) + +你需要自己管理 column offset 和 ``tmem_layout`` （一个 datapath D/F layout）。这正是下面 pool 会生成的序列。 + +Pool +~~~~ + +``T.TMEMPool`` 会封装上述所有工作：warp-uniform alloc/dealloc、column bump-allocation，以及 datapath layout： + +.. code-block:: python + + tmem_addr = pool.alloc((1,), "uint32") # pool = the kernel's smem pool + tmem_pool = T.TMEMPool(pool, total_cols=512, cta_group=cta_group, + tmem_addr=tmem_addr) + acc = tmem_pool.alloc((CTA_M, 512), "float32") # allocated_addr set for you + tmem_pool.commit() # emits tcgen05.alloc (one warp) + # ... use acc ... + tmem_pool.dealloc() # emits tcgen05.dealloc (one warp) + +完整示例见第三部分的 GEMM kernel。 + +Buffer API +---------- + +``Buffer`` 是 pointer 之上的 metadata（见上文 *声明 buffer*），因此它的大部分方法都是 *compile-time* reshape/reinterpret：它们改变 index arithmetic，或给你一个 pointer，本身不会生成运行时操作。常用方法如下： + +.. list-table:: + :header-rows: 1 + :widths: 34 66 + + * - 方法 + - 含义 + * - ``B.data`` + - 原始 data pointer（一个 ``Var``）；打印为 ``B_ptr`` + * - ``B.ptr_to([i, j])`` + - 指向某个 element 的 typed pointer（``address_of``）；打印为 ``&B_ptr[…]`` + * - ``B.vload([i], dtype="float32x4")`` / ``B.vstore([i], v)`` + - vectorized load / store；打印为 ``*(float4*)(B_ptr + …)`` + * - ``B.view(*shape, layout=…)`` + - 在新的 shape/layout 下重新解释同一块 storage（无 copy） + * - ``B.local(*shape, layout=…)`` + - ``local`` buffer 中属于调用 thread 的私有 register slice + * - ``B.permute(*dims)`` + - 轴被置换后的 view（transposed layout） + * - ``B.access_ptr(mask, …)`` + - masked access pointer（``tvm_access_ptr`` builtin），用于把一个 region 传给 intrinsic + +**Pointer：``ptr_to`` / ``data``。** ``ptr_to`` 用于把 element address 传给 intrinsic 或 inline function；``data`` 是 base pointer： + +.. code-block:: python + + B[tx] = T.cuda.func_call("ld", A.ptr_to([tx]), source_code=SRC, return_type="float32") + +.. code-block:: c++ + + B_ptr[tx] = ld(&A_ptr[tx]); // ptr_to([tx]) -> &A_ptr[tx]; A.data -> A_ptr + +**Vectorized access：``vload`` / ``vstore``。** 用一次 wide transfer 移动多个 element（另见 :doc:`data_types`）： + +.. code-block:: python + + B.vstore([tx * 4], A.vload([tx * 4], dtype="float32x4")) + +.. code-block:: c++ + + *(float4*)(B_ptr + tx * 4) = *(float4*)(A_ptr + tx * 4); + +**Reshape / reinterpret：``view`` / ``permute``。** 二者都是纯 metadata；data pointer 不变，只是 index arithmetic 不同。``A.view(64, 4)`` 会把 256-element buffer 看作 ``64×4``；``A.permute(1, 0)`` 会转置轴： + +.. code-block:: python + + A2 = A.view(64, 4); y = A2[tx, 0] + A2[tx, 3] # A2[tx, j] -> A_ptr[tx*4 + j] + At = A.permute(1, 0); z = At[i, j] # At[i, j] -> A_ptr[j*4 + i] + +.. code-block:: c++ + + A2_ptr[tx * 4] /* +3 */ // view: row-major 64x4 index + At_ptr[(j * 4) + i] // permute: swapped strides + +**Register：``local``。** 将 thread-axis ``local`` layout 分解成调用 thread 自己的扁平 register bundle（tile primitive 中大量使用）： + +.. code-block:: python + + R = T.alloc_buffer((32, 8), "float32", scope="local", layout=TileLayout(S[(32, 8) : (1 @ laneid, 1)])) + Rl = R.local(8) # this lane's 8 registers + +.. code-block:: c++ + + alignas(64) float Rl_ptr[8]; // the lane's private registers diff --git a/zh/tirx_guide/language_reference/cuda/control_flow.rst b/zh/tirx_guide/language_reference/cuda/control_flow.rst new file mode 100644 index 00000000..848c832f --- /dev/null +++ b/zh/tirx_guide/language_reference/cuda/control_flow.rst @@ -0,0 +1,91 @@ +Control flow +============ + +Control flow 包括 ``if``、loop family 和 ``while``；它们都会映射到直观的 CUDA。 + +if +-- + +Python ``if`` / ``else`` 会变成 CUDA ``if`` / ``else``。可以用 thread/lane 比较来 guard work，或者用 ``T.ptx.elect_sync()`` 选出一个 issuing thread： + +.. code-block:: python + + if tx < 128: + A[tx] = A[tx] * T.float32(2.0) + else: + A[tx] = A[tx] + T.float32(1.0) + + if T.ptx.elect_sync(): + ... # one elected lane (e.g. to issue TMA/MMA) + +.. code-block:: c++ + + if (((int)threadIdx.x) < 128) { + A_ptr[tx] = A_ptr[tx] * 2.0f; + } else { + A_ptr[tx] = A_ptr[tx] + 1.0f; + } + +如果需要 expression-level choice（没有 branch），使用 ``T.if_then_else(cond, a, b)``。它会 lower 成 ternary，因此不引入 control-flow divergence： + +.. code-block:: c++ + + O_ptr[tx] = (A_ptr[tx] > 0.0f) ? A_ptr[tx] : 0.0f; + +Uniform vs. divergent control flow +---------------------------------- + +``if tx < 128`` 这类 per-thread guard 对普通 work 没问题，但 **collective** operation 必须被它同步的所有线程 *uniformly* 到达。 + +例如，``T.cuda.cta_sync()`` 映射到 ``__syncthreads()`` ，需要 thread block 中所有线程到达。它绝不能放在 thread-divergent 或 warpgroup-divergent branch 内：如果放在 ``if wg_id == 0:`` 内，其他 warpgroup 永远不会到达，kernel 会 deadlock。当只有一个 warpgroup 需要同步时，使用 warpgroup-scoped ``T.cuda.warpgroup_sync(id)`` （见 :ref:`chap_gemm_advanced` 和 :doc:`threads_sync`）。 + +同样的注意事项适用于 barrier setup。``mbarrier`` 的 ``.init()`` 会 lower 成 single-thread guard（``if (threadIdx.x < 1)``）。如果把它嵌到另一个 divergent branch 中，barrier 可能保持未初始化，导致 unspecified launch failure。 + +loop +---- + +Loop 有四种形式；普通 Python ``range`` 会变成 ``T.serial``： + +- ``T.serial(n)``：顺序 loop（ptxas 仍可能 unroll 它）。 +- ``T.unroll(n)``：完全 unrolled（展开成 straight-line statement）。 +- ``T.vectorized(n)``：vectorized loop。 +- ``T.grid(*extents)``：嵌套 loop nest。 + +``break`` / ``continue`` 可以在 loop 内使用。 + +.. code-block:: python + + for i, j in T.grid(8, 8): + B[i, j] = T.max(A[i, j], T.float32(0.0)) + +.. code-block:: c++ + + for (int i = 0; i < 8; ++i) + for (int j = 0; j < 8; ++j) + B_ptr[i * 8 + j] = max(A_ptr[i * 8 + j], 0.0f); + +``T.unroll(4)`` 则展开成四条 straight-line statement，没有 loop。 + +while +----- + +``while`` loop 会运行到 condition 为 false。请使用 mutable scalar counter（见 :doc:`buffers`）： + +.. code-block:: python + + i: T.int32 = 0 + while i < 64: + A[i] = A[i] + T.float32(1.0) + i += 1 + +它会 lower 成带 early-exit ``break`` 的 ``while (1)`` （counter 是一个 one-element register buffer）： + +.. code-block:: c++ + + int i_ptr[1]; + i_ptr[0] = 0; + while (1) { + if (!(i_ptr[0] < 64)) { break; } + A_ptr[i_ptr[0]] = A_ptr[i_ptr[0]] + 1.0f; + i_ptr[0] = i_ptr[0] + 1; + } diff --git a/zh/tirx_guide/language_reference/cuda/data_types.rst b/zh/tirx_guide/language_reference/cuda/data_types.rst new file mode 100644 index 00000000..10f59951 --- /dev/null +++ b/zh/tirx_guide/language_reference/cuda/data_types.rst @@ -0,0 +1,84 @@ +Data types and expressions +========================== + +每个 TIRx expression 都带有一个低层 **dtype** 和一个高层 **type**。 + +Expression dtype +---------------- + +``PrimExpr`` 的 ``.dtype`` 是其 scalar（或 vector）element type，例如 ``float32``、``float16``、``bfloat16``、``int32``、``uint8``、``bool``、低精度 ``float8_e4m3fn`` / ``float4_e2m1fn``、``handle``（pointer），以及 ``float32x4`` 这样的 vector form。每种 dtype 都会打印成匹配的 CUDA type。下面示例跨多个 dtype 分配 local 和 shared buffer，并执行 vectorized ``float32x4`` load/store： + +.. code-block:: python + + @T.prim_func + def dtypes(A_ptr: T.handle, O_ptr: T.handle): + A = T.match_buffer(A_ptr, (256,), "float32") + O = T.match_buffer(O_ptr, (256,), "float32") + T.device_entry(); bx = T.cta_id([1]); tx = T.thread_id([64]) + f16 = T.alloc_local((1,), "float16") # register scalars ... + bf16 = T.alloc_local((1,), "bfloat16") + i32 = T.alloc_local((1,), "int32") + u8 = T.alloc_local((1,), "uint8") + b1 = T.alloc_local((1,), "bool") + sm = T.alloc_shared((64,), "float16") # ... and a shared tile + v = T.alloc_local((1,), "float32x4") # a vector-dtype register (float4) + v[0] = A.vload([tx * 4], dtype="float32x4") # vectorized load + O.vstore([tx * 4], v[0]) # vectorized store + # ... (use f16/bf16/i32/u8/b1/sm) ... + +它会 lower 成（生成 CUDA，省略部分）： + +.. code-block:: c++ + + half f16_ptr[1]; // float16 + nv_bfloat16 bf16_ptr[1]; // bfloat16 + int i32_ptr[1]; // int32 + uchar u8_ptr[1]; // uint8 + signed char b1_ptr[1]; // bool + __shared__ alignas(64) half sm_ptr[64]; // shared float16 + float4 v_ptr[1]; // float32x4 (vector) + v_ptr[0] = *(float4*)(A_ptr + tx * 4); // vectorized load + *(float4*)(O_ptr + tx * 4) = v_ptr[0]; // vectorized store + +Buffer 的 dtype 本身也可以是 **vector type**：``T.alloc_local((1,), "float32x4")`` 会直接声明一个 ``float4`` register（用 ``v[0]`` 索引），随后 ``float32x4`` 的 ``vload`` / ``vstore`` 会以一次 16-byte access 移动它。Vector dtype 不绑定到 ``vload``；任何 buffer 或 scalar 都可以携带它。 + +dtype -> CUDA 映射如下： + +.. list-table:: + :header-rows: 1 + :widths: 34 33 33 + + * - dtype → CUDA + - dtype → CUDA + - dtype → CUDA + * - ``float32`` → ``float`` + - ``float16`` → ``half`` + - ``bfloat16`` → ``nv_bfloat16`` + * - ``int32`` → ``int`` + - ``uint8`` → ``uchar`` + - ``bool`` → ``signed char`` + * - ``float32x4`` → ``float4`` + - ``handle`` → ``T*`` (pointer) + - (vector dtypes → CUDA vector types) + +dtype vs type +------------- + +``dtype`` 是 *low-level* 的，说明“哪些 bit”。另外，值还有一个高层 **type**：scalar 使用 ``PrimType(dtype)``，pointer 使用 ``PointerType(PrimType(dtype), scope)``。大多数 expression 是 scalar（``PrimType``）；type system 主要在 **pointer** 上重要。 + +Pointer（``handle``） +-------------------------- + +Buffer 的 ``data``，也就是 pointer，是 pointer type 的 ``Var``，并且是 immutable（pointer 不会被重新赋值）。这决定了你如何取得它： + +- ``T.alloc_buffer(...)`` 分配 storage，并定义它的 ``data`` pointer。 +- ``T.decl_buffer(..., data=ptr)`` 在已有 pointer ``Var`` ``ptr`` 上声明 buffer。 +- 如果要让 buffer 以 pointer **expression** 为 backing，例如 ``T.ptx.map_shared_rank``（PTX ``mapa``）给出的另一个 cluster CTA 的 shared address，必须先把这个 expression 绑定成 pointer ``Var``（``data`` 必须是 ``Var``，不能是 expression），使用 ``PointerType`` 的 ``T.let``： + + .. code-block:: python + + from tvm.ir.type import PointerType, PrimType + + ptr: T.let[T.Var(name="ptr", dtype=PointerType(PrimType("uint64")))] = \ + T.reinterpret("handle", T.ptx.map_shared_rank(mbar.ptr_to([0]), 0)) + remote_mbar = T.decl_buffer([1], "uint64", data=ptr, scope="shared") diff --git a/zh/tirx_guide/language_reference/cuda/parser_utils.rst b/zh/tirx_guide/language_reference/cuda/parser_utils.rst new file mode 100644 index 00000000..27aac421 --- /dev/null +++ b/zh/tirx_guide/language_reference/cuda/parser_utils.rst @@ -0,0 +1,52 @@ +Parser utilities +================ + +少数 helper 会在 **parse time** 起作用，也就是 TVMScript 被转换成 TIRx 的阶段。它们允许你 inline Python 计算出来的值，抽出可复用片段，并打包 parser-side state。 + +``T.meta_var`` — inline Python 值 +--------------------------------- + +``T.meta_var(x)`` 告诉 parser 把 ``x``，一个在 **Python** 中计算出来的值，当作 compile-time *meta* value，并直接 inline 到 IR 中，而不是把它解析成 script variable。它可以避免一次性的 local，也驱动 metaprogramming：对 meta value 的普通 Python ``for`` 会在 parser 中展开。 + +.. code-block:: python + + n = T.meta_var(4) # n is a Python int, inlined + for j in range(n): # unrolled at parse time + acc[0] = acc[0] + A[tx, j] + +``@T.inline`` — inline function +------------------------------- + +``@T.inline`` 定义一个函数，其 body 会在 parsing 阶段 **inline 到每个 call site**，生成代码中不会出现调用。它遵循 Python 的 lexical（LEGB）scope，并使用 late binding，因此参数会 shadow 外层变量： + +.. code-block:: python + + @T.inline + def add_into(acc, x): + acc[0] = acc[0] + x + + add_into(acc, A[tx, j]) # inlined -> acc[0] = acc[0] + A[tx, j] + +``@T.meta_class`` — parser-side state object +-------------------------------------------- + +``@T.meta_class`` 标记一个普通 Python class，其 **instance 是 parser meta value**：字段可以持有 buffer 和 scalar，因此你可以把相关 allocation 和 state 打包成一个对象，并在 kernel body 中使用它。 + +.. code-block:: python + + @T.meta_class + class State: + def __init__(self, smem): + self.acc = T.alloc_local([1], "float32") + self.buf = T.decl_buffer([64], "float16", smem, scope="shared.dyn") + + s = State(smem.data) + s.acc[0] = T.float32(0.0) # use its fields like ordinary buffers + # ... s.buf[i] ... + +这很适合把 kernel 的 pipeline state（barrier、accumulator、scratch view）分组，而不是把许多独立 local 在线程中传来传去。 + +``T.constexpr`` +--------------- + +``T.constexpr`` 标记 compile-time kernel parameter，它会由 ``@T.jit`` 的 ``.specialize(...)`` 固化。细节见 :ref:`chap_tirx_primer`。 diff --git a/zh/tirx_guide/language_reference/cuda/threads_sync.rst b/zh/tirx_guide/language_reference/cuda/threads_sync.rst new file mode 100644 index 00000000..24137b23 --- /dev/null +++ b/zh/tirx_guide/language_reference/cuda/threads_sync.rst @@ -0,0 +1,97 @@ +CUDA C++/PTX intrinsics +======================= + +当没有 tile primitive 覆盖你需要的操作时，有两种 escape hatch 可以直接触达硬件：**调用 backend intrinsic**（来自 ``tvm.backend.cuda`` 的 ``T.cuda.*`` / ``T.ptx.*`` namespace），或者 **inline raw CUDA** 源码。 + +Calling backend intrinsics +-------------------------- + +``T.cuda.*`` 和 ``T.ptx.*`` 直接暴露 CUDA backend 的 device intrinsic，包括 synchronization、mbarrier、reduction，以及 PTX data-movement / MMA 指令族： + +.. code-block:: python + + T.cuda.cta_sync() # block barrier (__syncthreads) + T.cuda.warp_sync() # __syncwarp + T.cuda.warpgroup_sync(8) # warpgroup barrier + T.cuda.cta_sum(val, num_warps, scratch.ptr_to([0])) # block-level reduction + + bar = T.alloc_shared((1,), "uint64") + T.ptx.mbarrier.init(bar.data, 1) # mbarrier for async completion + T.ptx.mbarrier.try_wait(bar.data, phase) + +一个完整可运行示例：通过 ``T.tvm_warp_shuffle_xor`` 做 warp all-reduce： + +.. code-block:: python + + @T.prim_func + def warp_reduce(A_ptr: T.handle): + A = T.match_buffer(A_ptr, (32,), "float32", align=16) + T.device_entry() + cta_id = T.cta_id([1]); warp_id = T.warp_id([1]); lane_id = T.lane_id([32]) + v = T.alloc_local((1,), "float32"); i = T.alloc_local((1,), "int32") + v[0] = T.float32(31 - lane_id) + i[0] = 16 + while i[0] >= 1: + v[0] += T.tvm_warp_shuffle_xor(0xFFFFFFFF, v[0], i[0], 32, 32) + i[0] = i[0] // 2 + A[lane_id] = v[0] + +Shuffle 会直接 lower 成 ``__shfl_xor_sync``： + +.. code-block:: c++ + + v_ptr[0] = v_ptr[0] + __shfl_xor_sync(0xFFFFFFFF, v_ptr[0], i_ptr[0], 32); + +``T.ptx.*`` / ``T.cuda.*`` 下还有其他指令族：``cp_async``（LDGSTS）、``cp_async.bulk.tensor``（TMA）、``ldmatrix`` / ``stmatrix``、``tcgen05.*``（Blackwell MMA）、``atomic_add``、``fence`` 等。完整 ``tvm.backend.cuda`` reference 请参阅 backend API reference。 + +Synchronization semantics +------------------------- + +GEMM 和 Flash Attention kernel 中反复出现四种同步机制。由于它们控制异步引擎和并行线程组，误用任何一种通常都会导致 silent corruption 或 deadlock。 + +**Mbarrier Phase。** Mbarrier 用一个内部 phase bit 跟踪 arrival。``T.ptx.mbarrier.try_wait(bar, phase)`` 会阻塞，直到 barrier 的内部 phase *不同于* 调用方提供的 ``phase`` 参数。因此，跨 loop iteration 复用 barrier 时，调用方必须在每次 wait 后翻转自己的本地 phase tracker（``phase ^= 1``）。如果不这样做，后续 wait 会立即返回，允许 engine 读取半写入的 memory。:ref:`chap_gemm_basics` 给出了完整 phase-tracking 表。 + +**Election。** ``T.ptx.elect_sync()`` 会在 *一个 warp 内的 active lane* 中选出一个，不一定是 lane 0，也不是每个 CTA 一个线程。要把 issuer 缩窄到恰好一个线程，必须配合 warp-level guard。:ref:`chap_gemm_basics` 中发起 ``Tx.gemm_async`` 和 ``tcgen05.commit`` 时使用的 pattern 是先 ``if warp_id == 0:``，再 ``if T.ptx.elect_sync():``。 + +**Named Warpgroup Barrier。** ``T.cuda.cta_sync()`` 映射到 ``__syncthreads()``，需要 *每个* CTA thread arrive。一旦 warpgroup specialization 到不同 code path，把 ``cta_sync()`` 放在 warpgroup branch 内就会 deadlock，因为其他 warpgroup 永远到不了。硬件提供 16 个 named barrier（ID 0 到 15）；``T.cuda.warpgroup_sync(10)`` 只同步一个 warpgroup 的线程。不同 warpgroup 使用不同 ID（例如 ``warpgroup_sync(wg_id + 10)``），避免撞到同一个硬件 barrier。见 :ref:`chap_gemm_advanced`。 + +**Fence。** Fence 会把 producer 的写入排在 consumer（通常是异步 engine）读取之前： + +.. list-table:: + :header-rows: 1 + :widths: 50 50 + + * - Fence + - 排序内容 + * - ``T.ptx.fence.proxy_async("shared::cta")`` + - 线程写入的 shared memory 在 async proxy（TMA store / MMA）读取之前可见 + * - ``T.ptx.fence.mbarrier_init()`` + - mbarrier initialization 在后续 arrival 或 wait 使用 barrier 之前完成 + * - ``T.ptx.tcgen05.fence.after_thread_sync()`` + - ``tcgen05`` writeback edge 上的保守 ordering fence（Steps 8 和 9 添加；TMA-to-MMA 路径不需要） + +Inlining raw CUDA +----------------- + +对于完全没有 intrinsic 的操作，可以用 ``T.cuda.func_call(name, *args, source_code=..., return_type=...)`` 从源码字符串注入一个 ``__device__`` function： + +.. code-block:: python + + SRC = r""" + __device__ __forceinline__ float my_relu(float x) { return x > 0.f ? x : 0.f; } + """ + + @T.prim_func + def k(A_ptr: T.handle, B_ptr: T.handle): + A = T.match_buffer(A_ptr, (256,), "float32") + B = T.match_buffer(B_ptr, (256,), "float32") + T.device_entry(); bx = T.cta_id([1]); tx = T.thread_id([256]) + B[tx] = T.cuda.func_call("my_relu", A[tx], source_code=SRC, return_type="float32") + +源代码会原样 emit，调用也会连上： + +.. code-block:: c++ + + __device__ __forceinline__ float my_relu(float x) { return x > 0.f ? x : 0.f; } + // ... + B_ptr[tx] = my_relu(A_ptr[tx]); diff --git a/zh/tirx_guide/language_reference/index.md b/zh/tirx_guide/language_reference/index.md index 57068790..87255c21 100644 --- a/zh/tirx_guide/language_reference/index.md +++ b/zh/tirx_guide/language_reference/index.md @@ -1,10 +1,14 @@ ---- -orphan: true ---- - (chap_language_reference)= # TIRx 语言参考 -> 翻译状态：待翻译。对应英文目录：`tirx_guide/language_reference/`。 +这是编写 TIRx device kernel 时可用的完整语言特性集合，从 {ref}`chap_tirx_primer` walkthrough 中拆出：parser utility、data type 和 expression、buffer 与 memory、control flow，以及 thread synchronization。需要精确语法或语义时，请查阅这些页面。 + +```{toctree} +:maxdepth: 1 -本页用于放置 TIRx 语言参考的中文内容。 +cuda/parser_utils +cuda/data_types +cuda/buffers +cuda/control_flow +cuda/threads_sync +```