NVIDIA / NVIDIA/cutlass

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

Open
#3,343 1 comment 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

inactive-30d
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.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.

Contributor guide

No contributing guide indexed for this repository

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.