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.

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 two major features: (1) a new FP8 current-scaling grouped-quantization path for Float8CurrentScalingQuantizer, including per-group amax computation (nvte_group_compute_amax_with_config), scale/scale-inv derivation (nvte_group_compute_scale_from_amax), and a graph-safe no-op flag; and (2) support for varying last-dim (and both-dims) grouped tensors in existing group-quantize kernels, extending the shape-representation enum to VARYING_LAST_DIM / VARYING_BOTH_DIMS and adding a new nvte_splits_to_offsets_2d helper.

• Adds group_amax_fp8.cuh with a load-balanced flat-scan amax kernel and a per-group scale-update kernel; both honour the no-op flag so CUDA-graph weight caching works correctly.
• Extends group_quantize_fp8.cuh with a 2D-grid dispatch for VARYING_LAST_DIM (flat rowwise fast-path + generic colwise/bidimensional fall-through) and adds the NVTE_DELAYED_TENSOR_SCALING case to the grouped-quantize dispatcher, wiring it to the new tensor-scaling kernel.
• Refactors make_grouped_tensor in Python to compute logical dims and per-tensor offsets for all four shape-representation cases, replacing the old assumption that last dim is always uniform; CUDA-graph-capture guard is unified and generalised.

Confidence Score: 5/5

Safe to merge; the end-to-end FP8 current-scaling grouped-quantization pipeline and the varying-last-dim dispatch are logically sound and correctly guarded against unsupported combinations.

The core additions — flat amax scan, scale derivation, 2D-grid group-quantize dispatch for VARYING_LAST_DIM, and the noop-flag plumbing — are all correctly wired together. Previously reported bugs (positional arg shift in grouped_mlp.py, noop_flag forwarding, error message accuracy) have been addressed. The remaining findings are documentation and minor interface-clarity issues that do not affect correctness.

group_amax_fp8.cuh (unused-parameter API surface) and grouped_tensor_storage.py (late failure for float8_current_scaling + VARYING_BOTH via Python path).

Important Files Changed

Filename	Overview
transformer_engine/common/cast/fp8/group_amax_fp8.cuh	New file implementing per-group amax and scale-update kernels for FP8 current scaling; flat-scan approach is correct but shared-memory allocation grows O(num_tensors) without an upper-bound check.
transformer_engine/common/cast/fp8/group_quantize_fp8.cuh	New file with FP8 per-tensor-scale group-quantize kernels; VARYING_LAST_DIM rowwise fast-path added; VARYING_BOTH_DIMS correctly rejected with NVTE_ERROR; COLWISE/BIDIMENSIONAL fall-through to generic 2D-grid kernel is logically sound.
transformer_engine/pytorch/csrc/extensions/cast.cpp	Adds grouped FP8 current-scaling amax + scale pipeline; guards for unsupported NVFP4+last_dims and VARYING_BOTH for current-scaling are present; allreduce and noop-flag wiring looks correct.
transformer_engine/pytorch/csrc/quantizer.cpp	All six quantizer subclasses updated to accept last_dims; build_grouped_tensor_offsets correctly handles all four dim-variation cases; Float8CurrentScalingQuantizer correctly shares a single scale_inv buffer for bidimensional on Blackwell.
transformer_engine/pytorch/tensor/storage/grouped_tensor_storage.py	make_grouped_tensor correctly handles all four shape-representation cases; CUDA-graph guard now unified and message is generic; missing early guard for float8_current_scaling + VARYING_BOTH_DIMS in Python path.
transformer_engine/common/util/splits_to_offsets.cu	New splits_to_offsets_2d_kernel is logically correct (inclusive prefix-scan with chunk_prefix carry); single-block launch limits throughput for very large num_tensors but is fine for typical MoE group counts.
transformer_engine/common/cast/core/grouped_layout.cuh	New file; decode_job / decode_block / find_tensor_from_offsets correctly handle all four ShapeRepresentation cases via flat element offsets; the 2D-grid excess-block pattern for VARYING_LAST_DIM is handled by per-block bounds checks.
transformer_engine/common/recipe/current_scaling.cu	group_compute_amax_impl correctly detects shape_rep from tensor metadata, gates offsets checks per case, and picks the right amax buffer; all four shape reps are handled.
transformer_engine/pytorch/csrc/extensions/pybind.cpp	last_dims inserted as keyword arg with default None; no positional-shift regressions remain.
transformer_engine/pytorch/ops/fused/grouped_mlp.py	Call-site updated to use keyword arguments for rowwise_amax and columnwise_amax, fixing the positional-argument shift introduced by the new last_dims parameter.

Sequence Diagram

sequenceDiagram
    participant PY as Python caller
    participant C as group_quantize (C++)
    participant CS as Float8CurrentScaling<br/>Quantizer::create_grouped_tensor
    participant AM as nvte_group_compute_amax<br/>_with_config
    participant SC as nvte_group_compute_scale<br/>_from_amax
    participant QK as nvte_group_quantize<br/>(group_quantize_fp8.cuh)

    PY->>C: "group_quantize(tensor, quantizer,<br/>num_tensors, first_dims, last_dims, noop_flag)"
    C->>CS: create_grouped_tensor(first_dims, last_dims, ...)
    CS-->>C: "grouped_output_cpp + grouped_output_py<br/>(amax/scale/scale_inv/data buffers)"
    C->>AM: "nvte_group_compute_amax_with_config<br/>(input, output, config, stream)"
    Note over AM: zero amax buffer,<br/>flat-scan over elements,<br/>atomicMax per tensor
    AM-->>C: amax[0..num_tensors-1] filled
    opt with_amax_reduction
        C->>C: NCCL AllReduce MAX on amax buffer
    end
    C->>SC: "nvte_group_compute_scale_from_amax<br/>(output, config, stream)"
    Note over SC: scale[i] = fp8_max / amax[i]<br/>scale_inv[i] = 1 / scale[i]
    SC-->>C: scale / scale_inv filled
    C->>QK: "nvte_group_quantize(input, output,<br/>quant_config, stream)"
    Note over QK: dispatch on shape_rep<br/>(SAME/VARYING_FIRST/VARYING_LAST)<br/>read scale_ptr[tensor_id] per element
    QK-->>C: FP8 quantized data written
    C-->>PY: grouped_output_py

Loading

_{Reviews (7): Last reviewed commit: "fix comment and add nvte check" | 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?