Current Scaling Group Quantization + Enabling Varying Last/Both Dims in Group Quantize

[vthumbe1503]

Description

Please include a brief summary of the changes, relevant motivation and context.

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:

• Change A
• Change B

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 grouped FP8 current-scaling quantization (amax-on-the-fly, per-group scale/scale_inv) and extends the existing grouped-quantize infrastructure to support varying last dim and varying both dims, in addition to the previously supported varying first dim.

• New CUDA kernels: flat vectorized grouped-amax kernel (group_amax_fp8.cuh) with 4-way unrolled load + atomicMaxFloat, per-group scale/scale_inv update kernel, and a new grouped FP8 cast kernel (group_quantize_fp8.cuh) with a Blackwell (sm100+) fast path using mul_cvt_4x.
• Infrastructure: nvte_splits_to_offsets_2d for bidimensional offset computation; all create_grouped_tensor overloads extended with last_dims; Python make_grouped_tensor updated to encode varying-both as (1, total_elements) logical shape; CUDA-graph-capture guard unified across all varying-dim paths.
• C++ / Python integration: group_quantize and nvfp4_group_quantize_with_amax accept the new last_dims / noop_flag kwargs; grouped_mlp.py updated to pass amax tensors as keyword args following the signature change.

Confidence Score: 4/5

Safe to merge for the FP8 current-scaling and MXFP8 varying-dim paths; the NVFP4 code path needs a missing guard before last_dims can be safely exposed to callers.

nvfp4_group_quantize_with_amax now accepts a last_dims argument, but the NVFP4 kernel was not updated to handle varying last dim. Unlike bgrad_group_quantize (which immediately rejects a non-None last_dims), there is no equivalent guard here. A caller passing a non-None last_dims would silently propagate the metadata into the NVFP4 dispatch path, which only knows about varying-first-dim layouts and would produce incorrect output.

transformer_engine/pytorch/csrc/extensions/cast.cpp — specifically the nvfp4_group_quantize_with_amax function body, which is missing a NVTE_CHECK(!last_dims.has_value(), …) guard.

Important Files Changed

Filename	Overview
transformer_engine/pytorch/csrc/extensions/cast.cpp	Adds grouped FP8 current-scaling quantize path: amax compute → optional NCCL reduce → scale/scale_inv kernel, plus last_dims and noop_flag threading through group_quantize / nvfp4_group_quantize_with_amax / bgrad_group_quantize. bgrad_group_quantize correctly guards against last_dims.has_value(), but nvfp4_group_quantize_with_amax lacks an equivalent guard.
transformer_engine/pytorch/csrc/extensions/pybind.cpp	Inserts last_dims=py::none() before required rowwise_amax/columnwise_amax in nvfp4_group_quantize_with_amax binding; the caller in grouped_mlp.py was updated to use keyword arguments, resolving the positional-arg shift flagged in a previous review thread.
transformer_engine/pytorch/csrc/quantizer.cpp	Adds last_dims parameter to all create_grouped_tensor overloads; introduces build_grouped_tensor_offsets 2D path (first+last, first-only, last-only); correctly adds is_varying_both scale allocation for MXFP8; FP8CurrentScaling guard rejects varying-both-dims; no issues found.
transformer_engine/pytorch/tensor/storage/grouped_tensor_storage.py	Removes hard-coded last_dims=None assumption; adds VARYING_LAST_DIM and VARYING_BOTH_DIMS offset/shape computation; unified CUDA-graph-capture guard; allocates scale and amax buffers for current-scaling Python path. Logic for varying-both encoding as (1, total_elements) is consistent with C++ path.
transformer_engine/common/cast/fp8/group_amax_fp8.cuh	New file: flat vectorized grouped amax kernel with 4-way unrolled inner loop, shared-memory offsets, atomicMaxFloat reduction; plus per-group scale/scale_inv kernel. Both kernels honor the noop flag for CUDA-graph-safe weight skipping.
transformer_engine/common/cast/fp8/group_quantize_fp8.cuh	New file: grouped FP8 current-scaling quantize kernels for SAME/VARYING_FIRST/VARYING_LAST shape representations; fast path for Blackwell (sm100+) with 4x-vectorized CVT; 2D grid dispatch with per-block bounds checks for varying last dim.
transformer_engine/common/recipe/current_scaling.cu	Adds group_compute_amax_impl that dispatches the new flat grouped-amax kernel; nvte_group_compute_amax_with_config and nvte_group_compute_scale_from_amax public APIs implemented correctly.
transformer_engine/common/util/splits_to_offsets.cu	Adds splits_to_offsets_2d_kernel that computes cumsum(first_dims[i] * last_dims[i]); single-block with chunked inner loop for num_tensors > kThreadsPerBlock; nvte_splits_to_offsets_2d API wrapper with full input validation.
transformer_engine/common/cast/dispatch/quantize.cuh	Adds NVTE_DELAYED_TENSOR_SCALING case to group_quantize_fwd_helper, dispatching to the new fp8::group_quantize; only reached from the FP8 current-scaling path since delayed FP8 is rejected upstream.
transformer_engine/pytorch/ops/fused/grouped_mlp.py	Switches rowwise_amax and columnwise_amax to keyword arguments in the nvfp4_group_quantize_with_amax call, correctly addressing the positional-arg shift introduced by inserting last_dims before them.

Sequence Diagram

sequenceDiagram
    participant PY as Python (group_quantize)
    participant CPP as cast.cpp
    participant AMAX as nvte_group_compute_amax_with_config
    participant NCCL as NCCL allreduce (optional)
    participant SCALE as nvte_group_compute_scale_from_amax
    participant QUANT as nvte_group_quantize (FP8 cast kernel)

    PY->>CPP: group_quantize(tensor, Float8CurrentScalingQuantizer, first_dims, last_dims, noop_flag)
    CPP->>CPP: create_grouped_tensor (allocates data, amax, scale, scale_inv)
    CPP->>AMAX: compute per-group amax from input data
    AMAX-->>CPP: amax[0..N-1] written
    alt "with_amax_reduction=True"
        CPP->>NCCL: allreduce(amax, MAX)
    end
    CPP->>SCALE: "derive scale = fp8_max / amax, scale_inv = 1/scale"
    SCALE-->>CPP: scale[0..N-1], scale_inv[0..N-1] written
    CPP->>QUANT: cast input to FP8 using per-group scale
    QUANT-->>CPP: grouped FP8 output tensor
    CPP-->>PY: grouped output Python object

Loading

Comments Outside Diff (1)

1. transformer_engine/pytorch/csrc/extensions/cast.cpp, line 383-395 (link)

   nvfp4_group_quantize_with_amax missing last_dims guard

   bgrad_group_quantize immediately rejects a non-None last_dims with an explicit NVTE_CHECK. nvfp4_group_quantize_with_amax has no equivalent check. The NVFP4 group-quantize kernel was not updated in this PR to handle VARYING_LAST_DIM or VARYING_BOTH_DIMS layouts. If a caller passes a non-None last_dims, the metadata is forwarded into the NVFP4 dispatch path where it would produce incorrect output or an internal assertion failure. Adding the same guard as bgrad_group_quantize makes the limitation explicit.

_{Reviews (5): Last reviewed commit: "Merge branch 'current_scaling_group_quan..." | Re-trigger Greptile}

[vthumbe1503]

/te-ci

[vthumbe1503]

Pipeline: 54747206