From 539829f1be076176387647aae5fec490b424e5f5 Mon Sep 17 00:00:00 2001 From: Hwanseo Choi Date: Mon, 5 Oct 2026 16:04:32 -0700 Subject: [PATCH 1/6] [PyTorch] Pass canonical layouts to cuDNN grouped GEMM in fused grouped MLP When the cuDNN frontend grouped GEMM GLU and quant wrappers advertise supports_canonical_layouts, the MXFP8 forward passes A as (sum(m), k), B as (num_groups, n, k), probabilities as 1-D, and the packed scale buffers as-is, and consumes the 2-D outputs directly. This drops the per-call view/permute chains on the host. NVFP4, SReLU, and older frontends keep the existing layouts. The FC2 quant call also skips the unused output amax when the wrapper advertises supports_optional_amax. Signed-off-by: Hwanseo Choi --- .../pytorch/ops/fused/grouped_mlp.py | 108 +++++++++++------- 1 file changed, 65 insertions(+), 43 deletions(-) diff --git a/transformer_engine/pytorch/ops/fused/grouped_mlp.py b/transformer_engine/pytorch/ops/fused/grouped_mlp.py index 108ce5311c..d466b58665 100644 --- a/transformer_engine/pytorch/ops/fused/grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/grouped_mlp.py @@ -1209,6 +1209,9 @@ def __init__( self.grouped_gemm_dactivation_kernel() raise RuntimeError(f"{self.__class__.__name__} is not supported on this system.") validate_grouped_mlp_dims(fc1, activation, fc2) + self._cudnn_canonical_layouts = getattr( + self.grouped_gemm_activation_kernel(), "supports_canonical_layouts", False + ) and getattr(self.grouped_gemm_quant_kernel(), "supports_canonical_layouts", False) if not is_glu_activation(activation): # grouped_gemm_srelu_wrapper_sm100 and grouped_gemm_dsrelu_wrapper_sm100 are # SReLU-specific and do not take GLU ``act_func`` selectors. @@ -1564,9 +1567,13 @@ def fuser_forward( # Data logical shape: (sum(m), k, 1) # Scale logical shape: (32 (block row), 4 (block row), # sum(m)/128, 4 (block col), k/128, 1) + # With canonical layouts, cuDNN takes A as (sum(m), k), B as + # (num_groups, n, k), and scales as the packed buffers as-is. + canonical = self._cudnn_canonical_layouts and not use_nvfp4 fc1_x_data = grouped_fc1_x.rowwise_data.view(dtype=data_dtype) fc1_x_data = fc1_x_data.view(in_shape[0], data_in_k) - fc1_x_data = fc1_x_data.unsqueeze(0).permute(1, 2, 0) + if not canonical: + fc1_x_data = fc1_x_data.unsqueeze(0).permute(1, 2, 0) fc1_x_scales = grouped_fc1_x.scale_inv fc1_x_scales = fc1_x_scales.view(dtype=scale_view_dtype) with_gemm_swizzled_scales = grouped_fc1_x._with_gemm_swizzled_scales @@ -1590,7 +1597,7 @@ def fuser_forward( 4, ) fc1_x_scales = fc1_x_scales.permute(3, 2, 1, 5, 4, 0) - else: + elif not canonical: fc1_x_scales = fc1_x_scales.view( 1, ceil_div(in_shape[0], 128), @@ -1611,7 +1618,9 @@ def fuser_forward( 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) + scales.detach() + .to(dtype=torch.float32 if use_nvfp4 else dtype) + .reshape((-1,) if canonical else (-1, 1, 1)) ) if use_nvfp4: # cuDNN receives NVFP4 block-scaled inputs without TE's per-group @@ -1753,17 +1762,18 @@ def fuser_forward( # Data logical shape: (n, k, num_groups) 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_w_scales.view(dtype=scale_view_dtype) - fc1_w_scales = fc1_w_scales.view( - num_groups, - 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) + if not canonical: + fc1_w_data = fc1_w_data.permute(1, 2, 0) + fc1_w_scales = fc1_w_scales.view( + num_groups, + 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 @@ -1790,17 +1800,18 @@ def fuser_forward( 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_k, k_sf_divisor), - 32, - 4, - 4, - ) - fc1_w_scales = fc1_w_scales.permute(3, 4, 1, 5, 2, 0) + if not canonical: + fc1_w_data = fc1_w_data.permute(1, 2, 0) + fc1_w_scales = fc1_w_scales.view( + 1, + ceil_div(fc1_weight_shape[0], 128), + ceil_div(fc1_weight_k, 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: @@ -1832,7 +1843,8 @@ def fuser_forward( raise RuntimeError(f"Unrecognized kernel variant ({kernel_impl})") activation_in = fc1_kernel_out["c_tensor"] - activation_in = activation_in.view(in_shape[0], fc1_weight_shape[0]) + if not canonical: + activation_in = activation_in.view(in_shape[0], fc1_weight_shape[0]) if fc2_is_dist: grouped_fc2_weight = materialize_weight_for_forward(grouped_fc2_weight) @@ -1897,13 +1909,14 @@ def fuser_forward( else: # Unpack MXFP8 output fc2_in_row_data = fc1_kernel_out["d_tensor"] - fc2_in_row_data = fc2_in_row_data.view(in_shape[0], fc2_weight_shape[1]) fc2_in_row_scale = fc1_kernel_out["sfd_row_tensor"] - fc2_in_row_scale = fc2_in_row_scale.permute(5, 2, 4, 0, 1, 3) fc2_in_col_data = fc1_kernel_out["d_col_tensor"] - fc2_in_col_data = fc2_in_col_data.view(in_shape[0], fc2_weight_shape[1]) fc2_in_col_scale = fc1_kernel_out["sfd_col_tensor"] - fc2_in_col_scale = fc2_in_col_scale.permute(5, 2, 4, 0, 1, 3) + if not canonical: + fc2_in_row_data = fc2_in_row_data.view(in_shape[0], fc2_weight_shape[1]) + fc2_in_row_scale = fc2_in_row_scale.permute(5, 2, 4, 0, 1, 3) + fc2_in_col_data = fc2_in_col_data.view(in_shape[0], fc2_weight_shape[1]) + fc2_in_col_scale = fc2_in_col_scale.permute(5, 2, 4, 0, 1, 3) grouped_fc2_x = GroupedTensorStorage( shape=(in_shape[0], fc2_weight_shape[1]), dtype=dtype, @@ -2082,7 +2095,9 @@ def fuser_forward( ) if fc2_scales is not None: fc2_scales_tensor = ( - fc2_scales.detach().to(dtype=torch.float32).reshape(-1, 1, 1) + fc2_scales.detach() + .to(dtype=torch.float32) + .reshape((-1,) if canonical else (-1, 1, 1)) ) fc2_quant_kwargs = { "a_tensor": fc1_kernel_out["d_tensor"], @@ -2100,6 +2115,9 @@ def fuser_forward( "use_dynamic_sched": True, } fc2_quant_kernel = self.grouped_gemm_quant_kernel() + if getattr(fc2_quant_kernel, "supports_optional_amax", False): + # The FC2 output amax is not used. + fc2_quant_kwargs["generate_amax"] = False if supports_single_group_runtime_offsets: fc2_quant_kwargs["use_single_group_runtime_offsets"] = use_dense_single_group @@ -2120,18 +2138,18 @@ def fuser_forward( 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) + if not canonical: + fc2_w_data = fc2_w_data.permute(1, 2, 0) + 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: @@ -2149,9 +2167,13 @@ def fuser_forward( 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["d_tensor"] = ( + fc2_out_buf + if canonical + else 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_kernel(**fc2_quant_kwargs) fc2_out = fc2_out_buf From edef648a19a9b83b68be2ae573b17f506e710db4 Mon Sep 17 00:00:00 2001 From: Hwanseo Choi Date: Tue, 6 Oct 2026 15:56:03 -0700 Subject: [PATCH 2/6] [PyTorch] Test and document canonical cuDNN layouts in fused grouped MLP Check that natural operand layouts reach the cuDNN frontend and give bitwise-identical outputs and gradients to the legacy views, and note the layout behavior in the fused op docstring. Signed-off-by: Hwanseo Choi --- tests/pytorch/test_grouped_mlp.py | 96 +++++++++++++++++++ .../pytorch/ops/fused/grouped_mlp.py | 6 +- 2 files changed, 101 insertions(+), 1 deletion(-) diff --git a/tests/pytorch/test_grouped_mlp.py b/tests/pytorch/test_grouped_mlp.py index 88af92f0bb..5c65248317 100644 --- a/tests/pytorch/test_grouped_mlp.py +++ b/tests/pytorch/test_grouped_mlp.py @@ -1762,6 +1762,102 @@ def test_grouped_mlp_glu_mxfp8_real_cudnn_fusion( ] assert traced_cudnn_grouped_glu_wrappers == expected_calls + def test_grouped_mlp_canonical_layouts_match_legacy(self, monkeypatch) -> None: + """Natural cuDNN operand layouts give bitwise-identical results to the legacy views.""" + from transformer_engine.common.recipe import Format, MXFP8BlockScaling + + maybe_skip_quantization("mxfp8") + monkeypatch.setenv("NVTE_CUTEDSL_FUSED_GROUPED_MLP", "1") + monkeypatch.setenv("NVTE_GROUPED_LINEAR_SINGLE_PARAM", "1") + fused_cls = te.ops.fused.GroupedMLP_CuTeGEMMGLU + if not fused_cls.is_supported(): + pytest.skip("cuDNN fused grouped MLP is not supported") + kernels = { + "grouped_gemm_activation_kernel": fused_cls.grouped_gemm_activation_kernel(), + "grouped_gemm_quant_kernel": fused_cls.grouped_gemm_quant_kernel(), + } + if not all( + getattr(kernel, "supports_canonical_layouts", False) for kernel in kernels.values() + ): + pytest.skip("Installed cuDNN frontend lacks canonical grouped GEMM layouts") + + recipe = MXFP8BlockScaling(fp8_format=Format.E4M3) + with te.quantized_model_init(enabled=True, recipe=recipe): + fc1 = te.ops.GroupedLinear( + 4, + 256, + 512, + bias=False, + device="cuda", + dtype=torch.bfloat16, + single_grouped_weight=True, + ) + fc2 = te.ops.GroupedLinear( + 4, + 256, + 256, + bias=False, + device="cuda", + dtype=torch.bfloat16, + single_grouped_weight=True, + ) + activation = te.ops.ScaledSwiGLU(glu_interleave_size=32) + x = torch.randn(1024, 256, device="cuda", dtype=torch.bfloat16) + probs = torch.rand(1024, device="cuda", dtype=torch.bfloat16) + splits = torch.full((4,), 256, dtype=torch.int64, device="cuda") + dy = torch.randn(1024, 256, device="cuda", dtype=torch.bfloat16) + + operands = ("a_tensor", "b_tensor", "sfa_tensor", "sfb_tensor", "prob_tensor", "d_tensor") + + def run(canonical): + layouts = [] + + def wrap(kernel): + # Hiding the capability attributes selects the legacy views. + if not canonical: + return lambda **kwargs: kernel(**kwargs) + + @functools.wraps(kernel) + def traced(**kwargs): + layouts.append( + { + name: (kwargs[name].dim(), kwargs[name].is_contiguous()) + for name in operands + if kwargs.get(name) is not None + } + ) + return kernel(**kwargs) + + return traced + + with monkeypatch.context() as patch: + for name, kernel in kernels.items(): + wrapped = wrap(kernel) + patch.setattr( + fused_cls, name, classmethod(lambda cls, wrapped=wrapped: wrapped) + ) + model = te.ops.Sequential(fc1, activation, fc2) + x_ = x.clone().requires_grad_() + probs_ = probs.clone().requires_grad_() + fc1.weight.grad = fc2.weight.grad = None + with te.autocast(enabled=True, recipe=recipe): + y = model(x_, splits, probs_, splits) + y.backward(dy) + forward_ops = model._module_groups[0]._forward_ops + assert len(forward_ops) == 1 and isinstance(forward_ops[0][0], fused_cls) + return [y, x_.grad, probs_.grad, fc1.weight.grad, fc2.weight.grad], layouts + + legacy, _ = run(canonical=False) + canonical, layouts = run(canonical=True) + # The forward FC1 and FC2 calls get every operand as its natural contiguous buffer. + for call in layouts[:2]: + assert all(contiguous for _, contiguous in call.values()) + assert call["a_tensor"][0] == 2 and call["b_tensor"][0] == 3 + assert call.get("prob_tensor", (1,))[0] == 1 + assert layouts[1]["d_tensor"][0] == 2 + for actual, expected in zip(canonical, legacy): + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + @pytest.mark.skipif(not nvfp4_available, reason=reason_for_no_nvfp4) @pytest.mark.parametrize("situ_betas", ((4.0, 25.0), (2.0, 8.0))) def test_grouped_mlp_situglu_nvfp4_real_cudnn_hadamard_fusion( diff --git a/transformer_engine/pytorch/ops/fused/grouped_mlp.py b/transformer_engine/pytorch/ops/fused/grouped_mlp.py index d466b58665..4577856df7 100644 --- a/transformer_engine/pytorch/ops/fused/grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/grouped_mlp.py @@ -1116,7 +1116,9 @@ class _GroupedMLP_CuTeGEMMBase(FusedOperation): MXFP8 uses CuTe DSL grouped-GEMM kernels via the cuDNN front-end. When caller output/grad_input buffers are provided, the GEMMs write into them - directly. + directly. When the front-end wrappers advertise ``supports_canonical_layouts``, + MXFP8 operands are passed in their natural layouts instead of kernel-facing + views. """ @@ -1209,6 +1211,8 @@ def __init__( self.grouped_gemm_dactivation_kernel() raise RuntimeError(f"{self.__class__.__name__} is not supported on this system.") validate_grouped_mlp_dims(fc1, activation, fc2) + # Newer cuDNN front-ends accept natural operand layouts, so the forward can + # skip building the 3-D/6-D kernel-facing views on every call. self._cudnn_canonical_layouts = getattr( self.grouped_gemm_activation_kernel(), "supports_canonical_layouts", False ) and getattr(self.grouped_gemm_quant_kernel(), "supports_canonical_layouts", False) From 96da078a22f17705ace6ad3b968c49e272a7cd93 Mon Sep 17 00:00:00 2001 From: Hwanseo Choi Date: Tue, 6 Oct 2026 15:21:34 -0700 Subject: [PATCH 3/6] [PyTorch] Pass canonical layouts to cuDNN grouped GEMM in fused grouped MLP backward When the cuDNN frontend dGLU wrapper advertises supports_canonical_layouts, the MXFP8 backward passes 2-D grad output, activation input, and FC1 dgrad operands, 1-D probabilities, packed scale buffers as-is, and the column-wise weights as (num_groups, k, n) with b_major="n". It consumes the 2-D dGLU outputs directly. wgrad passes flat scale buffers when the wgrad wrapper supports them. NVFP4, SReLU, and older frontends keep the existing layouts. Signed-off-by: Hwanseo Choi --- .../pytorch/ops/fused/grouped_mlp.py | 111 ++++++++++-------- 1 file changed, 64 insertions(+), 47 deletions(-) diff --git a/transformer_engine/pytorch/ops/fused/grouped_mlp.py b/transformer_engine/pytorch/ops/fused/grouped_mlp.py index 4577856df7..aa3478dad7 100644 --- a/transformer_engine/pytorch/ops/fused/grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/grouped_mlp.py @@ -741,12 +741,11 @@ def _cudnn_compute_wgrad( grouped_dy.columnwise_data.view(dtype=data_dtype).view(total_tokens, out_features).T ) b_tensor = grouped_x.columnwise_data.view(dtype=data_dtype).view(total_tokens, in_features) - sfa_tensor = grouped_dy.columnwise_scale_inv.view(sfa_leading_dim, -1).view( - dtype=scale_view_dtype - ) - sfb_tensor = grouped_x.columnwise_scale_inv.view(sfb_leading_dim, -1).view( - dtype=scale_view_dtype - ) + sfa_tensor = grouped_dy.columnwise_scale_inv.view(dtype=scale_view_dtype) + sfb_tensor = grouped_x.columnwise_scale_inv.view(dtype=scale_view_dtype) + if not getattr(wgrad_kernel_fn, "supports_canonical_layouts", False): + sfa_tensor = sfa_tensor.view(sfa_leading_dim, -1) + sfb_tensor = sfb_tensor.view(sfb_leading_dim, -1) common_wgrad_kwargs = { "a_tensor": a_tensor, @@ -1216,6 +1215,9 @@ def __init__( self._cudnn_canonical_layouts = getattr( self.grouped_gemm_activation_kernel(), "supports_canonical_layouts", False ) and getattr(self.grouped_gemm_quant_kernel(), "supports_canonical_layouts", False) + self._cudnn_canonical_backward = self._cudnn_canonical_layouts and getattr( + self.grouped_gemm_dactivation_kernel(), "supports_canonical_layouts", False + ) if not is_glu_activation(activation): # grouped_gemm_srelu_wrapper_sm100 and grouped_gemm_dsrelu_wrapper_sm100 are # SReLU-specific and do not take GLU ``act_func`` selectors. @@ -2436,9 +2438,11 @@ def fuser_backward( # Data logical shape: (sum(m), k, 1) # Scale logical shape: (32 (block row), 4 (block row), # sum(m)/128, 4 (block col), k/128, 1) + canonical = self._cudnn_canonical_backward and not use_nvfp4 fc2_dy_data = grouped_fc2_dy.rowwise_data.view(dtype=data_dtype) fc2_dy_data = fc2_dy_data.view(out_shape[0], data_k) - fc2_dy_data = fc2_dy_data.unsqueeze(0).permute(1, 2, 0) + if not canonical: + fc2_dy_data = fc2_dy_data.unsqueeze(0).permute(1, 2, 0) fc2_dy_scales = grouped_fc2_dy.scale_inv fc2_dy_scales = fc2_dy_scales.view(dtype=scale_view_dtype) with_gemm_swizzled_scales = grouped_fc2_dy._with_gemm_swizzled_scales @@ -2462,7 +2466,7 @@ def fuser_backward( 4, ) fc2_dy_scales = fc2_dy_scales.permute(3, 2, 1, 5, 4, 0) - else: + elif not canonical: fc2_dy_scales = fc2_dy_scales.view( 1, ceil_div(out_shape[0], 128), @@ -2504,7 +2508,7 @@ def fuser_backward( dscales_tensor = None if not unit_activation_scale: scales_f32 = scales.detach().to(dtype=torch.float32) - scales_tensor = scales_f32.reshape(-1, 1, 1) + scales_tensor = scales_f32.reshape((-1,) if canonical else (-1, 1, 1)) dscales_tensor = torch.zeros_like(scales_tensor) fc2_d_dtype = torch.bfloat16 if use_nvfp4 else torch.float8_e4m3fn @@ -2546,7 +2550,7 @@ def fuser_backward( fc2_dactivation_kwargs = { "a_tensor": fc2_dy_data, - "c_tensor": activation_in.unsqueeze(0).permute(1, 2, 0), + "c_tensor": activation_in if canonical else activation_in.unsqueeze(0).permute(1, 2, 0), "sfa_tensor": fc2_dy_scales, "padded_offsets": split_points, "alpha_tensor": fc2_alpha_tensor, @@ -2604,21 +2608,26 @@ def fuser_backward( fc2_w_data = fc2_weight_for_gemm.columnwise_data fc2_w_data = fc2_w_data.view(dtype=data_dtype) fc2_w_data = fc2_w_data.view(num_groups, 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 = fc2_weight_for_gemm.columnwise_scale_inv.view(dtype=scale_view_dtype) - fc2_w_scales = fc2_w_scales.view( - num_groups, - 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) - ) + if canonical: + fc2_dactivation_kwargs["b_major"] = "n" + else: + fc2_w_data = ( + fc2_w_data.permute(1, 2, 0) if use_nvfp4 else fc2_w_data.permute(2, 1, 0) + ) + fc2_w_scales = fc2_w_scales.view( + num_groups, + 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 @@ -2698,17 +2707,18 @@ def fuser_backward( else: fc1_dy_bf16 = None fc1_dy_row_data = fc2_dgrad_kernel_out["d_row_tensor"] - fc1_dy_row_data = fc1_dy_row_data.view(out_shape[0], fc1_weight_shape[0]) - # View scale in their actual swizzled shape - fc1_dy_row_scale = ( - fc2_dgrad_kernel_out["sfd_row_tensor"].permute(5, 2, 4, 0, 1, 3).view(-1) - ) + fc1_dy_row_scale = fc2_dgrad_kernel_out["sfd_row_tensor"] fc1_dy_col_data = fc2_dgrad_kernel_out["d_col_tensor"] - fc1_dy_col_data = fc1_dy_col_data.view(out_shape[0], fc1_weight_shape[0]) - # View scale in their actual swizzled shape - fc1_dy_col_scale = ( - fc2_dgrad_kernel_out["sfd_col_tensor"].permute(5, 2, 4, 0, 1, 3).view(-1) - ) + fc1_dy_col_scale = fc2_dgrad_kernel_out["sfd_col_tensor"] + if not canonical: + fc1_dy_row_data = fc1_dy_row_data.view(out_shape[0], fc1_weight_shape[0]) + # View scale in their actual swizzled shape + fc1_dy_row_scale = fc1_dy_row_scale.permute(5, 2, 4, 0, 1, 3) + fc1_dy_col_data = fc1_dy_col_data.view(out_shape[0], fc1_weight_shape[0]) + # View scale in their actual swizzled shape + fc1_dy_col_scale = fc1_dy_col_scale.permute(5, 2, 4, 0, 1, 3) + fc1_dy_row_scale = fc1_dy_row_scale.view(-1) + fc1_dy_col_scale = fc1_dy_col_scale.view(-1) grad_scales = fc2_dgrad_kernel_out["dprob_tensor"] if grad_scales is not None: grad_scales = grad_scales.view(-1) @@ -3052,19 +3062,22 @@ def fuser_backward( fc1_w_data = fc1_w_data.view( num_groups, fc1_weight_shape[0], fc1_weight_shape[1] ) - fc1_w_data = fc1_w_data.permute(2, 1, 0) fc1_w_scales = fc1_weight_for_gemm.columnwise_scale_inv.view( dtype=torch.float8_e8m0fnu ) - fc1_w_scales = fc1_w_scales.view( - num_groups, - ceil_div(fc1_weight_shape[1], 128), - ceil_div(fc1_weight_shape[0], 128), - MXFP8_BLOCK_SCALING_SIZE, - 4, - 4, - ) - fc1_w_scales = fc1_w_scales.permute(3, 4, 1, 5, 2, 0) + if canonical: + fc1_dgrad_kwargs["b_major"] = "n" + else: + fc1_w_data = fc1_w_data.permute(2, 1, 0) + fc1_w_scales = fc1_w_scales.view( + num_groups, + ceil_div(fc1_weight_shape[1], 128), + ceil_div(fc1_weight_shape[0], 128), + MXFP8_BLOCK_SCALING_SIZE, + 4, + 4, + ) + fc1_w_scales = fc1_w_scales.permute(3, 4, 1, 5, 2, 0) fc1_dgrad_kwargs["b_tensor"] = fc1_w_data fc1_dgrad_kwargs["sfb_tensor"] = fc1_w_scales @@ -3089,9 +3102,13 @@ def fuser_backward( grad_input_buffer = validate_or_alloc_output( grad_input_buffer, in_shape, dtype, device ) - fc1_dgrad_kwargs["d_tensor"] = grad_input_buffer.as_strided( - (out_shape[0], fc1_weight_shape[1], 1), - (fc1_weight_shape[1], 1, out_shape[0] * fc1_weight_shape[1]), + fc1_dgrad_kwargs["d_tensor"] = ( + grad_input_buffer + if canonical + else grad_input_buffer.as_strided( + (out_shape[0], fc1_weight_shape[1], 1), + (fc1_weight_shape[1], 1, out_shape[0] * fc1_weight_shape[1]), + ) ) fc1_dgrad_kernel(**fc1_dgrad_kwargs) grad_input = grad_input_buffer From 756f427eb77717ca3d97727a144e83be267fdfe1 Mon Sep 17 00:00:00 2001 From: Hwanseo Choi Date: Tue, 6 Oct 2026 16:02:28 -0700 Subject: [PATCH 4/6] [PyTorch] Test and document canonical cuDNN layouts in fused grouped MLP backward Extend the canonical-vs-legacy test to the dGLU and FC1 dgrad calls, and note the backward in the fused op docstring. Signed-off-by: Hwanseo Choi --- tests/pytorch/test_grouped_mlp.py | 22 +++++++++++++++---- .../pytorch/ops/fused/grouped_mlp.py | 4 +++- 2 files changed, 21 insertions(+), 5 deletions(-) diff --git a/tests/pytorch/test_grouped_mlp.py b/tests/pytorch/test_grouped_mlp.py index 5c65248317..85aec825b6 100644 --- a/tests/pytorch/test_grouped_mlp.py +++ b/tests/pytorch/test_grouped_mlp.py @@ -1775,6 +1775,7 @@ def test_grouped_mlp_canonical_layouts_match_legacy(self, monkeypatch) -> None: kernels = { "grouped_gemm_activation_kernel": fused_cls.grouped_gemm_activation_kernel(), "grouped_gemm_quant_kernel": fused_cls.grouped_gemm_quant_kernel(), + "grouped_gemm_dactivation_kernel": fused_cls.grouped_gemm_dactivation_kernel(), } if not all( getattr(kernel, "supports_canonical_layouts", False) for kernel in kernels.values() @@ -1807,7 +1808,16 @@ def test_grouped_mlp_canonical_layouts_match_legacy(self, monkeypatch) -> None: splits = torch.full((4,), 256, dtype=torch.int64, device="cuda") dy = torch.randn(1024, 256, device="cuda", dtype=torch.bfloat16) - operands = ("a_tensor", "b_tensor", "sfa_tensor", "sfb_tensor", "prob_tensor", "d_tensor") + operands = ( + "a_tensor", + "b_tensor", + "c_tensor", + "sfa_tensor", + "sfb_tensor", + "prob_tensor", + "dprob_tensor", + "d_tensor", + ) def run(canonical): layouts = [] @@ -1849,12 +1859,16 @@ def traced(**kwargs): legacy, _ = run(canonical=False) canonical, layouts = run(canonical=True) - # The forward FC1 and FC2 calls get every operand as its natural contiguous buffer. - for call in layouts[:2]: + # FC1 GLU, FC2 quant, FC2 dGLU, and FC1 dgrad quant get every operand as its + # natural contiguous buffer. + assert len(layouts) == 4 + for call in layouts: assert all(contiguous for _, contiguous in call.values()) assert call["a_tensor"][0] == 2 and call["b_tensor"][0] == 3 + assert call.get("c_tensor", (2,))[0] == 2 assert call.get("prob_tensor", (1,))[0] == 1 - assert layouts[1]["d_tensor"][0] == 2 + assert call.get("dprob_tensor", (1,))[0] == 1 + assert layouts[1]["d_tensor"][0] == 2 and layouts[3]["d_tensor"][0] == 2 for actual, expected in zip(canonical, legacy): torch.testing.assert_close(actual, expected, rtol=0, atol=0) diff --git a/transformer_engine/pytorch/ops/fused/grouped_mlp.py b/transformer_engine/pytorch/ops/fused/grouped_mlp.py index aa3478dad7..4b6a944a56 100644 --- a/transformer_engine/pytorch/ops/fused/grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/grouped_mlp.py @@ -1117,7 +1117,7 @@ class _GroupedMLP_CuTeGEMMBase(FusedOperation): caller output/grad_input buffers are provided, the GEMMs write into them directly. When the front-end wrappers advertise ``supports_canonical_layouts``, MXFP8 operands are passed in their natural layouts instead of kernel-facing - views. + views, in the forward and, with a supporting dGLU wrapper, the backward. """ @@ -2610,6 +2610,7 @@ def fuser_backward( fc2_w_data = fc2_w_data.view(num_groups, fc2_weight_shape[0], fc2_weight_k) fc2_w_scales = fc2_weight_for_gemm.columnwise_scale_inv.view(dtype=scale_view_dtype) if canonical: + # The column-wise weight is (num_groups, k, n): an n-major B for this GEMM. fc2_dactivation_kwargs["b_major"] = "n" else: fc2_w_data = ( @@ -3066,6 +3067,7 @@ def fuser_backward( dtype=torch.float8_e8m0fnu ) if canonical: + # The column-wise weight is (num_groups, k, n): an n-major B for this GEMM. fc1_dgrad_kwargs["b_major"] = "n" else: fc1_w_data = fc1_w_data.permute(2, 1, 0) From 76581d8accfc4cdcdf264b62ae128093ea12e70d Mon Sep 17 00:00:00 2001 From: Hwanseo Choi Date: Thu, 15 Oct 2026 00:03:49 -0700 Subject: [PATCH 5/6] [PyTorch] Keep the typed cuDNN weight views on the grouped weight The FC1 and FC2 weight buffers persist across steps, so the dtype and shape views cuDNN takes are built once per buffer instead of with two torch ops before each grouped GEMM launch (forward FC1 and FC2, backward FC2 dGLU and FC1 dgrad). A changed buffer is detected by its data pointer. Signed-off-by: Hwanseo Choi --- .../pytorch/ops/fused/grouped_mlp.py | 60 ++++++++++++++----- 1 file changed, 45 insertions(+), 15 deletions(-) diff --git a/transformer_engine/pytorch/ops/fused/grouped_mlp.py b/transformer_engine/pytorch/ops/fused/grouped_mlp.py index 4b6a944a56..a2fcf5a708 100644 --- a/transformer_engine/pytorch/ops/fused/grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/grouped_mlp.py @@ -521,6 +521,25 @@ def _single_quantized_tensor_from_grouped( ) +def cudnn_weight_view( + weight: GroupedTensor, + attr: str, + data: torch.Tensor, + dtype: torch.dtype, + shape: tuple[int, ...], +) -> torch.Tensor: + """``data.view(dtype=dtype).view(shape)``, kept on ``weight`` as ``attr``. + + Weight buffers persist across steps, so the typed view is built once per + buffer instead of with two torch ops before every kernel launch. + """ + view = getattr(weight, attr, None) + if view is None or view.data_ptr() != data.data_ptr(): + view = data.view(dtype=dtype).view(shape) + setattr(weight, attr, view) + return view + + def _pack_grouped_linear_bias_for_cudnn(linear_op: GroupedLinear) -> Optional[torch.Tensor]: """Bias layout expected by cuDNN grouped GEMM: shape (n, num_groups), stride (1, n).""" if not linear_op.has_bias: @@ -1766,8 +1785,13 @@ def fuser_forward( # Pack weight tensors for stacked kernel # Data actual shape: (num_groups, n, k) # Data logical shape: (n, k, num_groups) - 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 = cudnn_weight_view( + grouped_fc1_weight, + "_cudnn_rowwise_view", + fc1_w_data, + data_dtype, + (num_groups, fc1_weight_shape[0], fc1_weight_k), + ) fc1_w_scales = fc1_w_scales.view(dtype=scale_view_dtype) if not canonical: fc1_w_data = fc1_w_data.permute(1, 2, 0) @@ -2137,12 +2161,12 @@ def fuser_forward( 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 = cudnn_weight_view( + grouped_fc2_weight, + "_cudnn_rowwise_view", + fc2_weight_for_gemm.rowwise_data, + torch.float8_e4m3fn, + (num_groups, fc2_weight_shape[0], fc2_weight_shape[1]), ) fc2_w_scales = fc2_weight_for_gemm.scale_inv.view(dtype=torch.float8_e8m0fnu) if not canonical: @@ -2605,9 +2629,13 @@ def fuser_backward( # Pack weight tensors for stacked kernel # Data actual shape: (num_groups, k, n) # Data logical shape: (n, k, num_groups) - fc2_w_data = fc2_weight_for_gemm.columnwise_data - fc2_w_data = fc2_w_data.view(dtype=data_dtype) - fc2_w_data = fc2_w_data.view(num_groups, fc2_weight_shape[0], fc2_weight_k) + fc2_w_data = cudnn_weight_view( + grouped_fc2_weight, + "_cudnn_columnwise_view", + fc2_weight_for_gemm.columnwise_data, + data_dtype, + (num_groups, fc2_weight_shape[0], fc2_weight_k), + ) fc2_w_scales = fc2_weight_for_gemm.columnwise_scale_inv.view(dtype=scale_view_dtype) if canonical: # The column-wise weight is (num_groups, k, n): an n-major B for this GEMM. @@ -3058,10 +3086,12 @@ def fuser_backward( fc1_weight_for_gemm, rowwise=False, columnwise=True ) - fc1_w_data = fc1_weight_for_gemm.columnwise_data - fc1_w_data = fc1_w_data.view(dtype=torch.float8_e4m3fn) - fc1_w_data = fc1_w_data.view( - num_groups, fc1_weight_shape[0], fc1_weight_shape[1] + fc1_w_data = cudnn_weight_view( + grouped_fc1_weight, + "_cudnn_columnwise_view", + fc1_weight_for_gemm.columnwise_data, + torch.float8_e4m3fn, + (num_groups, fc1_weight_shape[0], fc1_weight_shape[1]), ) fc1_w_scales = fc1_weight_for_gemm.columnwise_scale_inv.view( dtype=torch.float8_e8m0fnu From 4a3dd5af308e99af9ea9c7c2539be1f7314d9f4a Mon Sep 17 00:00:00 2001 From: Hwanseo Choi Date: Thu, 15 Oct 2026 15:04:43 -0700 Subject: [PATCH 6/6] [PyTorch] Simplify the cached cuDNN weight view helper cudnn_weight_data takes the grouped weight and reads its row- or column-wise buffer itself, so call sites pass only the dtype and shape. Signed-off-by: Hwanseo Choi --- .../pytorch/ops/fused/grouped_mlp.py | 46 +++++++------------ 1 file changed, 16 insertions(+), 30 deletions(-) diff --git a/transformer_engine/pytorch/ops/fused/grouped_mlp.py b/transformer_engine/pytorch/ops/fused/grouped_mlp.py index a2fcf5a708..cf03c3dd73 100644 --- a/transformer_engine/pytorch/ops/fused/grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/grouped_mlp.py @@ -521,18 +521,16 @@ def _single_quantized_tensor_from_grouped( ) -def cudnn_weight_view( - weight: GroupedTensor, - attr: str, - data: torch.Tensor, - dtype: torch.dtype, - shape: tuple[int, ...], +def cudnn_weight_data( + weight: GroupedTensor, dtype: torch.dtype, shape: tuple[int, ...], columnwise: bool = False ) -> torch.Tensor: - """``data.view(dtype=dtype).view(shape)``, kept on ``weight`` as ``attr``. + """The weight's row- or column-wise data as the typed tensor cuDNN takes. - Weight buffers persist across steps, so the typed view is built once per - buffer instead of with two torch ops before every kernel launch. + The view is kept on ``weight``, so a persistent weight is viewed once + instead of with two torch ops before every kernel launch. """ + data = weight.columnwise_data if columnwise else weight.rowwise_data + attr = "_cudnn_columnwise_data" if columnwise else "_cudnn_rowwise_data" view = getattr(weight, attr, None) if view is None or view.data_ptr() != data.data_ptr(): view = data.view(dtype=dtype).view(shape) @@ -1771,7 +1769,6 @@ def fuser_forward( 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( @@ -1779,18 +1776,13 @@ def fuser_forward( 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 = cudnn_weight_view( - grouped_fc1_weight, - "_cudnn_rowwise_view", - fc1_w_data, - data_dtype, - (num_groups, fc1_weight_shape[0], fc1_weight_k), + fc1_w_data = cudnn_weight_data( + grouped_fc1_weight, data_dtype, (num_groups, fc1_weight_shape[0], fc1_weight_k) ) fc1_w_scales = fc1_w_scales.view(dtype=scale_view_dtype) if not canonical: @@ -2161,12 +2153,8 @@ def fuser_forward( columnwise=False, ) - fc2_w_data = cudnn_weight_view( - grouped_fc2_weight, - "_cudnn_rowwise_view", - fc2_weight_for_gemm.rowwise_data, - torch.float8_e4m3fn, - (num_groups, fc2_weight_shape[0], fc2_weight_shape[1]), + fc2_w_data = cudnn_weight_data( + grouped_fc2_weight, torch.float8_e4m3fn, (num_groups, *fc2_weight_shape) ) fc2_w_scales = fc2_weight_for_gemm.scale_inv.view(dtype=torch.float8_e8m0fnu) if not canonical: @@ -2629,12 +2617,11 @@ def fuser_backward( # Pack weight tensors for stacked kernel # Data actual shape: (num_groups, k, n) # Data logical shape: (n, k, num_groups) - fc2_w_data = cudnn_weight_view( + fc2_w_data = cudnn_weight_data( grouped_fc2_weight, - "_cudnn_columnwise_view", - fc2_weight_for_gemm.columnwise_data, data_dtype, (num_groups, fc2_weight_shape[0], fc2_weight_k), + columnwise=True, ) fc2_w_scales = fc2_weight_for_gemm.columnwise_scale_inv.view(dtype=scale_view_dtype) if canonical: @@ -3086,12 +3073,11 @@ def fuser_backward( fc1_weight_for_gemm, rowwise=False, columnwise=True ) - fc1_w_data = cudnn_weight_view( + fc1_w_data = cudnn_weight_data( grouped_fc1_weight, - "_cudnn_columnwise_view", - fc1_weight_for_gemm.columnwise_data, torch.float8_e4m3fn, - (num_groups, fc1_weight_shape[0], fc1_weight_shape[1]), + (num_groups, *fc1_weight_shape), + columnwise=True, ) fc1_w_scales = fc1_weight_for_gemm.columnwise_scale_inv.view( dtype=torch.float8_e8m0fnu