Optimizing GEMMs using `tcgen05` - Part 1

09 Sep, 2026

I had imagined this being one long post but it has gotten out of hand now. It's absurdly long already and I am still at Phase 4. So, will cover up the parts from Phase 5 in Part 2.

Since my last post on LBO/SBO/SWZ post, I spent some time understanding the tcgen05 instruction family and the extremely convoluted set of configurations it exposes. First things first, I am still not exactly at par with cuBLAS and have about 18% of the performance gap to cover (and hopefully one of the days, that hill shall be conquered). Having said that, it has been a lot of fun to iteratively build upon the vanilla mma instruction-based GEMM and add tcgen05 instructions to improve up on the memory and compute path.

Overall, I have structured this post to be in about five sections where each section briefly talks about the motivation for that optimization, the corresponding code snippet, and the added benefit in performance. The first phase scores about 714 TFLOP/s on a 4096 x 4096 x 4096 GEMM on a B300. From there I was able to push the performance to about 1452 TFLOP/s, which is about 82% of the peak cuBLAS can achieve at 1735 TFLOP/s.

Working with Claude and Grok has been super interesting because I was able to get a lot of rather naive questions out of the way by just reasoning with them and at times pointing them to the right resources (Thank You Colfax). Also shared some points which I got blatantly wrong and they eventually clicked and made a lot more sense.

If I were to summarize what I learned and the role tcgen05 plays - the tensor pipe's duty cycle is the number that matters and everything that can be done to push it, should be done. And to be fair, I was able to get within a couple of percent of cuBLASLt by phase four. The end-to-end performance still lingered around 73% then, because the epilogue was hiding it. So, closing the remaining gap is something I hope to cover soonish.

Feedback welcome on X or LinkedIn.

Why this worklog?

After I finally understood the shared memory descriptors (LBO, SBO, and everything related), I thought well, the hard part was potentially over but boy I was wrong. The first kernel that still used __pipeline_memcpy_async for loading into the shared memory with the tcgen05.mma instruction for the tensor-core MMAs was only a third of the cuBLAS speed.

Of the many things I learnt along the way, one thing I certainly appreciate more is how much you can learn and optimize by just knowing what ncu metrics to look for and how to interpret them. So, naturally the entire exercise was about identifying the critical path affecting the latency and resolving which would unlock the next tier in performance. If it helps, I have kept the questions and the ncu counters in. And if you find any gaps with any statements below, please don't hesitate to reach out.

What this worklog covers

The kernel goes through six rungs (I am not sure why both Grok and Claude love using the word rungs but NO, "phase" is more human-speak):

• Phase 0: cp.async double buffer, wait for every MMA

  All 128 threads copy the next tile with 16-byte cp.asyncs, thread 0 issues the two tcgen05.mma calls, and everyone waits for them, then for the copies, then at a __syncthreads.

• Phase 1: Add more asynchronous mem copies

  Instead of just one copy, add NSTAGES of mem copies to run ahead. The MMA wait moves from "right after issue" to "right before that slot is refilled."

• Phase 2: TMA loads, then warp specialization

  Finally, replace the cp.async mem copies with the TMA and also a good time to add warp-specialization.

• Phase 3: BN=256 and an 8-warp epilogue

  Do twice the tensor core work per fixed cost. But increasing BN from 128 to 256 also means more warps are needed to drain in the epilogue - hence the change from 4 to 8 warps per block.

• Phase 4: a 2-SM cluster (cta_group::2)

  The feature Nvidia likes a lot for some reason - Two CTAs now compute one 256×256 tile. Each loads half of B and the MMA reads both halves. This enables an axes of improvement I didn't expect immediately

• Phase 5: Epilogue through the Shared Memory and write TMA finally

  The scattered st.global becomes one bulk store per block.

Across the six rungs, throughput at 4096³ went 713 → 824 → 959 → 1153 → 1285 → 1451 TFLOP/s, 40% to 82% of cuBLASLt. Okay, going back to using phases now. I don't like rungs as much

A good northstar I found was to look at how far off we are from peak FLOPS for one mainloop iteration. One iteration is one BM x BN x BK stage with K=16 MMA instructions, which at peak is in the ballpark of 130 cycles of tensor-pipe time for 128 x 128 x 32 and 260 cycles for 128 x 256 x 32. Each section below is about where the rest of the cycles in an iteration went.

Another thing that was a bit counter-intuitive - warp stalls/active time mattered much less than I expected. In a tcgen05 kernel, the threads (and warps to an extent) are just the medium through which the instructions are dispatched to the asynchronous engines - TMA, Tensor Cores - and as long as the tensor pipes are busy, we are good.

What I am measuring. One dense GEMM, C[M,N] = A[M,K] × B[K,N], A row-major (K-major), B row-major (N-major), fp16 inputs, fp32 output, no epilogue beyond the store. Two shapes: 4096³, and a skinny M=1024, N=K=4096, which essentially shows the brittle nature of these optimizations. Every kernel is checked against cuBLASLt (atol 0.01, rtol 0.001). Inputs are L2-resident across timing iterations (A and B are 32 MB each at 4096³).

Contents

• Some Results First
• Constants and helpers
• Phase 0: the Serialized baseline
• Phase 1: Let's add some stages
• Appendix A: Master Results
• Appendix B: Per-phase Hardware Profile (for the curious)
• References

Continued in Part-2

• Phase 2: TMA and warp specialization
• Phase 3: BN=256 and the epilogue win
• Phase 4: a 2-SM cluster
• Phase 5: staged epilogue write
• Some metrics that seemed important but were useless
• Some things that I was just wrong about
• What's left

1. Some Results First

Kernel config: fp16 inputs, fp32 accumulation and output, BM=128, BK=32 throughout; BN=128 for phases 0–2 and 256 from phase 3; 4 warps per block for phases 0–2 and 8 from phase 3. Unless stated otherwise, M=N=K=4096.

¹ *TMA kernel alone gives about 959 TFLOP/s (NS=3). Warp specialization on top of it pushed it to 978 at 4096³, about +1.5% improvement, and about +30% at M=1024, where each SM has one block and shortening its chain is the only thing that helps (section 5.5).

Highligt of the changes across these five phases:

• Phase 0: the kernel from the LBO/SBO post with the MMA and a sort of crude epilogue attached.
• Phase 1: NSTAGES stages in dynamic shared memory, per-slot mbarriers, __pipeline_wait_prior(NSTAGES−1).
• Phase 2: Time for TMA aka cp.async.bulk.tensor from one thread; eventually evolving into warp specialization.
• Phase 3: On warp specialization, BN widened to 256 and number of warps per block from 4 to 8
• Phase 4: __cluster_dims__(2,1,1), tcgen05.mma.cta_group::2, each CTA loading half of B
• Phase 5: each warp writes its 32-row block into a swizzled staging tile in shared memory and issues TMA stores.

All numbers are 2·M·N·K over wall time: 137.4 GFLOP per launch at 4096³. Headline numbers are harness medians over 10 iterations after 5 warmups. Profiler durations quoted in the text are single launches under --cache-control none and run a few percent slower. The M=1024 column stops after rung 2 because BN=256 leaves that shape with 64 blocks for 160 SMs; that shape needs split-K, which this series does not do.

Fig. 1: Kernel progression from 713 to 1451 TFLOP/s at 4096³. Phases 1–2 remove/reorder waits and copy work from the mainloop, phases 3–4 push the work per fixed cost and essentially cut the bytes per FLOP, Phase 5 is the only change to the epilogue. cuBLASLt's 1762 and the 2500 TFLOP/s hardware peak are shown for scale._

Some useful B300 details to remember

B300 Max-Q, CC 10.3. 160 SMs at ~1.94 GHz in these runs · 228 KiB shared memory per SM · 2500 TFLOP/s dense fp16 tensor peak (8192 FLOP per SM-cycle, so a 128×128×16 tcgen05.mma is 64 cycles of tensor time and a 128×256×16 one 128).

Tensor Memory (TMEM) is 128 lanes × 512 columns of 32-bit cells per SM. A 128×128 fp32 accumulator takes 128 columns, so at most four such blocks can hold accumulators on one SM; at BN=256 it is two. This cap, not shared memory, decides residency for most of this post.

For measurements Compile with -Xptxas -v and check for spills (didn't see any yet). For counters, one profiled launch after the warmups:

ncu -k gemm_tcgen5_v0 --launch-skip 7 --launch-count 1 --cache-control none --metrics <list> ./bench 4096 4096 4096

Counters of interest: gpu__time_duration.sum; sm__ctas_active.avg ÷ sm__cycles_active.avg, which is the average number of resident blocks per SM; smsp__inst_issued.sum; the smsp__average_warps_issue_stalled_* ratios; sm__pipe_tensor_subpipe_hmma_cycles_active.avg.pct_of_peak_sustained_elapsed, tensor pipe duty cycle which tracked the runtime in everyphase; lts__t_sectors_srcunit_tex_op_read.sum and lts__t_{requests,sectors}_op_writelaunch__occupancy_limit_shared_mempair for bytes moved; andlts__throughput/l1tex__throughputaspct_of_peak` to identify whether memory was ever the bottleneck.

2. Constants and helpers

All the variants of the kernels use the same constants unless explicitly compiled with a different value (like BN, NSTAGES).

2.1 Tile and stages

constexpr int BM = 128;          // rows of C per block (per CTA in phase 4)
constexpr int BN = 128;          // 256 from phase 3
constexpr int BK = 32;           // K per stage: two m128nNk16 MMA calls
constexpr int WARPS_PER_BLOCK = 4;   // 8 from phase 3

#ifndef NSTAGES
#define NSTAGES 4                
#endif

constexpr size_t A_STAGE_BYTES = (size_t)BM * BK * sizeof(half);   // 8 KB
constexpr size_t B_STAGE_BYTES = (size_t)BK * BN * sizeof(half);   // 8 KB at BN=128, 16 KB at 256
constexpr size_t GEMM_SMEM_BYTES = NSTAGES * (A_STAGE_BYTES + B_STAGE_BYTES) + 1024;
constexpr uint32_t STAGE_TX_BYTES = (uint32_t)(A_STAGE_BYTES + B_STAGE_BYTES);   // what expect_tx arms

GEMM_SMEM_BYTES carries 1024 bytes of slack for rounding (aligning A and B at 1024 boundaries. It is what the launch passes as dynamic shared memory, after raising cudaFuncAttributeMaxDynamicSharedMemorySize.

2.2 Descriptor constants

The shared-memory descriptors are the ones derived in the LBO/SBO post. A is K-major with SW64 and B is N-major with SW128. B's byte distances change once, in phase 2, when the TMA comes into the picture: one copy can cover at most 64 columns (due to the 128-byte SW128 span), so a BN-wide tile is loaded as BN/64 slabs, each stored contiguously, and the atoms end up packed differently. The table in section 5.1 has both

constexpr uint32_t A_LBO = 16;        // K-major with swizzle: unused, field encodes as 1
constexpr uint32_t A_SBO = 512;       // next 8-row band: 8 x 64 B
constexpr uint32_t A_SWZ = 4;         // SW64
constexpr uint32_t A_K16_BYTES = 32;  // second MMA call: K[16:32] starts 32 B in

// phases 0-1, cp.async packing: the two N atoms of each K band sit side by side
constexpr uint32_t B_LBO = 1024;      // next 64-wide N atom
constexpr uint32_t B_SBO = 2048;      // next 8-row K band
constexpr uint32_t B_SWZ = 2;         // SW128
constexpr uint32_t B_K16_BYTES = 4096;

// phase 2 on, TMA slab packing: one 64-wide slab for all of BK, then the next slab
constexpr uint32_t B_LBO = 4096;      // next slab
constexpr uint32_t B_SBO = 1024;      // next K band inside a slab
constexpr uint32_t B_K16_BYTES = 2048;
constexpr uint32_t B_SLAB_BYTES = 4096;

2.3 The PTX wrappers

Every asynchronous op in this kernel goes through one of apis. mbar_wait is a try_wait.parity spin. tcgen05_commit makes an mbarrier track completion of all previously issued MMAs. mma_f16 is the m128nNk16 call with the accumulate flag as a predicate.

__device__ uint32_t smem_u32(const void* p) { return (uint32_t)__cvta_generic_to_shared(p); }

__host__ __device__ uint64_t make_smem_desc(uint32_t saddr, uint32_t lbo, uint32_t sbo, uint32_t swz) {
    return (uint64_t)((saddr >> 4) & 0x3fff) | ((uint64_t)((lbo >> 4) & 0x3fff) << 16)
         | ((uint64_t)((sbo >> 4) & 0x3fff) << 32) | (1ULL << 46) | ((uint64_t)swz << 61);
}

__device__ void make_ab_descs(half A_tile[BM][BK], half B_tile[BK][BN],
                              uint64_t* a0, uint64_t* a1, uint64_t* b0, uint64_t* b1) {
    uint32_t a_base = smem_u32(&A_tile[0][0]), b_base = smem_u32(&B_tile[0][0]);
    *a0 = make_smem_desc(a_base,                A_LBO, A_SBO, A_SWZ);
    *a1 = make_smem_desc(a_base + A_K16_BYTES,  A_LBO, A_SBO, A_SWZ);
    *b0 = make_smem_desc(b_base,                B_LBO, B_SBO, B_SWZ);
    *b1 = make_smem_desc(b_base + B_K16_BYTES,  B_LBO, B_SBO, B_SWZ);
}

__device__ void mbar_init(uint64_t* mbar, uint32_t arrival_count) {
    asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" :: "r"(smem_u32(mbar)), "r"(arrival_count) : "memory");
}

__device__ void mbar_wait(uint64_t* mbar, uint32_t parity) {
    uint32_t addr = smem_u32(mbar), done = 0;
    while (!done) {
        asm volatile("{\n\t.reg .pred p;\n\t"
                     "mbarrier.try_wait.parity.shared::cta.b64 p, [%1], %2;\n\t"
                     "selp.b32 %0, 1, 0, p;\n\t}"
                     : "=r"(done) : "r"(addr), "r"(parity) : "memory");
    }
}

__device__ void tmem_alloc(uint32_t* smem_slot, uint32_t ncols) {
    asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;"
                 :: "r"(smem_u32(smem_slot)), "r"(ncols));
}
__device__ void tmem_dealloc(uint32_t d_tmem, uint32_t ncols) {
    asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(d_tmem), "r"(ncols));
}
__device__ void tmem_relinquish_alloc_permit() {
    asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;");
}

__device__ void tcgen05_commit(uint64_t* mbar) {
    asm volatile("tcgen05.commit.cta_group::1.mbarrier::arrive::one.shared::cluster.b64 [%0];"
                 :: "r"(smem_u32(mbar)) : "memory");
}

__device__ void mma_f16(uint32_t d_tmem, uint64_t a_desc, uint64_t b_desc, uint32_t idesc, uint32_t accumulate) {
    asm volatile("{\n\t.reg .pred p;\n\tsetp.ne.b32 p, %4, 0;\n\t"
                 "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t}"
                 :: "r"(d_tmem), "l"(a_desc), "l"(b_desc), "r"(idesc), "r"(accumulate));
}

2.4 TMA and cluster helpers (phases 2 and 4)

These come in later but belong in the same list. Phase 5 adds a third tensor map, for C, with a {32, 16} fp32 box and SWIZZLE_128B; it is shown in section 8.2. tma_load_tile arms an mbarrier with the stage's byte count and issues one 2-D bulk-tensor copy for A and one per 64-wide slab of B. In phase 4 the same function takes the barrier's address in rank 0's shared memory instead, uses the .cta_group::2 form of the copy. The cluster helpers give a CTA its rank, a cluster-wide barrier, the address of a barrier as seen in rank 0's shared memory (mapa), and the two-CTA forms of mma and commit.

I would recommend spending some time trying to understand how the cluster .cta_group::2 functionality works with tensor cores, shared memory visibility and such as the mbarriers do get a bit complicated at first.

__device__ void tma_load_tile(const CUtensorMap& tmap_a, const CUtensorMap& tmap_b,
                              uint64_t* full_data_bar, half (*A_dst)[BK], half (*B_dst)[BN],
                              int k, int m_base, int n_base) {
    uint32_t bar = smem_u32(full_data_bar);
    asm volatile("mbarrier.arrive.expect_tx.shared::cta.b64 _, [%0], %1;" :: "r"(bar), "r"(STAGE_TX_BYTES) : "memory");
    asm volatile("cp.async.bulk.tensor.2d.shared::cluster.global.tile.mbarrier::complete_tx::bytes [%0], [%1, {%2, %3}], [%4];"
                 :: "r"(smem_u32(&A_dst[0][0])), "l"(&tmap_a), "r"(k), "r"(m_base), "r"(bar) : "memory");
    #pragma unroll
    for (int slab = 0; slab < BN / 64; slab++) {
        uint32_t bd = smem_u32((unsigned char*)&B_dst[0][0] + slab * B_SLAB_BYTES);
        asm volatile("cp.async.bulk.tensor.2d.shared::cluster.global.tile.mbarrier::complete_tx::bytes [%0], [%1, {%2, %3}], [%4];"
                     :: "r"(bd), "l"(&tmap_b), "r"(n_base + slab * 64), "r"(k), "r"(bar) : "memory");
    }
}

__device__ uint32_t cluster_rank() { uint32_t r; asm volatile("mov.u32 %0, %%cluster_ctarank;" : "=r"(r)); return r; }

__device__ void cluster_sync_all() {
    asm volatile("barrier.cluster.arrive.aligned;" ::: "memory");
    asm volatile("barrier.cluster.wait.aligned;"   ::: "memory");
}

// the same barrier, addressed in rank 0's shared memory
__device__ uint32_t mbar_rank0_addr(uint64_t* mbar) {
    uint32_t remote;
    asm volatile("mapa.shared::cluster.u32 %0, %1, 0;" : "=r"(remote) : "r"(smem_u32(mbar)));
    return remote;
}

// commit the pair's MMAs; the arrive lands on the same offset in both CTAs (mask 0b11)
__device__ void tcgen05_commit_mc(uint64_t* mbar) {
    asm volatile("tcgen05.commit.cta_group::2.mbarrier::arrive::one.shared::cluster.multicast::cluster.b64 [%0], %1;"
                 :: "r"(smem_u32(mbar)), "h"((unsigned short)0x3) : "memory");
}

__device__ void mma_f16_2sm(uint32_t d_tmem, uint64_t a_desc, uint64_t b_desc, uint32_t idesc, uint32_t accumulate) {
    asm volatile("{\n\t.reg .pred p;\n\tsetp.ne.b32 p, %4, 0;\n\t"
                 "tcgen05.mma.cta_group::2.kind::f16 [%0], %1, %2, %3, p;\n\t}"
                 :: "r"(d_tmem), "l"(a_desc), "l"(b_desc), "r"(idesc), "r"(accumulate));
}

2.5 The instruction descriptor and the epilogue mapping

While I already covered this in the previous post LBO-SBO-SWZ, here's a quick refresher on it. idesc is built from the tile: fp32 accumulate, B N-major, N and M encoded in units of 8 and 16. It changes twice, N (due to TMA) at phase 3 and M at phase 4 (the pair's M is 256).

uint32_t idesc = (1u << 4)                     // D = fp32
    | (1u << 16)                               // B N-major
    | ((uint32_t)BN >> 3 << 17)                // N
    | ((uint32_t)BM >> 4 << 24);               // M (2*BM in phase 4)

The epilogue reads TMEM with tcgen05.ld.32x32b. Warp w is pinned to TMEM lanes (w % 4) * 32 .. +31, so each lane holds one row of the accumulator, and one instruction hands it 8 consecutive columns. Keep the lane-equals-row detail in mind. It does not matter until phase 3, and it is the entire problem at phase 5 (This comes in handy when you have established that the MMA is decently saturated and it's time to look into other directions).

const uint32_t taddr = d_tmem + ((uint32_t)(warp_id * 32) << 16);   // lane in bits [31:16], column in [15:0]
for (int c = 0; c < BN; c += 8) {
    uint32_t r[8];
    asm volatile("tcgen05.ld.sync.aligned.32x32b.x8.b32 {%0,%1,%2,%3,%4,%5,%6,%7}, [%8];"
                 : "=r"(r[0]), "=r"(r[1]), "=r"(r[2]), "=r"(r[3]), "=r"(r[4]), "=r"(r[5]), "=r"(r[6]), "=r"(r[7])
                 : "r"(taddr + c));
    asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory");
    // two float4 stores of r[0..3], r[4..7] to C[out_row][out_col .. +8]
}

I think it's time look at some numbers.

3. Phase 0: the serialized baseline — 713 TFLOP/s at 4096³, 460 at M=1024

3.1 The loop

Two static stages, A_shared[2][128][32] and B_shared[2][32][128], 32 KB in all. All 128 threads copy the next tile with __pipeline_memcpy_async, applying the SW64 and SW128 permutations in software exactly as the LBO/SBO post describes. Thread 0 builds the descriptors, issues two MMAs, and commits to one mbarrier. Then everyone waits for the mbarrier, then for the copies, then at a barrier.

The main loop:

for (int k = BK; k < K; k += BK) {
    int next_stage = 1 - stage;
    // all 128 threads: 4 cp.async for A, 4 for B, into next_stage
    load_tile(A_shared[next_stage], B_shared[next_stage], ..., k, ...);
    __pipeline_commit();

    make_ab_descs(A_shared[stage], B_shared[stage], &a0, &a1, &b0, &b1);
    if (threadIdx.x == 0) {
        mma_f16(d_tmem, a0, b0, idesc, (k == BK) ? 0 : 1);   // K[0:16]
        mma_f16(d_tmem, a1, b1, idesc, 1);                    // K[16:32]
        tcgen05_commit(&mma_mbar);
    }
    mbar_wait(&mma_mbar, mma_parity);   // everyone waits for the MMA pair
    mma_parity ^= 1;
    __pipeline_wait_prior(0);           // and for the copies issued above
    __syncthreads();
    stage = next_stage;
}

One thing I did not know at the time, and I am noting it here because it belongs to this code. cp.async writes land through the generic proxy and tcgen05.mma reads shared memory through the async proxy, so a fence.proxy.async.shared::cta belongs between the __syncthreads and the MMA. There is none in phases 0 and 1. They pass correctness by timing, not by guarantee. Phase 2 removes the problem rather than patching it.

3.2 The numbers, and the first question

M=1024 N=4096 K=4096   kernel 0.075 ms  460 TFLOP/s   cuBLASLt 0.025 ms  1400   → 32.9%
M=4096 N=4096 K=4096   kernel 0.193 ms  713 TFLOP/s   cuBLASLt 0.078 ms  1762   → 40.5%

Same code, same tile, 1.55× apart. I asked the obvious thing first: if I change one thing, what should it be? The answer has to come from per-iteration arithmetic, because at this phase I had not run the profiler yet.

At M=1024 the loop runs K/BK = 128 iterations. 0.075 ms at 1.94 GHz is ~145k cycles. On an SM holding a single block (there are 64 of those, more on that below) that block runs the whole time, so one iteration costs ~1100 cycles. Of those, ~130 are the two MMAs.

Fig. 2: One phase-0 iteration. The tensor pipe is busy for ~130 of ~1100 cycles. The MMA round trip (issue, execute, commit, mbarrier flip, 128 threads wait and resolve) is paid every iteration because nothing is allowed to run until it completes, and the cp.async wait is exposed because the copies were issued only about ~400 cycles earlier and thus remain exposed in the sequence.

So the loop issues copies, issues MMAs, and then stops and drains everything it just started. However, this cost is almost fixed and doesn't scale with how much work each tcgen05.mma between these barriers. So, ideally "make the work between barriers bigger" should be the best strategy - within the confines of the resource bottlenecks.

3.3 The same kernel at 4096³: what 713 proved

At 4096³, only the grid differs from 1024 x 4096 x 4096, and the per-SM arithmetic works out exactly. 1024 blocks × 128 iterations is ~819 block-iterations per SM. 0.193 ms is ~374k cycles. So each SM retires one iteration every ~457 cycles against 130 cycles of expected tensor work, which is 28.5%, and that is the same as 713/2500. The per-block chain is still ~1100 cycles long.

Now with 4096 x 4096 x 4096 and the work quantum, TMEM is the bottleneck and (surprisingly) not the shared memory. 32 KB per block would allow seven blocks, but 128 accumulator columns each allows four. At M=1024 the grid is 8×32 = 256 blocks on 160 SMs, all resident from the start, which means 96 SMs holding two blocks and 64 holding one. At 4096³ it is 1024 blocks, 6.4 per SM.

Fig. 3: The same per-iteration chain on an SM with one resident block and with four. Left: 64 of the 160 SMs at M=1024; every wait is a wait for the whole SM. Right: typical at 4096³; one block's wait is another's MMA. Perfectly staggered chains would keep the pipe ~47% busy; the measured 28% is chains colliding and bunching. The 1.55× between the two runs is entirely this.

A good lesson to remember here is that it's generally better to have more independent CTAs per SM than blindly increasing the number of warps per CTA - primarily if you are an indiscriminate user of __syncthreads().

The kernel has four warps per block, so even at full residency that is 16 of 64 warp slots. More warps are of little help here since all the warps stall on the same __syncthreads() and the same mbarrier, so a 256-thread version just gives you eight warps that go idle together. Instead, having more CTAs to be colocated and independently executing on the SM allows better overall utilization here.

So, either doing more work per iteration (like pushing BK=128 and have eight MMAs per commit) or increase the number of blocks per SM overall - both tend to be viable paths for increasing the utilization and hiding the latency of memory operations. However, prior to heading into that direction, I also wanted to make sure we do what we can to just reduce the visible overhead of these outstanding copies, MMAs, and barriers and that is what phase 1 dives into.

4. Phase 1: let's add some stages — 713 → 824 TFLOP/s

As evident, the solution to the waits in Fig. 2 is to just have more memory loads run ahead prior to the first tcgen05.mma is issued such that the latter is less blocked on memory transfers to finish - latency hiding from first principles. The standard fix is thus a ring: NSTAGES buffers, copies issued NSTAGES−1 tiles ahead, and now the threads wait on before the overwriting the buffer (waiting for the corresponding mma to finish) instead of waiting immediately after issuing the mem copies.

4.1 Dynamic shared memory and the 1024-byte round-up

First the interesting tid-bit. Blackwell (and I think even older generations?) allow only up to 48KB of static shared memory allocation. So, if you want larger allocations, please use pointers. In this case, the allocation is with some gymnastics to align with the 1024-byte boundary, equivalent of which in static allocation is just using __align(1024)__:

extern __shared__ unsigned char smem_raw[];
unsigned char* smem_al = (unsigned char*)(((size_t)smem_raw + 1023) & ~(size_t)1023);
half (*A_shared)[BM][BK] = reinterpret_cast<half (*)[BM][BK]>(smem_al);
half (*B_shared)[BK][BN] = reinterpret_cast<half (*)[BK][BN]>(smem_al + NSTAGES * A_STAGE_BYTES);

Why 1024, and why alignment matters at all - well something that continued to hound me in bits and pieces. Here's one more stab at explaining it -

The reader, which is tcgen05.mma following SW64 in A's descriptor, computes its XOR key from absolute address bits: (addr >> 7) & 3. So, the shifting the base by 128 bytes means the hardware's key sequence shifts with it, consistently. The software store, on the other hand, computes the key from the logical row index (and doesn't have any notion of the base address) - (row >> 1) & 3. If not aligned properly, address bits [8:7] inadvertantly contribute to the reader address, i.e. we need to have a 512-byte alignment for SW64, and 1024 for B's SW128. Misaligned base would have every read fetch the wrong slot q ⊕ key ⊕ ((base >> 7) & 3) relative to where the write put it.

**Fig. 4: The writer and the reader compute the swizzle key from different things. Copy takes the key from (row >> 1) & 3. tcgen05.mma keys off the actual address, (addr >> 7) & 3. They agree only if the stage base is a multiple of 512 bytes for A; top panel. Let's say if the base is moved by 128 bytes and every reader key is one higher than the writer's, so every read lands one chunk over; bottom panel. This would cause copy loops to read wrong data. (The LBO/SBO post handles the same fact from the descriptor side with MBO)

smem_al is 1024-aligned for B's sake. A sits at offset 0 and B at NSTAGES × 8192, and every stage stride is a multiple of 1024, so both arrays keep their alignment in every slot. That is the invariant to protect. A padded stage, an odd BK, or an mbarrier array placed between A and B in the carve breaks B while A keeps working. I added a static_assert for it.

4.2 The ring

A more of a ring buffer approach where Tile t lives in slot t % NSTAGES. The prologue loads tiles 0..NSTAGES−2, and iteration t prefetches tile tp = t + NSTAGES − 1 into slot tp % NSTAGES before consuming tile t.

The main loop now looks like this:

__shared__ uint64_t mma_mbar[NSTAGES];          // one per slot
uint32_t mma_parity[NSTAGES] = {0};

for (int t = 0; t < T; t++) {
    const int s  = t % NSTAGES;                  // slot to consume
    const int tp = t + NSTAGES - 1;              // tile to prefetch
    if (tp < T) {
        const int slot = tp % NSTAGES;
        if (tp >= NSTAGES) {                     // slot has a previous occupant:
            mbar_wait(&mma_mbar[slot], mma_parity[slot]);   // wait for its MMA
            mma_parity[slot] ^= 1;
        }
        load_tile(A_shared[slot], B_shared[slot], ..., tp * BK, ...);
    }
    __pipeline_commit();
    __pipeline_wait_prior(NSTAGES - 1);          // oldest group done = tile t landed
    __syncthreads();

    make_ab_descs(A_shared[s], B_shared[s], &a0, &a1, &b0, &b1);
    if (threadIdx.x == 0) {
        mma_f16(d_tmem, a0, b0, idesc, (t == 0) ? 0 : 1);
        mma_f16(d_tmem, a1, b1, idesc, 1);
        tcgen05_commit(&mma_mbar[s]);           // no wait here any more
    }
}
mbar_wait(&mma_mbar[(T-1) % NSTAGES], mma_parity[(T-1) % NSTAGES]);   // drain

Things to remember and verify - After the commit in iteration t there are NSTAGES + t groups outstanding, so allowing NSTAGES−1 to remain pending means groups 0..t are complete, which is exactly tile t's data. The commit runs even when tp ≥ T and nothing was loaded (I have just accepted this). Slot tp % NSTAGES was last read by the MMA of tile tp − NSTAGES, which is the commit mma_mbar[slot] tracks - this ensure that load_tile doesn't overwrite what the tensor core is still reading."

4.3 Why it helped

With this, each iteration waits on loads NSTAGES−1 iterations away: wait_prior(NSTAGES−1) requires only the oldest outstanding group, which was issued iterations ago (and hopefully arrived in time). The MMA wait also moved off the MMA path: the tcgen05.commit is immediately called after instruction issue but the wait for it happens at the top of a later iteration, before the slot is refilled.

So the per-iteration critical path went from issue loads + drain MMA + drain loads + sync to issue loads + wait MMA + sync + issue MMA. The exposed load latency, which was the larger term, is gone. This is the same move as the k2 TMA ring in the grouped-GEMM post, built from cp.async groups and tcgen05.commit flags instead of expect_tx.

4.4 Tracing the MMA loop by hand

This tracing was just for my clarity but if it helps. This is for NSTAGES=4:

t=0: prefetch tile 3 → slot 3, wait copies of tile 0, issue MMA on slot 0
t=1: mbar_wait slot 0, prefetch tile 4 → slot 0, wait copies of tile 1, issue MMA on slot 1

At t=1 the loop waits on mma_mbar[0], and the MMA that arrives on that barrier was committed at the end of t=0, a few instructions earlier. The wait is one iteration deep, and it stays one iteration deep for any NSTAGES for MMA, because maximum-depth prefetch always refills the slot that was freed most recently.

Fig. 5: The four-slot ring, iteration by iteration. Blue is the tile being consumed, orange the tile issued this iteration, yellow tiles issued earlier and not yet consumed. The load wait is three iterations deep and resolves instantly; the MMA wait targets the commit from one iteration ago and is fully exposed. Both facts follow from (t + NSTAGES − 1) ≡ (t − 1) mod NSTAGES.

So phase 1 removed the exposed load latency and kept an exposed MMA round trip per iteration. The fix would be a reorder inside the body: wait for tile t's copies, issue MMA t, and only then wait on MMA t−1 and refill. I never applied it as a reorder. Phase 2 gives the two waits different warps instead.

One question I had to answer for myself before believing that either helps. There is one tensor pipe per SM for tcgen05 purposes, and a block's MMAs execute in issue order, so MMA t does not run in parallel with MMA t−1 no matter when it is issued. The benefit is not parallelism. It is that MMA t is already sitting in the queue when t−1 finishes, so the pipe goes straight from one to the next. In the wait-first order the pipe drains after every pair, and the whole commit → mbarrier → thread notices → refill → issue sequence, a few hundred cycles, runs with nothing queued. Same total work, no bubble.

4.5 What the profiler said

The sweep (4096³, harness): NS=2 808.8 · NS=3 801.2 · NS=4 627.3. For reference, A/B against phase 0: 725 → 824 at 4096³, 485 → 475 at M=1024.

Some ncu profiling numbers for the curios - At NS=2 sm__ctas_active / sm__cycles_active = 1,043,057 / 317,636 = 3.28 resident blocks per SM. The occupancy for NS=2: block limit 6 by shared memory, 37.5% theoretical / 23.8% achieved, 24 active warps per SM. At NS=4, this collapses to three blocks since each takes 65 KB, hence it made little sense to zoom in on that (at least for now).

At M=1024 the same change did nothing: 485 → 478. Looking at the profile, it was evident - long_scoreboard barely moved (2.24 → 2.03), duration barely moved (73.5 → 74.1 µs), and the scheduler report reads 0.28 waves, 1.65 active warps per scheduler, and no eligible warp on 74.5% of cycles. Most SMs hold a single block, that block spends most of its time waiting on memory, and there is nothing else on the SM to run in the meantime.

Phase 1 shortens the wait on the MMA, not the wait on memory. At 4096³ it made a difference, because three or four resident blocks were already covering each other's memory stalls and the MMA round trip was the next thing in line. At M=1024 the memory wait is still the biggest item, and a shorter MMA wait just means the block reaches the memory wait sooner. The number that would have to move at this shape is the number of blocks per SM, which is a grid problem (split-K), not a mainloop problem and not my problem - right now.

What's in Part 2

So far, till Phase 1, we have been able to see the micro optimizations that push the TFLOP/s from about 400 TFLOP/s in Grouped GEMM to about 980 TFLOP/s - which is about 60% of cuBLAS performance for a 4096 x 4096 x 4096 FP16 GEMM. The numbers for 1024 x 4096 x 4096 are slightly worse (somewhere in the vicinity of 50% of cuBLAS).

Part 2 should cover going the rest of the way, till up to 83% of the cuBLAS performance - unless there is a low-hanging optimization that pushes the perf to be at par with cuBLAS.

Appendix A: master results

Meanwhile, sharing these results here that will give you an exact idea about what to expect in the next post.

Appendix B: Per-phase Hardware Profile (for the curious)

All at 4096³, one launch, --cache-control none.

References

Technical references

• NVIDIA, Parallel Thread Execution ISA — cp.async, cp.async.bulk.tensor, mbarrier, tcgen05.mma / commit / alloc / ld, cta_group::2, mapa, multicast::cluster, proxies and fences.
• NVIDIA, CUDA Programming Guide — dynamic shared memory, clusters, cuTensorMapEncodeTiled.
• NVIDIA, Nsight Compute Profiling Guide — stall reason and throughput metric definitions.

Blogs

• Gaurav Jain, LBO, SBO and SWZ for a tcgen05 tile — the descriptors and copy loops this kernel starts from.
• Gaurav Jain, Grouped GEMM for Imbalanced Experts on Blackwell: A WIP Worklog — the k2 mbarrier ring phase 1 mirrors.
• Colfax Research, Optimization diaries: S/P ping-pong for FlashAttention-4 decode — the TMEM ping-pong that phase 6 will borrow.
• Thien Tran, tcgen05 for dummies.
• Aleksa Gordić, Inside NVIDIA GPUs: Anatomy of high performance matmul kernels.
• Simon Boehm, How to Optimize a CUDA Matmul Kernel for cuBLAS-like Performance: a Worklog.