From e1ebeb3fc1440fb08f0e25ac69b39dc3747eabc7 Mon Sep 17 00:00:00 2001 From: Siddhartha Raman Sundara Raman Date: Mon, 6 Jul 2026 23:23:02 -0700 Subject: [PATCH 01/14] Optimize single-group grouped MLP quantization Signed-off-by: Siddhartha Raman Sundara Raman --- .../pytorch/ops/fused/grouped_mlp.py | 194 +++++++++++++----- 1 file changed, 143 insertions(+), 51 deletions(-) diff --git a/transformer_engine/pytorch/ops/fused/grouped_mlp.py b/transformer_engine/pytorch/ops/fused/grouped_mlp.py index 83954a9b3de..a3b3c190e14 100644 --- a/transformer_engine/pytorch/ops/fused/grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/grouped_mlp.py @@ -29,7 +29,7 @@ from ...quantization import Recipe from ...tensor import NVFP4Quantizer, NVFP4Tensor, NVFP4TensorStorage, Quantizer from ...tensor.grouped_tensor import GroupedTensor -from ...tensor.mxfp8_tensor import MXFP8Quantizer +from ...tensor.mxfp8_tensor import MXFP8Quantizer, MXFP8Tensor from ...tensor.storage.grouped_tensor_storage import GroupedTensorStorage from ...triton.grouped_dbias_dscales import compute_grouped_dbias_dscales from ...utils import ( @@ -101,15 +101,15 @@ def _nvidia_cudnn_frontend_supports_wgrad() -> bool: return _cudnn_frontend_version_supported() -def _wrap_single_nvfp4_as_grouped( +def _wrap_single_quantized_as_grouped( tensor: torch.Tensor, - quantized: NVFP4Tensor | NVFP4TensorStorage, - quantizer: NVFP4Quantizer, + quantized: MXFP8Tensor | NVFP4Tensor | NVFP4TensorStorage, + quantizer: MXFP8Quantizer | NVFP4Quantizer, split_sizes: Optional[torch.Tensor], *, tensor_offsets: Optional[torch.Tensor] = None, ) -> GroupedTensor: - """Wrap a single NVFP4 tensor in GroupedTensor storage.""" + """Wrap a single quantized tensor in GroupedTensor storage.""" with_gemm_swizzled_scales = quantized._with_gemm_swizzled_scales if quantizer.optimize_for_gemm: tex.swizzle_scales_for_gemm_(quantized) @@ -119,8 +119,8 @@ def _wrap_single_nvfp4_as_grouped( rowwise_scale = quantized._rowwise_scale_inv columnwise_data = quantized._columnwise_data columnwise_scale = quantized._columnwise_scale_inv - amax = quantized._amax_rowwise - columnwise_amax = quantized._amax_columnwise + amax = getattr(quantized, "_amax_rowwise", None) + columnwise_amax = getattr(quantized, "_amax_columnwise", None) if split_sizes is None: split_sizes = torch.full((1,), tensor.shape[0], dtype=torch.int64, device=tensor.device) @@ -128,19 +128,18 @@ def _wrap_single_nvfp4_as_grouped( split_sizes = split_sizes.to(dtype=torch.int64, device=tensor.device) m_dim = tensor.shape[0] - if rowwise_data is not None: + if isinstance(quantizer, NVFP4Quantizer) and rowwise_data is not None: k_dim = rowwise_data.shape[-1] * 2 - elif columnwise_data is not None: + elif isinstance(quantizer, NVFP4Quantizer) and columnwise_data is not None: k_dim = columnwise_data.shape[0] else: k_dim = tensor.shape[-1] if tensor_offsets is None: - tensor_offsets = torch.cat( - [ - torch.zeros(1, dtype=torch.int64, device=tensor.device), - torch.cumsum(split_sizes * k_dim, dim=0), - ], + tensor_offsets = torch.tensor( + [0, m_dim * k_dim], + dtype=torch.int64, + device=tensor.device, ) return GroupedTensor( @@ -170,7 +169,16 @@ def _group_quantize_for_grouped_mlp( ) -> GroupedTensor: """Quantize into grouped storage.""" - if num_groups != 1 or not isinstance(quantizer, NVFP4Quantizer): + if num_groups != 1: + return tex.group_quantize( + tensor, + quantizer, + num_groups, + split_sizes, + tensor_offsets=tensor_offsets, + ) + + if not isinstance(quantizer, (MXFP8Quantizer, NVFP4Quantizer)): return tex.group_quantize( tensor, quantizer, @@ -180,7 +188,7 @@ def _group_quantize_for_grouped_mlp( ) quantized = tex.quantize(tensor, quantizer) - return _wrap_single_nvfp4_as_grouped( + return _wrap_single_quantized_as_grouped( tensor, quantized, quantizer, @@ -223,7 +231,7 @@ def _group_quantize_with_amax_for_grouped_mlp( quantized = tex.nvfp4_quantize_with_amax( tensor, quantizer, rowwise_amax.view(-1)[:1], columnwise_amax.view(-1)[:1] ) - return _wrap_single_nvfp4_as_grouped( + return _wrap_single_quantized_as_grouped( tensor, quantized, quantizer, @@ -253,22 +261,29 @@ def _nvfp4_amax( return torch.cat([amax.view(-1) for amax in amaxes], dim=0) -def _nvfp4_single_tensor_from_grouped( +def _single_quantized_tensor_from_grouped( grouped: GroupedTensor, - quantizer: Optional[NVFP4Quantizer] = None, + quantizer: Optional[MXFP8Quantizer | NVFP4Quantizer] = None, *, fp4_dtype: Optional[torch.dtype] = None, -) -> NVFP4Tensor: - """Build a single NVFP4Tensor view over a one-member grouped storage.""" +) -> MXFP8Tensor | NVFP4Tensor: + """Build a single quantized tensor view over a one-member grouped storage.""" if quantizer is None: quantizer = grouped.quantizer - if not isinstance(quantizer, NVFP4Quantizer): - raise TypeError("Expected an NVFP4 GroupedTensor.") + if not isinstance(quantizer, (MXFP8Quantizer, NVFP4Quantizer)): + raise TypeError("Expected an MXFP8 or NVFP4 GroupedTensor.") shape = tuple(grouped.logical_shape) + if isinstance(quantizer, NVFP4Quantizer): + data_shape = quantizer.convert_shape_for_fp4(shape) + columnwise_shape = quantizer.convert_shape_for_fp4(quantizer.get_columnwise_shape(shape)) + else: + data_shape = shape + columnwise_shape = quantizer.get_columnwise_shape(shape) + rowwise_data = None if grouped.rowwise_data is not None: - rowwise_data = grouped.rowwise_data.view(quantizer.convert_shape_for_fp4(shape)) + rowwise_data = grouped.rowwise_data.view(data_shape) rowwise_scale_inv = None if grouped.scale_inv is not None: @@ -276,10 +291,7 @@ def _nvfp4_single_tensor_from_grouped( columnwise_data = None if grouped.columnwise_data is not None: - columnwise_shape = quantizer.get_columnwise_shape(shape) - columnwise_data = grouped.columnwise_data.view( - quantizer.convert_shape_for_fp4(columnwise_shape) - ) + columnwise_data = grouped.columnwise_data.view(columnwise_shape) columnwise_scale_inv = None if grouped.columnwise_scale_inv is not None: @@ -287,6 +299,20 @@ def _nvfp4_single_tensor_from_grouped( quantizer.get_scale_shape(shape, True) ) + if isinstance(quantizer, MXFP8Quantizer): + return MXFP8Tensor( + shape=shape, + dtype=grouped.get_dtype(), + rowwise_data=rowwise_data, + rowwise_scale_inv=rowwise_scale_inv, + columnwise_data=columnwise_data, + columnwise_scale_inv=columnwise_scale_inv, + fp8_dtype=quantizer.dtype, + quantizer=quantizer, + requires_grad=False, + with_gemm_swizzled_scales=grouped._with_gemm_swizzled_scales, + ) + return NVFP4Tensor( shape=shape, dtype=grouped.get_dtype(), @@ -341,7 +367,44 @@ def _use_tmem_post_rht_amax() -> bool: return os.environ.get("NVTE_CUTEDSL_FUSED_GROUPED_MLP_FC1_GLU_RHT_AMAX_TMEM", "0") == "1" -def _nvfp4_single_group_wgrad_gemm( +def _single_group_split_metadata( + num_tokens: int, + device: torch.device, + *, + fc1_in_features: int, + fc2_in_features: int, + fc2_out_features: int, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Create split/offset metadata for num_groups=1 without a GPU prefix-sum kernel.""" + split_sizes = torch.tensor([num_tokens], dtype=torch.int64, device=device) + split_points = torch.tensor([num_tokens], dtype=torch.int32, device=device) + base_split_offsets = torch.tensor([0, num_tokens], dtype=torch.int64, device=device) + fc1_x_tensor_offsets = torch.tensor( + [0, num_tokens * fc1_in_features], + dtype=torch.int64, + device=device, + ) + fc2_x_tensor_offsets = torch.tensor( + [0, num_tokens * fc2_in_features], + dtype=torch.int64, + device=device, + ) + fc2_out_tensor_offsets = torch.tensor( + [0, num_tokens * fc2_out_features], + dtype=torch.int64, + device=device, + ) + return ( + split_sizes, + split_points, + base_split_offsets, + fc1_x_tensor_offsets, + fc2_x_tensor_offsets, + fc2_out_tensor_offsets, + ) + + +def _single_group_wgrad_gemm( grouped_x: GroupedTensor, grouped_dy: GroupedTensor, wgrad_output, @@ -349,9 +412,9 @@ def _nvfp4_single_group_wgrad_gemm( weight_shape: tuple[int, int], accumulate: bool, ) -> None: - """Run one-group NVFP4 wgrad with regular GEMM instead of grouped GEMM.""" - x_single = _nvfp4_single_tensor_from_grouped(grouped_x) - dy_single = _nvfp4_single_tensor_from_grouped(grouped_dy) + """Run one-group MXFP8/NVFP4 wgrad with regular GEMM instead of grouped GEMM.""" + x_single = _single_quantized_tensor_from_grouped(grouped_x) + dy_single = _single_quantized_tensor_from_grouped(grouped_dy) if isinstance(wgrad_output, GroupedTensor): out = wgrad_output.rowwise_data.view(1, *weight_shape)[0] else: @@ -620,11 +683,11 @@ def _compute_grad_params( num_groups == 1 and isinstance(grouped_x, GroupedTensor) and isinstance(grouped_dy, GroupedTensor) - and isinstance(grouped_x.quantizer, NVFP4Quantizer) - and isinstance(grouped_dy.quantizer, NVFP4Quantizer) + and isinstance(grouped_x.quantizer, (MXFP8Quantizer, NVFP4Quantizer)) + and isinstance(grouped_dy.quantizer, grouped_x.quantizer.__class__) ): gemm_fn = functools.partial( - _nvfp4_single_group_wgrad_gemm, + _single_group_wgrad_gemm, weight_shape=weight_shape, accumulate=accumulate_into_main_grad, ) @@ -886,6 +949,7 @@ def __init__( self._cudnn_geglu_alpha: float = activation._clamped.alpha self._cudnn_glu_clamp_max: float = activation._clamped.limit self._cudnn_glu_clamp_min: float = -activation._clamped.limit + self._single_group_split_metadata_cache = {} def fuser_forward( self, @@ -976,20 +1040,48 @@ def fuser_forward( raise ValueError(f"Expected {num_groups} splits, but got {int(split_sizes.numel())}.") # Prepare split metadata - split_sizes, ( - split_points, - base_split_offsets, - fc1_x_tensor_offsets, - fc2_x_tensor_offsets, - fc2_out_tensor_offsets, - ) = tex.splits_to_offsets_multi( - split_sizes, - device, - strides=[1, 1, fc1_weight_shape[1], fc2_weight_shape[1], fc2_weight_shape[0]], - include_leading_zero=[False, True, True, True, True], - dtypes=[torch.int32, torch.int64, torch.int64, torch.int64, torch.int64], - bulk_allocate=True, - ) + if num_groups == 1: + metadata_key = ( + device.type, + device.index, + in_shape[0], + fc1_weight_shape[1], + fc2_weight_shape[1], + fc2_weight_shape[0], + ) + if metadata_key not in self._single_group_split_metadata_cache: + self._single_group_split_metadata_cache[metadata_key] = ( + _single_group_split_metadata( + in_shape[0], + device, + fc1_in_features=fc1_weight_shape[1], + fc2_in_features=fc2_weight_shape[1], + fc2_out_features=fc2_weight_shape[0], + ) + ) + ( + split_sizes, + split_points, + base_split_offsets, + fc1_x_tensor_offsets, + fc2_x_tensor_offsets, + fc2_out_tensor_offsets, + ) = self._single_group_split_metadata_cache[metadata_key] + else: + split_sizes, ( + split_points, + base_split_offsets, + fc1_x_tensor_offsets, + fc2_x_tensor_offsets, + fc2_out_tensor_offsets, + ) = tex.splits_to_offsets_multi( + split_sizes, + device, + strides=[1, 1, fc1_weight_shape[1], fc2_weight_shape[1], fc2_weight_shape[0]], + include_leading_zero=[False, True, True, True, True], + dtypes=[torch.int32, torch.int64, torch.int64, torch.int64, torch.int64], + bulk_allocate=True, + ) # Extract per-row activation probabilities from the middle op. scales = basic_op_extra_inputs[1][0] @@ -1366,7 +1458,7 @@ def fuser_forward( fc2_w_single = grouped_fc2_weight.split_into_quantized_tensors()[0] else: fc2_w_single = grouped_fc2_weight[0] - fc2_x_single = _nvfp4_single_tensor_from_grouped( + fc2_x_single = _single_quantized_tensor_from_grouped( grouped_fc2_x, fc2_input_quantizer, fp4_dtype=fc2_w_single._fp4_dtype, @@ -2048,7 +2140,7 @@ def fuser_backward( fc1_w_single = grouped_fc1_weight.split_into_quantized_tensors()[0] else: fc1_w_single = grouped_fc1_weight[0] - fc1_dy_single = _nvfp4_single_tensor_from_grouped(grouped_fc1_dy) + fc1_dy_single = _single_quantized_tensor_from_grouped(grouped_fc1_dy) general_gemm( fc1_w_single, fc1_dy_single, From 719c199f31f73dfd8c663bed96db9ece46958d4d Mon Sep 17 00:00:00 2001 From: Siddhartha Raman Sundara Raman Date: Mon, 27 Jul 2026 07:29:22 -0700 Subject: [PATCH 02/14] Optimize single-group MXFP8 grouped MLP paths Signed-off-by: Siddhartha Raman Sundara Raman --- .../pytorch/ops/fused/grouped_mlp.py | 724 +++++++++++++----- 1 file changed, 523 insertions(+), 201 deletions(-) diff --git a/transformer_engine/pytorch/ops/fused/grouped_mlp.py b/transformer_engine/pytorch/ops/fused/grouped_mlp.py index a3b3c190e14..efa7bd2e0f0 100644 --- a/transformer_engine/pytorch/ops/fused/grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/grouped_mlp.py @@ -101,6 +101,22 @@ def _nvidia_cudnn_frontend_supports_wgrad() -> bool: return _cudnn_frontend_version_supported() +def _single_group_mlp_optimization_enabled(disable_env: str) -> bool: + """Whether to use one optimized single-group grouped-MLP path.""" + if int(os.getenv("NVTE_DISABLE_SINGLE_GROUP_MLP_OPT", "0")) != 0: + return False + return int(os.getenv(disable_env, "0")) == 0 + + +def _allocate_dprob_tensor(scales_tensor: torch.Tensor) -> torch.Tensor: + """Allocate the cuDNN DGLU dprob accumulator. + + Kept as a small helper so tests can inject nonzero initial contents and + verify cuDNN's zero-initialization requirement. + """ + return torch.zeros_like(scales_tensor) + + def _wrap_single_quantized_as_grouped( tensor: torch.Tensor, quantized: MXFP8Tensor | NVFP4Tensor | NVFP4TensorStorage, @@ -135,12 +151,11 @@ def _wrap_single_quantized_as_grouped( else: k_dim = tensor.shape[-1] - if tensor_offsets is None: - tensor_offsets = torch.tensor( - [0, m_dim * k_dim], - dtype=torch.int64, - device=tensor.device, - ) + if tensor_offsets is None and not _single_group_mlp_optimization_enabled( + "NVTE_DISABLE_SINGLE_GROUP_MLP_OFFSETLESS_WRAPPER_OPT" + ): + tensor_offsets = torch.zeros(2, dtype=torch.int64, device=tensor.device) + tensor_offsets[1] = m_dim * k_dim return GroupedTensor( shape=(m_dim, k_dim), @@ -178,7 +193,12 @@ def _group_quantize_for_grouped_mlp( tensor_offsets=tensor_offsets, ) - if not isinstance(quantizer, (MXFP8Quantizer, NVFP4Quantizer)): + if not isinstance(quantizer, (MXFP8Quantizer, NVFP4Quantizer)) or ( + isinstance(quantizer, MXFP8Quantizer) + and not _single_group_mlp_optimization_enabled( + "NVTE_DISABLE_SINGLE_GROUP_MLP_QUANT_OPT" + ) + ): return tex.group_quantize( tensor, quantizer, @@ -376,24 +396,18 @@ def _single_group_split_metadata( fc2_out_features: int, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: """Create split/offset metadata for num_groups=1 without a GPU prefix-sum kernel.""" - split_sizes = torch.tensor([num_tokens], dtype=torch.int64, device=device) - split_points = torch.tensor([num_tokens], dtype=torch.int32, device=device) - base_split_offsets = torch.tensor([0, num_tokens], dtype=torch.int64, device=device) - fc1_x_tensor_offsets = torch.tensor( - [0, num_tokens * fc1_in_features], - dtype=torch.int64, - device=device, - ) - fc2_x_tensor_offsets = torch.tensor( - [0, num_tokens * fc2_in_features], - dtype=torch.int64, - device=device, - ) - fc2_out_tensor_offsets = torch.tensor( - [0, num_tokens * fc2_out_features], - dtype=torch.int64, - device=device, - ) + split_sizes = torch.zeros(1, dtype=torch.int64, device=device) + split_sizes[0] = num_tokens + split_points = torch.zeros(1, dtype=torch.int32, device=device) + split_points[0] = num_tokens + base_split_offsets = torch.zeros(2, dtype=torch.int64, device=device) + base_split_offsets[1] = num_tokens + fc1_x_tensor_offsets = torch.zeros(2, dtype=torch.int64, device=device) + fc1_x_tensor_offsets[1] = num_tokens * fc1_in_features + fc2_x_tensor_offsets = torch.zeros(2, dtype=torch.int64, device=device) + fc2_x_tensor_offsets[1] = num_tokens * fc2_in_features + fc2_out_tensor_offsets = torch.zeros(2, dtype=torch.int64, device=device) + fc2_out_tensor_offsets[1] = num_tokens * fc2_out_features return ( split_sizes, split_points, @@ -431,6 +445,70 @@ def _single_group_wgrad_gemm( ) +def _single_group_fc2_gemm( + grouped_x: GroupedTensor, + grouped_weight, + input_quantizer: MXFP8Quantizer | NVFP4Quantizer, + out: torch.Tensor, + *, + single_grouped_weight: bool, + bias: Optional[torch.Tensor], + bias_scale: Optional[torch.Tensor], + dtype: torch.dtype, +) -> torch.Tensor: + """Run one-group MXFP8/NVFP4 FC2 with regular GEMM.""" + if single_grouped_weight: + weight = _single_quantized_tensor_from_grouped(grouped_weight) + else: + weight = grouped_weight[0] + + fp4_dtype = weight._fp4_dtype if isinstance(weight, NVFP4Tensor) else None + x = _single_quantized_tensor_from_grouped( + grouped_x, + input_quantizer, + fp4_dtype=fp4_dtype, + ) + general_gemm( + weight, + x, + out_dtype=dtype, + out=out, + layout="TN", + use_split_accumulator=False, + ) + + if bias is not None: + token_bias = bias.transpose(0, 1).contiguous().expand(out.shape[0], -1) + if bias_scale is not None: + out = out + token_bias * bias_scale.view(-1, 1) + else: + out = out + token_bias + return out + + +def _single_group_dgrad_gemm( + grouped_dy: GroupedTensor, + grouped_weight, + out: torch.Tensor, + *, + single_grouped_weight: bool, + dtype: torch.dtype, +) -> None: + """Run one-group MXFP8/NVFP4 dgrad with regular GEMM.""" + if single_grouped_weight: + weight = _single_quantized_tensor_from_grouped(grouped_weight) + else: + weight = grouped_weight[0] + dy = _single_quantized_tensor_from_grouped(grouped_dy) + general_gemm( + weight, + dy, + out_dtype=dtype, + out=out, + layout="NN", + ) + + def _cudnn_compute_wgrad( grouped_x: GroupedTensor, grouped_dy: GroupedTensor, @@ -665,7 +743,27 @@ def _compute_grad_params( raise RuntimeError( "distributed-weight fused grouped-MLP requires delay_wgrad_compute=False." ) - if cudnn_wgrad_kernel_fn is not None: + if ( + num_groups == 1 + and isinstance(grouped_x, (GroupedTensor, GroupedTensorStorage)) + and isinstance(grouped_dy, (GroupedTensor, GroupedTensorStorage)) + and ( + isinstance(grouped_x.quantizer, NVFP4Quantizer) + or ( + isinstance(grouped_x.quantizer, MXFP8Quantizer) + and _single_group_mlp_optimization_enabled( + "NVTE_DISABLE_SINGLE_GROUP_MLP_WGRAD_OPT" + ) + ) + ) + and isinstance(grouped_dy.quantizer, grouped_x.quantizer.__class__) + ): + gemm_fn = functools.partial( + _single_group_wgrad_gemm, + weight_shape=weight_shape, + accumulate=accumulate_into_main_grad, + ) + elif cudnn_wgrad_kernel_fn is not None: offsets = offsets if offsets.dtype == torch.int32 else offsets.to(dtype=torch.int32) gemm_fn = functools.partial( _cudnn_compute_wgrad, @@ -679,18 +777,6 @@ def _compute_grad_params( scale_view_dtype=scale_view_dtype, sf_vec_size=sf_vec_size, ) - elif ( - num_groups == 1 - and isinstance(grouped_x, GroupedTensor) - and isinstance(grouped_dy, GroupedTensor) - and isinstance(grouped_x.quantizer, (MXFP8Quantizer, NVFP4Quantizer)) - and isinstance(grouped_dy.quantizer, grouped_x.quantizer.__class__) - ): - gemm_fn = functools.partial( - _single_group_wgrad_gemm, - weight_shape=weight_shape, - accumulate=accumulate_into_main_grad, - ) else: gemm_fn = functools.partial( general_grouped_gemm_for_grouped_tensor, @@ -1039,8 +1125,60 @@ def fuser_forward( if int(split_sizes.numel()) != num_groups: raise ValueError(f"Expected {num_groups} splits, but got {int(split_sizes.numel())}.") + # Extract per-row activation probabilities from the middle op. + scales = basic_op_extra_inputs[1][0] + unit_activation_scale = bool( + getattr(self.basic_ops[1], "_grouped_mlp_unit_activation_scale", False) + ) + if unit_activation_scale and num_groups != 1: + raise ValueError( + "Unit activation scaling is only supported for a single-group grouped MLP." + ) + + # Shared experts have one dense group and all optimized kernels derive M + # from their runtime tensor shapes. Reuse the caller-owned split tensor + # as the ignored cuDNN padded-offset argument and omit tensor offsets + # entirely. This avoids both splits_to_offsets and cached CUDA pointers. + use_offsetless_metadata = ( + num_groups == 1 + and unit_activation_scale + and isinstance(fc1_input_quantizer, MXFP8Quantizer) + and _single_group_mlp_optimization_enabled( + "NVTE_DISABLE_SINGLE_GROUP_MLP_OFFSETLESS_METADATA_OPT" + ) + and _single_group_mlp_optimization_enabled( + "NVTE_DISABLE_SINGLE_GROUP_MLP_OFFSETLESS_WRAPPER_OPT" + ) + and _single_group_mlp_optimization_enabled( + "NVTE_DISABLE_SINGLE_GROUP_MLP_QUANT_OPT" + ) + and _single_group_mlp_optimization_enabled( + "NVTE_DISABLE_SINGLE_GROUP_MLP_MXFP8_FC2_GEMM_OPT" + ) + and _single_group_mlp_optimization_enabled( + "NVTE_DISABLE_SINGLE_GROUP_MLP_MXFP8_DGRAD_GEMM_OPT" + ) + and _single_group_mlp_optimization_enabled( + "NVTE_DISABLE_SINGLE_GROUP_MLP_WGRAD_OPT" + ) + ) + # Prepare split metadata - if num_groups == 1: + if use_offsetless_metadata: + # cuDNN requires an int32 tensor descriptor, although its + # use_single_group_runtime_offsets specialization never loads the + # pointer. This view aliases the live caller-owned [M] int64 tensor, + # has value M, and requires no allocation or CUDA kernel. + split_points = split_sizes.view(torch.int32)[:1] + # Backward saves this slot for the generic path. The optimized + # shared-expert path never consumes it. + base_split_offsets = split_sizes + fc1_x_tensor_offsets = None + fc2_x_tensor_offsets = None + fc2_out_tensor_offsets = None + elif num_groups == 1 and _single_group_mlp_optimization_enabled( + "NVTE_DISABLE_SINGLE_GROUP_MLP_METADATA_OPT" + ): metadata_key = ( device.type, device.index, @@ -1083,9 +1221,6 @@ def fuser_forward( bulk_allocate=True, ) - # Extract per-row activation probabilities from the middle op. - scales = basic_op_extra_inputs[1][0] - # Prepare FC1 grouped weight tensor for fused kernels. # - single_grouped_weight=True: op.weight is already a GroupedTensor # - single_grouped_weight=False: cute DSL kernel works with discrete weight tensors @@ -1161,8 +1296,26 @@ def fuser_forward( ): grouped_fc1_weight._with_gemm_swizzled_scales = False + # The canonical shared-expert TELinear saves its original BF16 input and + # creates the columnwise MXFP8 representation just before FC1 wgrad. + # Preserve that policy for the single-group fused path so fprop only + # pays for the rowwise representation required by the Rubin kernel. + defer_fc1_input_columnwise = ( + num_groups == 1 + and unit_activation_scale + and weight_requires_grad + and isinstance(fc1_input_quantizer, MXFP8Quantizer) + and not is_quantized_tensor(input_) + and _single_group_mlp_optimization_enabled( + "NVTE_DISABLE_SINGLE_GROUP_MLP_DEFER_FC1_COLUMNWISE_OPT" + ) + ) + # Group-quantize input tensor and convert dtypes if needed - fc1_input_quantizer.set_usage(rowwise=True, columnwise=weight_requires_grad) + fc1_input_quantizer.set_usage( + rowwise=True, + columnwise=weight_requires_grad and not defer_fc1_input_columnwise, + ) fc1_input_quantizer.optimize_for_gemm = True fc1_input_quantizer.internal = True input_quantizer = getattr(input_, "quantizer", None) @@ -1276,9 +1429,11 @@ def fuser_forward( fc2_bias_packed = _pack_grouped_linear_bias_for_cudnn(fc2_op) fc1_d_dtype = torch.bfloat16 if use_nvfp4 else torch.float8_e4m3fn - fc1_prob_tensor = ( - scales.detach().to(dtype=torch.float32 if use_nvfp4 else dtype).reshape(-1, 1, 1) - ) + fc1_prob_tensor = None + if not unit_activation_scale: + fc1_prob_tensor = ( + scales.detach().to(dtype=torch.float32 if use_nvfp4 else dtype).reshape(-1, 1, 1) + ) fc1_norm_const_tensor = None if use_nvfp4 else norm_const_tensor if use_nvfp4: nvfp4_fp4_max = 6.0 @@ -1338,6 +1493,7 @@ def fuser_forward( else: fc1_activation_kwargs["norm_const_tensor"] = fc1_norm_const_tensor fc1_activation_kwargs["discrete_col_sfd"] = not use_nvfp4 + fc1_activation_kwargs["use_single_group_runtime_offsets"] = num_groups == 1 if self._pass_geglu_runtime_params: fc1_activation_kwargs.update( linear_offset=self._cudnn_linear_offset, @@ -1349,16 +1505,35 @@ def fuser_forward( if fc1_op.single_grouped_weight: # Clone and swizzle scales for GEMM. fc1_weight_for_gemm = grouped_fc1_weight.copy() - tex.grouped_swizzle_for_gemm(fc1_weight_for_gemm, rowwise=True, columnwise=False) + use_single_group_weight_swizzle = ( + num_groups == 1 + and _single_group_mlp_optimization_enabled( + "NVTE_DISABLE_SINGLE_GROUP_MLP_FC1_WEIGHT_SWIZZLE_OPT" + ) + ) + if use_single_group_weight_swizzle: + fc1_weight_single = _single_quantized_tensor_from_grouped(fc1_weight_for_gemm) + fc1_weight_single._columnwise_data = None + fc1_weight_single._columnwise_scale_inv = None + tex.swizzle_scales_for_gemm_(fc1_weight_single) + fc1_w_data = fc1_weight_single._rowwise_data + fc1_w_scales = fc1_weight_single._rowwise_scale_inv + else: + tex.grouped_swizzle_for_gemm( + fc1_weight_for_gemm, + rowwise=True, + columnwise=False, + ) + fc1_w_data = fc1_weight_for_gemm.rowwise_data + fc1_w_scales = fc1_weight_for_gemm.scale_inv # Pack weight tensors for stacked kernel # Data actual shape: (num_groups, n, k) # Data logical shape: (n, k, num_groups) - fc1_w_data = fc1_weight_for_gemm.rowwise_data fc1_w_data = fc1_w_data.view(dtype=data_dtype) fc1_w_data = fc1_w_data.view(num_groups, fc1_weight_shape[0], fc1_weight_k) fc1_w_data = fc1_w_data.permute(1, 2, 0) - fc1_w_scales = fc1_weight_for_gemm.scale_inv.view(dtype=scale_view_dtype) + fc1_w_scales = fc1_w_scales.view(dtype=scale_view_dtype) fc1_w_scales = fc1_w_scales.view( num_groups, ceil_div(fc1_weight_shape[0], 128), @@ -1372,20 +1547,62 @@ def fuser_forward( fc1_activation_kwargs["b_tensor"] = fc1_w_data fc1_activation_kwargs["sfb_tensor"] = fc1_w_scales else: - # Discrete-weight kernel: per-expert data/scale pointers - fc1_b_ptrs, fc1_sfb_ptrs, _fc1_sfb_buffer = ( - tex.grouped_mlp_experimental.swizzle_scales_and_pack_ptrs_for_discrete_weights( - [w._rowwise_data for w in grouped_fc1_weight], - [w._rowwise_scale_inv for w in grouped_fc1_weight], - "nvfp4" if use_nvfp4 else "mxfp8_rowwise", - device, + use_single_discrete_weight = ( + num_groups == 1 + and _single_group_mlp_optimization_enabled( + "NVTE_DISABLE_SINGLE_GROUP_MLP_FC1_DISCRETE_WEIGHT_OPT" ) ) - fc1_activation_kwargs["b_ptrs"] = fc1_b_ptrs - fc1_activation_kwargs["sfb_ptrs"] = fc1_sfb_ptrs - fc1_activation_kwargs["n"] = fc1_weight_shape[0] - fc1_activation_kwargs["b_dtype"] = data_dtype - fc1_activation_kwargs["b_major"] = "k" + if use_single_discrete_weight: + fc1_weight_single = grouped_fc1_weight[0] + original_rowwise_scale = fc1_weight_single._rowwise_scale_inv + original_columnwise_data = fc1_weight_single._columnwise_data + original_columnwise_scale = fc1_weight_single._columnwise_scale_inv + original_swizzled = fc1_weight_single._with_gemm_swizzled_scales + fc1_weight_single._columnwise_data = None + fc1_weight_single._columnwise_scale_inv = None + tex.swizzle_scales_for_gemm_(fc1_weight_single) + swizzled_rowwise_scale = fc1_weight_single._rowwise_scale_inv + fc1_weight_single._rowwise_scale_inv = original_rowwise_scale + fc1_weight_single._columnwise_data = original_columnwise_data + fc1_weight_single._columnwise_scale_inv = original_columnwise_scale + fc1_weight_single._with_gemm_swizzled_scales = original_swizzled + + fc1_w_data = fc1_weight_single._rowwise_data.view(dtype=data_dtype) + fc1_w_data = fc1_w_data.view( + 1, + fc1_weight_shape[0], + fc1_weight_k, + ) + fc1_w_data = fc1_w_data.permute(1, 2, 0) + fc1_w_scales = swizzled_rowwise_scale.view(dtype=scale_view_dtype) + fc1_w_scales = fc1_w_scales.view( + 1, + ceil_div(fc1_weight_shape[0], 128), + ceil_div(fc1_weight_shape[1], k_sf_divisor), + 32, + 4, + 4, + ) + fc1_w_scales = fc1_w_scales.permute(3, 4, 1, 5, 2, 0) + fc1_activation_kwargs["b_tensor"] = fc1_w_data + fc1_activation_kwargs["sfb_tensor"] = fc1_w_scales + else: + # Discrete-weight kernel: per-expert data/scale pointers + fc1_b_ptrs, fc1_sfb_ptrs, _fc1_sfb_buffer = ( + tex.grouped_mlp_experimental + .swizzle_scales_and_pack_ptrs_for_discrete_weights( + [w._rowwise_data for w in grouped_fc1_weight], + [w._rowwise_scale_inv for w in grouped_fc1_weight], + "nvfp4" if use_nvfp4 else "mxfp8_rowwise", + device, + ) + ) + fc1_activation_kwargs["b_ptrs"] = fc1_b_ptrs + fc1_activation_kwargs["sfb_ptrs"] = fc1_sfb_ptrs + fc1_activation_kwargs["n"] = fc1_weight_shape[0] + fc1_activation_kwargs["b_dtype"] = data_dtype + fc1_activation_kwargs["b_major"] = "k" if use_fc1_act_hadamard: fc1_kernel_out = self.grouped_gemm_act_hadamard_kernel()(**fc1_activation_kwargs) @@ -1454,31 +1671,16 @@ def fuser_forward( and grouped_fc2_x.columnwise_data is not None and grouped_fc2_x.columnwise_scale_inv is not None ): - if fc2_op.single_grouped_weight: - fc2_w_single = grouped_fc2_weight.split_into_quantized_tensors()[0] - else: - fc2_w_single = grouped_fc2_weight[0] - fc2_x_single = _single_quantized_tensor_from_grouped( + fc2_out_buf = _single_group_fc2_gemm( grouped_fc2_x, + grouped_fc2_weight, fc2_input_quantizer, - fp4_dtype=fc2_w_single._fp4_dtype, - ) - general_gemm( - fc2_w_single, - fc2_x_single, - out_dtype=dtype, - out=fc2_out_buf, - layout="TN", - use_split_accumulator=False, + fc2_out_buf, + single_grouped_weight=fc2_op.single_grouped_weight, + bias=fc2_bias_packed, + bias_scale=fc2_scales, + dtype=dtype, ) - if fc2_bias_packed is not None: - token_bias = ( - fc2_bias_packed.transpose(0, 1).contiguous().expand(in_shape[0], -1) - ) - if fc2_scales is not None: - fc2_out_buf += token_bias * fc2_scales.view(-1, 1) - else: - fc2_out_buf += token_bias else: fc2_out_grouped = GroupedTensorStorage( shape=(in_shape[0], fc2_weight_shape[0]), @@ -1523,77 +1725,110 @@ def fuser_forward( with_gemm_swizzled_scales=True, ) - fc2_scales_tensor = ( - fc2_scales.detach().to(dtype=torch.float32).reshape(-1, 1, 1) - if fc2_scales is not None - else torch.ones((in_shape[0], 1, 1), dtype=torch.float32, device=device) + use_single_group_dense_fc2 = ( + num_groups == 1 + and _single_group_mlp_optimization_enabled( + "NVTE_DISABLE_SINGLE_GROUP_MLP_MXFP8_FC2_GEMM_OPT" + ) ) - fc2_quant_kwargs = { - "a_tensor": fc1_kernel_out["d_tensor"], - "sfa_tensor": fc1_kernel_out["sfd_row_tensor"], - "padded_offsets": split_points, - "alpha_tensor": alpha_tensor, - "bias_tensor": fc2_bias_packed, - "norm_const_tensor": None, - "prob_tensor": fc2_scales_tensor, - "acc_dtype": torch.float32, - "d_dtype": dtype, - "cd_major": "n", - "sf_vec_size": MXFP8_BLOCK_SCALING_SIZE, - "current_stream": current_stream, - "use_dynamic_sched": True, - } - - if fc2_op.single_grouped_weight: - # Clone and swizzle scales for GEMM (original stays unmodified for save_for_backward) - fc2_weight_for_gemm = grouped_fc2_weight.copy() - tex.grouped_swizzle_for_gemm(fc2_weight_for_gemm, rowwise=True, columnwise=False) - - fc2_w_data = fc2_weight_for_gemm.rowwise_data - fc2_w_data = fc2_w_data.view(dtype=torch.float8_e4m3fn) - fc2_w_data = fc2_w_data.view(num_groups, fc2_weight_shape[0], fc2_weight_shape[1]) - fc2_w_data = fc2_w_data.permute(1, 2, 0) - - fc2_w_scales = fc2_weight_for_gemm.scale_inv.view(dtype=torch.float8_e8m0fnu) - fc2_w_scales = fc2_w_scales.view( - num_groups, - ceil_div(fc2_weight_shape[0], 128), - ceil_div(fc2_weight_shape[1], 128), - MXFP8_BLOCK_SCALING_SIZE, - 4, - 4, + 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( + grouped_fc2_x, + grouped_fc2_weight, + fc2_input_quantizer, + fc2_out_buf, + single_grouped_weight=fc2_op.single_grouped_weight, + bias=fc2_bias_packed, + bias_scale=fc2_scales, + dtype=dtype, ) - fc2_w_scales = fc2_w_scales.permute(3, 4, 1, 5, 2, 0) - fc2_quant_kwargs["b_tensor"] = fc2_w_data - fc2_quant_kwargs["sfb_tensor"] = fc2_w_scales else: - fc2_b_ptrs, fc2_sfb_ptrs, _fc2_sfb_buffer = ( - tex.grouped_mlp_experimental.swizzle_scales_and_pack_ptrs_for_discrete_weights( - [w._rowwise_data for w in grouped_fc2_weight], - [w._rowwise_scale_inv for w in grouped_fc2_weight], - "nvfp4" if use_nvfp4 else "mxfp8_rowwise", - device, + fc2_scales_tensor = ( + fc2_scales.detach().to(dtype=torch.float32).reshape(-1, 1, 1) + if fc2_scales is not None + else torch.ones((in_shape[0], 1, 1), dtype=torch.float32, device=device) + ) + fc2_quant_kwargs = { + "a_tensor": fc1_kernel_out["d_tensor"], + "sfa_tensor": fc1_kernel_out["sfd_row_tensor"], + "padded_offsets": split_points, + "alpha_tensor": alpha_tensor, + "bias_tensor": fc2_bias_packed, + "norm_const_tensor": None, + "prob_tensor": fc2_scales_tensor, + "acc_dtype": torch.float32, + "d_dtype": dtype, + "cd_major": "n", + "sf_vec_size": MXFP8_BLOCK_SCALING_SIZE, + "current_stream": current_stream, + "use_dynamic_sched": True, + "use_single_group_runtime_offsets": num_groups == 1, + } + + if fc2_op.single_grouped_weight: + # Clone and swizzle scales for GEMM + # (original stays unmodified for save_for_backward). + fc2_weight_for_gemm = grouped_fc2_weight.copy() + tex.grouped_swizzle_for_gemm( + fc2_weight_for_gemm, + rowwise=True, + columnwise=False, ) + + fc2_w_data = fc2_weight_for_gemm.rowwise_data + fc2_w_data = fc2_w_data.view(dtype=torch.float8_e4m3fn) + fc2_w_data = fc2_w_data.view( + num_groups, + fc2_weight_shape[0], + fc2_weight_shape[1], + ) + fc2_w_data = fc2_w_data.permute(1, 2, 0) + + fc2_w_scales = fc2_weight_for_gemm.scale_inv.view( + dtype=torch.float8_e8m0fnu + ) + fc2_w_scales = fc2_w_scales.view( + num_groups, + ceil_div(fc2_weight_shape[0], 128), + ceil_div(fc2_weight_shape[1], 128), + MXFP8_BLOCK_SCALING_SIZE, + 4, + 4, + ) + fc2_w_scales = fc2_w_scales.permute(3, 4, 1, 5, 2, 0) + fc2_quant_kwargs["b_tensor"] = fc2_w_data + fc2_quant_kwargs["sfb_tensor"] = fc2_w_scales + else: + fc2_b_ptrs, fc2_sfb_ptrs, _fc2_sfb_buffer = ( + tex.grouped_mlp_experimental + .swizzle_scales_and_pack_ptrs_for_discrete_weights( + [w._rowwise_data for w in grouped_fc2_weight], + [w._rowwise_scale_inv for w in grouped_fc2_weight], + "mxfp8_rowwise", + device, + ) + ) + fc2_quant_kwargs["b_ptrs"] = fc2_b_ptrs + fc2_quant_kwargs["sfb_ptrs"] = fc2_sfb_ptrs + fc2_quant_kwargs["n"] = fc2_weight_shape[0] + fc2_quant_kwargs["b_dtype"] = torch.float8_e4m3fn + fc2_quant_kwargs["b_major"] = "k" + + fc2_quant_kwargs["d_tensor"] = fc2_out_buf.as_strided( + (in_shape[0], fc2_weight_shape[0], 1), + (fc2_weight_shape[0], 1, in_shape[0] * fc2_weight_shape[0]), ) - fc2_quant_kwargs["b_ptrs"] = fc2_b_ptrs - fc2_quant_kwargs["sfb_ptrs"] = fc2_sfb_ptrs - fc2_quant_kwargs["n"] = fc2_weight_shape[0] - fc2_quant_kwargs["b_dtype"] = torch.float8_e4m3fn - fc2_quant_kwargs["b_major"] = "k" - - # Always allocate the output (the caller's buffer if provided, else a fresh one) and - # pass it as the kernel's d_tensor, so the kernel writes in place and the call is uniform. - output_buffer = validate_or_alloc_output(output_buffer, fc2_out_shape, dtype, device) - fc2_quant_kwargs["d_tensor"] = output_buffer.as_strided( - (in_shape[0], fc2_weight_shape[0], 1), - (fc2_weight_shape[0], 1, in_shape[0] * fc2_weight_shape[0]), - ) - self.grouped_gemm_quant_kernel()(**fc2_quant_kwargs) - fc2_out = output_buffer + self.grouped_gemm_quant_kernel()(**fc2_quant_kwargs) + fc2_out = fc2_out_buf # Save state for backward pass if requires_grad: - mark_grouped_tensor(grouped_fc1_x, activation_in, scales, grouped_fc2_x) + # The deferred path saves the caller-owned BF16 input in the slot + # normally occupied by the columnwise grouped tensor. Backward + # converts it immediately before FC1 wgrad. + saved_fc1_x = input_ if defer_fc1_input_columnwise else grouped_fc1_x + mark_grouped_tensor(saved_fc1_x, activation_in, scales, grouped_fc2_x) activation_op = self.basic_ops[1] cpu_offloading = is_cpu_offload_enabled() activation_is_srelu = isinstance(activation_op, ScaledSReLU) @@ -1612,14 +1847,14 @@ def fuser_forward( # MXFP8 wgrad only needs columnwise tiles. NVFP4 generic GEMM fallbacks # need the full grouped tensor state, including rowwise data and amax. if not use_nvfp4: - for grouped_fc_x in (grouped_fc1_x, saved_grouped_fc2_x): - if grouped_fc_x is not None: + for grouped_fc_x in (saved_fc1_x, saved_grouped_fc2_x): + if isinstance(grouped_fc_x, (GroupedTensor, GroupedTensorStorage)): grouped_fc_x.rowwise_data = None grouped_fc_x.scale_inv = None if cpu_offloading: activation_tensors = [ - t for t in (grouped_fc1_x, activation_in, saved_grouped_fc2_x) if t is not None + t for t in (saved_fc1_x, activation_in, saved_grouped_fc2_x) if t is not None ] start_offload(*activation_tensors) mark_activation_offload(*activation_tensors) @@ -1641,7 +1876,7 @@ def fuser_forward( split_sizes, base_split_offsets, split_points, - grouped_fc1_x, + saved_fc1_x, *fc1_weight_tensors, activation_in, scales, @@ -1654,6 +1889,8 @@ def fuser_forward( fc1_ctx.dtype = dtype fc1_ctx.input_requires_grad = input_requires_grad fc1_ctx.weight_requires_grad = weight_requires_grad + fc1_ctx.unit_activation_scale = unit_activation_scale + fc1_ctx.defer_fc1_input_columnwise = defer_fc1_input_columnwise fc2_ctx.input_quantizers = [fc2_input_quantizer] fc2_ctx.grad_output_quantizers = [fc2_grad_output_quantizer] @@ -1733,6 +1970,18 @@ def fuser_backward( if not fc1_ctx.weight_requires_grad: grouped_fc1_x = None + elif bool(getattr(fc1_ctx, "defer_fc1_input_columnwise", False)): + fc1_input_quantizer = fc1_ctx.input_quantizers[0] + fc1_input_quantizer.set_usage(rowwise=False, columnwise=True) + fc1_input_quantizer.optimize_for_gemm = True + fc1_input_quantizer.internal = True + grouped_fc1_x = _group_quantize_for_grouped_mlp( + grouped_fc1_x, + fc1_input_quantizer, + num_groups, + split_sizes, + tensor_offsets=None, + ) if not fc2_ctx.weight_requires_grad: grouped_fc2_x = None @@ -1772,7 +2021,11 @@ def fuser_backward( fc2_grad_output_quantizer, num_groups, split_sizes, - tensor_offsets=base_split_offsets * fc2_weight_shape[0], + tensor_offsets=( + None + if num_groups == 1 + else base_split_offsets * fc2_weight_shape[0] + ), ) use_nvfp4 = ( @@ -1838,9 +2091,14 @@ def fuser_backward( norm_const_tensor = get_cached_ones_tensor(1, torch.float32, device) current_stream = torch.cuda.current_stream().cuda_stream - scales_f32 = scales.detach().to(dtype=torch.float32) - scales_tensor = scales_f32.reshape(-1, 1, 1) - dscales_tensor = torch.zeros_like(scales_tensor) + unit_activation_scale = bool(getattr(fc1_ctx, "unit_activation_scale", False)) + scales_f32 = None + scales_tensor = None + dscales_tensor = None + if not unit_activation_scale: + scales_f32 = scales.detach().to(dtype=torch.float32) + scales_tensor = scales_f32.reshape(-1, 1, 1) + dscales_tensor = _allocate_dprob_tensor(scales_tensor) fc2_d_dtype = torch.bfloat16 if use_nvfp4 else torch.float8_e4m3fn if use_nvfp4: @@ -1886,6 +2144,7 @@ def fuser_backward( "current_stream": current_stream, "discrete_col_sfd": not use_nvfp4, "use_dynamic_sched": True, + "use_single_group_runtime_offsets": num_groups == 1, } if self._cudnn_dact_func is not None: fc2_dactivation_kwargs["beta_tensor"] = fc2_beta_tensor @@ -1933,19 +2192,69 @@ def fuser_backward( fc2_dactivation_kwargs["b_tensor"] = fc2_w_data fc2_dactivation_kwargs["sfb_tensor"] = fc2_w_scales else: - fc2_b_ptrs, fc2_sfb_ptrs, _fc2_sfb_buffer = ( - tex.grouped_mlp_experimental.swizzle_scales_and_pack_ptrs_for_discrete_weights( - [w._columnwise_data for w in grouped_fc2_weight], - [w._columnwise_scale_inv for w in grouped_fc2_weight], - "nvfp4" if use_nvfp4 else "mxfp8_columnwise", - device, + use_single_discrete_weight = ( + num_groups == 1 + and _single_group_mlp_optimization_enabled( + "NVTE_DISABLE_SINGLE_GROUP_MLP_FC2_DGRAD_DISCRETE_WEIGHT_OPT" ) ) - fc2_dactivation_kwargs["b_ptrs"] = fc2_b_ptrs - fc2_dactivation_kwargs["sfb_ptrs"] = fc2_sfb_ptrs - fc2_dactivation_kwargs["n"] = fc2_weight_shape[1] - fc2_dactivation_kwargs["b_dtype"] = data_dtype - fc2_dactivation_kwargs["b_major"] = "k" if use_nvfp4 else "n" + if use_single_discrete_weight: + fc2_weight_single = grouped_fc2_weight[0] + original_rowwise_data = fc2_weight_single._rowwise_data + original_rowwise_scale = fc2_weight_single._rowwise_scale_inv + original_columnwise_scale = fc2_weight_single._columnwise_scale_inv + original_swizzled = fc2_weight_single._with_gemm_swizzled_scales + fc2_weight_single._rowwise_data = None + fc2_weight_single._rowwise_scale_inv = None + tex.swizzle_scales_for_gemm_(fc2_weight_single) + swizzled_columnwise_scale = fc2_weight_single._columnwise_scale_inv + fc2_weight_single._rowwise_data = original_rowwise_data + fc2_weight_single._rowwise_scale_inv = original_rowwise_scale + fc2_weight_single._columnwise_scale_inv = original_columnwise_scale + fc2_weight_single._with_gemm_swizzled_scales = original_swizzled + + fc2_w_data = fc2_weight_single._columnwise_data.view(dtype=data_dtype) + fc2_w_data = fc2_w_data.view( + 1, + fc2_weight_shape[0], + fc2_weight_k, + ) + fc2_w_data = ( + fc2_w_data.permute(1, 2, 0) + if use_nvfp4 + else fc2_w_data.permute(2, 1, 0) + ) + fc2_w_scales = swizzled_columnwise_scale.view(dtype=scale_view_dtype) + fc2_w_scales = fc2_w_scales.view( + 1, + ceil_div(fc2_weight_shape[1], k_sf_divisor), + ceil_div(fc2_weight_shape[0], 128), + 32, + 4, + 4, + ) + fc2_w_scales = ( + fc2_w_scales.permute(3, 4, 2, 5, 1, 0) + if use_nvfp4 + else fc2_w_scales.permute(3, 4, 1, 5, 2, 0) + ) + fc2_dactivation_kwargs["b_tensor"] = fc2_w_data + fc2_dactivation_kwargs["sfb_tensor"] = fc2_w_scales + else: + fc2_b_ptrs, fc2_sfb_ptrs, _fc2_sfb_buffer = ( + tex.grouped_mlp_experimental + .swizzle_scales_and_pack_ptrs_for_discrete_weights( + [w._columnwise_data for w in grouped_fc2_weight], + [w._columnwise_scale_inv for w in grouped_fc2_weight], + "nvfp4" if use_nvfp4 else "mxfp8_columnwise", + device, + ) + ) + fc2_dactivation_kwargs["b_ptrs"] = fc2_b_ptrs + fc2_dactivation_kwargs["sfb_ptrs"] = fc2_sfb_ptrs + fc2_dactivation_kwargs["n"] = fc2_weight_shape[1] + fc2_dactivation_kwargs["b_dtype"] = data_dtype + fc2_dactivation_kwargs["b_major"] = "k" if use_nvfp4 else "n" fc2_dgrad_kernel_out = self.grouped_gemm_dactivation_kernel()(**fc2_dactivation_kwargs) @@ -1970,7 +2279,9 @@ def fuser_backward( fc1_dy_col_scale = ( fc2_dgrad_kernel_out["sfd_col_tensor"].permute(5, 2, 4, 0, 1, 3).view(-1) ) - grad_scales = fc2_dgrad_kernel_out["dprob_tensor"].view(-1) + grad_scales = fc2_dgrad_kernel_out["dprob_tensor"] + if grad_scales is not None: + grad_scales = grad_scales.view(-1) if recompute_fc2_x_from_dsrelu: d_srelu_tensor = fc2_dgrad_kernel_out.get("d_srelu_tensor") @@ -2058,7 +2369,9 @@ def fuser_backward( fc1_bias_grads = [dbias_2d[group_idx] for group_idx in range(num_groups)] # FC1 grad output for dgrad and wgrad GEMMs - fc1_dy_tensor_offsets = base_split_offsets * fc1_weight_shape[0] + fc1_dy_tensor_offsets = ( + None if num_groups == 1 else base_split_offsets * fc1_weight_shape[0] + ) fc1_grad_output_quantizer = fc1_ctx.grad_output_quantizers[0] if use_nvfp4: fc1_grad_output_quantizer.set_usage( @@ -2133,38 +2446,43 @@ def fuser_backward( if is_distributed_weight(fc1_leader): grouped_fc1_weight = materialize_weight_for_backward(fc1_leader) - if use_nvfp4: - grad_input = validate_or_alloc_output(grad_input_buffer, in_shape, dtype, device) - if num_groups == 1: - if fc1_op.single_grouped_weight: - fc1_w_single = grouped_fc1_weight.split_into_quantized_tensors()[0] - else: - fc1_w_single = grouped_fc1_weight[0] - fc1_dy_single = _single_quantized_tensor_from_grouped(grouped_fc1_dy) - general_gemm( - fc1_w_single, - fc1_dy_single, - out_dtype=dtype, - out=grad_input, - layout="NN", - ) - else: - fc1_x_tensor_offsets = base_split_offsets * fc1_weight_shape[1] - grouped_grad_input = GroupedTensor( - shape=(out_shape[0], fc1_weight_shape[1]), - dtype=dtype, - num_tensors=num_groups, - quantizer=None, - data=grad_input.view(-1), - first_dims=split_sizes, - tensor_offsets=fc1_x_tensor_offsets, - ) - general_grouped_gemm_for_grouped_tensor( - grouped_fc1_weight, - grouped_fc1_dy, - grouped_grad_input, - layout="NN", - ) + use_single_group_dense_dgrad = num_groups == 1 and ( + use_nvfp4 + or _single_group_mlp_optimization_enabled( + "NVTE_DISABLE_SINGLE_GROUP_MLP_MXFP8_DGRAD_GEMM_OPT" + ) + ) + if use_single_group_dense_dgrad: + grad_input = validate_or_alloc_output( + grad_input_buffer, in_shape, dtype, device + ) + _single_group_dgrad_gemm( + grouped_fc1_dy, + grouped_fc1_weight, + grad_input, + single_grouped_weight=fc1_op.single_grouped_weight, + dtype=dtype, + ) + elif use_nvfp4: + grad_input = validate_or_alloc_output( + grad_input_buffer, in_shape, dtype, device + ) + fc1_x_tensor_offsets = base_split_offsets * fc1_weight_shape[1] + grouped_grad_input = GroupedTensor( + shape=(out_shape[0], fc1_weight_shape[1]), + dtype=dtype, + num_tensors=num_groups, + quantizer=None, + data=grad_input.view(-1), + first_dims=split_sizes, + tensor_offsets=fc1_x_tensor_offsets, + ) + general_grouped_gemm_for_grouped_tensor( + grouped_fc1_weight, + grouped_fc1_dy, + grouped_grad_input, + layout="NN", + ) else: fc1_dgrad_a_data = fc2_dgrad_kernel_out["d_row_tensor"] fc1_dgrad_a_scales = fc2_dgrad_kernel_out["sfd_row_tensor"] @@ -2185,6 +2503,7 @@ def fuser_backward( "current_stream": current_stream, "discrete_col_sfd": True, "use_dynamic_sched": True, + "use_single_group_runtime_offsets": num_groups == 1, } if fc1_op.single_grouped_weight: @@ -2275,10 +2594,13 @@ def fuser_backward( grouped_fc1_x.columnwise_data, grouped_fc1_x.scale_inv, grouped_fc1_x.columnwise_scale_inv, - ) + ) fc2_grad_extra = (None, None) if fc2_op._scale_bias else (None,) - activation_grad_extra = (grad_scales,) if grad_scales is not None else () + if unit_activation_scale: + activation_grad_extra = (None,) + else: + activation_grad_extra = (grad_scales,) if grad_scales is not None else () return ( grad_input, [fc1_grad_params, (), fc2_grad_params], From c56f26bd8897e36c397454e9507533651ce5ce3c Mon Sep 17 00:00:00 2001 From: Siddhartha Raman Sundara Raman Date: Mon, 27 Jul 2026 07:33:36 -0700 Subject: [PATCH 03/14] Test single-group MXFP8 grouped MLP Signed-off-by: Siddhartha Raman Sundara Raman --- tests/pytorch/test_fusible_ops.py | 22 ++++++++++++++++++- .../pytorch/ops/fused/grouped_mlp.py | 4 ++-- 2 files changed, 23 insertions(+), 3 deletions(-) diff --git a/tests/pytorch/test_fusible_ops.py b/tests/pytorch/test_fusible_ops.py index 66857d81253..db4cb6dfb1e 100644 --- a/tests/pytorch/test_fusible_ops.py +++ b/tests/pytorch/test_fusible_ops.py @@ -3558,7 +3558,10 @@ def test_grouped_mlp( """GroupedLinear + scaled activation + GroupedLinear""" # Split sizes - split_sizes = [split_alignment * (i) for i in range(group_size)] + if group_size == 1: + split_sizes = [split_alignment] + else: + split_sizes = [split_alignment * i for i in range(group_size)] random.shuffle(split_sizes) split_sizes = torch.tensor(split_sizes, dtype=torch.int, device=device) @@ -3749,6 +3752,23 @@ def test_grouped_mlp_nvfp4_rht_srelu( activation="scaled_srelu", ) + @pytest.mark.parametrize("bias", (False, True)) + def test_grouped_mlp_single_group_mxfp8( + self, + *, + bias: bool, + device: torch.device = "cuda", + ) -> None: + """Single-group GroupedLinear + ScaledSwiGLU + GroupedLinear with MXFP8.""" + + self.test_grouped_mlp( + group_size=1, + bias=bias, + dtype=torch.bfloat16, + quantization="mxfp8", + device=device, + ) + class TestCustomOps: """Test with ops that are defined externally""" diff --git a/transformer_engine/pytorch/ops/fused/grouped_mlp.py b/transformer_engine/pytorch/ops/fused/grouped_mlp.py index efa7bd2e0f0..d286dada229 100644 --- a/transformer_engine/pytorch/ops/fused/grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/grouped_mlp.py @@ -480,9 +480,9 @@ def _single_group_fc2_gemm( if bias is not None: token_bias = bias.transpose(0, 1).contiguous().expand(out.shape[0], -1) if bias_scale is not None: - out = out + token_bias * bias_scale.view(-1, 1) + out.add_(token_bias * bias_scale.view(-1, 1)) else: - out = out + token_bias + out.add_(token_bias) return out From 47f13d671e6353272448a5ca26fce7951b25329e Mon Sep 17 00:00:00 2001 From: Siddhartha Raman Sundara Raman Date: Mon, 27 Jul 2026 07:41:29 -0700 Subject: [PATCH 04/14] Enable single-group grouped MLP optimizations by default Signed-off-by: Siddhartha Raman Sundara Raman --- .../pytorch/ops/fused/grouped_mlp.py | 101 ++---------------- 1 file changed, 9 insertions(+), 92 deletions(-) diff --git a/transformer_engine/pytorch/ops/fused/grouped_mlp.py b/transformer_engine/pytorch/ops/fused/grouped_mlp.py index d286dada229..0b2f4c6bdf8 100644 --- a/transformer_engine/pytorch/ops/fused/grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/grouped_mlp.py @@ -101,22 +101,6 @@ def _nvidia_cudnn_frontend_supports_wgrad() -> bool: return _cudnn_frontend_version_supported() -def _single_group_mlp_optimization_enabled(disable_env: str) -> bool: - """Whether to use one optimized single-group grouped-MLP path.""" - if int(os.getenv("NVTE_DISABLE_SINGLE_GROUP_MLP_OPT", "0")) != 0: - return False - return int(os.getenv(disable_env, "0")) == 0 - - -def _allocate_dprob_tensor(scales_tensor: torch.Tensor) -> torch.Tensor: - """Allocate the cuDNN DGLU dprob accumulator. - - Kept as a small helper so tests can inject nonzero initial contents and - verify cuDNN's zero-initialization requirement. - """ - return torch.zeros_like(scales_tensor) - - def _wrap_single_quantized_as_grouped( tensor: torch.Tensor, quantized: MXFP8Tensor | NVFP4Tensor | NVFP4TensorStorage, @@ -151,12 +135,6 @@ def _wrap_single_quantized_as_grouped( else: k_dim = tensor.shape[-1] - if tensor_offsets is None and not _single_group_mlp_optimization_enabled( - "NVTE_DISABLE_SINGLE_GROUP_MLP_OFFSETLESS_WRAPPER_OPT" - ): - tensor_offsets = torch.zeros(2, dtype=torch.int64, device=tensor.device) - tensor_offsets[1] = m_dim * k_dim - return GroupedTensor( shape=(m_dim, k_dim), dtype=tensor.dtype, @@ -193,12 +171,7 @@ def _group_quantize_for_grouped_mlp( tensor_offsets=tensor_offsets, ) - if not isinstance(quantizer, (MXFP8Quantizer, NVFP4Quantizer)) or ( - isinstance(quantizer, MXFP8Quantizer) - and not _single_group_mlp_optimization_enabled( - "NVTE_DISABLE_SINGLE_GROUP_MLP_QUANT_OPT" - ) - ): + if not isinstance(quantizer, (MXFP8Quantizer, NVFP4Quantizer)): return tex.group_quantize( tensor, quantizer, @@ -747,15 +720,7 @@ def _compute_grad_params( num_groups == 1 and isinstance(grouped_x, (GroupedTensor, GroupedTensorStorage)) and isinstance(grouped_dy, (GroupedTensor, GroupedTensorStorage)) - and ( - isinstance(grouped_x.quantizer, NVFP4Quantizer) - or ( - isinstance(grouped_x.quantizer, MXFP8Quantizer) - and _single_group_mlp_optimization_enabled( - "NVTE_DISABLE_SINGLE_GROUP_MLP_WGRAD_OPT" - ) - ) - ) + and isinstance(grouped_x.quantizer, (MXFP8Quantizer, NVFP4Quantizer)) and isinstance(grouped_dy.quantizer, grouped_x.quantizer.__class__) ): gemm_fn = functools.partial( @@ -1143,24 +1108,6 @@ def fuser_forward( num_groups == 1 and unit_activation_scale and isinstance(fc1_input_quantizer, MXFP8Quantizer) - and _single_group_mlp_optimization_enabled( - "NVTE_DISABLE_SINGLE_GROUP_MLP_OFFSETLESS_METADATA_OPT" - ) - and _single_group_mlp_optimization_enabled( - "NVTE_DISABLE_SINGLE_GROUP_MLP_OFFSETLESS_WRAPPER_OPT" - ) - and _single_group_mlp_optimization_enabled( - "NVTE_DISABLE_SINGLE_GROUP_MLP_QUANT_OPT" - ) - and _single_group_mlp_optimization_enabled( - "NVTE_DISABLE_SINGLE_GROUP_MLP_MXFP8_FC2_GEMM_OPT" - ) - and _single_group_mlp_optimization_enabled( - "NVTE_DISABLE_SINGLE_GROUP_MLP_MXFP8_DGRAD_GEMM_OPT" - ) - and _single_group_mlp_optimization_enabled( - "NVTE_DISABLE_SINGLE_GROUP_MLP_WGRAD_OPT" - ) ) # Prepare split metadata @@ -1176,9 +1123,7 @@ def fuser_forward( fc1_x_tensor_offsets = None fc2_x_tensor_offsets = None fc2_out_tensor_offsets = None - elif num_groups == 1 and _single_group_mlp_optimization_enabled( - "NVTE_DISABLE_SINGLE_GROUP_MLP_METADATA_OPT" - ): + elif num_groups == 1: metadata_key = ( device.type, device.index, @@ -1306,9 +1251,6 @@ def fuser_forward( and weight_requires_grad and isinstance(fc1_input_quantizer, MXFP8Quantizer) and not is_quantized_tensor(input_) - and _single_group_mlp_optimization_enabled( - "NVTE_DISABLE_SINGLE_GROUP_MLP_DEFER_FC1_COLUMNWISE_OPT" - ) ) # Group-quantize input tensor and convert dtypes if needed @@ -1505,12 +1447,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 - and _single_group_mlp_optimization_enabled( - "NVTE_DISABLE_SINGLE_GROUP_MLP_FC1_WEIGHT_SWIZZLE_OPT" - ) - ) + use_single_group_weight_swizzle = num_groups == 1 if use_single_group_weight_swizzle: fc1_weight_single = _single_quantized_tensor_from_grouped(fc1_weight_for_gemm) fc1_weight_single._columnwise_data = None @@ -1547,12 +1484,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 - and _single_group_mlp_optimization_enabled( - "NVTE_DISABLE_SINGLE_GROUP_MLP_FC1_DISCRETE_WEIGHT_OPT" - ) - ) + use_single_discrete_weight = num_groups == 1 if use_single_discrete_weight: fc1_weight_single = grouped_fc1_weight[0] original_rowwise_scale = fc1_weight_single._rowwise_scale_inv @@ -1725,12 +1657,7 @@ def fuser_forward( with_gemm_swizzled_scales=True, ) - use_single_group_dense_fc2 = ( - num_groups == 1 - and _single_group_mlp_optimization_enabled( - "NVTE_DISABLE_SINGLE_GROUP_MLP_MXFP8_FC2_GEMM_OPT" - ) - ) + use_single_group_dense_fc2 = num_groups == 1 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( @@ -2098,7 +2025,7 @@ def fuser_backward( if not unit_activation_scale: scales_f32 = scales.detach().to(dtype=torch.float32) scales_tensor = scales_f32.reshape(-1, 1, 1) - dscales_tensor = _allocate_dprob_tensor(scales_tensor) + dscales_tensor = torch.zeros_like(scales_tensor) fc2_d_dtype = torch.bfloat16 if use_nvfp4 else torch.float8_e4m3fn if use_nvfp4: @@ -2192,12 +2119,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 - and _single_group_mlp_optimization_enabled( - "NVTE_DISABLE_SINGLE_GROUP_MLP_FC2_DGRAD_DISCRETE_WEIGHT_OPT" - ) - ) + use_single_discrete_weight = num_groups == 1 if use_single_discrete_weight: fc2_weight_single = grouped_fc2_weight[0] original_rowwise_data = fc2_weight_single._rowwise_data @@ -2446,12 +2368,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 and ( - use_nvfp4 - or _single_group_mlp_optimization_enabled( - "NVTE_DISABLE_SINGLE_GROUP_MLP_MXFP8_DGRAD_GEMM_OPT" - ) - ) + use_single_group_dense_dgrad = num_groups == 1 if use_single_group_dense_dgrad: grad_input = validate_or_alloc_output( grad_input_buffer, in_shape, dtype, device From a745d3e16f73c216ce0553878a56571b0cecc737 Mon Sep 17 00:00:00 2001 From: Siddhartha Raman Sundara Raman Date: Mon, 27 Jul 2026 08:11:52 -0700 Subject: [PATCH 05/14] Simplify single-group quantizer guard Signed-off-by: Siddhartha Raman Sundara Raman --- transformer_engine/pytorch/ops/fused/grouped_mlp.py | 11 +---------- 1 file changed, 1 insertion(+), 10 deletions(-) diff --git a/transformer_engine/pytorch/ops/fused/grouped_mlp.py b/transformer_engine/pytorch/ops/fused/grouped_mlp.py index 0b2f4c6bdf8..e3de93ec2b9 100644 --- a/transformer_engine/pytorch/ops/fused/grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/grouped_mlp.py @@ -162,16 +162,7 @@ def _group_quantize_for_grouped_mlp( ) -> GroupedTensor: """Quantize into grouped storage.""" - if num_groups != 1: - return tex.group_quantize( - tensor, - quantizer, - num_groups, - split_sizes, - tensor_offsets=tensor_offsets, - ) - - if not isinstance(quantizer, (MXFP8Quantizer, NVFP4Quantizer)): + if num_groups != 1 or not isinstance(quantizer, (MXFP8Quantizer, NVFP4Quantizer)): return tex.group_quantize( tensor, quantizer, From 19154e68aa3bcf9c8540ae12cb2727bce8569de2 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 27 Jul 2026 15:34:56 +0000 Subject: [PATCH 06/14] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- .../pytorch/ops/fused/grouped_mlp.py | 31 ++++++------------- 1 file changed, 9 insertions(+), 22 deletions(-) diff --git a/transformer_engine/pytorch/ops/fused/grouped_mlp.py b/transformer_engine/pytorch/ops/fused/grouped_mlp.py index e3de93ec2b9..cbc4bab986f 100644 --- a/transformer_engine/pytorch/ops/fused/grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/grouped_mlp.py @@ -1513,8 +1513,7 @@ def fuser_forward( else: # Discrete-weight kernel: per-expert data/scale pointers fc1_b_ptrs, fc1_sfb_ptrs, _fc1_sfb_buffer = ( - tex.grouped_mlp_experimental - .swizzle_scales_and_pack_ptrs_for_discrete_weights( + tex.grouped_mlp_experimental.swizzle_scales_and_pack_ptrs_for_discrete_weights( [w._rowwise_data for w in grouped_fc1_weight], [w._rowwise_scale_inv for w in grouped_fc1_weight], "nvfp4" if use_nvfp4 else "mxfp8_rowwise", @@ -1703,9 +1702,7 @@ def fuser_forward( ) fc2_w_data = fc2_w_data.permute(1, 2, 0) - fc2_w_scales = fc2_weight_for_gemm.scale_inv.view( - dtype=torch.float8_e8m0fnu - ) + fc2_w_scales = fc2_weight_for_gemm.scale_inv.view(dtype=torch.float8_e8m0fnu) fc2_w_scales = fc2_w_scales.view( num_groups, ceil_div(fc2_weight_shape[0], 128), @@ -1719,8 +1716,7 @@ def fuser_forward( fc2_quant_kwargs["sfb_tensor"] = fc2_w_scales else: fc2_b_ptrs, fc2_sfb_ptrs, _fc2_sfb_buffer = ( - tex.grouped_mlp_experimental - .swizzle_scales_and_pack_ptrs_for_discrete_weights( + tex.grouped_mlp_experimental.swizzle_scales_and_pack_ptrs_for_discrete_weights( [w._rowwise_data for w in grouped_fc2_weight], [w._rowwise_scale_inv for w in grouped_fc2_weight], "mxfp8_rowwise", @@ -1940,9 +1936,7 @@ def fuser_backward( num_groups, split_sizes, tensor_offsets=( - None - if num_groups == 1 - else base_split_offsets * fc2_weight_shape[0] + None if num_groups == 1 else base_split_offsets * fc2_weight_shape[0] ), ) @@ -2133,9 +2127,7 @@ def fuser_backward( fc2_weight_k, ) fc2_w_data = ( - fc2_w_data.permute(1, 2, 0) - if use_nvfp4 - else fc2_w_data.permute(2, 1, 0) + fc2_w_data.permute(1, 2, 0) if use_nvfp4 else fc2_w_data.permute(2, 1, 0) ) fc2_w_scales = swizzled_columnwise_scale.view(dtype=scale_view_dtype) fc2_w_scales = fc2_w_scales.view( @@ -2155,8 +2147,7 @@ def fuser_backward( fc2_dactivation_kwargs["sfb_tensor"] = fc2_w_scales else: fc2_b_ptrs, fc2_sfb_ptrs, _fc2_sfb_buffer = ( - tex.grouped_mlp_experimental - .swizzle_scales_and_pack_ptrs_for_discrete_weights( + tex.grouped_mlp_experimental.swizzle_scales_and_pack_ptrs_for_discrete_weights( [w._columnwise_data for w in grouped_fc2_weight], [w._columnwise_scale_inv for w in grouped_fc2_weight], "nvfp4" if use_nvfp4 else "mxfp8_columnwise", @@ -2361,9 +2352,7 @@ def fuser_backward( use_single_group_dense_dgrad = num_groups == 1 if use_single_group_dense_dgrad: - grad_input = validate_or_alloc_output( - grad_input_buffer, in_shape, dtype, device - ) + grad_input = validate_or_alloc_output(grad_input_buffer, in_shape, dtype, device) _single_group_dgrad_gemm( grouped_fc1_dy, grouped_fc1_weight, @@ -2372,9 +2361,7 @@ def fuser_backward( dtype=dtype, ) elif use_nvfp4: - grad_input = validate_or_alloc_output( - grad_input_buffer, in_shape, dtype, device - ) + grad_input = validate_or_alloc_output(grad_input_buffer, in_shape, dtype, device) fc1_x_tensor_offsets = base_split_offsets * fc1_weight_shape[1] grouped_grad_input = GroupedTensor( shape=(out_shape[0], fc1_weight_shape[1]), @@ -2502,7 +2489,7 @@ def fuser_backward( grouped_fc1_x.columnwise_data, grouped_fc1_x.scale_inv, grouped_fc1_x.columnwise_scale_inv, - ) + ) fc2_grad_extra = (None, None) if fc2_op._scale_bias else (None,) if unit_activation_scale: From 65f4693a2247ea9eaee7f2275e0c9a045e0ea741 Mon Sep 17 00:00:00 2001 From: Siddhartha Raman Sundara Raman Date: Mon, 27 Jul 2026 20:30:53 -0700 Subject: [PATCH 07/14] Generate FC1 columnwise input in forward Signed-off-by: Siddhartha Raman Sundara Raman --- .../pytorch/ops/fused/grouped_mlp.py | 32 ++----------------- 1 file changed, 2 insertions(+), 30 deletions(-) diff --git a/transformer_engine/pytorch/ops/fused/grouped_mlp.py b/transformer_engine/pytorch/ops/fused/grouped_mlp.py index 8c9b345e16b..cd13a27b02e 100644 --- a/transformer_engine/pytorch/ops/fused/grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/grouped_mlp.py @@ -1240,22 +1240,10 @@ def fuser_forward( ): grouped_fc1_weight._with_gemm_swizzled_scales = False - # The canonical shared-expert TELinear saves its original BF16 input and - # creates the columnwise MXFP8 representation just before FC1 wgrad. - # Preserve that policy for the single-group fused path so fprop only - # pays for the rowwise representation required by the Rubin kernel. - defer_fc1_input_columnwise = ( - num_groups == 1 - and unit_activation_scale - and weight_requires_grad - and isinstance(fc1_input_quantizer, MXFP8Quantizer) - and not is_quantized_tensor(input_) - ) - # Group-quantize input tensor and convert dtypes if needed fc1_input_quantizer.set_usage( rowwise=True, - columnwise=weight_requires_grad and not defer_fc1_input_columnwise, + columnwise=weight_requires_grad, ) fc1_input_quantizer.optimize_for_gemm = True fc1_input_quantizer.internal = True @@ -1746,10 +1734,7 @@ def fuser_forward( # Save state for backward pass if requires_grad: - # The deferred path saves the caller-owned BF16 input in the slot - # normally occupied by the columnwise grouped tensor. Backward - # converts it immediately before FC1 wgrad. - saved_fc1_x = input_ if defer_fc1_input_columnwise else grouped_fc1_x + saved_fc1_x = grouped_fc1_x mark_grouped_tensor(saved_fc1_x, activation_in, scales, grouped_fc2_x) activation_op = self.basic_ops[1] cpu_offloading = is_cpu_offload_enabled() @@ -1812,7 +1797,6 @@ def fuser_forward( fc1_ctx.input_requires_grad = input_requires_grad fc1_ctx.weight_requires_grad = weight_requires_grad fc1_ctx.unit_activation_scale = unit_activation_scale - fc1_ctx.defer_fc1_input_columnwise = defer_fc1_input_columnwise fc2_ctx.input_quantizers = [fc2_input_quantizer] fc2_ctx.grad_output_quantizers = [fc2_grad_output_quantizer] @@ -1892,18 +1876,6 @@ def fuser_backward( if not fc1_ctx.weight_requires_grad: grouped_fc1_x = None - elif bool(getattr(fc1_ctx, "defer_fc1_input_columnwise", False)): - fc1_input_quantizer = fc1_ctx.input_quantizers[0] - fc1_input_quantizer.set_usage(rowwise=False, columnwise=True) - fc1_input_quantizer.optimize_for_gemm = True - fc1_input_quantizer.internal = True - grouped_fc1_x = _group_quantize_for_grouped_mlp( - grouped_fc1_x, - fc1_input_quantizer, - num_groups, - split_sizes, - tensor_offsets=None, - ) if not fc2_ctx.weight_requires_grad: grouped_fc2_x = None From f24c7c3a41e3afd85578f4242ac146ba875b3fb8 Mon Sep 17 00:00:00 2001 From: Siddhartha Raman Sundara Raman Date: Mon, 27 Jul 2026 22:33:21 -0500 Subject: [PATCH 08/14] Update transformer_engine/pytorch/ops/fused/grouped_mlp.py Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> Signed-off-by: Siddhartha Raman Sundara Raman --- transformer_engine/pytorch/ops/fused/grouped_mlp.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/transformer_engine/pytorch/ops/fused/grouped_mlp.py b/transformer_engine/pytorch/ops/fused/grouped_mlp.py index cd13a27b02e..d161a13925b 100644 --- a/transformer_engine/pytorch/ops/fused/grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/grouped_mlp.py @@ -1087,9 +1087,7 @@ def fuser_forward( getattr(self.basic_ops[1], "_grouped_mlp_unit_activation_scale", False) ) if unit_activation_scale and num_groups != 1: - raise ValueError( - "Unit activation scaling is only supported for a single-group grouped MLP." - ) + unit_activation_scale = False # Shared experts have one dense group and all optimized kernels derive M # from their runtime tensor shapes. Reuse the caller-owned split tensor From 566b0cf7de59994297a8844ace240d91620382aa Mon Sep 17 00:00:00 2001 From: Siddhartha Raman Sundara Raman Date: Mon, 27 Jul 2026 20:34:51 -0700 Subject: [PATCH 09/14] Move grouped MLP fusion test Signed-off-by: Siddhartha Raman Sundara Raman --- tests/pytorch/test_fusible_ops.py | 23 +---------------------- tests/pytorch/test_grouped_mlp.py | 21 ++++++++++++++++++++- 2 files changed, 21 insertions(+), 23 deletions(-) diff --git a/tests/pytorch/test_fusible_ops.py b/tests/pytorch/test_fusible_ops.py index db4cb6dfb1e..7ca443b2b0e 100644 --- a/tests/pytorch/test_fusible_ops.py +++ b/tests/pytorch/test_fusible_ops.py @@ -3558,10 +3558,7 @@ def test_grouped_mlp( """GroupedLinear + scaled activation + GroupedLinear""" # Split sizes - if group_size == 1: - split_sizes = [split_alignment] - else: - split_sizes = [split_alignment * i for i in range(group_size)] + split_sizes = [split_alignment * (i) for i in range(group_size)] random.shuffle(split_sizes) split_sizes = torch.tensor(split_sizes, dtype=torch.int, device=device) @@ -3752,24 +3749,6 @@ def test_grouped_mlp_nvfp4_rht_srelu( activation="scaled_srelu", ) - @pytest.mark.parametrize("bias", (False, True)) - def test_grouped_mlp_single_group_mxfp8( - self, - *, - bias: bool, - device: torch.device = "cuda", - ) -> None: - """Single-group GroupedLinear + ScaledSwiGLU + GroupedLinear with MXFP8.""" - - self.test_grouped_mlp( - group_size=1, - bias=bias, - dtype=torch.bfloat16, - quantization="mxfp8", - device=device, - ) - - class TestCustomOps: """Test with ops that are defined externally""" diff --git a/tests/pytorch/test_grouped_mlp.py b/tests/pytorch/test_grouped_mlp.py index cadfd6c1f6e..386ee62e23a 100644 --- a/tests/pytorch/test_grouped_mlp.py +++ b/tests/pytorch/test_grouped_mlp.py @@ -728,7 +728,10 @@ def test_grouped_mlp( """GroupedLinear + scaled activation + GroupedLinear""" # Split sizes - split_sizes = [split_alignment * (i) for i in range(group_size)] + if group_size == 1: + split_sizes = [split_alignment] + else: + split_sizes = [split_alignment * i for i in range(group_size)] random.shuffle(split_sizes) split_sizes = torch.tensor(split_sizes, dtype=torch.int, device=device) @@ -1131,6 +1134,22 @@ def test_grouped_mlp_fp16( activation=activation, ) + @pytest.mark.parametrize("bias", (False, True)) + def test_grouped_mlp_single_group_mxfp8( + self, + *, + bias: bool, + ) -> None: + """Single-group GroupedLinear + ScaledSwiGLU + GroupedLinear with MXFP8.""" + self.test_grouped_mlp( + group_size=1, + bias=bias, + hidden_size=128, + quantization="mxfp8", + single_grouped_weight=False, + activation="scaled_swiglu", + ) + @pytest.mark.skipif(not mxfp8_available, reason=reason_for_no_mxfp8) def test_single_grouped_weight_eval_preserves_columnwise_usage( self, From fa6fcbe6d5f479f6f3a1e40c8adc3607e94eec9d Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Tue, 28 Jul 2026 03:36:16 +0000 Subject: [PATCH 10/14] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- tests/pytorch/test_fusible_ops.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/pytorch/test_fusible_ops.py b/tests/pytorch/test_fusible_ops.py index 7ca443b2b0e..66857d81253 100644 --- a/tests/pytorch/test_fusible_ops.py +++ b/tests/pytorch/test_fusible_ops.py @@ -3749,6 +3749,7 @@ def test_grouped_mlp_nvfp4_rht_srelu( activation="scaled_srelu", ) + class TestCustomOps: """Test with ops that are defined externally""" From 70a1293324153d828e80c7c7ac754fa2819929f1 Mon Sep 17 00:00:00 2001 From: Siddhartha Raman Sundara Raman Date: Mon, 27 Jul 2026 20:44:31 -0700 Subject: [PATCH 11/14] Support older cuDNN grouped MLP APIs Signed-off-by: Siddhartha Raman Sundara Raman --- tests/pytorch/test_grouped_mlp.py | 9 +++++ .../pytorch/ops/fused/grouped_mlp.py | 36 ++++++++++++++----- 2 files changed, 37 insertions(+), 8 deletions(-) diff --git a/tests/pytorch/test_grouped_mlp.py b/tests/pytorch/test_grouped_mlp.py index 386ee62e23a..ac125ce74ea 100644 --- a/tests/pytorch/test_grouped_mlp.py +++ b/tests/pytorch/test_grouped_mlp.py @@ -15,6 +15,7 @@ import torch import transformer_engine.pytorch as te +import transformer_engine.pytorch.ops.fused.grouped_mlp as grouped_mlp_module from transformer_engine.pytorch.ops.fused.grouped_mlp import ( _cudnn_frontend_supports_grouped_gemm_srelu, _cudnn_frontend_version_supported, @@ -1135,12 +1136,20 @@ def test_grouped_mlp_fp16( ) @pytest.mark.parametrize("bias", (False, True)) + @pytest.mark.parametrize("runtime_offsets_supported", (False, True)) def test_grouped_mlp_single_group_mxfp8( self, + monkeypatch, *, bias: bool, + runtime_offsets_supported: bool, ) -> None: """Single-group GroupedLinear + ScaledSwiGLU + GroupedLinear with MXFP8.""" + monkeypatch.setattr( + grouped_mlp_module, + "_cudnn_frontend_supports_single_group_runtime_offsets", + lambda: runtime_offsets_supported, + ) self.test_grouped_mlp( group_size=1, bias=bias, diff --git a/transformer_engine/pytorch/ops/fused/grouped_mlp.py b/transformer_engine/pytorch/ops/fused/grouped_mlp.py index d161a13925b..fa6449bf854 100644 --- a/transformer_engine/pytorch/ops/fused/grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/grouped_mlp.py @@ -101,6 +101,11 @@ def _nvidia_cudnn_frontend_supports_wgrad() -> bool: return _cudnn_frontend_version_supported() +def _cudnn_frontend_supports_single_group_runtime_offsets() -> bool: + """Check cuDNN FE min version for single-group runtime offsets.""" + return _cudnn_frontend_version_at_least("1.27.0") + + def _wrap_single_quantized_as_grouped( tensor: torch.Tensor, quantized: MXFP8Tensor | NVFP4Tensor | NVFP4TensorStorage, @@ -1089,14 +1094,22 @@ def fuser_forward( if unit_activation_scale and num_groups != 1: unit_activation_scale = False + activation_kernel = self.grouped_gemm_activation_kernel() + supports_single_group_runtime_offsets = ( + _cudnn_frontend_supports_single_group_runtime_offsets() + ) + # Shared experts have one dense group and all optimized kernels derive M # from their runtime tensor shapes. Reuse the caller-owned split tensor # as the ignored cuDNN padded-offset argument and omit tensor offsets # entirely. This avoids both splits_to_offsets and cached CUDA pointers. + # Older cuDNN frontends do not expose this specialization, so retain the + # real single-group metadata path for compatibility. use_offsetless_metadata = ( num_groups == 1 and unit_activation_scale and isinstance(fc1_input_quantizer, MXFP8Quantizer) + and supports_single_group_runtime_offsets ) # Prepare split metadata @@ -1420,7 +1433,8 @@ def fuser_forward( else: fc1_activation_kwargs["norm_const_tensor"] = fc1_norm_const_tensor fc1_activation_kwargs["discrete_col_sfd"] = not use_nvfp4 - fc1_activation_kwargs["use_single_group_runtime_offsets"] = num_groups == 1 + if supports_single_group_runtime_offsets: + fc1_activation_kwargs["use_single_group_runtime_offsets"] = num_groups == 1 if self._pass_geglu_runtime_params: fc1_activation_kwargs.update( linear_offset=self._cudnn_linear_offset, @@ -1523,7 +1537,7 @@ def fuser_forward( if use_fc1_act_hadamard: fc1_kernel_out = self.grouped_gemm_act_hadamard_kernel()(**fc1_activation_kwargs) else: - fc1_kernel_out = self.grouped_gemm_activation_kernel()(**fc1_activation_kwargs) + fc1_kernel_out = activation_kernel(**fc1_activation_kwargs) if fc2_is_dist: grouped_fc2_weight = materialize_weight_for_forward(grouped_fc2_weight) @@ -1674,8 +1688,10 @@ def fuser_forward( "sf_vec_size": MXFP8_BLOCK_SCALING_SIZE, "current_stream": current_stream, "use_dynamic_sched": True, - "use_single_group_runtime_offsets": num_groups == 1, } + 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 if fc2_op.single_grouped_weight: # Clone and swizzle scales for GEMM @@ -1727,7 +1743,7 @@ def fuser_forward( (in_shape[0], fc2_weight_shape[0], 1), (fc2_weight_shape[0], 1, in_shape[0] * fc2_weight_shape[0]), ) - self.grouped_gemm_quant_kernel()(**fc2_quant_kwargs) + fc2_quant_kernel(**fc2_quant_kwargs) fc2_out = fc2_out_buf # Save state for backward pass @@ -2034,8 +2050,10 @@ def fuser_backward( "current_stream": current_stream, "discrete_col_sfd": not use_nvfp4, "use_dynamic_sched": True, - "use_single_group_runtime_offsets": num_groups == 1, } + dactivation_kernel = self.grouped_gemm_dactivation_kernel() + if _cudnn_frontend_supports_single_group_runtime_offsets(): + fc2_dactivation_kwargs["use_single_group_runtime_offsets"] = num_groups == 1 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 @@ -2138,7 +2156,7 @@ def fuser_backward( fc2_dactivation_kwargs["b_dtype"] = data_dtype fc2_dactivation_kwargs["b_major"] = "k" if use_nvfp4 else "n" - fc2_dgrad_kernel_out = self.grouped_gemm_dactivation_kernel()(**fc2_dactivation_kwargs) + fc2_dgrad_kernel_out = dactivation_kernel(**fc2_dactivation_kwargs) if use_nvfp4: fc1_dy_bf16 = fc2_dgrad_kernel_out["d_row_tensor"] @@ -2376,8 +2394,10 @@ def fuser_backward( "current_stream": current_stream, "discrete_col_sfd": True, "use_dynamic_sched": True, - "use_single_group_runtime_offsets": num_groups == 1, } + fc1_dgrad_kernel = self.grouped_gemm_quant_kernel() + if _cudnn_frontend_supports_single_group_runtime_offsets(): + fc1_dgrad_kwargs["use_single_group_runtime_offsets"] = num_groups == 1 if fc1_op.single_grouped_weight: # Clone and swizzle scales for GEMM @@ -2432,7 +2452,7 @@ def fuser_backward( (out_shape[0], fc1_weight_shape[1], 1), (fc1_weight_shape[1], 1, out_shape[0] * fc1_weight_shape[1]), ) - self.grouped_gemm_quant_kernel()(**fc1_dgrad_kwargs) + fc1_dgrad_kernel(**fc1_dgrad_kwargs) grad_input = grad_input_buffer # FC1 wgrad GEMM From 8d7f03076e976955783b5882a7dbb7216c1e6367 Mon Sep 17 00:00:00 2001 From: Siddhartha Raman Sundara Raman Date: Tue, 28 Jul 2026 08:43:21 -0700 Subject: [PATCH 12/14] Use live offsets for older cuDNN frontends Signed-off-by: Siddhartha Raman Sundara Raman --- .../pytorch/ops/fused/grouped_mlp.py | 64 +------------------ 1 file changed, 2 insertions(+), 62 deletions(-) diff --git a/transformer_engine/pytorch/ops/fused/grouped_mlp.py b/transformer_engine/pytorch/ops/fused/grouped_mlp.py index fa6449bf854..33cf5e98f05 100644 --- a/transformer_engine/pytorch/ops/fused/grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/grouped_mlp.py @@ -356,37 +356,6 @@ def _use_tmem_post_rht_amax() -> bool: return os.environ.get("NVTE_CUTEDSL_FUSED_GROUPED_MLP_FC1_GLU_RHT_AMAX_TMEM", "0") == "1" -def _single_group_split_metadata( - num_tokens: int, - device: torch.device, - *, - fc1_in_features: int, - fc2_in_features: int, - fc2_out_features: int, -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - """Create split/offset metadata for num_groups=1 without a GPU prefix-sum kernel.""" - split_sizes = torch.zeros(1, dtype=torch.int64, device=device) - split_sizes[0] = num_tokens - split_points = torch.zeros(1, dtype=torch.int32, device=device) - split_points[0] = num_tokens - base_split_offsets = torch.zeros(2, dtype=torch.int64, device=device) - base_split_offsets[1] = num_tokens - fc1_x_tensor_offsets = torch.zeros(2, dtype=torch.int64, device=device) - fc1_x_tensor_offsets[1] = num_tokens * fc1_in_features - fc2_x_tensor_offsets = torch.zeros(2, dtype=torch.int64, device=device) - fc2_x_tensor_offsets[1] = num_tokens * fc2_in_features - fc2_out_tensor_offsets = torch.zeros(2, dtype=torch.int64, device=device) - fc2_out_tensor_offsets[1] = num_tokens * fc2_out_features - return ( - split_sizes, - split_points, - base_split_offsets, - fc1_x_tensor_offsets, - fc2_x_tensor_offsets, - fc2_out_tensor_offsets, - ) - - def _single_group_wgrad_gemm( grouped_x: GroupedTensor, grouped_dy: GroupedTensor, @@ -996,8 +965,6 @@ def __init__( self._cudnn_geglu_alpha: float = activation._clamped.alpha self._cudnn_glu_clamp_max: float = activation._clamped.limit self._cudnn_glu_clamp_min: float = -activation._clamped.limit - self._single_group_split_metadata_cache = {} - def fuser_forward( self, basic_op_ctxs: list[OperationContext], @@ -1103,8 +1070,8 @@ def fuser_forward( # from their runtime tensor shapes. Reuse the caller-owned split tensor # as the ignored cuDNN padded-offset argument and omit tensor offsets # entirely. This avoids both splits_to_offsets and cached CUDA pointers. - # Older cuDNN frontends do not expose this specialization, so retain the - # real single-group metadata path for compatibility. + # 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 and unit_activation_scale @@ -1125,33 +1092,6 @@ def fuser_forward( fc1_x_tensor_offsets = None fc2_x_tensor_offsets = None fc2_out_tensor_offsets = None - elif num_groups == 1: - metadata_key = ( - device.type, - device.index, - in_shape[0], - fc1_weight_shape[1], - fc2_weight_shape[1], - fc2_weight_shape[0], - ) - if metadata_key not in self._single_group_split_metadata_cache: - self._single_group_split_metadata_cache[metadata_key] = ( - _single_group_split_metadata( - in_shape[0], - device, - fc1_in_features=fc1_weight_shape[1], - fc2_in_features=fc2_weight_shape[1], - fc2_out_features=fc2_weight_shape[0], - ) - ) - ( - split_sizes, - split_points, - base_split_offsets, - fc1_x_tensor_offsets, - fc2_x_tensor_offsets, - fc2_out_tensor_offsets, - ) = self._single_group_split_metadata_cache[metadata_key] else: split_sizes, ( split_points, From e45c19604eba2a4b3e5309120c279e11abf8fe98 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Tue, 28 Jul 2026 15:45:27 +0000 Subject: [PATCH 13/14] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- transformer_engine/pytorch/ops/fused/grouped_mlp.py | 1 + 1 file changed, 1 insertion(+) diff --git a/transformer_engine/pytorch/ops/fused/grouped_mlp.py b/transformer_engine/pytorch/ops/fused/grouped_mlp.py index 33cf5e98f05..97ea587bf6a 100644 --- a/transformer_engine/pytorch/ops/fused/grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/grouped_mlp.py @@ -965,6 +965,7 @@ def __init__( self._cudnn_geglu_alpha: float = activation._clamped.alpha self._cudnn_glu_clamp_max: float = activation._clamped.limit self._cudnn_glu_clamp_min: float = -activation._clamped.limit + def fuser_forward( self, basic_op_ctxs: list[OperationContext], From 7ec5e7443186ce45712d57ca7d278a1e7100fdad Mon Sep 17 00:00:00 2001 From: Siddhartha Raman Sundara Raman Date: Tue, 28 Jul 2026 16:20:42 -0700 Subject: [PATCH 14/14] Skip runtime-offset test with older cuDNN frontend Signed-off-by: Siddhartha Raman Sundara Raman --- tests/pytorch/test_grouped_mlp.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/tests/pytorch/test_grouped_mlp.py b/tests/pytorch/test_grouped_mlp.py index ac125ce74ea..d195eb2f78d 100644 --- a/tests/pytorch/test_grouped_mlp.py +++ b/tests/pytorch/test_grouped_mlp.py @@ -1145,6 +1145,11 @@ def test_grouped_mlp_single_group_mxfp8( runtime_offsets_supported: bool, ) -> None: """Single-group GroupedLinear + ScaledSwiGLU + GroupedLinear with MXFP8.""" + if ( + runtime_offsets_supported + and not grouped_mlp_module._cudnn_frontend_supports_single_group_runtime_offsets() + ): + pytest.skip("Requires cuDNN frontend >= 1.27.0") monkeypatch.setattr( grouped_mlp_module, "_cudnn_frontend_supports_single_group_runtime_offsets",