Add a backwards linear function to be used with the fused mla q up-proj

transformer_engine/pytorch/attention/fused_mla_q_uproj.py

@@ -142,6 +142,34 @@ def run(
         # 2nd return is the activation to save for wgrad: MXFP8 (fp8 path) or bf16 (16-bit path).
         return query, x_saved
 
+    @classmethod
+    def backward_linear(
+        cls,
+        grad_output,
+        x_saved,
+        w_q,
+        act_dtype,
+        wgrad_store,
+        fuse_wgrad_accumulation,
+        tp_group,
+        sequence_parallel,
+        **kwargs,
+    ):
+        """Linear backward for the fused Q up-proj — delegates to :func:`~transformer_engine.pytorch.module.linear.backward_linear`."""
+        from ..module.linear import backward_linear as _bwd
+
+        return _bwd(
+            grad_output,
+            x_saved,
+            w_q,
+            act_dtype,
+            wgrad_store,
+            fuse_wgrad_accumulation,
+            tp_group,
+            sequence_parallel,
+            **kwargs,
+        )
+
     @classmethod
     def wrap_mxfp8(
         cls,

transformer_engine/pytorch/module/linear.py

@@ -786,6 +786,92 @@ def _linear_setup_ctx(
     return (saved_inputmat, wt_save, saved_weight, saved_bias)
 
 
+def backward_linear(
+    grad_output: torch.Tensor,
+    x_saved,
+    w_q,
+    act_dtype: torch.dtype,
+    wgrad_store,
+    fuse_wgrad_accumulation: bool,
+    tp_group,
+    sequence_parallel: bool,
+    *,
+    use_bias: bool = False,
+    requires_dgrad: bool = True,
+    requires_wgrad: bool = True,
+    parallel_mode: str = "column",
+    backward_input_needs_gather: bool = False,
+) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor], Optional[torch.Tensor]]:
+    """Linear backward for fused operations that bypass TE's autograd chain.
+
+    Wraps :func:`_linear_backward` with a simplified interface for callers
+    (e.g. Megatron's fused MLA Q up-proj) that run their own forward kernel
+    and need to delegate the projection backward to TE.
+
+    Args:
+        grad_output: upstream gradient (e.g. post-RoPE-backward) ``[tokens, out_features]``.
+        x_saved: activation saved from the forward (``MXFP8Tensor`` or bf16).
+        w_q: weight (``MXFP8Tensor`` for FP8 path, bf16 tensor otherwise).
+        act_dtype: output dtype for the dgrad tensor.
+        wgrad_store: optional deferred weight-grad store.
+        fuse_wgrad_accumulation: accumulate wgrad directly into ``w_q.main_grad``.
+        tp_group: tensor-parallel process group (or ``None``).
+        sequence_parallel: whether sequence parallelism is active.
+        use_bias: compute a bias gradient (default ``False``).
+        requires_dgrad: compute dgrad (default ``True``).
+        requires_wgrad: compute wgrad (default ``True``).
+        parallel_mode: cuBLAS parallel mode (default ``"column"``).
+        backward_input_needs_gather: all-gather ``x_saved`` before the wgrad
+            GEMM (default ``False`` — assumes fused forward pre-gathers).
+
+    Returns:
+        ``(dgrad, wgrad, grad_bias)`` — ``wgrad`` is a typed dummy when
+        ``fuse_wgrad_accumulation=True``; ``grad_bias`` is ``None`` when
+        ``use_bias=False``.
+    """
+    tp_size = get_distributed_world_size(tp_group) if tp_group is not None else 1
+    fp8 = isinstance(w_q, QuantizedTensor)
+
+    grad_output_quantizer = None
+    if fp8:
+        grad_output_quantizer = MXFP8Quantizer(
+            fp8_dtype=tex.DType.kFloat8E4M3, rowwise=True, columnwise=True
+        )
+        grad_output_quantizer.optimize_for_gemm = True
+
+    bwd_args = LinearBwdArgs(
+        grad_output=grad_output,
+        inputmat=x_saved,
+        weight_fp8=w_q,
+        saved_weight=w_q,
+        bias=None,
+        grad_output_quantizer=grad_output_quantizer,
+        use_bias=use_bias,
+        requires_dgrad=requires_dgrad,
+        requires_wgrad=requires_wgrad,
+        inp_shape=x_saved.shape,
+        activation_dtype=act_dtype,
+        fp8=fp8,
+        dgrad_use_split_accumulator=_2X_ACC_DGRAD,
+        wgrad_use_split_accumulator=_2X_ACC_WGRAD,
+        is_weight_param_quantized=fp8,
+        parallel_mode=parallel_mode,
+        tp_group=tp_group,
+        tp_size=tp_size,
+        tensor_parallel=tp_size > 1,
+        sequence_parallel=sequence_parallel,
+        backward_input_needs_gather=backward_input_needs_gather,
+        is_fsdp2=False,
+        fuse_wgrad_accumulation=fuse_wgrad_accumulation,
+        wgrad_store=wgrad_store,
+        origin_weight_ref=weakref.ref(w_q) if fuse_wgrad_accumulation else None,
+        main_grad_func=(lambda: w_q.main_grad) if fuse_wgrad_accumulation else None,
+    )
+
+    wgrad, dgrad, grad_bias = _linear_backward(bwd_args)
+    return dgrad, wgrad, grad_bias
+
+
 def _linear_backward(args: LinearBwdArgs) -> Tuple[Union[torch.Tensor, None], ...]:
     """Backward implementation for the linear layer.