[Common][PyTorch]Current Scaling Group Quantization + Enabling Varying Last/Both Dims in Group Quantize

[vthumbe1503]

Description

• Current Scaling Group Quantization added in Common/Pytorch for Varying First/Last Dims
• Enables Varying Last/Both Dims for Group Quantize from Pytorch C++/Python Infrastructure(Wasnt present before for MXFP8 group quantize)
• Enables No_op flag from group_quantize Pytorch infra.(This was present previously only in Pytorch single tensor quantization infra)
• Benchmarking Scripts/Tests to benchmark Group Quantize Current Scaling for corner cases like overallocation, non-uniform group sizes etc. Tests are added for both Common(CPP) and Pytorch(Python)

Performance from Benchmarking Script for Current Scaling Group Quantize.
NOTE:

1. overalloc means the grouped tensor is overallocated 4 times the actual size. But this doesnt matter to the kernel since kernel reads the offsets from the device and decides the CTA to elements mapping.
2. zipf, heavy, mild etc represent the amount of group imbalance. As seen below, the kernel doesnt deteriorate in performance if group imbalance is present in Blackwell. For Hopper, group_imbalance affects the perf of columnwise cast kernel at the moment due to the use of 2D grid. This will be fixed in a followup PR.

GB200

H100

Fixes # (issue)

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

Please list the changes introduced in this PR:

• Core Changes

  • common.cuh has been seperated out into two files -- grouped_layout.cuh and grouped_tma.cuh. The idea is that, grouped_layout.cuh has generic utilities commonly used in grouped tensor kernels. And grouped_tma.cuh has the arch specific TMA changes. This is done so that current scaling can still be non-arch specific and it can use just the grouped_layout.cuh since we didnt need TMA for current scaling.
  • Dispatch has the change to enable calling group quantize for NVTE_DELAYED_SCALING which is the enum used for both delayed scaling and current scaling in the legacy infra(Note: Delayed Scaling hasnt enabled in the Pytorch Infra in this PR)
  • FP8 : Kernels for Computing Grouped Amax and grouped Cast to Per tensor FP8 are added. Both hit close to 6TB/s to all use-cases on Blackwell and 3 TB/s for H100 rowwise cast kernel and 2.5 TB/s for H100 columnwise cast kernel.

• Other Common Changes

  • NVTE APIs for computing grouped amax and computing grouped scale from grouped amax added which uses the kernels implemented in core.
  • nvte_splits_to_offsets_2d --> NVTE API for computing for computing offsets for VARYING_BOTH_DIMS added. Since we now enable enabling varying last/both dims from pytorch infra.

• Pytorch Group Quantize API Changes

  • Current Scaling Support added that uses the NVTE APIs from common for computing grouped amax and performing a grouped cast
  • Accepts last_dims in the API that enables Varying Last Dims/Both Dims in the Pytorch infrastructure
  • Accepts no_op flag for cuda graph when wanting skip_fp8_weight_update. (Would be needed for single_grouped_weight correctness with cuda graph when skip_weight_update is enabled --> essentially when we use group_quantize on the weights)

• Benchmarking Scripts/Tests added for Current Scaling for Varying First/Last Dims and to handle overallocation of grouped tensors. Varying All Dims is not supported yet.

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

[vthumbe1503]

/te-ci pytorch

[greptile-apps]

Greptile Summary

This PR adds current-scaling group quantization (per-group amax → scale → FP8 cast) for both rowwise and columnwise layouts, and extends the existing group-quantize infrastructure in PyTorch and Common to support tensors with a varying last dimension (VARYING_LAST_DIM) and varying both dimensions (VARYING_BOTH_DIMS), in addition to the already-supported VARYING_FIRST_DIM case.

• New CUDA kernels (group_amax_fp8.cuh, group_quantize_fp8.cuh): flat vectorized per-group amax, scale derivation, and FP8 cast with shape-representation dispatch; Blackwell fast-path uses hardware-accelerated 4-element scaled FP8 conversion.
• PyTorch / C++ plumbing: last_dims parameter added throughout create_grouped_tensor overloads, build_grouped_tensor_offsets handles 1D/2D split-to-offset cases, and splits_to_offsets_2d_kernel computes cumulative element offsets for VARYING_BOTH_DIMS on device.
• Shared scale_inv buffer for FP8 current scaling: on Blackwell (is_non_tn_fp8_gemm_supported), a single scale_inv tensor is aliased for both rowwise and (suppressed) columnwise paths, matching the C++ grouped tensor wrapper behavior.

Confidence Score: 4/5

The current-scaling FP8 path is well-guarded and the new kernels handle all four shape representations correctly; the one unguarded combination (MXFP8 + last_dims) would silently produce wrong quantized output if triggered.

All new code paths for FP8 current scaling, the new CUDA kernels, the 2D offset kernel, and the Python storage changes are internally consistent. NVFP4 and bgrad_group_quantize both explicitly reject a non-None last_dims, but the MXFP8 dispatch in group_quantize carries no equivalent guard. Passing last_dims with an MXFP8 quantizer causes create_grouped_tensor to size the scale buffer using the total (summed) logical last dim rather than per-tensor dims, then the MXFP8 cast kernel operates on the mismatched scale layout without error.

transformer_engine/pytorch/csrc/extensions/cast.cpp — the MXFP8_GROUPED_QUANTIZE switch case needs a guard matching the one already present for NVFP4.

Important Files Changed

Filename	Overview
transformer_engine/pytorch/csrc/extensions/cast.cpp	Adds current-scaling group-quantize dispatch, noop_flag forwarding, and last_dims plumbing; NVFP4/bgrad cases guard against last_dims but MXFP8 case is missing the equivalent guard, allowing wrong scale buffer sizing to silently propagate.
transformer_engine/pytorch/csrc/quantizer.cpp	Adds last_dims parameter to all create_grouped_tensor overloads, updates build_grouped_tensor_offsets for 1D/2D varying cases, and implements shared scale_inv for FP8 current scaling on Blackwell; logic is sound.
transformer_engine/pytorch/tensor/storage/grouped_tensor_storage.py	Extends make_grouped_tensor to handle VARYING_LAST_DIM and VARYING_BOTH_DIMS, unifies the graph-capture guard, and adds shared scale_inv for FP8 current scaling; the logical_shape conventions for varying dims are correct.
transformer_engine/common/cast/fp8/group_quantize_fp8.cuh	New file implementing per-group FP8 cast kernels for all four shape representations; fast-path specializations for rowwise-flat and same-shape transpose; VARYING_LAST_DIM handled correctly via per-tensor col bounds check in the 2D grid.
transformer_engine/common/cast/fp8/group_amax_fp8.cuh	New file with flat vectorized per-group amax kernel and scale/scale_inv derivation kernel; handles all shape representations via device-side offset arrays; shared-memory usage scales with num_tensors.
transformer_engine/common/util/splits_to_offsets.cu	Adds splits_to_offsets_2d_kernel for VARYING_BOTH_DIMS element-offset computation; single-block chunked inclusive scan with correct chunk_prefix bookkeeping across arbitrary num_tensors.
transformer_engine/pytorch/ops/fused/grouped_mlp.py	Converts rowwise_amax and columnwise_amax to keyword arguments in the nvfp4_group_quantize_with_amax call site to accommodate the new last_dims positional parameter.
transformer_engine/common/recipe/current_scaling.cu	Adds nvte_group_compute_amax_with_config API; shape detection logic matches the four representation cases correctly.
transformer_engine/pytorch/csrc/extensions/pybind.cpp	Adds last_dims (default None) and noop_flag (default None) to group_quantize binding, and last_dims to nvfp4_group_quantize_with_amax and bgrad_group_quantize bindings.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A["group_quantize(tensor, quantizer, num_tensors,\nfirst_dims, last_dims, tensor_offsets, noop_flag)"] --> B{Quantizer type?}
    B -->|NVFP4| C["NVTE_CHECK: !last_dims ✓\nnvfp4 group quantize impl"]
    B -->|MXFP8| D["MXFP8_GROUPED_QUANTIZE\n❌ no last_dims guard\nnvte_group_quantize(MXFP8)"]
    B -->|Float8CurrentScaling| E["compute_grouped_fp8_current_scaling_\namax_and_scale()"]
    B -->|Other| F["NVTE_ERROR"]
    E --> E1["nvte_group_compute_amax_with_config\n(zero + flat vectorized amax kernel)"]
    E1 --> E2{"with_amax_reduction?"}
    E2 -->|Yes| E3["NCCL allreduce MAX on amax buffer"]
    E2 -->|No| E4
    E3 --> E4["nvte_group_compute_scale_from_amax\n(scale = fp8_max / amax)"]
    E4 --> E5["nvte_group_quantize\n(FP8 cast kernel)"]
    subgraph "Shape Representations"
        SR1["SAME_BOTH_DIMS\nflat rowwise fast path"]
        SR2["VARYING_FIRST_DIM\naligned flat rowwise / generic 2D colwise"]
        SR3["VARYING_LAST_DIM\nper-row flat rowwise / generic 2D colwise"]
        SR4["VARYING_BOTH_DIMS\n(blocked in current scaling)"]
    end
    E5 --> SR1 & SR2 & SR3 & SR4

Loading

%%{init: {'theme': 'base', 'themeVariables': {"darkMode": true, "background": "#0d1117", "primaryColor": "#21262d", "primaryTextColor": "#e6edf3", "primaryBorderColor": "#8b949e", "lineColor": "#8b949e", "textColor": "#e6edf3", "edgeLabelBackground": "#161b22", "actorBkg": "#21262d", "actorBorder": "#8b949e", "actorTextColor": "#e6edf3", "actorLineColor": "#8b949e", "signalColor": "#8b949e", "signalTextColor": "#e6edf3", "noteBkgColor": "#373320", "noteBorderColor": "#d4a72c", "noteTextColor": "#f0e6c0", "labelBoxBkgColor": "#21262d", "labelBoxBorderColor": "#8b949e", "labelTextColor": "#e6edf3", "loopTextColor": "#e6edf3", "activationBkgColor": "#30363d", "activationBorderColor": "#8b949e"}}}%%
flowchart TD
    A["group_quantize(tensor, quantizer, num_tensors,\nfirst_dims, last_dims, tensor_offsets, noop_flag)"] --> B{Quantizer type?}
    B -->|NVFP4| C["NVTE_CHECK: !last_dims ✓\nnvfp4 group quantize impl"]
    B -->|MXFP8| D["MXFP8_GROUPED_QUANTIZE\n❌ no last_dims guard\nnvte_group_quantize(MXFP8)"]
    B -->|Float8CurrentScaling| E["compute_grouped_fp8_current_scaling_\namax_and_scale()"]
    B -->|Other| F["NVTE_ERROR"]
    E --> E1["nvte_group_compute_amax_with_config\n(zero + flat vectorized amax kernel)"]
    E1 --> E2{"with_amax_reduction?"}
    E2 -->|Yes| E3["NCCL allreduce MAX on amax buffer"]
    E2 -->|No| E4
    E3 --> E4["nvte_group_compute_scale_from_amax\n(scale = fp8_max / amax)"]
    E4 --> E5["nvte_group_quantize\n(FP8 cast kernel)"]
    subgraph "Shape Representations"
        SR1["SAME_BOTH_DIMS\nflat rowwise fast path"]
        SR2["VARYING_FIRST_DIM\naligned flat rowwise / generic 2D colwise"]
        SR3["VARYING_LAST_DIM\nper-row flat rowwise / generic 2D colwise"]
        SR4["VARYING_BOTH_DIMS\n(blocked in current scaling)"]
    end
    E5 --> SR1 & SR2 & SR3 & SR4

Loading

_{Reviews (9): Last reviewed commit: "add CPP tests for current scaling group ..." | Re-trigger Greptile}

[vthumbe1503]

/te-ci

[vthumbe1503]

Pipeline: 54747206

[vthumbe1503]

/te-ci

[vthumbe1503]

/te-ci

[Oleg-Goncharov]

Hi @vthumbe1503, could you please also add C++ unit tests covering various shapes, similar to grouped MXFP8?

[vthumbe1503]

/te-ci

[greptile-apps]

MXFP8 case missing last_dims guard — produces silently wrong output

The NVFP4_GROUPED_QUANTIZE case correctly rejects a non-None last_dims at line 341. The MXFP8_GROUPED_QUANTIZE case has no equivalent check. If a caller passes last_dims with an MXFP8 quantizer, MXFP8Quantizer::create_grouped_tensor will have already been called (before the switch) with is_varying_both = false, causing it to compute the total rowwise/columnwise scale buffer size via get_scale_shape({logical_first_dim, sum(last_dims)}, …) — treating all tensors as having the same last dimension equal to the sum, instead of computing per-tensor scale counts. The MXFP8 kernel then reads scale offsets that were sized for the wrong shape, producing wrong quantized output with no error raised. bgrad_group_quantize documents this exact constraint ("MXFP8 dbias kernel requires a constant last dimension"), so the same guard should be applied here.