Optimizing GEMMs using tcgen05 - Part 2

24 Sep, 2026

Quick recap about Part 1. It had a bunch of code and Phase 1 which put the performance 824 TFLOP/s on a 4096 x 4096 x 4096 GEMM on a B300, which is about 47% of cuBLASLt (1762 TFLOP/s). In this part, I cover what takes it to about 1550 TFLOP/s, 88% of cuBLASLt.

One feedback I got from Part 1 was that it was too long - so trying to keep this one short. Good thing is I don't need to go over the PTX wrappers and other bookkeeping since they are already defined in Part 1 (Give it a read here if you haven't already).

Fig. 1: The phases in this part at 4096³. Phase 1 was covered in Part1. cuBLASLt is the reference, and the dashed box is the B300's 2500 TFLOP/s peak.

Contents

• Phase 2: TMA and warp specialization — 824 → 978 TFLOP/s
• Phases 3 and 4: BN=256 and a 2-SM cluster — 978 → 1285 TFLOP/s
• Phases 5 and 6: Optimize the Epilogue, then BK=64 — 1285 → 1550 TFLOP/s
• What's left
• Appendix C: numbers for this part

5. Phase 2: TMA and warp specialization — 824 → 978 TFLOP/s

In this, I begin to move from using the threads doing the copies to now using the TMA engine. When using cp.async per thread, the kernel had almost about 153 instructions issued per stage and I had to ensure the swizzle and address math for each of them were aligned as expected by the MMA. With TMA a lot of that changes.

5.1 TMA does the copying

I was excited getting this out of the way since this is where I got a good grip on what CUtensorMap means. Well the TLDR is it is a means to describe the global tensor on the host by specifying the shape, strides, box size, and swizzle. Now, only one thread is needed to program the TMA unit. Each instruction passes the tensor map and the coordinates of one box, and the TMA unit walks this tensor map, generates the addresses, applies the swizzle on the way into shared memory and zero-fills anything out of bounds. Since this is an asynchronous operation, a barrier keeps track of whenever the transfers are complete - the TMA engine decrements the mbarrier's transaction count (complete_tx), so the barrier's phase completes once all the bytes armed by expect_tx have arrived.

__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);
    // arm the barrier: this phase completes once STAGE_TX_BYTES have landed
    asm volatile("mbarrier.arrive.expect_tx.shared::cta.b64 _, [%0], %1;"
                 :: "r"(bar), "r"(STAGE_TX_BYTES) : "memory");
    // A: one BK x BM box
    uint32_t ad = smem_u32(&A_dst[0][0]);
    asm volatile("cp.async.bulk.tensor.2d.shared::cluster.global.tile.mbarrier::complete_tx::bytes"
                 " [%0], [%1, {%2, %3}], [%4];"
                 :: "r"(ad), "l"(&tmap_a), "r"(k), "r"(m_base), "r"(bar) : "memory");
    // B: SWIZZLE_128B caps the box at 128 B = 64 halfs, so B comes in as 64-column slabs
    #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");
    }
}

A quick thing to note here is that with SWIZZLE_128B the TMA loads are capped at 64 columns (128 bytes in fp16) - which needs to be reflected accordingly in the B CUtensorMap. Checkout B's descriptor (section 2.2 of Part 1):

In the first version, I retained the entire TMA mainloop to be on thread 0: prefetch, both waits, MMA issue and commit. The other 127 threads do nothing but wait at a __syncthreads for the epilogue. The NSTAGES sweep gave NS=2 850, NS=3 959, NS=4 692. (NSTAGES impacts several resources like shared memory used and thus can impact how many CTAs can be scheduled on the SM - max 4 per SM)

The memory side barely changed (7% fewer sectors, same long scoreboard). The gain comes from the work in front of each MMA issue: about 150 instructions per warp plus a block-wide barrier became four instructions from one thread (expect_tx and three copies).

5.2 Warp specialization

Obvious optimization from here was to reduce the scope of tasks one thread does and divide the responsibilities with one or more threads. With one thread running everything, the wait for an MMA to free a slot and the wait for a slot's data sit in the same instruction stream, so one blocks the other. mbarrier.try_wait is a per-thread instruction, and mbar_wait spins on it until the barrier's phase completes, which blocks any subsequent tasks on this particular thread. But since waiting for a slot to tma_load_tile is independent of an MMA on a different slot, it's natural to split them. Thread 32 (warp 1) is the producer and only waits on mma_mbar. Thread 0 (warp 0) is the consumer and, inside the loop, only waits on full_data_bar.

if (threadIdx.x == 32) {                              // producer, warp 1
    uint32_t mma_parity[NSTAGES] = {0};
    for (int t = 0; t < T; t++) {
        const int s = t % NSTAGES;
        if (t >= NSTAGES) {                           // slot s was read by the MMA of k-tile t - NSTAGES
            mbar_wait(&mma_mbar[s], mma_parity[s]);
            mma_parity[s] ^= 1;
        }
        tma_load_tile(tmap_a, tmap_b, &full_data_bar[s], A_shared[s],
                      B_shared[s], t * BK, m_base, n_base);
    }
}
if (threadIdx.x == 0) {                               // consumer, warp 0
    uint32_t full_data_parity[NSTAGES] = {0};
    for (int t = 0; t < T; t++) {
        const int s = t % NSTAGES;
        mbar_wait(&full_data_bar[s], full_data_parity[s]);
        full_data_parity[s] ^= 1;
        asm volatile("tcgen05.fence::after_thread_sync;" ::: "memory");
        uint64_t a_desc0, a_desc1, b_desc0, b_desc1;
        make_ab_descs(A_shared[s], B_shared[s], &a_desc0, &a_desc1, &b_desc0, &b_desc1);
        mma_f16(d_tmem, a_desc0, b_desc0, idesc, t == 0 ? 0 : 1);
        mma_f16(d_tmem, a_desc1, b_desc1, idesc, 1);
        tcgen05_commit(&mma_mbar[s]);                 // slot s is free once these MMAs finish
    }
    mbar_wait(&mma_mbar[(T-1) % NSTAGES], ((T - 1) / NSTAGES) & 1);   // drain: last MMA done
}
__syncthreads();

Threads are now responsible only for instruction dispatch. Compute and data movement are owned by the engines: the tensor core and the TMA unit.

Fig. 2: Threads only dispatch instructions from Phase 2 on. Left: Phase 1, where all 128 threads issue cp.async and meet at a __syncthreads before thread 0 can issue the MMAs. Right: Phase 2, where the producer sets the barrier and issues three TMA copies, the TMA unit does the address generation and the swizzle, and the consumer only waits for the bytes and issues MMAs. The commit on mma_mbar[s] triggers the producer refill slot s.

With the same build flags this went from 964 to 978 at 4096³, about +1.5%, and about +30% at M=1024. At 4096³ there are ~3.3 blocks resident per SM, so while one block's thread 0 waited on its own MMA, the other blocks' MMAs were already keeping the tensor pipe busy. At M=1024 each SM holds a single block, so shortening that block's wait shows up directly in runtime.

6. Phases 3 and 4: BN=256 and a 2-SM cluster — 978 → 1285 TFLOP/s

6.1 Phase 3: BN=256 and the 8-warp epilogue — 978 → 1153 TFLOP/s

The first thing to change here was just push BN from 128 to 256. This is because a single tcgen05.mma supports N=256 at M=128, so the loop stays the same. Each stage is still two MMAs, but they are now 128×256×16, which is about 260 cycles of tensor work per stage instead of 130. But cost for the barrier, expect_tx and commit doesn't change and there is 2x the work that happens. A is also read once per 256 columns of C instead of once per 128.

However, now the accumulator takes 256 of TMEM's 512 columns (this means only two CTAs can fit on the SM at most). One downstream consequence of this is the need to increase the number of warps - 4 to 8 in this case to drain the TMEM at the same time. tcgen05.ld lets warp w access only TMEM lanes 32*(w%4) to 32*(w%4)+31, so warps w and w+4 share the same 32 rows and split the columns:

const int wq = warp_id % 4;              // TMEM lane quarter = 32 rows of C
const int wg = warp_id / 4;              // column half
constexpr int HALF_N = BN / 2;
const uint32_t taddr = d_tmem + ((uint32_t)(wq * 32) << 16) + (uint32_t)(wg * HALF_N);

B's descriptor constants stay the same. BN=256 is four 64-column slabs instead of two, and the slab loop in tma_load_tile already handles that.

The sweep gave NS=2 - 825, NS=3 - 1079, NS=4 1153, and NS=5 720 TFLOP/s. At NS=5 a block needs 124 KB of shared memory, only one block fits per SM, and ncu shows no eligible warp on 92.9% of cycles. Against an 8-warp BN=128 build, L2 read sectors went from 49.3 M to 36.6 M, a ratio of 0.74, where the byte math (A traffic halves, B stays the same) predicts 0.75.

6.2 Phase 4: a 2-SM cluster — 1153 → 1285 TFLOP/s

On Blackwell, two SMs in a cluster can run one MMA together with tcgen05.mma.cta_group::2. The 2-SM MMA took a minute since there are a couple of nuanced implementation details like which instructions have to be issued on both the SMs vs which are issued by only one - designated as the leader. The leader CTA issues it with M=256. Each SM computes its own 128 rows of D into its own TMEM, reading A from its own shared memory and B from both CTAs' shared memory. So each CTA only loads its 128 rows of A and half of B (BNH = 128) - i.e. a third fewer bytes are read from the memory when compared to BN=256 on a single SM but same amount of work between two barriers.

The kernel gets __cluster_dims__(2, 1, 1), the TMEM alloc/dealloc/relinquish switch to cta_group::2, and the M field of idesc becomes 2 * BM. The mainloop:

if (threadIdx.x == 32) {                              // producer, warp 1 of BOTH CTAs
    uint32_t mma_parity[NSTAGES] = {0};
    uint32_t bar0[NSTAGES];                           // full_data_bar[s] in rank 0's smem
    for (int s = 0; s < NSTAGES; s++) bar0[s] = mbar_rank0_addr(&full_data_bar[s]);   // mapa
    for (int t = 0; t < T; t++) {
        const int s = t % NSTAGES;
        if (t >= NSTAGES) { mbar_wait(&mma_mbar[s], mma_parity[s]); mma_parity[s] ^= 1; }
        if (rank == 0)                                // one arrival, armed for both CTAs' bytes
            asm volatile("mbarrier.arrive.expect_tx.shared::cta.b64 _, [%0], %1;"
                         :: "r"(smem_u32(&full_data_bar[s])), "r"(2 * STAGE_TX_BYTES) : "memory");
        tma_load_tile(tmap_a, tmap_b, bar0[s], A_shared[s], B_shared[s], t * BK, m_base, n_base);
    }
}
if (threadIdx.x == 0 && rank == 0) {                  // consumer, leader only
    uint32_t full_data_parity[NSTAGES] = {0};
    for (int t = 0; t < T; t++) {
        const int s = t % NSTAGES;
        mbar_wait(&full_data_bar[s], full_data_parity[s]); full_data_parity[s] ^= 1;
        asm volatile("tcgen05.fence::after_thread_sync;" ::: "memory");
        uint64_t a_desc0, a_desc1, b_desc0, b_desc1;
        make_ab_descs(A_shared[s], B_shared[s], &a_desc0, &a_desc1, &b_desc0, &b_desc1);
        mma_f16_2sm(d_tmem, a_desc0, b_desc0, idesc, t == 0 ? 0 : 1);
        mma_f16_2sm(d_tmem, a_desc1, b_desc1, idesc, 1);
        tcgen05_commit_mc(t == T - 1 ? &done_mbar : &mma_mbar[s]);   // multicast to both CTAs
    }
}
if (threadIdx.x == 0) mbar_wait(&done_mbar, 0);      // both ranks
__syncthreads();

tma_load_tile changes too. It takes the address of the barrier in rank 0's shared memory and only rank 0 sets the barrier (above, for both CTAs), loads BNH / 64 slabs, and its copies carry .cta_group::2, which lets rank 1's copies signal a barrier in rank 0's shared memory.

Fig. 3: One 256×256 tile on two SMs. Each CTA loads its own A rows and its half of B. The leader issues one tcgen05.mma.cta_group::2 that reads both B halves and writes 128 rows into each SM's TMEM. Both CTAs' copies complete on rank 0's barrier, and the commit is multicast to both CTAs.

The two barriers are set up differently. For the loads, both CTAs' copies have to complete on the one barrier that rank 0's consumer waits on. So rank 1 gets the address of rank 0's full_data_bar through mapa, and rank 0 arms it for both CTAs' bytes (2 * STAGE_TX_BYTES). For the MMA, one instruction reads shared memory in both CTAs, and both producers need to know when their slot is free. So the commit is multicast to mma_mbar[s] at the same offset in both CTAs.

Because rank0 is responsible for issuing MMAs on both the CTAs and the MMAs run asynchronously, both ranks need to know when it's done and it can start reading out the TMEM and that's where done_mbar comes into the picture.

The sweep gave NS=2 797, NS=3 1061, NS=4 1173, NS=5 1285, NS=6 1261. As you would notice, bringing two SMs into the fold also allowed me to build deeper stages for TMA + MMA overlap.

Why it helps: per SM per stage it is 16 KB instead of 24 KB for the same 128×256×32 of tensor work. So a 5-deep ring (83 KB) still fits two blocks per SM, where Phase 3 dropped to one block per SM at NS=5. The performance improvement is a consequence of the deeper ring, which allows more overlap between compute and memory. At equal depth (NS=4) in matched sweeps the two perform the same, 1173 vs 1151, and L2 sits at 40–48% of peak. Waiting on both SMs' copies makes each stage's round trip longer, since the barrier completes only after the slower SM's bytes land. The smaller stages allow a ring deep enough to cover that latency.

7. Phases 5 and 6: the epilogue, then BK=64 — 1285 → 1550 TFLOP/s

7.1 SKIP_STORES completely

Making MMAs fast has its advantages but it also has its cons. It can very quickly make the stores to global memory the bottleneck. To concretely establish that, I skipped the stores completely (using SKIP_STORES) and the TFLOP/s pumped up to 1728 TFLOP/s from the 1285 TFLOP/s obtained in the previous section. This is almost on par with what cuBLASLt achieves.

7.2 Phase 5: staged epilogue write — 1285 → 1451 TFLOP/s

tcgen05.ld.32x32b gives lane i TMEM lane i, which is row i of the tile. So one warp's float4 store writes 16 B into 32 different rows of C, 16 KB apart, which is 32 separate L2 requests of one sector each. The fix here is to write each warp's fragment into shared memory lane-per-row, swizzled according to SW128, and hand it to a TMA store with a {32, 16} box (16 rows of one full 128-byte line each). C gets its own tensor map:

static CUtensorMap hb_tmap_c(const float* C, int M, int N, int K) {
    (void)K;
    CUtensorMap m{};
    uint64_t dim[2]    = {(uint64_t)N, (uint64_t)M};
    uint64_t stride[1] = {(uint64_t)N * sizeof(float)};
    uint32_t box[2]    = {32, 16};        // 32 fp32 = 128 B = the SW128 span, 16 rows
    uint32_t estr[2]   = {1, 1};
    CUresult r = cuTensorMapEncodeTiled(&m, CU_TENSOR_MAP_DATA_TYPE_FLOAT32, 2,
        (void*)C, dim, stride, box, estr,
        CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B,
        CU_TENSOR_MAP_L2_PROMOTION_L2_128B, CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE);
    if (r != CUDA_SUCCESS) { fprintf(stderr, "tmap_c failed: %d\n", (int)r); exit(1); }
    return m;
}

And the epilogue, per warp, with one 2 KB slice per warp placed after the input ring (C_STAGE_BYTES is added to the dynamic shared memory size):

float* slice = C_stage + warp_id * (C_SLICE_BYTES / sizeof(float));   // 2 KB per warp, after the ring
const uint32_t slice_addr = smem_u32(slice);
float* row_base = slice + (lane_id % 16) * 32;                         // slice row = lane % 16

for (int chunk = 0; chunk < HALF_N / 32; chunk++) {
    uint32_t r[4][8];                                 // 32 columns: 4 x tcgen05.ld.32x32b.x8, one wait
    ...
    for (int half = 0; half < 2; half++) {            // the 32 rows leave as two 16-row boxes
        if (lane_id == 0)                             // the 2 KB slice is free once TMA has read it
            asm volatile("cp.async.bulk.wait_group.read 0;" ::: "memory");
        __syncwarp();
        if (lane_id / 16 == half) {
            for (int q = 0; q < 4; q++)
                for (int h = 0; h < 2; h++) {
                    const int atom = q * 2 + h;
                    const int phys = atom ^ (lane_id & 7);     // SW128
                    *(float4*)(row_base + phys * 4) = make_float4(
                        __uint_as_float(r[q][h * 4 + 0]), __uint_as_float(r[q][h * 4 + 1]),
                        __uint_as_float(r[q][h * 4 + 2]), __uint_as_float(r[q][h * 4 + 3]));
                }
        }
        asm volatile("fence.proxy.async.shared::cta;" ::: "memory");   // generic smem writes -> visible to TMA
        __syncwarp();
        if (lane_id == 0) {
            asm volatile("cp.async.bulk.tensor.2d.global.shared::cta.tile.bulk_group"
                         " [%0, {%1, %2}], [%3];"
                         :: "l"(&tmap_c), "r"(col_base + chunk * 32),
                            "r"(out_row_base + half * 16), "r"(slice_addr) : "memory");
            asm volatile("cp.async.bulk.commit_group;" ::: "memory");
        }
        __syncwarp();
    }
}

Fig. 4: One epilogue Store. Phase 4 writes 16 B into each of 32 lines. Phase 5 stages the fragment in shared memory and TMA writes full lines. Phase 6 (section 7.3) writes the same full lines with a 256-bit store after a transpose through shared memory.

The same 64MB for C is written with 6.9x fewer L2 write requests and this directly translates into performance gains. The epilogue runs after a block's last MMA, so shortening it raises tensor duty directly: 54.4% to 62.3%, and 1285 to 1451.

7.3 Phase 6: BK=64 and st.global.v8 — 1451 → 1550 TFLOP/s

I think the overall phenomenon throughout remains the same - do more work for less instructions. The TMA loads allow a loading 128 B wide rows (aligned with the SMEM width). So, taking advantage of that, I bumped up BK=64. Thus, A's row in a stage is now 64 halfs = 128 B, which is exactly the SW128 span, so A moves from SW64 to SW128. B's slab is now 64 K-rows deep, so LBO doubles. This also allows doing four MMAs instead of erstwhile two.

#define BK 64
constexpr size_t A_STAGE_BYTES = (size_t)BM * BK * sizeof(half);    // 16 KB
constexpr size_t B_STAGE_BYTES = (size_t)BK * BNH * sizeof(half);   // 16 KB
constexpr uint32_t A_SBO = 1024;          // 8 rows x 128 B (was 512)
constexpr uint32_t A_SWZ = 2;             // SW128 (was 4, SW64)
constexpr uint32_t B_LBO = 8192;          // one slab = 64 K-rows x 128 B (was 4096)
constexpr uint32_t B_SLAB_BYTES = 8192;
// tmap_a: box {BK, BM} with CU_TENSOR_MAP_SWIZZLE_128B

// consumer: descriptors at a_base + k*32 and b_base + k*2048, k = 0..3
mma_f16_2sm(d_tmem, a_desc0, b_desc0, idesc, t == 0 ? 0 : 1);
mma_f16_2sm(d_tmem, a_desc1, b_desc1, idesc, 1);
mma_f16_2sm(d_tmem, a_desc2, b_desc2, idesc, 1);
mma_f16_2sm(d_tmem, a_desc3, b_desc3, idesc, 1);
tcgen05_commit_mc(t == T - 1 ? &done_mbar : &mma_mbar[s]);

Fig. 5: BK=32 vs BK=64 in one CTA's stage. Colors mark which K=16 MMA reads which bytes. The K=16 steps (+32 B in A, +2048 B in B) stay the same; SBO for A and LBO for B double.

Why it helps: the fixed cost of a stage (the barrier wait, expect_tx, the copies, the descriptors, the commit) is now paid once per ~520 cycles of MMA instead of once per ~260, and each output tile takes 64 stages instead of 128. The price is 32 KB per stage per CTA, so two CTAs per SM fit only up to NS=3.

The st.global.v8 epilogue. Blackwell has 256-bit global stores. Each warp copies its 32 x 32 fp32 chunk from TMEM into a padded shared memory tile, lane-per-row, then reads it back transposed so that 4 lanes cover one row's 128 B. Each warp's st.global.v8.f32 then writes 8 full lines. The tile reuses the input ring, which is idle once done_mbar has fired. This removes Phase 5's dedicated 16 KB of C staging, the TMA store, the bulk-group waits and the proxy fence.

constexpr int STG_PITCH = 36;             // 32 floats + 4 of padding per row
float* stg = reinterpret_cast<float*>(smem_al) + warp_id * (32 * STG_PITCH);   // 4.5 KB per warp, in the idle ring

for (int c0 = 0; c0 < HALF_N; c0 += 32) {
    uint32_t r[32];                       // 4 x tcgen05.ld.32x32b.x8 at taddr + c0 + q*8, one wait
    ...
    // lane i writes its row
    float4* my_row = reinterpret_cast<float4*>(stg + lane_id * STG_PITCH);
    for (int q = 0; q < 8; q++)
        my_row[q] = make_float4(__uint_as_float(r[q * 4 + 0]), __uint_as_float(r[q * 4 + 1]),
                                __uint_as_float(r[q * 4 + 2]), __uint_as_float(r[q * 4 + 3]));
    __syncwarp();
    // read back transposed: 4 lanes per row, 8 fp32 each, 8 rows per store
    for (int i = 0; i < 4; i++) {
        const int row = i * 8 + (lane_id >> 2);
        const int cc  = (lane_id & 3) * 8;
        const float4* src = reinterpret_cast<const float4*>(stg + row * STG_PITCH + cc);
        const float4 v0 = src[0], v1 = src[1];
        const int out_row = out_row_base + row;
        const int out_col = col_base + c0 + cc;
        if (out_row < M && out_col + 8 <= N)
            asm volatile("st.global.v8.f32 [%0], {%1,%2,%3,%4,%5,%6,%7,%8};"
                         :: "l"(&C[out_row * N + out_col]),
                            "f"(v0.x), "f"(v0.y), "f"(v0.z), "f"(v0.w),
                            "f"(v1.x), "f"(v1.y), "f"(v1.z), "f"(v1.w) : "memory");
    }
    __syncwarp();
}

A pitch of 36 floats (32 plus 4 of padding) makes both the row-wise writes and the transposed reads conflict-free.

Together these took 1451 to about 1550 at 4096³. To be honest, I think I am still a bit weak on this entire epilogue optimization but as a sanity check I ran some numbers - 1550 TFLOP/s is 88.7 µs, or ~172k cycles at 1.94 GHz. 512 CTAs × 64 k-tiles over 160 SMs is ~205 k-tiles per SM, so one retires every ~840 cycles against ~520 of MMA, which is 62%, the same as 1550/2500.

8. What's left

1550 is 88% of cuBLASLt. This next part is not me. Fable and Grok believe I should be able to get to cuBLASLt perf with these but I don't trust them and thus haven't tried them yet.

Fig. 6: The final kernel, one cluster over time. The input ring keeps the loads ahead of the MMAs, done_mbar releases the epilogue, and the tensor core is idle for this cluster during its epilogue.

Things I am yet to read more about but should perhaps be a perf booster:

• A persistent cluster with two TMEM accumulators (2 × 256 = 512 columns). The MMA thread alternates accumulators per tile, and the epilogue warps drain tile i while tile i+1 accumulates into the other one.
• The tail. Once the kernel is persistent with one cluster per SM pair, 256 output tiles over 80 pairs is 3.2 tiles per pair, so the last round has work for only 16 of the 80 pairs. Stream-K for the leftover tiles is something I have seen used often.
• Descriptor hoisting. Each slot's descriptors are the same every time the slot comes around, so they can be built once instead of every stage.

Appendix C: numbers for this part

All at 4096 x 4096 x 4096, fp16 in, fp32 out, B300.