[draft] NVFP4 block-16 scale support for SM90 mixed-input grouped GEMM

[changjonathanc]

Summary

Adds optional NVFP4 (e2m1 data + e4m3 scale) block-16 scaling to the SM90 warp-specialized mixed-input collectives, so weight-scaled W4A8 grouped GEMM works when the scale block size (16) is smaller than the GMMA K tile. +69/-12 across two headers; all new behaviour is gated by UseNvfp4Block16Scales/Broadcast, which is false for every pre-existing instantiation — so existing kernels are byte-for-byte unaffected.

Currently this patch is carried out-of-tree in Forge as full-file cutlass_overrides/ copies (shadowing a second stock CUTLASS 4.4.2 submodule). That copies ~2,700 lines of stock CUTLASS to deliver these ~85 patched lines. Landing it in the fork lets Colonels drop both the overrides and the second submodule and build against this fork alone (matching how the existing _C extension already builds).

Changes

sm90_mma_array_tma_gmma_rs_warpspecialized_mixed_input.hpp

• UseNvfp4Block16Scales / ScaleAtomM / ScaleAtomK gate the new path (e2m1 A, e4m3 scale, TileK % 16 == 0); SmemLayoutAtomScale gains multiple scale columns per K tile in that case.
• Relax the size<1>(SmemLayoutAtomScale)==1 static_assert for that path only.
• Clamp grouped-GEMM init M/N/K to the tile shape; set an explicit StrideScale.
• Relax the can_implement chunk-size check to also allow TileK % chunk_size == 0 for NVFP4.

mixed_input_utils.hpp

• UseNvfp4Block16ScaleBroadcast + get_mma_smem_layout_scale() build a broadcast MMA view over the compact scale columns (stride-0 across the 16-row atom).
• Refresh scales every k-block (not only k_block==0) when there are multiple scale columns; route scale/zero smem tensors through the broadcast layout.

Validation

Built against this branch (CUTLASS 4.3.5 base) on H200 with the Colonels W4A8 MoE extension (no cutlass_overrides, no second submodule). On a real NVFP4 laguna-xs checkpoint, W4A8 prompt logprobs are bit-identical to the previous build that used stock 4.4.2 + the override copies (mean|Δ|=0, max|Δ|=0, Pearson=1.0), and large-batch MoE speedup is preserved (1.3–1.7× vs Marlin at M=16k–32k). So 4.3.5 is sufficient — no CUTLASS version bump required.

🤖 Generated with Claude Code

[changjonathanc]

Superseded by the 4.4.2-based PR: building W4A8 against the 4.3.5 base is a ~5.3x perf regression vs 4.4.2 (measured on H200). Redone on a v4.4.2 base.