[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 new issue was established.

Summary

This PR reduces repeated NVFP4 work when the fused grouped MLP rebuilds its activation during backward.

• Recomputed FC2 inputs reuse forward amax values for quantization.
• The fused FC1 path matches its amax calculation to the randomized transform.

No actionable new issues were established. No previous review threads were supplied. This review used source inspection; GPU tests were not run.

Diagram

%%{init: {'theme': 'neutral'}}%%
flowchart TD
  A[FC1 and activation] --> B[FC2 input and amax]
  B --> C[Row-wise NVFP4 quantization]
  C --> D[FC2 forward]
  B --> E[Save amax and pre-activation]
  E --> F[Backward rebuilds FC2 input]
  F --> G[Column-wise quantization with saved amax]
  G --> H[FC2 weight gradient]

Loading

_{Reviews (11) · Last reviewed commit: "Merge branch 'main' into xren/fused_grou..." · Reviewed by Greptile}

[timmoon10]

Secret envvar configs are an antipattern. This optimization is a strict improvement, so we should just do it when suported.

Suggested change

	def _fc2_x_amax_reuse_enabled() -> bool:
	"""Whether the backward reuses the forward NVFP4 amax of a recomputed fc2 input.
	The grouped dSReLU kernel regenerates fc2_x bit-exactly from the saved pre-activation,
	so the per-group amax computed in the forward is also the amax of the regenerated
	tensor. Reusing it turns the backward re-quantization into a single cast kernel.
	``NVTE_FUSED_GROUPED_MLP_REUSE_FC2_X_AMAX=0`` disables the reuse.
	"""
	return int(os.environ.get("NVTE_FUSED_GROUPED_MLP_REUSE_FC2_X_AMAX", "1")) > 0

[xrennvidia]

Fixed.

[timmoon10]

This is overly defensive. transformer_engine_torch is internal implementation, and it's expected to match the Python code the same way Python files are expected to match each other.

Suggested change

	amax_fn = getattr(tex, "nvfp4_group_hadamard_amax", None)
	if (
	amax_fn is None
	or not isinstance(quantizer, NVFP4Quantizer)
	or split_sizes is None
	or tensor_offsets is None
	):
	return None
	return amax_fn(tensor, quantizer, num_groups, split_sizes, tensor_offsets=tensor_offsets)
	if (
	not isinstance(quantizer, NVFP4Quantizer)
	or split_sizes is None
	or tensor_offsets is None
	):
	return None
	return tex.nvfp4_group_hadamard_amax(tensor, quantizer, num_groups, split_sizes, tensor_offsets=tensor_offsets)

[xrennvidia]

Fixed

[timmoon10]

This variable is no longer telling you whether "the activation is SReLU", but now is doing something convoluted depending on the cuDNN FE version. Claude likes doing this kind of thing because it narrowly accomplishes its task with a single line change, and technical debt is a problem for someone else. If this variable no longer has a logically coherent meaning, it would be better to get rid of it and make the logic more explicit.

Where we choose the cuDNN kernel:

TransformerEngine/transformer_engine/pytorch/ops/fused/grouped_mlp.py

Lines 1761 to 1762 in 5eace13

	elif activation_is_srelu and _cudnn_frontend_supports_grouped_gemm_srelu_hadamard():
	kernel_impl = "gemm_act_rht_amax"

elif isinstance(activation_op, ScaledSReLU) and _cudnn_frontend_supports_grouped_gemm_srelu_hadamard():
    kernel_impl = "gemm_act_rht_amax"
elif isinstance(activation_op, ScaledTanhSReLU) and _cudnn_frontend_supports_grouped_gemm_srelu_tanh_hadamard():
    kernel_impl = "gemm_act_rht_amax"

Where we configure args to pass to the cuDNN kernel:

TransformerEngine/transformer_engine/pytorch/ops/fused/grouped_mlp.py

Lines 1895 to 1897 in 5eace13

	if activation_is_srelu:
	if kernel_impl in ("gemm_act_rht_amax", "gemm_act_rht_quant"):
	fc1_activation_kwargs["act_func"] = "srelu"

if isinstance(activation_op, (ScaledSReLU, ScaledTanhSReLU)):
    if kernel_impl in ("gemm_act_rht_amax", "gemm_act_rht_quant"):
        fc1_activation_kwargs["act_func"] = "srelu"

[xrennvidia]

Fixed

[timmoon10]

Nit: Logically this isn't specific to SReLU. If the GLU dgrad kernel supported recompute we could use the same logic.

Suggested change

	# Decide up front whether the backward rebuilds fc2_x from the dSReLU kernel
	# (which applies the tanh clamp when regenerating fc2_x). When it does, the
	# forward only needs row-wise fc2_x for the fc2 GEMM; the column-wise copy
	# for wgrad is produced in the backward from the recomputed tensor.
	recompute_srelu_fc2_x = (
	# Decide up front whether the backward rebuilds fc2_x from the activation
	# backward kernel. If so, no need to compute column-wise data in forward.
	recompute_fc2_x = (

[xrennvidia]

Fixed

[timmoon10]

Tensors should be saved with save_for_backward, not attached directly to the autograd context. Otherwise PyTorch's autograd can do weird things (see PyTorch autograd context save_for_backward).

TransformerEngine/transformer_engine/pytorch/ops/fused/grouped_mlp.py

Line 2451 in 5eace13

fc1_ctx.save_for_backward(

[xrennvidia]

Fixed