Skip to content

[BUG] cutlass_library does not emit grouped NVFP4 (ue4m3) block-scaled GEMM for SM120, although example 79d runs that kernel #3343

Description

@waynehacking8

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.

Metadata

Metadata

Assignees

No one assigned

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions