6. Register Reallocation, Multicast, and a Persistent Kernel

V5 keeps the tensor cores busy inside each consumer group, but it still moves more data per FLOP than it needs to, and every CTA pays for its own pipeline fill and epilogue. The obvious response — make the tile bigger — runs into the constraint that has shaped every version so far: the accumulator lives in registers, and with fp32 accumulation V5’s 128 x 256 tile already uses half the register file.

This version keeps V0–V5’s fp32 accumulation, so it stays numerically like-for-like with cuBLAS, and makes four changes:

1. Register reallocation with setmaxnreg(): the producer warp group gives up most of its registers so the two consumer groups can hold a 64 x 320 fp32 accumulator each, for a 128 x 320 tile.

2. A 2-CTA cluster with TMA multicast: two CTAs that need the same B tile each load half of it and multicast it to both, cutting L2 traffic.

3. A persistent tile loop: one CTA per SM walks over output tiles, so the producer prefetches the next tile while the consumers run the epilogue.

4. A shared-memory TMA epilogue, with each consumer group staging its results through its own slot for coalesced bulk stores.

Together they take the kernel past cuBLAS on this shape.

The Full Kernel

MatmulWGMMAV6 — full kernel (including Pipeline class)

class Pipeline(tilus.Class):
    def __init__(
        self,
        num_stages: int,
        producer_arrive_count: int = 1,
        consumer_arrive_count: int = 1,
    ):
        self.num_stages: int = num_stages
        self.empty_barriers = self.mbarrier.alloc(
            [consumer_arrive_count for _ in range(num_stages)]
        )
        self.full_barriers = self.mbarrier.alloc(
            [producer_arrive_count for _ in range(num_stages)]
        )
        self.producer_stage: int32 = 0
        self.consumer_stage: int32 = 0
        self.producer_phase: uint32 = self.mbarrier.producer_initial_phase
        self.consumer_phase: uint32 = self.mbarrier.consumer_initial_phase

    def producer_acquire(self):
        self.mbarrier.wait(
            barrier=self.empty_barriers[self.producer_stage],
            phase=self.producer_phase,
            sem="relaxed",
            scope="cta",
        )

    def producer_barrier(self) -> RegisterTensor:
        return self.full_barriers[self.producer_stage]

    def producer_advance(self):
        self.producer_stage = (self.producer_stage + 1) % self.num_stages
        self.producer_phase = self.producer_phase ^ (self.producer_stage == 0)

    def consumer_acquire(self):
        self.mbarrier.wait(
            barrier=self.full_barriers[self.consumer_stage],
            phase=self.consumer_phase,
            sem="relaxed",
            scope="cta",
        )

    def consumer_barrier(self) -> RegisterTensor:
        return self.empty_barriers[self.consumer_stage]

    def consumer_advance(self):
        self.consumer_stage = (self.consumer_stage + 1) % self.num_stages
        self.consumer_phase = self.consumer_phase ^ (self.consumer_stage == 0)

    def prev_consumer_barrier(self) -> RegisterTensor:
        prev_stage = (self.consumer_stage + (self.num_stages - 1)) % self.num_stages
        return self.empty_barriers[prev_stage]


@tilus.autotune("num_stages", [3])
@tilus.autotune("block_m, block_n", [[128, 320]])
@tilus.autotune("block_k", [64])
@tilus.autotune("swizzle_size", [8])
class MatmulWGMMAV6(tilus.Script):
    def __init__(self, num_stages, block_m, block_n, block_k, swizzle_size):
        super().__init__()
        self.num_stages = num_stages
        self.block_m = block_m
        self.block_n = block_n
        self.block_k = block_k
        self.swizzle_size = swizzle_size
        # Persistent kernel: launch one CTA per SM.
        self.num_sms = torch.cuda.get_device_properties(
            torch.cuda.current_device()
        ).multi_processor_count

    def compute_block_coord(
        self, linear_idx: int32, num_m_blocks: int32, num_n_blocks: int
    ):
        swizzle_size = self.swizzle_size
        tiles_per_group = num_m_blocks * swizzle_size
        group_idx, in_group_idx = self.fast_divmod(linear_idx, tiles_per_group)
        first_n = group_idx * swizzle_size
        m_block: int32 = 0
        n_block: int32 = 0
        remainder = num_n_blocks - num_n_blocks // swizzle_size * swizzle_size
        last_group_width = remainder if remainder > 0 else swizzle_size
        if first_n + swizzle_size <= num_n_blocks:
            m_block, r = self.fast_divmod(in_group_idx, swizzle_size)
            n_block = first_n + r
        else:
            m_block, r = self.fast_divmod(in_group_idx, last_group_width)
            n_block = first_n + r
        return m_block, n_block

    def tile_offsets(self, tile: int32, num_m_pairs: int32, num_n_blocks: int):
        # A tile index names a pair of vertically adjacent CTA tiles; each CTA of
        # the cluster takes the one matching its rank.
        m_pair, n_block = self.compute_block_coord(tile, num_m_pairs, num_n_blocks)
        offset_m: int32 = (m_pair * 2 + self.cluster.blockRank) * self.block_m
        offset_n: int32 = n_block * self.block_n
        return offset_m, offset_n

    def release_stage(self, barrier: RegisterTensor):
        # Both CTAs multicast half of B into this stage, so both producers must
        # see it released before either refills it.
        for rank in range(2):
            self.mbarrier.arrive_and_expect_tx_remote(
                barrier, transaction_bytes=0, target_rank=rank, sem="relaxed"
            )

    def mma_stage(
        self,
        sa: SharedTensor,
        sb: SharedTensor,
        tma_pipe: Pipeline,
        acc0: RegisterTensor,
        acc1: RegisterTensor,
        consumer_idx: int,
    ):
        # Two m64n160k16 WGMMA streams cover the 320 columns; one commit group.
        self.wgmma.fence()
        self.wgmma.mma(
            sa[tma_pipe.consumer_stage, consumer_idx],
            sb[tma_pipe.consumer_stage, 0].transpose(),
            acc0,
        )
        self.wgmma.mma(
            sa[tma_pipe.consumer_stage, consumer_idx],
            sb[tma_pipe.consumer_stage, 1].transpose(),
            acc1,
        )
        self.wgmma.commit_group()

    def store_half(
        self,
        sc: SharedTensor,
        gc: GlobalTensor,
        acc: RegisterTensor,
        offset_m: int32,
        offset_n: int32,
    ):
        self.store_shared(sc, self.cast(acc, dtype=float16))
        self.fence.proxy_async(space="shared")
        self.sync()
        with self.single_thread():
            self.tma.shared_to_global(sc, gc, offsets=[offset_m, offset_n])
            self.tma.commit_group()

    def wait_store_read(self):
        # The staging slot may be rewritten once the TMA engine has read it.
        with self.single_thread():
            self.tma.wait_group(n=0, read=True)
        self.sync()

    def consume(
        self,
        sa: SharedTensor,
        sb: SharedTensor,
        sc: SharedTensor,
        gc: GlobalTensor,
        tma_pipe: Pipeline,
        consumer_idx: int,
        k_size: int,
        num_m_pairs: int32,
        num_n_blocks: int,
    ):
        self.setmaxnreg(232, "inc")
        half_n = self.block_n // 2
        num_tiles = num_m_pairs * num_n_blocks
        for tile in self.range(self.blockIdx.x // 2, num_tiles, self.gridDim.x // 2):
            offset_m, offset_n = self.tile_offsets(tile, num_m_pairs, num_n_blocks)
            acc0 = self.register_tensor(dtype=float32, shape=[64, half_n], init=0.0)
            acc1 = self.register_tensor(dtype=float32, shape=[64, half_n], init=0.0)

            tma_pipe.consumer_acquire()
            self.mma_stage(sa, sb, tma_pipe, acc0, acc1, consumer_idx)
            tma_pipe.consumer_advance()
            for offset_k in self.range(
                self.block_k, k_size, self.block_k, unroll=self.num_stages
            ):
                tma_pipe.consumer_acquire()
                self.mma_stage(sa, sb, tma_pipe, acc0, acc1, consumer_idx)
                self.wgmma.wait_group(1)
                with self.single_thread():
                    self.release_stage(tma_pipe.prev_consumer_barrier())
                tma_pipe.consumer_advance()
            self.wgmma.wait_group(0)
            with self.single_thread():
                self.release_stage(tma_pipe.prev_consumer_barrier())

            # Epilogue: each group stages its two 64x160 halves in turn through
            # its own slot of sc, so the two groups never wait on each other.
            row = offset_m + consumer_idx * 64
            self.wait_store_read()
            self.store_half(sc[consumer_idx], gc, acc0, row, offset_n)
            self.wait_store_read()
            self.store_half(sc[consumer_idx], gc, acc1, row, offset_n + half_n)

        with self.single_thread():
            self.tma.wait_group(n=0)

    def produce(
        self,
        ga: GlobalTensor,
        gb: GlobalTensor,
        sa: SharedTensor,
        sb: SharedTensor,
        tma_pipe: Pipeline,
        k_size: int,
        num_m_pairs: int32,
        num_n_blocks: int,
    ):
        half_n = self.block_n // 2
        cta_rank = self.cluster.blockRank
        num_tiles = num_m_pairs * num_n_blocks
        for tile in self.range(self.blockIdx.x // 2, num_tiles, self.gridDim.x // 2):
            offset_m, offset_n = self.tile_offsets(tile, num_m_pairs, num_n_blocks)
            for offset_k in self.range(0, k_size, self.block_k, unroll=self.num_stages):
                tma_pipe.producer_acquire()
                with self.single_thread():
                    # Expect the whole stage: this CTA's A plus both halves of
                    # B, one loaded here and one multicast by the peer.
                    self.mbarrier.arrive_and_expect_tx(
                        tma_pipe.producer_barrier(),
                        transaction_bytes=sa[tma_pipe.producer_stage].nbytes
                        + sb[tma_pipe.producer_stage].nbytes,
                    )
                self.tma.global_to_shared(
                    src=ga,
                    dst=sa[tma_pipe.producer_stage, 0],
                    offsets=[offset_m, offset_k],
                    mbarrier=tma_pipe.producer_barrier(),
                )
                self.tma.global_to_shared(
                    src=ga,
                    dst=sa[tma_pipe.producer_stage, 1],
                    offsets=[offset_m + 64, offset_k],
                    mbarrier=tma_pipe.producer_barrier(),
                )
                self.tma.global_to_shared(
                    src=gb,
                    dst=sb[tma_pipe.producer_stage, cta_rank],
                    offsets=[offset_n + cta_rank * half_n, offset_k],
                    mbarrier=tma_pipe.producer_barrier(),
                    multicast_mask=0b11,
                )
                tma_pipe.producer_advance()

        for _ in self.range(self.num_stages):
            tma_pipe.producer_acquire()
            tma_pipe.producer_advance()

    def __call__(
        self,
        m_size: int32,
        n_size: int,
        k_size: int,
        a_ptr: ~float16,
        b_ptr: ~float16,
        c_ptr: ~float16,
    ):
        num_stages = self.num_stages
        block_m, block_n, block_k = self.block_m, self.block_n, self.block_k
        half_n = block_n // 2

        num_m_pairs = cdiv(cdiv(m_size, block_m), 2)
        num_n_blocks = cdiv(n_size, block_n)
        self.attrs.blocks = self.num_sms // 2 * 2
        self.attrs.cluster_blocks = 2
        self.attrs.warps = 12  # 2 consumer WGs + 1 producer WG

        ga = self.global_view(a_ptr, dtype=float16, shape=[m_size, k_size])
        gb = self.global_view(b_ptr, dtype=float16, shape=[n_size, k_size])
        gc = self.global_view(c_ptr, dtype=float16, shape=[m_size, n_size])
        # A as one 64-row slab per consumer; B as two 160-row halves (TMA boxes
        # are at most 256 rows); sc as one 64x160 staging slot per consumer.
        sa = self.shared_tensor(dtype=float16, shape=[num_stages, 2, 64, block_k])
        sb = self.shared_tensor(dtype=float16, shape=[num_stages, 2, half_n, block_k])
        sc = self.shared_tensor(dtype=float16, shape=[2, 64, half_n])

        # Each stage is released by both consumers of both CTAs.
        tma_pipe = Pipeline(num_stages, producer_arrive_count=1, consumer_arrive_count=4)
        # The peer must not multicast into this CTA before its barriers exist.
        self.cluster.sync()

        # The producer branch is placed after the consumers: with it first,
        # ptxas serializes the WGMMAs around the setmaxnreg in the divergent path.
        with self.thread_group(thread_begin=0, num_threads=128):
            self.consume(sa, sb, sc, gc, tma_pipe, 0, k_size, num_m_pairs, num_n_blocks)

        with self.thread_group(thread_begin=128, num_threads=128):
            self.consume(sa, sb, sc, gc, tma_pipe, 1, k_size, num_m_pairs, num_n_blocks)

        with self.thread_group(thread_begin=256, num_threads=128):
            self.setmaxnreg(40, "dec")
            with self.thread_group(thread_begin=0, num_threads=32):
                self.produce(ga, gb, sa, sb, tma_pipe, k_size, num_m_pairs, num_n_blocks)

        # Keep both CTAs resident until neither will touch the other's memory.
        self.cluster.sync()

What Changed from V5

	V5	V6
Output tile	128 x 256	128 x 320
Accumulator dtype	fp32	fp32 (unchanged)
Warp groups	2 consumers + 1 producer warp	2 consumers + 1 producer warp group
Registers per thread	Uniform	232 per consumer thread, 40 per producer thread (setmaxnreg)
WGMMA shape	m64n256k16	m64n160k16, two per stage
Grid	One CTA per output tile	Persistent: one CTA per SM, in 2-CTA clusters
B tile loads	Each CTA loads its own	Each CTA loads half and multicasts it to its cluster peer
Epilogue	store_global from registers	Staged through shared memory, bulk TMA store
Pipeline depth	4 stages	3 stages (the larger tile costs more shared memory per stage)
New instructions		setmaxnreg(), store_shared(), arrive_and_expect_tx_remote(), shared_to_global(), commit_group(), wait_group()

The Register Budget Problem

An accumulator of m x n distributed over a 128-thread warp group costs m * n / 128 registers per thread in fp32. A wider tile is the most direct way to move less data per FLOP — 128 x 320 reads 7% fewer operand bytes per FLOP than 128 x 256 — but each consumer would then hold a 64 x 320 accumulator: 160 registers per thread, before operands, addresses, and loop state.

The register file is 65,536 registers per SM, and by default every thread of the CTA gets the same share. With two consumer warp groups and a producer warp group (384 threads) that share is 168 registers, which the consumers cannot fit into: compiled at 168, this kernel spills 3.2 KB per thread and runs five times slower. But the producer only issues TMA loads and needs a few dozen registers.

setmaxnreg() lets warp groups trade registers at run time (PTX setmaxnreg). The producer releases down to 40 and the consumers claim 232, which is exactly the same total:

setmaxnreg moves registers from the producer warp group to the consumers. The total is unchanged; only the split is.

The producer warp group releases registers; one of its warps issues TMA

with self.thread_group(thread_begin=256, num_threads=128):
    self.setmaxnreg(40, "dec")
    with self.thread_group(thread_begin=0, num_threads=32):
        self.produce(ga, gb, sa, sb, tma_pipe, k_size, num_m_pairs, num_n_blocks)

Each consumer starts with self.setmaxnreg(232, "inc"), which blocks until the producer’s release has made enough registers available. The counts must add up: if the consumers ask for more than the producer releases, the inc never completes and the kernel hangs.

Important

The order of the thread groups in the source matters to ptxas. With the producer branch first, ptxas serializes every WGMMA in the kernel (“wgmma serialized due to program dependence on compiler-inserted WG.AR in divergent path”), which cost 5% here. Placing the consumer branches first avoids it. The branches run concurrently either way, so the order changes nothing else.

The 320 columns are split into two m64n160k16 WGMMAs per stage, because a WGMMA is at most 256 wide and B is loaded as two 160-row halves (a TMA box is also at most 256 rows). Both go into one commit group, so V5’s wait_group(1) overlap is unchanged:

One pipeline stage: two 160-wide WGMMA streams

def mma_stage(
    self,
    sa: SharedTensor,
    sb: SharedTensor,
    tma_pipe: Pipeline,
    acc0: RegisterTensor,
    acc1: RegisterTensor,
    consumer_idx: int,
):
    # Two m64n160k16 WGMMA streams cover the 320 columns; one commit group.
    self.wgmma.fence()
    self.wgmma.mma(
        sa[tma_pipe.consumer_stage, consumer_idx],
        sb[tma_pipe.consumer_stage, 0].transpose(),
        acc0,
    )
    self.wgmma.mma(
        sa[tma_pipe.consumer_stage, consumer_idx],
        sb[tma_pipe.consumer_stage, 1].transpose(),
        acc1,
    )
    self.wgmma.commit_group()

Cluster Multicast

Adjacent tiles in M need the same B tile. V6 launches CTAs in clusters of two and gives the two CTAs of a cluster vertically adjacent tiles. Each loads its own A, but only one half of B, multicast to both CTAs (multicast_mask=0b11):

Producer: A for this CTA, half of B for both

def produce(
    self,
    ga: GlobalTensor,
    gb: GlobalTensor,
    sa: SharedTensor,
    sb: SharedTensor,
    tma_pipe: Pipeline,
    k_size: int,
    num_m_pairs: int32,
    num_n_blocks: int,
):
    half_n = self.block_n // 2
    cta_rank = self.cluster.blockRank
    num_tiles = num_m_pairs * num_n_blocks
    for tile in self.range(self.blockIdx.x // 2, num_tiles, self.gridDim.x // 2):
        offset_m, offset_n = self.tile_offsets(tile, num_m_pairs, num_n_blocks)
        for offset_k in self.range(0, k_size, self.block_k, unroll=self.num_stages):
            tma_pipe.producer_acquire()
            with self.single_thread():
                # Expect the whole stage: this CTA's A plus both halves of
                # B, one loaded here and one multicast by the peer.
                self.mbarrier.arrive_and_expect_tx(
                    tma_pipe.producer_barrier(),
                    transaction_bytes=sa[tma_pipe.producer_stage].nbytes
                    + sb[tma_pipe.producer_stage].nbytes,
                )
            self.tma.global_to_shared(
                src=ga,
                dst=sa[tma_pipe.producer_stage, 0],
                offsets=[offset_m, offset_k],
                mbarrier=tma_pipe.producer_barrier(),
            )
            self.tma.global_to_shared(
                src=ga,
                dst=sa[tma_pipe.producer_stage, 1],
                offsets=[offset_m + 64, offset_k],
                mbarrier=tma_pipe.producer_barrier(),
            )
            self.tma.global_to_shared(
                src=gb,
                dst=sb[tma_pipe.producer_stage, cta_rank],
                offsets=[offset_n + cta_rank * half_n, offset_k],
                mbarrier=tma_pipe.producer_barrier(),
                multicast_mask=0b11,
            )
            tma_pipe.producer_advance()

    for _ in self.range(self.num_stages):
        tma_pipe.producer_acquire()
        tma_pipe.producer_advance()

That cuts each CTA’s operand traffic per stage from 448 rows to 288. A multicast load signals the barrier at the same shared memory offset in every destination CTA, so each CTA’s arrive_and_expect_tx still expects the whole stage.

The cost is that the pipeline now spans two CTAs. A producer may refill a stage only once both CTAs’ consumers have finished with it, because its multicast writes into its peer’s shared memory too. Each consumer therefore releases a stage to both CTAs, and the empty barriers count four arrivals:

Releasing a stage to both producers of the cluster

def release_stage(self, barrier: RegisterTensor):
    # Both CTAs multicast half of B into this stage, so both producers must
    # see it released before either refills it.
    for rank in range(2):
        self.mbarrier.arrive_and_expect_tx_remote(
            barrier, transaction_bytes=0, target_rank=rank, sem="relaxed"
        )

Note

sem="relaxed" matters. The default release semantics at cluster scope made each arrive wait for prior memory operations to become visible cluster wide, which made the kernel twice as slow. The consumer only reads the stage, and its WGMMA reads have completed by the time it arrives, so there is nothing to publish.

self.cluster.sync() brackets the kernel: before the first multicast, so the peer’s barriers exist, and at the end, so neither CTA exits while the other may still write into its shared memory.

Persistent Tile Loop

V0–V5 launch one CTA per output tile. Each CTA fills the pipeline, computes, runs its epilogue, and exits, and the tensor cores idle through the fill. V6 launches one CTA per SM and loops over tiles instead. Both CTAs of a cluster walk the same sequence of tile pairs, which keeps their multicasts in lockstep:

Mapping a tile-pair index to this CTA’s tile

def tile_offsets(self, tile: int32, num_m_pairs: int32, num_n_blocks: int):
    # A tile index names a pair of vertically adjacent CTA tiles; each CTA of
    # the cluster takes the one matching its rank.
    m_pair, n_block = self.compute_block_coord(tile, num_m_pairs, num_n_blocks)
    offset_m: int32 = (m_pair * 2 + self.cluster.blockRank) * self.block_m
    offset_n: int32 = n_block * self.block_n
    return offset_m, offset_n

The Pipeline stage and phase simply carry over from one tile to the next. The producer runs ahead into the next tile’s loads as soon as stages free up, so by the time the consumers finish an epilogue, the next tile’s first stages are already in shared memory.

Shared-Memory Epilogue

Each consumer writes its 64 x 320 result through shared memory and a bulk TMA store rather than per-thread store_global. A full staging buffer for both groups would be 80 KB on top of the 168 KB ring buffer, past the 228 KB budget, so each group gets one 64 x 160 slot and stages its two halves in turn:

Consumer epilogue

# Epilogue: each group stages its two 64x160 halves in turn through
# its own slot of sc, so the two groups never wait on each other.
row = offset_m + consumer_idx * 64
self.wait_store_read()
self.store_half(sc[consumer_idx], gc, acc0, row, offset_n)
self.wait_store_read()
self.store_half(sc[consumer_idx], gc, acc1, row, offset_n + half_n)

Staging one half and issuing its TMA store

def store_half(
    self,
    sc: SharedTensor,
    gc: GlobalTensor,
    acc: RegisterTensor,
    offset_m: int32,
    offset_n: int32,
):
    self.store_shared(sc, self.cast(acc, dtype=float16))
    self.fence.proxy_async(space="shared")
    self.sync()
    with self.single_thread():
        self.tma.shared_to_global(sc, gc, offsets=[offset_m, offset_n])
        self.tma.commit_group()

def wait_store_read(self):
    # The staging slot may be rewritten once the TMA engine has read it.
    with self.single_thread():
        self.tma.wait_group(n=0, read=True)
    self.sync()

Because each group has its own slot, the two groups never wait on each other. The only wait is tma.wait_group(n=0, read=True) before a slot is rewritten, and it waits only for the TMA engine to finish reading shared memory, not for the global writes to land.

Important

The fence.proxy_async(space="shared") is not optional. store_shared() writes through the generic proxy, while 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.

Note

An earlier version of this tutorial took a different route to a larger tile: it accumulated in fp16 (wgmma.mma_async...f16.f16.f16), which halves the accumulator and fits a 256 x 256 tile without register reallocation. That kernel ran at 1.37 ms, but cuBLAS accumulates in fp32, so the comparison was not like-for-like, and the precision cost was real: over a K=8192 reduction with unit-variance outputs, the mean absolute error rose from 0.00014 to 0.0023, and the maximum reached 0.047.

Performance

V6 reaches 756 TFLOPS (1.455 ms) against cuBLAS at 752 TFLOPS (1.463 ms) — 101% of cuBLAS, and 10% ahead of V5. Against the cuBLAS timing taken in the same process, V6 wins every run:

Run	V6	cuBLAS	V6 / cuBLAS
1	1.4548 ms	1.4697 ms	101.0%
2	1.4602 ms	1.4632 ms	100.2%
3	1.4482 ms	1.4616 ms	100.9%

The margin is small, and both kernels vary by 1–2% run to run, so quote the range rather than a single run. That the two are this close is not a coincidence: cuBLAS’s kernel for this shape, nvjet_sm90_hsh_320x128_64x3_1x2, is a 320 x 128 tile with block_k 64, 3 stages, and a 1 x 2 cluster, launched at 168 registers per thread — the same design. Under Nsight Compute, V6’s tensor pipe is active 97.9% of active cycles against cuBLAS’s 97.1%, and V6 finishes in 1.548 ms against 1.560 ms at the profiler’s pinned clock.

Both kernels also share the same remaining loss: 1,664 tiles over 132 SMs is 12.6 rounds, so the last round runs with 40% of the SMs idle.

Warning

The ratio is also sensitive to how long the timed window is, because an 8192³ fp16 GEMM saturates the board’s power budget. At 100 timed iterations instead of this tutorial’s 30, the H100 throttles partway through the window. 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 complete source is at examples/hopper_matmul/matmul_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 rebalanced the register file for a 128 x 320 tile, shared B across a 2-CTA cluster, and made the kernel persistent (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 had to move registers from the warp group that does not need them to the ones that do.

Caution

The result reported here is specific to this shape (M=N=K=8192), dtype (fp16 inputs and outputs, fp32 accumulation), GPU (H100 SXM), and benchmark methodology. It is not a claim about other GEMM shapes: at 4096³, for example, the 320-wide tile leaves the last of only four rounds mostly idle, and V6 trails cuBLAS. The autotune spaces checked into the examples are pinned to a single configuration each, tuned for exactly this workload.