[xrennvidia]

Description

With activation recompute in the fused NVFP4 GroupedMLP, the backward rebuilds the fc2 input (fc2_x) with the dSReLU kernel and re-quantizes it column-wise for wgrad. The forward nevertheless quantized fc2_x both row- and column-wise and discarded the column-wise copy, and the backward re-quantization ran a full group Hadamard-amax pass even though the regenerated tensor is bit-identical to the forward one. This PR removes that redundant work and fixes an amax mismatch in the Hadamard-fused fc1 kernel path with randomized Hadamard transform (RHT).

Measured on GB300 with a fused grouped MLP (8 experts, NVFP4 + RHT): the group quantize kernel for the fc2 input drops from ~570us to ~330us per fused call on a 130B hybrid MoE proxy, and the forward and backward each launch one fewer Hadamard-amax kernel per fused call where the Hadamard-fused fc1 kernel is used. Recompute on and off give bitwise-equal output and gradients for SReLU, tanh-SReLU and SwiGLU.

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

• Decide recompute_srelu_fc2_x before quantizing the fc2 input and request only row-wise NVFP4 data when the backward will rebuild the tensor; mark only the tensor actually saved for backward for paged stashing.
• Keep the forward (pre-RHT, post-RHT) amax of fc2_x on the fc2 context and quantize the recomputed tensor with nvfp4_group_quantize_with_amax, dropping the backward amax pass. Add a nvfp4_group_hadamard_amax binding so the plain SReLU kernel path computes both amax values in the single forward pass. NVTE_FUSED_GROUPED_MLP_REUSE_FC2_X_AMAX=0 disables the reuse.
• Route ScaledTanhSReLU through the Hadamard-fused fc1 kernel when the cuDNN frontend grouped_gemm_glu_hadamard_wrapper_sm100 accepts tanh_clamp_scale (detected by signature), so the forward amax pass disappears for tanh-SReLU as well.
• Pass the quantizer's RHT sign mask (rht_sign_mask) to the Hadamard-fused kernel when the cuDNN frontend accepts it. Previously the kernel computed the post-RHT amax for the plain transform, so the NVFP4 global scale was off by the sign pattern (up to ~17% too large for SReLU-shaped activations). On frontends without the argument, the fused kernel is used only when the mask is zero.
• tests/pytorch/test_grouped_mlp.py: the recompute test asserts a bitwise-equal fc2 weight gradient for tanh-SReLU with NVFP4; plain SReLU keeps its tolerance until the cuDNN frontend fix for the Hadamard-fused kernel is required.

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]

Confidence Score: 5/5

[Medium impact] The PR appears safe to merge; no actionable defects were found.

Summary

The PR reduces repeated NVFP4 work when the fused grouped MLP rebuilds its second-layer input during backward.

• Recomputed FC2 inputs reuse forward NVFP4 amax values during backward.
• Tanh-SReLU uses the fused Hadamard path when cuDNN supports its clamp.
• Fused FC1 kernels measure amax using the quantizer’s randomized Hadamard transform.

No actionable issues were found. xrennvidia explicitly documented the existing plain SReLU rounding difference on older frontends; this review does not report it as a new issue.

Diagram

%%{init: {'theme': 'neutral'}}%%
flowchart TD
  A[FC1 and activation] --> B[Get pre-RHT and post-RHT amax]
  B --> C[Quantize row-wise for FC2]
  B --> D[Keep amax for backward]
  C --> E[FC2 forward]
  E --> F[Backward rebuilds FC2 input]
  D --> G[Quantize rebuilt input column-wise]
  F --> G
  G --> H[FC2 weight gradient]

Loading

_{Reviews (1) · Last reviewed commit: "[pre-commit.ci] auto fixes from pre-comm..." · Reviewed by Greptile}