5. Overlapping WGMMA Groups

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 V4 untouched.

The Full Kernel

MatmulWGMMAV5 — 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]


# block_m must be >= 128 so each WG's WGMMA M = block_m/2 >= 64.
# Best of the original 96-schedule search on H100.
@tilus.autotune("num_stages", [4])
@tilus.autotune("block_m, block_n", [[128, 256]])
@tilus.autotune("block_k", [64])
@tilus.autotune("swizzle_size", [4])
class MatmulWGMMAV5(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

    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 __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
        block_m_half = block_m // 2

        num_m_blocks = cdiv(m_size, block_m)
        num_n_blocks = cdiv(n_size, block_n)
        self.attrs.blocks = num_m_blocks * num_n_blocks
        self.attrs.warps = 9  # 1 producer + 2 consumer WGs (4 warps each)

        m_block, n_block = self.compute_block_coord(
            self.blockIdx.x, num_m_blocks, num_n_blocks
        )
        offset_m: int32 = m_block * block_m
        offset_n: int32 = n_block * block_n

        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])
        # Per-WG A slab: index as sa[stage, wg_idx].
        sa = self.shared_tensor(
            dtype=float16, shape=[num_stages, 2, block_m_half, block_k]
        )
        sb = self.shared_tensor(dtype=float16, shape=[num_stages, block_n, block_k])

        tma_pipe = Pipeline(num_stages, producer_arrive_count=1, consumer_arrive_count=2)

        with self.thread_group(thread_begin=256, num_threads=32):  # TMA producer
            for offset_k in self.range(0, k_size, block_k, unroll=num_stages):
                tma_pipe.producer_acquire()
                with self.single_thread():
                    self.mbarrier.arrive_and_expect_tx(
                        tma_pipe.producer_barrier(),
                        transaction_bytes=sa[tma_pipe.producer_stage, 0].nbytes
                        + sa[tma_pipe.producer_stage, 1].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 + block_m_half, offset_k],
                    mbarrier=tma_pipe.producer_barrier(),
                )
                self.tma.global_to_shared(
                    src=gb,
                    dst=sb[tma_pipe.producer_stage],
                    offsets=[offset_n, offset_k],
                    mbarrier=tma_pipe.producer_barrier(),
                )
                tma_pipe.producer_advance()

            for _ in self.range(min(num_stages, cdiv(k_size, block_k))):
                tma_pipe.producer_acquire()
                tma_pipe.producer_advance()

        with self.thread_group(thread_begin=0, num_threads=128):  # consumer WG0
            acc0 = self.register_tensor(
                dtype=float32, shape=[block_m_half, block_n], init=0.0
            )
            tma_pipe.consumer_acquire()
            self.wgmma.fence()
            self.wgmma.mma(
                sa[tma_pipe.consumer_stage, 0],
                sb[tma_pipe.consumer_stage].transpose(),
                acc0,
            )
            self.wgmma.commit_group()
            tma_pipe.consumer_advance()

            for offset_k in self.range(block_k, k_size, block_k, unroll=num_stages):
                tma_pipe.consumer_acquire()
                self.wgmma.fence()
                self.wgmma.mma(
                    sa[tma_pipe.consumer_stage, 0],
                    sb[tma_pipe.consumer_stage].transpose(),
                    acc0,
                )
                self.wgmma.commit_group()
                self.wgmma.wait_group(1)
                with self.single_thread():
                    self.mbarrier.arrive(tma_pipe.prev_consumer_barrier())
                tma_pipe.consumer_advance()

            self.wgmma.wait_group(0)
            with self.single_thread():
                self.mbarrier.arrive(tma_pipe.prev_consumer_barrier())

            casted0 = self.cast(acc0, dtype=float16)
            self.store_global(gc, casted0, offsets=[offset_m, offset_n])

        with self.thread_group(thread_begin=128, num_threads=128):  # consumer WG1
            acc1 = self.register_tensor(
                dtype=float32, shape=[block_m_half, block_n], init=0.0
            )
            tma_pipe.consumer_acquire()
            self.wgmma.fence()
            self.wgmma.mma(
                sa[tma_pipe.consumer_stage, 1],
                sb[tma_pipe.consumer_stage].transpose(),
                acc1,
            )
            self.wgmma.commit_group()
            tma_pipe.consumer_advance()

            for offset_k in self.range(block_k, k_size, block_k, unroll=num_stages):
                tma_pipe.consumer_acquire()
                self.wgmma.fence()
                self.wgmma.mma(
                    sa[tma_pipe.consumer_stage, 1],
                    sb[tma_pipe.consumer_stage].transpose(),
                    acc1,
                )
                self.wgmma.commit_group()
                self.wgmma.wait_group(1)
                with self.single_thread():
                    self.mbarrier.arrive(tma_pipe.prev_consumer_barrier())
                tma_pipe.consumer_advance()

            self.wgmma.wait_group(0)
            with self.single_thread():
                self.mbarrier.arrive(tma_pipe.prev_consumer_barrier())

            casted1 = self.cast(acc1, dtype=float16)
            self.store_global(gc, casted1, offsets=[offset_m + block_m_half, offset_n])

What Changed from V4

	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

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: wgmma.commit_group() closes a group over the MMAs issued since the last commit, groups complete in order, and 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:

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():

Releasing the stage one behind the current one

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

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:

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

TMA producer warp

with self.thread_group(thread_begin=256, num_threads=32):  # TMA producer
    for offset_k in self.range(0, k_size, block_k, unroll=num_stages):
        tma_pipe.producer_acquire()
        with self.single_thread():
            self.mbarrier.arrive_and_expect_tx(
                tma_pipe.producer_barrier(),
                transaction_bytes=sa[tma_pipe.producer_stage, 0].nbytes
                + sa[tma_pipe.producer_stage, 1].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 + block_m_half, offset_k],
            mbarrier=tma_pipe.producer_barrier(),
        )
        self.tma.global_to_shared(
            src=gb,
            dst=sb[tma_pipe.producer_stage],
            offsets=[offset_n, offset_k],
            mbarrier=tma_pipe.producer_barrier(),
        )
        tma_pipe.producer_advance()

    for _ in self.range(min(num_stages, cdiv(k_size, block_k))):
        tma_pipe.producer_acquire()
        tma_pipe.producer_advance()

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

Consumer warp group 0

with self.thread_group(thread_begin=0, num_threads=128):  # consumer WG0
    acc0 = self.register_tensor(
        dtype=float32, shape=[block_m_half, block_n], init=0.0
    )
    tma_pipe.consumer_acquire()
    self.wgmma.fence()
    self.wgmma.mma(
        sa[tma_pipe.consumer_stage, 0],
        sb[tma_pipe.consumer_stage].transpose(),
        acc0,
    )
    self.wgmma.commit_group()
    tma_pipe.consumer_advance()

    for offset_k in self.range(block_k, k_size, block_k, unroll=num_stages):
        tma_pipe.consumer_acquire()
        self.wgmma.fence()
        self.wgmma.mma(
            sa[tma_pipe.consumer_stage, 0],
            sb[tma_pipe.consumer_stage].transpose(),
            acc0,
        )
        self.wgmma.commit_group()
        self.wgmma.wait_group(1)
        with self.single_thread():
            self.mbarrier.arrive(tma_pipe.prev_consumer_barrier())
        tma_pipe.consumer_advance()

    self.wgmma.wait_group(0)
    with self.single_thread():
        self.mbarrier.arrive(tma_pipe.prev_consumer_barrier())

    casted0 = self.cast(acc0, dtype=float16)
    self.store_global(gc, casted0, offsets=[offset_m, offset_n])

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

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; V6 gives that up. The complete source is at examples/hopper_matmul/matmul_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 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.