From c19f0df7b64e6d4f339195a6ac3ac2ca9911a6c4 Mon Sep 17 00:00:00 2001 From: Chase Block Date: Tue, 7 Jul 2026 16:16:46 -0700 Subject: [PATCH 01/14] Add support for fused Q Up-Proj GEMM/RoPE/Quant. This commit add support for fusing the GEMM in the Q Up Proj step of DeepseekV3 training with the following RoPE and MXFP8 quantization operations. This uses a custom kernel from cudnn_frontend, and supports both 16-bit projection and mxfp8 projection. Signed-off-by: Chase Block --- transformer_engine/pytorch/__init__.py | 1 + .../pytorch/attention/__init__.py | 2 + .../dot_product_attention/backends.py | 14 +- .../dot_product_attention.py | 23 +++ .../attention/dot_product_attention/utils.py | 154 ++++++++++++++++++ .../pytorch/attention/fused_mla_q_uproj.py | 142 ++++++++++++++++ 6 files changed, 334 insertions(+), 2 deletions(-) create mode 100644 transformer_engine/pytorch/attention/fused_mla_q_uproj.py diff --git a/transformer_engine/pytorch/__init__.py b/transformer_engine/pytorch/__init__.py index 06db28ee27..0977b939df 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..55bd27aaca 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 4cc4cab1b8..116e1d7d3c 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 @@ -1115,6 +1116,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, @@ -1585,6 +1587,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, @@ -1958,6 +1979,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, @@ -1995,6 +2017,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 9ee6ad0101..48959ba525 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py @@ -2923,6 +2923,160 @@ def _build_outputs(scale_list, alignment): return result, "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_only only supports bshd/sbhd, got {src_format!r}." + ) + _s_dim = {"bshd": 1, "sbhd": 0} + _d_dim = {"bshd": 3, "sbhd": 3} + + fp8_tensors = [] + for tensor, quantizer in tensor_quantizer_pairs: + original_shape = tensor.shape + rs_shape = list(original_shape) + 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 + if src_format == "bshd": + t2d = tensor.view(*tensor.shape[:2], -1) + else: + t2d = tensor.view(tensor.shape[0], -1) + orig_optimize = quantizer.optimize_for_gemm + quantizer.optimize_for_gemm = False + fp8_2d = quantizer(t2d) + quantizer.optimize_for_gemm = orig_optimize + # Re-wrap with the original 4D SBHD 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_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}." + ) + + rs_list = [t._rowwise_scale_inv for t in fp8_tensors] + cs_list = [t._columnwise_scale_inv for t in fp8_tensors] + + def _align_up(x, a): + return ((x + a - 1) // a) * a + + def _bhsd_shape(src_4d, d_pad): + if src_format == "sbhd": + S, B, H, _ = src_4d.shape + else: + B, S, H, _ = src_4d.shape + return (B, H, S, d_pad) + + def _build_outputs(scale_list, alignment): + entries = [] + total = 0 + for s in scale_list: + if s is None: + entries.append(None) + continue + d_pad = _align_up(s.shape[-1], alignment) + shape = _bhsd_shape(s, d_pad) + numel = 1 + for dim in shape: + numel *= dim + entries.append((total, numel, shape)) + total += numel + if total == 0: + return [None] * len(scale_list) + device = next(s for s in scale_list if s is not None).device + 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] + + rs_outs = _build_outputs(rs_list, 4) + cs_outs = _build_outputs(cs_list, 128) + + rs_permuted = tex.multi_tensor_transpose_to_bhsd( + 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 + ) + + for t, rp, cp in zip(fp8_tensors, rs_permuted, cs_permuted): + 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 + + 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 + + def combine_and_quantize( qkv_layout, q, 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..6623281d07 --- /dev/null +++ b/transformer_engine/pytorch/attention/fused_mla_q_uproj.py @@ -0,0 +1,142 @@ +# 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 + +import torch +import transformer_engine_torch as tex + +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 + + +class FusedMLAQUpProjRopeQuant: + """Wrapper for the cuDNN fused MLA Q up-proj + per-head RoPE + MXFP8 quantize kernel (v4). + + - 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). + from cudnn import gemm_proj_rope_mxfp8_wrapper_sm100 + return gemm_proj_rope_mxfp8_wrapper_sm100 + + @classmethod + @functools.lru_cache(maxsize=None) + def is_supported(cls) -> bool: + if int(os.environ.get("NVTE_FUSED_MLA_Q_UPROJ", "1")) <= 0: + return False + if get_device_compute_capability()[0] < 10: + return False + try: + cls._kernel() + except ImportError: + return False + return True + + @classmethod + def warmup(cls, *, num_heads: int, q_lora_rank: int, q_head_dim: int, qk_pos_emb_head_dim: int, tokens: int) -> None: + """Pre-compile the fused kernel during model init using dummy FP8 tensors.""" + if not cls.is_supported(): + return + head_dim = q_head_dim + qk_pos_emb_head_dim + dev = torch.cuda.current_device() + x_code = torch.zeros(tokens, q_lora_rank, dtype=torch.float8_e4m3fn, device=dev) + x_scale = torch.zeros(tokens, q_lora_rank // MXFP8_BLOCK_SCALING_SIZE, dtype=torch.uint8, device=dev) + w_code = torch.zeros(num_heads * head_dim, q_lora_rank, dtype=torch.float8_e4m3fn, device=dev) + w_scale = torch.zeros(num_heads * head_dim, q_lora_rank // MXFP8_BLOCK_SCALING_SIZE, dtype=torch.uint8, device=dev) + cos = torch.zeros(tokens, qk_pos_emb_head_dim, dtype=torch.bfloat16, device=dev) + sin = torch.zeros(tokens, qk_pos_emb_head_dim, dtype=torch.bfloat16, device=dev) + cls._kernel()(x_code, w_code, cos, sin, x_scale=x_scale, w_scale=w_scale, w_out_in=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): + # ---- 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: + 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, + ) From cd730aafdaef0a1e49ae67bc72edff39fefcdbea Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 3 Aug 2026 15:42:20 +0000 Subject: [PATCH 02/14] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- .../dot_product_attention/backends.py | 2 +- .../attention/dot_product_attention/utils.py | 14 +-- .../pytorch/attention/fused_mla_q_uproj.py | 86 ++++++++++++++----- 3 files changed, 72 insertions(+), 30 deletions(-) diff --git a/transformer_engine/pytorch/attention/dot_product_attention/backends.py b/transformer_engine/pytorch/attention/dot_product_attention/backends.py index 55bd27aaca..aeac501ff8 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/backends.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/backends.py @@ -1430,7 +1430,7 @@ def forward( 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 + 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, diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py index 48959ba525..17002b370e 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py @@ -2944,9 +2944,10 @@ def mxfp8_quantize_only(tensor_quantizer_pairs, src_format): """ if not tensor_quantizer_pairs: return [] - assert src_format in ("bshd", "sbhd"), ( - f"mxfp8_quantize_only only supports bshd/sbhd, got {src_format!r}." - ) + assert src_format in ( + "bshd", + "sbhd", + ), f"mxfp8_quantize_only only supports bshd/sbhd, got {src_format!r}." _s_dim = {"bshd": 1, "sbhd": 0} _d_dim = {"bshd": 3, "sbhd": 3} @@ -3020,9 +3021,10 @@ def mxfp8_transpose_swizzle(fp8_tensors, src_format): if not fp8_tensors: return - assert src_format in ("bshd", "sbhd"), ( - f"mxfp8_transpose_swizzle only supports bshd/sbhd, got {src_format!r}." - ) + assert src_format in ( + "bshd", + "sbhd", + ), f"mxfp8_transpose_swizzle only supports bshd/sbhd, got {src_format!r}." rs_list = [t._rowwise_scale_inv for t in fp8_tensors] cs_list = [t._columnwise_scale_inv for t in fp8_tensors] diff --git a/transformer_engine/pytorch/attention/fused_mla_q_uproj.py b/transformer_engine/pytorch/attention/fused_mla_q_uproj.py index 6623281d07..6d187886ef 100644 --- a/transformer_engine/pytorch/attention/fused_mla_q_uproj.py +++ b/transformer_engine/pytorch/attention/fused_mla_q_uproj.py @@ -20,9 +20,9 @@ class FusedMLAQUpProjRopeQuant: """Wrapper for the cuDNN fused MLA Q up-proj + per-head RoPE + MXFP8 quantize kernel (v4). - - 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. + - 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 @@ -32,6 +32,7 @@ def _kernel(cls): # lazy-import registration (which would require overlaying cudnn/__init__.py and # could revert atomicrmw fixes present in the container's version). from cudnn import gemm_proj_rope_mxfp8_wrapper_sm100 + return gemm_proj_rope_mxfp8_wrapper_sm100 @classmethod @@ -48,16 +49,33 @@ def is_supported(cls) -> bool: return True @classmethod - def warmup(cls, *, num_heads: int, q_lora_rank: int, q_head_dim: int, qk_pos_emb_head_dim: int, tokens: int) -> None: + def warmup( + cls, + *, + num_heads: int, + q_lora_rank: int, + q_head_dim: int, + qk_pos_emb_head_dim: int, + tokens: int, + ) -> None: """Pre-compile the fused kernel during model init using dummy FP8 tensors.""" if not cls.is_supported(): return head_dim = q_head_dim + qk_pos_emb_head_dim dev = torch.cuda.current_device() - x_code = torch.zeros(tokens, q_lora_rank, dtype=torch.float8_e4m3fn, device=dev) - x_scale = torch.zeros(tokens, q_lora_rank // MXFP8_BLOCK_SCALING_SIZE, dtype=torch.uint8, device=dev) - w_code = torch.zeros(num_heads * head_dim, q_lora_rank, dtype=torch.float8_e4m3fn, device=dev) - w_scale = torch.zeros(num_heads * head_dim, q_lora_rank // MXFP8_BLOCK_SCALING_SIZE, dtype=torch.uint8, device=dev) + x_code = torch.zeros(tokens, q_lora_rank, dtype=torch.float8_e4m3fn, device=dev) + x_scale = torch.zeros( + tokens, q_lora_rank // MXFP8_BLOCK_SCALING_SIZE, dtype=torch.uint8, device=dev + ) + w_code = torch.zeros( + num_heads * head_dim, q_lora_rank, dtype=torch.float8_e4m3fn, device=dev + ) + w_scale = torch.zeros( + num_heads * head_dim, + q_lora_rank // MXFP8_BLOCK_SCALING_SIZE, + dtype=torch.uint8, + device=dev, + ) cos = torch.zeros(tokens, qk_pos_emb_head_dim, dtype=torch.bfloat16, device=dev) sin = torch.zeros(tokens, qk_pos_emb_head_dim, dtype=torch.bfloat16, device=dev) cls._kernel()(x_code, w_code, cos, sin, x_scale=x_scale, w_scale=w_scale, w_out_in=True) @@ -66,7 +84,7 @@ def warmup(cls, *, num_heads: int, q_lora_rank: int, q_head_dim: int, qk_pos_emb def run( cls, x: torch.Tensor, - w, # MXFP8Tensor (primary FP8 param) or bf16 torch.Tensor + w, # MXFP8Tensor (primary FP8 param) or bf16 torch.Tensor cos: torch.Tensor, sin: torch.Tensor, s: int, @@ -78,6 +96,7 @@ def run( """ from cuda.bindings import driver as cuda + stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) wrapper = cls._kernel() @@ -85,17 +104,28 @@ def run( # ---- 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_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 + 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) + 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. @@ -107,11 +137,16 @@ def run( x_saved = x nh = out["out_fp8_row"].shape[1] - d = out["out_fp8_row"].shape[2] + 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, + 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 @@ -119,9 +154,14 @@ def run( @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, + 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: blk = MXFP8_BLOCK_SCALING_SIZE # Both rowwise and columnwise Q are required: From d0e3915621770f4a0a9bb72b3effb31124723896 Mon Sep 17 00:00:00 2001 From: Chase Block Date: Wed, 5 Aug 2026 08:21:11 -0700 Subject: [PATCH 03/14] Remove unused function from fused mla q uproj, add error handling Signed-off-by: Chase Block --- .../pytorch/attention/fused_mla_q_uproj.py | 49 +++++-------------- 1 file changed, 11 insertions(+), 38 deletions(-) diff --git a/transformer_engine/pytorch/attention/fused_mla_q_uproj.py b/transformer_engine/pytorch/attention/fused_mla_q_uproj.py index 6d187886ef..5787e10b43 100644 --- a/transformer_engine/pytorch/attention/fused_mla_q_uproj.py +++ b/transformer_engine/pytorch/attention/fused_mla_q_uproj.py @@ -18,7 +18,7 @@ class FusedMLAQUpProjRopeQuant: - """Wrapper for the cuDNN fused MLA Q up-proj + per-head RoPE + MXFP8 quantize kernel (v4). + """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) @@ -31,9 +31,12 @@ 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). - from cudnn import gemm_proj_rope_mxfp8_wrapper_sm100 + try: + from cudnn import gemm_proj_rope_mxfp8_wrapper_sm100 - return gemm_proj_rope_mxfp8_wrapper_sm100 + return gemm_proj_rope_mxfp8_wrapper_sm100 + except ImportError: + return None @classmethod @functools.lru_cache(maxsize=None) @@ -42,44 +45,10 @@ def is_supported(cls) -> bool: return False if get_device_compute_capability()[0] < 10: return False - try: - cls._kernel() - except ImportError: + if cls._kernel() is None: return False return True - @classmethod - def warmup( - cls, - *, - num_heads: int, - q_lora_rank: int, - q_head_dim: int, - qk_pos_emb_head_dim: int, - tokens: int, - ) -> None: - """Pre-compile the fused kernel during model init using dummy FP8 tensors.""" - if not cls.is_supported(): - return - head_dim = q_head_dim + qk_pos_emb_head_dim - dev = torch.cuda.current_device() - x_code = torch.zeros(tokens, q_lora_rank, dtype=torch.float8_e4m3fn, device=dev) - x_scale = torch.zeros( - tokens, q_lora_rank // MXFP8_BLOCK_SCALING_SIZE, dtype=torch.uint8, device=dev - ) - w_code = torch.zeros( - num_heads * head_dim, q_lora_rank, dtype=torch.float8_e4m3fn, device=dev - ) - w_scale = torch.zeros( - num_heads * head_dim, - q_lora_rank // MXFP8_BLOCK_SCALING_SIZE, - dtype=torch.uint8, - device=dev, - ) - cos = torch.zeros(tokens, qk_pos_emb_head_dim, dtype=torch.bfloat16, device=dev) - sin = torch.zeros(tokens, qk_pos_emb_head_dim, dtype=torch.bfloat16, device=dev) - cls._kernel()(x_code, w_code, cos, sin, x_scale=x_scale, w_scale=w_scale, w_out_in=True) - @classmethod def run( cls, @@ -101,6 +70,10 @@ def run( wrapper = cls._kernel() if isinstance(w, QuantizedTensor): + assert isinstance(w, MXFP8Tensor), ( + f"FusedMLAQUpProjRopeQuant expects an MXFP8Tensor weight (MXFP8BlockScaling recipe), " + f"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). From 9854af357b3359175a2c30dac3213dfa9c5d4e9a Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Wed, 5 Aug 2026 15:22:38 +0000 Subject: [PATCH 04/14] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- transformer_engine/pytorch/attention/fused_mla_q_uproj.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/transformer_engine/pytorch/attention/fused_mla_q_uproj.py b/transformer_engine/pytorch/attention/fused_mla_q_uproj.py index 5787e10b43..f02c97c71a 100644 --- a/transformer_engine/pytorch/attention/fused_mla_q_uproj.py +++ b/transformer_engine/pytorch/attention/fused_mla_q_uproj.py @@ -71,8 +71,9 @@ def run( if isinstance(w, QuantizedTensor): assert isinstance(w, MXFP8Tensor), ( - f"FusedMLAQUpProjRopeQuant expects an MXFP8Tensor weight (MXFP8BlockScaling recipe), " - f"got {type(w).__name__}. Use the unfused path for other quantization recipes." + "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 From 5f90db78cdf1404c1ba6a4140b5b2d1528e04e91 Mon Sep 17 00:00:00 2001 From: Chase Block Date: Wed, 5 Aug 2026 08:22:48 -0700 Subject: [PATCH 05/14] Adjust comment SBHD/BSHD Signed-off-by: Chase Block --- .../pytorch/attention/dot_product_attention/utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py index 17002b370e..ed20dcd0b8 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py @@ -2966,7 +2966,7 @@ def mxfp8_quantize_only(tensor_quantizer_pairs, src_format): quantizer.optimize_for_gemm = False fp8_2d = quantizer(t2d) quantizer.optimize_for_gemm = orig_optimize - # Re-wrap with the original 4D SBHD shape so that shape[-1] equals the per-head + # 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( From e2ebeacf25da0c0c5339aac9ec7dc41ca766f665 Mon Sep 17 00:00:00 2001 From: Chase Block Date: Wed, 5 Aug 2026 11:43:30 -0700 Subject: [PATCH 06/14] Add cudnn_frontend version check; eliminate code dup in mxfp8_quantize_fast_path. Signed-off-by: Chase Block --- .../attention/dot_product_attention/utils.py | 131 +----------------- .../pytorch/attention/fused_mla_q_uproj.py | 16 +++ 2 files changed, 19 insertions(+), 128 deletions(-) diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py index bdf5757a6d..6eb3ce54f1 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py @@ -2786,135 +2786,10 @@ def mxfp8_quantize_fast_path(tensor_quantizer_pairs, src_format): """ if not tensor_quantizer_pairs: return [], src_format - assert src_format in ( - "bshd", - "sbhd", - ), f"mxfp8_quantize_fast_path only supports bshd/sbhd, got {src_format!r}." - _s_dim = {"bshd": 1, "sbhd": 0} - _d_dim = {"bshd": 3, "sbhd": 3} - - fp8_tensors = [] - for tensor, quantizer in tensor_quantizer_pairs: - original_shape = tensor.shape - rs_shape = list(original_shape) - 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) - else: - tensor = tensor.view(tensor.shape[0], -1) - - # quantize - orig_optimize = quantizer.optimize_for_gemm - quantizer.optimize_for_gemm = False - fp8_tensor = quantizer(tensor) - 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 - ) - fp8_tensors.append(fp8_tensor) - - # ---- 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] - - def _align_up(x, a): - return ((x + a - 1) // a) * a - - def _bhsd_shape(src_4d, d_pad): - if src_format == "sbhd": - S, B, H, _ = src_4d.shape - else: - B, S, H, _ = src_4d.shape - return (B, H, S, d_pad) - - def _build_outputs(scale_list, alignment): - entries = [] - total = 0 - for s in scale_list: - if s is None: - entries.append(None) - continue - d_pad = _align_up(s.shape[-1], alignment) - shape = _bhsd_shape(s, d_pad) - numel = 1 - for dim in shape: - numel *= dim - entries.append((total, numel, shape)) - total += numel - if total == 0: - return [None] * len(scale_list) - device = next(s for s in scale_list if s is not None).device - 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, - ) - cs_permuted = tex.multi_tensor_transpose_to_bhsd( - 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, - ) - ) - - # 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: - t._with_gemm_swizzled_scales = True - return result, "bhsd" + 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): diff --git a/transformer_engine/pytorch/attention/fused_mla_q_uproj.py b/transformer_engine/pytorch/attention/fused_mla_q_uproj.py index f02c97c71a..0ecf6c3bbf 100644 --- a/transformer_engine/pytorch/attention/fused_mla_q_uproj.py +++ b/transformer_engine/pytorch/attention/fused_mla_q_uproj.py @@ -7,15 +7,29 @@ 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. @@ -43,6 +57,8 @@ def _kernel(cls): def is_supported(cls) -> bool: 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: From 417d23364584e856dd877fd024c634faa7a78554 Mon Sep 17 00:00:00 2001 From: Chase Block Date: Wed, 5 Aug 2026 12:08:33 -0700 Subject: [PATCH 07/14] Add missing docstrings Signed-off-by: Chase Block --- transformer_engine/pytorch/attention/fused_mla_q_uproj.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/transformer_engine/pytorch/attention/fused_mla_q_uproj.py b/transformer_engine/pytorch/attention/fused_mla_q_uproj.py index 0ecf6c3bbf..cd62e16f3a 100644 --- a/transformer_engine/pytorch/attention/fused_mla_q_uproj.py +++ b/transformer_engine/pytorch/attention/fused_mla_q_uproj.py @@ -55,6 +55,7 @@ def _kernel(cls): @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(): @@ -153,6 +154,8 @@ def wrap_mxfp8( 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 From dc89f019774982906bcada4714866efcadd94e1e Mon Sep 17 00:00:00 2001 From: Chase Block Date: Wed, 5 Aug 2026 12:49:19 -0700 Subject: [PATCH 08/14] Add unit test for q up proj fusion Signed-off-by: Chase Block --- .../attention/test_fused_mla_q_uproj.py | 197 ++++++++++++++++++ 1 file changed, 197 insertions(+) create mode 100644 tests/pytorch/attention/test_fused_mla_q_uproj.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..5950467307 --- /dev/null +++ b/tests/pytorch/attention/test_fused_mla_q_uproj.py @@ -0,0 +1,197 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Unit tests for FusedMLAQUpProjRopeQuant. + +Validates the fused MXFP8 Q up-proj + per-head RoPE + dual-direction MXFP8 quantize kernel +against an unfused reference: MXFP8 quantize inputs -> bf16 GEMM -> RoPE -> MXFP8 quantize. + +DSv3 671B MLA dimensions throughout; tokens must be a multiple of TILE_M (128). + +Run: + pytest tests/pytorch/attention/test_fused_mla_q_uproj.py -v +""" + +import pytest +import torch + +import transformer_engine.pytorch as te +import transformer_engine_torch as tex +from transformer_engine.pytorch.attention import FusedMLAQUpProjRopeQuant +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 # K dimension of the up-proj GEMM +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 _build_mxfp8_weight(proj_dim: int, k: int, device: torch.device) -> MXFP8Tensor: + """Quantize a random bf16 weight to MXFP8Tensor (rowwise only, matching the primary param).""" + w_bf16 = torch.randn(proj_dim, k, dtype=torch.bfloat16, device=device) + quantizer = MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3, rowwise=True, columnwise=False) + return quantizer(w_bf16) + + +def _dequantize_fused_output(query: MXFP8Tensor, s: int, b: int) -> torch.Tensor: + """Dequantize the rowwise fused output to bf16 [s, b, nh, head_dim]. + + query._rowwise_data is a plain float8_e4m3fn tensor — TE's C++ dequantize + kernel requires 2D layout, so we reshape before calling MXFP8Tensor.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. + + We compute in bf16 (dequantizing x and w) rather than re-quantizing the output, + so we compare against the fused kernel's dequantized output directly. + The fused kernel's MXFP8 GEMM + output quantize will naturally introduce some error. + """ + 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) + + # Per-head RoPE on the trailing HEAD_DIM_ROPE features, interleaved convention. + 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 = q_rope[..., 0::2] + x2 = q_rope[..., 1::2] + x_left = x1 * cos_[..., :half] - x2 * sin_[..., :half] + x_right = x2 * cos_[..., half:] + x1 * sin_[..., half:] + q_rope_out = torch.cat([x_left, x_right], dim=-1) + return torch.cat([q_nope, q_rope_out], dim=-1) + + +@pytest.mark.skipif(not fused_supported, reason=reason_not_supported) +@pytest.mark.parametrize( + "s, b", + [ + (128, 1), # minimum tile (TILE_M=128) + (256, 1), + (128, 2), + ], + ids=["s128_b1", "s256_b1", "s128_b2"], +) +def test_fused_mla_q_uproj_output_shapes(s: int, b: int) -> None: + """Fused kernel output tensors must have the shapes wrap_mxfp8 promises.""" + device = torch.device("cuda") + tokens = s * b + torch.manual_seed(SEED) + + x = torch.randn(tokens, Q_LORA_RANK, dtype=torch.bfloat16, device=device) + w = _build_mxfp8_weight(PROJ_DIM, Q_LORA_RANK, device) + cos, sin = _build_rope_tables(tokens, device) + + query, x_saved = FusedMLAQUpProjRopeQuant.run(x, w, cos, sin, s, b) + + assert isinstance(query, MXFP8Tensor), f"Expected MXFP8Tensor, got {type(query)}" + assert query.shape == (s, b, NUM_HEADS, HEAD_DIM), ( + f"Expected query shape {(s, b, NUM_HEADS, HEAD_DIM)}, got {query.shape}" + ) + assert query._rowwise_data is not None, "Missing rowwise data" + assert query._columnwise_data is not None, "Missing columnwise data" + + blk = 32 + assert query._rowwise_scale_inv.shape == (s, b, NUM_HEADS, HEAD_DIM // blk) + assert query._columnwise_scale_inv.shape == (s // blk, b, NUM_HEADS, HEAD_DIM) + + +@pytest.mark.skipif(not fused_supported, reason=reason_not_supported) +@pytest.mark.parametrize("tokens", [128, 256], ids=["tokens128", "tokens256"]) +def test_fused_mla_q_uproj_numerics(tokens: int) -> None: + """Fused kernel output must be close to the unfused MXFP8-precision reference. + + Tolerance reflects the double quantization noise (activation + output quantize). + """ + 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) + w = _build_mxfp8_weight(PROJ_DIM, Q_LORA_RANK, device) + cos, sin = _build_rope_tables(tokens, device) + + query, _ = FusedMLAQUpProjRopeQuant.run(x, w, cos, sin, s, b) + + fused_dq = _dequantize_fused_output(query, s, b) + ref_dq = _reference_q_uproj(x, w, cos, sin, s, b) + + # FP8 GEMM + output quantize introduce ~10-20% relative error. + torch.testing.assert_close(fused_dq, ref_dq, atol=0.5, rtol=0.1) + + +@pytest.mark.skipif(not fused_supported, reason=reason_not_supported) +def test_fused_mla_q_uproj_x_saved_is_mxfp8() -> None: + """In the MXFP8 weight path, x_saved must be an MXFP8Tensor with columnwise data.""" + s, b = 128, 1 + device = torch.device("cuda") + torch.manual_seed(SEED) + + x = torch.randn(s * b, Q_LORA_RANK, dtype=torch.bfloat16, device=device) + w = _build_mxfp8_weight(PROJ_DIM, Q_LORA_RANK, device) + cos, sin = _build_rope_tables(s * b, device) + + _, x_saved = FusedMLAQUpProjRopeQuant.run(x, w, cos, sin, s, b) + + assert isinstance(x_saved, MXFP8Tensor), ( + f"x_saved should be MXFP8Tensor for MXFP8 weight path, got {type(x_saved)}" + ) + 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" + + +def _build_rope_tables( + tokens: int, device: torch.device +) -> tuple[torch.Tensor, torch.Tensor]: + """Plain cos/sin tables for HEAD_DIM_ROPE, interleaved [tokens, rope_dim] bf16.""" + inv_freq = 1.0 / ( + 10000 ** (torch.arange(0, HEAD_DIM_ROPE, 2, dtype=torch.float32, device=device) + / HEAD_DIM_ROPE) + ) + t = torch.arange(tokens, dtype=torch.float32, device=device) + freqs = torch.outer(t, inv_freq) + freqs = torch.cat([freqs, freqs], dim=-1) + return freqs.cos().to(torch.bfloat16), freqs.sin().to(torch.bfloat16) From 24e0b0316d754458ed357ae5836da3dc5784ab75 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Wed, 5 Aug 2026 19:50:32 +0000 Subject: [PATCH 09/14] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- .../attention/test_fused_mla_q_uproj.py | 47 +++++++++++-------- 1 file changed, 27 insertions(+), 20 deletions(-) diff --git a/tests/pytorch/attention/test_fused_mla_q_uproj.py b/tests/pytorch/attention/test_fused_mla_q_uproj.py index 5950467307..349ce73bae 100644 --- a/tests/pytorch/attention/test_fused_mla_q_uproj.py +++ b/tests/pytorch/attention/test_fused_mla_q_uproj.py @@ -26,15 +26,21 @@ HEAD_DIM_NOPE = 128 HEAD_DIM_ROPE = 64 HEAD_DIM = HEAD_DIM_NOPE + HEAD_DIM_ROPE # 192 -Q_LORA_RANK = 1536 # K dimension of the up-proj GEMM -PROJ_DIM = NUM_HEADS * HEAD_DIM # 24576 +Q_LORA_RANK = 1536 # K dimension of the up-proj GEMM +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)") + (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)" + ), + ) ) @@ -98,7 +104,7 @@ def _reference_q_uproj( half = HEAD_DIM_ROPE // 2 x1 = q_rope[..., 0::2] x2 = q_rope[..., 1::2] - x_left = x1 * cos_[..., :half] - x2 * sin_[..., :half] + x_left = x1 * cos_[..., :half] - x2 * sin_[..., :half] x_right = x2 * cos_[..., half:] + x1 * sin_[..., half:] q_rope_out = torch.cat([x_left, x_right], dim=-1) return torch.cat([q_nope, q_rope_out], dim=-1) @@ -108,7 +114,7 @@ def _reference_q_uproj( @pytest.mark.parametrize( "s, b", [ - (128, 1), # minimum tile (TILE_M=128) + (128, 1), # minimum tile (TILE_M=128) (256, 1), (128, 2), ], @@ -127,14 +133,17 @@ def test_fused_mla_q_uproj_output_shapes(s: int, b: int) -> None: query, x_saved = FusedMLAQUpProjRopeQuant.run(x, w, cos, sin, s, b) assert isinstance(query, MXFP8Tensor), f"Expected MXFP8Tensor, got {type(query)}" - assert query.shape == (s, b, NUM_HEADS, HEAD_DIM), ( - f"Expected query shape {(s, b, NUM_HEADS, HEAD_DIM)}, got {query.shape}" - ) - assert query._rowwise_data is not None, "Missing rowwise data" + assert query.shape == ( + s, + b, + NUM_HEADS, + HEAD_DIM, + ), f"Expected query shape {(s, b, NUM_HEADS, HEAD_DIM)}, got {query.shape}" + assert query._rowwise_data is not None, "Missing rowwise data" assert query._columnwise_data is not None, "Missing columnwise data" blk = 32 - assert query._rowwise_scale_inv.shape == (s, b, NUM_HEADS, HEAD_DIM // blk) + assert query._rowwise_scale_inv.shape == (s, b, NUM_HEADS, HEAD_DIM // blk) assert query._columnwise_scale_inv.shape == (s // blk, b, NUM_HEADS, HEAD_DIM) @@ -176,20 +185,18 @@ def test_fused_mla_q_uproj_x_saved_is_mxfp8() -> None: _, x_saved = FusedMLAQUpProjRopeQuant.run(x, w, cos, sin, s, b) - assert isinstance(x_saved, MXFP8Tensor), ( - f"x_saved should be MXFP8Tensor for MXFP8 weight path, got {type(x_saved)}" - ) + assert isinstance( + x_saved, MXFP8Tensor + ), f"x_saved should be MXFP8Tensor for MXFP8 weight path, got {type(x_saved)}" 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" -def _build_rope_tables( - tokens: int, device: torch.device -) -> tuple[torch.Tensor, torch.Tensor]: +def _build_rope_tables(tokens: int, device: torch.device) -> tuple[torch.Tensor, torch.Tensor]: """Plain cos/sin tables for HEAD_DIM_ROPE, interleaved [tokens, rope_dim] bf16.""" inv_freq = 1.0 / ( - 10000 ** (torch.arange(0, HEAD_DIM_ROPE, 2, dtype=torch.float32, device=device) - / HEAD_DIM_ROPE) + 10000 + ** (torch.arange(0, HEAD_DIM_ROPE, 2, dtype=torch.float32, device=device) / HEAD_DIM_ROPE) ) t = torch.arange(tokens, dtype=torch.float32, device=device) freqs = torch.outer(t, inv_freq) From e20f595aeb83e5e298a22774160035f473bed0eb Mon Sep 17 00:00:00 2001 From: Chase Block Date: Wed, 5 Aug 2026 13:23:18 -0700 Subject: [PATCH 10/14] Add test for bwd path of fused gemm+rope+quant. Signed-off-by: Chase Block --- .../attention/test_fused_mla_q_uproj.py | 58 +++++++++++++++++++ 1 file changed, 58 insertions(+) diff --git a/tests/pytorch/attention/test_fused_mla_q_uproj.py b/tests/pytorch/attention/test_fused_mla_q_uproj.py index 349ce73bae..93f001cd55 100644 --- a/tests/pytorch/attention/test_fused_mla_q_uproj.py +++ b/tests/pytorch/attention/test_fused_mla_q_uproj.py @@ -19,6 +19,7 @@ import transformer_engine.pytorch as te 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 @@ -192,6 +193,63 @@ def test_fused_mla_q_uproj_x_saved_is_mxfp8() -> None: 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) +@pytest.mark.parametrize("tokens", [128, 256], ids=["tokens128", "tokens256"]) +def test_fused_mla_q_uproj_backward(tokens: int) -> None: + """Backward pass (dgrad + wgrad) via general_gemm must match a bf16 reference. + + Mirrors the exact backward path in _FusedMLAQUpProjFunction.backward(): + - grad_output quantized to MXFP8 with optimize_for_gemm=True + - dgrad: general_gemm(w[colwise], gy[rowwise], layout="NN") + - wgrad: general_gemm(x_saved[colwise], gy[colwise], layout="NT") + """ + 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), so quantize with both directions here. + w_bf16 = torch.randn(PROJ_DIM, Q_LORA_RANK, dtype=torch.bfloat16, device=device) + w = MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3, rowwise=True, columnwise=True)(w_bf16) + cos, sin = _build_rope_tables(tokens, device) + + _, x_saved = FusedMLAQUpProjRopeQuant.run(x, w, cos, sin, s, b) + + # Synthetic grad output [tokens, PROJ_DIM] matching the post-RoPE GEMM output shape. + grad_output = torch.randn(tokens, PROJ_DIM, dtype=torch.bfloat16, device=device) + + # Quantize grad output with optimize_for_gemm=True (pre-swizzled scales), matching + # the backward path in _FusedMLAQUpProjFunction.backward(). + 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) + + # dgrad: grad_input = grad_output @ W (A=W[colwise], B=gy[rowwise], NN) + grad_x = general_gemm( + w, gy, layout="NN", grad=True, out_dtype=torch.bfloat16, use_split_accumulator=True + )[0] + + # wgrad: grad_weight = x_saved^T @ grad_output (A=x_saved[colwise], B=gy[colwise], NT) + grad_w = general_gemm( + x_saved, gy, layout="NT", grad=True, out_dtype=torch.bfloat16, use_split_accumulator=True + )[0] + + # bf16 reference (dequantized inputs, no FP8 GEMM noise) + x_dq = x_saved.dequantize().to(torch.bfloat16) + w_dq = w.dequantize().to(torch.bfloat16) + gy_dq = gy.dequantize().to(torch.bfloat16) + + ref_grad_x = gy_dq @ w_dq # [tokens, Q_LORA_RANK] + ref_grad_w = gy_dq.t() @ x_dq # [PROJ_DIM, Q_LORA_RANK] + + torch.testing.assert_close(grad_x, ref_grad_x, atol=0.5, rtol=0.1) + torch.testing.assert_close(grad_w, ref_grad_w, atol=0.5, rtol=0.1) + + def _build_rope_tables(tokens: int, device: torch.device) -> tuple[torch.Tensor, torch.Tensor]: """Plain cos/sin tables for HEAD_DIM_ROPE, interleaved [tokens, rope_dim] bf16.""" inv_freq = 1.0 / ( From 97708298073562f26f3df25f197b7a2e4b7a0673 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Wed, 5 Aug 2026 20:24:51 +0000 Subject: [PATCH 11/14] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- tests/pytorch/attention/test_fused_mla_q_uproj.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/pytorch/attention/test_fused_mla_q_uproj.py b/tests/pytorch/attention/test_fused_mla_q_uproj.py index 93f001cd55..919c9a98f6 100644 --- a/tests/pytorch/attention/test_fused_mla_q_uproj.py +++ b/tests/pytorch/attention/test_fused_mla_q_uproj.py @@ -243,8 +243,8 @@ def test_fused_mla_q_uproj_backward(tokens: int) -> None: w_dq = w.dequantize().to(torch.bfloat16) gy_dq = gy.dequantize().to(torch.bfloat16) - ref_grad_x = gy_dq @ w_dq # [tokens, Q_LORA_RANK] - ref_grad_w = gy_dq.t() @ x_dq # [PROJ_DIM, Q_LORA_RANK] + ref_grad_x = gy_dq @ w_dq # [tokens, Q_LORA_RANK] + ref_grad_w = gy_dq.t() @ x_dq # [PROJ_DIM, Q_LORA_RANK] torch.testing.assert_close(grad_x, ref_grad_x, atol=0.5, rtol=0.1) torch.testing.assert_close(grad_w, ref_grad_w, atol=0.5, rtol=0.1) From a14039ab115b5bfbbf655c49b0a3047e678c8c3e Mon Sep 17 00:00:00 2001 From: Chase Block Date: Wed, 5 Aug 2026 13:49:22 -0700 Subject: [PATCH 12/14] Consolidate fused q proj tests and add to unit test list Signed-off-by: Chase Block --- qa/L0_pytorch_unittest/test.sh | 1 + .../attention/test_fused_mla_q_uproj.py | 184 ++++-------------- 2 files changed, 40 insertions(+), 145 deletions(-) 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 index 919c9a98f6..9c0502830a 100644 --- a/tests/pytorch/attention/test_fused_mla_q_uproj.py +++ b/tests/pytorch/attention/test_fused_mla_q_uproj.py @@ -4,11 +4,6 @@ """Unit tests for FusedMLAQUpProjRopeQuant. -Validates the fused MXFP8 Q up-proj + per-head RoPE + dual-direction MXFP8 quantize kernel -against an unfused reference: MXFP8 quantize inputs -> bf16 GEMM -> RoPE -> MXFP8 quantize. - -DSv3 671B MLA dimensions throughout; tokens must be a multiple of TILE_M (128). - Run: pytest tests/pytorch/attention/test_fused_mla_q_uproj.py -v """ @@ -16,7 +11,7 @@ import pytest import torch -import transformer_engine.pytorch as te +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 @@ -27,7 +22,7 @@ HEAD_DIM_NOPE = 128 HEAD_DIM_ROPE = 64 HEAD_DIM = HEAD_DIM_NOPE + HEAD_DIM_ROPE # 192 -Q_LORA_RANK = 1536 # K dimension of the up-proj GEMM +Q_LORA_RANK = 1536 PROJ_DIM = NUM_HEADS * HEAD_DIM # 24576 SEED = 42 @@ -37,26 +32,16 @@ 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)" - ), + "FusedMLAQUpProjRopeQuant.is_supported() returned False " + "(SM100+, cudnn-frontend >= 1.27.0, and NVTE_FUSED_MLA_Q_UPROJ=1 required)", ) ) -def _build_mxfp8_weight(proj_dim: int, k: int, device: torch.device) -> MXFP8Tensor: - """Quantize a random bf16 weight to MXFP8Tensor (rowwise only, matching the primary param).""" - w_bf16 = torch.randn(proj_dim, k, dtype=torch.bfloat16, device=device) - quantizer = MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3, rowwise=True, columnwise=False) - return quantizer(w_bf16) - - def _dequantize_fused_output(query: MXFP8Tensor, s: int, b: int) -> torch.Tensor: """Dequantize the rowwise fused output to bf16 [s, b, nh, head_dim]. - query._rowwise_data is a plain float8_e4m3fn tensor — TE's C++ dequantize - kernel requires 2D layout, so we reshape before calling MXFP8Tensor.dequantize(). + TE's C++ dequantize kernel requires 2D layout, so reshape before calling dequantize(). """ tokens = s * b q_2d = MXFP8Tensor( @@ -82,78 +67,46 @@ def _reference_q_uproj( s: int, b: int, ) -> torch.Tensor: - """Unfused bf16 reference: dequantize-then-GEMM, RoPE. Returns [s, b, nh, head_dim] bf16. - - We compute in bf16 (dequantizing x and w) rather than re-quantizing the output, - so we compare against the fused kernel's dequantized output directly. - The fused kernel's MXFP8 GEMM + output quantize will naturally introduce some error. - """ + """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) - # Per-head RoPE on the trailing HEAD_DIM_ROPE features, interleaved convention. 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 = q_rope[..., 0::2] - x2 = q_rope[..., 1::2] - x_left = x1 * cos_[..., :half] - x2 * sin_[..., :half] - x_right = x2 * cos_[..., half:] + x1 * sin_[..., half:] - q_rope_out = torch.cat([x_left, x_right], dim=-1) + 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) -@pytest.mark.skipif(not fused_supported, reason=reason_not_supported) -@pytest.mark.parametrize( - "s, b", - [ - (128, 1), # minimum tile (TILE_M=128) - (256, 1), - (128, 2), - ], - ids=["s128_b1", "s256_b1", "s128_b2"], -) -def test_fused_mla_q_uproj_output_shapes(s: int, b: int) -> None: - """Fused kernel output tensors must have the shapes wrap_mxfp8 promises.""" - device = torch.device("cuda") - tokens = s * b - torch.manual_seed(SEED) - - x = torch.randn(tokens, Q_LORA_RANK, dtype=torch.bfloat16, device=device) - w = _build_mxfp8_weight(PROJ_DIM, Q_LORA_RANK, device) - cos, sin = _build_rope_tables(tokens, device) - - query, x_saved = FusedMLAQUpProjRopeQuant.run(x, w, cos, sin, s, b) - - assert isinstance(query, MXFP8Tensor), f"Expected MXFP8Tensor, got {type(query)}" - assert query.shape == ( - s, - b, - NUM_HEADS, - HEAD_DIM, - ), f"Expected query shape {(s, b, NUM_HEADS, HEAD_DIM)}, got {query.shape}" - assert query._rowwise_data is not None, "Missing rowwise data" - assert query._columnwise_data is not None, "Missing columnwise data" - - blk = 32 - assert query._rowwise_scale_inv.shape == (s, b, NUM_HEADS, HEAD_DIM // blk) - assert query._columnwise_scale_inv.shape == (s // blk, b, NUM_HEADS, HEAD_DIM) +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", [128, 256], ids=["tokens128", "tokens256"]) -def test_fused_mla_q_uproj_numerics(tokens: int) -> None: - """Fused kernel output must be close to the unfused MXFP8-precision reference. +@pytest.mark.parametrize("tokens", [256]) +def test_fused_mla_q_uproj(tokens: int) -> None: + """Forward numerics + x_saved properties + backward dgrad/wgrad numerics. - Tolerance reflects the double quantization noise (activation + output quantize). + Forward is checked first so a broken kernel is caught before the backward assertions. """ s, b = tokens, 1 device = torch.device("cuda") @@ -161,102 +114,43 @@ def test_fused_mla_q_uproj_numerics(tokens: int) -> None: torch.cuda.manual_seed(SEED) x = torch.randn(tokens, Q_LORA_RANK, dtype=torch.bfloat16, device=device) - w = _build_mxfp8_weight(PROJ_DIM, Q_LORA_RANK, 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, _ = FusedMLAQUpProjRopeQuant.run(x, w, cos, sin, s, b) + 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) - - # FP8 GEMM + output quantize introduce ~10-20% relative error. torch.testing.assert_close(fused_dq, ref_dq, atol=0.5, rtol=0.1) - -@pytest.mark.skipif(not fused_supported, reason=reason_not_supported) -def test_fused_mla_q_uproj_x_saved_is_mxfp8() -> None: - """In the MXFP8 weight path, x_saved must be an MXFP8Tensor with columnwise data.""" - s, b = 128, 1 - device = torch.device("cuda") - torch.manual_seed(SEED) - - x = torch.randn(s * b, Q_LORA_RANK, dtype=torch.bfloat16, device=device) - w = _build_mxfp8_weight(PROJ_DIM, Q_LORA_RANK, device) - cos, sin = _build_rope_tables(s * b, device) - - _, x_saved = FusedMLAQUpProjRopeQuant.run(x, w, cos, sin, s, b) - - assert isinstance( - x_saved, MXFP8Tensor - ), f"x_saved should be MXFP8Tensor for MXFP8 weight path, got {type(x_saved)}" + # --- 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" - -@pytest.mark.skipif(not fused_supported, reason=reason_not_supported) -@pytest.mark.parametrize("tokens", [128, 256], ids=["tokens128", "tokens256"]) -def test_fused_mla_q_uproj_backward(tokens: int) -> None: - """Backward pass (dgrad + wgrad) via general_gemm must match a bf16 reference. - - Mirrors the exact backward path in _FusedMLAQUpProjFunction.backward(): - - grad_output quantized to MXFP8 with optimize_for_gemm=True - - dgrad: general_gemm(w[colwise], gy[rowwise], layout="NN") - - wgrad: general_gemm(x_saved[colwise], gy[colwise], layout="NT") - """ - 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), so quantize with both directions here. - w_bf16 = torch.randn(PROJ_DIM, Q_LORA_RANK, dtype=torch.bfloat16, device=device) - w = MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3, rowwise=True, columnwise=True)(w_bf16) - cos, sin = _build_rope_tables(tokens, device) - - _, x_saved = FusedMLAQUpProjRopeQuant.run(x, w, cos, sin, s, b) - - # Synthetic grad output [tokens, PROJ_DIM] matching the post-RoPE GEMM output shape. + # --- Backward: dgrad + wgrad --- grad_output = torch.randn(tokens, PROJ_DIM, dtype=torch.bfloat16, device=device) - - # Quantize grad output with optimize_for_gemm=True (pre-swizzled scales), matching - # the backward path in _FusedMLAQUpProjFunction.backward(). 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) - # dgrad: grad_input = grad_output @ W (A=W[colwise], B=gy[rowwise], NN) grad_x = general_gemm( w, gy, layout="NN", grad=True, out_dtype=torch.bfloat16, use_split_accumulator=True )[0] - - # wgrad: grad_weight = x_saved^T @ grad_output (A=x_saved[colwise], B=gy[colwise], NT) grad_w = general_gemm( x_saved, gy, layout="NT", grad=True, out_dtype=torch.bfloat16, use_split_accumulator=True )[0] - # bf16 reference (dequantized inputs, no FP8 GEMM noise) - x_dq = x_saved.dequantize().to(torch.bfloat16) - w_dq = w.dequantize().to(torch.bfloat16) + x_dq = x_saved.dequantize().to(torch.bfloat16) + w_dq = w.dequantize().to(torch.bfloat16) gy_dq = gy.dequantize().to(torch.bfloat16) - ref_grad_x = gy_dq @ w_dq # [tokens, Q_LORA_RANK] - ref_grad_w = gy_dq.t() @ x_dq # [PROJ_DIM, Q_LORA_RANK] - - torch.testing.assert_close(grad_x, ref_grad_x, atol=0.5, rtol=0.1) - torch.testing.assert_close(grad_w, ref_grad_w, atol=0.5, rtol=0.1) - - -def _build_rope_tables(tokens: int, device: torch.device) -> tuple[torch.Tensor, torch.Tensor]: - """Plain cos/sin tables for HEAD_DIM_ROPE, interleaved [tokens, rope_dim] bf16.""" - inv_freq = 1.0 / ( - 10000 - ** (torch.arange(0, HEAD_DIM_ROPE, 2, dtype=torch.float32, device=device) / HEAD_DIM_ROPE) - ) - t = torch.arange(tokens, dtype=torch.float32, device=device) - freqs = torch.outer(t, inv_freq) - freqs = torch.cat([freqs, freqs], dim=-1) - return freqs.cos().to(torch.bfloat16), freqs.sin().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) From d6785bfe291c53fb3730b767ebf3a61ec3354c66 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Wed, 5 Aug 2026 20:52:47 +0000 Subject: [PATCH 13/14] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- .../attention/test_fused_mla_q_uproj.py | 19 ++++++++++++------- 1 file changed, 12 insertions(+), 7 deletions(-) diff --git a/tests/pytorch/attention/test_fused_mla_q_uproj.py b/tests/pytorch/attention/test_fused_mla_q_uproj.py index 9c0502830a..13654deda6 100644 --- a/tests/pytorch/attention/test_fused_mla_q_uproj.py +++ b/tests/pytorch/attention/test_fused_mla_q_uproj.py @@ -32,8 +32,10 @@ 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)", + ( + "FusedMLAQUpProjRopeQuant.is_supported() returned False " + "(SM100+, cudnn-frontend >= 1.27.0, and NVTE_FUSED_MLA_Q_UPROJ=1 required)" + ), ) ) @@ -83,8 +85,11 @@ def _reference_q_uproj( 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 + [ + x1 * cos_[..., :half] - x2 * sin_[..., :half], + x2 * cos_[..., half:] + x1 * sin_[..., half:], + ], + dim=-1, ) return torch.cat([q_nope, q_rope_out], dim=-1) @@ -148,9 +153,9 @@ def test_fused_mla_q_uproj(tokens: int) -> None: 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) + 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_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) From 11c0f1bafa301186e3a83151e75b5985abc681c9 Mon Sep 17 00:00:00 2001 From: Chase Block Date: Thu, 6 Aug 2026 07:48:54 -0700 Subject: [PATCH 14/14] Remove backward test from fused mla q uproj. These tests really belong in Megatron. Signed-off-by: Chase Block --- .../attention/test_fused_mla_q_uproj.py | 36 ++++--------------- .../pytorch/attention/fused_mla_q_uproj.py | 2 +- 2 files changed, 7 insertions(+), 31 deletions(-) diff --git a/tests/pytorch/attention/test_fused_mla_q_uproj.py b/tests/pytorch/attention/test_fused_mla_q_uproj.py index 13654deda6..2c061607fc 100644 --- a/tests/pytorch/attention/test_fused_mla_q_uproj.py +++ b/tests/pytorch/attention/test_fused_mla_q_uproj.py @@ -14,7 +14,6 @@ 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 @@ -109,9 +108,10 @@ 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 + x_saved properties + backward dgrad/wgrad numerics. + """Forward numerics and x_saved properties for FusedMLAQUpProjRopeQuant.run(). - Forward is checked first so a broken kernel is caught before the backward assertions. + Full forward+backward autograd testing (via _FusedMLAQUpProjFunction) lives in + Megatron-Core. """ s, b = tokens, 1 device = torch.device("cuda") @@ -119,43 +119,19 @@ def test_fused_mla_q_uproj(tokens: int) -> None: 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)( + w = MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3, rowwise=True, columnwise=False)( 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 --- + # Forward numerics: FP8 GEMM + output quantize introduce ~10% relative error. 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 --- + # x_saved: must be MXFP8 with only columnwise data retained for wgrad. 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/attention/fused_mla_q_uproj.py b/transformer_engine/pytorch/attention/fused_mla_q_uproj.py index cd62e16f3a..c176985254 100644 --- a/transformer_engine/pytorch/attention/fused_mla_q_uproj.py +++ b/transformer_engine/pytorch/attention/fused_mla_q_uproj.py @@ -83,7 +83,7 @@ def run( from cuda.bindings import driver as cuda - stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) + stream = cuda.CUstream(torch.cuda.current_stream(x.device).cuda_stream) wrapper = cls._kernel() if isinstance(w, QuantizedTensor):