Description
GenerateSM120_TensorOp_fp4_UMMA_gemm_with_block_scaled (python/cutlass_library/generator.py, ~L11441) skips grouped NVFP4 (ue4m3) block-scaled GEMM tile descriptions for SM120, behind this comment:
# grouped schedules only support ue8m0 (MXF4); NVF4 (ue4m3) grouped requires
# NVF4-specific PtrArray schedule tags not yet available
That comment appears to be stale:
to_grouped_schedule() (python/cutlass_library/library.py, L1051-1052) already maps
Nvf4TmaWarpSpecialized{Cooperative,Pingpong}Sm120 → PtrArrayTmaWarpSpecialized{Cooperative,Pingpong}.
- Example
examples/79_blackwell_geforce_gemm/79d_blackwell_geforce_nvfp4_grouped_gemm.cu instantiates exactly this kernel — ArchTag = cutlass::arch::Sm120, e2m1 operands, ue4m3 SF, KernelPtrArrayTmaWarpSpecialized{Cooperative,Pingpong} — and runs correctly.
Because the grouped schedules collapse to PtrArray*, is_nvf4(kernel_schedule) is False in the grouped branch, so only the ue8m0 (MXF4) grouped tiles are emitted and the NVFP4 grouped tiles are dropped.
Effect
The library / profiler / autotuner can select grouped MXFP4 on SM120 but not grouped NVFP4, so frameworks running NVFP4 MoE on consumer Blackwell (RTX PRO 6000 / RTX 5090) don't get these grouped kernels and fall back to slower paths (related: #2800).
Evidence (RTX PRO 6000, SM120, CUDA 13.0)
- Built + ran example 79d: both schedules
Disposition: Passed, 303 / 253 TFLOPS.
- Driving the grouped generator path directly (
GenerateSM120_TensorOp_fp4_UMMA_gemm_with_block_scaled(..., gemm_kind=GemmKind.GroupedBlockScaledUniversal3x)):
- today: 14 grouped fp4 ops, 0 NVFP4 (ue4m3) grouped;
- after adding an
is_grouped && ue4m3 clause: 28 ops, 14 NVFP4 grouped (e.g. cutlass3x_sm120_bstensorop_gemm_grouped_ue4m3xe2m1_ue4m3xe2m1_f32_void_f32_128x32x128_1x1x1_0_tnt_align32_warpspecialized_pingpong).
Proposed fix
In the grouped branch, key on the SF type instead of is_nvf4(kernel_schedule) (which only matches the non-grouped Sm120 tags):
if (is_grouped_schedule and math_inst.element_scale_factor == DataType.ue8m0) or \
+ (is_grouped_schedule and math_inst.element_scale_factor == DataType.ue4m3) or \
(not is_grouped_schedule and math_inst.element_scale_factor == DataType.ue4m3 and is_nvf4(kernel_schedule)) or \
(not is_grouped_schedule and math_inst.element_scale_factor == DataType.ue8m0 and not is_nvf4(kernel_schedule)):
Happy to send this as a PR (with the stale comment corrected). I verified emission (0 → 14) and the kernel template via example 79d, but did not run a full cutlass_profiler build of all emitted tile variants locally — worth confirming in CI that every tile compiles (some SM120 tiles may hit Stages≥2 constraints; if so they can be pruned in the same clause).
cc @depaulmillz — this is the SM12x grouped / block-scaled generator area.
Description
GenerateSM120_TensorOp_fp4_UMMA_gemm_with_block_scaled(python/cutlass_library/generator.py, ~L11441) skips grouped NVFP4 (ue4m3) block-scaled GEMM tile descriptions for SM120, behind this comment:That comment appears to be stale:
to_grouped_schedule()(python/cutlass_library/library.py, L1051-1052) already mapsNvf4TmaWarpSpecialized{Cooperative,Pingpong}Sm120 → PtrArrayTmaWarpSpecialized{Cooperative,Pingpong}.examples/79_blackwell_geforce_gemm/79d_blackwell_geforce_nvfp4_grouped_gemm.cuinstantiates exactly this kernel —ArchTag = cutlass::arch::Sm120,e2m1operands,ue4m3SF,KernelPtrArrayTmaWarpSpecialized{Cooperative,Pingpong}— and runs correctly.Because the grouped schedules collapse to
PtrArray*,is_nvf4(kernel_schedule)isFalsein the grouped branch, so only theue8m0(MXF4) grouped tiles are emitted and the NVFP4 grouped tiles are dropped.Effect
The library / profiler / autotuner can select grouped MXFP4 on SM120 but not grouped NVFP4, so frameworks running NVFP4 MoE on consumer Blackwell (RTX PRO 6000 / RTX 5090) don't get these grouped kernels and fall back to slower paths (related: #2800).
Evidence (RTX PRO 6000, SM120, CUDA 13.0)
Disposition: Passed, 303 / 253 TFLOPS.GenerateSM120_TensorOp_fp4_UMMA_gemm_with_block_scaled(..., gemm_kind=GemmKind.GroupedBlockScaledUniversal3x)):is_grouped && ue4m3clause: 28 ops, 14 NVFP4 grouped (e.g.cutlass3x_sm120_bstensorop_gemm_grouped_ue4m3xe2m1_ue4m3xe2m1_f32_void_f32_128x32x128_1x1x1_0_tnt_align32_warpspecialized_pingpong).Proposed fix
In the grouped branch, key on the SF type instead of
is_nvf4(kernel_schedule)(which only matches the non-grouped Sm120 tags):if (is_grouped_schedule and math_inst.element_scale_factor == DataType.ue8m0) or \ + (is_grouped_schedule and math_inst.element_scale_factor == DataType.ue4m3) or \ (not is_grouped_schedule and math_inst.element_scale_factor == DataType.ue4m3 and is_nvf4(kernel_schedule)) or \ (not is_grouped_schedule and math_inst.element_scale_factor == DataType.ue8m0 and not is_nvf4(kernel_schedule)):Happy to send this as a PR (with the stale comment corrected). I verified emission (0 → 14) and the kernel template via example 79d, but did not run a full
cutlass_profilerbuild of all emitted tile variants locally — worth confirming in CI that every tile compiles (some SM120 tiles may hit Stages≥2 constraints; if so they can be pruned in the same clause).cc @depaulmillz — this is the SM12x grouped / block-scaled generator area.