From 068e1119c1065d4d97eb42c2c055916a04754d43 Mon Sep 17 00:00:00 2001 From: Wentao Guo Date: Thu, 3 Sep 2026 17:58:23 -0400 Subject: [PATCH] [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. Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: Wentao Guo --- .../pytorch/ops/fused/grouped_mlp.py | 48 +++++++++++++------ 1 file changed, 34 insertions(+), 14 deletions(-) diff --git a/transformer_engine/pytorch/ops/fused/grouped_mlp.py b/transformer_engine/pytorch/ops/fused/grouped_mlp.py index dca1119800..2a50d37909 100644 --- a/transformer_engine/pytorch/ops/fused/grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/grouped_mlp.py @@ -195,10 +195,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, @@ -226,6 +227,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): @@ -235,9 +237,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, @@ -646,6 +649,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. @@ -713,7 +717,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)) @@ -1055,6 +1059,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 @@ -1125,7 +1132,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 @@ -1206,6 +1213,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)] @@ -1243,6 +1251,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)] @@ -1290,6 +1299,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( @@ -1421,7 +1431,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, @@ -1438,7 +1448,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 @@ -1475,7 +1485,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 @@ -1577,6 +1587,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( @@ -1585,11 +1596,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 ): @@ -1647,7 +1659,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( @@ -1688,7 +1700,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 @@ -1853,6 +1865,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 @@ -1959,6 +1974,7 @@ def fuser_backward( num_groups, split_sizes, tensor_offsets=fc2_out_tensor_offsets, + use_dense_single_group=use_dense_single_group, ) use_nvfp4 = ( @@ -2080,7 +2096,7 @@ def fuser_backward( } dactivation_kernel = self.grouped_gemm_dactivation_kernel() 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 @@ -2132,7 +2148,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 @@ -2237,6 +2253,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") @@ -2314,6 +2331,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( @@ -2350,6 +2368,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 @@ -2375,7 +2394,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( @@ -2429,7 +2448,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 @@ -2506,6 +2525,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