From b33a69eaf9c55d216cf0098319a019a00c831788 Mon Sep 17 00:00:00 2001 From: William Zhang Date: Fri, 14 Aug 2026 22:03:24 -0400 Subject: [PATCH] [Docs] Add Hopper matmul tutorial; fix benchmark protocol Adds the V0-V6 Hopper matmul tutorial (docs/source/tutorials/matmul-hopper/) and corrects the benchmark methodology the numbers rest on. benchmark.py, aligned with examples/blackwell_matmul/benchmark.py: - cuBLAS is timed last, under the same 3 s cooldown and iteration count as every version. Previously it was timed first, on the coldest GPU of the run, which flattered whichever entry followed it. That bias is why V0 was recorded at 3.35 ms: its best schedule measures 3.51 ms in isolation. - Both operands are scaled by k**0.25 so C has unit variance, replacing the asymmetric one-operand scaling. atol and rtol are now both meaningful. - assert_close argument order corrected to (actual, expected). - repeat stays at 30 rather than Blackwell's 100, deliberately. An 8192^3 fp16 GEMM is ~1.5 ms on an H100, so 100 iterations is ~150 ms of power-capped work and the board throttles inside the timed window: three fresh processes give v6/cuBLAS = 108.1/106.5/107.8 % at repeat=30 but 99.0/98.9/110.4 % at repeat=100. Added --repeat to compare. - New "accumulate" column, because v6 accumulates in fp16 while every other version and cuBLAS accumulate in fp32 -- which is what its looser atol was silently encoding. Numbers re-measured on an idle H100 SXM, median of three fresh processes, and cross-checked on a second GPU. All prose, tables, and plot data updated: v0 3.60 ms / 305 TF v4 1.71 ms / 642 TF v1 2.04 ms / 540 TF v5 1.62 ms / 678 TF v2 2.12 ms / 518 TF v6 1.37 ms / 800 TF (fp16 accumulate) v3 1.91 ms / 575 TF cuBLAS 1.48 ms / 742 TF v6 leads cuBLAS by 6-8% on wall clock and 2% under NCU replay, now stated with the fp16-vs-fp32 accumulation caveat in the summary table, the V6 warning box, and the closing caution. Tutorial corrections found while verifying against the source and the hardware: - v0: autotuner picks block_k=64, not 32 (confirmed fastest of the five schedules that compile). - v3: the TMA calls are inside single_thread, not issued at warp scope -- that is V4's form. - v4: tile rasterization is used by V4 itself, not "turned on by V5". - v5: a 2-stage ring buffer does not deadlock. It runs correctly at 2.24 ms. The 4th stage is a tie with 3 (1.62 vs 1.60 ms), not a requirement. - v1/v2/v3: plot captions still claimed NCU-derived TFLOPS and a cuBLAS-estimated peak; both plots use event timing and the published peak. - v6: 4x128x128 fp32 accumulator registers is the whole 64K SM register file, so fp32 at 256x256 does not fit at all. .gitignore: ignore .venv. Docs build clean with -W. Cold-cache tests/examples/test_examples.py -k hopper_matmul: 7 passed. Co-Authored-By: Claude Opus 5 (1M context) --- .gitignore | 1 + docs/source/index.rst | 1 + .../tutorials/matmul-hopper/__init__.rst | 113 +++++ .../matmul-hopper/figures/v0_block_tiling.svg | 68 +++ .../matmul-hopper/figures/v0_data_flow.svg | 55 +++ .../matmul-hopper/figures/v1_mma_vs_wgmma.svg | 66 +++ .../matmul-hopper/figures/v2_pipeline.svg | 69 +++ .../figures/v3_warp_specialization.svg | 79 ++++ .../figures/v4_pipeline_class.svg | 131 ++++++ .../figures/v4_tile_rasterization.svg | 190 ++++++++ .../matmul-hopper/figures/v4_tile_split.svg | 48 ++ .../figures/v5_wgmma_overlap.svg | 90 ++++ .../figures/v6_tile_partition.svg | 68 +++ .../tutorials/matmul-hopper/plots/plot_all.py | 7 + .../matmul-hopper/plots/plot_perf.py | 152 +++++++ .../tutorials/matmul-hopper/plots/plot_v0.py | 9 + .../tutorials/matmul-hopper/plots/plot_v1.py | 9 + .../tutorials/matmul-hopper/plots/plot_v2.py | 9 + .../tutorials/matmul-hopper/plots/plot_v3.py | 9 + .../tutorials/matmul-hopper/plots/plot_v4.py | 9 + .../tutorials/matmul-hopper/plots/plot_v5.py | 9 + .../tutorials/matmul-hopper/plots/plot_v6.py | 9 + docs/source/tutorials/matmul-hopper/v0.rst | 417 ++++++++++++++++++ docs/source/tutorials/matmul-hopper/v1.rst | 220 +++++++++ docs/source/tutorials/matmul-hopper/v2.rst | 282 ++++++++++++ docs/source/tutorials/matmul-hopper/v3.rst | 260 +++++++++++ docs/source/tutorials/matmul-hopper/v4.rst | 344 +++++++++++++++ docs/source/tutorials/matmul-hopper/v5.rst | 245 ++++++++++ docs/source/tutorials/matmul-hopper/v6.rst | 340 ++++++++++++++ examples/hopper_matmul/benchmark.py | 105 +++-- examples/hopper_matmul/matmul_v6.py | 5 +- 31 files changed, 3390 insertions(+), 29 deletions(-) create mode 100644 docs/source/tutorials/matmul-hopper/__init__.rst create mode 100644 docs/source/tutorials/matmul-hopper/figures/v0_block_tiling.svg create mode 100644 docs/source/tutorials/matmul-hopper/figures/v0_data_flow.svg create mode 100644 docs/source/tutorials/matmul-hopper/figures/v1_mma_vs_wgmma.svg create mode 100644 docs/source/tutorials/matmul-hopper/figures/v2_pipeline.svg create mode 100644 docs/source/tutorials/matmul-hopper/figures/v3_warp_specialization.svg create mode 100644 docs/source/tutorials/matmul-hopper/figures/v4_pipeline_class.svg create mode 100644 docs/source/tutorials/matmul-hopper/figures/v4_tile_rasterization.svg create mode 100644 docs/source/tutorials/matmul-hopper/figures/v4_tile_split.svg create mode 100644 docs/source/tutorials/matmul-hopper/figures/v5_wgmma_overlap.svg create mode 100644 docs/source/tutorials/matmul-hopper/figures/v6_tile_partition.svg create mode 100644 docs/source/tutorials/matmul-hopper/plots/plot_all.py create mode 100644 docs/source/tutorials/matmul-hopper/plots/plot_perf.py create mode 100644 docs/source/tutorials/matmul-hopper/plots/plot_v0.py create mode 100644 docs/source/tutorials/matmul-hopper/plots/plot_v1.py create mode 100644 docs/source/tutorials/matmul-hopper/plots/plot_v2.py create mode 100644 docs/source/tutorials/matmul-hopper/plots/plot_v3.py create mode 100644 docs/source/tutorials/matmul-hopper/plots/plot_v4.py create mode 100644 docs/source/tutorials/matmul-hopper/plots/plot_v5.py create mode 100644 docs/source/tutorials/matmul-hopper/plots/plot_v6.py create mode 100644 docs/source/tutorials/matmul-hopper/v0.rst create mode 100644 docs/source/tutorials/matmul-hopper/v1.rst create mode 100644 docs/source/tutorials/matmul-hopper/v2.rst create mode 100644 docs/source/tutorials/matmul-hopper/v3.rst create mode 100644 docs/source/tutorials/matmul-hopper/v4.rst create mode 100644 docs/source/tutorials/matmul-hopper/v5.rst create mode 100644 docs/source/tutorials/matmul-hopper/v6.rst diff --git a/.gitignore b/.gitignore index 30ef3539..35e2c14e 100644 --- a/.gitignore +++ b/.gitignore @@ -1,6 +1,7 @@ # Byte-compiled / optimized / DLL files __pycache__/ *.py[cod] +.venv *$py.class # Virtual environment diff --git a/docs/source/index.rst b/docs/source/index.rst index 3202b65b..57e4e1e6 100644 --- a/docs/source/index.rst +++ b/docs/source/index.rst @@ -21,6 +21,7 @@ Additional features include automatic tuning, caching, and a Pythonic interface :caption: Tutorials tutorials/matmul-ampere/__init__ + tutorials/matmul-hopper/__init__ tutorials/matmul-blackwell/__init__ .. toctree:: diff --git a/docs/source/tutorials/matmul-hopper/__init__.rst b/docs/source/tutorials/matmul-hopper/__init__.rst new file mode 100644 index 00000000..64febc60 --- /dev/null +++ b/docs/source/tutorials/matmul-hopper/__init__.rst @@ -0,0 +1,113 @@ +Matmul (Hopper) +=============== + +This tutorial shows how to implement a high-performance matrix multiplication kernel +(C = A x B\ :sup:`T`) targeting **NVIDIA Hopper GPUs** using **Tilus**. + +Starting from a minimal working kernel, each version introduces one new Hopper feature +or optimization technique. By the final version, the kernel exceeds vendor-library +performance for this shape. The figure below shows the progression: V0 starts at +~305 TFLOPS with a minimal kernel that pushes every operand through the register file, +and each optimization closes the gap to cuBLAS, with V6 passing it at ~800 TFLOPS. +All kernels and the benchmark script to reproduce the result can be found at +:github:`examples/hopper_matmul/`. + +.. note:: + + V6 is the one version that does not compute the same thing as the rest: it + accumulates in fp16 inside the WGMMA, while V0--V5 and cuBLAS accumulate in + fp32. That is what buys it the register budget for a ``256 x 256`` tile, and + it costs about an order of magnitude in accumulated error. :doc:`V6 ` + quantifies the trade; V5 is the fastest version that is numerically + like-for-like with cuBLAS. + +.. plot:: tutorials/matmul-hopper/plots/plot_all.py + + Hopper matmul performance on H100 SXM (M=N=K=8192, fp16). Latency is + CUDA-event timed, median of three fresh processes. Peak is the published + dense FP16 tensor core throughput of the H100 SXM. + +The progression is not perfectly monotonic: V2 introduces multi-stage pipelining but +measures slightly slower than V1, because a ring buffer alone does not create overlap +when all threads still meet at a block-wide barrier. :doc:`V3 ` supplies the +missing half and both changes pay off together. The :doc:`V2 ` page works through +this in detail --- it is the most instructive step in the series. + +.. list-table:: Summary (H100 SXM, M=N=K=8192, fp16) + :header-rows: 1 + :widths: 8 30 10 14 12 12 + + * - Version + - Optimization + - Accum. + - Latency + - TFLOPS + - Tensor pipe + * - :doc:`V0 ` + - TMA loads, register-staged ``mma.sync`` + - fp32 + - 3.60 ms + - 305 + - 59% + * - :doc:`V1 ` + - WGMMA from shared memory + - fp32 + - 2.04 ms + - 540 + - 67% + * - :doc:`V2 ` + - Multi-stage software pipelining + - fp32 + - 2.12 ms + - 518 + - 71% + * - :doc:`V3 ` + - Warp specialization + - fp32 + - 1.91 ms + - 575 + - 75% + * - :doc:`V4 ` + - Two consumers, ``Pipeline`` class, tile rasterization + - fp32 + - 1.71 ms + - 642 + - 80% + * - :doc:`V5 ` + - Overlapped WGMMA groups + - fp32 + - 1.62 ms + - 678 + - 88% + * - :doc:`V6 ` + - Four consumers, fp16 accumulation, TMA epilogue + - **fp16** + - **1.37 ms** + - **800** + - 93% + * - cuBLAS + - ``nvjet_sm90_hsh_320x128_64x3_1x2_h_bz_coopB_TNT`` + - fp32 + - 1.48 ms + - 742 + - 94% + +Tensor pipe utilization is from Nsight Compute +(``sm__pipe_tensor_cycles_active.avg.pct_of_peak_sustained_elapsed``); the latency and +TFLOPS columns are CUDA-event timings, median of three fresh processes. Reproduce +with:: + + python examples/hopper_matmul/benchmark.py --size 8192 8192 8192 + python examples/hopper_matmul/benchmark.py --ncu --size 8192 8192 8192 + +.. toctree:: + :maxdepth: 1 + :caption: Versions + + v0 + v1 + v2 + v3 + v4 + v5 + v6 diff --git a/docs/source/tutorials/matmul-hopper/figures/v0_block_tiling.svg b/docs/source/tutorials/matmul-hopper/figures/v0_block_tiling.svg new file mode 100644 index 00000000..005ca033 --- /dev/null +++ b/docs/source/tutorials/matmul-hopper/figures/v0_block_tiling.svg @@ -0,0 +1,68 @@ + + + + + + + + + + + + + A (M, K) + + + + + + block_m + + block_k + + M + K + + + + BT (K, N) + + + + + + block_k + + block_n + + K + N + + + + C (M, N) + + + + + + + + + + + block_m + + block_n + + M + N + + + + Each thread block computes one tile + + + + A x BT = C + diff --git a/docs/source/tutorials/matmul-hopper/figures/v0_data_flow.svg b/docs/source/tutorials/matmul-hopper/figures/v0_data_flow.svg new file mode 100644 index 00000000..37bb3e77 --- /dev/null +++ b/docs/source/tutorials/matmul-hopper/figures/v0_data_flow.svg @@ -0,0 +1,55 @@ + + + + + + Global + A, B tiles + + + + Shared + sa, sb + + + + Registers + a, b, acc + + + + Tensor Core + mma.sync + + + + Global + C + + + + + + + + + tma. + global_to_shared + + + load_ + shared + + + + dot + (acc) + + + store_global + + + + operands round-trip through the register file + V1 removes this: WGMMA reads shared memory directly + diff --git a/docs/source/tutorials/matmul-hopper/figures/v1_mma_vs_wgmma.svg b/docs/source/tutorials/matmul-hopper/figures/v1_mma_vs_wgmma.svg new file mode 100644 index 00000000..ab1129e6 --- /dev/null +++ b/docs/source/tutorials/matmul-hopper/figures/v1_mma_vs_wgmma.svg @@ -0,0 +1,66 @@ + + + + + + + + + V0: mma.sync — operands staged in registers + + + Shared Memory + sa [128, 64] + sb [256, 64] + 48 KB per K-chunk + + + Register File + a, b fragments + acc (fp32) + all 48 KB passes through + + + Tensor Core + mma.sync + m16n8k16 + per-warp, synchronous + + + ldmatrix x N + + registers + + + + + + V1: wgmma.mma_async — operands read from shared memory + + + Shared Memory + sa [128, 64] + sb [256, 64] + read in place + + + Register File + acc (fp32) only + operand registers freed + + + Tensor Core + wgmma.mma_async + up to 64 x 256 x 16 + per-warp-group, async + + + + 64-bit shared memory descriptor + + + + accumulate + + no ldmatrix, no operand registers, one instruction per warp group + diff --git a/docs/source/tutorials/matmul-hopper/figures/v2_pipeline.svg b/docs/source/tutorials/matmul-hopper/figures/v2_pipeline.svg new file mode 100644 index 00000000..ac0fae00 --- /dev/null +++ b/docs/source/tutorials/matmul-hopper/figures/v2_pipeline.svg @@ -0,0 +1,69 @@ + + + V1: single stage — load and compute are serial + + TMA + WGMMA + + + Load k=0 + + MMA k=0 + + + Load k=1 + + MMA k=1 + + + Load k=2 + + MMA k=2 + + + each engine idle + half the time + + + + V2: multi-stage ring buffer — loads run ahead of compute + + TMA + WGMMA + + + + Load k=0 + + Load k=1 + + Load k=2 + + Load k=3 + + Load k=4 + + + + MMA k=0 + + MMA k=1 + + MMA k=2 + + MMA k=3 + + MMA k=4 + + + time + + + + prefill + + both engines busy + + In V2 the two rows still belong to the same 128 threads, separated by + block-wide syncs — V3 gives each row its own warps. + diff --git a/docs/source/tutorials/matmul-hopper/figures/v3_warp_specialization.svg b/docs/source/tutorials/matmul-hopper/figures/v3_warp_specialization.svg new file mode 100644 index 00000000..b3912a42 --- /dev/null +++ b/docs/source/tutorials/matmul-hopper/figures/v3_warp_specialization.svg @@ -0,0 +1,79 @@ + + + + + + + + + + + + V2: one warp group alternates roles, gated by __syncthreads() + + warps 0-3 + + issue TMA + + sync + + WGMMA + + sync + + issue TMA + + sync + + WGMMA + + the tensor cores cannot be issued to while these warps are loading or synchronizing + + + + + V3: separate warps, coupled only by mbarriers + + + warp 4 + producer + + TMA k=0 + + TMA k=1 + + TMA k=2 + + TMA k=3 + + TMA k=4 + + TMA k=5 + + + warps 0-3 + consumer + + MMA k=0 + + MMA k=1 + + MMA k=2 + + MMA k=3 + + MMA k=4 + + + + + + consumer_barriers (stage full) + + + + producer_barriers (stage empty) + + the producer runs num_stages tiles ahead; neither warp ever waits on the other's code + no __syncthreads() anywhere in the main loop + diff --git a/docs/source/tutorials/matmul-hopper/figures/v4_pipeline_class.svg b/docs/source/tutorials/matmul-hopper/figures/v4_pipeline_class.svg new file mode 100644 index 00000000..ba5df981 --- /dev/null +++ b/docs/source/tutorials/matmul-hopper/figures/v4_pipeline_class.svg @@ -0,0 +1,131 @@ + + + + + + + + + Async Pipeline (stages = 5) + + + + Producer + + + producer_acquire() + wait for empty slot + + + producer_barrier() + get full_barrier to signal when done + + + producer_advance() + move to next slot + + + producer_stage → slot 3 + + + + Consumer + + + consumer_acquire() + wait for filled slot + + + consumer_barrier() + get empty_barrier to signal when done + + + consumer_advance() + move to next slot + + + consumer_stage → slot 1 + + + + Ring Buffer + + + slot + full + empty + state + + + + 0 + ✗ + ✓ + empty + + + + 1 + ✓ + ✗ + consuming + + + + 2 + ✓ + ✗ + full + + + + 3 + ✗ + ✓ + producing + + + + 4 + ✗ + ✓ + empty + + + + + + + + + slots cycle: 0 → 1 → 2 → 3 → 4 → 0 → ... + phase flips each time a pointer wraps around + + + + + Initial state: all slots are empty. Each slot is in one of four states: producing, full, consuming, or empty. + + full ✓ + = producer has filled this slot + + empty ✓ + = consumer has consumed this slot + + full ✗ + = not yet filled by producer + + empty ✗ + = not yet consumed by consumer + + + + Producer waits on empty barrier (slot free?), signals full barrier (data ready). + + + Consumer waits on full barrier (data ready?), signals empty barrier (slot freed). + + + Both advance their stage pointer independently, cycling through the ring buffer. + + diff --git a/docs/source/tutorials/matmul-hopper/figures/v4_tile_rasterization.svg b/docs/source/tutorials/matmul-hopper/figures/v4_tile_rasterization.svg new file mode 100644 index 00000000..994cfb5a --- /dev/null +++ b/docs/source/tutorials/matmul-hopper/figures/v4_tile_rasterization.svg @@ -0,0 +1,190 @@ + + + + Column-major (wave = 16 blocks) + + + + + + + + + + + + + + + + + + + + + + + + + 0 + + 1 + + 2 + + 3 + + 4 + + 5 + + 6 + + 7 + + + 8 + + 9 + + 10 + + 11 + + 12 + + 13 + + 14 + + 15 + + + + + + + + + + + + + + + + + + Active: 8 A rows + 2 B cols = 10 tiles + + + + Swizzled, S=4 (wave = 16 blocks) + + + + + + + + + + + + + + + + + + + + + + + + + 0 + + 1 + + 2 + + 3 + + + 4 + + 5 + + 6 + + 7 + + + 8 + + 9 + + 10 + + 11 + + + 12 + + 13 + + 14 + + 15 + + + + + + + + + + + + + + + + + + + + + + + + + Active: 4 A rows + 4 B cols = 8 tiles + + + + + Legend + + + = active tile (one thread block) + + + = inactive tile + + + = active A row (needs A tile in L2) + + + = active B col (needs B tile in L2) + + + + Same 16 active blocks, but swizzle reduces L2 working set from + 10 to + 8 unique A+B tiles + + diff --git a/docs/source/tutorials/matmul-hopper/figures/v4_tile_split.svg b/docs/source/tutorials/matmul-hopper/figures/v4_tile_split.svg new file mode 100644 index 00000000..64b96efb --- /dev/null +++ b/docs/source/tutorials/matmul-hopper/figures/v4_tile_split.svg @@ -0,0 +1,48 @@ + + + + + + + + One 128 x 256 output tile, two consumer warp groups + + + A (shared) + + sa[stage, 0] + 64 x 64 + + sa[stage, 1] + 64 x 64 + + + B (shared, shared by both) + + sb[stage] + 256 x 64 + + + C tile (registers) + + consumer WG0 — acc0 + rows 0-63, threads 0-127 + + consumer WG1 — acc1 + rows 64-127, threads 128-255 + + + + + + + wgmma + wgmma + + + + producer warp (threads 256-287): 3 TMA loads per stage + + Different A rows per group, the same B tile for both — so widening the tile + in M costs no extra B traffic, and each group's accumulator is half as large. + diff --git a/docs/source/tutorials/matmul-hopper/figures/v5_wgmma_overlap.svg b/docs/source/tutorials/matmul-hopper/figures/v5_wgmma_overlap.svg new file mode 100644 index 00000000..83e3d5c6 --- /dev/null +++ b/docs/source/tutorials/matmul-hopper/figures/v5_wgmma_overlap.svg @@ -0,0 +1,90 @@ + + + V4: wait_group(0) — the tensor core pipeline drains every K-tile + + warp group + tensor core + + + + issue 0 + + wait_group(0) + + release + + + issue 1 + + wait_group(0) + + release + + + issue 2 + + wait_group(0) + + + + MMA 0 + + MMA 1 + + MMA 2 + + + + idle + + idle + + drains between MMAs + + + + + V5: wait_group(1) — the next MMA is issued before waiting on the previous + + warp group + tensor core + + + issue 0 + + issue 1 + + wait_group(1) + + release 0 + + + issue 2 + + wait_group(1) + + release 1 + + + issue 3 + + wait_group(1) + + + + MMA 0 + + MMA 1 + + MMA 2 + + MMA 3 + + MMA 4 + + no gaps + + When wait_group(1) returns, only the previous MMA has completed — the current one + is still reading shared memory. The stage release therefore lags one iteration behind, + which is why the consumer holds two stages and the ring buffer grows to four. + diff --git a/docs/source/tutorials/matmul-hopper/figures/v6_tile_partition.svg b/docs/source/tutorials/matmul-hopper/figures/v6_tile_partition.svg new file mode 100644 index 00000000..3fba8e4a --- /dev/null +++ b/docs/source/tutorials/matmul-hopper/figures/v6_tile_partition.svg @@ -0,0 +1,68 @@ + + + + + + + + 256 x 256 tile across four consumer warp groups + + + A (shared) + + sa[stage, 0] + + sa[stage, 1] + + sa[stage, 2] + + sa[stage, 3] + 4 x (64 x 64) + + + B (shared) + + sb[stage] + 256 x 64 + read by all four + + + C tile (fp16 accumulators, in registers) + + WG0 — rows 0-63 + 64 x 256 fp16 = 64 regs/thread + + WG1 — rows 64-127 + + WG2 — rows 128-191 + + WG3 — rows 192-255 + + + + + + + + + Epilogue: four quarters take turns through one shared buffer + + + WG0..WG3 + + + sc [64, 256] + one quarter at a time + + + Global C + bulk TMA store + + + store_shared + + producer warp + + epilogue_ready / + epilogue_free + diff --git a/docs/source/tutorials/matmul-hopper/plots/plot_all.py b/docs/source/tutorials/matmul-hopper/plots/plot_all.py new file mode 100644 index 00000000..cd070021 --- /dev/null +++ b/docs/source/tutorials/matmul-hopper/plots/plot_all.py @@ -0,0 +1,7 @@ +import os +import sys + +sys.path.insert(0, os.path.join(os.getcwd(), "tutorials", "matmul-hopper", "plots")) +from plot_perf import plot_performance + +plot_performance() diff --git a/docs/source/tutorials/matmul-hopper/plots/plot_perf.py b/docs/source/tutorials/matmul-hopper/plots/plot_perf.py new file mode 100644 index 00000000..fc329049 --- /dev/null +++ b/docs/source/tutorials/matmul-hopper/plots/plot_perf.py @@ -0,0 +1,152 @@ +"""Shared performance plotting for Hopper matmul tutorials. + +Benchmark data collected on an H100 80GB HBM3 (SXM) with +``examples/hopper_matmul/benchmark.py`` at M=N=K=8192, fp16. + +Latencies are CUDA-event timings (5 warmup + 30 timed iterations), taken as the +median of three fresh processes. TFLOPS = 2*M*N*K / latency. + +Note that V6 accumulates in fp16 while V0-V5 and cuBLAS accumulate in fp32; see +the tutorial's V6 page for what that costs in accuracy. +""" + +import matplotlib.pyplot as plt +import numpy as np + +# --- Benchmark data (H100 SXM, M=N=K=8192, fp16, CUDA-event timing) --- + +VERSIONS = ["V0", "V1", "V2", "V3", "V4", "V5", "V6"] + +# Optimization label for each version +LABELS = [ + "TMA + MMA", + "WGMMA", + "Pipelining", + "Warp Spec.", + "2 Consumers", + "WGMMA Overlap", + "4 Consumers", +] + +# Median latency in ms over three fresh processes +LATENCY_MS = [3.600, 2.037, 2.124, 1.913, 1.714, 1.623, 1.374] +CUBLAS_LATENCY_MS = 1.481 + +_FLOP = 2 * 8192**3 + + +def _tflops(ms): + return _FLOP / (ms * 1e-3) * 1e-12 + + +TFLOPS = [_tflops(ms) for ms in LATENCY_MS] +CUBLAS_TFLOPS = _tflops(CUBLAS_LATENCY_MS) + +# Published dense FP16 tensor core throughput of the H100 SXM. +PEAK_TFLOPS = 989.4 + + +def plot_performance(up_to_version: int | None = None): + """Plot TFLOPS for tutorial versions. + + Parameters + ---------- + up_to_version : int or None + If given, highlight V0 through V{up_to_version} (solid line) and + show remaining versions as dashed (preview). Labels are only shown + for the highlighted versions. + If None, show all versions as solid with labels. + """ + n_total = len(VERSIONS) + if up_to_version is not None: + n_solid = up_to_version + 1 + else: + n_solid = n_total + + x = np.arange(n_total) + + fig, ax = plt.subplots(figsize=(max(6.0, 1.25 * n_total + 1.2), 4.2)) + + # Solid line: current and past versions + ax.plot(x[:n_solid], TFLOPS[:n_solid], "o-", color="#5B9BD5", linewidth=2.2, markersize=8, zorder=4) + + # Dashed line: future versions (preview) + if n_solid < n_total: + x_dash = x[n_solid - 1 :] + y_dash = TFLOPS[n_solid - 1 :] + ax.plot(x_dash, y_dash, "o--", color="#5B9BD5", linewidth=1.2, markersize=5, alpha=0.35, zorder=3) + + # cuBLAS reference line + ax.axhline(y=CUBLAS_TFLOPS, color="#E07B39", linewidth=1.5, linestyle="--", zorder=2) + + # Peak TFLOPS reference line + ax.axhline(y=PEAK_TFLOPS, color="#888888", linewidth=1.5, linestyle="--", zorder=2) + + # Inline labels for reference lines (left side, bold) + ax.text( + 0.02, + PEAK_TFLOPS + PEAK_TFLOPS * 0.012, + f"Peak ({PEAK_TFLOPS:.0f} TFLOPS)", + ha="left", + va="bottom", + fontsize=9.5, + fontweight="bold", + color="#888888", + transform=ax.get_yaxis_transform(), + ) + ax.text( + 0.02, + CUBLAS_TFLOPS - PEAK_TFLOPS * 0.012, + f"cuBLAS ({CUBLAS_TFLOPS:.0f} TFLOPS)", + ha="left", + va="top", + fontsize=9.5, + fontweight="bold", + color="#E07B39", + transform=ax.get_yaxis_transform(), + ) + + ax.set_ylabel("TFLOPS", fontsize=11) + ax.set_xticks(x) + ax.set_xticklabels(VERSIONS, fontsize=10) + ax.tick_params(axis="y", labelsize=10) + # Extra room on the right so the final version's label stays inside the axes + ax.set_xlim(-0.3, n_total - 0.5 + 0.85) + ax.set_ylim(0, PEAK_TFLOPS * 1.22) + + # Labels near solid points + for i in range(n_solid): + yi = TFLOPS[i] + + # TFLOPS value above the point, with white background to avoid + # overlap with reference lines + ax.annotate( + f"{yi:.0f}", + xy=(x[i], yi), + xytext=(0, 10), + textcoords="offset points", + ha="center", + va="bottom", + fontsize=9.5, + color="#333", + bbox=dict(boxstyle="round,pad=0.15", fc="white", ec="none", alpha=0.85), + zorder=5, + ) + + # Optimization label below and to the right of the point + ax.annotate( + LABELS[i], + xy=(x[i], yi), + xytext=(4, -14), + textcoords="offset points", + ha="left", + va="top", + fontsize=9, + color="#666", + style="italic", + ) + + ax.grid(axis="y", alpha=0.3, zorder=0) + fig.tight_layout() + + return fig diff --git a/docs/source/tutorials/matmul-hopper/plots/plot_v0.py b/docs/source/tutorials/matmul-hopper/plots/plot_v0.py new file mode 100644 index 00000000..d2ec4c37 --- /dev/null +++ b/docs/source/tutorials/matmul-hopper/plots/plot_v0.py @@ -0,0 +1,9 @@ +import os +import sys + +# plot_basedir is set to docs/source/ in conf.py, and plot_directive +# sets the cwd to plot_basedir before running the script +sys.path.insert(0, os.path.join(os.getcwd(), "tutorials", "matmul-hopper", "plots")) +from plot_perf import plot_performance + +plot_performance(up_to_version=0) diff --git a/docs/source/tutorials/matmul-hopper/plots/plot_v1.py b/docs/source/tutorials/matmul-hopper/plots/plot_v1.py new file mode 100644 index 00000000..52e8bcfc --- /dev/null +++ b/docs/source/tutorials/matmul-hopper/plots/plot_v1.py @@ -0,0 +1,9 @@ +import os +import sys + +# plot_basedir is set to docs/source/ in conf.py, and plot_directive +# sets the cwd to plot_basedir before running the script +sys.path.insert(0, os.path.join(os.getcwd(), "tutorials", "matmul-hopper", "plots")) +from plot_perf import plot_performance + +plot_performance(up_to_version=1) diff --git a/docs/source/tutorials/matmul-hopper/plots/plot_v2.py b/docs/source/tutorials/matmul-hopper/plots/plot_v2.py new file mode 100644 index 00000000..53d5e866 --- /dev/null +++ b/docs/source/tutorials/matmul-hopper/plots/plot_v2.py @@ -0,0 +1,9 @@ +import os +import sys + +# plot_basedir is set to docs/source/ in conf.py, and plot_directive +# sets the cwd to plot_basedir before running the script +sys.path.insert(0, os.path.join(os.getcwd(), "tutorials", "matmul-hopper", "plots")) +from plot_perf import plot_performance + +plot_performance(up_to_version=2) diff --git a/docs/source/tutorials/matmul-hopper/plots/plot_v3.py b/docs/source/tutorials/matmul-hopper/plots/plot_v3.py new file mode 100644 index 00000000..482e23f4 --- /dev/null +++ b/docs/source/tutorials/matmul-hopper/plots/plot_v3.py @@ -0,0 +1,9 @@ +import os +import sys + +# plot_basedir is set to docs/source/ in conf.py, and plot_directive +# sets the cwd to plot_basedir before running the script +sys.path.insert(0, os.path.join(os.getcwd(), "tutorials", "matmul-hopper", "plots")) +from plot_perf import plot_performance + +plot_performance(up_to_version=3) diff --git a/docs/source/tutorials/matmul-hopper/plots/plot_v4.py b/docs/source/tutorials/matmul-hopper/plots/plot_v4.py new file mode 100644 index 00000000..8dc50ee4 --- /dev/null +++ b/docs/source/tutorials/matmul-hopper/plots/plot_v4.py @@ -0,0 +1,9 @@ +import os +import sys + +# plot_basedir is set to docs/source/ in conf.py, and plot_directive +# sets the cwd to plot_basedir before running the script +sys.path.insert(0, os.path.join(os.getcwd(), "tutorials", "matmul-hopper", "plots")) +from plot_perf import plot_performance + +plot_performance(up_to_version=4) diff --git a/docs/source/tutorials/matmul-hopper/plots/plot_v5.py b/docs/source/tutorials/matmul-hopper/plots/plot_v5.py new file mode 100644 index 00000000..b2632308 --- /dev/null +++ b/docs/source/tutorials/matmul-hopper/plots/plot_v5.py @@ -0,0 +1,9 @@ +import os +import sys + +# plot_basedir is set to docs/source/ in conf.py, and plot_directive +# sets the cwd to plot_basedir before running the script +sys.path.insert(0, os.path.join(os.getcwd(), "tutorials", "matmul-hopper", "plots")) +from plot_perf import plot_performance + +plot_performance(up_to_version=5) diff --git a/docs/source/tutorials/matmul-hopper/plots/plot_v6.py b/docs/source/tutorials/matmul-hopper/plots/plot_v6.py new file mode 100644 index 00000000..eb2ea13b --- /dev/null +++ b/docs/source/tutorials/matmul-hopper/plots/plot_v6.py @@ -0,0 +1,9 @@ +import os +import sys + +# plot_basedir is set to docs/source/ in conf.py, and plot_directive +# sets the cwd to plot_basedir before running the script +sys.path.insert(0, os.path.join(os.getcwd(), "tutorials", "matmul-hopper", "plots")) +from plot_perf import plot_performance + +plot_performance(up_to_version=6) diff --git a/docs/source/tutorials/matmul-hopper/v0.rst b/docs/source/tutorials/matmul-hopper/v0.rst new file mode 100644 index 00000000..0a510d88 --- /dev/null +++ b/docs/source/tutorials/matmul-hopper/v0.rst @@ -0,0 +1,417 @@ +.. _tutorial_hopper_matmul_v0: + +0. A Minimal Hopper Matmul +========================== + +This first version implements a minimal but correct matrix multiplication kernel +on Hopper GPUs. It introduces two key Hopper features: +**TMA** (Tensor Memory Access, :doc:`tma `) +for bulk data movement, and **asynchronous barriers** +(:doc:`mbarrier `) for tracking when that +movement completes. + +The tensor cores are still driven the Ampere way --- operands are staged into +registers and multiplied with the classic ``mma.sync`` instruction. That is the +piece we replace in V1. The kernel is not yet fast, but it establishes the +foundation for everything that follows. + + +The Full Kernel +--------------- + +Before diving into the details, here is the complete kernel so you can see the +big picture. We will explain each part in the sections that follow. + +.. hint:: + :class: margin + + To view the generated CUDA source code, check the cache directory. + See :doc:`/programming-guides/cache` for details. + +.. literalinclude:: ../../../../examples/hopper_matmul/matmul_v0.py + :language: python + :start-at: @tilus.autotune + :end-at: self.store_global(gc, casted_acc, offsets=[offset_m, offset_n]) + :caption: MatmulTMA --- full kernel + + +Block Tiling +------------ + +We compute :math:`C = A \times B^T` where A is (M, K) and B is (N, K). +The output matrix C is (M, N). + +Each thread block is responsible for computing one ``block_m x block_n`` tile of +C. The K dimension is iterated in chunks of ``block_k``. + +.. figure:: figures/v0_block_tiling.svg + :width: 100% + :align: center + + Block tiling of the matmul. Each thread block computes one output tile. + The hatched regions show the full slices of A and B\ :sup:`T` that participate + in computing the highlighted C tile. + +.. note:: + + **Data layout: K-major.** Hopper tensor cores expect operands in shared memory + with K contiguous (or MN-contiguous). This tutorial uses K-major throughout, so + A is ``[M, K]`` and B is ``[N, K]``. The MMA expects logical shapes ``[M, K]`` + and ``[K, N]``, which is why we call ``b.transpose()`` --- a view operation + that reinterprets the layout without moving data. + + +Data Flow +--------- + +Triton also uses Hopper hardware features like TMA and WGMMA, but manages them +automatically through compiler passes. Tilus opens the black box: you control +memory placement, data movement, and synchronization directly, which is necessary +for achieving peak performance. The kernel moves data through three memory levels: + +.. figure:: figures/v0_data_flow.svg + :width: 100% + :align: center + + Data flow in the kernel: Global Memory |rarr| Shared Memory |rarr| Registers |rarr| Global Memory. + +.. |rarr| unicode:: U+2192 + +1. **Global** |rarr| **Shared**: + :meth:`tma.global_to_shared() ` + loads tiles of A and B from global memory into shared memory asynchronously, + using the dedicated TMA hardware engine. +2. **Shared** |rarr| **Register**: :meth:`~tilus.Script.load_shared` reads the + staged tiles into per-thread registers, laid out to match what the tensor core + MMA instruction expects. +3. **Register** |rarr| **Register**: :meth:`~tilus.Script.dot` multiplies the two + register tiles and accumulates into an fp32 register accumulator. This lowers + to the classic ``mma.sync`` tensor core instruction. +4. **Register** |rarr| **Global**: :meth:`~tilus.Script.store_global` writes the + final result back to global memory. + +Note that steps 2 and 3 are where Hopper leaves performance behind: every operand +element makes a round trip through the register file before reaching the tensor +core. V1 removes that round trip entirely. + + +TMA: Tensor Memory Access +------------------------- + +TMA is a hardware unit introduced on Hopper that asynchronously copies a +multi-dimensional tile between global and shared memory. Compared to the +Ampere-era ``cp.async`` path (where every thread issues its own 16-byte copy): + +- **Fewer instructions**: one TMA call replaces hundreds of per-thread copy + instructions. +- **No thread occupation**: the TMA engine operates independently; the issuing + thread can proceed to other work. +- **Built-in address generation**: TMA handles multi-dimensional indexing and + shared-memory swizzling internally, so no registers are burned on address + math. + +In Tilus, TMA loads are issued via +:meth:`tma.global_to_shared() `. +The instruction takes a global tensor ``src``, a shared tensor ``dst``, +``offsets`` into the global tensor, and an ``mbarrier`` for completion tracking. +The tile shape and swizzle pattern are derived from the shared tensor, and Tilus +builds the required tensor map descriptor for you. + +For more details, see :doc:`/python-api/instruction-groups/tma`. + + +Asynchronous Barriers (mbarrier) +-------------------------------- + +In Triton, synchronization is handled implicitly. On Hopper, many operations are +**asynchronous**: the instruction returns immediately and the work completes in +the background. This enables overlapping data movement with computation, but +requires explicit tracking of when operations finish. This is the role of the +**mbarrier** (memory barrier, see :doc:`/python-api/instruction-groups/mbarrier`). + +.. figure:: /python-api/instruction-groups/figures/mbarrier_state.svg + :width: 88% + :align: center + + An mbarrier tracks pending arrivals and a phase bit. + +An mbarrier is a **64-bit synchronization object in shared memory** that tracks: + +- **Pending arrivals**: how many threads still need to signal they are done. + Each :meth:`mbarrier.arrive() ` + call decrements this count. +- **Pending transactions** (tx-count): how many bytes of asynchronous transfer + are still outstanding. +- **Phase** (1 bit): flips between 0 and 1 each time a phase completes. + +A phase completes when both pending arrivals and tx-count reach zero. At that +point, the hardware automatically flips the phase bit and resets the counters +for the next phase. + +**Wait** checks the phase: +:meth:`mbarrier.wait(barrier, phase=p) ` +blocks until the barrier's current phase differs from ``p``. When the phase has +flipped, the tracked operations are guaranteed to have completed. + +**Why flip the phase?** The same barrier is reused across loop iterations. The +phase bit distinguishes "this iteration completed" from "the previous iteration +completed." After each wait, the caller flips its local phase variable +(``phase ^= 1``) so the next wait targets the new phase: + +.. code-block:: python + + phase: uint32 = 0 # start expecting phase 0 + for ...: + ... # issue async work on the barrier + self.mbarrier.wait(barrier, phase=phase) # wait for current phase + phase ^= 1 # next iteration waits for the other phase + + +Tracking TMA Completion with tx-count +-------------------------------------- + +TMA loads are tracked through the mbarrier's **tx-count** (transaction byte +count) rather than through arrivals. The flow is: + +1. A single thread calls + :meth:`mbarrier.arrive_and_expect_tx() ` + to declare how many bytes the upcoming TMA transfers will deliver. This both + arrives at the barrier (decrementing pending arrivals) and increases the + barrier's tx-count. +2. :meth:`tma.global_to_shared() ` + is issued. When the TMA engine completes a transfer, the hardware + automatically decrements the barrier's tx-count by the number of bytes + delivered. +3. :meth:`mbarrier.wait() ` + blocks until both pending arrivals **and** tx-count reach zero --- meaning the + declaration has been made and all TMA data has landed in shared memory. + +.. note:: + + The ``transaction_bytes`` must exactly match the total bytes that will be + transferred by the subsequent TMA calls. In our case, that is + ``sa.nbytes + sb.nbytes``, the combined size of the two shared tiles + (see :attr:`SharedTensor.nbytes `). + + +Thread Groups +------------- + +By default, every instruction in a Tilus kernel operates on the **entire thread +block**: the ``__call__`` body defines the behavior of all threads in the block +collectively. However, efficient matrix multiplication kernels on Hopper require +different warps to perform different jobs and collaborate with each other +asynchronously. To narrow the execution scope to a subset of threads, Tilus +provides :doc:`thread groups `. + +A thread group selects a subset of threads within the block using +:meth:`~tilus.Script.thread_group`. For example: + +.. code-block:: python + + with self.thread_group(thread_begin=0, num_threads=32): + # Only threads 0-31 (one warp) execute this + ... + + with self.thread_group(thread_begin=32, num_threads=32): + # Only threads 32-63 execute this + ... + +Tilus also provides shortcuts for common patterns: +:meth:`~tilus.Script.single_thread` for one thread, +:meth:`~tilus.Script.single_warp` for one warp (32 threads), and +:meth:`~tilus.Script.warp_group` for a full warp group (4 warps). + +Note that Tilus does not expose ``threadIdx`` to the user. There is no way to +write ``if threadIdx.x < 32`` in a Tilus program. Instead, use +:meth:`~tilus.Script.thread_group` and its shortcuts to narrow the execution +scope. + +Every Tilus instruction has a requirement on the thread group it can execute in. +Some instructions work in any thread group, while others require a single thread, +a single warp, or a warp group. V0 uses only the simplest case: +:meth:`~tilus.Script.single_thread`, so that ``arrive_and_expect_tx`` counts one +arrival and one byte declaration instead of 128 of each: + +.. literalinclude:: ../../../../examples/hopper_matmul/matmul_v0.py + :language: python + :start-at: with self.single_thread(): + :end-at: self.mbarrier.wait(tma_barrier, phase=phase) + :dedent: 12 + +For more details, see :doc:`/programming-guides/thread-group`. + + +Walkthrough +----------- + +With the key Hopper features covered above (TMA, asynchronous barriers, and +thread groups), let us now walk through the kernel code in detail. + +A Tilus kernel is defined as a subclass of :class:`~tilus.Script`. The +``__init__`` method stores compile-time hyperparameters (tile sizes), and +``__call__`` describes the kernel logic. For more on the script structure, see +:doc:`/programming-guides/tilus-script`. + +The ``@tilus.autotune`` decorators define a search space for compile-time +hyperparameters. Tilus benchmarks all combinations and picks the fastest +configuration automatically. For more on autotuning, see +:doc:`/programming-guides/autotuning`. + + +Kernel Setup +~~~~~~~~~~~~ + +.. literalinclude:: ../../../../examples/hopper_matmul/matmul_v0.py + :language: python + :start-at: self.attrs.blocks = [ + :end-at: phase: uint32 = 0 + :dedent: 8 + :caption: Kernel setup + +- :attr:`self.attrs.blocks ` sets the grid + dimensions: ``ceil(M / block_m) x ceil(N / block_n)`` thread blocks. +- :attr:`self.attrs.warps ` sets the number + of warps per block. Here we use 4 warps (128 threads) --- one warp group, the + granularity the Hopper tensor core will require from V1 onward. +- ``offset_m`` and ``offset_n`` are the output tile offsets, computed from the + block index (:attr:`~tilus.Script.blockIdx`). +- :meth:`~tilus.Script.global_view` interprets the raw pointers as 2D global + memory tensors with the given dtype and shape. +- :meth:`~tilus.Script.shared_tensor` allocates shared memory tiles for staging + A and B data. +- :meth:`~tilus.Script.register_tensor` allocates the fp32 accumulator, + distributed across the 128 threads of the block. A ``128 x 256`` fp32 + accumulator costs 256 registers per thread --- a real constraint on Hopper, + since the accumulator lives in the same register file as everything else. +- :meth:`mbarrier.alloc() ` + allocates one mbarrier with an expected arrival count of 1, because exactly one + thread will call ``arrive_and_expect_tx`` on it each iteration. + + +Main Loop +~~~~~~~~~ + +.. literalinclude:: ../../../../examples/hopper_matmul/matmul_v0.py + :language: python + :start-at: for offset_k in range(0, k_size, block_k): + :end-at: phase ^= 1 + :dedent: 8 + :caption: Main loop + +In each iteration: + +- :meth:`~tilus.Script.single_thread` narrows the scope so that + :meth:`mbarrier.arrive_and_expect_tx() ` + declares the expected bytes exactly once. The two + :meth:`tma.global_to_shared() ` + calls then load the A and B tiles for this K-chunk, and + :meth:`mbarrier.wait() ` + blocks until both have landed. +- :meth:`~tilus.Script.sync` after the ``single_thread`` block is what makes the + data visible to the *other* 127 threads: only thread 0 executed the wait, so + the remaining threads need a block-wide barrier before they may read shared + memory. This lowers to ``__syncthreads()``. +- :meth:`~tilus.Script.load_shared` reads the two shared tiles into register + tensors, and :meth:`~tilus.Script.dot` accumulates + ``acc += a @ b.transpose()`` on the tensor cores. +- The second :meth:`~tilus.Script.sync` guards the *other* direction: the next + iteration's TMA will overwrite ``sa`` and ``sb``, so every thread must be done + reading them before thread 0 is allowed to issue the next load. +- ``phase ^= 1`` flips the local phase so the next ``mbarrier.wait`` targets + the new phase of the reused barrier. + +Note how much of the loop is *waiting*. The TMA runs, then everyone waits; the +MMA runs, then everyone waits again. Nothing overlaps. That is the theme of the +next several versions. + + +Epilogue +~~~~~~~~ + +.. literalinclude:: ../../../../examples/hopper_matmul/matmul_v0.py + :language: python + :start-at: casted_acc = self.cast(acc, dtype=float16) + :end-at: self.store_global(gc, casted_acc, offsets=[offset_m, offset_n]) + :dedent: 8 + :caption: Epilogue + +After the loop, :meth:`~tilus.Script.cast` converts the fp32 accumulator to fp16 +and :meth:`~tilus.Script.store_global` writes the result to global memory +directly from registers. + +.. note:: + + You might expect a :meth:`~tilus.Script.free_shared` on ``sa`` and ``sb`` + here, and the kernel deliberately does not have one. Nothing is allocated + after this point, so freeing would reclaim no shared memory --- but it would + return those slots to the allocator's free list. Barriers are placed *after* + the whole function has been emitted, at which point the allocator sees only + that static free list and not the fact that the TMA engine writes ``sa`` and + ``sb`` throughout the loop above. Placing an mbarrier inside a live TMA + destination corrupts the barrier's state, and the failure is a subtle one: + the kernel still runs, still produces mostly-correct output, and drops NaNs + into a fraction of a percent of the result on some launches and not others. + + +Running the Kernel +------------------ + +``MatmulTMA()`` creates a kernel template. Compilation happens on the first call. + +Note the two different integer annotations in the function signature: + +- ``int32`` (e.g., ``m_size: int32``): a **runtime** parameter. The value is + passed to the GPU kernel as an argument and can change between calls without + recompilation. +- ``int`` (e.g., ``n_size: int``, ``k_size: int``): a **compile-time + constant**. The value is baked into the generated CUDA code, so a new value + triggers JIT recompilation and autotuning. + +Making ``n_size`` and ``k_size`` compile-time constants allows the compiler to +specialize the kernel (e.g., unroll loops, compute constant addresses). For more +details, see :doc:`/programming-guides/tilus-script`. + +Once compiled, subsequent calls with the same compile-time values dispatch +directly to the GPU. + +.. literalinclude:: ../../../../examples/hopper_matmul/matmul_v0.py + :language: python + :start-at: def main + :end-at: print(df) + :caption: Launch, verify, and benchmark + + +Performance +----------- + +This minimal kernel reaches **305 TFLOPS** (3.60 ms), about 41% of cuBLAS. The +autotuner settles on a small ``64 x 128`` tile with ``block_k=64``. Only 5 of the +15 candidates in the search space compile at all: a ``128 x 256`` fp32 +accumulator needs 256 registers per thread on its own, past the 255-register +limit, and the operand fragments ``mma.sync`` requires have to fit alongside it. +Nsight Compute reports only 59% tensor pipe utilization --- the tensor cores +spend most of their time waiting on the shared-to-register traffic feeding them. +The complete source is at :github:`examples/hopper_matmul/matmul_v0.py`. + +.. plot:: tutorials/matmul-hopper/plots/plot_v0.py + + Hopper matmul performance on H100 SXM (M=N=K=8192, fp16). Latency is + CUDA-event timed, median of three fresh processes. Peak is the published + dense FP16 tensor core throughput of the H100 SXM. + + +What's Next +----------- + +This kernel works but is far from optimal. The main bottleneck is the +**MMA path**: :meth:`~tilus.Script.dot` lowers to ``mma.sync``, which requires +every operand fragment to be loaded from shared memory into registers first. +That costs instructions, register file bandwidth, and registers --- all of which +compete with the accumulator that already dominates the register budget. + +In :doc:`the next version `, we replace it with **WGMMA** (warp-group MMA), +Hopper's asynchronous tensor core instruction. WGMMA reads its operands +**directly from shared memory** via a descriptor, so the ``load_shared`` step +disappears completely, and a single instruction issued by one warp group covers a +much larger tile. diff --git a/docs/source/tutorials/matmul-hopper/v1.rst b/docs/source/tutorials/matmul-hopper/v1.rst new file mode 100644 index 00000000..f48af63e --- /dev/null +++ b/docs/source/tutorials/matmul-hopper/v1.rst @@ -0,0 +1,220 @@ +.. _tutorial_hopper_matmul_v1: + +1. WGMMA: Hopper's Asynchronous Tensor Core +============================================ + +V0 drove the tensor cores through :meth:`~tilus.Script.dot`, which lowers to the +Ampere-era ``mma.sync`` instruction. Every operand fragment had to be copied from +shared memory into registers first, by explicit ``load_shared`` calls, before the +tensor core could see it. + +This version replaces that with **WGMMA** (Warp Group Matrix Multiply-Accumulate, +:doc:`wgmma `), Hopper's native tensor core +instruction. WGMMA is **asynchronous** and reads its ``A`` and ``B`` operands +**directly from shared memory** through a descriptor, so the register round trip +disappears entirely. A single WGMMA instruction, issued cooperatively by a warp +group (4 warps, 128 threads), covers a tile up to ``64 x 256 x 16``. + +The change is small in code --- three lines swapped for four --- but it is the +single most important instruction on Hopper, and every later version builds on +its asynchronous protocol. + + +The Full Kernel +--------------- + +.. literalinclude:: ../../../../examples/hopper_matmul/matmul_v1.py + :language: python + :start-at: @tilus.autotune + :end-at: self.store_global(gc, casted_acc, offsets=[offset_m, offset_n]) + :caption: MatmulWGMMA --- full kernel + + +What Changed from V0 +-------------------- + +The kernel structure is unchanged --- same block tiling, same TMA loads, same +single-stage loop. Only the compute step differs. + +.. list-table:: + :header-rows: 1 + :widths: 15 40 40 + + * - + - V0 + - V1 + * - **MMA instruction** + - :meth:`~tilus.Script.dot` (``mma.sync``, synchronous) + - :meth:`wgmma.mma() ` (``wgmma.mma_async``, asynchronous) + * - **Operand source** + - Registers (staged via :meth:`~tilus.Script.load_shared`) + - Shared memory, read directly by the tensor core + * - **Accumulator** + - fp32 registers + - fp32 registers (unchanged) + * - **Issuing scope** + - All threads + - One warp group (4 warps), collectively + * - **Completion** + - Implicit (instruction retires in order) + - ``commit_group`` + ``wait_group`` + * - **New instructions** + - + - :meth:`~tilus.lang.instructions.wgmma.WgmmaInstructionGroup.fence`, + :meth:`~tilus.lang.instructions.wgmma.WgmmaInstructionGroup.mma`, + :meth:`~tilus.lang.instructions.wgmma.WgmmaInstructionGroup.commit_group`, + :meth:`~tilus.lang.instructions.wgmma.WgmmaInstructionGroup.wait_group` + + +Why Operands from Shared Memory Matter +-------------------------------------- + +.. figure:: figures/v1_mma_vs_wgmma.svg + :width: 100% + :align: center + + ``mma.sync`` (V0) stages both operands through the register file. WGMMA (V1) + hands the tensor core a shared-memory descriptor instead, and only the + accumulator stays in registers. + +Consider a ``128 x 256 x 64`` block tile in fp16. With ``mma.sync``, the A and B +data for one K-chunk is ``(128 + 256) x 64 x 2 = 48 KB``, and all of it must pass +through the register file on its way to the tensor core --- every iteration, for +every block. That traffic costs three things: + +- **Instructions**: hundreds of ``ldmatrix``/``LDS`` operations per K-chunk, all + issued by the same warps that are supposed to be feeding the tensor core. +- **Register file bandwidth**: shared with the accumulator writes the tensor core + is already performing. +- **Registers**: operand fragments need somewhere to live, and on Hopper the fp32 + accumulator alone can occupy 256 registers per thread. + +WGMMA removes all three at once. The instruction takes a **descriptor** --- a +64-bit value encoding the shared memory base address, the leading/stride byte +offsets, and the swizzle mode --- and the tensor core walks shared memory itself. +In Tilus you never construct the descriptor by hand; passing a +:class:`~tilus.ir.tensor.SharedTensor` to +:meth:`wgmma.mma() ` is +enough, and the compiler derives the encoding from the tensor's layout. + +.. note:: + + WGMMA can also take its ``A`` operand from registers (``B`` must always come + from shared memory). That variant is useful when A is produced on the fly, but + for matmul the shared-memory form is what you want. + + +The WGMMA Protocol +------------------ + +WGMMA is asynchronous: :meth:`wgmma.mma() ` +returns immediately and the tensor core keeps working in the background. It also +reads shared memory and writes registers *outside* the normal instruction +ordering, so the hardware needs to be told where the boundaries are. Hopper +defines a strict four-step protocol: + +.. code-block:: python + + self.wgmma.fence() # 1. prior writes to operands/accumulator are visible + self.wgmma.mma(sa, sb.transpose(), acc) # 2. issue (may be called many times) + self.wgmma.commit_group() # 3. bundle all issued MMAs into one commit group + self.wgmma.wait_group(0) # 4. wait until at most 0 groups remain pending + +1. :meth:`wgmma.fence() ` + establishes ordering between generic memory accesses and the asynchronous + tensor core. It guarantees that the shared memory written by TMA, and the + accumulator registers written by any previous non-WGMMA instruction, are + visible to the MMA about to be issued. +2. :meth:`wgmma.mma() ` + computes ``d = a @ b + d``. A ``[block_m, block_k]`` by ``[block_k, block_n]`` + product is decomposed by the compiler into the hardware's native + ``64 x N x 16`` shapes and issued as a sequence of instructions. +3. :meth:`wgmma.commit_group() ` + closes a *commit group* over every MMA issued since the last commit. Groups + complete in order. +4. :meth:`wgmma.wait_group(n) ` + blocks until at most ``n`` commit groups are still pending. ``wait_group(0)`` + waits for everything. + +V1 uses ``wait_group(0)`` immediately after committing, which throws away the +asynchrony --- the warp group issues one MMA and stands still until it finishes. +That is deliberate: it keeps V1 a one-line change in behavior from V0. Keeping +groups in flight with ``wait_group(1)`` is what :doc:`V5 ` does once there is +a pipeline deep enough to feed it. + +.. note:: + + All four instructions must be executed by a **full warp group** --- 4 + consecutive warps, 128 threads, starting at a warp-group-aligned index. In V1 + the whole block is one warp group (``warps = 4``), so the plain block scope + satisfies this. From :doc:`V3 ` onward, where the block contains warps + with different jobs, WGMMA is issued inside an explicit + :meth:`~tilus.Script.thread_group`. + + +Walkthrough +----------- + +Setup and epilogue are identical to V0. Only the compute half of the main loop +changes. + +Main Loop +~~~~~~~~~ + +.. literalinclude:: ../../../../examples/hopper_matmul/matmul_v1.py + :language: python + :start-at: for offset_k in range(0, k_size, block_k): + :end-at: phase ^= 1 + :dedent: 8 + :caption: Main loop + +**Load phase** (unchanged from V0): one thread declares the transaction bytes, +two :meth:`tma.global_to_shared() ` +calls fetch the A and B tiles, and the ``mbarrier.wait`` plus +:meth:`~tilus.Script.sync` make the data visible block-wide. + +**Compute phase**: where V0 had ``load_shared`` twice followed by +:meth:`~tilus.Script.dot`, V1 has the four-step WGMMA sequence operating on +``sa`` and ``sb`` --- the shared tensors themselves. ``sb.transpose()`` is a view +that swaps the logical axes of the ``[block_n, block_k]`` tile into the +``[block_k, block_n]`` shape the MMA expects; no data is moved, and the transpose +is absorbed into the descriptor's stride encoding. + +The trailing :meth:`~tilus.Script.sync` still guards the shared buffers against +the next iteration's TMA. Note that it is only correct because +``wait_group(0)`` has already retired the MMA --- with an in-flight WGMMA, a +plain ``__syncthreads()`` would say nothing about whether the tensor core is +still reading ``sa``. + + +Performance +----------- + +Removing the register round trip is worth **1.8x**: 540 TFLOPS (2.04 ms), up from +V0's 305. Tensor pipe utilization rises from 59% to 67%. Freeing the operand +registers also lets the autotuner move up to a ``128 x 128`` tile with +``block_k=64``, twice V0's tile area, which is itself part of the gain. + +Note what did *not* change: the kernel is still load-then-compute with nothing +overlapping, so it remains far from cuBLAS. +The complete source is at :github:`examples/hopper_matmul/matmul_v1.py`. + +.. plot:: tutorials/matmul-hopper/plots/plot_v1.py + + Hopper matmul performance on H100 SXM (M=N=K=8192, fp16). Latency is + CUDA-event timed, median of three fresh processes. Peak is the published + dense FP16 tensor core throughput of the H100 SXM. + + +What's Next +----------- + +V1 is still **single-stage**: the loop waits for TMA to complete before issuing +the MMA, then waits for the MMA before starting the next TMA. Load and compute +are fully serialized, so the TMA engine idles during compute and the tensor cores +idle during load. We now have the right instruction, driven in the wrong shape. + +In :doc:`the next version `, we introduce **multi-stage software pipelining** +--- shared memory becomes a ring buffer with one barrier per stage, and the TMA +for iteration *i+1* is issued before waiting on iteration *i*, so loading and +computing finally overlap. diff --git a/docs/source/tutorials/matmul-hopper/v2.rst b/docs/source/tutorials/matmul-hopper/v2.rst new file mode 100644 index 00000000..5bbef28c --- /dev/null +++ b/docs/source/tutorials/matmul-hopper/v2.rst @@ -0,0 +1,282 @@ +.. _tutorial_hopper_matmul_v2: + +2. Multi-Stage Software Pipelining +=================================== + +In :doc:`V1 `, each loop iteration first waits for TMA to finish loading data, +then issues the WGMMA. Load and compute are fully serialized --- the TMA engine +sits idle during the MMA, and the tensor cores sit idle during the load. Both of +Hopper's asynchronous engines spend most of their time waiting for the other. + +This version introduces **multi-stage software pipelining**: shared memory is +divided into multiple stages (a ring buffer), and the kernel prefills several +stages before entering the main loop. In each iteration of the main loop, the TMA +loads data for a future iteration while the tensor cores process data from a +previously loaded stage. + +Pipelining is the right idea, and every later version keeps it. But V2 is also +the one version in this series that ends up *slower* than its predecessor, and +understanding why is more instructive than the speedup would have been: a ring +buffer indexed at runtime costs more than the overlap it buys, and the block-wide +``__syncthreads()`` around each stage puts a hard floor under how much overlap is +achievable at all. :doc:`V3 ` fixes both. + +If you have used Triton, this is similar to Triton's ``num_stages`` parameter --- +but here you control the pipelining explicitly: allocating per-stage buffers, +issuing prefill loads, and managing phase tracking yourself. + + +The Full Kernel +--------------- + +.. literalinclude:: ../../../../examples/hopper_matmul/matmul_v2.py + :language: python + :start-at: @tilus.autotune + :end-at: self.store_global(gc, casted_acc, offsets=[offset_m, offset_n]) + :caption: MatmulWGMMAV2 --- full kernel + + +What Changed from V1 +-------------------- + +.. list-table:: + :header-rows: 1 + :widths: 15 40 40 + + * - + - V1 + - V2 + * - **Shared memory** + - Single stage: ``[block_m, block_k]`` + - Multi-stage ring buffer: ``[num_stages, block_m, block_k]`` + * - **TMA barriers** + - 1 barrier + - 1 barrier **per stage** + * - **Phase tracking** + - Single ``phase`` scalar + - Per-stage ``phase`` register tensor + * - **Loop structure** + - Load then compute, serial + - Prefill stages, then overlap load and compute + * - **New parameter** + - --- + - ``num_stages`` (autotuned: 2, 3, or 4) + + +Software Pipelining +------------------- + +.. figure:: figures/v2_pipeline.svg + :width: 100% + :align: center + + Top: V1 serializes load and compute. Bottom: V2 overlaps them using a + multi-stage ring buffer. + +The idea is simple: if we have ``S`` stages of shared memory, we can have up to +``S`` TMA loads in flight while one stage is being consumed by the tensor cores. +The kernel proceeds in two phases: + +1. **Prefill** --- Before the main loop, issue TMA loads for the first ``S`` + K-tiles. These loads run asynchronously; the kernel does not wait for them. +2. **Main loop** --- Each iteration does three things: + + - **Wait**: block on the current stage's barrier until its TMA has landed. + - **Compute**: run WGMMA on the current stage's data. + - **Preload**: issue a TMA load for K-tile ``iter + S`` into the stage that + was just consumed. + + The stage index advances modulo ``num_stages``, cycling through the ring + buffer. + +The crucial reordering compared to V1 is that the *preload* for a future tile is +issued while the tensor cores still have work queued behind them. By the time the +loop comes back around to that stage, its data has already arrived, and the wait +costs nothing. + + +Multi-Stage Shared Memory +------------------------- + +In V1, shared tensors had shape ``[block_m, block_k]`` --- a single buffer that +was overwritten every iteration. In V2, shared tensors gain a leading stage +dimension: + +.. code-block:: python + + sa = self.shared_tensor(dtype=float16, shape=[self.num_stages, block_m, block_k]) + sb = self.shared_tensor(dtype=float16, shape=[self.num_stages, block_n, block_k]) + +Each stage ``sa[i]`` / ``sb[i]`` is an independent buffer. TMA writes to one stage +while WGMMA reads from another, without conflicts. This is also why ``num_stages`` +must be autotuned rather than simply maximized: the ring buffer is +``num_stages * (block_m + block_n) * block_k * 2`` bytes and has to fit in the +228 KB of shared memory an H100 SM can give a single block. Deeper pipelines +hide more latency, but force smaller tiles. + + +Per-Stage Barriers and Phase Tracking +------------------------------------- + +Each stage has its own mbarrier so that its TMA completion is tracked +independently: + +.. code-block:: python + + tma_barriers = self.mbarrier.alloc(counts=[1 for _ in range(self.num_stages)]) + phase = self.register_tensor(dtype=uint32, shape=[self.num_stages], init=0) + +V2 keeps a **per-stage phase**, held in a small register tensor, and flips +``phase[stage]`` each time that stage is consumed. This is the most direct way to +express the ring buffer: each barrier alternates between "filled" and "consumed" +on its own schedule, and the phase array simply remembers where each one is. + +.. hint:: + :class: margin + + :doc:`V3 ` replaces this with a single per-role phase scalar that flips on + wrap-around, which the compiler can keep in one register instead of + ``num_stages`` of them. + + +Loop Unrolling and Stage Indices +-------------------------------- + +There is a subtlety with a ring buffer: ``stage = iter % self.num_stages`` is a +runtime value, so every ``sa[stage]`` access needs an address computation, and +the compiler cannot see which barrier a given wait refers to. If instead the loop +body is unrolled by ``num_stages``, each unrolled copy has a *constant* stage +index --- the modulo folds away, addresses become compile-time offsets, and the +barrier waits resolve to specific barriers. + +V2 uses Python's ``range()`` and pays that cost. From :doc:`V3 ` onward the +loops switch to :meth:`self.range() ` with +``unroll=num_stages``: + +.. code-block:: python + + for offset_k in self.range(0, k_size, block_k, unroll=self.num_stages): + +Both are lowered to the same loop statement internally; ``self.range`` just +carries the extra unroll hint. + + +Walkthrough +----------- + +Prefill +~~~~~~~ + +.. literalinclude:: ../../../../examples/hopper_matmul/matmul_v2.py + :language: python + :start-at: for stage in range(max_num_stages): + :end-before: for iter in range(num_iters): + :dedent: 8 + :caption: Prefill: load the first num_stages tiles + +Before the main loop, one TMA load is issued per stage without waiting. Each +targets stage ``i`` and signals ``tma_barriers[i]``. ``max_num_stages`` guards +the case where the K loop is shorter than the pipeline depth --- with +``k_size / block_k < num_stages`` there is simply not enough work to fill every +stage, and issuing loads past the end of K would read out of bounds. + + +Main Loop +~~~~~~~~~ + +.. literalinclude:: ../../../../examples/hopper_matmul/matmul_v2.py + :language: python + :start-at: for iter in range(num_iters): + :end-before: # sa/sb are deliberately not freed + :dedent: 8 + :caption: Main loop: overlap preload and compute + +In each iteration: + +- **Wait** (on ``stage``): + :meth:`mbarrier.wait() ` + blocks until this stage's TMA data has arrived, using that stage's own phase. + The following :meth:`~tilus.Script.sync` publishes the arrival to the whole + block, since only one thread waited. + +- **Compute** (from ``stage``): the WGMMA sequence from V1, reading + ``sa[stage]`` and ``sb[stage]``. ``phase[stage] ^= 1`` prepares that stage's + barrier for its next cycle. + +- **Preload** (into ``preload_stage``): if K-tile ``iter + num_stages`` exists, + issue its TMA into the stage that was just freed. The guard + ``preload_iter < num_iters`` stops the pipeline from running past the end of K + in the final iterations, letting it drain naturally. + +The trailing :meth:`~tilus.Script.sync` closes the iteration: it must come +*after* the preload has been issued, so the loads for later stages are already in +flight when the next iteration begins. + +.. note:: + + Correctness here still leans on ``wgmma.wait_group(0)`` inside the loop. The + tensor cores fully retire stage ``i``'s MMA before the code reaches the point + where stage ``i`` is reused as a preload target, so a plain block-wide + ``sync`` is enough to protect the buffer. Once :doc:`V5 ` keeps a WGMMA + group in flight across iterations, that reasoning breaks and an explicit + producer-consumer handshake becomes mandatory. + + +Performance +----------- + +V2 measures **518 TFLOPS** (2.12 ms) --- about 4% *slower* than V1's 540. The +autotuner picks a 2-stage pipeline on the same ``128 x 128`` tile with +``block_k=64`` that V1 chose, so this is a clean like-for-like comparison, and the +pipelining genuinely does not pay for itself here. Nsight Compute shows where the +overlap went: DRAM throughput jumps from 25% to 68%, while tensor pipe +utilization moves only from 67% to 71%. The ring buffer is keeping the memory +system busy, and almost none of that is reaching the tensor cores. + +.. note:: + + V1 and V2 are close enough that the two measurement methods disagree on the + order: under Nsight Compute's replay clock V2 profiles slightly *faster* than + V1 (2.12 ms vs 2.15 ms), while CUDA-event timing at full boost clock puts it + slower. Wall clock is the ranking authority throughout this tutorial; the NCU + columns are used only to explain *why*. + +Three costs eat the gain: + +- **Runtime stage indexing.** The loop is a plain ``range()``, so + ``stage = iter % num_stages`` is a runtime value. Every ``sa[stage]`` access + needs address arithmetic, and ``phase[stage]`` is a register tensor indexed by + a runtime value --- which the compiler cannot keep in registers. +- **Two block-wide syncs per iteration.** These are unchanged from V1, and they + serialize the very phases the ring buffer is trying to overlap. +- **Only two stages.** Deeper pipelines were available in the search space but + lose to shallower ones, because at this tile size the extra shared memory does + not buy proportionally more latency hiding. + +The lesson is that a ring buffer alone does not create overlap --- it only creates +the *opportunity* for it. As long as all 128 threads must meet at a barrier +between loading and computing, the opportunity goes unused. +The complete source is at :github:`examples/hopper_matmul/matmul_v2.py`. + +.. plot:: tutorials/matmul-hopper/plots/plot_v2.py + + Hopper matmul performance on H100 SXM (M=N=K=8192, fp16). Latency is + CUDA-event timed, median of three fresh processes. Peak is the published + dense FP16 tensor core throughput of the H100 SXM. + + +What's Next +----------- + +V2 overlaps TMA loads with WGMMA compute across iterations, but there is still a +structural limitation: **every thread does every job**. The same 128 threads +issue the TMA, wait on the barrier, run the MMA, and wait for it --- separated by +``__syncthreads()`` calls that force the whole block into lockstep at each +transition. The tensor cores cannot run ahead, because the warps that would issue +the next MMA are parked in a block-wide barrier. + +In :doc:`the next version `, we split the block by role: a dedicated +**producer warp** that does nothing but issue TMA loads, and a **consumer warp +group** that does nothing but run WGMMA. They communicate through a pair of +producer/consumer barriers instead of ``__syncthreads()``, so each can run at its +own pace. diff --git a/docs/source/tutorials/matmul-hopper/v3.rst b/docs/source/tutorials/matmul-hopper/v3.rst new file mode 100644 index 00000000..7524ca96 --- /dev/null +++ b/docs/source/tutorials/matmul-hopper/v3.rst @@ -0,0 +1,260 @@ +.. _tutorial_hopper_matmul_v3: + +3. Warp Specialization +====================== + +:doc:`V2 ` overlaps TMA and WGMMA across iterations, but every warp in the +block does every job. The same 128 threads issue the TMA, wait on its barrier, +run the MMA, and wait for it --- with a ``__syncthreads()`` at each transition +that forces the whole block into lockstep. The tensor cores cannot run ahead of +the loader, because the warps that would issue the next MMA are sitting in a +block-wide barrier. + +This version introduces **warp specialization**: warps are given *different jobs* +and run *different code*. One warp becomes a dedicated **producer** that does +nothing but issue TMA loads; the remaining four warps become a **consumer** warp +group that does nothing but run WGMMA. They never meet at a ``__syncthreads()``; +instead they hand stages back and forth through a pair of mbarriers. + +Triton also performs warp specialization internally, but as a compiler pass with +no user-level control. In Tilus you explicitly assign roles to warps and define +how they communicate. + + +The Full Kernel +--------------- + +.. literalinclude:: ../../../../examples/hopper_matmul/matmul_v3.py + :language: python + :start-at: @tilus.autotune + :end-at: self.store_global(gc, casted_acc, offsets=[offset_m, offset_n]) + :caption: MatmulWGMMAV3 --- full kernel + + +What Changed from V2 +-------------------- + +.. list-table:: + :header-rows: 1 + :widths: 15 40 40 + + * - + - V2 + - V3 + * - **Warp structure** + - 4 warps, all doing everything + - 5 warps: 1 TMA producer + 4-warp consumer group + * - **Barriers** + - TMA barrier per stage + - ``consumer_barriers`` + ``producer_barriers`` per stage + * - **Synchronization** + - ``__syncthreads()`` twice per iteration + - None in the loop --- only mbarrier handshakes + * - **Phase tracking** + - Per-stage phase array + - Per-role phase array, one per participant + * - **Prefill** + - Explicit prefill loop + - Implicit: the producer runs ahead on its own + * - **Loops** + - ``range()`` + - :meth:`self.range() ` with ``unroll=num_stages`` + * - **New instructions** + - + - :meth:`~tilus.Script.thread_group`, + :meth:`~tilus.lang.instructions.mbarrier.BarrierInstructionGroup.arrive` + + +Why a Separate Producer Warp +---------------------------- + +.. figure:: figures/v3_warp_specialization.svg + :width: 100% + :align: center + + V2 alternates roles inside one warp group, gated by ``__syncthreads()``. + V3 gives the TMA its own warp, so the producer can run several K-tiles ahead + of the consumer. + +A TMA load is issued by a *single thread* --- the rest of the warp contributes +nothing to it. In V2 that thread is part of the same warp group that runs WGMMA, +so issuing the next load means the warp group is not issuing MMA, and the +block-wide sync means no other warp can cover for it. + +Splitting the roles fixes both problems: + +- **TMA warp** (threads 128--159): loops over K-tiles issuing loads back-to-back. + Before filling a stage, it waits only on ``producer_barriers[stage]`` to + confirm the consumer is finished with that slot. +- **Consumer warp group** (threads 0--127): loops over K-tiles issuing WGMMA + back-to-back. Before each MMA it waits only on ``consumer_barriers[stage]`` to + confirm the data has landed. + +Neither ever waits for the other's *code* --- only for a specific stage's data +dependency. The producer naturally runs ``num_stages`` tiles ahead, so the +consumer's wait is usually already satisfied when it arrives. + +.. note:: + + ``warps = 5`` is not a typo, and the ordering matters. The consumer group + occupies warps 0--3 because WGMMA requires a **warp-group-aligned** span of + four consecutive warps; warp 4 is left over for the producer. Putting the + producer first would push the consumer to warps 1--4, which is not a valid + warp group. + + +Producer-Consumer Barriers +-------------------------- + +V2 used one barrier per stage to signal "TMA has landed". That is only half the +handshake --- it says when a stage becomes *full*, but nothing about when it +becomes *empty* again, which V2 got for free from ``__syncthreads()``. Without +the block-wide sync, both directions must be explicit: + +.. code-block:: python + + consumer_barriers = self.mbarrier.alloc(counts=[1 for _ in range(self.num_stages)]) + producer_barriers = self.mbarrier.alloc(counts=[128 for _ in range(self.num_stages)]) + +- ``consumer_barriers[i]``: signaled by the TMA engine's tx-count when stage + ``i`` has been filled. The consumer waits on these. Arrival count is **1**, + since a single thread declares the transaction bytes. +- ``producer_barriers[i]``: signaled when the consumer has finished reading stage + ``i``. The producer waits on these. Arrival count is **128**, because every + thread of the consumer warp group executes + :meth:`mbarrier.arrive() ` + after ``wgmma.wait_group(0)``. + +The **initial phases** are what make the pipeline start correctly: + +- ``producer_phases`` starts at **1**. All mbarriers begin at hardware phase 0, so + a wait expecting phase 1 does not match and passes immediately. That is exactly + right: every stage starts empty, and the producer should begin filling without + blocking. +- ``consumer_phases`` starts at **0**, which *does* match, so the consumer blocks + until the producer's first load actually completes. + +.. hint:: + :class: margin + + Tilus exposes these two values as + ``self.mbarrier.producer_initial_phase`` and + ``self.mbarrier.consumer_initial_phase``, which :doc:`V4 ` uses instead of + hard-coded literals. + + +Draining the Pipeline +--------------------- + +The producer's main loop exits after issuing the last K-tile, but at that moment +up to ``num_stages`` loads are still in flight and the consumer is still working +through them. If the producer warp simply exits, its threads leave the block +while the consumer is still arriving on ``producer_barriers`` --- so V3 adds a +drain loop that consumes the outstanding empty-signals without issuing anything: + +.. literalinclude:: ../../../../examples/hopper_matmul/matmul_v3.py + :language: python + :start-at: for _ in self.range(min(self.num_stages, cdiv(k_size, self.block_k))): + :end-at: stage = (stage + 1) % self.num_stages + :dedent: 12 + :caption: Producer drain loop + +The ``min(...)`` handles the short-K case for the same reason as V2's +``max_num_stages``: when there are fewer K-tiles than stages, fewer stages were +ever filled, so fewer signals will arrive. + + +Walkthrough +----------- + +TMA Warp (Producer) +~~~~~~~~~~~~~~~~~~~ + +.. literalinclude:: ../../../../examples/hopper_matmul/matmul_v3.py + :language: python + :start-at: with self.thread_group(thread_begin=128, num_threads=32): + :end-before: with self.thread_group(thread_begin=0, num_threads=128): + :dedent: 8 + :caption: TMA warp + +Each iteration: + +- :meth:`mbarrier.wait() ` + on ``producer_barriers[stage]`` blocks until the consumer has released this + stage, then the local phase for that stage flips. +- Inside :meth:`~tilus.Script.single_thread`, + :meth:`mbarrier.arrive_and_expect_tx() ` + declares the bytes for both tiles on ``consumer_barriers[stage]``. +- Two :meth:`tma.global_to_shared() ` + calls load A and B into ``sa[stage]`` / ``sb[stage]``. In V3 these sit inside + the same ``single_thread`` block as the declaration --- the simplest thing that + works. :doc:`V4 ` moves them out to warp scope, which is the form the later + versions use. +- The stage index advances modulo ``num_stages``. + +Note there is no explicit prefill loop as in V2. The producer simply starts +running, and because ``producer_phases`` starts at 1, its first +``num_stages`` waits all pass immediately. + + +Consumer Warp Group +~~~~~~~~~~~~~~~~~~~ + +.. literalinclude:: ../../../../examples/hopper_matmul/matmul_v3.py + :language: python + :start-at: with self.thread_group(thread_begin=0, num_threads=128): + :end-at: self.store_global(gc, casted_acc, offsets=[offset_m, offset_n]) + :dedent: 8 + :caption: Consumer warp group + +The consumer runs the matching loop: + +- :meth:`mbarrier.wait() ` + on ``consumer_barriers[stage]`` blocks until the TMA data has arrived. +- The WGMMA sequence computes on ``sa[stage]`` and ``sb[stage]``. +- :meth:`mbarrier.arrive() ` + on ``producer_barriers[stage]`` releases the stage. It comes *after* + ``wgmma.wait_group(0)``, which is what makes the release safe: the tensor cores + have finished reading shared memory, so the producer may overwrite it. +- The epilogue runs entirely within the consumer group, which is convenient --- + the accumulator lives in these 128 threads' registers, so no data movement is + needed to reach the ``store_global``. + + +Performance +----------- + +Warp specialization lifts the kernel to **575 TFLOPS** (1.91 ms), 11% ahead of V2 +and 6% ahead of V1. The autotuner chooses the *same* configuration as V2 --- 2 +stages, ``128 x 128``, ``block_k=64`` --- so the entire gain comes from the +restructuring: removing the block-wide syncs, unrolling the ring buffer so stage +indices become constants, and letting the producer run ahead on its own warp. +Tensor pipe utilization rises to 75%, and DRAM throughput settles at 61%. + +This is also where V2's investment finally pays off. Pipelining and warp +specialization are complementary: the ring buffer provides the slots, and warp +specialization provides the independent execution needed to fill them. +The complete source is at :github:`examples/hopper_matmul/matmul_v3.py`. + +.. plot:: tutorials/matmul-hopper/plots/plot_v3.py + + Hopper matmul performance on H100 SXM (M=N=K=8192, fp16). Latency is + CUDA-event timed, median of three fresh processes. Peak is the published + dense FP16 tensor core throughput of the H100 SXM. + + +What's Next +----------- + +V3 achieves true overlap between TMA and WGMMA. The remaining bottleneck is on +the compute side: there is exactly **one** consumer warp group, and it issues one +MMA and immediately waits for it. Between the ``wait_group(0)`` and the next +``mbarrier.wait``, the tensor core pipeline has nothing queued and drains. +Feeding it faster is not a matter of loading faster --- it needs *more +independent MMA work* available at any instant. + +In :doc:`the next version `, we split the output tile across **two consumer +warp groups**, each owning half the rows of C, so two independent WGMMA streams +share the same loaded B tile. We also refactor the barrier bookkeeping into a +reusable ``Pipeline`` class, since the number of barriers, phases, and stage +counters is about to grow. diff --git a/docs/source/tutorials/matmul-hopper/v4.rst b/docs/source/tutorials/matmul-hopper/v4.rst new file mode 100644 index 00000000..a1a18fe9 --- /dev/null +++ b/docs/source/tutorials/matmul-hopper/v4.rst @@ -0,0 +1,344 @@ +.. _tutorial_hopper_matmul_v4: + +4. Pipeline Abstraction and Two Consumer Warp Groups +===================================================== + +:doc:`V3 ` decoupled loading from computing, but the compute side is still +narrow: **one** consumer warp group issues one WGMMA and immediately waits for +it. Between ``wgmma.wait_group(0)`` and the next ``mbarrier.wait``, the tensor +core pipeline has nothing queued and drains. More bandwidth will not help --- the +kernel needs more *independent MMA work* available at any instant. + +This version adds two things: + +1. **Two consumer warp groups** --- the output tile is split by rows, and each + warp group computes its own half against the same shared B tile. Two + independent WGMMA streams now feed the tensor cores. +2. **Pipeline abstraction** --- the barrier, phase, and stage bookkeeping from V3 + is encapsulated in a reusable ``Pipeline`` class built on ``tilus.Class``. With + three participants instead of two, and more pipelines coming in later + versions, the inline bookkeeping has outgrown its welcome. + +The kernel also gains **tile rasterization**: a 1D grid remapped so that +concurrently running blocks share B tiles in L2. V4 turns it on with +``swizzle_size=4``, and :doc:`V5 ` and :doc:`V6 ` keep it. + + +The Full Kernel +--------------- + +.. literalinclude:: ../../../../examples/hopper_matmul/matmul_v4.py + :language: python + :start-at: class Pipeline + :end-at: offsets=[offset_m + block_m_half, offset_n], + :caption: MatmulWGMMAV4 --- full kernel (including Pipeline class) + + +What Changed from V3 +-------------------- + +.. list-table:: + :header-rows: 1 + :widths: 15 40 40 + + * - + - V3 + - V4 + * - **Warp structure** + - 5 warps: 1 producer + 1 consumer group + - 9 warps: 1 producer + **2** consumer groups + * - **Output tile** + - One ``block_m x block_n`` accumulator in one warp group + - Split by rows: each group owns ``block_m/2 x block_n`` + * - **A in shared memory** + - ``[stages, block_m, block_k]`` + - ``[stages, 2, block_m/2, block_k]`` --- one slab per consumer + * - **B in shared memory** + - ``[stages, block_n, block_k]`` + - Unchanged --- **shared by both** consumer groups + * - **Barrier management** + - Manual barriers, phases, and stage indices + - ``Pipeline`` class (``tilus.Class``) encapsulates the bookkeeping + * - **Empty-barrier arrivals** + - 128 (every consumer thread) + - 2 (one elected thread per consumer group) + * - **Pipeline depth** + - 2 stages + - 3 stages + * - **Grid layout** + - 2D grid + - 1D grid with swizzled rasterization (``swizzle_size=4``) + * - **New instructions** + - + - :meth:`~tilus.Script.fast_divmod`, ``tilus.Class`` + + +Two Consumer Warp Groups +------------------------ + +.. figure:: figures/v4_tile_split.svg + :width: 100% + :align: center + + The ``block_m x block_n`` output tile is split by rows across two consumer + warp groups. Each loads its own A slab; both read the same B tile. + +A WGMMA instruction is issued by one warp group and its accumulator lives in that +group's registers. To get two MMAs in flight, we need two warp groups, and they +need separate accumulators --- so the natural split is by **rows of C**: + +- Consumer WG0 (threads 0--127) computes rows ``[0, block_m/2)`` of the tile. +- Consumer WG1 (threads 128--255) computes rows ``[block_m/2, block_m)``. +- The producer warp (threads 256--287) feeds both. + +The split has a pleasant property for memory traffic: the two halves need +**different rows of A** but the **same columns of B**. So A is stored as two +slabs, ``sa[stage, 0]`` and ``sa[stage, 1]``, one per consumer, while ``sb`` stays +a single tile that both groups read. Splitting the accumulator across two warp +groups also halves the per-thread register pressure of the accumulator, which is +what allows the tuned tile to grow from ``128 x 128`` (V3) to ``128 x 256``. + +.. note:: + + ``warps = 9`` again reflects warp-group alignment: consumers occupy warps + 0--3 and 4--7 (both warp-group aligned), leaving warp 8 for the producer. + +Because two warp groups now share each stage, the empty-barrier arrival count +changes. In V3 all 128 consumer threads arrived; here each group elects a single +thread: + +.. literalinclude:: ../../../../examples/hopper_matmul/matmul_v4.py + :language: python + :start-at: tma_pipe = Pipeline(num_stages, producer_arrive_count=1, consumer_arrive_count=2) + :end-at: tma_pipe = Pipeline(num_stages, producer_arrive_count=1, consumer_arrive_count=2) + :dedent: 8 + :caption: Two arrivals per stage, one per consumer group + +``consumer_arrive_count=2`` means the producer may refill a stage only after +*both* consumer groups have released it. Electing one thread per group (rather +than letting all 256 arrive) turns 256 barrier updates per K-tile into 2. + + +Pipeline Abstraction +-------------------- + +On Hopper, mbarriers are the mechanism for tracking asynchronous work, and shared +memory is the buffer for data in transit. When producer and consumer run at +different speeds --- always, in practice --- a **pipeline** decouples them. + +A pipeline has three components: + +1. **Producer** --- generates data and writes it into a buffer slot when one is + available. +2. **Consumer** --- reads data from a slot when one is filled. +3. **Ring buffer** --- a fixed number of slots (``num_stages``) that producer and + consumer cycle through independently. + +Each slot carries two mbarriers: + +- **full barrier** --- signaled when the producer has filled the slot. Consumers + wait on this. +- **empty barrier** --- signaled when the consumers have drained the slot. The + producer waits on this. + +Producer and consumer each keep a **stage pointer** and a **phase variable**, and +advance through the ring independently, synchronized only by barrier signals. + +.. figure:: figures/v4_pipeline_class.svg + :width: 100% + :align: center + + A 5-stage pipeline. The producer is filling slot 3 while the consumers drain + slot 1; slot 2 is full and waiting, slots 0 and 4 are empty. The check marks + indicate whether each slot's full/empty mbarrier has completed. + +V3 managed all of this inline. The ``Pipeline`` class below packages it behind a +small API. Note that this is not a built-in part of Tilus --- it is assembled +from ordinary instructions (``mbarrier.alloc``, ``mbarrier.wait``, ...) as a +user-level helper. Managing the barriers by hand, as in V3, remains perfectly +valid. + +.. literalinclude:: ../../../../examples/hopper_matmul/matmul_v4.py + :language: python + :start-at: class Pipeline + :end-at: self.consumer_phase = self.consumer_phase ^ (self.consumer_stage == 0) + :caption: Pipeline class + +``Pipeline`` inherits from ``tilus.Class``, which behaves like +:class:`~tilus.Script` but for helper objects that are not kernels themselves: it +can allocate barriers and shared tensors and use any Tilus instruction. Two +details are worth pointing out: + +- The phases are initialized from ``self.mbarrier.producer_initial_phase`` and + ``self.mbarrier.consumer_initial_phase`` rather than the literals ``1`` and + ``0`` that V3 used. +- ``producer_advance`` / ``consumer_advance`` flip the phase **on wrap-around** + (``phase ^= (stage == 0)``) rather than keeping a per-stage array as V2 did. + One scalar per role replaces ``num_stages`` registers, and after the loop is + unrolled by ``num_stages`` the compiler resolves each stage index to a + constant. + +The kernel-side usage reads cleanly: + +.. code-block:: python + + tma_pipe.producer_acquire() # wait for an empty slot + # ... issue TMA loads against tma_pipe.producer_barrier() ... + tma_pipe.producer_advance() + + tma_pipe.consumer_acquire() # wait for a full slot + # ... issue WGMMA on tma_pipe.consumer_stage ... + self.mbarrier.arrive(tma_pipe.consumer_barrier()) # release the slot + tma_pipe.consumer_advance() + + +Tile Rasterization +------------------ + +V4 also introduces the grid-remapping machinery that :doc:`V5 ` relies on. +Each output tile (m, n) needs a row-strip of A and a column-strip of B. A rows +are unique per tile, but **B columns are shared by every tile in the same +N-column** --- so B traffic can be served from L2 if the tiles that share it run +at the same time. + +The question is how to order tiles so that the set of A rows and B columns +touched by the concurrently running blocks --- the **L2 working set** --- stays +small. + +.. figure:: figures/v4_tile_rasterization.svg + :width: 100% + :align: center + + An 8 x 8 tile grid with a wave of 16 active blocks. Orange bars mark active + A rows; blue bars mark active B columns. Swizzling yields a smaller working + set (8 vs 10 strips) for the same number of active blocks. + +With a plain 2D grid, ``blockIdx.x`` walks down M first, so a wave of 16 blocks +fills two full columns: 8 A rows plus 2 B columns = 10 strips resident. Grouping +the same 16 blocks into a 4 x 4 square touches 4 A rows plus 4 B columns = 8 +strips --- 20% less L2 pressure. The mapping divides the N axis into groups of +``swizzle_size`` columns and assigns tiles within a group in row-major order: + +.. literalinclude:: ../../../../examples/hopper_matmul/matmul_v4.py + :language: python + :start-at: def compute_block_coord + :end-at: return m_block, n_block + :dedent: 4 + :caption: Tile rasterization with swizzle grouping + +When ``num_n_blocks`` is not divisible by ``swizzle_size``, the final group is +narrower; ``last_group_width`` handles that case so the mapping stays a bijection. + +.. hint:: + + Integer division and modulo are expensive on GPUs. For compile-time constant + divisors (like ``swizzle_size``) the compiler emits a multiply and shift + automatically. For **grid-constant** divisors --- the same for every block, but + not known at compile time, like ``tiles_per_group`` --- + :meth:`~tilus.Script.fast_divmod` precomputes a magic number once per launch + and uses integer multiply-shift instead of the compiler's floating-point + fallback. + +V4's tuned configuration selects ``swizzle_size=4``. The kernel also keeps a +``swizzle_size=1`` bypass that launches a plain 2D grid with no remapping at all; +because ``swizzle_size`` is a compile-time autotune constant, whichever branch +applies is resolved while tracing, so the unused one costs nothing. + + +Pipeline Depth +-------------- + +V4 also deepens the ring buffer from two stages to three, and this matters more +than it looks. The consumer here is still **synchronous** --- ``wait_group(0)`` +after every commit --- so it stops dead at the end of each K-tile. With only two +stages the producer can be at most one tile ahead, and every consumer stall is a +producer stall shortly after. Measured on an H100, a 2-stage V4 runs at +1.90 ms --- statistically indistinguishable from V3's 1.91 ms, i.e. the two +consumer warp groups buy nothing at all. Three stages gives the producer enough +slack to stay ahead of the drain, and the same kernel drops to 1.71 ms. + +The obvious next question --- why not four stages, or five --- is what +:doc:`V5 ` answers: past three, depth alone stops helping, and what the +kernel needs instead is for the consumer to stop draining. + + +Walkthrough +----------- + +Producer Warp +~~~~~~~~~~~~~ + +.. literalinclude:: ../../../../examples/hopper_matmul/matmul_v4.py + :language: python + :start-at: with self.thread_group(thread_begin=256, num_threads=32): # TMA producer warp + :end-before: with self.thread_group(thread_begin=0, num_threads=128): # consumer WG0 + :dedent: 8 + :caption: TMA producer warp + +The structure matches V3's producer, now expressed through the Pipeline API. +Three TMA loads are issued per stage instead of two: one per A slab, plus B. The +``arrive_and_expect_tx`` declares all three tiles' bytes at once, so a single +barrier tracks the whole stage. + +Note the placement of :meth:`~tilus.Script.single_thread`: it wraps only the +``arrive_and_expect_tx``, not the TMA calls. The transaction-byte declaration +must happen exactly once, but the loads themselves are issued at warp +granularity. + + +Consumer Warp Groups +~~~~~~~~~~~~~~~~~~~~ + +.. literalinclude:: ../../../../examples/hopper_matmul/matmul_v4.py + :language: python + :start-at: with self.thread_group(thread_begin=0, num_threads=128): # consumer WG0 + :end-before: with self.thread_group(thread_begin=128, num_threads=128): # consumer WG1 + :dedent: 8 + :caption: Consumer warp group 0 + +Each consumer group runs the same loop against its own A slab (``consumer_idx`` +0 or 1) and its own accumulator, then writes its half of C with +:meth:`~tilus.Script.store_global`. The two groups are otherwise identical, and +the second differs only in the slab index and the row offset of its store. + +The MMA is still **synchronous** --- ``wait_group(0)`` after every commit --- so +each group finishes its MMA before releasing the stage. The parallelism gained +here comes from having *two* groups doing that at once, not from overlapping +within a group. Overlapping within a group is V5's job, and it needs a deeper +pipeline to be safe. + + +Performance +----------- + +V4 reaches **642 TFLOPS** (1.71 ms), **12% ahead of V3**, and wins in every one of +three fresh processes. Two counters explain it. Tensor pipe utilization rises +from 75% to 80% --- the second consumer group supplies the independent MMA work +V3 could not. And DRAM throughput falls from 61% to 27%, because both groups read +the same B tile and the 4-wide raster keeps those tiles resident in L2. V4 turns +a memory-hungry kernel into a comfortably compute-bound one. + +The three ingredients are not separable: two consumer groups without the extra +pipeline stage measure no faster than V3 at all, and the raster only pays once +the kernel is issuing enough MMA to be sensitive to B-tile latency. +The complete source is at :github:`examples/hopper_matmul/matmul_v4.py`. + +.. plot:: tutorials/matmul-hopper/plots/plot_v4.py + + Hopper matmul performance on H100 SXM (M=N=K=8192, fp16). Latency is + CUDA-event timed, median of three fresh processes. Peak is the published + dense FP16 tensor core throughput of the H100 SXM. + + +What's Next +----------- + +V4 doubles the number of independent MMA streams, but each stream is still +strictly serial: issue, commit, **wait**, release, repeat. The tensor cores drain +between every K-tile of every group, and at 80% tensor utilization those drains +are now the largest remaining gap. Adding pipeline stages cannot close it --- a +deeper buffer feeds a consumer that keeps stopping. + +In :doc:`the next version `, the consumer stops stopping: we keep a WGMMA +group **in flight** across iterations with ``wait_group(1)``, issuing K-tile +*i+1*'s MMA before waiting on K-tile *i*'s. diff --git a/docs/source/tutorials/matmul-hopper/v5.rst b/docs/source/tutorials/matmul-hopper/v5.rst new file mode 100644 index 00000000..cfde038e --- /dev/null +++ b/docs/source/tutorials/matmul-hopper/v5.rst @@ -0,0 +1,245 @@ +.. _tutorial_hopper_matmul_v5: + +5. Overlapping WGMMA Groups +============================ + +:doc:`V4 ` runs two consumer warp groups, but each one is strictly serial: +issue, commit, **wait**, release, repeat. Every K-tile, both groups stop at +``wgmma.wait_group(0)`` until the tensor cores are completely done. The tensor +core pipeline therefore drains once per K-tile per group, and the barrier +handshake that follows sits squarely on the critical path. + +WGMMA is asynchronous precisely so this is avoidable. This version keeps **one +WGMMA group in flight at all times**: the consumer issues K-tile *i+1*'s MMA and +only then waits for K-tile *i* to finish, using ``wait_group(1)`` instead of +``wait_group(0)``. While the tensor cores work on tile *i+1*, the warp group is +free to release stage *i*, wait on the next barrier, and issue again. + +Keeping an MMA in flight has a consequence: the shared memory it reads is still +live. The stage release must therefore lag one iteration behind, and the pipeline +must be deep enough to absorb that lag. Everything else carries over from +:doc:`V4 ` untouched. + + +The Full Kernel +--------------- + +.. literalinclude:: ../../../../examples/hopper_matmul/matmul_v5.py + :language: python + :start-at: class Pipeline + :end-at: self.store_global(gc, casted1, offsets=[offset_m + block_m_half, offset_n]) + :caption: MatmulWGMMAV5 --- full kernel (including Pipeline class) + + +What Changed from V4 +-------------------- + +.. list-table:: + :header-rows: 1 + :widths: 15 40 40 + + * - + - V4 + - V5 + * - **MMA completion** + - ``wait_group(0)`` --- drain after every commit + - ``wait_group(1)`` --- one group stays in flight + * - **Loop shape** + - Uniform loop over all K-tiles + - Prologue MMA, steady-state loop, epilogue drain + * - **Stage release** + - Current stage, after the MMA completes + - **Previous** stage, via ``prev_consumer_barrier()`` + * - **Pipeline depth** + - 3 stages + - 4 stages (3 also works; see below) + * - **Rasterization** + - ``swizzle_size=4`` + - unchanged + * - **New Pipeline method** + - + - ``prev_consumer_barrier()`` + + +Keeping a WGMMA Group in Flight +------------------------------- + +.. figure:: figures/v5_wgmma_overlap.svg + :width: 100% + :align: center + + ``wait_group(0)`` drains the tensor core pipeline every K-tile. + ``wait_group(1)`` allows the next MMA to be issued first, so the tensor cores + always have work queued. + +Recall the WGMMA protocol: :meth:`wgmma.commit_group() ` +closes a group over the MMAs issued since the last commit, groups complete in +order, and :meth:`wgmma.wait_group(n) ` +blocks until at most ``n`` groups remain pending. + +``wait_group(1)`` says: *"let one group still be running."* Restructuring the +loop around that gives: + +.. code-block:: text + + prologue: acquire stage 0, fence, mma(0), commit # 1 group pending + steady: acquire stage i, fence, mma(i), commit # 2 groups pending + wait_group(1) # mma(i-1) is done + release stage i-1 + epilogue: wait_group(0) # mma(last) is done + release last stage + +The MMA for tile *i* is issued **before** the wait for tile *i-1*. From the +tensor cores' perspective there is no gap: the moment tile *i-1* retires, tile +*i* is already queued behind it. + +The price is that the release must shift. When ``wait_group(1)`` returns, only +tile *i-1*'s MMA has certainly completed --- tile *i*'s is still reading +``sa[stage_i]`` and ``sb[stage_i]``. Releasing the *current* stage here would let +the producer overwrite shared memory that the tensor cores are actively reading. +So V5 adds ``prev_consumer_barrier()``: + +.. literalinclude:: ../../../../examples/hopper_matmul/matmul_v5.py + :language: python + :start-at: def prev_consumer_barrier(self) -> RegisterTensor: + :end-at: return self.empty_barriers[prev_stage] + :dedent: 4 + :caption: Releasing the stage one behind the current one + +Because the consumer now holds two stages at once (one being read by the tensor +cores, one just acquired), the ring buffer effectively loses a slot. A 2-stage +buffer still runs correctly --- the release of stage *i-1* always precedes the +acquire of stage *i+1* --- but it leaves the producer no slack at all, and the +kernel falls to 2.24 ms, well behind V4. Measured across the depths that fit: + +.. list-table:: + :header-rows: 1 + :widths: 20 20 60 + + * - ``num_stages`` + - Latency + - + * - 2 + - 2.24 ms + - correct, but the producer can never run ahead + * - 3 + - 1.60 ms + - enough slack for the overlap to pay off + * - 4 + - 1.62 ms + - what the checked-in config uses; a tie with 3 + * - 5 + - --- + - does not fit: ``5 x 48 KB`` exceeds the 228 KB limit + +So three stages is where the overlap starts working, and the fourth is free +rather than necessary. The kernel ships with four. + +.. note:: + + This is the point where the informal reasoning of :doc:`V2 ` --- "the MMA + has retired, so a block-wide ``sync`` protects the buffer" --- stops being + valid. With an MMA in flight, no ``__syncthreads()`` tells you anything about + what the tensor cores are still reading. Only the WGMMA group counter does, + which is why the empty-barrier arrival is placed immediately after + ``wait_group(1)`` and refers to the previous stage. + + +Everything Else Is Unchanged +---------------------------- + +Worth stating explicitly, because it makes the attribution clean: V5 keeps V4's +two consumer warp groups, its ``128 x 256`` tile, its ``Pipeline`` class, and its +``swizzle_size=4`` rasterization exactly as they were. The only differences are +``wait_group(1)`` in place of ``wait_group(0)``, the lagging stage release that +requires, and the pipeline depth that makes the lag comfortable. + +So the speedup measured below is attributable to the overlap alone. + + +Walkthrough +----------- + +Producer Warp +~~~~~~~~~~~~~ + +.. literalinclude:: ../../../../examples/hopper_matmul/matmul_v5.py + :language: python + :start-at: with self.thread_group(thread_begin=256, num_threads=32): # TMA producer + :end-before: with self.thread_group(thread_begin=0, num_threads=128): # consumer WG0 + :dedent: 8 + :caption: TMA producer warp + +Unchanged from V4 apart from the deeper ring buffer: acquire an empty stage, +declare the transaction bytes for both A slabs and B, issue three TMA loads, +advance. The drain loop at the end absorbs the trailing empty-signals so the warp +does not exit while consumers are still releasing stages. + + +Consumer Warp Group +~~~~~~~~~~~~~~~~~~~ + +.. literalinclude:: ../../../../examples/hopper_matmul/matmul_v5.py + :language: python + :start-at: with self.thread_group(thread_begin=0, num_threads=128): # consumer WG0 + :end-before: with self.thread_group(thread_begin=128, num_threads=128): # consumer WG1 + :dedent: 8 + :caption: Consumer warp group 0 + +The three-part structure is explicit in the code: + +- **Prologue** --- acquire stage 0, fence, MMA, commit, advance. No wait: this + first group is deliberately left in flight. +- **Steady state** --- the loop starts at ``block_k`` rather than 0, because tile + 0 was already issued. Each iteration acquires the next stage, issues and + commits its MMA, then ``wait_group(1)`` retires the *previous* MMA, and one + elected thread arrives on ``prev_consumer_barrier()``. +- **Epilogue** --- ``wait_group(0)`` retires the final MMA, its stage is released, + and the accumulator is cast to fp16 and stored. + +Consumer WG1 is identical except that it reads A slab 1 and stores to the lower +half of the output tile. + +.. note:: + :class: margin + + :meth:`~tilus.Script.single_thread` elects exactly one thread of the warp + group to arrive, matching the pipeline's ``consumer_arrive_count=2`` --- one + arrival per consumer group, not per thread. + + +Performance +----------- + +V5 reaches **678 TFLOPS** (1.62 ms), 6% ahead of V4 and 91% of cuBLAS. Tensor +pipe utilization climbs from 80% to 88%, which is exactly the metric this change +targets --- the tensor cores now almost always have a queued group to start on the +cycle the previous one retires. Since everything else is inherited unchanged from +V4, the gain is attributable to the overlap alone: holding V4's three stages and +changing only ``wait_group(0)`` to ``wait_group(1)`` already gets 1.60 ms. + +V5 is also the last version that is numerically like-for-like with cuBLAS. It +accumulates in fp32, as V0--V4 do; :doc:`V6 ` gives that up. +The complete source is at :github:`examples/hopper_matmul/matmul_v5.py`. + +.. plot:: tutorials/matmul-hopper/plots/plot_v5.py + + Hopper matmul performance on H100 SXM (M=N=K=8192, fp16). Latency is + CUDA-event timed, median of three fresh processes. Peak is the published + dense FP16 tensor core throughput of the H100 SXM. + + +What's Next +----------- + +V5 keeps the tensor cores fed within each of its two consumer groups, and Nsight +Compute confirms it: tensor pipe utilization reaches about 88%, up from 80% in +V4. The remaining headroom is in two places. First, the tile is still +``128 x 256``, so pipeline overhead is amortized over a relatively small amount of +compute. Second, the epilogue is a plain per-thread ``store_global`` from +registers, issued by both consumer groups at the same time at the very end. + +In :doc:`the final version `, the tile grows to ``256 x 256`` split across +**four** consumer warp groups, the accumulator switches to native fp16 WGMMA +accumulation to fit the register budget, and the epilogue routes through a shared +memory buffer so results leave via a bulk TMA store. diff --git a/docs/source/tutorials/matmul-hopper/v6.rst b/docs/source/tutorials/matmul-hopper/v6.rst new file mode 100644 index 00000000..0c456b6b --- /dev/null +++ b/docs/source/tutorials/matmul-hopper/v6.rst @@ -0,0 +1,340 @@ +.. _tutorial_hopper_matmul_v6: + +6. Four Consumers, FP16 Accumulation, and a TMA Epilogue +========================================================= + +:doc:`V5 ` keeps the tensor cores busy inside each consumer group, but two +costs remain fixed per output tile: the pipeline prologue and drain, and the +epilogue. With a ``128 x 256`` tile there is only so much compute to amortize +them over. The obvious response --- make the tile bigger --- runs into the +constraint that has shaped every version so far: **the accumulator lives in +registers**. + +This version breaks that deadlock with three changes that only work together: + +1. **Four consumer warp groups** on a ``256 x 256`` tile, so each group owns a + ``64 x 256`` quarter and four independent WGMMA streams are in flight. +2. **Native fp16 WGMMA accumulation**, halving accumulator register cost so the + larger tile fits at all. +3. **A shared-memory TMA epilogue**, where the four groups take turns staging + their quarter through one shared buffer for a bulk TMA store. + +Together with an 8-wide raster group, this is the version that passes cuBLAS. + + +The Full Kernel +--------------- + +.. literalinclude:: ../../../../examples/hopper_matmul/matmul_v6.py + :language: python + :start-at: class Pipeline + :end-at: sa, sb, sc, tma_pipe, epilogue_ready, epilogue_free, 3, k_size + :caption: MatmulWGMMAV6 --- full kernel (including Pipeline class) + + +What Changed from V5 +-------------------- + +.. list-table:: + :header-rows: 1 + :widths: 15 40 40 + + * - + - V5 + - V6 + * - **Output tile** + - 128 x 256 + - **256 x 256** + * - **Consumer groups** + - 2, each owning a 64 x 256 half + - **4**, each owning a 64 x 256 quarter + * - **Warps** + - 9 (1 producer + 2 groups) + - **17** (1 producer + 4 groups) + * - **Accumulator dtype** + - fp32 (cast to fp16 in the epilogue) + - **fp16**, accumulated natively by WGMMA + * - **Epilogue** + - ``store_global`` from registers, per group + - Serialized through one shared buffer, bulk TMA store + * - **Epilogue issuer** + - Each consumer group + - The **producer warp**, after its loads are done + * - **Rasterization** + - ``swizzle_size=4`` + - ``swizzle_size=8`` + * - **Pipeline depth** + - 4 stages + - 3 stages (the larger tile costs more shared memory per stage) + * - **New instructions** + - + - :meth:`~tilus.Script.store_shared`, + :meth:`~tilus.lang.instructions.fence.FenceInstructionGroup.proxy_async`, + :meth:`~tilus.lang.instructions.tma.TmaInstructionGroup.shared_to_global`, + :meth:`~tilus.lang.instructions.tma.TmaInstructionGroup.commit_group`, + :meth:`~tilus.lang.instructions.tma.TmaInstructionGroup.wait_group` + + +The Register Budget Problem +--------------------------- + +.. figure:: figures/v6_tile_partition.svg + :width: 100% + :align: center + + A 256 x 256 output tile split by rows across four consumer warp groups. Each + group loads its own A slab; all four read the same B tile. + +A CUDA thread can hold at most 255 registers. An accumulator of ``m x n`` +distributed over a 128-thread warp group costs ``m * n / 128`` registers per +thread in fp32. For V5's ``64 x 256`` per-group accumulator that is **128 +registers** --- already half the budget, before operands, addresses, and loop +state. + +Scaling to a ``256 x 256`` tile with four groups keeps each group's share at +``64 x 256``, so fp32 would still cost 128 registers per thread. Per-thread that +is legal, but there are now 512 consumer threads, and ``512 x 128 = 65536`` is +the entire register file of an SM --- with nothing left for operands, addresses, +loop state, or the producer warp. V6 instead accumulates in **fp16**: + +.. literalinclude:: ../../../../examples/hopper_matmul/matmul_v6.py + :language: python + :start-at: acc = self.register_tensor( + :end-at: ) + :dedent: 8 + :caption: fp16 accumulator + +WGMMA supports fp16 accumulation natively for fp16 operands +(``wgmma.mma_async...f16.f16.f16``), so this is not a cast --- the tensor cores +accumulate in half precision throughout. Two fp16 values pack into one 32-bit +register, halving the accumulator to **64 registers per thread** and leaving room +for the larger tile and the epilogue. + +.. warning:: + + fp16 accumulation trades precision for capacity, and the trade is real. Over a + K=8192 reduction with unit-variance outputs, the measured absolute error + against cuBLAS has mean 0.0023, p99 0.0117, and a maximum of 0.047 --- which is + why the benchmark checks V6 with ``atol=5e-2`` while earlier versions use + ``1e-2``. Against an fp32 reference, V5 on the same inputs has mean error + 0.00014 and maximum 0.0020, so this is roughly an order of magnitude, not a + rounding detail. It also means the headline comparison below is not + like-for-like: cuBLAS is accumulating in fp32. For inference-style workloads + the trade is typically fine; for training or ill-conditioned inputs, prefer + the fp32 accumulation of :doc:`V5 `. + +The B tile is shared by all four groups, so widening the tile in M costs no extra +B traffic --- the arithmetic intensity of the block improves, which is the whole +point. + + +Serialized Shared-Memory Epilogue +--------------------------------- + +With four groups each holding a ``64 x 256`` fp16 quarter, writing results out +becomes its own problem. Four simultaneous ``store_global`` calls from registers +produce many small, poorly coalesced transactions. Routing through TMA instead +requires the data to be in shared memory --- but a full ``256 x 256`` fp16 +staging buffer would be 128 KB, competing with the pipeline's ring buffer for the +same 228 KB budget. + +V6 allocates **one quarter-sized buffer** and has the groups take turns: + +.. literalinclude:: ../../../../examples/hopper_matmul/matmul_v6.py + :language: python + :start-at: sc = self.shared_tensor(dtype=float16, shape=[block_m_slice, block_n]) + :end-at: epilogue_free = self.mbarrier.alloc([1, 1, 1]) + :dedent: 8 + :caption: Shared-memory and barrier allocation, including the epilogue handoff + +Two barrier arrays sequence the handoff: + +- ``epilogue_ready[i]`` --- consumer *i* has finished writing its quarter into + ``sc``. Arrival count 128: every thread of the group participates in + :meth:`~tilus.Script.store_shared`. +- ``epilogue_free[i]`` --- the buffer has been drained after consumer *i*, so + consumer *i+1* may write. Arrival count 1, and there are only three of them: + the last consumer needs no successor. + +Each consumer therefore waits for its predecessor to clear the buffer before +writing, except consumer 0 which finds it free: + +.. literalinclude:: ../../../../examples/hopper_matmul/matmul_v6.py + :language: python + :start-at: if consumer_idx > 0: + :end-at: self.mbarrier.arrive(epilogue_ready[consumer_idx]) + :dedent: 8 + :caption: Consumer side of the epilogue handoff + +The drain side runs on the **producer warp**, which by this point has finished +all its TMA loads and would otherwise be idle: + +.. literalinclude:: ../../../../examples/hopper_matmul/matmul_v6.py + :language: python + :start-at: def store_epilogue + :end-at: self.mbarrier.arrive(epilogue_free[consumer_idx]) + :dedent: 4 + :caption: Producer-side epilogue: shared memory to global via TMA + +The sequence per quarter is: wait for the quarter to be staged, fence, issue the +bulk TMA store, wait for it, then release the buffer. + +.. important:: + + The :meth:`fence.proxy_async(space="shared") ` + is not optional. :meth:`~tilus.Script.store_shared` writes through the + **generic proxy** (the ordinary load/store path), while + :meth:`tma.shared_to_global() ` + reads through the **async proxy** used by the TMA engine. Without a + ``fence.proxy.async.shared::cta`` between them, the TMA engine may read stale + data. + +``tma.wait_group(n=0, read=True)`` waits only for the TMA engine to finish +**reading** shared memory --- enough to hand the buffer to the next consumer --- +rather than for the global writes to become visible, which nothing downstream +needs. + +.. note:: + + Global-to-shared TMA reports completion through **mbarrier tx-count**; + shared-to-global TMA uses **commit_group + wait_group** instead. See + `cp.async.bulk `__ + in the PTX documentation. + +Note also the ordering of the four TMA loads in the producer's main loop: slabs +0, 2, 3, then 1. Slab 1's load is issued last so that consumer 1 --- the first +group that has to *wait* for the shared buffer --- is the least likely to be +blocked on its input as well. + + +Walkthrough +----------- + +Producer Warp +~~~~~~~~~~~~~ + +.. literalinclude:: ../../../../examples/hopper_matmul/matmul_v6.py + :language: python + :start-at: with self.thread_group(thread_begin=512, num_threads=32): + :end-before: with self.thread_group(thread_begin=0, num_threads=128): + :dedent: 8 + :caption: Producer warp: K-loop, drain, then the epilogue + +The producer has three phases. It fills the pipeline over the K loop with five +TMA loads per stage (four A slabs plus B), drains the outstanding empty-signals, +and then serves all four epilogue quarters in order. + + +Consumer Warp Groups +~~~~~~~~~~~~~~~~~~~~ + +.. literalinclude:: ../../../../examples/hopper_matmul/matmul_v6.py + :language: python + :start-at: def consume_tile + :end-at: self.mbarrier.arrive(epilogue_ready[consumer_idx]) + :dedent: 4 + :caption: consume_tile --- shared by all four consumer groups + +All four groups run the same ``consume_tile`` method, parameterized by +``consumer_idx``. The K-loop is V5's overlapped structure verbatim --- prologue +MMA, steady-state loop with ``wait_group(1)`` and a lagging stage release, final +``wait_group(0)`` --- followed by the epilogue handoff. The four call sites differ +only in their thread range and index: + +.. literalinclude:: ../../../../examples/hopper_matmul/matmul_v6.py + :language: python + :start-at: with self.thread_group(thread_begin=0, num_threads=128): + :end-at: sa, sb, sc, tma_pipe, epilogue_ready, epilogue_free, 3, k_size + :dedent: 8 + :caption: Four consumer warp groups + +``consumer_arrive_count=4`` on the pipeline reflects that all four groups must +release a stage before the producer may refill it. + + +Performance +----------- + +V6 reaches **800 TFLOPS** (1.37 ms) against cuBLAS at 742 TFLOPS (1.48 ms), and is +18% ahead of V5. Nsight Compute puts tensor pipe utilization at 93.2%, essentially +level with cuBLAS's 93.8%, and DRAM throughput at 20%, the lowest of any version: +the ``256 x 256`` tile with a shared B tile has made the kernel almost entirely +compute bound. + +Comparing against the cuBLAS timing taken in the same process, V6 wins every run, +by a **median of 7.8%**: + +.. list-table:: + :header-rows: 1 + :widths: 12 20 20 24 + + * - Run + - V6 + - cuBLAS + - V6 faster by + * - 1 + - 1.3886 ms + - 1.5014 ms + - 8.1% + * - 2 + - 1.3742 ms + - 1.4630 ms + - 6.5% + * - 3 + - 1.3742 ms + - 1.4809 ms + - 7.8% + +Quote the range, not a single run. Both kernels vary by 1--2% run to run, so the +ratio moves with whichever samples you happen to draw. + +.. warning:: + + The margin is also sensitive to how long the timed window is, because an + 8192\ :sup:`3` fp16 GEMM saturates the board's power budget. At the 30 timed + iterations this tutorial uses, V6 leads by 6--8%; at 100 iterations the H100 + throttles partway through the window and the same three measurements read + 99%, 99%, and 110% of cuBLAS. The benchmark script takes a 3-second cooldown + before every measurement, cuBLAS included, and exposes ``--repeat`` so this + can be checked directly. Run it on an idle GPU. + +The advantage survives the change of measurement method: under Nsight Compute's +replay profiling V6 comes in at 1.53 ms against cuBLAS's 1.56 ms --- still ahead, +but by 2% rather than 8%. Replay pins the SM clock near 1.41 GHz against a +1.98 GHz maximum, and at reduced clock the memory system is comparatively +over-provisioned, which is where most of V6's wall-clock advantage comes from. +The complete source is at :github:`examples/hopper_matmul/matmul_v6.py`. + +.. plot:: tutorials/matmul-hopper/plots/plot_v6.py + + Hopper matmul performance on H100 SXM (M=N=K=8192, fp16). Latency is + CUDA-event timed, median of three fresh processes. Peak is the published + dense FP16 tensor core throughput of the H100 SXM. + + +Summary +------- + +Starting from a minimal TMA-fed kernel that pushed every operand through the +register file (V0), we replaced the MMA path with Hopper's asynchronous +shared-memory WGMMA (V1), overlapped loading and computing with a multi-stage +ring buffer (V2), separated the two into dedicated producer and consumer warps +(V3), doubled the independent MMA streams and factored the bookkeeping into a +``Pipeline`` class (V4), kept a WGMMA group permanently in flight (V5), and +finally widened the tile to ``256 x 256`` across four consumer groups with fp16 +accumulation and a bulk TMA epilogue (V6). + +Two themes run through the whole series. The first is that Hopper's engines --- +TMA, the tensor cores, and the SM's own instruction issue --- are independent, and +performance comes from arranging for all of them to have work queued at all +times. The second is that the register file is the binding constraint on how +large a tile a Hopper kernel can hold, which is why the final step needed both +more warp groups and a narrower accumulator. + +.. caution:: + + The result reported here is specific to this shape (M=N=K=8192), dtype + (fp16), accumulation precision (fp16, where cuBLAS uses fp32), GPU (H100 SXM), + and benchmark methodology. It is not a claim that this kernel beats cuBLAS + across GEMM shapes; the autotune spaces checked into the examples are pinned + to a single configuration each, tuned for exactly this workload. diff --git a/examples/hopper_matmul/benchmark.py b/examples/hopper_matmul/benchmark.py index ec3e9f5b..deb24097 100644 --- a/examples/hopper_matmul/benchmark.py +++ b/examples/hopper_matmul/benchmark.py @@ -119,58 +119,101 @@ def parse_ncu_report(report_path: str) -> list[tuple[str, dict]]: return [(k, per_kernel[k]) for k in kernel_order] -def benchmark_all(versions: list[str], m_size: int, n_size: int, k_size: int): - """Benchmark all versions using benchmark_func (event-loop timing).""" - import math - +# Timing protocol, following examples/blackwell_matmul/benchmark.py: CUDA-event +# timing via benchmark_func (median of `REPEAT` iterations, L2 flushed before +# each), and a `COOLDOWN_S` pause before every measurement -- including cuBLAS, +# which is timed last like any other entry -- so nothing is measured at a +# thermal state the rest did not see. +# +# REPEAT deliberately differs from the Blackwell script's 100. An 8192^3 fp16 +# GEMM takes ~1.5 ms on an H100, so 100 back-to-back iterations is ~150 ms of +# power-capped tensor core work and the board throttles partway through the +# timed window. Measured over three fresh processes at this shape, repeat=30 +# gives v6/cuBLAS = 108.1 / 106.5 / 107.8 %, while repeat=100 gives +# 99.0 / 98.9 / 110.4 % -- same kernels, but the median lands on whichever side +# of the throttle transition the run happened to sit. Pass --repeat to compare. +WARMUP = 5 +REPEAT = 30 +COOLDOWN_S = 3 + +# v6 accumulates in fp16 inside the WGMMA (see matmul_v6.py); every other version +# and cuBLAS accumulate in fp32. Measured at 8192^3 with unit-variance outputs, +# that raises the mean absolute error from 1.4e-4 to 2.3e-3, so v6 needs its own +# tolerance. This is a real precision difference, not a benchmarking artifact -- +# see the warning in docs/source/tutorials/matmul-hopper/v6.rst. +FP16_ACCUMULATE_VERSIONS = {"v6"} + + +def benchmark_all( + versions: list[str], + m_size: int, + n_size: int, + k_size: int, + repeat: int = REPEAT, +): + """Benchmark all versions and cuBLAS using benchmark_func (CUDA-event timing).""" import pandas import torch from tilus.utils import benchmark_func - headers = ["version", "latency (ms)", "tflops", "% of cublas"] + headers = ["version", "accumulate", "latency (ms)", "tflops", "% of cublas"] rows = [] - # Scale only one operand. Scaling both made even grossly permuted outputs - # small enough to pass the old absolute tolerance. - a = torch.randn(m_size, k_size, dtype=torch.float16, device="cuda") / math.sqrt( - k_size - ) - b = torch.randn(n_size, k_size, dtype=torch.float16, device="cuda") + # Scale both operands by k**0.25 so that C has unit variance: each output + # element sums k products of two N(0, k**-0.5) values. That keeps atol and + # rtol comparable in the correctness check below -- with unscaled operands C + # has standard deviation sqrt(k), and with both operands scaled by 1/sqrt(k) + # it has 1/sqrt(k), which makes an absolute tolerance meaningless in either + # direction. + scale = k_size**0.25 + a = torch.randn(m_size, k_size, dtype=torch.float16, device="cuda") / scale + b = torch.randn(n_size, k_size, dtype=torch.float16, device="cuda") / scale c_ref = torch.empty(m_size, n_size, dtype=torch.float16, device="cuda") c_tilus = torch.empty(m_size, n_size, dtype=torch.float16, device="cuda") - cublas_lat = benchmark_func( - lambda: torch.matmul(a, b.T, out=c_ref), warmup=5, repeat=30 - ) - def tf(ms): return 2 * m_size * n_size * k_size / ms * 1e-9 - cublas_tf = tf(cublas_lat) - for name in versions: + acc_dtype = "fp16" if name in FP16_ACCUMULATE_VERSIONS else "fp32" try: matmul = _load_version(name)() matmul(m_size, n_size, k_size, a, b, c_tilus) torch.cuda.synchronize() - atol = 5e-2 if name == "v6" else 1e-2 - torch.testing.assert_close(c_ref, c_tilus, atol=atol, rtol=1e-2) + torch.matmul(a, b.T, out=c_ref) + torch.cuda.synchronize() + atol = 5e-2 if acc_dtype == "fp16" else 1e-2 + torch.testing.assert_close(c_tilus, c_ref, atol=atol, rtol=1e-2) + time.sleep(COOLDOWN_S) t = benchmark_func( lambda: matmul(m_size, n_size, k_size, a, b, c_tilus), - warmup=5, - repeat=30, + warmup=WARMUP, + repeat=repeat, ) - rows.append([f"tilus_{name}", t, tf(t), tf(t) / cublas_tf * 100.0]) - time.sleep(1) + rows.append([f"tilus_{name}", acc_dtype, t, tf(t), float("nan")]) except Exception as e: - print(f" tilus_{name} ERROR: {e}") - rows.append([f"tilus_{name}", float("nan"), float("nan"), float("nan")]) + print(f" tilus_{name} ERROR: {type(e).__name__}: {e}") + rows.append( + [f"tilus_{name}", acc_dtype, float("nan"), float("nan"), float("nan")] + ) + + # cuBLAS last, under the same cooldown and iteration count as every version. + time.sleep(COOLDOWN_S) + cublas_lat = benchmark_func( + lambda: torch.matmul(a, b.T, out=c_ref), warmup=WARMUP, repeat=repeat + ) + cublas_tf = tf(cublas_lat) + rows.append(["cublas", "fp32", cublas_lat, cublas_tf, 100.0]) - rows.append(["cublas", cublas_lat, cublas_tf, 100.0]) + for row in rows[:-1]: + row[4] = row[3] / cublas_tf * 100.0 df = pandas.DataFrame(rows, columns=headers) - print(f"\nBenchmark results (m={m_size}, n={n_size}, k={k_size}):") + print( + f"\nBenchmark results (m={m_size}, n={n_size}, k={k_size}, " + f"warmup={WARMUP}, repeat={repeat}):" + ) print(df.to_string(index=False)) @@ -255,13 +298,19 @@ def main(): metavar=("M", "N", "K"), help="Workload size M N K (default: 8192 8192 8192)", ) + parser.add_argument( + "--repeat", + type=int, + default=REPEAT, + help=f"Timed iterations per measurement (default: {REPEAT})", + ) args = parser.parse_args() m_size, n_size, k_size = args.size if args.ncu: ncu_profile_all(args.versions, m_size, n_size, k_size) else: - benchmark_all(args.versions, m_size, n_size, k_size) + benchmark_all(args.versions, m_size, n_size, k_size, repeat=args.repeat) if __name__ == "__main__": diff --git a/examples/hopper_matmul/matmul_v6.py b/examples/hopper_matmul/matmul_v6.py index 9ac775df..b5de21f3 100644 --- a/examples/hopper_matmul/matmul_v6.py +++ b/examples/hopper_matmul/matmul_v6.py @@ -312,7 +312,10 @@ def main(): matmul(m, n, k, a, b, c_actual) torch.cuda.synchronize() - torch.testing.assert_close(c_expect, c_actual, atol=5e-2, rtol=1e-2) + # Looser atol than v0-v5: this kernel accumulates in fp16 inside the + # WGMMA rather than fp32, which costs roughly 5x in absolute error over + # a K=8192 reduction. + torch.testing.assert_close(c_actual, c_expect, atol=5e-2, rtol=1e-2) for name, func in [ ("torch", lambda: torch.matmul(a, b.T, out=c_expect)),