[Common] Group NVFP4 Quantize Kernels

[Oleg-Goncharov]

Description

This PR adds optimized grouped NVFP4 1D quantization support for BF16 grouped tensors.

The implementation enables grouped and MoE-style workloads to quantize multiple tensors with different shapes without launching the existing NVFP4 kernel separately for each tensor. It supports rowwise NVFP4 output together with an optional transposed columnwise output and the corresponding scaling factors.

All four grouped tensor shape representations are supported:

• SAME_BOTH_DIMS
• VARYING_FIRST_DIM
• VARYING_LAST_DIM
• VARYING_BOTH_DIMS

Type of change

• Documentation change (change only to the documentation, either a fix or a new content)
• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Infra/Build change
• Code refactoring

Changes

• Added grouped NVFP4 1D quantization dispatch through nvte_group_quantize.
• Added a tuned grouped NVFP4 kernel for BF16 input and FP4 E2M1 output.
• Added fused generation of rowwise output and optional transposed columnwise output.
• Added support for per-tensor rowwise and columnwise amax values and compact E4M3 scaling factors.
• Added scheduling and work mapping for all four ShapeRepresentation layouts.
• Added per-tensor metadata for dimensions, data offsets, and rowwise/columnwise scale offsets.
• Extended grouped TMA descriptor handling to support bit-addressed 4-bit output data and transposed tensor layouts.
• Added handling for zero-sized tensors without creating invalid TMA descriptors.
• Refactored rowwise, columnwise, and auxiliary NVFP4 scaling routines into a shared implementation used by both grouped and non-grouped kernels.
• Added DefaultCastConfig and shape-specific CastConfig specializations, allowing kernel parameters to be tuned independently for each grouped layout.
• Added C++ operator tests covering different tensor counts, shapes, and all grouped shape representations.

Checklist:

• I have read and followed the contributing guidelines
• The functionality is complete
• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes

[greptile-apps]

Greptile Summary

This PR adds grouped BF16-to-NVFP4 quantization with optional transposed output and shared scaling logic.

• Adds grouped dispatch and tuned kernels for all four grouped shape representations.
• Extends grouped layouts and TMA descriptors for FP4 output and transposed tensors.
• Adds grouped operator coverage and restores the existing quantization-mode dispatch branches.

Confidence Score: 5/5

The PR appears safe to merge because the previously reported quantization dispatch regression is fixed at current HEAD.

No blocking failure remains.

Important Files Changed

Filename	Overview
transformer_engine/common/cast/dispatch/quantize.cuh	Adds grouped NVFP4 dispatch while retaining active delayed-scaling, MXFP8, and block-scaling branches.
transformer_engine/common/cast/nvfp4/group_quantize_transpose_nvfp4.cuh	Introduces the grouped NVFP4 quantize-and-transpose dispatch path.
transformer_engine/common/cast/nvfp4/specialized/group_quantize_transpose_nvfp4_tuned_1D.cuh	Implements tuned shape-specific grouped scheduling and fused rowwise/columnwise output generation.
transformer_engine/common/cast/core/grouped_layout.cuh	Extends grouped tensor metadata and layout handling needed by varying-shape NVFP4 workloads.
transformer_engine/common/cast/core/grouped_tma.cuh	Extends grouped TMA descriptor handling for bit-addressed FP4 and transposed layouts.
tests/cpp/operator/test_cast_nvfp4_transpose_grouped.cu	Adds numerical operator coverage for grouped NVFP4 quantization across the supported shape representations.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart LR
  API[nvte_group_quantize] --> Dispatch[Scaling-mode dispatch]
  Dispatch --> NVFP4[Grouped NVFP4 kernel]
  NVFP4 --> Layout[Shape-specific work mapping]
  Layout --> Row[Rowwise FP4 output and scales]
  Layout --> Col[Optional transposed FP4 output and scales]

Loading

_{Reviews (4): Last reviewed commit: "Merge branch 'main' into pr_nvfp4_group_..." | Re-trigger Greptile}

[Oleg-Goncharov]

/te-ci

[Oleg-Goncharov]

/te-ci