diff --git a/qa/L0_pytorch_unittest/test.sh b/qa/L0_pytorch_unittest/test.sh index 6a41515602..759432857a 100644 --- a/qa/L0_pytorch_unittest/test.sh +++ b/qa/L0_pytorch_unittest/test.sh @@ -60,6 +60,7 @@ NVTE_ALLOW_UNSAFE_PICKLE_EXTRA_STATE=1 python3 -m pytest --tb=auto --junitxml=$X python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_flex_attention.xml $TE_PATH/tests/pytorch/attention/test_flex_attention.py || test_fail "test_flex_attention.py" NVTE_ALLOW_UNSAFE_PICKLE_EXTRA_STATE=1 NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_attention_deterministic.xml $TE_PATH/tests/pytorch/attention/test_attention.py || test_fail "NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 test_attention.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_linear_mxfp8_attention.xml $TE_PATH/tests/pytorch/attention/test_linear_mxfp8_attention.py || test_fail "test_linear_mxfp8_attention.py" +python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_fused_mla_q_uproj.xml $TE_PATH/tests/pytorch/attention/test_fused_mla_q_uproj.py || test_fail "test_fused_mla_q_uproj.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_kv_cache.xml $TE_PATH/tests/pytorch/attention/test_kv_cache.py || test_fail "test_kv_cache.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_cu_seqlens_cache.xml $TE_PATH/tests/pytorch/attention/test_cu_seqlens_cache.py || test_fail "test_cu_seqlens_cache.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_hf_integration.xml $TE_PATH/tests/pytorch/test_hf_integration.py || test_fail "test_hf_integration.py" diff --git a/tests/pytorch/attention/test_fused_mla_q_uproj.py b/tests/pytorch/attention/test_fused_mla_q_uproj.py new file mode 100644 index 0000000000..13654deda6 --- /dev/null +++ b/tests/pytorch/attention/test_fused_mla_q_uproj.py @@ -0,0 +1,161 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Unit tests for FusedMLAQUpProjRopeQuant. + +Run: + pytest tests/pytorch/attention/test_fused_mla_q_uproj.py -v +""" + +import pytest +import torch + +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.cpp_extensions import general_gemm +from transformer_engine.pytorch.tensor.mxfp8_tensor import MXFP8Quantizer, MXFP8Tensor + +# DSv3 671B MLA dims +NUM_HEADS = 128 +HEAD_DIM_NOPE = 128 +HEAD_DIM_ROPE = 64 +HEAD_DIM = HEAD_DIM_NOPE + HEAD_DIM_ROPE # 192 +Q_LORA_RANK = 1536 +PROJ_DIM = NUM_HEADS * HEAD_DIM # 24576 + +SEED = 42 + +fused_supported, reason_not_supported = ( + (True, "") + if FusedMLAQUpProjRopeQuant.is_supported() + else ( + False, + ( + "FusedMLAQUpProjRopeQuant.is_supported() returned False " + "(SM100+, cudnn-frontend >= 1.27.0, and NVTE_FUSED_MLA_Q_UPROJ=1 required)" + ), + ) +) + + +def _dequantize_fused_output(query: MXFP8Tensor, s: int, b: int) -> torch.Tensor: + """Dequantize the rowwise fused output to bf16 [s, b, nh, head_dim]. + + TE's C++ dequantize kernel requires 2D layout, so reshape before calling dequantize(). + """ + tokens = s * b + q_2d = MXFP8Tensor( + shape=(tokens, PROJ_DIM), + dtype=torch.bfloat16, + rowwise_data=query._rowwise_data.view(tokens, PROJ_DIM), + rowwise_scale_inv=query._rowwise_scale_inv.view(tokens, PROJ_DIM // 32), + columnwise_data=None, + columnwise_scale_inv=None, + quantizer=query._quantizer, + requires_grad=False, + fp8_dtype=query._fp8_dtype, + with_gemm_swizzled_scales=False, + ) + return q_2d.dequantize().to(torch.bfloat16).view(s, b, NUM_HEADS, HEAD_DIM) + + +def _reference_q_uproj( + x: torch.Tensor, + w_mxfp8: MXFP8Tensor, + cos: torch.Tensor, + sin: torch.Tensor, + s: int, + b: int, +) -> torch.Tensor: + """Unfused bf16 reference: dequantize-then-GEMM + RoPE. Returns [s, b, nh, head_dim] bf16.""" + x_dq = ( + MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3, rowwise=True, columnwise=False)(x) + .dequantize() + .to(torch.bfloat16) + ) + w_dq = w_mxfp8.dequantize().to(torch.bfloat16) + out = (x_dq @ w_dq.t()).view(s, b, NUM_HEADS, HEAD_DIM) + + q_nope = out[..., :HEAD_DIM_NOPE] + q_rope = out[..., HEAD_DIM_NOPE:] + cos_ = cos[:, None, None, :].to(q_rope.dtype) + sin_ = sin[:, None, None, :].to(q_rope.dtype) + half = HEAD_DIM_ROPE // 2 + x1, x2 = q_rope[..., 0::2], q_rope[..., 1::2] + q_rope_out = torch.cat( + [ + x1 * cos_[..., :half] - x2 * sin_[..., :half], + x2 * cos_[..., half:] + x1 * sin_[..., half:], + ], + dim=-1, + ) + return torch.cat([q_nope, q_rope_out], dim=-1) + + +def _build_rope_tables(tokens: int, device: torch.device) -> tuple[torch.Tensor, torch.Tensor]: + inv_freq = 1.0 / ( + 10000 + ** (torch.arange(0, HEAD_DIM_ROPE, 2, dtype=torch.float32, device=device) / HEAD_DIM_ROPE) + ) + freqs = torch.cat( + [torch.outer(torch.arange(tokens, device=device, dtype=torch.float32), inv_freq)] * 2, + dim=-1, + ) + return freqs.cos().to(torch.bfloat16), freqs.sin().to(torch.bfloat16) + + +@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 + x_saved properties + backward dgrad/wgrad numerics. + + Forward is checked first so a broken kernel is caught before the backward assertions. + """ + s, b = tokens, 1 + device = torch.device("cuda") + torch.manual_seed(SEED) + torch.cuda.manual_seed(SEED) + + x = torch.randn(tokens, Q_LORA_RANK, dtype=torch.bfloat16, device=device) + # Backward needs columnwise data on w for the dgrad GEMM (general_gemm layout="NN" + # unwraps A via the columnwise direction). + w = MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3, rowwise=True, columnwise=True)( + torch.randn(PROJ_DIM, Q_LORA_RANK, dtype=torch.bfloat16, device=device) + ) + cos, sin = _build_rope_tables(tokens, device) + + query, x_saved = FusedMLAQUpProjRopeQuant.run(x, w, cos, sin, s, b) + + # --- Forward numerics --- + fused_dq = _dequantize_fused_output(query, s, b) + ref_dq = _reference_q_uproj(x, w, cos, sin, s, b) + torch.testing.assert_close(fused_dq, ref_dq, atol=0.5, rtol=0.1) + + # --- x_saved properties --- + 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" + + # --- Backward: dgrad + wgrad --- + grad_output = torch.randn(tokens, PROJ_DIM, dtype=torch.bfloat16, device=device) + grad_output_quantizer = MXFP8Quantizer( + fp8_dtype=tex.DType.kFloat8E4M3, rowwise=True, columnwise=True + ) + grad_output_quantizer.optimize_for_gemm = True + gy = grad_output_quantizer(grad_output) + + grad_x = general_gemm( + w, gy, layout="NN", grad=True, out_dtype=torch.bfloat16, use_split_accumulator=True + )[0] + grad_w = general_gemm( + x_saved, gy, layout="NT", grad=True, out_dtype=torch.bfloat16, use_split_accumulator=True + )[0] + + x_dq = x_saved.dequantize().to(torch.bfloat16) + w_dq = w.dequantize().to(torch.bfloat16) + gy_dq = gy.dequantize().to(torch.bfloat16) + + torch.testing.assert_close(grad_x, gy_dq @ w_dq, atol=0.5, rtol=0.1) + torch.testing.assert_close(grad_w, gy_dq.t() @ x_dq, atol=0.5, rtol=0.1) diff --git a/transformer_engine/pytorch/__init__.py b/transformer_engine/pytorch/__init__.py index 4c6b7fc67a..2b1803bfb2 100644 --- a/transformer_engine/pytorch/__init__.py +++ b/transformer_engine/pytorch/__init__.py @@ -29,6 +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 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 c4c2aa3e72..f6e4f0b37f 100644 --- a/transformer_engine/pytorch/attention/__init__.py +++ b/transformer_engine/pytorch/attention/__init__.py @@ -5,12 +5,14 @@ """Python interface for attention""" from .dot_product_attention import DotProductAttention +from .fused_mla_q_uproj import FusedMLAQUpProjRopeQuant from .multi_head_attention import MultiheadAttention from .inference import InferenceParams from .rope import RotaryPositionEmbedding __all__ = [ "DotProductAttention", + "FusedMLAQUpProjRopeQuant", "MultiheadAttention", "InferenceParams", "RotaryPositionEmbedding", diff --git a/transformer_engine/pytorch/attention/dot_product_attention/backends.py b/transformer_engine/pytorch/attention/dot_product_attention/backends.py index 891a5c661d..aeac501ff8 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/backends.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/backends.py @@ -1363,6 +1363,7 @@ def forward( deterministic, softmax_offset, fp8_output, + bf16_backward, layer_number, return_max_logit, packed_qkv=None, @@ -1427,6 +1428,9 @@ def forward( # fp8_dtype = tex.DType.kFloat8E4M3 if is_input_fp8: q_fp8, k_fp8, v_fp8 = q, k, v + + if fp8_recipe.mxfp8(): + qkv_scale_inv_format = "bhsd" # Same as what combine_and_quantize would give else: q_fp8, k_fp8, v_fp8, qkv_layout, qkv_scale_inv_format = combine_and_quantize( qkv_layout, @@ -1602,6 +1606,8 @@ def forward( ctx.is_input_fp8 = is_input_fp8 ctx.is_output_fp8 = is_output_fp8 + # Return dQ/dK/dV in bf16 even if is_input_fp8 + ctx.bf16_backward = bf16_backward tensors_to_save, tensor_objects = prepare_for_saving( *fp8_tensors, @@ -1860,7 +1866,8 @@ def backward(ctx, d_out, *_args): # dq, dk, dv: torch.Tensor; dtype = torch.float16 or torch.bfloat16 dq, dk, dv = dq_, dk_, dv_ is_quantized_tensor = isinstance(dq_, QuantizedTensorStorage) - if is_quantized_tensor and not ctx.is_input_fp8: + + if is_quantized_tensor and (not ctx.is_input_fp8 or ctx.bf16_backward): # return in F16 dq, dk, dv = combine_and_dequantize( ctx.dqkv_layout, @@ -1869,7 +1876,7 @@ def backward(ctx, d_out, *_args): dv_, src_nominal_dtype=dq_.dtype, ) - if not is_quantized_tensor and ctx.is_input_fp8: + if not is_quantized_tensor and ctx.is_input_fp8 and not ctx.bf16_backward: # return in FP8 dq, dk, dv, _, _ = combine_and_quantize( ctx.dqkv_layout, dq_, dk_, dv_, ctx.dQKV_quantizer @@ -1968,6 +1975,7 @@ def backward(ctx, d_out, *_args): None, None, # packed_qkv None, # packed_kv + None, ) @@ -2064,6 +2072,7 @@ def forward( score_mod_bprop_tensors: Optional[Dict[str, torch.Tensor]] = None, packed_qkv: Optional[torch.Tensor] = None, packed_kv: Optional[torch.Tensor] = None, + bf16_backward: bool = False, ) -> torch.Tensor: """fused attention fprop""" assert ( @@ -2280,6 +2289,7 @@ def forward( self.deterministic, softmax_offset, fp8_output, + bf16_backward, self.layer_number, self.return_max_logit, packed_qkv, diff --git a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py index 11d6500ab7..e6f50eb0da 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py @@ -32,6 +32,7 @@ Float8BlockScalingRecipeState, ) from transformer_engine.pytorch.tensor.storage.float8_tensor_storage import Float8TensorStorage +from transformer_engine.pytorch.tensor.storage.mxfp8_tensor_storage import MXFP8TensorStorage from transformer_engine.pytorch.module.base import TransformerEngineBaseModule from transformer_engine.pytorch.export import is_in_onnx_export_mode from transformer_engine.pytorch.constants import AttnMaskTypes, AttnTypes, dist_group_type, DType @@ -1347,6 +1348,7 @@ def forward( inference_params: Optional[InferenceParams] = None, pad_between_seqs: Optional[bool] = None, fp8_output: Optional[bool] = False, + bf16_backward: Optional[bool] = False, num_splits: Optional[int] = 1, score_mod: Optional[Callable] = None, score_mod_bprop: Optional[Callable] = None, @@ -1817,6 +1819,25 @@ def forward( qkv_format=qkv_format, inference_params=inference_params, ) + elif all( + isinstance(x, MXFP8TensorStorage) for x in [query_layer, key_layer, value_layer] + ): + # Pre-quantized MXFP8 q/k/v: the wrapper has no real storage, so run + # layout detection on the underlying rowwise data (mirrors the Float8 path). + ( + qkv_layout, + query_layer._rowwise_data, + key_layer._rowwise_data, + value_layer._rowwise_data, + q_format, + kv_format, + ) = dpa_utils.get_qkv_layout( + query_layer._rowwise_data, + key_layer._rowwise_data, + value_layer._rowwise_data, + qkv_format=qkv_format, + inference_params=inference_params, + ) else: ( qkv_layout, @@ -2190,6 +2211,7 @@ def forward( fp8_output=fp8_output, packed_qkv=qkv_layer, packed_kv=kv_layer, + bf16_backward=bf16_backward, ) return self.fused_attention( query_layer, @@ -2227,6 +2249,7 @@ def forward( score_mod_bprop_tensors=score_mod_bprop_tensors, packed_qkv=qkv_layer, packed_kv=kv_layer, + bf16_backward=bf16_backward, ) if use_unfused_attention: diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py index d99ee4c8e8..6eb3ce54f1 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py @@ -2786,10 +2786,37 @@ def mxfp8_quantize_fast_path(tensor_quantizer_pairs, src_format): """ if not tensor_quantizer_pairs: return [], src_format + + fp8_tensors = mxfp8_quantize_only(tensor_quantizer_pairs, src_format) + mxfp8_transpose_swizzle(fp8_tensors, src_format) + return fp8_tensors, "bhsd" + + +def mxfp8_quantize_only(tensor_quantizer_pairs, src_format): + """Phase 1 of mxfp8_quantize_fast_path: quantize only, no BHSD transpose or GEMM swizzle. + + Returns MXFP8Tensors with data and scale_invs reshaped to src_format layout. + Call mxfp8_transpose_swizzle to complete the BHSD permute + swizzle when ready + (e.g. after pre-quantized tensors from fused kernels are also available). + + Parameters + ---------- + tensor_quantizer_pairs : list of (torch.Tensor, MXFP8Quantizer) + Same contract as mxfp8_quantize_fast_path. + src_format : str + ``"bshd"`` or ``"sbhd"``. + + Returns + ------- + fp8_tensors : list of MXFP8Tensor + Data and scale_invs in src_format layout; NOT yet BHSD-permuted or swizzled. + """ + if not tensor_quantizer_pairs: + return [] assert src_format in ( "bshd", "sbhd", - ), f"mxfp8_quantize_fast_path only supports bshd/sbhd, got {src_format!r}." + ), f"mxfp8_quantize_only only supports bshd/sbhd, got {src_format!r}." _s_dim = {"bshd": 1, "sbhd": 0} _d_dim = {"bshd": 3, "sbhd": 3} @@ -2800,45 +2827,74 @@ def mxfp8_quantize_fast_path(tensor_quantizer_pairs, src_format): rs_shape[_d_dim[src_format]] //= MXFP8_BLOCK_SCALING_SIZE cs_shape = list(original_shape) cs_shape[_s_dim[src_format]] //= MXFP8_BLOCK_SCALING_SIZE - - # view tensor as 2D for quantization - # BSHD -> (B*S, H*D) - # SBHD -> (S, B*H*D) if src_format == "bshd": - tensor = tensor.view(*tensor.shape[:2], -1) + t2d = tensor.view(*tensor.shape[:2], -1) else: - tensor = tensor.view(tensor.shape[0], -1) - - # quantize + t2d = tensor.view(tensor.shape[0], -1) orig_optimize = quantizer.optimize_for_gemm quantizer.optimize_for_gemm = False - fp8_tensor = quantizer(tensor) + fp8_2d = quantizer(t2d) quantizer.optimize_for_gemm = orig_optimize - - # reshape rowwise/columnwise data to original shape - fp8_tensor._rowwise_data = ( - fp8_tensor._rowwise_data.view(original_shape) - if fp8_tensor._rowwise_data is not None - else None - ) - fp8_tensor._columnwise_data = ( - fp8_tensor._columnwise_data.view(original_shape) - if fp8_tensor._columnwise_data is not None - else None - ) - fp8_tensor._rowwise_scale_inv = ( - fp8_tensor._rowwise_scale_inv.view(rs_shape) - if fp8_tensor._rowwise_scale_inv is not None - else None - ) - fp8_tensor._columnwise_scale_inv = ( - fp8_tensor._columnwise_scale_inv.view(cs_shape) - if fp8_tensor._columnwise_scale_inv is not None - else None + # Re-wrap with the original 4D SBHD/BSHD shape so that shape[-1] equals the per-head + # dimension (matching Q's wrapper shape) and fused_attn_bwd produces 4D dkv that + # matches key/value's expected gradient shape in _KFQuantizeKVForAttn.backward. + fp8_t = MXFP8Tensor( + shape=original_shape, + dtype=tensor.dtype, + rowwise_data=( + fp8_2d._rowwise_data.view(original_shape) + if fp8_2d._rowwise_data is not None + else None + ), + rowwise_scale_inv=( + fp8_2d._rowwise_scale_inv.view(rs_shape) + if fp8_2d._rowwise_scale_inv is not None + else None + ), + columnwise_data=( + fp8_2d._columnwise_data.view(original_shape) + if fp8_2d._columnwise_data is not None + else None + ), + columnwise_scale_inv=( + fp8_2d._columnwise_scale_inv.view(cs_shape) + if fp8_2d._columnwise_scale_inv is not None + else None + ), + quantizer=quantizer, + requires_grad=False, + fp8_dtype=fp8_2d._fp8_dtype, + with_gemm_swizzled_scales=False, ) - fp8_tensors.append(fp8_tensor) + fp8_tensors.append(fp8_t) + return fp8_tensors + + +def mxfp8_transpose_swizzle(fp8_tensors, src_format): + """Phase 2 of mxfp8_quantize_fast_path: batched BHSD-transpose + GEMM-swizzle. + + For tensors whose data is already quantized (e.g. from a fused GEMM+quant kernel + or from mxfp8_quantize_only), permutes each tensor's scale_invs from src_format to + BHSD and applies the GEMM swizzle in-place. Complements mxfp8_quantize_only to + allow pre-quantized tensors (like a fused-kernel Q) to be processed in the same + batched operation as freshly quantized K/V. + + Parameters + ---------- + fp8_tensors : list of MXFP8Tensor + Tensors with _rowwise_scale_inv / _columnwise_scale_inv in src_format layout. + Modified in-place: scale_invs are replaced with BHSD-permuted, swizzled versions. + src_format : str + ``"bshd"`` or ``"sbhd"``. + """ + if not fp8_tensors: + return + + assert src_format in ( + "bshd", + "sbhd", + ), f"mxfp8_transpose_swizzle only supports bshd/sbhd, got {src_format!r}." - # ---- Pad + permute + swizzle scale_inv to BHSD ---- rs_list = [t._rowwise_scale_inv for t in fp8_tensors] cs_list = [t._columnwise_scale_inv for t in fp8_tensors] @@ -2872,50 +2928,25 @@ def _build_outputs(scale_list, alignment): buf = torch.empty(total, dtype=torch.uint8, device=device) return [buf[e[0] : e[0] + e[1]].view(e[2]) if e is not None else None for e in entries] - # allocate buffers with padding in mind rs_outs = _build_outputs(rs_list, 4) cs_outs = _build_outputs(cs_list, 128) - # permute scale_invs to BHSD; batched rs_permuted = tex.multi_tensor_transpose_to_bhsd( - rs_list, - original_format=src_format, - outputs=rs_outs, + rs_list, original_format=src_format, outputs=rs_outs ) cs_permuted = tex.multi_tensor_transpose_to_bhsd( - cs_list, - original_format=src_format, - outputs=cs_outs, + cs_list, original_format=src_format, outputs=cs_outs ) - # build output tensors - result = [] for t, rp, cp in zip(fp8_tensors, rs_permuted, cs_permuted): - rp = rp.view(-1, rp.shape[-1]) if rp is not None else None - cp = cp.view(-1, cp.shape[-1]) if cp is not None else None - result.append( - MXFP8Tensor( - shape=t.shape, - dtype=t.dtype, - rowwise_data=t._rowwise_data, - rowwise_scale_inv=rp, - columnwise_data=t._columnwise_data, - columnwise_scale_inv=cp, - quantizer=t._quantizer, - requires_grad=False, - fp8_dtype=t._fp8_dtype, - with_gemm_swizzled_scales=t._with_gemm_swizzled_scales, - ) - ) + t._rowwise_scale_inv = rp.view(-1, rp.shape[-1]) if rp is not None else None + t._columnwise_scale_inv = cp.view(-1, cp.shape[-1]) if cp is not None else None - # swizzle in place; batched - tex.multi_tensor_swizzle_scales_for_gemm_unchecked_(result, True, False) - tex.multi_tensor_swizzle_scales_for_gemm_unchecked_(result, False, True) - for t in result: + tex.multi_tensor_swizzle_scales_for_gemm_unchecked_(fp8_tensors, True, False) + tex.multi_tensor_swizzle_scales_for_gemm_unchecked_(fp8_tensors, False, True) + for t in fp8_tensors: t._with_gemm_swizzled_scales = True - return result, "bhsd" - def combine_and_quantize( qkv_layout, diff --git a/transformer_engine/pytorch/attention/fused_mla_q_uproj.py b/transformer_engine/pytorch/attention/fused_mla_q_uproj.py new file mode 100644 index 0000000000..cd62e16f3a --- /dev/null +++ b/transformer_engine/pytorch/attention/fused_mla_q_uproj.py @@ -0,0 +1,175 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Fused MLA Q up-projection + per-head RoPE + MXFP8 quantize.""" + +from __future__ import annotations +import functools +import os +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 + +from ..constants import MXFP8_BLOCK_SCALING_SIZE +from ..quantized_tensor import QuantizedTensor +from ..tensor.mxfp8_tensor import MXFP8Quantizer, MXFP8Tensor +from ..utils import get_device_compute_capability + +_CUDNN_FRONTEND_MIN_VERSION = "1.27.0" + + +def _cudnn_frontend_version_supported() -> bool: + """Check that the installed nvidia-cudnn-frontend meets the minimum version.""" + try: + return PkgVersion(get_pkg_version("nvidia-cudnn-frontend")) >= PkgVersion( + _CUDNN_FRONTEND_MIN_VERSION + ) + except PackageNotFoundError: + return False + + +class FusedMLAQUpProjRopeQuant: + """Wrapper for the cuDNN fused MLA Q up-proj + per-head RoPE + MXFP8 quantize kernel. + + - If w is already a QuantizedTensor (primary FP8 parameter in MXFP8BlockScaling recipe), + this performs an MXFP8 GEMM within the fusion (and quantizes the input if necessary) + - Otherwise (plain BF16 weight), x and w are passed as-is to the BF16 kernel variant. + """ + + @classmethod + @functools.lru_cache(maxsize=None) + def _kernel(cls): + # Import directly from the subpackage to avoid depending on cudnn/__init__.py + # lazy-import registration (which would require overlaying cudnn/__init__.py and + # could revert atomicrmw fixes present in the container's version). + try: + from cudnn import gemm_proj_rope_mxfp8_wrapper_sm100 + + return gemm_proj_rope_mxfp8_wrapper_sm100 + except ImportError: + return None + + @classmethod + @functools.lru_cache(maxsize=None) + def is_supported(cls) -> bool: + """Whether the cuDNN FE fused gemm rope quant wrapper is available""" + if int(os.environ.get("NVTE_FUSED_MLA_Q_UPROJ", "1")) <= 0: + return False + if not _cudnn_frontend_version_supported(): + return False + if get_device_compute_capability()[0] < 10: + return False + if cls._kernel() is None: + return False + return True + + @classmethod + def run( + cls, + x: torch.Tensor, + w, # MXFP8Tensor (primary FP8 param) or bf16 torch.Tensor + cos: torch.Tensor, + sin: torch.Tensor, + s: int, + b: int, + ) -> "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. + """ + + from cuda.bindings import driver as cuda + + stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) + wrapper = cls._kernel() + + if isinstance(w, QuantizedTensor): + assert isinstance(w, MXFP8Tensor), ( + "FusedMLAQUpProjRopeQuant expects an MXFP8Tensor weight (MXFP8BlockScaling" + 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 + ) + 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 + + # Primary FP8 parameter: already quantized; use its rowwise FP8 codes + E8M0 scales. + w.update_usage(rowwise_usage=True, columnwise_usage=None) + w_code = w._rowwise_data.view(torch.float8_e4m3fn) # [N, K] + w_scale = w._rowwise_scale_inv # [N, K//32] uint8 + + out = wrapper( + x_code, + w_code, + cos, + sin, + x_scale=x_scale, + w_scale=w_scale, + w_out_in=True, + stream=stream, + ) + + # Drop rowwise data now. + # Only columnwise x is needed for the FP8 wgrad in backward. + x_mxfp8.update_usage(rowwise_usage=False, columnwise_usage=True) + x_saved = x_mxfp8 + else: + # ---- 16-bit projection: bf16 GEMM inputs -> bf16in (the projection stays bf16) ---- + out = wrapper(x, w, cos, sin, w_out_in=True, stream=stream) + x_saved = x + + nh = out["out_fp8_row"].shape[1] + d = out["out_fp8_row"].shape[2] + query = cls.wrap_mxfp8( + out["out_fp8_row"], + out["out_scales_row"], + out["out_fp8_col"], + out["out_scales_col"], + s, + b, + nh, + d, + ) + # 2nd return is the activation to save for wgrad: MXFP8 (fp8 path) or bf16 (16-bit path). + return query, x_saved + + @classmethod + def wrap_mxfp8( + cls, + fp8_row: torch.Tensor, + scales_row: torch.Tensor, + fp8_col: torch.Tensor, + scales_col: torch.Tensor, + s: int, + b: int, + nh: int, + d: int, + ) -> MXFP8Tensor: + """Wrap raw data and scale tensors into an MXFP8Tensor""" + + blk = MXFP8_BLOCK_SCALING_SIZE + # Both rowwise and columnwise Q are required: + # - Forward QK^T uses rowwise + # - cuDNN backward (fused_attn_fp8_bwd_impl) requires columnwise for dK gradient + quantizer = MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3, rowwise=True, columnwise=True) + return MXFP8Tensor( + shape=(s, b, nh, d), + dtype=torch.bfloat16, + rowwise_data=fp8_row.view(s, b, nh, d), + rowwise_scale_inv=scales_row.view(s, b, nh, d // blk), + columnwise_data=fp8_col.view(s, b, nh, d), + columnwise_scale_inv=scales_col.view(s // blk, b, nh, d), + quantizer=quantizer, + requires_grad=False, + fp8_dtype=tex.DType.kFloat8E4M3, + with_gemm_swizzled_scales=False, + )