diff --git a/tests/pytorch/test_grouped_mlp.py b/tests/pytorch/test_grouped_mlp.py index 88af92f0bb..85aec825b6 100644 --- a/tests/pytorch/test_grouped_mlp.py +++ b/tests/pytorch/test_grouped_mlp.py @@ -1762,6 +1762,116 @@ 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(), + "grouped_gemm_dactivation_kernel": fused_cls.grouped_gemm_dactivation_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", + "c_tensor", + "sfa_tensor", + "sfb_tensor", + "prob_tensor", + "dprob_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) + # 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 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) + @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 108ce5311c..cf03c3dd73 100644 --- a/transformer_engine/pytorch/ops/fused/grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/grouped_mlp.py @@ -521,6 +521,23 @@ def _single_quantized_tensor_from_grouped( ) +def cudnn_weight_data( + weight: GroupedTensor, dtype: torch.dtype, shape: tuple[int, ...], columnwise: bool = False +) -> torch.Tensor: + """The weight's row- or column-wise data as the typed tensor cuDNN takes. + + 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) + 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: @@ -741,12 +758,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, @@ -1116,7 +1132,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, in the forward and, with a supporting dGLU wrapper, the backward. """ @@ -1209,6 +1227,14 @@ 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) + 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. @@ -1564,9 +1590,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 +1620,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 +1641,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 @@ -1737,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( @@ -1745,25 +1776,26 @@ 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 = 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_data = cudnn_weight_data( + grouped_fc1_weight, data_dtype, (num_groups, fc1_weight_shape[0], fc1_weight_k) ) - fc1_w_scales = fc1_w_scales.permute(3, 4, 1, 5, 2, 0) + fc1_w_scales = fc1_w_scales.view(dtype=scale_view_dtype) + 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 +1822,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 +1865,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 +1931,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 +2117,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 +2137,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 @@ -2113,25 +2153,21 @@ 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_data( + grouped_fc2_weight, torch.float8_e4m3fn, (num_groups, *fc2_weight_shape) ) - 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 +2185,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 @@ -2410,9 +2450,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 @@ -2436,7 +2478,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), @@ -2478,7 +2520,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 @@ -2520,7 +2562,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, @@ -2575,24 +2617,33 @@ 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 = 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) + fc2_w_data = cudnn_weight_data( + grouped_fc2_weight, + 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: + # 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 = ( + 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 @@ -2672,17 +2723,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) @@ -3021,24 +3073,29 @@ 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_data( + grouped_fc1_weight, + torch.float8_e4m3fn, + (num_groups, *fc1_weight_shape), + columnwise=True, ) - 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: + # 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) + 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 @@ -3063,9 +3120,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