From ed22445edc1311eff673a9d117d7f1493313fd85 Mon Sep 17 00:00:00 2001 From: Chase Block Date: Fri, 7 Aug 2026 13:43:50 -0700 Subject: [PATCH 1/7] Add a backwards linear function to be used with the fused mla q up-proj Signed-off-by: Chase Block --- .../pytorch/attention/fused_mla_q_uproj.py | 8 ++ transformer_engine/pytorch/module/linear.py | 87 +++++++++++++++++++ 2 files changed, 95 insertions(+) diff --git a/transformer_engine/pytorch/attention/fused_mla_q_uproj.py b/transformer_engine/pytorch/attention/fused_mla_q_uproj.py index c176985254..4d8ab3c207 100644 --- a/transformer_engine/pytorch/attention/fused_mla_q_uproj.py +++ b/transformer_engine/pytorch/attention/fused_mla_q_uproj.py @@ -142,6 +142,14 @@ 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, diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py index 56622db5e6..be793e3f23 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -786,6 +786,93 @@ 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]]: + """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)`` — ``wgrad`` is a typed dummy when + ``fuse_wgrad_accumulation=True``. + """ + import weakref + + 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, _ = _linear_backward(bwd_args) + return dgrad, wgrad + + def _linear_backward(args: LinearBwdArgs) -> Tuple[Union[torch.Tensor, None], ...]: """Backward implementation for the linear layer. From 346249c26e7c7c182375a9e66e98ef574e004bbe Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 7 Aug 2026 21:16:08 +0000 Subject: [PATCH 2/7] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- .../pytorch/attention/fused_mla_q_uproj.py | 28 ++++++++++++++++--- 1 file changed, 24 insertions(+), 4 deletions(-) diff --git a/transformer_engine/pytorch/attention/fused_mla_q_uproj.py b/transformer_engine/pytorch/attention/fused_mla_q_uproj.py index 4d8ab3c207..d09d2871e9 100644 --- a/transformer_engine/pytorch/attention/fused_mla_q_uproj.py +++ b/transformer_engine/pytorch/attention/fused_mla_q_uproj.py @@ -143,12 +143,32 @@ def run( 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): + 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) + + return _bwd( + grad_output, + x_saved, + w_q, + act_dtype, + wgrad_store, + fuse_wgrad_accumulation, + tp_group, + sequence_parallel, + **kwargs, + ) @classmethod def wrap_mxfp8( From 554d64438ffa659234126c6097e825948e3ade58 Mon Sep 17 00:00:00 2001 From: Chase Block Date: Fri, 7 Aug 2026 14:26:05 -0700 Subject: [PATCH 3/7] Remove redundant import in backward_linear Signed-off-by: Chase Block --- transformer_engine/pytorch/module/linear.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py index be793e3f23..a577ca7c3e 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -828,8 +828,6 @@ def backward_linear( ``(dgrad, wgrad)`` — ``wgrad`` is a typed dummy when ``fuse_wgrad_accumulation=True``. """ - import weakref - tp_size = get_distributed_world_size(tp_group) if tp_group is not None else 1 fp8 = isinstance(w_q, QuantizedTensor) From 050a4ed07aeac1209a11b752e340d8a0440a602b Mon Sep 17 00:00:00 2001 From: Chase Block Date: Mon, 10 Aug 2026 08:06:05 -0700 Subject: [PATCH 4/7] Handle biad gradient in lin bwd wrapper. Signed-off-by: Chase Block --- transformer_engine/pytorch/module/linear.py | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py index a577ca7c3e..1950096fb2 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -801,7 +801,7 @@ def backward_linear( requires_wgrad: bool = True, parallel_mode: str = "column", backward_input_needs_gather: bool = False, -) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor]]: +) -> 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 @@ -825,8 +825,9 @@ def backward_linear( GEMM (default ``False`` — assumes fused forward pre-gathers). Returns: - ``(dgrad, wgrad)`` — ``wgrad`` is a typed dummy when - ``fuse_wgrad_accumulation=True``. + ``(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) @@ -867,8 +868,8 @@ def backward_linear( main_grad_func=(lambda: w_q.main_grad) if fuse_wgrad_accumulation else None, ) - wgrad, dgrad, _ = _linear_backward(bwd_args) - return dgrad, wgrad + wgrad, dgrad, grad_bias = _linear_backward(bwd_args) + return dgrad, wgrad, grad_bias def _linear_backward(args: LinearBwdArgs) -> Tuple[Union[torch.Tensor, None], ...]: From 9e19715810d657648b2d936c1dbdc39bb8c9f908 Mon Sep 17 00:00:00 2001 From: Chase Block Date: Tue, 11 Aug 2026 07:36:38 -0700 Subject: [PATCH 5/7] Move linear backward function to fused_mla_q_uproj.py Signed-off-by: Chase Block --- .../pytorch/attention/fused_mla_q_uproj.py | 54 +++++++++--- transformer_engine/pytorch/module/linear.py | 85 ------------------- 2 files changed, 40 insertions(+), 99 deletions(-) diff --git a/transformer_engine/pytorch/attention/fused_mla_q_uproj.py b/transformer_engine/pytorch/attention/fused_mla_q_uproj.py index d09d2871e9..d13a7baf96 100644 --- a/transformer_engine/pytorch/attention/fused_mla_q_uproj.py +++ b/transformer_engine/pytorch/attention/fused_mla_q_uproj.py @@ -7,6 +7,7 @@ from __future__ import annotations import functools import os +import weakref from importlib.metadata import PackageNotFoundError, version as get_pkg_version import torch @@ -14,6 +15,7 @@ from packaging.version import Version as PkgVersion from ..constants import MXFP8_BLOCK_SCALING_SIZE +from ..distributed import get_distributed_world_size from ..quantized_tensor import QuantizedTensor from ..tensor.mxfp8_tensor import MXFP8Quantizer, MXFP8Tensor from ..utils import get_device_compute_capability @@ -153,23 +155,47 @@ def backward_linear( 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, + """Linear backward for the fused Q up-proj.""" + from ..module.linear import LinearBwdArgs, _linear_backward, _2X_ACC_DGRAD, _2X_ACC_WGRAD + + 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, + grad_output_quantizer=grad_output_quantizer, + 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="column", + tp_group=tp_group, + tp_size=tp_size, + tensor_parallel=tp_size > 1, + sequence_parallel=sequence_parallel, + 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 + @classmethod def wrap_mxfp8( cls, diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py index 1950096fb2..99b321826b 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -786,91 +786,6 @@ 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. From 9da364fec53b498542751211dab17177d4b3139e Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Tue, 11 Aug 2026 14:40:10 +0000 Subject: [PATCH 6/7] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- transformer_engine/pytorch/module/linear.py | 1 - 1 file changed, 1 deletion(-) diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py index 99b321826b..56622db5e6 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -786,7 +786,6 @@ def _linear_setup_ctx( return (saved_inputmat, wt_save, saved_weight, saved_bias) - def _linear_backward(args: LinearBwdArgs) -> Tuple[Union[torch.Tensor, None], ...]: """Backward implementation for the linear layer. From 470c1251eb60008c039d6a2b32f26cff91bdafc7 Mon Sep 17 00:00:00 2001 From: Sudhakar Singh Date: Tue, 11 Aug 2026 15:15:42 -0700 Subject: [PATCH 7/7] Add fused MLA Q up-projection backward Port the RoPE and projection backward needed by the fused Q up-projection path. Keep the one-use Triton kernel colocated and optional, and request both gradients because this temporary wrapper returns both. Cover the real MXFP8 autograd path so dgrad and wgrad are verified end to end. Signed-off-by: Sudhakar Singh --- .../attention/test_fused_mla_q_uproj.py | 92 ++++++- .../pytorch/attention/fused_mla_q_uproj.py | 224 ++++++++++++++++++ 2 files changed, 311 insertions(+), 5 deletions(-) diff --git a/tests/pytorch/attention/test_fused_mla_q_uproj.py b/tests/pytorch/attention/test_fused_mla_q_uproj.py index 2c061607fc..52c59e885b 100644 --- a/tests/pytorch/attention/test_fused_mla_q_uproj.py +++ b/tests/pytorch/attention/test_fused_mla_q_uproj.py @@ -14,6 +14,8 @@ import transformer_engine.pytorch # registers transformer_engine_torch import transformer_engine_torch as tex from transformer_engine.pytorch.attention import FusedMLAQUpProjRopeQuant +from transformer_engine.pytorch.attention.fused_mla_q_uproj import _FusedMLAQUpProjFunction +from transformer_engine.pytorch.cpp_extensions import general_gemm as _fused_general_gemm from transformer_engine.pytorch.tensor.mxfp8_tensor import MXFP8Quantizer, MXFP8Tensor # DSv3 671B MLA dims @@ -108,11 +110,7 @@ def _build_rope_tables(tokens: int, device: torch.device) -> tuple[torch.Tensor, @pytest.mark.skipif(not fused_supported, reason=reason_not_supported) @pytest.mark.parametrize("tokens", [256]) def test_fused_mla_q_uproj(tokens: int) -> None: - """Forward numerics and x_saved properties for FusedMLAQUpProjRopeQuant.run(). - - Full forward+backward autograd testing (via _FusedMLAQUpProjFunction) lives in - Megatron-Core. - """ + """Forward numerics and x_saved properties for FusedMLAQUpProjRopeQuant.run().""" s, b = tokens, 1 device = torch.device("cuda") torch.manual_seed(SEED) @@ -135,3 +133,87 @@ def test_fused_mla_q_uproj(tokens: int) -> None: assert isinstance(x_saved, MXFP8Tensor) assert x_saved._columnwise_data is not None, "x_saved must retain columnwise data for wgrad" assert x_saved._rowwise_data is None, "x_saved rowwise data should be dropped after forward" + + +@pytest.mark.skipif(not fused_supported, reason=reason_not_supported) +def test_fused_mla_q_uproj_autograd() -> None: + """The real autograd path must produce correct input and weight gradients.""" + import triton + + from transformer_engine.pytorch.attention.fused_mla_q_uproj import rotary_bwd_q_kernel + + tokens, s, b = 256, 256, 1 + device = torch.device("cuda") + torch.manual_seed(SEED) + + x = torch.randn(s, b, Q_LORA_RANK, dtype=torch.bfloat16, device=device, requires_grad=True) + w_bf16 = torch.randn( + PROJ_DIM, Q_LORA_RANK, dtype=torch.bfloat16, device=device, requires_grad=True + ) + w = MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3, rowwise=True, columnwise=True)(w_bf16) + cos, sin = _build_rope_tables(tokens, device) + cos_flat = cos.reshape(s, -1).contiguous() + sin_flat = sin.reshape(s, -1).contiguous() + _, x_saved = FusedMLAQUpProjRopeQuant.run( + x.detach().reshape(tokens, Q_LORA_RANK), w.detach(), cos_flat, sin_flat, s, b + ) + grad_out = torch.randn(s, b, NUM_HEADS, HEAD_DIM, dtype=torch.bfloat16, device=device) + + query = _FusedMLAQUpProjFunction.apply( + x, + w, + cos[:, None, None, :], + sin[:, None, None, :], + None, + False, + NUM_HEADS, + HEAD_DIM, + HEAD_DIM_NOPE, + HEAD_DIM_ROPE, + s, + b, + None, + False, + ) + assert query.requires_grad + assert query.grad_fn is not None + torch.autograd.backward(query, grad_out.clone()) + assert x.grad is not None + assert w_bf16.grad is not None + + dq3 = grad_out.reshape(tokens, NUM_HEADS, HEAD_DIM).clone().contiguous() + grid = lambda META: (tokens, triton.cdiv(NUM_HEADS, META["BLOCK_H"])) + rotary_bwd_q_kernel[grid]( + dq3, + cos_flat, + sin_flat, + HEAD_DIM_NOPE, + HEAD_DIM_ROPE, + NUM_HEADS, + 1, + None, + None, + dq3.stride(0), + dq3.stride(1), + 0, + 1, + ) + dq2d = dq3.reshape(tokens, PROJ_DIM).contiguous() + gy_quantizer = MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3, rowwise=True, columnwise=True) + gy_quantizer.optimize_for_gemm = True + gy = gy_quantizer(dq2d) + + w.update_usage(rowwise_usage=True, columnwise_usage=True) + grad_x_ref = _fused_general_gemm( + w, gy, layout="NN", grad=True, out_dtype=torch.bfloat16, use_split_accumulator=True + )[0] + grad_w_ref = _fused_general_gemm( + x_saved, + gy, + layout="NT", + grad=True, + out_dtype=torch.bfloat16, + use_split_accumulator=True, + )[0] + torch.testing.assert_close(x.grad.reshape(tokens, Q_LORA_RANK), grad_x_ref, atol=0.5, rtol=0.1) + torch.testing.assert_close(w_bf16.grad, grad_w_ref, atol=0.5, rtol=0.1) diff --git a/transformer_engine/pytorch/attention/fused_mla_q_uproj.py b/transformer_engine/pytorch/attention/fused_mla_q_uproj.py index d13a7baf96..2d25064c38 100644 --- a/transformer_engine/pytorch/attention/fused_mla_q_uproj.py +++ b/transformer_engine/pytorch/attention/fused_mla_q_uproj.py @@ -14,6 +14,13 @@ import transformer_engine_torch as tex from packaging.version import Version as PkgVersion +try: + import triton + import triton.language as tl +except ImportError: + triton = None + tl = None + from ..constants import MXFP8_BLOCK_SCALING_SIZE from ..distributed import get_distributed_world_size from ..quantized_tensor import QuantizedTensor @@ -33,6 +40,110 @@ def _cudnn_frontend_version_supported() -> bool: return False +if triton is not None: + + @triton.jit + def _get_thd_token_idx(cu_seqlens, pid_m, seq_num, cp_rank, cp_size): + token_idx = -1 + this_seq_len = 0 + seq_idx = 0 + last_cum_seqlen = tl.load(cu_seqlens) // cp_size + while seq_idx < seq_num: + cur_cum_seqlen = tl.load(cu_seqlens + seq_idx + 1) // cp_size + if token_idx == -1 and cur_cum_seqlen > pid_m: + token_idx = pid_m - last_cum_seqlen + this_seq_len = cur_cum_seqlen - last_cum_seqlen + last_cum_seqlen = cur_cum_seqlen + seq_idx += 1 + if cp_size > 1: + if token_idx < this_seq_len // 2: + token_idx = token_idx + cp_rank * this_seq_len // 2 + else: + token_idx = (token_idx - this_seq_len // 2) + ( + 2 * cp_size - cp_rank - 1 + ) * this_seq_len // 2 + return token_idx + + @triton.autotune( + configs=[ + triton.Config({"BLOCK_H": 1}), + triton.Config({"BLOCK_H": 2}), + triton.Config({"BLOCK_H": 4}), + triton.Config({"BLOCK_H": 8}), + triton.Config({"BLOCK_H": 16}), + triton.Config({"BLOCK_H": 32}), + triton.Config({"BLOCK_H": 64}), + triton.Config({"BLOCK_H": 128}), + ], + key=["emb_dim", "head_num"], + restore_value=["DO"], + ) + @triton.jit + def rotary_bwd_q_kernel( + DO, + COS, + SIN, + qk_head_dim, + emb_dim: tl.constexpr, + head_num: tl.constexpr, + batch_size, + seq_num, + cu_seqlens_q, + stride_x_seq, + stride_x_nheads, + cp_rank, + cp_size, + BLOCK_H: tl.constexpr, + ): + """ + Triton kernel of the backward pass for applying YARN RoPE to MLA's query. + This kernel inplace modifies the input tensor DO. + + Input: + DO: [seq_len, batch_size, head_num, qk_head_dim + emb_dim] + or [total_seq_len, head_num, qk_head_dim + emb_dim] + COS/SIN: [max_seq_len, emb_dim] + + batch_size, seq_num, and cu_seqlens_q are the same as in the forward pass + """ + pid_m = tl.program_id(axis=0) + pid_head = tl.program_id(axis=1) + + if cu_seqlens_q is None: + token_idx = pid_m // batch_size + else: + token_idx = _get_thd_token_idx(cu_seqlens_q, pid_m, seq_num, cp_rank, cp_size) + + cos_left = tl.load(COS + token_idx * emb_dim + tl.arange(0, emb_dim // 2)) + sin_left = tl.load(SIN + token_idx * emb_dim + tl.arange(0, emb_dim // 2)) + cos_right = tl.load(COS + token_idx * emb_dim + emb_dim // 2 + tl.arange(0, emb_dim // 2)) + sin_right = tl.load(SIN + token_idx * emb_dim + emb_dim // 2 + tl.arange(0, emb_dim // 2)) + cos_left = cos_left.expand_dims(0).broadcast_to(BLOCK_H, emb_dim // 2) + sin_left = sin_left.expand_dims(0).broadcast_to(BLOCK_H, emb_dim // 2) + cos_right = cos_right.expand_dims(0).broadcast_to(BLOCK_H, emb_dim // 2) + sin_right = sin_right.expand_dims(0).broadcast_to(BLOCK_H, emb_dim // 2) + + DO = DO + pid_m * stride_x_seq + pid_head * BLOCK_H * stride_x_nheads + + x_off = tl.arange(0, BLOCK_H)[:, None] * stride_x_nheads + qk_head_dim + mask = x_off < head_num * stride_x_nheads + x_left_off = x_off + tl.arange(0, emb_dim // 2)[None, :] + x_right_off = x_left_off + emb_dim // 2 + x_left = tl.load(DO + x_left_off, mask=mask) + x_right = tl.load(DO + x_right_off, mask=mask) + + x_1 = x_left * cos_left + x_right * sin_right + x_2 = -x_left * sin_left + x_right * cos_right + + x_1_off = x_off + tl.arange(0, emb_dim // 2)[None, :] * 2 + x_2_off = x_1_off + 1 + tl.store(DO + x_1_off, x_1, mask=mask) + tl.store(DO + x_2_off, x_2, mask=mask) + +else: + rotary_bwd_q_kernel = None + + class FusedMLAQUpProjRopeQuant: """Wrapper for the cuDNN fused MLA Q up-proj + per-head RoPE + MXFP8 quantize kernel. @@ -178,6 +289,9 @@ def backward_linear( inp_shape=x_saved.shape, activation_dtype=act_dtype, fp8=fp8, + # This temporary fused API always computes both projection gradients. + requires_dgrad=True, + requires_wgrad=True, dgrad_use_split_accumulator=_2X_ACC_DGRAD, wgrad_use_split_accumulator=_2X_ACC_WGRAD, is_weight_param_quantized=fp8, @@ -227,3 +341,113 @@ def wrap_mxfp8( fp8_dtype=tex.DType.kFloat8E4M3, with_gemm_swizzled_scales=False, ) + + +class _FusedMLAQUpProjFunction(torch.autograd.Function): + """Fused Q up-proj: q_normed -> (GEMM + per-head RoPE + MXFP8) -> MXFP8Tensor Q.""" + + @staticmethod + def forward( + ctx, + q_normed, # [s, b, q_lora_rank] bf16 (post-layernorm) + w_q, # [nh*q_head_dim, q_lora_rank] FP8 QuantizedTensor or bf16 (TE out×in layout) + cos, # [s, 1, 1, rope_dim] + sin, # [s, 1, 1, rope_dim] + wgrad_store, + fuse_wgrad_accumulation, + nh, + q_head_dim, + qk_head_dim, + qk_pos_emb_head_dim, + s, + b, + tp_group, # tensor-parallel process group + sequence_parallel, # True if sequence parallelism is active + ): + """Run the fused gemm + rope + mxfp8 quantization""" + + tokens = s * b + x = q_normed.detach().reshape(tokens, -1).contiguous() + + # Reshape [s, 1, 1, rope_dim] -> [s*b, rope_dim] bf16 as required by the KF kernel. + def _flat(t): + t = t.reshape(s, -1) + if b > 1: + t = t.unsqueeze(1).expand(s, b, t.shape[-1]).reshape(tokens, -1) + return t.to(torch.bfloat16).contiguous() + + cos, sin = _flat(cos), _flat(sin) + query, x_saved = FusedMLAQUpProjRopeQuant.run(x, w_q.detach(), cos, sin, s, b) + + ctx.save_for_backward(x_saved, w_q, cos, sin) + ctx.wgrad_store = wgrad_store + ctx.fuse_wgrad_accumulation = fuse_wgrad_accumulation + ctx.act_dtype = q_normed.dtype + ctx.dims = (nh, q_head_dim, qk_head_dim, qk_pos_emb_head_dim, s, b) + ctx.tp_group = tp_group + ctx.sequence_parallel = sequence_parallel + return query + + @staticmethod + def backward(ctx, dq): + """Backward is unfused and matches the typical backward pass""" + if rotary_bwd_q_kernel is None: + raise RuntimeError("Fused MLA Q up-projection backward requires Triton") + + x_saved, w_q, cos, sin = ctx.saved_tensors + nh, q_head_dim, qk_head_dim, qk_pos_emb_head_dim, s, b = ctx.dims + tokens = s * b + act_dtype = ctx.act_dtype + + # --- RoPE backward (unchanged: bf16, same rotary_bwd_q_kernel as the unfused path) --- + dq3 = dq.reshape(tokens, nh, q_head_dim).contiguous() + grid = lambda META: (tokens, triton.cdiv(nh, META["BLOCK_H"])) + rotary_bwd_q_kernel[grid]( + dq3, + cos.contiguous(), + sin.contiguous(), + qk_head_dim, + qk_pos_emb_head_dim, + nh, + 1, + None, + None, + dq3.stride(0), + dq3.stride(1), + 0, + 1, + ) + # grad w.r.t. the (pre-RoPE) up-proj GEMM output; bf16. + dq2d = dq3.reshape(tokens, nh * q_head_dim).contiguous() + + # Delegate the projection backward to TE's _linear_backward (via backward_linear) + grad_x, ret_grad_w, _ = FusedMLAQUpProjRopeQuant.backward_linear( + grad_output=dq2d, + x_saved=x_saved, + w_q=w_q, + act_dtype=act_dtype, + wgrad_store=ctx.wgrad_store, + fuse_wgrad_accumulation=ctx.fuse_wgrad_accumulation, + tp_group=ctx.tp_group, + sequence_parallel=ctx.sequence_parallel, + ) + grad_x = grad_x.reshape(s, b, -1) + + # grads for: q_normed, w_q, cos, sin, then 10 non-tensor args + # (including tp_group, sequence_parallel) + return ( + grad_x, + ret_grad_w, + None, + None, + None, + None, + None, + None, + None, + None, + None, + None, + None, + None, + )