diff --git a/transformer_engine/pytorch/__init__.py b/transformer_engine/pytorch/__init__.py index 2b1803bfb2..90003122a8 100644 --- a/transformer_engine/pytorch/__init__.py +++ b/transformer_engine/pytorch/__init__.py @@ -29,7 +29,7 @@ from transformer_engine.pytorch.module import destroy_ub from transformer_engine.pytorch.module import UserBufferQuantizationMode from transformer_engine.pytorch.attention import DotProductAttention -from transformer_engine.pytorch.attention import FusedMLAQUpProjRopeQuant +from transformer_engine.pytorch.attention import FusedMLAQUpProjFunction, FusedMLAQUpProjRopeQuant from transformer_engine.pytorch.attention import MultiheadAttention from transformer_engine.pytorch.attention import InferenceParams from transformer_engine.pytorch.attention import RotaryPositionEmbedding diff --git a/transformer_engine/pytorch/attention/__init__.py b/transformer_engine/pytorch/attention/__init__.py index f6e4f0b37f..f171884116 100644 --- a/transformer_engine/pytorch/attention/__init__.py +++ b/transformer_engine/pytorch/attention/__init__.py @@ -5,13 +5,14 @@ """Python interface for attention""" from .dot_product_attention import DotProductAttention -from .fused_mla_q_uproj import FusedMLAQUpProjRopeQuant +from .fused_mla_q_uproj import FusedMLAQUpProjFunction, FusedMLAQUpProjRopeQuant from .multi_head_attention import MultiheadAttention from .inference import InferenceParams from .rope import RotaryPositionEmbedding __all__ = [ "DotProductAttention", + "FusedMLAQUpProjFunction", "FusedMLAQUpProjRopeQuant", "MultiheadAttention", "InferenceParams", diff --git a/transformer_engine/pytorch/attention/fused_mla_q_uproj.py b/transformer_engine/pytorch/attention/fused_mla_q_uproj.py index c176985254..5cb61301b8 100644 --- a/transformer_engine/pytorch/attention/fused_mla_q_uproj.py +++ b/transformer_engine/pytorch/attention/fused_mla_q_uproj.py @@ -2,18 +2,31 @@ # # See LICENSE for license information. -"""Fused MLA Q up-projection + per-head RoPE + MXFP8 quantize.""" +"""Fused MLA Q up-projection + per-head RoPE + MXFP8 quantize. + +The MLA RoPE Triton kernels are taken from: +Megatron-LM megatron/core/fusions/fused_mla_yarn_rope_apply.py +""" 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 +44,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. @@ -69,7 +186,7 @@ def is_supported(cls) -> bool: @classmethod def run( cls, - x: torch.Tensor, + x, # MXFP8Tensor when w is MXFP8 (already quantized by the norm), else bf16 Tensor w, # MXFP8Tensor (primary FP8 param) or bf16 torch.Tensor cos: torch.Tensor, sin: torch.Tensor, @@ -78,7 +195,10 @@ def run( ) -> "tuple[MXFP8Tensor, torch.Tensor]": """Run the fused kernel; return (Q MXFP8Tensor, activation saved for the wgrad backward). - The kernel precision is selected by the weight precision. + The kernel precision is selected by the weight precision. On the FP8 path ``x`` must + arrive already MXFP8-quantized: the caller's normalization emits MXFP8 straight from + its FP32 accumulator, exactly as `TELayerNormColumnParallelLinear` does, and quantizing + again here would round a second time and change the GEMM's input bytes. """ from cuda.bindings import driver as cuda @@ -92,13 +212,17 @@ def run( f" recipe), got {type(w).__name__}. Use the unfused path for other quantization" " recipes." ) - # ---- FP8 projection: MXFP8-cast x (both usages) + reuse w's fp8 codes -> mxfp8in ---- - # Quantize x with both rowwise (for the forward GEMM) and columnwise (for the FP8 - # wgrad in backward, matching the unfused path). - x_quantizer = MXFP8Quantizer( - fp8_dtype=tex.DType.kFloat8E4M3, rowwise=True, columnwise=True + # ---- FP8 projection: MXFP8 x straight from the norm + w's fp8 codes -> mxfp8in ---- + # x carries both usages: rowwise feeds this GEMM, columnwise the FP8 wgrad in backward. + x_mxfp8 = x + assert isinstance(x_mxfp8, QuantizedTensor) and hasattr(x_mxfp8, "_rowwise_data"), ( + "FusedMLAQUpProjRopeQuant needs an MXFP8-quantized input on the FP8 path, got" + f" {type(x_mxfp8).__name__}; have the normalization emit MXFP8 directly." + ) + assert not x_mxfp8._with_gemm_swizzled_scales, ( + "x scales must be unswizzled: the cuDNN kernel reads them as a plain" + " [tokens, K//32] array." ) - x_mxfp8 = x_quantizer(x) x_code = x_mxfp8._rowwise_data.view(torch.float8_e4m3fn) # [tokens, K] x_scale = x_mxfp8._rowwise_scale_inv # [tokens, K//32] uint8 @@ -142,6 +266,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 +352,153 @@ def wrap_mxfp8( fp8_dtype=tex.DType.kFloat8E4M3, with_gemm_swizzled_scales=False, ) + + +class FusedMLAQUpProjFunction(torch.autograd.Function): + """Fused Q up-proj: q -> (layernorm + MXFP8 quant) -> (GEMM + per-head RoPE + MXFP8) -> MXFP8Tensor Q.""" + + @staticmethod + def forward( + ctx, + x, # [s, b, q_lora_rank] pre-norm input + gamma, # [q_lora_rank] q_layernorm weight + 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 + eps, # norm epsilon + normalization, # "RMSNorm" + zero_centered_gamma, + ): + """Run the normalization (quantized output) then the fused gemm + rope + mxfp8""" + from ..module._common import apply_normalization + + tokens = s * b + tp_size = get_distributed_world_size(tp_group) if tp_group is not None else 1 + if tp_size > 1: + raise RuntimeError( + "FusedMLAQUpProjFunction does not support tensor parallelism (TP>1): " + "the backward dgrad is reduce-scattered over TP ranks but the caller " + "passes a pre-gathered full-sequence input. Use TP=1 or the unfused path." + ) + if normalization != "RMSNorm": + raise RuntimeError( + "FusedMLAQUpProjFunction supports RMSNorm only, got " + f"{normalization}; LayerNorm would also need mu saved in forward." + ) + x2d = x.detach().reshape(tokens, -1).contiguous() + fp8 = isinstance(w_q, QuantizedTensor) + + # Matches LayerNormLinear's input quantizer with one deliberate difference: + # optimize_for_gemm stays off, because the fused GEMM+RoPE+Quant cuDNN kernel reads the rowwise scales as a + # plain [tokens, K//32] array. + x_quantizer = None + if fp8: + x_quantizer = MXFP8Quantizer( + fp8_dtype=tex.DType.kFloat8E4M3, rowwise=True, columnwise=True + ) + + ln_out, _, rsigma = apply_normalization( + x2d, + None, + gamma, + None, + eps, + x_quantizer, + x2d.dtype, + normalization, + int(os.getenv("NVTE_FWD_LAYERNORM_SM_MARGIN", "0")), + zero_centered_gamma, + ) + + # 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(ln_out, w_q.detach(), cos, sin, s, b) + + ctx.save_for_backward(x2d, rsigma, gamma, x_saved, w_q, cos, sin) + ctx.wgrad_store = wgrad_store + ctx.fuse_wgrad_accumulation = fuse_wgrad_accumulation + ctx.act_dtype = x.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 + ctx.normalization = normalization + ctx.zero_centered_gamma = zero_centered_gamma + 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") + + x2d, rsigma, gamma, 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() + + def grid(META): + return (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_ln_out, 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, + ) + + # --- Norm backward, on the rsigma this forward saved --- + bwd_sm_margin = int(os.getenv("NVTE_BWD_LAYERNORM_SM_MARGIN", "0")) + grad_x, dgamma = tex.rmsnorm_bwd( + grad_ln_out.reshape(x2d.shape), + x2d, + rsigma, + gamma, + bwd_sm_margin, + ctx.zero_centered_gamma, + ) + grad_x = grad_x.reshape(s, b, -1) + + # grads for: x, gamma, w_q, then cos, sin and the 13 non-tensor args + return (grad_x, dgamma, ret_grad_w) + (None,) * 15