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)),