[CuTeDSL] Add SM120 MXF4/NVFP4 native-TMA path

[alecco]

Using this PR code, QuACK can match examples/79_blackwell_geforce_gemm/79a_blackwell_geforce_nvfp4_bf16_gemm.cu at 95% consistently. Dao-AILab/quack#145

Summary

This PR adds the CuTe DSL plumbing for the SM120 MXF4/NVFP4 native-TMA path:

• add SM120 MXF4/NVFP4 warp MMA lowering
• teach cute.gemm to lower explicit (operand_fragment, scale_fragment) bundles for the narrow SM120 MXF4/NVFP4 MMA case
• add SM120 E4M3 FP8 scale-fragment helpers
• add position-independent swizzle tensor support for SMEM copy partitioning
• add an already-elected TMA producer acquire helper
• add cutlass.utils.gemm.sm120 native-TMA layout helpers for MXF4/NVFP4
• add a fixed native-TMA microtile example and focused SM120 tests

The intent is to match the SM120 block-scaled structure rather than force SM120 through the SM100 tcgen05 / TMEM model. SM120 uses native A/B TMA, native FP8 scale TMA, register A/B fragments, explicit FP8 scale fragments, and warp-level MXF4/NVFP4 MMA.

Details

Warp MMA lowering

Adds the SM120 MXF4/NVFP4 MMA path for:

• A/B: Float4E2M1FN
• scales: Float8E4M3FN
• accumulator: Float32
• instruction shape: m16n8k64
• scale vector size: 16

The cute.gemm integration is intentionally narrow. It handles only the SM120 MXF4/NVFP4 (fragment, scale) bundle case and otherwise leaves the existing variadic GEMM path unchanged.

Native TMA helpers

Adds SM120 MXF4/NVFP4 helper utilities for:

• logical A/B GMEM layouts
• interleaved scale GMEM layouts
• native A/B TMA atoms
• native FP8 scale TMA atoms
• packed and unpack A/B SMEM view construction
• FP8 scale SMEM view construction
• format-aware TMA transaction byte helpers
• the 128x128x128 tiled MMA layout used by the SM120 path

The default A/B SMEM format is the packed path. The unpack path remains available explicitly and has compile coverage.

Pipeline / copy support

Adds:

• as_position_independent_swizzle_tensor(...) so copy partitioning can use a non-swizzled logical layout while keeping the swizzle on the pointer
• producer_acquire_already_elected(...) for TMA producer code that has already entered an elected-lane region

These are needed by the SM120 native-TMA path and also keep the helper layering close to the existing Blackwell/CuTe DSL copy pipeline style.

Example

Adds a minimal SM120 MXF4/NVFP4 native-TMA microtile example.

The example intentionally computes one fixed 16x8 BF16 output microtile from a 128x128x128 native-TMA tile. It is not a full GEMM kernel. It demonstrates the native-TMA plumbing:

• build A/B native TMA atoms
• build SFA/SFB native FP8 scale TMA atoms
• issue one A and one B 3D TMA load, plus one SFA and one SFB 2D scale TMA load
• execute two K64 MXF4/NVFP4 MMA instructions
• store one BF16 output microtile

Tests / coverage

This PR adds SM120 tests covering:

• direct mma_mxf4nvf4(...) execution
• bundled cute.gemm(...) parity
• full-K K64 accumulation
• distinct nonzero C accumulation across K64 halves
• SFA/SFB scale-fragment mapping
• per-scale-column behavior
• row-varying SFA behavior
• negative dtype/layout validation
• plain F16/BF16 cute.gemm compile coverage to ensure non-bundled GEMM dispatch is unaffected
• native TMA layout helper validation
• packed/unpack byte helper behavior
• unpack TMA atom compile coverage
• native-TMA microtile PTX and runtime result checks

The microtile test checks for the expected instruction shape:

• two MXF4/NVFP4 MMA instructions
• two A/B 3D TMA ops
• two scale 2D TMA ops
• expected BF16 output value for the fixed uniform-scale setup

Notes

SM120 is intentionally handled differently from SM100. There is no tcgen05 / TMEM path here. The implementation is centered on native TMA, register fragments, explicit FP8 scale fragments, and warp-level MXF4/NVFP4 MMA.

Env

# Install the CUDA 13 CuTe DSL runtime wheel set matching this branch.
# This file currently contains: nvidia-cutlass-dsl[cu13]==4.5.1
python -m pip install -r python/CuTeDSL/requirements-cu13.txt

# Prepare the editable CuTe DSL source tree from the matching wheel.
# This copies generated _mlir Python/extension files and writes VERSION.EDITABLE.
python python/CuTeDSL/prep_editable_install.py

# Install CuTe DSL in editable/dev mode.
python -m pip install -e python/CuTeDSL

Expected verification:

python -m pip show nvidia-cutlass-dsl nvidia-cutlass-dsl-libs-base nvidia-cutlass-dsl-libs-cu13 cuda-python

Local working environment currently shows:

nvidia-cutlass-dsl            4.5.1.dev0, editable: ${HOME}/Soft/cutlass/python/CuTeDSL
nvidia-cutlass-dsl-libs-base  4.5.1
nvidia-cutlass-dsl-libs-cu13  4.5.1
cuda-python                   13.2.0
cuda-toolkit                  13.0.2
torch                         2.11.0
triton                        3.6.0

For running SM120 tests/benchmarks:

export CUTE_DSL_ARCH=sm_120a
export CUTE_DSL_CACHE_DIR=/data/agent/CuTeDSL/cache

If the runtime library is not found automatically, set:

export CUTE_DSL_LIBS=/path/to/site-packages/nvidia_cutlass_dsl/lib/libcute_dsl_runtime.so

On my local install that path is:

${HOME}/.local/lib/python3.14/site-packages/nvidia_cutlass_dsl/lib/libcute_dsl_runtime.so