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 c176985254..2d25064c38 100644 --- a/transformer_engine/pytorch/attention/fused_mla_q_uproj.py +++ b/transformer_engine/pytorch/attention/fused_mla_q_uproj.py @@ -7,13 +7,22 @@ from __future__ import annotations import functools import os +import weakref from importlib.metadata import PackageNotFoundError, version as get_pkg_version import torch 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 from ..tensor.mxfp8_tensor import MXFP8Quantizer, MXFP8Tensor from ..utils import get_device_compute_capability @@ -31,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. @@ -142,6 +255,61 @@ 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, + ): + """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, + # 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, + 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, @@ -173,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, + )