Enable NVFP4 grouped MLP GLU RHT amax path

[sraman-rgb]

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

[greptile-apps]

Greptile Summary

This PR enables the NVFP4 grouped MLP GLU RHT amax path by adding two new C++ entry points (nvfp4_quantize_with_amax and nvfp4_group_quantize_with_amax) that accept externally precomputed per-group amaxes and handle distributed allreduce before quantization. A new fused grouped_gemm_glu_hadamard_kernel is wired into ForwardGroupedMLP_CuTeGEMMGLU so the FC1 GLU+RHT pass can emit amaxes directly to FC2 input quantization.

• New C++ quantize-with-amax functions: allreduce is correctly placed before cast kernels in both paths; empty-rank distributed scenarios are handled by the compute_amax=false branch in quantize_impl.
• quantize_impl refactor: reduce_amaxes lambda fires at a single site before all cast kernels, fixing ordering for both the existing and new paths.
• Python integration: _group_quantize_with_amax_for_grouped_mlp dispatches to the correct C++ path; _use_tmem_post_rht_amax is properly cached with lru_cache.

Confidence Score: 5/5

The new quantize-with-amax paths correctly place allreduce before cast kernels in both the single-tensor and grouped variants, and empty-rank distributed scenarios are handled properly for all new code paths.

All new C++ and Python code introduced in this PR correctly implements the amax-reduction-before-quantization ordering. The compute_amax=false branch in quantize_impl calls reduce_amaxes() before the cast kernels for both non-empty and empty inputs. The grouped path calls allreduce_nvfp4_amax_tensors explicitly before group_quantize_nvfp4_impl. Pre-existing behavior for the regular quantize path (empty-rank early exit without allreduce) is unchanged.

No files require special attention; the most complex logic is in cast.cpp and quantizer.cpp, both of which correctly sequence allreduce before quantization in all new code paths.

Important Files Changed

Filename	Overview
transformer_engine/pytorch/csrc/extensions/cast.cpp	Adds nvfp4_quantize_with_amax and nvfp4_group_quantize_with_amax; allreduce is correctly placed before the cast in both paths, and empty-rank handling is correct for the new compute_amax=false path.
transformer_engine/pytorch/csrc/quantizer.cpp	Refactors quantize_impl to accept compute_amax flag; reduce_amaxes lambda is now called before cast kernels for both flag values. Empty ranks with compute_amax=false now participate in the collective.
transformer_engine/pytorch/ops/_common.py	Adds _group_quantize_with_amax_for_grouped_mlp and extracts _wrap_single_nvfp4_as_grouped; the num_groups==1 path now uses nvfp4_quantize_with_amax with correct allreduce before cast.
transformer_engine/pytorch/ops/fused/forward_grouped_mlp.py	Adds grouped_gemm_glu_hadamard_kernel to ForwardGroupedMLP_CuTeGEMMGLU and integrates the new fused GLU+RHT+amax kernel path; _use_tmem_post_rht_amax is properly cached via lru_cache.
transformer_engine/pytorch/csrc/common.h	Moves quantize_impl from private to public on NVFP4Quantizer to allow cast.cpp to call it directly with compute_amax=false.
transformer_engine/pytorch/csrc/extensions.h	Declares two new public entry points nvfp4_quantize_with_amax and nvfp4_group_quantize_with_amax; signatures are consistent with implementations.
transformer_engine/pytorch/csrc/extensions/pybind.cpp	Exposes the two new quantization-with-amax functions to Python via pybind11; argument names and default values match the C++ signatures.

_{Reviews (8): Last reviewed commit: "Merge branch 'main' into nvfp4-grouped-m..." | Re-trigger Greptile}

[vthumbe1503]

Mostly LGTM. Left a few comments on code duplication and other minor issues.

[timmoon10]

/te-ci pytorch

[timmoon10]

LGTM, pending CI