[BUG] cutlass_library does not emit grouped NVFP4 (ue4m3) block-scaled GEMM for SM120, although example 79d runs that kernel
Nobody has claimed this yet.
- Dominant language
- C++
- Stars
- 10.5k
- Forks
- 2.1k
- Avg merge
- 3d 11h
- Merged PRs (30d)
- 7
Description
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.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) 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 && 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_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.
Contributor guide
No contributing guide indexed for this repository
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start in python/cutlass_library/generator.py at GenerateSM120_TensorOp_fp4_UMMA_gemm_with_block_scaled and compare the grouped condition with to_grouped_schedule() in python/cutlass_library/library.py. Run the direct grouped generator path and example 79d_blackwell_geforce_nvfp4_grouped_gemm.cu. Done means grouped ue4m3 variants are emitted and the relevant SM120 tile variants compile and run.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- cpp, python
- Domain
- performance, tooling
- Issue type
- Bug
- Difficulty
- 3/5
- Estimated time
- 1-2 days
- Activity status
- Quiet
- Clarity
- Clearly specified
- Newbie friendliness
- 76/100