Porting Block-Scaled GEMM Kernels from Blackwell to Rubin

If you have written a persistent block-scaled GEMM for first-generation Blackwell (sm_100a) or Blackwell Ultra (sm_103a), you know the tcgen05 execution model. One thread issues Tensor Memory Accelerator (TMA) loads into 128-byte-swizzled shared memory (SMEM). The block scale factors (SFA, SFB) are copied from SMEM into Tensor Memory (TMEM) with tcgen05.cp (UTCCP). tcgen05.mma accumulates FP32 products in TMEM, and a four-warp epilogue drains TMEM to registers with tcgen05.ld.

A Blackwell kernel ported unchanged to Rubin (sm_107a) leaves performance unused. In two places it goes wrong: an FP8 instruction with M = 64 is rejected when K = 64, and a cp.async.cg cache hint can silently stop applying.

Everything below is read from the dedicated sm_107a kernels in FlashInfer at commit 518f4c7f0, and from their sm_103a and sm_100a counterparts where I compare against those. I don’t have a Rubin part, so nothing here is measured. Where a claim rests on a comment in the code rather than on the code itself, I say so.

Rubin changes four things in hardware:

A wider Tensor Core K. Each block-scaled MMA instruction covers 64 bytes along K: K = 128 for FP4, and K = 64 for FP8 and for mixed FP8 × FP4. That is twice sm_100a for both, and 4/3 of sm_103a for FP4. A new instruction-descriptor bit, bit 3, selects it (section 1).
Packed 4-bit operands in mixed precision. tcgen05.mma.kind::mxf8f6f4 reads FP4 directly from packed SMEM, so the FP4 × FP8 path no longer needs Blackwell’s U4_UNPACK_U8 TMA unpacking (section 2).
More on-chip memory per SM. SMEM grows from 228 KiB to 328 KiB, and TMEM from 512 columns (256 KiB) to 576 columns (288 KiB) (section 3).
New quantization formats. There is an unsigned 8-bit UE5M3 scale-factor format for FP4, and an FP8 conversion (cvt ... .scaled::n1::ue8m0) that applies the block scale inside the cvt (section 6).

FlashInfer’s Rubin kernels build three techniques on top of these:

B-operand reuse at M = 512 (Bkeep-Breuse), which the extra TMEM makes possible, with scale-factor copies interleaved per K-block (section 4).
A swap-AB MoE kernel, which patches the instruction’s N at runtime to match the number of routed tokens (section 5).
A mixed-cluster persistent scheduler, which fills the SMs left over by large clusters with smaller ones and shares out work without hardware Cluster Launch Control (section 7).

Section 8 collects the changes into a checklist.

1. The wider K, and why it removes `sm_103a`’s circular descriptors#

FlashInfer names each K mode by the instruction K in units of 32 elements:

Kind	`sm_100a`	`sm_103a`	`sm_107a`
FP4 (`mxf4`, `mxf4nvf4`)	2xFP4: K = 64 (32 bytes)	3xFP4: K = 96 (48 bytes)	4xFP4: K = 128 (64 bytes)
FP8 and mixed (`mxf8f6f4`)	1xFP8: K = 32 (32 bytes)	1xFP8: K = 32 (32 bytes)	2xFP8: K = 64 (64 bytes)

The relationship that matters is between the instruction’s K extent in bytes and the 128-byte TMA swizzle atom (K_SW128).

1.1 Why `sm_103a`’s K = 96 was awkward#

On sm_100a, an FP4 instruction reads 32 bytes along K, and an FP8 instruction also reads 32 bytes. Four instructions fit exactly in one 128-byte swizzle atom.

On sm_103a, the FP4 instruction reads 48 bytes, and 128 is not a multiple of 48. The smallest span that holds a whole number of both is lcm(48, 128) = 384 bytes, which is three 128-byte TMA buffers and eight MMA instructions:

MMA 0 reads bytes 0 to 47 of buffer 0.
MMA 1 reads bytes 48 to 95 of buffer 0.
MMA 2 needs 48 bytes starting at byte 96, so it reads bytes 96 to 127 of buffer 0 and then bytes 0 to 15 of buffer 1.

FlashInfer’s sm_103a kernel (dense_blockscaled_gemm_sm103.py) handles that crossing with a circular SMEM descriptor. make_umma_smem_desc takes both the current buffer and the next one (next_src). In the PTX descriptor format this is bit 52, the leading-dimension stride mode, set to absolute. The leading-byte-offset field then holds the address of the next buffer rather than an offset, so the Tensor Core can continue reading there partway through an instruction.

Two more costs follow. A buffer can be released as soon as the last MMA that touches it has been issued: buffer 0 after MMA 2, while the K-tile is still running. So the sm_103a mainloop releases A/B buffers in the middle of a K-tile. The scale factors also don’t follow the three-buffer period of A and B: they advance by four SF buffers per K-tile with sf_vec_size = 16 and by two with sf_vec_size = 32 (sf_buffers_per_tile_k). So that kernel loads them from a separate TMA warp (tma_sf_warp_id, beside tma_ab_warp_id), with its own barrier ring.

1.2 How 64 bytes per instruction fixes it#

On sm_107a, every block-scaled instruction reads 64 bytes along K, which has three consequences:

Two instructions fill one swizzle atom. An MMA tile of K = 256 for FP4, or K = 128 for FP8 and mixed, is two instructions and exactly 128 bytes.
Scale factors line up with whole chunks. Scale factors are stored in 512-byte basic chunks, each covering 128 rows × 4 K-scales. With sf_vec_size = 32, one K-tile has 4 scales per row (MXFP8 at K = 128) or 8 (MXFP4 at K = 256): one or two chunks along K. With sf_vec_size = 16 (NVFP4 at K = 256), it has 16 scales per row, which is four chunks.
One TMA warp again. A, B, SFA and SFB all advance by one stage per 128-byte K-tile, and no instruction crosses a buffer boundary. The sm_107a kernels drop the circular descriptors and the mid-tile releases, and load all four operands from a single TMA producer warp.

1.3 The instruction descriptor on `sm_107a`#

If you emit tcgen05.mma through inline PTX, for example to set N at runtime as in section 5, you build the 32-bit instruction descriptor (idesc) yourself. These are the fields FlashInfer’s Rubin MoE kernel sets (dynamic_mainloop.py):

Bits	Field	Notes
0–2	(not set)	Left at 0 by this code.
3	`k_size_upper`	New on `sm_107a`; see below.
4–5	`b_sf_id`	Which 32-bit slot of the SFB TMEM word holds this instruction’s scales.
6	(not set)	Left at 0.
7–9	`a_format`	Element type of A; the codes depend on the kind (section 2.3).
10–12	`b_format`	Element type of B.
13–14	negate A / negate B	Set to 0.
15	`a_major`	K-major or MN-major.
16	`b_major`
17–22	`n_dim`	Instruction N divided by 8.
23–25	`scale_format`	Scale-factor type.
26	`sfa_layout`	Set to 0.
27–28	`m_dim`	Instruction M divided by 128.
29–30	`a_sf_id`	Slot of the SFA TMEM word.
31	`k_size`

On sm_100a, bit 31 alone selected the K size. On sm_107a, the 2x K mode is selected by bits 3 and 31 together, and the right combination depends on the kind:

Kind	Instruction K	Bit 3 (`k_size_upper`)	Bit 31 (`k_size`)
`mxf4`, `mxf4nvf4`	128	1	0
`mxf8f6f4`	64	0	1

There is also a new constraint on M. On sm_100a, a 1-CTA FP8 MMA could have M = 64. On sm_107a, whenever K = 64, the instruction M must be 128 with .cta_group::1 or 256 with .cta_group::2. FlashInfer’s sm_107a dense FP8 GEMM documents this limit, and its MoE descriptor builder raises an error for any other M.

2. Packed FP4 in mixed precision, and the format-code trap#

FP4 weights with FP8 activations (for example MXFP4 × MXFP8, with Float4E2M1FN for A, Float8E4M3FN for B and sf_vec_size = 32) run on the .kind::mxf8f6f4 datapath.

2.1 What `sm_100a` and `sm_103a` had to do#

On sm_100a and sm_103a, .kind::mxf8f6f4 reads both operands as one byte per element whenever their widths differ. Packed FP4 weights can’t go into SMEM as they are. The TMA descriptor has to use U4_UNPACK_U8 (.b4x16_p64), which expands each 4-bit value into its own byte.

That doubles the SMEM footprint of sA, from half a byte per element to one byte. It also halves the number of K elements in one 128-byte swizzle atom, from 256 to 128.

2.2 What `sm_107a` does instead#

On sm_107a, the same datapath reads FP4 directly from packed SMEM. FlashInfer picks the SMEM element type like this (tcgen05_smem_alloc_type, shown here simplified):

def get_smem_alloc_dtype(dtype, peer_dtype, arch_number: int):
    # sm_100a/sm_103a need 1-byte containers when the two operand widths differ.
    # sm_107a reads mixed FP4 from packed SMEM.
    needs_mixed_width_unpack = (arch_number in (100, 103)) and (dtype.width != peer_dtype.width)
    # A 6-bit operand is always unpacked: TMA has no packed 6-bit format.
    needs_unpack = needs_mixed_width_unpack or (6 in (dtype.width, peer_dtype.width))
    return Int8 if (needs_unpack and dtype.width < 8) else dtype

The second condition is the FP6 case. A 6-bit operand is unpacked to one byte on every architecture, Rubin included, because there is no packed 6-bit TMA format.

For a per-CTA tile of M = 128 and K = 128, sA per stage shrinks from 16 KiB (128 × 128 × 1 byte) to 8 KiB (128 × 128 × 0.5 bytes). How many extra stages that buys depends on the rest of the stage. Take a 256 KiB mainloop budget and an FP8 B tile, and ignore the small scale-factor buffers:

B tile	Stage on `sm_100a`	Stage on `sm_107a`	Stages in 256 KiB
N = 128 (16 KiB)	32 KiB	24 KiB	8 → 10
N = 256 (32 KiB)	48 KiB	40 KiB	5 → 6

2.3 The kind-dependent format codes#

Bits 7–9 (a_format) and 10–12 (b_format) of idesc give the element types. The hardware decodes them against a different table for each kind, and the tables overlap:

Element type	Code under `mxf4` / `mxf4nvf4`	Code under `mxf8f6f4`
`Float8E4M3FN`	not allowed	0
`Float8E5M2`	not allowed	1
`Float4E2M1FN`	1	5

The trap is code 1. In a pure-FP4 instruction it means E2M1. If you reuse it for the FP4 operand of a mixed mxf8f6f4 instruction, the hardware reads it as Float8E5M2. The descriptor is still well formed, so there is nothing to fault on, and the mistake shows up only as wrong results. FlashInfer avoids it by looking the code up per kind (_instruction_descriptor_format_code). The only way to catch it is to compare against a reference GEMM.

3. 576 TMEM columns and 328 KiB of SMEM#

Rubin enlarges both SM-local memories:

TMEM: 512 columns (256 KiB) → 576 columns (288 KiB) per SM.
SMEM: 228 KiB → 328 KiB per SM.

3.1 Two disjoint accumulators at N = 256#

In a block-scaled tcgen05.mma, the FP32 accumulator D and the scale factors SFA and SFB must all be in TMEM at the same time. Take a 2-CTA tile of 256 × 256 × 256, so each CTA owns 128 rows (cta_tile_m = 128, cta_tile_n = 256, tile_k = 256):

One FP32 accumulator stage takes 256 columns.
Two stages, so the MMA warp can accumulate tile $t+1$ while the epilogue drains tile $t$ , take 512 columns.
The scale factors take:

\text{Cols}_{\text{SFA}} = \frac{\text{tile\_k}}{\text{sf\_vec\_size}}, \qquad \text{Cols}_{\text{SFB}} = \frac{\text{tile\_n}}{128} \times \frac{\text{tile\_k}}{\text{sf\_vec\_size}}

That is 16 + 32 = 48 columns with sf_vec_size = 16, or 8 + 16 = 24 with sf_vec_size = 32.

On sm_100a and sm_103a, 512 + 48 = 560 columns doesn’t fit in 512. Blackwell kernels therefore overlap the two accumulator stages: stage 0 at columns 0–255 and stage 1 at 208–463, sharing columns 208–255. The epilogue then has to traverse subtiles alternately forwards and backwards. It also has to release the shared columns early, in the middle of the epilogue (fence_view_async_tmem_load() followed by consumer_release()).

On sm_107a, the same 560 columns fit in 576. Above 512, tcgen05.alloc allocates in multiples of 32 columns (544 or 576) rather than powers of two. FlashInfer’s rounding rule (_round_tmem_allocation_columns):

def round_tmem_alloc_cols(used_cols: int, capacity_cols: int = 576) -> int:
    if used_cols <= 512:
        cols = max(32, 1 << (used_cols - 1).bit_length())
    else:
        cols = ((used_cols + 31) // 32) * 32  # 544 or 576
    if cols > capacity_cols:
        raise ValueError(f"TMEM plan needs {cols} columns, exceeding {capacity_cols}.")
    return cols

The two accumulator stages sit at columns 0–255 and 256–511 without overlapping, and the scale factors at 512–559. The epilogue no longer needs reverse traversal or early releases.

3.2 When the overlap comes back: SFA written by transform warps#

FlashInfer’s Rubin MoE FC1 kernel (blockscaled_contiguous_gather_grouped_gemm_swiglu_fusion.py) turns the overlap back on at N = 256. That kernel is a grouped GEMM over experts whose A rows are gathered by token index rather than loaded as one contiguous block.

The reason is the SFA layout. When token rows A[token_id] and SFA[token_id] are gathered from global memory, SFA is stored as a plain row-major (M, K_sf) array. For K = 256 and sf_vec_size = 16, that’s 16 one-byte scales, 16 bytes per token row. It is not in the swizzled 512-byte MMA layout. UTCCP (tcgen05.cp) can only copy from the swizzled layout, so the kernel can’t use it for SFA. It writes SFA into TMEM from registers instead.

The kernel has 20 warps:

Warps	Role
0–3	Epilogue (`tcgen05.ld`, SwiGLU, optional quantization, TMA store)
4–7	Gather A (`cp.async.cg` or TMA gather4)
8	MMA
9	TMA for B and SFB
10	Scheduler
11	2-CTA relay for A when A uses `cp.async` (idle otherwise)
12–15	Gather SFA (`cp.async.cg`)
16–19	SFA transform (`ld.shared` + `tcgen05.st`)

The SFA path uses two of these groups:

Warps 12–15 gather. Each of their 128 threads issues one 16-byte cp.async.cg.shared.global, one token row’s 16 scales, into a row-major SMEM buffer sSFA[M_per_cta, 16, stage].
Warps 16–19 transform. Each lane l does four 16-byte shared loads (ld.shared.v4.b32) from rows 32 * ci + l, ci = 0 to 3, one per 32-row M-block. It then issues four tcgen05.st.sync.aligned.32x32b.x4.b32, one per K-group gi = 0 to 3. Each store writes four TMEM columns, one per M-block. A fence_view_async_tmem_store() follows the stores.

// One K-group gi: 32 lanes x 4 columns, one column per M-block (ci = 0..3),
// at columns 4*gi .. 4*gi + 3 of this SFA stage.
tcgen05.st.sync.aligned.32x32b.x4.b32 [tmem_addr], {r0, r1, r2, r3};

All four transform warps write the same data. A warp’s tcgen05.st can only reach its own 32-lane quarter of TMEM. The MMA expects the same scale factors in all four quarters, which is what UTCCP’s 4x32x128b copy produces by broadcasting. Without that broadcast, each warp writes the full SFA into its own quarter.

The transform warps also need several TMEM stages. With a single 16-column SFA stage, they couldn’t write K-tile $k+1$ until the MMA had finished reading K-tile $k$ . The kernel allocates four stages (num_sfa_tmem_stage = 4, 64 columns), so the transform warps can fill up to three K-tiles ahead of the one the MMA is reading. They also issue the shared loads before producer_acquire on the transform pipeline, so the load latency overlaps the MMA that still holds the next TMEM slot.

This is what costs the TMEM. Four SFA stages (64 columns) plus SFB (32) is 96 columns, and two disjoint accumulators at N = 256 would need 512 + 96 = 608 > 576. The kernel overlaps the accumulator stages by one epilogue subtile (epi_tile_n = 64): stage 0 at 0–255 and stage 1 at 192–447, 448 columns in all. Then:

\underbrace{448}_{\text{overlapped acc}} + \underbrace{64}_{\text{4-stage SFA}} + \underbrace{32}_{\text{SFB}} = 544 \le 576

3.3 The 248 KiB limit for `cp.async.cg` destinations#

The same kernel orders its shared-storage struct with the cp.async.cg targets first:

@cute.struct
class SharedStorage:
    # mbarriers and small metadata first
    ...
    sC: cute.struct.Align[...]    # epilogue staging
    # sSFA (and sA, when A is loaded with cp.async) before sB / sSFB
    sSFA: cute.struct.Align[...]  # cp.async.cg target: keep below 248 KiB
    sA: cute.struct.Align[...]
    sB: cute.struct.Align[...]    # TMA target: no such limit
    sSFB: cute.struct.Align[...]  # TMA target: no such limit

The comment in the source gives the reason: SFA is placed low “so SFA gets a low SMEM offset: CpAsync128.CG destination addr stays under the 248KB threshold (above which the .CG cache-mode hint would degrade to a plain CpAsync).”

The cp.async destination address is a runtime register value, so this can’t be done by the assembler. It has to be how the hardware treats the .cg hint for destinations at 248 KiB (0x3E000) or above. I couldn’t find it in the PTX ISA. The comment is the only source, and I haven’t measured it.

On earlier architectures the question never came up, because SMEM was smaller than 248 KiB: 164 KiB on sm_80, 228 KiB on sm_90, sm_100 and sm_103. On sm_107a, with 328 KiB, A and B across 8–10 stages can easily take more than 250 KiB. Declared after them, sSFA would land above the limit, and its streaming scale-factor loads would lose the L1 bypass, filling L1 with data that is read once.

TMA (cp.async.bulk) has no such limit, so the rule is to put cp.async.cg destinations low in the struct and TMA destinations above them. FlashInfer also keeps sC low, ahead of sSFA. sC isn’t a cp.async target, so if sC plus the cp.async buffers ever exceeded 248 KiB, you could move sC above sB instead.

4. B-operand reuse up to M = 512, and interleaved UTCCP#

Operands read from SMEM pass through a small collector buffer inside the Tensor Core before they reach the multipliers. Consider an MMA tile whose M is twice the instruction’s M (mma_tiler_m = 2 * mma_inst_m): 256 with a 1-CTA instruction of M = 128, or 512 with a 2-CTA instruction of M = 256, which is 256 rows per CTA. Each K-block then runs two instructions, one for each M-half of A, and both multiply the same B and SFB subtile.

4.1 Why the extra TMEM enables it#

FlashInfer builds two TiledMma objects:

tiled_mma_bkeep: a_collector_op = CollectorOp.DISCARD, b_collector_op = CollectorOp.FILL
tiled_mma_breuse: a_collector_op = CollectorOp.DISCARD, b_collector_op = CollectorOp.LASTUSE

The keep instruction reads B from SMEM and keeps it in the collector buffer. The reuse instruction reads B from the collector buffer, so it doesn’t read B from SMEM at all.

The cost is TMEM. Each CTA now holds two accumulator subtiles, one per M-half, so the accumulator columns double:

\text{Cols}_{\text{acc}} = \text{cta\_tile\_n} \times \text{num\_acc\_stage} \times 2

At cta_tile_n = 256, FlashInfer drops to one accumulator stage (num_acc_stage = 1), which gives 512 columns. On sm_100a and sm_103a that’s all of TMEM, with nothing left for scale factors. On sm_107a it leaves columns 512–575. With sf_vec_size = 16, SFA takes 32 columns (two M-halves × 16) and SFB 32, so 512 + 32 + 32 = 576 fills TMEM exactly. The larger SMEM holds the doubled sA and sSFA per stage without reducing the number of AB stages much.

The single accumulator stage has a cost: the epilogue of one tile no longer overlaps the mainloop of the next. FlashInfer uses one stage only when B-reuse is on and the tile N is 192 or 256; otherwise it keeps two.

4.2 The M-mode permutation for 2-CTA plus B-reuse#

With both use_2cta_instrs = True (instruction M = 256) and enable_breuse = True (mma_tiler_m = 512), FlashInfer gives the TiledMma a custom M layout:

m_layout = cute.make_layout(
    shape=(mma_inst_shape[0] // 2, 2, 2),                   # (128, 2 keep/reuse, 2 CTAs)
    stride=(1, mma_inst_shape[0], mma_inst_shape[0] // 2),  # (1, 256, 128)
)
permutation_mnk = (m_layout, mma_inst_shape[1], mma_inst_shape[2])

Without it, a 2-CTA TiledMma gives rows 0–255 of the 512-row tile to CTA 0 and rows 256–511 to CTA 1. The kernel splits the tile into the keep and reuse halves by indexing the M mode that separates them (mode 1). With the default layout, the keep instruction would get rows 0–127 on CTA 0 and rows 256–383 on CTA 1. One 256-row instruction would then cover two separate 128-row pieces.

The permutation gives the keep/reuse mode a stride of 256 and the CTA mode a stride of 128:

The keep instruction gets rows 0–127 on CTA 0 and 128–255 on CTA 1, so rows 0–255 together.
The reuse instruction gets rows 256–383 on CTA 0 and 384–511 on CTA 1, so rows 256–511 together.

4.3 Interleaved UTCCP for 4xFP4 with B-reuse#

In a 4xFP4 kernel with B-reuse, each AB stage has two K-blocks. Each K-block has two SFA tiles, one per M-half, and one SFB tile:

                    ┌─────┐
                    │ B0  │ (k_block 0)
                    ├─────┤
                    │ B1  │ (k_block 1)
                    └─────┘
  ┌─────┬─────┐     ┌──────────────┐
  │ A0  │ A2  │ ──► │ MMA0 (keep)  │
  ├─────┼─────┤     ├──────────────┤
  │ A1  │ A3  │ ──► │ MMA1 (reuse) │
  └─────┴─────┘     └──────────────┘

The plain order issues all six UTCCP copies (A0, A1, A2, A3, B0, B1) before the K-block loop. The first MMA needs only A0 and B0, but it is issued behind all six.

FlashInfer’s interleaved order (_mainloop_s2t_interleaved_copies, used for FP4 × FP4 with B-reuse) issues three copies per K-block, then that K-block’s two MMAs:

UTCCP of the keep SFA (A0 for K-block 0, A2 for K-block 1)
UTCCP of SFB (B0, then B1)
UTCCP of the reuse SFA (A1, then A3)
MMA0 (keep)
MMA1 (reuse)

The first MMA is now issued behind three copies instead of six. K-block 1’s three copies are issued after K-block 0’s two MMAs, so they can proceed while those MMAs execute.

5. A swap-AB MoE kernel: runtime N, split pipelines, and a TMEM transpose#

This section is about FlashInfer’s Rubin MegaMoE kernel (next_cutedsl_megamoe), which runs a whole expert-parallel MoE layer, communication included, in one kernel. Its techniques aren’t all Rubin-specific. Patching N at runtime and splitting the pipeline would work on sm_100a too. What is Rubin-specific is the descriptor layout from section 1.3 and the compact SFB layout below.

In an MoE layer, the expert weights are large and fixed, while the number of tokens routed to each expert is small and varies, often 15 to 90. To keep the weights on the 128- or 256-row M side of tcgen05.mma and put the token count on N, which has finer granularity, the kernel swaps the operands (swap_ab = True):

A (M): expert weights, [M = features, K]
B (N): routed tokens, [N = tokens, K]

5.1 Setting the instruction’s N at runtime#

An expert’s tile may be 128 or 256 tokens wide but hold only 23 real tokens. A static N = 256 instruction would spend about 91% of its work on padding (233 of 256 columns); a static N = 128 would spend 82%.

Instead, the MMA warp writes n_dim (bits 17–22) on every tile. It rounds the token count up to a multiple of 16 and divides by 8. It also writes the two-bit scale-factor slot IDs (bits 29–30 for a_sf_id, bits 4–5 for b_sf_id) from bits 30–31 of the TMEM addresses (compute_idesc):

def compute_idesc(static_base: int, valid_tokens: Int32, sfa_tmem_i32: Int32, sfb_tmem_i32: Int32) -> Int32:
    n_dim_val = ((valid_tokens + 15) & -16) >> 3
    desc = Int32(static_base) | (n_dim_val << 17)
    # Mask after each shift: Int32 is signed, so >> fills the high bits with copies of bit 31.
    desc |= ((sfa_tmem_i32 & 0xC0000000) >> (30 - 29)) & (0x3 << 29)
    desc |= ((sfb_tmem_i32 & 0xC0000000) >> (30 - 4)) & (0x3 << 4)
    return desc

The masks after the shifts matter. When bit 31 of a TMEM address is set, an arithmetic right shift of the signed value fills everything above the result with ones, and without the mask that would overwrite the whole descriptor.

Rounding to 16 matches the 2-CTA rule. With instruction M = 256, N must be a multiple of 16, which is 8 per CTA. A 1-CTA instruction allows multiples of 8, so 16 is conservative there. A tile with no tokens gives N = 0, which isn’t legal; compute_idesc doesn’t check for it, so the caller must not issue an MMA for an empty tile.

A smaller runtime N changes how the Tensor Core addresses both sB in SMEM and tCtSFB in TMEM, which needs two adjustments:

The TMA-B offset on the second CTA. With .cta_group::2, the instruction’s N is split evenly across the two CTAs: CTA 0 holds tokens 0 to $N_\text{inst}/2 - 1$ , and CTA 1 holds the rest. tma_partition is built for the static mma_tiler_n (for example 256), so it fixes CTA 1’s global offset at +128.

With 23 tokens, $N_\text{inst} = 32$ , and the MMA expects CTA 1 to hold tokens 16–31. Loading from +128 would fetch padding and never load tokens 16–31. So on the non-leader CTA, the TMA-B warp shifts its B view first (compute_non_leader_cta_load_shift):
$\text{load\_shift} = \frac{\text{align16}(\text{valid\_tokens})}{2} - \frac{\text{mma\_tiler\_n}}{2}$
Here that is 32/2 − 128 = −112, so CTA 1 loads from +16.
The compact SFB layout for N ≤ 128. For NVFP4 (sf_vec_size = 16), one instruction of K = 128 needs 8 scales per row, which is two 4-scale chunks along K. A tile with N up to 256 needs two chunks along N, so each instruction uses four 512-byte chunks. The MMA tile, two instructions of K = 128, uses eight, matching the 32 SFB columns from section 3.1 (eight chunks × 4 columns). Per instruction, SMEM orders the four chunks as follows (the Rest_Tiler mode, the mode that indexes the chunks):

SMEM index Rows K
0 N 0–127 K 0–63
1 N 128–255 K 0–63
2 N 0–127 K 64–127
3 N 128–255 K 64–127

With $N_\text{inst} > 128$ , the MMA expects all four in TMEM at indices 0 to 3. With $N_\text{inst} \le 128$ , it switches to a compact layout with only the first 128 rows: (N 0–127, K 0–63) at TMEM index 0 and (N 0–127, K 64–127) at index 1, not index 2. So for small tiles, the MMA warp copies SMEM index 0 to TMEM 0 and SMEM index 2 to TMEM 1.

SMEM index	Rows	K
0	N 0–127	K 0–63
1	N 128–255	K 0–63
2	N 0–127	K 64–127
3	N 128–255	K 64–127

5.2 Separate pipeline depths for weights and tokens#

With swap-AB, A (the weights) has cta_tile_m = 128, and B (the tokens) has cta_tile_n of 128 or 256. Going from 128 to 256 tokens leaves the bytes of A and SFA per stage unchanged, but doubles those of B and SFB. In a single shared AB pipeline, the number of stages drops accordingly.

The weights and the tokens don’t need the same depth. The weights stream from HBM and need many stages to hide DRAM latency. The tokens are few and usually already in L2.

So the kernel can split the mainloop into two PipelineTmaUmma rings. The token pipeline (B and SFB) gets a fixed six stages (asymmetric_token_stages = 6). The weight pipeline (A and SFA) gets whatever the budget has left:

\text{weight stages} = \left\lfloor \frac{\text{budget} - 6 \times \text{token stage bytes}}{\text{weight stage bytes}} \right\rfloor

Each pipeline needs its own full/empty mbarrier pair per stage, so the split pays for a second set of barriers. The MMA warp also waits on two barriers per stage instead of one. FlashInfer refuses the split unless it gives the weights more stages than the shared plan would (num_weight_ab_stages > shared_stages).

5.3 Transposing in TMEM in the epilogue#

With swap-AB, the accumulator holds features (M = 128) along the TMEM lanes and tokens (N = 64 per subtile) along the columns. The FC1 accumulator interleaves the gate and up projections in blocks of 64 features. SwiGLU combines each gate/up pair into one output, which halves the feature count.

To quantize the output in 16- or 32-feature blocks and TMA-store it as [token, feature], the epilogue has to transpose each [features, 64 tokens] subtile.

Transposing 128 × 64 FP32 values in SMEM would take 32 KiB of scratch, one or two mainloop stages. Instead, the epilogue uses the accumulator subtile it is draining as scratch, through tcgen05.ld and tcgen05.st. The MMA warp can’t write that TMEM stage until the epilogue calls acc_pipeline.consumer_release.

That reasoning holds only while no other stage shares those columns. In a kernel with overlapping accumulator stages, like the one in section 3.2, the shared columns would have to stay unreleased until the transpose is done. I haven’t checked whether FlashInfer ever combines the two.

The steps are:

Load. Each warp loads gate and up for tokens 0–31 and 32–63 into four 16-register FP32 arrays (Ld16x64bOp.x16).
Transpose and SwiGLU, interleaved.
- The activation is alpha_swiglu_clamp. It dequantizes gate and up by the expert’s global scale alpha, clamps them in the gpt-oss style (gate from above only, up symmetrically), and then computes up * gate * sigmoid(gate).
- For tokens 0–31, the warp applies SwiGLU first. That folds 32 gate and up values per thread into 16 outputs, so only a half-size transpose (16 × 32 → 32 × 16, TmemTranspose16x32) is left. On entry, register elem_idx of lane lane_idx holds TMEM lane 2 * elem_idx + (lane_idx // 2) % 2 and column (lane_idx % 2) * 8 + lane_idx // 4. After the four rounds, it holds lane lane_idx and column elem_idx, so each thread owns one full row.
- For tokens 32–63, the warp transposes first, a full 32 × 32 in four rounds of TMEM stores and loads (St16x128bOp.x8, then Ld16x64bOp.x16 + St32x32bOp.x16, then Ld16x256bOp.x2 + St32x32bOp.x16, then Ld16x128bOp.x4). These rounds are interleaved with the SwiGLU arithmetic for tokens 0–31.
- After the fourth round, thread lane_idx of each warp holds 16 contiguous output features for token lane_idx and 16 for token lane_idx + 32: 32 FP32 registers.
The amax exchange for 32-element blocks. Each warp owns 16 output features, so a 16-element NVFP4 block lies entirely within one thread. A 32-element MXFP8 block spans warps w and w ^ 1, which must exchange half-block amax values.

The exchange buffer holds one 4-byte slot per thread per 32-token chunk: 2 chunks × 4 warps × 32 threads × 4 bytes = 1024 bytes. The epilogue doesn’t allocate this separately. It borrows it from the current subtile’s FC1 output staging buffer in SMEM, which isn’t written until the quantized values are stored after the exchange barrier.

6. New quantization formats: `UE5M3` scales and a scaled FP8 `cvt`#

6.1 `UE5M3` scale factors for FP4#

On sm_100a and sm_103a, block-scaled FP4 supports two scale-factor formats:

UE4M3 (Float8E4M3FN): 4 exponent bits and 3 mantissa bits, with a maximum of 448 and a dynamic range of about $10^5$ . That range is too narrow to cover both outliers and the overall magnitude of a tensor. NVFP4 therefore pairs it with a per-tensor FP32 global scale (alpha, and norm_const when quantizing outputs), which MoE kernels carry per expert.
UE8M0 (Float8E8M0FNU): 8 exponent bits and no mantissa, giving the full FP32 exponent range (about $10^{76}$ from smallest to largest) without a global scale. Scales are then powers of two, so the step between representable scales is a factor of 2.

On sm_107a, tcgen05.mma also accepts FloatNV8E5M3FNU (UE5M3), with sf_vec_size 16 or 32. Block scales are never negative, since they are amax / 6, so E4M3’s sign bit is wasted. UE5M3 uses it as a fifth exponent bit. The dynamic range grows from about $10^5$ to about $10^{10}$ , and the 3 mantissa bits stay.

6.2 FP8 conversion with the scale applied: `cvt.rn.satfinite.scaled::n1::ue8m0`#

To quantize FP32 values to MXFP8 (E4M3 or E5M2 with UE8M0 scales) on sm_100a, a kernel computes the scale byte with cvt.rp.satfinite.ue8m0x2.f32, and then has to:

Widen the ue8m0 byte back to BF16 and FP32.
Compute the reciprocal with rcp.approx.ftz.f32 on the MUFU.
Clamp infinity (fmin(..., Fp32Max)) and mask zero (fmin(scale_f32 * 1e30, 1.0)), so that 0 × ∞ doesn’t produce NaN.
Multiply all 32 values by the reciprocal.
Convert the 32 scaled values to e4m3x2.

On sm_107a, cvt takes the scale byte directly:

// Converts four FP32 inputs ($1..$4), each divided by the UE8M0 scale byte s,
// to four E4M3 values packed in $0.
cvt.u8.u32 s, $5;
cvt.rn.satfinite.scaled::n1::ue8m0.e4m3x2.f32 lo, $2, $1, s;
cvt.rn.satfinite.scaled::n1::ue8m0.e4m3x2.f32 hi, $4, $3, s;
mov.b32 $0, {lo, hi};

Dividing an FP32 value by a power-of-two scale $2^{s-127}$ amounts to subtracting $s - 127$ from its exponent before rounding to FP8. The converter still has to handle results that fall into FP8’s subnormal range and saturate those that overflow. Doing the division inside the converter removes the MUFU reciprocal, the NaN guards, and the 32 FP32 multiplies per block.

7. Mixed-cluster persistent scheduling without CLC#

Suppose the best cluster shape for TMA multicast is (4, 1), 4 CTAs, or (4, 2), 8 CTAs. Clusters are placed within a GPC, so if the SMs available for clusters in a GPC aren’t a multiple of the cluster size, a launch that uses only that shape leaves the remainder idle.

Launching with a preferred and a fallback cluster shape isn’t new on Rubin. CUDA 12.8 added cudaLaunchAttributePreferredClusterDimension for Blackwell. sm_100 also has hardware Cluster Launch Control (CLC, clusterlaunchcontrol.try_cancel) for dynamic persistent scheduling.

What FlashInfer’s Rubin MoE kernel adds is a persistent scheduler for such a mixed launch that doesn’t use CLC: work IDs come from an atomic counter in global memory (non_clc_mixed_cga.py). The source doesn’t say why it avoids CLC. I don’t know whether that’s because sm_107a lacks it or because the kernel prefers its own counter.

The host passes both shapes at launch (from FlashInfer’s sm_107a dense GEMM):

kernel.launch(
    grid=preferred_grid,
    block=[threads_per_cta, 1, 1],
    cluster=(*preferred_cluster_shape_mn, 1),          # e.g. (4, 1, 1)
    fallback_cluster=(*fallback_cluster_shape_mn, 1),  # e.g. (2, 1, 1)
    smem_merge_branch_allocs=True,
    min_blocks_per_mp=1,
    stream=stream,
)

The kernel has two code paths, one per cluster shape, and only one runs in each cluster. smem_merge_branch_allocs=True lets the two paths share the same SMEM allocation. Three conditions have to hold for this to work.

7.1 Size the grid with floor division#

The grid’s x and y extents are one preferred cluster. Its z extent is the number of preferred-sized clusters to launch. The hardware runs as many preferred clusters as fit and fills the rest with fallback clusters.

The host computes two grids. fallback_grid is the persistent grid for a launch using only fallback clusters, sized from the fallback cluster’s occupancy. preferred_grid is the same for the preferred shape. The z extent comes from the fallback grid’s CTA count:

max_ctas_for_fallback_cluster = fallback_grid[0] * fallback_grid[1] * fallback_grid[2]
max_preferred_cluster_count = max_ctas_for_fallback_cluster // preferred_cluster_size_mn
preferred_grid = (preferred_grid[0], preferred_grid[1], max_preferred_cluster_count)

This has to be floor division. With ceil_div, when the CTA count isn’t a multiple of the preferred cluster size, the launch asks for more CTAs than one wave can hold. The hardware then runs a second, partial wave, and the code comments call that tail a significant regression.

7.2 Two TMA descriptors and runtime masks#

cpasync.tma_partition bakes the cluster’s M/N split into the TMA descriptor. A (4, 1) cluster and a (2, 1) cluster can’t share a descriptor for the multicast operands, B and SFB:

On the host, the kernel builds a (preferred, fallback) pair of TMA atoms for each multicast operand.
On the device, each CTA reads its cluster’s runtime shape (cute.arch.block_in_cluster_dim()) and decides whether it is in a fallback cluster. The MoE kernel compares only the M extent with the fallback M. That works only if the two shapes differ in M; the dense GEMM compares all three extents with the preferred shape, which is safer.
From that flag, the CTA picks the matching TMA atom. Its pipeline (PipelineTmaUmmaMixedCga, from FlashInfer’s name for this launch, “mixed-CGA”) sets the empty barrier’s expected arrival count to match the CTAs that consume each multicast stage (fewer in the fallback cluster). It also ANDs the multicast mask with the fallback pattern, so a fallback cluster never multicasts to CTAs it doesn’t have. For a (2, 1) fallback in a (4, 1) preferred layout, that pattern is 0x0003:

mov.b16 mask, PREFERRED_PATTERN;
@!is_preferred and.b16 mask, mask, FALLBACK_PATTERN;
shl.b16 mask, mask, offset;

7.3 Fallback clusters pair up for each work ID#

With a dynamic scheduler, each work ID is one preferred-cluster tile, for example 4 × 128 = 512 rows for (4, 1). A (2, 1) fallback cluster covers only 256 rows, half a work ID (split_factor = 4 / 2 = 2). Two fallback clusters therefore form a fixed group and split each work ID. The member with in_group_idx = 0 takes the first fallback_m CTA rows of the tile, and the member with in_group_idx = 1 the next fallback_m.

Separate clusters can’t reach each other’s distributed shared memory, so the group coordinates through a 64-bit token in global memory (work_id_claim.py):

Registration. At startup, each fallback cluster’s leader CTA atomically increments a global registration_counter to get fallback_ordinal. That fixes its group (fallback_ordinal // split_factor) and its place in the group (fallback_ordinal % split_factor) for the rest of the launch.
Claiming a work ID.
- Only the leader of member 0 claims work IDs from the global counter.
- Before it claims, lanes 0 to split_factor − 1 load the group’s token slots with acquire semantics, until every slot holds the previous token. That shows every partner has read the previous work ID. Slot 0 is member 0’s own slot and matches trivially.
- Member 0 then writes (next_generation << 32) | claimed_work_id to slot 0 and broadcasts the work ID to its own cluster through distributed shared memory.
- Each other member polls slot 0 until the token changes. It echoes the token into its own slot, broadcasts the work ID to its cluster, and offsets its CTA M coordinate by in_group_idx * fallback_m.

The work ID travels inside the token, so a partner reads the new generation and the work ID in one 64-bit load. No other data is published through the token, so the plain (relaxed) stores and partner loads are enough. When the claimed ID runs past the end of the work list, every member of the group gets the same ID and exits together. The registration counter and the token slots live in a device workspace that is reset after each launch (reset="tail_reset").

Two assumptions hold the scheme together:

Fallback clusters come in whole groups. The host rounds the fallback cluster count down to a multiple of split_factor, and the scheduler’s constructor rejects any other count. An odd number of (2, 1) fallback clusters, which would leave one member spinning forever, isn’t allowed.
All the fallback clusters run at once. The host sizes the launch from occupancy queries to fit in one wave. If some fallback clusters didn’t start until others finished, a member could wait forever for a partner that hasn’t started.

8. Checklist: porting an `sm_100a` or `sm_103a` kernel to `sm_107a`#

The items are mostly independent, so the order doesn’t matter.

Use the 2x K shapes. For 4xFP4, set mma_inst_shape_k = 128 and mma_tiler_k = 256. For 2xFP8 and mixed mxf8f6f4, set mma_inst_shape_k = 64 and mma_tiler_k = 128. Remove sm_103a’s circular SMEM descriptors and mid-tile releases, and merge the separate AB and SF TMA warps into one.
Update the K-size bits in idesc. Set bit 3 = 1, bit 31 = 0 for FP4 (K = 128), and bit 3 = 0, bit 31 = 1 for mxf8f6f4 (K = 64). With K = 64, use M = 128 with .cta_group::1 or M = 256 with .cta_group::2.
Drop U4_UNPACK_U8 for FP4 × FP8. Allocate the FP4 SMEM buffer as packed Float4E2M1FN. Use format code 5, not 1, for Float4E2M1FN under .kind::mxf8f6f4. FP6 operands still need unpacking.
Use the 576 TMEM columns. Allocate in multiples of 32 above 512, and check against capacity. At N = 256, turn off overlapping accumulators (512 + 48 = 560 ≤ 576), or spend the extra columns on B-reuse or on more SFA stages.
Keep cp.async.cg destinations below 248 KiB. Put sSFA, and sA when A is loaded with cp.async, low in the shared-storage struct, with the TMA-loaded sB and sSFB above.
Use B-reuse at M = 512. Build keep and reuse TiledMma objects with CollectorOp.FILL and CollectorOp.LASTUSE. Apply the (128, 2, 2) M permutation in 2-CTA mode. Expect one accumulator stage at N = 256.
Interleave the UTCCP copies with B-reuse. Issue the keep SFA, SFB and reuse SFA copies for each K-block just before that K-block’s two MMAs, not all at the start.
For swap-AB MoE kernels: set N in idesc per tile (rounded to 16), mask the scale-factor ID shifts, shift the non-leader CTA’s TMA-B view in 2-CTA mode, and remap SFB to the compact layout when N ≤ 128. Consider separate weight and token pipelines.
Use UE5M3 and .scaled::n1::ue8m0. Use FloatNV8E5M3FNU scale factors for FP4 when you need range without a global scale. In MXFP8 quantization, replace rcp.approx and the FP32 multiplies with cvt.rn.satfinite.scaled::n1::ue8m0.{e4m3x2,e5m2x2}.f32.
For mixed-cluster persistent launches: size the grid with floor division, build a preferred and a fallback TMA descriptor for each multicast operand, and mask multicast and barrier counts at runtime. If the scheduler is dynamic, group fallback clusters in whole multiples of the split factor.

1. The wider K, and why it removes sm_103a’s circular descriptors#

1.1 Why sm_103a’s K = 96 was awkward#