[PyTorch] Fix NaN expert weight gradients at num_groups == 1 with SReLU (#3485)

[PyTorch] Take the num_groups == 1 dense shortcuts only when the kernels agree

At num_groups == 1 the fused grouped MLP may quantize and scale-swizzle the token
buffer as ONE dense tensor over `tensor.shape[0]` rows and read it back the same
way. That is coherent only if the cuDNN kernels also write their intermediates
densely over `tensor.shape[0]`, which is what `use_single_group_runtime_offsets`
asks of them -- they then derive M from the runtime tensor shapes rather than from
the offsets.

The two halves of that decision are currently gated differently. The request to
the kernels is already guarded, because it can be refused:

    if supports_single_group_runtime_offsets:
        fc1_activation_kwargs["use_single_group_runtime_offsets"] = num_groups == 1

but its siblings -- `use_single_group_weight_swizzle`, `use_single_discrete_weight`,
`use_single_group_dense_fc2`, `use_single_group_dense_dgrad`, the dense quantize
and the plain wgrad GEMM -- are gated on `num_groups == 1` alone. So when the
specialization is unavailable, `ScaledSReLU` at any version or any activation on a
cuDNN frontend older than 1.27.0, the kernels pack their output to
`sum(split_sizes)` rows while TE still reads everything densely.

Because the columnwise swizzled MXFP8 scale layout is `[k/128][m/128][32][4][4]`,
with the m-tile as an INNER stride, an m-extent mismatch misindexes every k-tile
past the first: consumer k-tile j reads producer k-tile j*T_c/T_p. The HEAD rows
-- the ones that do belong to the group -- pick up the wrong scale factors, and
high k-tiles read scale memory the producer never wrote. The data is read
correctly; only the scales move, so the error is order 1 and often NaN. Crucially
this is value-independent: zeroing the padded tail does not prevent it, because
the defect is in the indexing rather than in the padding's contents, so a caller
cannot work around it.

Callers may legally pass a token buffer padded past `sum(split_sizes)` -- an MoE
expert-parallel fixed-capacity receive buffer always does, and
`test_grouped_linear_cuda_graph_safe` documents the padded tail as "intentionally
outside every group" and "uninitialized".

Measured on B300 with MXFP8 and a native ScaledSReLU at num_groups == 1, input
padded 512 -> 768 rows, FC1 main_grad against the unpadded result:

    before   zeroed tail    406528 / 1048576 non-finite
    before   poisoned tail  406528 / 1048576 non-finite   (identical -- not the values)
    after    either tail    0.0000e+00, bit-identical
    control  num_groups=2   clean before and after

Every corrected arm also sits at 1.70e-02 against an unfused reference, the MXFP8
fused-vs-unfused noise floor.

Gate the whole shortcut family on the same predicate the kernels use, applying the
guard TE already writes at the sites that were missing it. No new helper: the
binding calls `_cudnn_frontend_supports_single_group_runtime_offsets` directly,
as this file already does at three other sites, and the flag is named
`use_dense_single_group` to match its neighbours. It deliberately does not claim
the group spans the buffer -- that is unknowable here -- only that producer and
consumer will agree on the extent. It is a host-side check
on an activation type and a package version, so it costs nothing and stays CUDA
graph capturable, and it leaves the GLU activations -- which do get the
specialization -- byte-for-byte on their existing fast path. That matters:
removing the shortcuts outright costs a shared expert +6% to +30% of the fused MLP
step (+5-10% at hidden 4096 / ffn 14336, rising to +28-30% at 1024/4096 with few
tokens), measured by CUDA-graph replay against num_groups >= 2 null controls, and
Megatron's FusedSharedExpertMLP builds exactly this path.

In the two grouped-quantize helpers the condition is just `use_dense_single_group`:
`num_groups == 1` is already folded into that flag at both binding sites, and the
`isinstance(quantizer, (MXFP8Quantizer, NVFP4Quantizer))` test was dead because
`fuse_grouped_mlp_ops` only builds this fusion for an MXFP8 or NVFP4 recipe.

All the shortcuts have to move together: quantization layout and kernel indexing
must agree, so enabling a subset pairs a dense consumer with a packed producer and
yields NaN rather than a slowdown. Hence one flag threaded explicitly through the
grouped-quantize helpers and `_compute_grad_params` rather than a predicate per
site.

What this does NOT do: where the specialization is available, producer and
consumer agree and the residual exposure is that the plain wgrad GEMM still
contracts over `logical_shape[0]` rows, summing any padding past
`sum(split_sizes)` into the weight gradient. That part is value-dependent -- a
zeroed tail contributes an all-zero outer product and is bit-exact, verified over
10 launches -- so a caller that zeroes its padding, or passes none, is unaffected.
Closing that gap needs either a device-side masked zeroing of TE's own
intermediate buffers or an explicit per-op dense-single-group contract; neither is
attempted here.

The condition cannot be evaluated at run time: `split_sizes` is device-resident,
so testing `sum(split_sizes) == tensor.shape[0]` forces a device-to-host sync, and
that test exists to keep this flow sync-free and graph-capturable.

tests/pytorch/test_grouped_mlp.py is at baseline: 6 failed, 48 passed, the 6 being
pre-existing nvfp4_rht and cuBLAS-version failures also present on main.

Signed-off-by: Wentao Guo <wg0420@princeton.edu>
Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
Co-authored-by: vthumbe1503 <vthumbe@nvidia.com>

Author: 3 people

transformer_engine/pytorch/ops/fused/grouped_mlp.py

@@ -239,10 +239,11 @@ def _group_quantize_for_grouped_mlp(
     split_sizes: Optional[torch.Tensor],
     *,
     tensor_offsets: Optional[torch.Tensor] = None,
+    use_dense_single_group: bool = False,
 ) -> GroupedTensor:
     """Quantize into grouped storage."""
 
-    if num_groups != 1 or not isinstance(quantizer, (MXFP8Quantizer, NVFP4Quantizer)):
+    if not use_dense_single_group:
         return tex.group_quantize(
             tensor,
             quantizer,
@@ -270,6 +271,7 @@ def _group_quantize_with_amax_for_grouped_mlp(
     columnwise_amax: torch.Tensor,
     *,
     tensor_offsets: Optional[torch.Tensor] = None,
+    use_dense_single_group: bool = False,
 ) -> GroupedTensor:
     """Quantize with precomputed NVFP4 amaxes into grouped storage."""
     if not isinstance(quantizer, NVFP4Quantizer):
@@ -279,9 +281,10 @@ def _group_quantize_with_amax_for_grouped_mlp(
             num_groups,
             split_sizes,
             tensor_offsets=tensor_offsets,
+            use_dense_single_group=use_dense_single_group,
         )
 
-    if num_groups != 1:
+    if not use_dense_single_group:
         return tex.nvfp4_group_quantize_with_amax(
             tensor,
             quantizer,
@@ -690,6 +693,7 @@ def _compute_grad_params(
     scale_view_dtype,
     sf_vec_size,
     offsets,
+    use_dense_single_group,
 ):
     """Compute weight gradients and build grad_params for a GroupedLinear layer.
     Returns the grad_params list in parameter registration order.
@@ -757,7 +761,7 @@ def _compute_grad_params(
                 "distributed-weight fused grouped-MLP requires delay_wgrad_compute=False."
             )
         if (
-            num_groups == 1
+            use_dense_single_group
             and isinstance(grouped_x, (GroupedTensor, GroupedTensorStorage))
             and isinstance(grouped_dy, (GroupedTensor, GroupedTensorStorage))
             and isinstance(grouped_x.quantizer, (MXFP8Quantizer, NVFP4Quantizer))
@@ -1125,6 +1129,9 @@ def fuser_forward(
             raise ValueError(f"Unsupported input shape for fused grouped MLP ({in_shape=}).")
 
         num_groups = fc1_op.num_groups
+        use_dense_single_group = num_groups == 1 and (
+            _cudnn_frontend_supports_single_group_runtime_offsets(type(activation_op))
+        )
         fc1_weight_param = fc1_op.weight if fc1_op.single_grouped_weight else fc1_op.weight0
         fc2_weight_param = fc2_op.weight if fc2_op.single_grouped_weight else fc2_op.weight0
 
@@ -1195,7 +1202,7 @@ def fuser_forward(
         # Older cuDNN frontends do not expose this specialization, so use the
         # live generic offset calculation rather than caching CUDA metadata.
         use_offsetless_metadata = (
-            num_groups == 1
+            use_dense_single_group
             and unit_activation_scale
             and isinstance(fc1_input_quantizer, MXFP8Quantizer)
             and supports_single_group_runtime_offsets
@@ -1276,6 +1283,7 @@ def fuser_forward(
                     fc1_weight_quantizer,
                     num_groups,
                     None,
+                    use_dense_single_group=use_dense_single_group,
                 )
         else:
             fc1_weights = [getattr(fc1_op, f"weight{idx}") for idx in range(num_groups)]
@@ -1313,6 +1321,7 @@ def fuser_forward(
                     fc2_weight_quantizer,
                     num_groups,
                     None,
+                    use_dense_single_group=use_dense_single_group,
                 )
         else:
             fc2_weights = [getattr(fc2_op, f"weight{idx}") for idx in range(num_groups)]
@@ -1360,6 +1369,7 @@ def fuser_forward(
                 num_groups,
                 split_sizes,
                 tensor_offsets=fc1_x_tensor_offsets,
+                use_dense_single_group=use_dense_single_group,
             )
 
         use_nvfp4 = isinstance(fc1_input_quantizer, NVFP4Quantizer) or isinstance(
@@ -1495,7 +1505,7 @@ def fuser_forward(
             fc1_activation_kwargs["norm_const_tensor"] = fc1_norm_const_tensor
             fc1_activation_kwargs["discrete_col_sfd"] = not use_nvfp4
             if supports_single_group_runtime_offsets:
-                fc1_activation_kwargs["use_single_group_runtime_offsets"] = num_groups == 1
+                fc1_activation_kwargs["use_single_group_runtime_offsets"] = use_dense_single_group
         if self._pass_geglu_runtime_params:
             fc1_activation_kwargs.update(
                 linear_offset=self._cudnn_linear_offset,
@@ -1514,7 +1524,7 @@ def fuser_forward(
         if fc1_op.single_grouped_weight:
             # Clone and swizzle scales for GEMM.
             fc1_weight_for_gemm = grouped_fc1_weight.copy()
-            use_single_group_weight_swizzle = num_groups == 1
+            use_single_group_weight_swizzle = use_dense_single_group
             if use_single_group_weight_swizzle:
                 fc1_weight_single = _single_quantized_tensor_from_grouped(fc1_weight_for_gemm)
                 fc1_weight_single._columnwise_data = None
@@ -1551,7 +1561,7 @@ def fuser_forward(
             fc1_activation_kwargs["b_tensor"] = fc1_w_data
             fc1_activation_kwargs["sfb_tensor"] = fc1_w_scales
         else:
-            use_single_discrete_weight = num_groups == 1
+            use_single_discrete_weight = use_dense_single_group
             if use_single_discrete_weight:
                 fc1_weight_single = grouped_fc1_weight[0]
                 original_rowwise_scale = fc1_weight_single._rowwise_scale_inv
@@ -1653,6 +1663,7 @@ def fuser_forward(
                     fc1_kernel_out["amax_tensor"].view(-1),
                     fc1_kernel_out["post_rht_amax_tensor"].view(-1),
                     tensor_offsets=fc2_x_tensor_offsets,
+                    use_dense_single_group=use_dense_single_group,
                 )
             else:
                 grouped_fc2_x = _group_quantize_for_grouped_mlp(
@@ -1661,11 +1672,12 @@ def fuser_forward(
                     num_groups,
                     split_sizes,
                     tensor_offsets=fc2_x_tensor_offsets,
+                    use_dense_single_group=use_dense_single_group,
                 )
 
             fc2_out_buf = validate_or_alloc_output(output_buffer, fc2_out_shape, dtype, device)
             if (
-                num_groups == 1
+                use_dense_single_group
                 and grouped_fc2_x.columnwise_data is not None
                 and grouped_fc2_x.columnwise_scale_inv is not None
             ):
@@ -1723,7 +1735,7 @@ def fuser_forward(
                 with_gemm_swizzled_scales=True,
             )
 
-            use_single_group_dense_fc2 = num_groups == 1
+            use_single_group_dense_fc2 = use_dense_single_group
             fc2_out_buf = validate_or_alloc_output(output_buffer, fc2_out_shape, dtype, device)
             if use_single_group_dense_fc2:
                 fc2_out = _single_group_fc2_gemm(
@@ -1764,7 +1776,7 @@ def fuser_forward(
                 }
                 fc2_quant_kernel = self.grouped_gemm_quant_kernel()
                 if supports_single_group_runtime_offsets:
-                    fc2_quant_kwargs["use_single_group_runtime_offsets"] = num_groups == 1
+                    fc2_quant_kwargs["use_single_group_runtime_offsets"] = use_dense_single_group
 
                 if fc2_op.single_grouped_weight:
                     # Clone and swizzle scales for GEMM
@@ -1935,6 +1947,9 @@ def fuser_backward(
             grad_output = grad_output.reshape(-1, fc2_weight_shape[0])
         out_shape = list(grad_output.size())
         num_groups = fc1_op.num_groups
+        use_dense_single_group = num_groups == 1 and (
+            _cudnn_frontend_supports_single_group_runtime_offsets(type(activation_op))
+        )
         fc1_weight_param = fc1_op.weight if fc1_op.single_grouped_weight else fc1_op.weight0
         fc2_weight_param = fc2_op.weight if fc2_op.single_grouped_weight else fc2_op.weight0
         device = fc1_weight_param.device
@@ -2041,6 +2056,7 @@ def fuser_backward(
                     num_groups,
                     split_sizes,
                     tensor_offsets=fc2_out_tensor_offsets,
+                    use_dense_single_group=use_dense_single_group,
                 )
 
         use_nvfp4 = (
@@ -2194,7 +2210,7 @@ def fuser_backward(
             # Never passed to a wrapper that would reject it -- the check above raises first.
             fc2_dactivation_kwargs["deterministic"] = True
         if _cudnn_frontend_supports_single_group_runtime_offsets(type(activation_op)):
-            fc2_dactivation_kwargs["use_single_group_runtime_offsets"] = num_groups == 1
+            fc2_dactivation_kwargs["use_single_group_runtime_offsets"] = use_dense_single_group
         if self._cudnn_dact_func is not None:
             fc2_dactivation_kwargs["beta_tensor"] = fc2_beta_tensor
             fc2_dactivation_kwargs["act_func"] = self._cudnn_dact_func
@@ -2248,7 +2264,7 @@ def fuser_backward(
             fc2_dactivation_kwargs["b_tensor"] = fc2_w_data
             fc2_dactivation_kwargs["sfb_tensor"] = fc2_w_scales
         else:
-            use_single_discrete_weight = num_groups == 1
+            use_single_discrete_weight = use_dense_single_group
             if use_single_discrete_weight:
                 fc2_weight_single = grouped_fc2_weight[0]
                 original_rowwise_data = fc2_weight_single._rowwise_data
@@ -2353,6 +2369,7 @@ def fuser_backward(
                     num_groups,
                     split_sizes,
                     tensor_offsets=fc2_x_tensor_offsets,
+                    use_dense_single_group=use_dense_single_group,
                 )
             else:
                 sfd_col_d_srelu_tensor = fc2_dgrad_kernel_out.get("sfd_col_d_srelu_tensor")
@@ -2430,6 +2447,7 @@ def fuser_backward(
                 num_groups,
                 split_sizes,
                 tensor_offsets=fc1_dy_tensor_offsets,
+                use_dense_single_group=use_dense_single_group,
             )
         else:
             grouped_fc1_dy = GroupedTensor(
@@ -2466,6 +2484,7 @@ def fuser_backward(
             scale_view_dtype=scale_view_dtype,
             sf_vec_size=sf_vec_size,
             offsets=split_points,
+            use_dense_single_group=use_dense_single_group,
         )
 
         # Clear FC2 input tensor if possible
@@ -2491,7 +2510,7 @@ def fuser_backward(
             if is_distributed_weight(fc1_leader):
                 grouped_fc1_weight = materialize_weight_for_backward(fc1_leader)
 
-            use_single_group_dense_dgrad = num_groups == 1
+            use_single_group_dense_dgrad = use_dense_single_group
             if use_single_group_dense_dgrad:
                 grad_input = validate_or_alloc_output(grad_input_buffer, in_shape, dtype, device)
                 _single_group_dgrad_gemm(
@@ -2545,7 +2564,7 @@ def fuser_backward(
                 }
                 fc1_dgrad_kernel = self.grouped_gemm_quant_kernel()
                 if _cudnn_frontend_supports_single_group_runtime_offsets(type(activation_op)):
-                    fc1_dgrad_kwargs["use_single_group_runtime_offsets"] = num_groups == 1
+                    fc1_dgrad_kwargs["use_single_group_runtime_offsets"] = use_dense_single_group
 
                 if fc1_op.single_grouped_weight:
                     # Clone and swizzle scales for GEMM
@@ -2622,6 +2641,7 @@ def fuser_backward(
             scale_view_dtype=scale_view_dtype,
             sf_vec_size=sf_vec_size,
             offsets=split_points,
+            use_dense_single_group=use_dense_single_group,
         )
 
         # Clear FC1 input tensor if possible