Porting Block-Scaled GEMM Kernels from Blackwell to Rubin

FlashInfer at commit 518f4c7f0 now includes dedicated block-scaled GEMM kernels for NVIDIA's Rubin architecture (sm_107a), documenting four hardware changes from Blackwell: a wider Tensor Core K (K = 128 for FP4 and K = 64 for FP8 and mixed FP8 × FP4, selected by instruction-descriptor bit 3), packed 4-bit operands in mixed precision that remove Blackwell's U4_UNPACK_U8 TMA unpacking, SMEM growth from 228 KiB to 328 KiB and TMEM growth from 512 columns (256 KiB) to 576 columns (288 KiB), and new quantization formats including an unsigned 8-bit UE5M3 scale-factor format for FP4. The write-up notes that a Blackwell kernel ported unchanged to sm_107a leaves performance unused, failing in two places: an FP8 instruction with M = 64 is rejected when K = 64, and a cp.async.cg cache hint can silently stop applying. The author states no Rubin hardware was available, so none of the porting claims are measured.

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 https://github.com/flashinfer-ai/flashinfer/tree/518f4c7f01b1b97fe57add94d46ddc1a80344e51 , 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: 1. 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 . 2. 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 . 3. 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 . 4. 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 1-the-wider-k-and-why-it-removes-sm 103as-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 11-why-sm 103as-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 12-how-64-bytes-per-instruction-fixes-it On sm 107a , every block-scaled instruction reads 64 bytes along K, which has three consequences: 1. 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. 2. 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. 3. 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 13-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 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 21-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 22-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 : python 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 23-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 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 31-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 while the epilogue drains tile , take 512 columns. - The scale factors take: 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 : php 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 32-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 until the MMA had finished reading K-tile . 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: 3.3 The 248 KiB limit for cp.async.cg destinations 33-the-248-kib-limit-for-cpasynccg-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 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 41-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: 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 42-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 43-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: 1. UTCCP of the keep SFA A0 for K-block 0, A2 for K-block 1 2. UTCCP of SFB B0, then B1 3. UTCCP of the reuse SFA A1, then A3 4. MMA0 keep 5. 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 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 51-setting-the-instructions-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 : python 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: 1. 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 , 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, , 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 :Here that is 32/2 − 128 = −112, so CTA 1 loads from +16. 2. 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 , the MMA expects all four in TMEM at indices 0 to 3. With , 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. 5.2 Separate pipeline depths for weights and tokens 52-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: 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 53-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: 1. Load. Each warp loads gate and up for tokens 0–31 and 32–63 into four 16-register FP32 arrays Ld16x64bOp.x16 . 2. 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. 3. The activation is 4. 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-new-quantization-formats-ue5m3-scales-and-a-scaled-fp8-cvt 6.1 UE5M3 scale factors for FP4 61-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 . 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 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 to about , and the 3 mantissa bits stay. 6.2 FP8 conversion with the scale applied: cvt.rn.satfinite.scaled::n1::ue8m0 62-fp8-conversion-with-the-scale-applied-cvtrnsatfinitescaledn1ue8m0 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: 1. Widen the ue8m0 byte back to BF16 and FP32. 2. Compute the reciprocal with rcp.approx.ftz.f32 on the MUFU. 3. Clamp infinity fmin ..., Fp32Max and mask zero fmin scale f32 1e30, 1.0 , so that 0 × ∞ doesn’t produce NaN. 4. Multiply all 32 values by the reciprocal. 5. 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 amounts to subtracting 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 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 71-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 72-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 73-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 : 1. 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. 2. 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 8-checklist-porting-an-sm 100a-or-sm 103a-kernel-to-sm 107a The items are mostly independent, so the order doesn’t matter. 1. 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. 2. 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 . 3. 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. 4. 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. 5. 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. 6. 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. 7. 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. 8. 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. 9. 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 . 10. 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.