From c19f0df7b64e6d4f339195a6ac3ac2ca9911a6c4 Mon Sep 17 00:00:00 2001 From: Chase Block Date: Tue, 7 Jul 2026 16:16:46 -0700 Subject: [PATCH 1/5] 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 2/5] [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 3/5] 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 4/5] [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 5/5] 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(