From 4cba8bc35b798db72be82378dec80712e68cd16a Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Tue, 21 Apr 2026 15:42:43 +0200 Subject: [PATCH 01/11] [PyTorch] Make UnfusedDotProductAttention compatible with torch.compile + CUDA graphs Refactor TE custom kernels used by the unfused attention path so that `torch.compile(fullgraph=True, mode="reduce-overhead")` can trace the forward and backward and capture them into CUDA graphs without graph breaks. - softmax.py / softmax.cpp: register all `scaled_*_softmax_{forward,backward}` kernels as `torch.library.custom_op`s with fake impls and an autograd binding that mirrors the previous `torch.autograd.Function`s. The C++ backward kernels now allocate a fresh output buffer instead of writing in-place into `output_grad`, so the ops no longer alias their inputs (required by `torch.library.custom_op` and inductor cudagraph trees). - utils.py: convert `ConvertTHDtoBSHD` / `ConvertBSHDtoTHD` to `torch.library.custom_op`s, with thin wrapper classes that keep the existing `.apply(...)` callsite syntax. Drop the `int(cu_seqlens[-1].item())` from the hot path of `ConvertBSHDtoTHD.apply` -- under `torch.compile` it created an unbacked SymInt, which made the Inductor partitioner emit `None` placeholders for output buffers and caused `cudagraph_trees` to assert. `num_tokens` is now passed in by the caller as a regular (Sym)Int. - backends.py: in the THD branch of unfused DPA, capture `total_tokens_q = query_layer.shape[0]` before overwriting `query_layer` with the BSHD form, and thread it back into `ConvertBSHDtoTHD.apply` at the end of the forward. - test_torch_compile.py: add `test_unfused_dpa_torch_compile`, parametrized over qkv layouts (`bshd_bshd_bshd`, `sbhd_sbhd_sbhd`, `thd_thd_thd`, `bs3hd`, `sbh3d`), that compiles `UnfusedDotProductAttention.forward` directly with `fullgraph=True, mode="reduce-overhead"` and runs forward+backward several times so the CUDA graphs are recorded and replayed. Signed-off-by: Pawel Gadzinski Made-with: Cursor --- tests/pytorch/test_torch_compile.py | 163 +++++++++ .../dot_product_attention/backends.py | 7 + .../dot_product_attention/softmax.py | 324 ++++++++++++------ .../attention/dot_product_attention/utils.py | 174 +++++++--- .../pytorch/csrc/extensions/softmax.cpp | 44 ++- 5 files changed, 546 insertions(+), 166 deletions(-) diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py index 309a5d124e..ca515aaebe 100644 --- a/tests/pytorch/test_torch_compile.py +++ b/tests/pytorch/test_torch_compile.py @@ -3,6 +3,7 @@ # See LICENSE for license information. import abc +import os import pytest import torch @@ -33,6 +34,9 @@ is_nvfp4_available, ) from utils import recipe_id +from transformer_engine.pytorch.attention.dot_product_attention.backends import ( + UnfusedDotProductAttention, +) fp8_available, reason_for_no_fp8 = is_fp8_available(return_reason=True) mxfp8_available, reason_for_no_mxfp8 = is_mxfp8_available(return_reason=True) @@ -363,3 +367,162 @@ def fn(inp): out = compiled(inp) out.sum().backward() + + + +_UNFUSED_DPA_CONFIG = dict( + batch_size=2, + num_heads=4, + head_dim=64, + max_seqlen_q=128, + max_seqlen_kv=128, +) + + +def _make_unfused_attention(dtype: torch.dtype) -> UnfusedDotProductAttention: + cfg = _UNFUSED_DPA_CONFIG + softmax_scale = cfg["head_dim"] ** -0.5 + module = UnfusedDotProductAttention( + softmax_scale=softmax_scale, + attention_type="self", + attention_dropout=0.0, + layer_number=1, + softmax_type="vanilla", + return_max_logit=False, + ) + return module.to(dtype=dtype, device="cuda") + + +_EMPTY_ALIBI_CACHE = { + "_num_heads": None, + "_alibi_slopes": None, + "_max_seqlen_q": None, + "_max_seqlen_kv": None, + "_bottom_right_alignment": True, + "_alibi_bias": None, + "_alibi_slopes_require_update": False, + "_alibi_bias_require_update": False, +} + + +def _make_unfused_qkv(qkv_layout: str, dtype: torch.dtype, requires_grad: bool = True): + """Build (q, k, v) tensors matching `qkv_layout`. Returns also the + extra kwargs (`cu_seqlens_*`, `max_seqlen_*`) that the unfused module + needs for `thd` layouts (empty dict otherwise).""" + cfg = _UNFUSED_DPA_CONFIG + b, s_q, s_kv = cfg["batch_size"], cfg["max_seqlen_q"], cfg["max_seqlen_kv"] + h, d = cfg["num_heads"], cfg["head_dim"] + qkv_format = "".join(c for c in qkv_layout.split("_")[0] if c.isalpha()) + + extra: dict = {} + + def _separate(shape): + return tuple( + torch.randn(shape, dtype=dtype, device="cuda", requires_grad=requires_grad) + for _ in range(3) + ) + + if qkv_layout == "bshd_bshd_bshd": + q, k, v = _separate((b, s_q, h, d)) + elif qkv_layout == "sbhd_sbhd_sbhd": + q, k, v = _separate((s_q, b, h, d)) + elif qkv_layout == "thd_thd_thd": + # All sequences in the batch have the maximum length; no padding. + cu = torch.arange(0, (b + 1) * s_q, step=s_q, dtype=torch.int32, device="cuda") + q, k, v = _separate((b * s_q, h, d)) + extra = dict( + cu_seqlens_q=cu, + cu_seqlens_kv=cu, + max_seqlen_q=s_q, + max_seqlen_kv=s_kv, + ) + elif qkv_layout == "bs3hd": + # Packed: shape (b, s, 3, h, d), q/k/v are views along dim=-3. + qkv = torch.randn( + (b, s_q, 3, h, d), + dtype=dtype, + device="cuda", + requires_grad=requires_grad, + ) + q, k, v = qkv[:, :, 0], qkv[:, :, 1], qkv[:, :, 2] + # q/k/v are non-leaf views; retain their grads so the assertions in + # the test (`q.grad is not None` etc.) work for packed layouts. + if requires_grad: + for t in (q, k, v): + t.retain_grad() + elif qkv_layout == "sbh3d": + # Packed: shape (s, b, h, 3, d), q/k/v are views along dim=-2. + qkv = torch.randn( + (s_q, b, h, 3, d), + dtype=dtype, + device="cuda", + requires_grad=requires_grad, + ) + q, k, v = qkv[:, :, :, 0], qkv[:, :, :, 1], qkv[:, :, :, 2] + if requires_grad: + for t in (q, k, v): + t.retain_grad() + else: + raise ValueError(f"Unsupported qkv_layout in test: {qkv_layout}") + + return q, k, v, extra, qkv_format + + +def _call_unfused( + module: UnfusedDotProductAttention, + qkv_layout: str, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + extra: dict, +) -> torch.Tensor: + return module( + _EMPTY_ALIBI_CACHE, + q, + k, + v, + qkv_layout=qkv_layout, + attn_mask_type="causal", + **extra, + ) + + +@pytest.mark.parametrize( + "qkv_layout", + [ + "bshd_bshd_bshd", + "sbhd_sbhd_sbhd", + "thd_thd_thd", + "bs3hd", + "sbh3d", + ], +) +def test_unfused_dpa_torch_compile(qkv_layout): + """Compile UnfusedDotProductAttention.forward with + `torch.compile(fullgraph=True, mode="reduce-overhead")` for several + qkv layouts. + + - `fullgraph=True` makes the test fail on any graph break inside the + unfused attention path. + - `mode="reduce-overhead"` uses the inductor cudagraphs backend, so + forward+backward are captured into CUDA graphs and replayed on + subsequent iterations.""" + dtype = torch.bfloat16 + + module = _make_unfused_attention(dtype) + + def fn(q, k, v, extra): + return _call_unfused(module, qkv_layout, q, k, v, extra) + + torch._dynamo.reset() + compiled = torch.compile(fn, fullgraph=True, mode="reduce-overhead") + + for _ in range(3): + q, k, v, extra, _ = _make_unfused_qkv(qkv_layout, dtype, requires_grad=True) + out = compiled(q, k, v, extra) + out.sum().backward() + torch.cuda.synchronize() + assert torch.isfinite(out).all() + assert q.grad is not None + assert k.grad is not None + assert v.grad is not None diff --git a/transformer_engine/pytorch/attention/dot_product_attention/backends.py b/transformer_engine/pytorch/attention/dot_product_attention/backends.py index 8f42983553..e5180eab8b 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/backends.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/backends.py @@ -425,6 +425,12 @@ def forward( if qkv_format == "thd": assert cu_seqlens_q is not None and cu_seqlens_kv is not None assert max_seqlen_q is not None and max_seqlen_kv is not None + # Capture token counts as plain (Sym)Ints from the THD inputs *before* + # we overwrite `query_layer` with the BSHD form. We thread `total_tokens_q` + # back into `ConvertBSHDtoTHD.apply` below so that it does not need to + # call `cu_seqlens_q[-1].item()` itself -- which under `torch.compile` + # would create an unbacked SymInt and break Inductor + cudagraphs. + total_tokens_q = query_layer.shape[0] query_layer = ConvertTHDtoBSHD.apply( query_layer, cu_seqlens_q, @@ -708,6 +714,7 @@ def forward( context_layer = ConvertBSHDtoTHD.apply( context_layer, cu_seqlens_q, + total_tokens_q, ) # [tq, h, d] --> [tq, hd] diff --git a/transformer_engine/pytorch/attention/dot_product_attention/softmax.py b/transformer_engine/pytorch/attention/dot_product_attention/softmax.py index 74d9583ce5..1207b67985 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/softmax.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/softmax.py @@ -4,7 +4,7 @@ """Fused scaled masked softmax functions""" import os -from typing import Callable, Tuple, Union, Optional +from typing import Callable, Optional import torch from torch import nn import transformer_engine_torch as tex @@ -18,127 +18,249 @@ _default_causal_mask = {} -def _get_default_causal_mask(mask_type: str, sq: int, sk: int) -> torch.Tensor: - """Return the causal upper triangular mask for softmax input""" +def _scale_to_tensor(scale: float) -> torch.Tensor: + """Wrap a Python float in a 0-D tensor expected by the tex kernels.""" + return torch.tensor([scale])[0] - def _get_mask(): - diagonal_offset = sk - sq + 1 if "bottom_right" in mask_type else 1 - return torch.triu( - torch.ones(sq, sk, dtype=torch.bool, device="cuda"), diagonal=diagonal_offset - ) - if is_in_onnx_export_mode(): - return _get_mask() - matrix_identifiers = (mask_type, sq, sk) - if matrix_identifiers not in _default_causal_mask: - _default_causal_mask[matrix_identifiers] = _get_mask() - return _default_causal_mask[matrix_identifiers] +# ----------------------------- ScaledSoftmax ------------------------------- -class ScaledUpperTriangMaskedSoftmax(torch.autograd.Function): - """ - Fused operation which performs following three operations in sequence - 1. Scale the tensor. - 2. Apply upper triangular mask (typically used in gpt models). - 3. Perform softmax. - """ +@torch.library.custom_op("te_softmax::scaled_softmax_fwd", mutates_args=()) +def scaled_softmax_forward(inputs: torch.Tensor, scale: float) -> torch.Tensor: + """Forward pass for ScaledSoftmax.""" + return tex.scaled_softmax_forward(inputs, _scale_to_tensor(scale)) - @staticmethod - def forward(ctx, inputs: torch.Tensor, scale: float) -> torch.Tensor: - """ScaledUpperTriangMaskedSoftmax fwd""" - scale_t = torch.tensor([scale]) - softmax_results = tex.scaled_upper_triang_masked_softmax_forward(inputs, scale_t[0]) - ctx.save_for_backward(softmax_results, scale_t) - return softmax_results +@scaled_softmax_forward.register_fake +def _scaled_softmax_forward_fake(inputs: torch.Tensor, scale: float) -> torch.Tensor: + del scale + return torch.empty_like(inputs) - @staticmethod - def backward(ctx, output_grads: torch.Tensor) -> Tuple[Union[torch.Tensor, None], ...]: - """ScaledUpperTriangMaskedSoftmax bwd""" - softmax_results, scale_t = ctx.saved_tensors - input_grads = tex.scaled_upper_triang_masked_softmax_backward( - output_grads, softmax_results, scale_t[0] - ) - return input_grads, None +@torch.library.custom_op("te_softmax::scaled_softmax_bwd", mutates_args=()) +def scaled_softmax_backward( + output_grads: torch.Tensor, softmax_results: torch.Tensor, scale: float +) -> torch.Tensor: + """Backward pass for ScaledSoftmax.""" + return tex.scaled_softmax_backward(output_grads, softmax_results, _scale_to_tensor(scale)) -class ScaledAlignedCausalMaskedSoftmax(torch.autograd.Function): - """ - Fused operation which performs following three operations in sequence - 1. Scale the tensor. - 2. Apply causal mask aligned to the bottom right corner of the input matrix - 3. Perform softmax. - """ +@scaled_softmax_backward.register_fake +def _scaled_softmax_backward_fake( + output_grads: torch.Tensor, softmax_results: torch.Tensor, scale: float +) -> torch.Tensor: + del softmax_results, scale + return torch.empty_like(output_grads) - @staticmethod - def forward(ctx, inputs: torch.Tensor, scale: float) -> torch.Tensor: - """ScaledAlignedCausalMaskedSoftmax fwd""" - scale_t = torch.tensor([scale]) - softmax_results = tex.scaled_aligned_causal_masked_softmax_forward(inputs, scale_t[0]) - ctx.save_for_backward(softmax_results, scale_t) - return softmax_results - @staticmethod - def backward(ctx, output_grads: torch.Tensor) -> Tuple[Union[torch.Tensor, None], ...]: - """ScaledAlignedCausalMaskedSoftmax bwd""" - softmax_results, scale_t = ctx.saved_tensors - input_grads = tex.scaled_aligned_causal_masked_softmax_backward( - output_grads, softmax_results, scale_t[0] - ) +def _scaled_softmax_setup_context(ctx, inputs, output): + _inp, scale = inputs + ctx.scale = scale + ctx.save_for_backward(output) - return input_grads, None +def _scaled_softmax_backward_wrapper(ctx, grad_output): + (softmax_results,) = ctx.saved_tensors + grad_inputs = torch.ops.te_softmax.scaled_softmax_bwd(grad_output, softmax_results, ctx.scale) + return grad_inputs, None -class ScaledMaskedSoftmax(torch.autograd.Function): - """ - Fused operation which performs following three operations in sequence - 1. Scale the tensor. - 2. Apply the mask. - 3. Perform softmax. - """ - @staticmethod - def forward(ctx, inputs: torch.Tensor, mask: torch.Tensor, scale: float) -> torch.Tensor: - """ScaledMaskedSoftmax fwd""" - scale_t = torch.tensor([scale]) +scaled_softmax_forward.register_autograd( + _scaled_softmax_backward_wrapper, + setup_context=_scaled_softmax_setup_context, +) - softmax_results = tex.scaled_masked_softmax_forward(inputs, mask, scale_t[0]) - ctx.save_for_backward(softmax_results, scale_t) - return softmax_results - @staticmethod - def backward(ctx, output_grads: torch.Tensor) -> Tuple[Union[torch.Tensor, None], ...]: - """ScaledMaskedSoftmax bwd""" - softmax_results, scale_t = ctx.saved_tensors +# --------------------------- ScaledMaskedSoftmax --------------------------- - input_grads = tex.scaled_masked_softmax_backward(output_grads, softmax_results, scale_t[0]) - return input_grads, None, None +@torch.library.custom_op("te_softmax::scaled_masked_softmax_fwd", mutates_args=()) +def scaled_masked_softmax_forward( + inputs: torch.Tensor, mask: torch.Tensor, scale: float +) -> torch.Tensor: + """Forward pass for ScaledMaskedSoftmax.""" + return tex.scaled_masked_softmax_forward(inputs, mask, _scale_to_tensor(scale)) -class ScaledSoftmax(torch.autograd.Function): - """ - Fused operation which performs following two operations in sequence - 1. Scale the tensor. - 2. Perform softmax. - """ - @staticmethod - def forward(ctx, inputs: torch.Tensor, scale: float) -> torch.Tensor: - """ScaledSoftmax fwd""" - scale_t = torch.tensor([scale]) +@scaled_masked_softmax_forward.register_fake +def _scaled_masked_softmax_forward_fake( + inputs: torch.Tensor, mask: torch.Tensor, scale: float +) -> torch.Tensor: + del mask, scale + return torch.empty_like(inputs) - softmax_results = tex.scaled_softmax_forward(inputs, scale_t[0]) - ctx.save_for_backward(softmax_results, scale_t) - return softmax_results - @staticmethod - def backward(ctx, output_grads: torch.Tensor) -> Tuple[Union[torch.Tensor, None], ...]: - """ScaledSoftmax bwd""" - softmax_results, scale_t = ctx.saved_tensors +@torch.library.custom_op("te_softmax::scaled_masked_softmax_bwd", mutates_args=()) +def scaled_masked_softmax_backward( + output_grads: torch.Tensor, softmax_results: torch.Tensor, scale: float +) -> torch.Tensor: + """Backward pass for ScaledMaskedSoftmax.""" + return tex.scaled_masked_softmax_backward( + output_grads, softmax_results, _scale_to_tensor(scale) + ) + + +@scaled_masked_softmax_backward.register_fake +def _scaled_masked_softmax_backward_fake( + output_grads: torch.Tensor, softmax_results: torch.Tensor, scale: float +) -> torch.Tensor: + del softmax_results, scale + return torch.empty_like(output_grads) + + +def _scaled_masked_softmax_setup_context(ctx, inputs, output): + _inp, _mask, scale = inputs + ctx.scale = scale + ctx.save_for_backward(output) + + +def _scaled_masked_softmax_backward_wrapper(ctx, grad_output): + (softmax_results,) = ctx.saved_tensors + grad_inputs = torch.ops.te_softmax.scaled_masked_softmax_bwd( + grad_output, softmax_results, ctx.scale + ) + return grad_inputs, None, None + + +scaled_masked_softmax_forward.register_autograd( + _scaled_masked_softmax_backward_wrapper, + setup_context=_scaled_masked_softmax_setup_context, +) + + +# ---------------------- ScaledUpperTriangMaskedSoftmax ---------------------- + + +@torch.library.custom_op("te_softmax::scaled_upper_triang_masked_softmax_fwd", mutates_args=()) +def scaled_upper_triang_masked_softmax_forward( + inputs: torch.Tensor, scale: float +) -> torch.Tensor: + """Forward pass for ScaledUpperTriangMaskedSoftmax.""" + return tex.scaled_upper_triang_masked_softmax_forward(inputs, _scale_to_tensor(scale)) + + +@scaled_upper_triang_masked_softmax_forward.register_fake +def _scaled_upper_triang_masked_softmax_forward_fake( + inputs: torch.Tensor, scale: float +) -> torch.Tensor: + del scale + return torch.empty_like(inputs) - input_grads = tex.scaled_softmax_backward(output_grads, softmax_results, scale_t[0]) - return input_grads, None, None + +@torch.library.custom_op("te_softmax::scaled_upper_triang_masked_softmax_bwd", mutates_args=()) +def scaled_upper_triang_masked_softmax_backward( + output_grads: torch.Tensor, softmax_results: torch.Tensor, scale: float +) -> torch.Tensor: + """Backward pass for ScaledUpperTriangMaskedSoftmax.""" + return tex.scaled_upper_triang_masked_softmax_backward( + output_grads, softmax_results, _scale_to_tensor(scale) + ) + + +@scaled_upper_triang_masked_softmax_backward.register_fake +def _scaled_upper_triang_masked_softmax_backward_fake( + output_grads: torch.Tensor, softmax_results: torch.Tensor, scale: float +) -> torch.Tensor: + del softmax_results, scale + return torch.empty_like(output_grads) + + +def _scaled_upper_triang_masked_softmax_setup_context(ctx, inputs, output): + _inp, scale = inputs + ctx.scale = scale + ctx.save_for_backward(output) + + +def _scaled_upper_triang_masked_softmax_backward_wrapper(ctx, grad_output): + (softmax_results,) = ctx.saved_tensors + grad_inputs = torch.ops.te_softmax.scaled_upper_triang_masked_softmax_bwd( + grad_output, softmax_results, ctx.scale + ) + return grad_inputs, None + + +scaled_upper_triang_masked_softmax_forward.register_autograd( + _scaled_upper_triang_masked_softmax_backward_wrapper, + setup_context=_scaled_upper_triang_masked_softmax_setup_context, +) + + +# -------------------- ScaledAlignedCausalMaskedSoftmax --------------------- + + +@torch.library.custom_op("te_softmax::scaled_aligned_causal_masked_softmax_fwd", mutates_args=()) +def scaled_aligned_causal_masked_softmax_forward( + inputs: torch.Tensor, scale: float +) -> torch.Tensor: + """Forward pass for ScaledAlignedCausalMaskedSoftmax.""" + return tex.scaled_aligned_causal_masked_softmax_forward(inputs, _scale_to_tensor(scale)) + + +@scaled_aligned_causal_masked_softmax_forward.register_fake +def _scaled_aligned_causal_masked_softmax_forward_fake( + inputs: torch.Tensor, scale: float +) -> torch.Tensor: + del scale + return torch.empty_like(inputs) + + +@torch.library.custom_op("te_softmax::scaled_aligned_causal_masked_softmax_bwd", mutates_args=()) +def scaled_aligned_causal_masked_softmax_backward( + output_grads: torch.Tensor, softmax_results: torch.Tensor, scale: float +) -> torch.Tensor: + """Backward pass for ScaledAlignedCausalMaskedSoftmax.""" + return tex.scaled_aligned_causal_masked_softmax_backward( + output_grads, softmax_results, _scale_to_tensor(scale) + ) + + +@scaled_aligned_causal_masked_softmax_backward.register_fake +def _scaled_aligned_causal_masked_softmax_backward_fake( + output_grads: torch.Tensor, softmax_results: torch.Tensor, scale: float +) -> torch.Tensor: + del softmax_results, scale + return torch.empty_like(output_grads) + + +def _scaled_aligned_causal_masked_softmax_setup_context(ctx, inputs, output): + _inp, scale = inputs + ctx.scale = scale + ctx.save_for_backward(output) + + +def _scaled_aligned_causal_masked_softmax_backward_wrapper(ctx, grad_output): + (softmax_results,) = ctx.saved_tensors + grad_inputs = torch.ops.te_softmax.scaled_aligned_causal_masked_softmax_bwd( + grad_output, softmax_results, ctx.scale + ) + return grad_inputs, None + + +scaled_aligned_causal_masked_softmax_forward.register_autograd( + _scaled_aligned_causal_masked_softmax_backward_wrapper, + setup_context=_scaled_aligned_causal_masked_softmax_setup_context, +) + + +_default_causal_mask = {} + + +def _get_default_causal_mask(mask_type: str, sq: int, sk: int) -> torch.Tensor: + """Return the causal upper triangular mask for softmax input""" + + def _get_mask(): + diagonal_offset = sk - sq + 1 if "bottom_right" in mask_type else 1 + return torch.triu( + torch.ones(sq, sk, dtype=torch.bool, device="cuda"), diagonal=diagonal_offset + ) + + if is_in_onnx_export_mode(): + return _get_mask() + matrix_identifiers = (mask_type, sq, sk) + if matrix_identifiers not in _default_causal_mask: + _default_causal_mask[matrix_identifiers] = _get_mask() + return _default_causal_mask[matrix_identifiers] class FusedScaleMaskSoftmax(nn.Module): @@ -234,16 +356,16 @@ def forward_fused_softmax( padding, padding_causal, padding_causal_bottom_right | ScaledMaskedSoftmax arbitrary ([1, 1, sq, sk] or [b, 1, sq, sk]) | ScaledMaskedSoftmax """ - scale = 1.0 if scale is None else scale + scale = 1.0 if scale is None else float(scale) # Disable for now until unalignment bug is fixed. # if self.attn_mask_type in ["causal", "causal_bottom_right"]: - # return ScaledAlignedCausalMaskedSoftmax.apply(inp, scale) + # return torch.ops.te_softmax.scaled_aligned_causal_masked_softmax_fwd(inp, scale) # input is 4D tensor (1, 1, sq, sk) or (b, 1, sq, sk) if mask is not None and self.attn_mask_type != "no_mask": - return ScaledMaskedSoftmax.apply(inp, mask, scale) - return ScaledSoftmax.apply(inp, scale) + return torch.ops.te_softmax.scaled_masked_softmax_fwd(inp, mask, scale) + return torch.ops.te_softmax.scaled_softmax_fwd(inp, scale) def forward_torch_softmax( self, inp: torch.Tensor, mask: torch.Tensor, scale: Optional[float] = None diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py index 8a07d7af79..baac46e232 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py @@ -2138,76 +2138,144 @@ def backward(ctx, grad_output): return None, None, _pack_tensor(indices, grad_output) -class ConvertTHDtoBSHD(torch.autograd.Function): +# --------------------------------------------------------------------------- +# THD <-> BSHD conversions exposed as `torch.library` custom ops so that +# `torch.compile` can trace them. Backward of each direction is the other +# direction, wired through `register_autograd` + `setup_context`, mirroring +# the pattern in `transformer_engine/pytorch/permutation.py`. +# --------------------------------------------------------------------------- + + +@torch.library.custom_op("te_attention::convert_thd_to_bshd", mutates_args=()) +def _convert_thd_to_bshd_op( + thd_tensor: torch.Tensor, + cu_seqlens: torch.Tensor, + batch_size: int, + max_seqlen: int, +) -> torch.Tensor: + """Forward pass for THD->BSHD conversion.""" + if not thd_tensor.is_contiguous(): + thd_tensor = thd_tensor.contiguous() + return tex.convert_thd_to_bshd(thd_tensor, cu_seqlens, batch_size, max_seqlen) + + +@_convert_thd_to_bshd_op.register_fake +def _convert_thd_to_bshd_fake( + thd_tensor: torch.Tensor, + cu_seqlens: torch.Tensor, + batch_size: int, + max_seqlen: int, +) -> torch.Tensor: + del cu_seqlens + h, d = thd_tensor.shape[1], thd_tensor.shape[2] + return torch.empty( + (batch_size, max_seqlen, h, d), dtype=thd_tensor.dtype, device=thd_tensor.device + ) + + +@torch.library.custom_op("te_attention::convert_bshd_to_thd", mutates_args=()) +def _convert_bshd_to_thd_op( + bshd_tensor: torch.Tensor, + cu_seqlens: torch.Tensor, + num_tokens: int, +) -> torch.Tensor: + """Forward pass for BSHD->THD conversion.""" + if not bshd_tensor.is_contiguous(): + bshd_tensor = bshd_tensor.contiguous() + return tex.convert_bshd_to_thd(bshd_tensor, cu_seqlens, num_tokens) + + +@_convert_bshd_to_thd_op.register_fake +def _convert_bshd_to_thd_fake( + bshd_tensor: torch.Tensor, + cu_seqlens: torch.Tensor, + num_tokens: int, +) -> torch.Tensor: + del cu_seqlens + h, d = bshd_tensor.shape[2], bshd_tensor.shape[3] + return torch.empty( + (num_tokens, h, d), dtype=bshd_tensor.dtype, device=bshd_tensor.device + ) + + +def _convert_thd_to_bshd_setup_context(ctx, inputs, output): + thd_tensor, cu_seqlens, _batch_size, _max_seqlen = inputs + ctx.save_for_backward(cu_seqlens) + ctx.num_tokens = thd_tensor.size(0) + + +def _convert_thd_to_bshd_backward_wrapper(ctx, grad_bshd): + (cu_seqlens,) = ctx.saved_tensors + grad_thd = torch.ops.te_attention.convert_bshd_to_thd( + grad_bshd, cu_seqlens, ctx.num_tokens + ) + return grad_thd, None, None, None + + +_convert_thd_to_bshd_op.register_autograd( + _convert_thd_to_bshd_backward_wrapper, + setup_context=_convert_thd_to_bshd_setup_context, +) + + +def _convert_bshd_to_thd_setup_context(ctx, inputs, output): + bshd_tensor, cu_seqlens, _num_tokens = inputs + ctx.save_for_backward(cu_seqlens) + ctx.batch_size = bshd_tensor.size(0) + ctx.max_seqlen = bshd_tensor.size(1) + + +def _convert_bshd_to_thd_backward_wrapper(ctx, grad_thd): + (cu_seqlens,) = ctx.saved_tensors + grad_bshd = torch.ops.te_attention.convert_thd_to_bshd( + grad_thd, cu_seqlens, ctx.batch_size, ctx.max_seqlen + ) + return grad_bshd, None, None + + +_convert_bshd_to_thd_op.register_autograd( + _convert_bshd_to_thd_backward_wrapper, + setup_context=_convert_bshd_to_thd_setup_context, +) + + +class ConvertTHDtoBSHD: """ Convert a tensor from qkv_format = thd to qkv_format = bshd. + + Thin wrapper around the `te_attention::convert_thd_to_bshd` custom op, + kept so that callsites can continue to use the `.apply(...)` syntax. """ @staticmethod - def forward(ctx, thd_tensor, cu_seqlens, max_seqlen): + def apply(thd_tensor, cu_seqlens, max_seqlen): # pylint: disable=missing-function-docstring batch_size = cu_seqlens.shape[0] - 1 - if not thd_tensor.is_contiguous(): - thd_tensor = thd_tensor.contiguous() - bshd_tensor = tex.convert_thd_to_bshd( - thd_tensor, - cu_seqlens, - batch_size, - max_seqlen, - ) - ctx.save_for_backward(cu_seqlens) - ctx.num_tokens = thd_tensor.shape[0] - return bshd_tensor - - @staticmethod - def backward(ctx, bshd_tensor): - # pylint: disable=missing-function-docstring - (cu_seqlens,) = ctx.saved_tensors - if not bshd_tensor.is_contiguous(): - bshd_tensor = bshd_tensor.contiguous() - thd_tensor = tex.convert_bshd_to_thd( - bshd_tensor, - cu_seqlens, - ctx.num_tokens, + return torch.ops.te_attention.convert_thd_to_bshd( + thd_tensor, cu_seqlens, batch_size, max_seqlen ) - return thd_tensor, None, None -class ConvertBSHDtoTHD(torch.autograd.Function): +class ConvertBSHDtoTHD: """ Convert a tensor from qkv_format = bshd to qkv_format = thd. - """ - @staticmethod - def forward(ctx, bshd_tensor, cu_seqlens): - # pylint: disable=missing-function-docstring - num_tokens = cu_seqlens[-1] - max_seqlen = bshd_tensor.shape[1] - if not bshd_tensor.is_contiguous(): - bshd_tensor = bshd_tensor.contiguous() - thd_tensor = tex.convert_bshd_to_thd( - bshd_tensor, - cu_seqlens, - num_tokens, - ) - ctx.save_for_backward(cu_seqlens) - ctx.max_seqlen = max_seqlen - return thd_tensor + Thin wrapper around the `te_attention::convert_bshd_to_thd` custom op, + kept so that callsites can continue to use the `.apply(...)` syntax. + """ @staticmethod - def backward(ctx, thd_tensor): + def apply(bshd_tensor, cu_seqlens, num_tokens): # pylint: disable=missing-function-docstring - (cu_seqlens,) = ctx.saved_tensors - batch_size = cu_seqlens.shape[0] - 1 - if not thd_tensor.is_contiguous(): - thd_tensor = thd_tensor.contiguous() - bshd_tensor = tex.convert_thd_to_bshd( - thd_tensor, - cu_seqlens, - batch_size, - ctx.max_seqlen, + # `num_tokens` (== `cu_seqlens[-1]`) is taken as an explicit argument + # rather than being read from `cu_seqlens[-1].item()` here, so that + # `torch.compile` does not introduce an unbacked SymInt for it. With an + # unbacked SymInt the Inductor partitioner emits `None` placeholders for + # output buffers, which `cudagraph_trees` then refuses (only + # `Tensor / int / Generator` allowed as graph inputs). + return torch.ops.te_attention.convert_bshd_to_thd( + bshd_tensor, cu_seqlens, num_tokens ) - return bshd_tensor, None def get_qkv_format( diff --git a/transformer_engine/pytorch/csrc/extensions/softmax.cpp b/transformer_engine/pytorch/csrc/extensions/softmax.cpp index 3bb6a5e7b3..3c771ae5b6 100644 --- a/transformer_engine/pytorch/csrc/extensions/softmax.cpp +++ b/transformer_engine/pytorch/csrc/extensions/softmax.cpp @@ -52,15 +52,20 @@ at::Tensor scaled_softmax_backward(at::Tensor output_grad_, at::Tensor softmax_r (softmax_results.scalar_type() == at::ScalarType::BFloat16), "Only fp16 and bf16 are supported"); + // Allocate a fresh output buffer so the op does not alias / mutate its + // inputs (required by `torch.library.custom_op`). + auto input_grads = + torch::empty(output_grads.sizes(), output_grads.options().requires_grad(false)); + auto output_grads_cu = makeTransformerEngineTensor(output_grads); auto softmax_results_cu = makeTransformerEngineTensor(softmax_results); + auto input_grads_cu = makeTransformerEngineTensor(input_grads); - // Produce gradients in place. nvte_scaled_softmax_backward(output_grads_cu.data(), softmax_results_cu.data(), - output_grads_cu.data(), scale_factor, + input_grads_cu.data(), scale_factor, at::cuda::getCurrentCUDAStream()); - return output_grads; + return input_grads; } at::Tensor scaled_masked_softmax_forward(at::Tensor input, at::Tensor mask, float scale_factor) { @@ -115,15 +120,20 @@ at::Tensor scaled_masked_softmax_backward(at::Tensor output_grad_, at::Tensor so (softmax_results.scalar_type() == at::ScalarType::BFloat16), "Only fp16 and bf16 are supported"); + // Allocate a fresh output buffer so the op does not alias / mutate its + // inputs (required by `torch.library.custom_op`). + auto input_grads = + torch::empty(output_grads.sizes(), output_grads.options().requires_grad(false)); + auto output_grads_cu = makeTransformerEngineTensor(output_grads); auto softmax_results_cu = makeTransformerEngineTensor(softmax_results); + auto input_grads_cu = makeTransformerEngineTensor(input_grads); - // Produce gradients in place. nvte_scaled_softmax_backward(output_grads_cu.data(), softmax_results_cu.data(), - output_grads_cu.data(), scale_factor, + input_grads_cu.data(), scale_factor, at::cuda::getCurrentCUDAStream()); - return output_grads; + return input_grads; } at::Tensor scaled_upper_triang_masked_softmax_forward(at::Tensor input, float scale_factor) { @@ -167,15 +177,20 @@ at::Tensor scaled_upper_triang_masked_softmax_backward(at::Tensor output_grads_, TORCH_CHECK(output_grads.size(1) == output_grads.size(2)); + // Allocate a fresh output buffer so the op does not alias / mutate its + // inputs (required by `torch.library.custom_op`). + auto input_grads = + torch::empty(output_grads.sizes(), output_grads.options().requires_grad(false)); + auto output_grads_cu = makeTransformerEngineTensor(output_grads); auto softmax_results_cu = makeTransformerEngineTensor(softmax_results); + auto input_grads_cu = makeTransformerEngineTensor(input_grads); - // Produce gradients in place. nvte_scaled_upper_triang_masked_softmax_backward( - output_grads_cu.data(), softmax_results_cu.data(), output_grads_cu.data(), scale_factor, + output_grads_cu.data(), softmax_results_cu.data(), input_grads_cu.data(), scale_factor, at::cuda::getCurrentCUDAStream()); - return output_grads; + return input_grads; } at::Tensor scaled_aligned_causal_masked_softmax_forward(at::Tensor input, float scale_factor) { @@ -223,15 +238,20 @@ at::Tensor scaled_aligned_causal_masked_softmax_backward(at::Tensor output_grad_ (softmax_results.scalar_type() == at::ScalarType::BFloat16), "Only fp16 and bf16 are supported"); + // Allocate a fresh output buffer so the op does not alias / mutate its + // inputs (required by `torch.library.custom_op`). + auto input_grads = + torch::empty(output_grads.sizes(), output_grads.options().requires_grad(false)); + auto output_grads_cu = makeTransformerEngineTensor(output_grads); auto softmax_results_cu = makeTransformerEngineTensor(softmax_results); + auto input_grads_cu = makeTransformerEngineTensor(input_grads); - // Produce gradients in place. nvte_scaled_aligned_causal_masked_softmax_backward( - output_grads_cu.data(), softmax_results_cu.data(), output_grads_cu.data(), scale_factor, + output_grads_cu.data(), softmax_results_cu.data(), input_grads_cu.data(), scale_factor, at::cuda::getCurrentCUDAStream()); - return output_grads; + return input_grads; } } // namespace transformer_engine::pytorch From 51901503cdc545ba0268377f4902d67652a66eeb Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Wed, 8 Jul 2026 17:23:40 +0200 Subject: [PATCH 02/11] [PyTorch] torch.compile support for the FP8-emulation path of UnfusedDotProductAttention Make the FP8-emulation path (NVTE_UnfusedDPA_Emulate_FP8=1) of UnfusedDotProductAttention traceable by torch.compile(fullgraph=True). - backends.py: register the quantize+dequantize roundtrips used by FP8EmulationFunc as torch.library custom ops (te_fp8_emu::roundtrip_ and te_fp8_emu::roundtrip_qkv_) taking the quantizer as a value-opaque argument, with fake impls for tracing. Ops are registered only for the value-opaque quantizer classes (Float8CurrentScalingQuantizer, MXFP8Quantizer); Float8Quantizer (delayed scaling) carries scale/amax tensor state, is not value-opaque, and deliberately keeps the plain eager path -- FP8 emulation with delayed scaling is not supported under torch.compile. - backends.py: dispatch helpers `_fp8_emu_roundtrip{,_qkv}` key on `type(quantizer).__qualname__` so they stay traceable for opaque quantizer arguments; FP8EmulationFunc forward/backward now call them (onnx_forward unchanged). - backends.py: the joint q/k/v roundtrip clones any output whose storage is shared with an input or another output, checking storage identity directly -- the dequantized q/k/v can be views into one combined buffer, and view metadata (`_base`) is not populated under the torch-dispatch mode AOTAutograd runs custom ops with, so a `_base`-guarded clone triggered the custom-op aliasing deprecation warning under torch.compile. - UnfusedDotProductAttention.forward: only query FP8GlobalStateManager.get_fp8_recipe() when fp8_meta["local_recipes"] is absent. - test_torch_compile.py: add test_unfused_dpa_fp8_emulation_torch_compile (current scaling + mxfp8, sbhd/bshd layouts; compiled fullgraph forward+backward must match eager) and test_unfused_dpa_fp8_emulation_delayed_scaling_eager guarding the eager delayed-scaling path after the refactor. Signed-off-by: Pawel Gadzinski --- tests/pytorch/test_torch_compile.py | 126 ++++++++++++ .../dot_product_attention/backends.py | 179 ++++++++++++++---- 2 files changed, 268 insertions(+), 37 deletions(-) diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py index 0115370da1..31ea2fcb71 100644 --- a/tests/pytorch/test_torch_compile.py +++ b/tests/pytorch/test_torch_compile.py @@ -552,6 +552,132 @@ def fn(q, k, v, extra): assert v.grad is not None +def _make_unfused_fp8_emulation_quantizers(recipe_name: str): + """Build the ``quantizers`` dict + ``fp8_meta`` that + ``UnfusedDotProductAttention.forward(fp8=True, ...)`` expects, mimicking + what ``DotProductAttention.init_fp8_metadata`` produces for the recipe: + S/dP always get delayed-scaling quantizers (the forward reverts them to + current scaling when the recipe is CS).""" + from transformer_engine.pytorch.cpp_extensions.fused_attn import ( + META_QKV, + META_O, + META_S, + META_DQKV, + META_DO, + META_DP, + ) + from transformer_engine.pytorch.tensor.float8_tensor import Float8Quantizer + from transformer_engine.pytorch.tensor.mxfp8_tensor import MXFP8Quantizer + + def ds(dtype): + return Float8Quantizer( + scale=torch.ones(1, device="cuda"), + amax=torch.zeros(1, device="cuda"), + fp8_dtype=dtype, + ) + + def cs(dtype): + return Float8CurrentScalingQuantizer(fp8_dtype=dtype, device=torch.device("cuda")) + + def mx(dtype): + return MXFP8Quantizer(fp8_dtype=dtype) + + fwd_dtype, bwd_dtype = tex.DType.kFloat8E4M3, tex.DType.kFloat8E5M2 + if recipe_name == "delayed": + fp8_recipe = recipe.DelayedScaling() + fwd = {META_QKV: ds(fwd_dtype), META_O: ds(fwd_dtype), META_S: ds(fwd_dtype)} + bwd = {META_DQKV: ds(bwd_dtype), META_DO: ds(bwd_dtype), META_DP: ds(bwd_dtype)} + elif recipe_name == "current": + fp8_recipe = recipe.Float8CurrentScaling() + fwd = {META_QKV: cs(fwd_dtype), META_O: cs(fwd_dtype), META_S: ds(fwd_dtype)} + bwd = {META_DQKV: cs(bwd_dtype), META_DO: cs(bwd_dtype), META_DP: ds(bwd_dtype)} + elif recipe_name == "mxfp8": + fp8_recipe = recipe.MXFP8BlockScaling() + fwd = {META_QKV: mx(fwd_dtype), META_O: mx(fwd_dtype), META_S: mx(fwd_dtype)} + bwd = {META_DQKV: mx(bwd_dtype), META_DO: mx(bwd_dtype), META_DP: mx(bwd_dtype)} + else: + raise ValueError(recipe_name) + + quantizers = { + "scaling_fwd": [fwd.get(i) for i in range(max(fwd) + 1)], + "scaling_bwd": [bwd.get(i) for i in range(max(bwd) + 1)], + } + fp8_meta = {"local_recipes": [fp8_recipe]} + return quantizers, fp8_meta + + +def _run_unfused_fp8_emulation(qkv_layout: str, recipe_name: str, do_compile: bool): + """Run unfused DPA with FP8 emulation eagerly, and optionally compare + against torch.compile(fullgraph=True) on fresh leaf copies of the same + inputs.""" + dtype = torch.bfloat16 + module = _make_unfused_attention(dtype) + quantizers, fp8_meta = _make_unfused_fp8_emulation_quantizers(recipe_name) + + def fn(q, k, v, extra): + return module( + _EMPTY_ALIBI_CACHE, + q, + k, + v, + qkv_layout=qkv_layout, + attn_mask_type="causal", + fp8=True, + fp8_meta=fp8_meta, + quantizers=quantizers, + **extra, + ) + + q, k, v, extra, _ = _make_unfused_qkv(qkv_layout, dtype, requires_grad=True) + out_ref = fn(q, k, v, extra) + out_ref.sum().backward() + assert torch.isfinite(out_ref).all() + assert q.grad is not None and k.grad is not None and v.grad is not None + + if not do_compile: + return + + torch._dynamo.reset() + compiled = torch.compile(fn, fullgraph=True) + for _ in range(2): + q2, k2, v2 = (x.detach().clone().requires_grad_() for x in (q, k, v)) + out = compiled(q2, k2, v2, extra) + out.sum().backward() + torch.testing.assert_close(out, out_ref) + torch.testing.assert_close(q2.grad, q.grad) + torch.testing.assert_close(k2.grad, k.grad) + torch.testing.assert_close(v2.grad, v.grad) + + +@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) +@pytest.mark.parametrize("qkv_layout", ["sbhd_sbhd_sbhd", "bshd_bshd_bshd"]) +@pytest.mark.parametrize( + "recipe_name", + [ + pytest.param("current", id="float8_current_scaling"), + pytest.param( + "mxfp8", + id="mxfp8", + marks=pytest.mark.skipif(not mxfp8_available, reason=reason_for_no_mxfp8), + ), + ], +) +def test_unfused_dpa_fp8_emulation_torch_compile(recipe_name, qkv_layout): + """FP8-emulation path of UnfusedDotProductAttention under + torch.compile(fullgraph=True): the te_fp8_emu::* custom ops take the + value-opaque quantizers in-graph; compiled forward+backward must match + eager.""" + _run_unfused_fp8_emulation(qkv_layout, recipe_name, do_compile=True) + + +@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) +def test_unfused_dpa_fp8_emulation_delayed_scaling_eager(): + """Delayed scaling is not supported under torch.compile (Float8Quantizer + carries tensor state and is not value-opaque); guard that the eager + emulation path still works after the custom-op refactor.""" + _run_unfused_fp8_emulation("sbhd_sbhd_sbhd", "delayed", do_compile=False) + + # --------------------------------------------------------------------------- # Value-opaque quantizers # --------------------------------------------------------------------------- diff --git a/transformer_engine/pytorch/attention/dot_product_attention/backends.py b/transformer_engine/pytorch/attention/dot_product_attention/backends.py index ef718cee6c..e235332a0c 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/backends.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/backends.py @@ -198,6 +198,142 @@ def _qkv_quantizer_type(qkv_quantizer): raise TypeError(f"Unsupported FP8 attention QKV quantizer: {type(qkv_quantizer).__name__}") +# --------------------------------------------------------------------------- +# FP8 emulation roundtrips (quantize + dequantize) +# +# The roundtrips used by FP8EmulationFunc are registered as torch.library +# custom ops taking the quantizer as a value-opaque argument (see +# transformer_engine.pytorch.dynamo.quantizer_opaque), so that the FP8 +# emulation path of UnfusedDotProductAttention traces under torch.compile +# without graph breaks. Ops exist only for the value-opaque quantizer classes +# (Float8CurrentScalingQuantizer, MXFP8Quantizer); Float8Quantizer (delayed +# scaling) carries scale/amax tensor state, is not value-opaque, and keeps the +# plain eager path -- FP8 emulation with delayed scaling is not supported +# under torch.compile. +# --------------------------------------------------------------------------- + + +def _fp8_emu_roundtrip_eager(tensor: torch.Tensor, quantizer) -> torch.Tensor: + """Quantize + dequantize a single tensor to emulate FP8.""" + t_fp8 = quantizer(tensor) + return t_fp8.dequantize(dtype=tensor.dtype) + + +def _fp8_emu_roundtrip_qkv_eager(q, k, v, quantizer, qkv_layout): + """Quantize + dequantize q/k/v jointly (matching fused-attention packing) + to emulate FP8. Returns tensors in the same (sbhd) layout as the inputs.""" + q, k, v = q.contiguous(), k.contiguous(), v.contiguous() + q_fp8, k_fp8, v_fp8, new_qkv_layout, _ = combine_and_quantize( + qkv_layout, + q, + k, + v, + quantizer, + keep_same_data_and_scale_inv_format=True, + ) + tensors = combine_and_dequantize( + new_qkv_layout, q_fp8, k_fp8, v_fp8, src_nominal_dtype=q.dtype + ) + if isinstance(quantizer, MXFP8Quantizer): + # bhsd_bhsd_bhsd after combine_and_quantize; permute back to sbhd_sbhd_sbhd + tensors = [x.permute(2, 0, 1, 3).contiguous() for x in tensors] + # The dequantized q/k/v can be views into one combined buffer, but a custom + # op must not return tensors that alias its inputs or each other: clone any + # output whose storage was already seen. Storage identity is checked + # directly (rather than view metadata such as ``_base``, which is not + # populated under the torch-dispatch modes torch.compile runs the op with). + storages = {x.untyped_storage()._cdata for x in (q, k, v)} + outputs = [] + for x in tensors: + if x.untyped_storage()._cdata in storages: + x = x.clone() + storages.add(x.untyped_storage()._cdata) + outputs.append(x) + return outputs[0], outputs[1], outputs[2] + + +# quantizer class qualname -> (roundtrip op, roundtrip_qkv op) +_FP8_EMU_OPS: Dict[str, Tuple[Callable, Callable]] = {} +# Keep the Library object alive: its destructor would deregister the ops. +_FP8_EMU_LIB = None + + +def _register_fp8_emulation_ops() -> None: + """Register the te_fp8_emu::* custom ops for value-opaque quantizer classes. + + No-op on PyTorch builds without the opaque-object API: the eager fallback + in `_fp8_emu_roundtrip` / `_fp8_emu_roundtrip_qkv` keeps emulation working, + only torch.compile support is lost. + """ + global _FP8_EMU_LIB + try: + # pylint: disable=import-outside-toplevel + from torch._library.opaque_object import get_opaque_type_name + except (ImportError, AttributeError): + return + + lib = _FP8_EMU_LIB = torch.library.Library("te_fp8_emu", "DEF") + for qcls in (Float8CurrentScalingQuantizer, MXFP8Quantizer): + try: + opaque_name = get_opaque_type_name(qcls) + except Exception: # pylint: disable=broad-except + # Class not registered as an opaque type on this PyTorch build. + continue + + op_name = f"roundtrip_{qcls.__name__}" + qkv_op_name = f"roundtrip_qkv_{qcls.__name__}" + lib.define(f"{op_name}(Tensor x, {opaque_name} quantizer) -> Tensor") + lib.define( + f"{qkv_op_name}(Tensor q, Tensor k, Tensor v, {opaque_name} quantizer," + " str qkv_layout) -> (Tensor, Tensor, Tensor)" + ) + + torch.library.impl(f"te_fp8_emu::{op_name}", "CompositeExplicitAutograd", lib=lib)( + _fp8_emu_roundtrip_eager + ) + torch.library.register_fake(f"te_fp8_emu::{op_name}", lib=lib)( + lambda x, quantizer: torch.empty_like(x, memory_format=torch.contiguous_format) + ) + torch.library.impl(f"te_fp8_emu::{qkv_op_name}", "CompositeExplicitAutograd", lib=lib)( + _fp8_emu_roundtrip_qkv_eager + ) + torch.library.register_fake(f"te_fp8_emu::{qkv_op_name}", lib=lib)( + lambda q, k, v, quantizer, qkv_layout: tuple( + torch.empty_like(x, memory_format=torch.contiguous_format) for x in (q, k, v) + ) + ) + + _FP8_EMU_OPS[qcls.__qualname__] = ( + getattr(torch.ops.te_fp8_emu, op_name), + getattr(torch.ops.te_fp8_emu, qkv_op_name), + ) + + +_register_fp8_emulation_ops() + + +def _fp8_emu_roundtrip(tensor: torch.Tensor, quantizer) -> torch.Tensor: + """Single-tensor FP8 emulation roundtrip; uses the custom op when the + quantizer class has one (torch.compile-traceable), eager otherwise. + + Dispatch keys on ``type(quantizer).__qualname__`` rather than the class + itself so it stays traceable for opaque quantizer arguments (same trick + as ``is_value_opaque_quantizer``).""" + ops = _FP8_EMU_OPS.get(type(quantizer).__qualname__) + if ops is not None: + return ops[0](tensor, quantizer) + return _fp8_emu_roundtrip_eager(tensor, quantizer) + + +def _fp8_emu_roundtrip_qkv(q, k, v, quantizer, qkv_layout): + """Joint q/k/v FP8 emulation roundtrip; custom op when available, eager + otherwise.""" + ops = _FP8_EMU_OPS.get(type(quantizer).__qualname__) + if ops is not None: + return ops[1](q, k, v, quantizer, qkv_layout) + return _fp8_emu_roundtrip_qkv_eager(q, k, v, quantizer, qkv_layout) + + class FP8EmulationFunc(torch.autograd.Function): """ Emulate the effects of FP8 quantization on tensors. Used in UnfusedDotProductAttention as follows: @@ -214,28 +350,11 @@ def forward(ctx, tensor1, tensor2, tensor3, quantizer, quantizer_name, qkv_layou ) if quantizer_name == "QKV_quantizer": - query_layer, key_layer, value_layer = [ - x.contiguous() for x in [tensor1, tensor2, tensor3] - ] # always in sbhd_sbhd_sbhd shape at this point - q_fp8, k_fp8, v_fp8, qkv_layout, _ = combine_and_quantize( - qkv_layout, - query_layer, - key_layer, - value_layer, - quantizer, - keep_same_data_and_scale_inv_format=True, - ) - tensors = combine_and_dequantize( - qkv_layout, q_fp8, k_fp8, v_fp8, src_nominal_dtype=query_layer.dtype - ) - if isinstance(quantizer, MXFP8Quantizer): - # bhsd_bhsd_bhsd after combine_and_quantize; permute back to sbhd_sbhd_sbhd - tensors = [x.permute(2, 0, 1, 3).contiguous() for x in tensors] + tensors = _fp8_emu_roundtrip_qkv(tensor1, tensor2, tensor3, quantizer, qkv_layout) elif quantizer_name in ["S_quantizer", "O_quantizer"]: if quantizer is not None: - t_fp8 = quantizer(tensor1) - tensors = (t_fp8.dequantize(dtype=tensor1.dtype), tensor2, tensor3) + tensors = (_fp8_emu_roundtrip(tensor1, quantizer), tensor2, tensor3) else: tensors = (tensor1, tensor2, tensor3) else: @@ -250,27 +369,12 @@ def backward(ctx, grad1, grad2, grad3): # pylint: disable=missing-function-docstring if ctx.quantizer_name in ["dO_quantizer", "dP_quantizer"]: if ctx.quantizer is not None: - dt_fp8 = ctx.quantizer(grad1) - tensors = dt_fp8.dequantize(dtype=grad1.dtype), grad2, grad3 + tensors = _fp8_emu_roundtrip(grad1, ctx.quantizer), grad2, grad3 else: tensors = grad1, grad2, grad3 elif ctx.quantizer_name == "dQKV_quantizer": - query_grad, key_grad, value_grad = [x.contiguous() for x in [grad1, grad2, grad3]] # always in sbhd_sbhd_sbhd shape at this point - dq_fp8, dk_fp8, dv_fp8, new_qkv_layout, _ = combine_and_quantize( - ctx.qkv_layout, - query_grad, - key_grad, - value_grad, - ctx.quantizer, - keep_same_data_and_scale_inv_format=True, - ) - tensors = combine_and_dequantize( - new_qkv_layout, dq_fp8, dk_fp8, dv_fp8, src_nominal_dtype=query_grad.dtype - ) - if isinstance(ctx.quantizer, MXFP8Quantizer): - # bhsd_bhsd_bhsd after combine_and_quantize; permute back to sbhd_sbhd_sbhd - tensors = [x.permute(2, 0, 1, 3).contiguous() for x in tensors] + tensors = _fp8_emu_roundtrip_qkv(grad1, grad2, grad3, ctx.quantizer, ctx.qkv_layout) else: tensors = grad1, grad2, grad3 return tensors[0], tensors[1], tensors[2], None, None, None @@ -515,9 +619,10 @@ def forward( if fp8: # get fp8 recipe for DPA - fp8_recipe = FP8GlobalStateManager.get_fp8_recipe() if fp8_meta is not None and fp8_meta.get("local_recipes", None) is not None: fp8_recipe = fp8_meta["local_recipes"][0] + else: + fp8_recipe = FP8GlobalStateManager.get_fp8_recipe() # get quantizers from DPA; all Nones if not fp8 QKV_quantizer, O_quantizer, S_quantizer, dQKV_quantizer, dO_quantizer, dP_quantizer = ( dpa_utils.get_attention_quantizers(fp8, quantizers) From 15cbcd14b9406ef9174508d9bb5a4b523ff774bf Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Wed, 8 Jul 2026 17:45:53 +0200 Subject: [PATCH 03/11] [PyTorch] Run UnfusedDotProductAttention as an eager island when fp8_output=True With fp8_output=True the backend returns a Float8Tensor -- a tensor subclass that cannot cross a torch.compile graph boundary -- so the forward dispatches to a torch._dynamo.disable'd wrapper, the same mechanism DotProductAttention and FusedAttention use module-wide. With fp8_output=False the dispatcher is resolved at trace time and adds no graph break. Signed-off-by: Pawel Gadzinski --- .../attention/dot_product_attention/backends.py | 17 ++++++++++++++++- 1 file changed, 16 insertions(+), 1 deletion(-) diff --git a/transformer_engine/pytorch/attention/dot_product_attention/backends.py b/transformer_engine/pytorch/attention/dot_product_attention/backends.py index e235332a0c..1ab5cd4591 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/backends.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/backends.py @@ -469,7 +469,22 @@ def fast_setattr(self, name: str, value: Any) -> None: """Fast attribute set for non-parameter fields.""" self.__dict__[name] = value - def forward( + def forward(self, *args, fp8_output: bool = False, **kwargs) -> torch.Tensor: + """Unfused attention fprop; see `_forward` for the argument list. + + ``fp8_output=True`` returns a Float8Tensor, which cannot cross a + torch.compile graph boundary -- run the backend as an eager island. + """ + if fp8_output: + return self._forward_eager(*args, fp8_output=True, **kwargs) + return self._forward(*args, fp8_output=False, **kwargs) + + @no_torch_dynamo() + def _forward_eager(self, *args, **kwargs) -> torch.Tensor: + """Eager-only (dynamo-disabled) wrapper around `_forward`.""" + return self._forward(*args, **kwargs) + + def _forward( self, _alibi_cache: Dict[str, Any], query_layer: torch.Tensor, From 8bd9881dbb6519ed47d388ab1ecb1ac3b1c0481b Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Wed, 8 Jul 2026 17:49:59 +0200 Subject: [PATCH 04/11] [PyTorch] Test FP8-emulation compile path also with reduce-overhead (cudagraphs) Parametrize test_unfused_dpa_fp8_emulation_torch_compile over compile mode (default, reduce-overhead), run 3 iterations so the CUDA graphs are recorded and replayed. The te_fp8_emu roundtrip ops for current scaling are pure (no mutated args), so inductor cudagraphs capture them; verified no cudagraph skips with TORCH_LOGS=cudagraphs. Signed-off-by: Pawel Gadzinski --- tests/pytorch/test_torch_compile.py | 21 ++++++++++++++------- 1 file changed, 14 insertions(+), 7 deletions(-) diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py index 31ea2fcb71..71943d3a95 100644 --- a/tests/pytorch/test_torch_compile.py +++ b/tests/pytorch/test_torch_compile.py @@ -606,10 +606,13 @@ def mx(dtype): return quantizers, fp8_meta -def _run_unfused_fp8_emulation(qkv_layout: str, recipe_name: str, do_compile: bool): +def _run_unfused_fp8_emulation( + qkv_layout: str, recipe_name: str, do_compile: bool, compile_mode: str = "default" +): """Run unfused DPA with FP8 emulation eagerly, and optionally compare against torch.compile(fullgraph=True) on fresh leaf copies of the same - inputs.""" + inputs. ``compile_mode="reduce-overhead"`` additionally captures the + compiled forward+backward into CUDA graphs.""" dtype = torch.bfloat16 module = _make_unfused_attention(dtype) quantizers, fp8_meta = _make_unfused_fp8_emulation_quantizers(recipe_name) @@ -638,11 +641,13 @@ def fn(q, k, v, extra): return torch._dynamo.reset() - compiled = torch.compile(fn, fullgraph=True) - for _ in range(2): + mode = None if compile_mode == "default" else compile_mode + compiled = torch.compile(fn, fullgraph=True, mode=mode) + for _ in range(3): q2, k2, v2 = (x.detach().clone().requires_grad_() for x in (q, k, v)) out = compiled(q2, k2, v2, extra) out.sum().backward() + torch.cuda.synchronize() torch.testing.assert_close(out, out_ref) torch.testing.assert_close(q2.grad, q.grad) torch.testing.assert_close(k2.grad, k.grad) @@ -662,12 +667,14 @@ def fn(q, k, v, extra): ), ], ) -def test_unfused_dpa_fp8_emulation_torch_compile(recipe_name, qkv_layout): +@pytest.mark.parametrize("compile_mode", ["default", "reduce-overhead"]) +def test_unfused_dpa_fp8_emulation_torch_compile(recipe_name, qkv_layout, compile_mode): """FP8-emulation path of UnfusedDotProductAttention under torch.compile(fullgraph=True): the te_fp8_emu::* custom ops take the value-opaque quantizers in-graph; compiled forward+backward must match - eager.""" - _run_unfused_fp8_emulation(qkv_layout, recipe_name, do_compile=True) + eager. With "reduce-overhead" the graphs are captured and replayed via + inductor cudagraphs.""" + _run_unfused_fp8_emulation(qkv_layout, recipe_name, do_compile=True, compile_mode=compile_mode) @pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) From ac1f2b0303dab8897ab0f7eaf6f6054a148dad5b Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Thu, 9 Jul 2026 13:57:58 +0200 Subject: [PATCH 05/11] [PyTorch] Drop FP8 torch.compile support in UnfusedDotProductAttention; run FP8 as an eager island FP8 in the unfused backend (emulation and Float8Tensor output) is not supported under torch.compile: the forward dispatcher routes fp8=True and/or fp8_output=True to a torch._dynamo.disable'd wrapper, same as DotProductAttention does module-wide. Remove the FP8-emulation compile tests. The te_fp8_emu::* custom ops taking value-opaque quantizers stay as the eager implementation of FP8EmulationFunc. Signed-off-by: Pawel Gadzinski --- tests/pytorch/test_torch_compile.py | 133 ------------------ .../dot_product_attention/backends.py | 15 +- 2 files changed, 7 insertions(+), 141 deletions(-) diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py index 71943d3a95..0115370da1 100644 --- a/tests/pytorch/test_torch_compile.py +++ b/tests/pytorch/test_torch_compile.py @@ -552,139 +552,6 @@ def fn(q, k, v, extra): assert v.grad is not None -def _make_unfused_fp8_emulation_quantizers(recipe_name: str): - """Build the ``quantizers`` dict + ``fp8_meta`` that - ``UnfusedDotProductAttention.forward(fp8=True, ...)`` expects, mimicking - what ``DotProductAttention.init_fp8_metadata`` produces for the recipe: - S/dP always get delayed-scaling quantizers (the forward reverts them to - current scaling when the recipe is CS).""" - from transformer_engine.pytorch.cpp_extensions.fused_attn import ( - META_QKV, - META_O, - META_S, - META_DQKV, - META_DO, - META_DP, - ) - from transformer_engine.pytorch.tensor.float8_tensor import Float8Quantizer - from transformer_engine.pytorch.tensor.mxfp8_tensor import MXFP8Quantizer - - def ds(dtype): - return Float8Quantizer( - scale=torch.ones(1, device="cuda"), - amax=torch.zeros(1, device="cuda"), - fp8_dtype=dtype, - ) - - def cs(dtype): - return Float8CurrentScalingQuantizer(fp8_dtype=dtype, device=torch.device("cuda")) - - def mx(dtype): - return MXFP8Quantizer(fp8_dtype=dtype) - - fwd_dtype, bwd_dtype = tex.DType.kFloat8E4M3, tex.DType.kFloat8E5M2 - if recipe_name == "delayed": - fp8_recipe = recipe.DelayedScaling() - fwd = {META_QKV: ds(fwd_dtype), META_O: ds(fwd_dtype), META_S: ds(fwd_dtype)} - bwd = {META_DQKV: ds(bwd_dtype), META_DO: ds(bwd_dtype), META_DP: ds(bwd_dtype)} - elif recipe_name == "current": - fp8_recipe = recipe.Float8CurrentScaling() - fwd = {META_QKV: cs(fwd_dtype), META_O: cs(fwd_dtype), META_S: ds(fwd_dtype)} - bwd = {META_DQKV: cs(bwd_dtype), META_DO: cs(bwd_dtype), META_DP: ds(bwd_dtype)} - elif recipe_name == "mxfp8": - fp8_recipe = recipe.MXFP8BlockScaling() - fwd = {META_QKV: mx(fwd_dtype), META_O: mx(fwd_dtype), META_S: mx(fwd_dtype)} - bwd = {META_DQKV: mx(bwd_dtype), META_DO: mx(bwd_dtype), META_DP: mx(bwd_dtype)} - else: - raise ValueError(recipe_name) - - quantizers = { - "scaling_fwd": [fwd.get(i) for i in range(max(fwd) + 1)], - "scaling_bwd": [bwd.get(i) for i in range(max(bwd) + 1)], - } - fp8_meta = {"local_recipes": [fp8_recipe]} - return quantizers, fp8_meta - - -def _run_unfused_fp8_emulation( - qkv_layout: str, recipe_name: str, do_compile: bool, compile_mode: str = "default" -): - """Run unfused DPA with FP8 emulation eagerly, and optionally compare - against torch.compile(fullgraph=True) on fresh leaf copies of the same - inputs. ``compile_mode="reduce-overhead"`` additionally captures the - compiled forward+backward into CUDA graphs.""" - dtype = torch.bfloat16 - module = _make_unfused_attention(dtype) - quantizers, fp8_meta = _make_unfused_fp8_emulation_quantizers(recipe_name) - - def fn(q, k, v, extra): - return module( - _EMPTY_ALIBI_CACHE, - q, - k, - v, - qkv_layout=qkv_layout, - attn_mask_type="causal", - fp8=True, - fp8_meta=fp8_meta, - quantizers=quantizers, - **extra, - ) - - q, k, v, extra, _ = _make_unfused_qkv(qkv_layout, dtype, requires_grad=True) - out_ref = fn(q, k, v, extra) - out_ref.sum().backward() - assert torch.isfinite(out_ref).all() - assert q.grad is not None and k.grad is not None and v.grad is not None - - if not do_compile: - return - - torch._dynamo.reset() - mode = None if compile_mode == "default" else compile_mode - compiled = torch.compile(fn, fullgraph=True, mode=mode) - for _ in range(3): - q2, k2, v2 = (x.detach().clone().requires_grad_() for x in (q, k, v)) - out = compiled(q2, k2, v2, extra) - out.sum().backward() - torch.cuda.synchronize() - torch.testing.assert_close(out, out_ref) - torch.testing.assert_close(q2.grad, q.grad) - torch.testing.assert_close(k2.grad, k.grad) - torch.testing.assert_close(v2.grad, v.grad) - - -@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) -@pytest.mark.parametrize("qkv_layout", ["sbhd_sbhd_sbhd", "bshd_bshd_bshd"]) -@pytest.mark.parametrize( - "recipe_name", - [ - pytest.param("current", id="float8_current_scaling"), - pytest.param( - "mxfp8", - id="mxfp8", - marks=pytest.mark.skipif(not mxfp8_available, reason=reason_for_no_mxfp8), - ), - ], -) -@pytest.mark.parametrize("compile_mode", ["default", "reduce-overhead"]) -def test_unfused_dpa_fp8_emulation_torch_compile(recipe_name, qkv_layout, compile_mode): - """FP8-emulation path of UnfusedDotProductAttention under - torch.compile(fullgraph=True): the te_fp8_emu::* custom ops take the - value-opaque quantizers in-graph; compiled forward+backward must match - eager. With "reduce-overhead" the graphs are captured and replayed via - inductor cudagraphs.""" - _run_unfused_fp8_emulation(qkv_layout, recipe_name, do_compile=True, compile_mode=compile_mode) - - -@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) -def test_unfused_dpa_fp8_emulation_delayed_scaling_eager(): - """Delayed scaling is not supported under torch.compile (Float8Quantizer - carries tensor state and is not value-opaque); guard that the eager - emulation path still works after the custom-op refactor.""" - _run_unfused_fp8_emulation("sbhd_sbhd_sbhd", "delayed", do_compile=False) - - # --------------------------------------------------------------------------- # Value-opaque quantizers # --------------------------------------------------------------------------- diff --git a/transformer_engine/pytorch/attention/dot_product_attention/backends.py b/transformer_engine/pytorch/attention/dot_product_attention/backends.py index 1ab5cd4591..16f52b8da7 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/backends.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/backends.py @@ -469,15 +469,15 @@ def fast_setattr(self, name: str, value: Any) -> None: """Fast attribute set for non-parameter fields.""" self.__dict__[name] = value - def forward(self, *args, fp8_output: bool = False, **kwargs) -> torch.Tensor: + def forward(self, *args, fp8: bool = False, fp8_output: bool = False, **kwargs) -> torch.Tensor: """Unfused attention fprop; see `_forward` for the argument list. - ``fp8_output=True`` returns a Float8Tensor, which cannot cross a - torch.compile graph boundary -- run the backend as an eager island. + FP8 (emulation and/or Float8Tensor output) is not supported under + torch.compile -- run the backend as an eager island in that case. """ - if fp8_output: - return self._forward_eager(*args, fp8_output=True, **kwargs) - return self._forward(*args, fp8_output=False, **kwargs) + if fp8 or fp8_output: + return self._forward_eager(*args, fp8=fp8, fp8_output=fp8_output, **kwargs) + return self._forward(*args, fp8=False, fp8_output=False, **kwargs) @no_torch_dynamo() def _forward_eager(self, *args, **kwargs) -> torch.Tensor: @@ -634,10 +634,9 @@ def _forward( if fp8: # get fp8 recipe for DPA + fp8_recipe = FP8GlobalStateManager.get_fp8_recipe() if fp8_meta is not None and fp8_meta.get("local_recipes", None) is not None: fp8_recipe = fp8_meta["local_recipes"][0] - else: - fp8_recipe = FP8GlobalStateManager.get_fp8_recipe() # get quantizers from DPA; all Nones if not fp8 QKV_quantizer, O_quantizer, S_quantizer, dQKV_quantizer, dO_quantizer, dP_quantizer = ( dpa_utils.get_attention_quantizers(fp8, quantizers) From fab4dd76810fb23df343d5362b0235a4c76c9487 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Thu, 9 Jul 2026 14:12:00 +0200 Subject: [PATCH 06/11] [PyTorch] Remove te_fp8_emu custom ops; restore upstream FP8EmulationFunc The ops existed solely to make the FP8-emulation path traceable by torch.compile; since FP8 in the unfused backend now always runs as an eager island, they are dead machinery (plus import-time registration and output clones the plain eager path never needed). Signed-off-by: Pawel Gadzinski --- .../dot_product_attention/backends.py | 176 ++++-------------- 1 file changed, 36 insertions(+), 140 deletions(-) diff --git a/transformer_engine/pytorch/attention/dot_product_attention/backends.py b/transformer_engine/pytorch/attention/dot_product_attention/backends.py index 16f52b8da7..e0401f2ebf 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/backends.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/backends.py @@ -198,142 +198,6 @@ def _qkv_quantizer_type(qkv_quantizer): raise TypeError(f"Unsupported FP8 attention QKV quantizer: {type(qkv_quantizer).__name__}") -# --------------------------------------------------------------------------- -# FP8 emulation roundtrips (quantize + dequantize) -# -# The roundtrips used by FP8EmulationFunc are registered as torch.library -# custom ops taking the quantizer as a value-opaque argument (see -# transformer_engine.pytorch.dynamo.quantizer_opaque), so that the FP8 -# emulation path of UnfusedDotProductAttention traces under torch.compile -# without graph breaks. Ops exist only for the value-opaque quantizer classes -# (Float8CurrentScalingQuantizer, MXFP8Quantizer); Float8Quantizer (delayed -# scaling) carries scale/amax tensor state, is not value-opaque, and keeps the -# plain eager path -- FP8 emulation with delayed scaling is not supported -# under torch.compile. -# --------------------------------------------------------------------------- - - -def _fp8_emu_roundtrip_eager(tensor: torch.Tensor, quantizer) -> torch.Tensor: - """Quantize + dequantize a single tensor to emulate FP8.""" - t_fp8 = quantizer(tensor) - return t_fp8.dequantize(dtype=tensor.dtype) - - -def _fp8_emu_roundtrip_qkv_eager(q, k, v, quantizer, qkv_layout): - """Quantize + dequantize q/k/v jointly (matching fused-attention packing) - to emulate FP8. Returns tensors in the same (sbhd) layout as the inputs.""" - q, k, v = q.contiguous(), k.contiguous(), v.contiguous() - q_fp8, k_fp8, v_fp8, new_qkv_layout, _ = combine_and_quantize( - qkv_layout, - q, - k, - v, - quantizer, - keep_same_data_and_scale_inv_format=True, - ) - tensors = combine_and_dequantize( - new_qkv_layout, q_fp8, k_fp8, v_fp8, src_nominal_dtype=q.dtype - ) - if isinstance(quantizer, MXFP8Quantizer): - # bhsd_bhsd_bhsd after combine_and_quantize; permute back to sbhd_sbhd_sbhd - tensors = [x.permute(2, 0, 1, 3).contiguous() for x in tensors] - # The dequantized q/k/v can be views into one combined buffer, but a custom - # op must not return tensors that alias its inputs or each other: clone any - # output whose storage was already seen. Storage identity is checked - # directly (rather than view metadata such as ``_base``, which is not - # populated under the torch-dispatch modes torch.compile runs the op with). - storages = {x.untyped_storage()._cdata for x in (q, k, v)} - outputs = [] - for x in tensors: - if x.untyped_storage()._cdata in storages: - x = x.clone() - storages.add(x.untyped_storage()._cdata) - outputs.append(x) - return outputs[0], outputs[1], outputs[2] - - -# quantizer class qualname -> (roundtrip op, roundtrip_qkv op) -_FP8_EMU_OPS: Dict[str, Tuple[Callable, Callable]] = {} -# Keep the Library object alive: its destructor would deregister the ops. -_FP8_EMU_LIB = None - - -def _register_fp8_emulation_ops() -> None: - """Register the te_fp8_emu::* custom ops for value-opaque quantizer classes. - - No-op on PyTorch builds without the opaque-object API: the eager fallback - in `_fp8_emu_roundtrip` / `_fp8_emu_roundtrip_qkv` keeps emulation working, - only torch.compile support is lost. - """ - global _FP8_EMU_LIB - try: - # pylint: disable=import-outside-toplevel - from torch._library.opaque_object import get_opaque_type_name - except (ImportError, AttributeError): - return - - lib = _FP8_EMU_LIB = torch.library.Library("te_fp8_emu", "DEF") - for qcls in (Float8CurrentScalingQuantizer, MXFP8Quantizer): - try: - opaque_name = get_opaque_type_name(qcls) - except Exception: # pylint: disable=broad-except - # Class not registered as an opaque type on this PyTorch build. - continue - - op_name = f"roundtrip_{qcls.__name__}" - qkv_op_name = f"roundtrip_qkv_{qcls.__name__}" - lib.define(f"{op_name}(Tensor x, {opaque_name} quantizer) -> Tensor") - lib.define( - f"{qkv_op_name}(Tensor q, Tensor k, Tensor v, {opaque_name} quantizer," - " str qkv_layout) -> (Tensor, Tensor, Tensor)" - ) - - torch.library.impl(f"te_fp8_emu::{op_name}", "CompositeExplicitAutograd", lib=lib)( - _fp8_emu_roundtrip_eager - ) - torch.library.register_fake(f"te_fp8_emu::{op_name}", lib=lib)( - lambda x, quantizer: torch.empty_like(x, memory_format=torch.contiguous_format) - ) - torch.library.impl(f"te_fp8_emu::{qkv_op_name}", "CompositeExplicitAutograd", lib=lib)( - _fp8_emu_roundtrip_qkv_eager - ) - torch.library.register_fake(f"te_fp8_emu::{qkv_op_name}", lib=lib)( - lambda q, k, v, quantizer, qkv_layout: tuple( - torch.empty_like(x, memory_format=torch.contiguous_format) for x in (q, k, v) - ) - ) - - _FP8_EMU_OPS[qcls.__qualname__] = ( - getattr(torch.ops.te_fp8_emu, op_name), - getattr(torch.ops.te_fp8_emu, qkv_op_name), - ) - - -_register_fp8_emulation_ops() - - -def _fp8_emu_roundtrip(tensor: torch.Tensor, quantizer) -> torch.Tensor: - """Single-tensor FP8 emulation roundtrip; uses the custom op when the - quantizer class has one (torch.compile-traceable), eager otherwise. - - Dispatch keys on ``type(quantizer).__qualname__`` rather than the class - itself so it stays traceable for opaque quantizer arguments (same trick - as ``is_value_opaque_quantizer``).""" - ops = _FP8_EMU_OPS.get(type(quantizer).__qualname__) - if ops is not None: - return ops[0](tensor, quantizer) - return _fp8_emu_roundtrip_eager(tensor, quantizer) - - -def _fp8_emu_roundtrip_qkv(q, k, v, quantizer, qkv_layout): - """Joint q/k/v FP8 emulation roundtrip; custom op when available, eager - otherwise.""" - ops = _FP8_EMU_OPS.get(type(quantizer).__qualname__) - if ops is not None: - return ops[1](q, k, v, quantizer, qkv_layout) - return _fp8_emu_roundtrip_qkv_eager(q, k, v, quantizer, qkv_layout) - - class FP8EmulationFunc(torch.autograd.Function): """ Emulate the effects of FP8 quantization on tensors. Used in UnfusedDotProductAttention as follows: @@ -350,11 +214,28 @@ def forward(ctx, tensor1, tensor2, tensor3, quantizer, quantizer_name, qkv_layou ) if quantizer_name == "QKV_quantizer": + query_layer, key_layer, value_layer = [ + x.contiguous() for x in [tensor1, tensor2, tensor3] + ] # always in sbhd_sbhd_sbhd shape at this point - tensors = _fp8_emu_roundtrip_qkv(tensor1, tensor2, tensor3, quantizer, qkv_layout) + q_fp8, k_fp8, v_fp8, qkv_layout, _ = combine_and_quantize( + qkv_layout, + query_layer, + key_layer, + value_layer, + quantizer, + keep_same_data_and_scale_inv_format=True, + ) + tensors = combine_and_dequantize( + qkv_layout, q_fp8, k_fp8, v_fp8, src_nominal_dtype=query_layer.dtype + ) + if isinstance(quantizer, MXFP8Quantizer): + # bhsd_bhsd_bhsd after combine_and_quantize; permute back to sbhd_sbhd_sbhd + tensors = [x.permute(2, 0, 1, 3).contiguous() for x in tensors] elif quantizer_name in ["S_quantizer", "O_quantizer"]: if quantizer is not None: - tensors = (_fp8_emu_roundtrip(tensor1, quantizer), tensor2, tensor3) + t_fp8 = quantizer(tensor1) + tensors = (t_fp8.dequantize(dtype=tensor1.dtype), tensor2, tensor3) else: tensors = (tensor1, tensor2, tensor3) else: @@ -369,12 +250,27 @@ def backward(ctx, grad1, grad2, grad3): # pylint: disable=missing-function-docstring if ctx.quantizer_name in ["dO_quantizer", "dP_quantizer"]: if ctx.quantizer is not None: - tensors = _fp8_emu_roundtrip(grad1, ctx.quantizer), grad2, grad3 + dt_fp8 = ctx.quantizer(grad1) + tensors = dt_fp8.dequantize(dtype=grad1.dtype), grad2, grad3 else: tensors = grad1, grad2, grad3 elif ctx.quantizer_name == "dQKV_quantizer": + query_grad, key_grad, value_grad = [x.contiguous() for x in [grad1, grad2, grad3]] # always in sbhd_sbhd_sbhd shape at this point - tensors = _fp8_emu_roundtrip_qkv(grad1, grad2, grad3, ctx.quantizer, ctx.qkv_layout) + dq_fp8, dk_fp8, dv_fp8, new_qkv_layout, _ = combine_and_quantize( + ctx.qkv_layout, + query_grad, + key_grad, + value_grad, + ctx.quantizer, + keep_same_data_and_scale_inv_format=True, + ) + tensors = combine_and_dequantize( + new_qkv_layout, dq_fp8, dk_fp8, dv_fp8, src_nominal_dtype=query_grad.dtype + ) + if isinstance(ctx.quantizer, MXFP8Quantizer): + # bhsd_bhsd_bhsd after combine_and_quantize; permute back to sbhd_sbhd_sbhd + tensors = [x.permute(2, 0, 1, 3).contiguous() for x in tensors] else: tensors = grad1, grad2, grad3 return tensors[0], tensors[1], tensors[2], None, None, None From b310d53eec8cf0c4018b0c4e2ce67861bb6c19f8 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Thu, 9 Jul 2026 14:17:09 +0200 Subject: [PATCH 07/11] [PyTorch] Shorten the total_tokens_q comment Signed-off-by: Pawel Gadzinski --- .../pytorch/attention/dot_product_attention/backends.py | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/transformer_engine/pytorch/attention/dot_product_attention/backends.py b/transformer_engine/pytorch/attention/dot_product_attention/backends.py index e0401f2ebf..1cd875ab19 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/backends.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/backends.py @@ -440,11 +440,9 @@ def _forward( if qkv_format == "thd": assert cu_seqlens_q is not None and cu_seqlens_kv is not None assert max_seqlen_q is not None and max_seqlen_kv is not None - # Capture token counts as plain (Sym)Ints from the THD inputs *before* - # we overwrite `query_layer` with the BSHD form. We thread `total_tokens_q` - # back into `ConvertBSHDtoTHD.apply` below so that it does not need to - # call `cu_seqlens_q[-1].item()` itself -- which under `torch.compile` - # would create an unbacked SymInt and break Inductor + cudagraphs. + # Token count for ConvertBSHDtoTHD below; deriving it there via + # cu_seqlens_q[-1].item() syncs with the GPU and breaks + # torch.compile+cudagraphs (unbacked SymInt). total_tokens_q = query_layer.shape[0] query_layer = ConvertTHDtoBSHD.apply( query_layer, From 1c7a5a8e5f104cea3aebaa1c071d8c685872927b Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Thu, 9 Jul 2026 14:23:40 +0200 Subject: [PATCH 08/11] [PyTorch] Drop _scale_to_tensor; pass the softmax scale as a plain float The tex softmax kernels take 'float scale_factor' directly. The 0-D tensor wrapping was a leftover of the old autograd.Function idiom, where the float had to be a tensor only to fit save_for_backward; the custom ops keep the scale on ctx as a plain attribute. Signed-off-by: Pawel Gadzinski --- .../dot_product_attention/softmax.py | 21 +++++++------------ 1 file changed, 8 insertions(+), 13 deletions(-) diff --git a/transformer_engine/pytorch/attention/dot_product_attention/softmax.py b/transformer_engine/pytorch/attention/dot_product_attention/softmax.py index 1207b67985..3b75df2f64 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/softmax.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/softmax.py @@ -18,18 +18,13 @@ _default_causal_mask = {} -def _scale_to_tensor(scale: float) -> torch.Tensor: - """Wrap a Python float in a 0-D tensor expected by the tex kernels.""" - return torch.tensor([scale])[0] - - # ----------------------------- ScaledSoftmax ------------------------------- @torch.library.custom_op("te_softmax::scaled_softmax_fwd", mutates_args=()) def scaled_softmax_forward(inputs: torch.Tensor, scale: float) -> torch.Tensor: """Forward pass for ScaledSoftmax.""" - return tex.scaled_softmax_forward(inputs, _scale_to_tensor(scale)) + return tex.scaled_softmax_forward(inputs, scale) @scaled_softmax_forward.register_fake @@ -43,7 +38,7 @@ def scaled_softmax_backward( output_grads: torch.Tensor, softmax_results: torch.Tensor, scale: float ) -> torch.Tensor: """Backward pass for ScaledSoftmax.""" - return tex.scaled_softmax_backward(output_grads, softmax_results, _scale_to_tensor(scale)) + return tex.scaled_softmax_backward(output_grads, softmax_results, scale) @scaled_softmax_backward.register_fake @@ -80,7 +75,7 @@ def scaled_masked_softmax_forward( inputs: torch.Tensor, mask: torch.Tensor, scale: float ) -> torch.Tensor: """Forward pass for ScaledMaskedSoftmax.""" - return tex.scaled_masked_softmax_forward(inputs, mask, _scale_to_tensor(scale)) + return tex.scaled_masked_softmax_forward(inputs, mask, scale) @scaled_masked_softmax_forward.register_fake @@ -97,7 +92,7 @@ def scaled_masked_softmax_backward( ) -> torch.Tensor: """Backward pass for ScaledMaskedSoftmax.""" return tex.scaled_masked_softmax_backward( - output_grads, softmax_results, _scale_to_tensor(scale) + output_grads, softmax_results, scale ) @@ -137,7 +132,7 @@ def scaled_upper_triang_masked_softmax_forward( inputs: torch.Tensor, scale: float ) -> torch.Tensor: """Forward pass for ScaledUpperTriangMaskedSoftmax.""" - return tex.scaled_upper_triang_masked_softmax_forward(inputs, _scale_to_tensor(scale)) + return tex.scaled_upper_triang_masked_softmax_forward(inputs, scale) @scaled_upper_triang_masked_softmax_forward.register_fake @@ -154,7 +149,7 @@ def scaled_upper_triang_masked_softmax_backward( ) -> torch.Tensor: """Backward pass for ScaledUpperTriangMaskedSoftmax.""" return tex.scaled_upper_triang_masked_softmax_backward( - output_grads, softmax_results, _scale_to_tensor(scale) + output_grads, softmax_results, scale ) @@ -194,7 +189,7 @@ def scaled_aligned_causal_masked_softmax_forward( inputs: torch.Tensor, scale: float ) -> torch.Tensor: """Forward pass for ScaledAlignedCausalMaskedSoftmax.""" - return tex.scaled_aligned_causal_masked_softmax_forward(inputs, _scale_to_tensor(scale)) + return tex.scaled_aligned_causal_masked_softmax_forward(inputs, scale) @scaled_aligned_causal_masked_softmax_forward.register_fake @@ -211,7 +206,7 @@ def scaled_aligned_causal_masked_softmax_backward( ) -> torch.Tensor: """Backward pass for ScaledAlignedCausalMaskedSoftmax.""" return tex.scaled_aligned_causal_masked_softmax_backward( - output_grads, softmax_results, _scale_to_tensor(scale) + output_grads, softmax_results, scale ) From c5a2dff80893b4828ef2867e2f503e477632f5a5 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Thu, 9 Jul 2026 14:28:39 +0200 Subject: [PATCH 09/11] [PyTorch] Drop redundant num_tokens comment (rationale documented at the callsite) Signed-off-by: Pawel Gadzinski --- .../pytorch/attention/dot_product_attention/utils.py | 6 ------ 1 file changed, 6 deletions(-) diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py index 747728a3ef..883af16d11 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py @@ -2271,12 +2271,6 @@ class ConvertBSHDtoTHD: @staticmethod def apply(bshd_tensor, cu_seqlens, num_tokens): # pylint: disable=missing-function-docstring - # `num_tokens` (== `cu_seqlens[-1]`) is taken as an explicit argument - # rather than being read from `cu_seqlens[-1].item()` here, so that - # `torch.compile` does not introduce an unbacked SymInt for it. With an - # unbacked SymInt the Inductor partitioner emits `None` placeholders for - # output buffers, which `cudagraph_trees` then refuses (only - # `Tensor / int / Generator` allowed as graph inputs). return torch.ops.te_attention.convert_bshd_to_thd( bshd_tensor, cu_seqlens, num_tokens ) From 2de9691fa83c0b9fd56e1c0fe11dede6a76358ff Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Thu, 9 Jul 2026 14:35:42 +0200 Subject: [PATCH 10/11] [PyTorch] Review cleanups: black formatting, drop unused import/duplicate dict, silence W0613 - run black over the four changed files (earlier commits skipped pre-commit) - drop unused 'import os' in test_torch_compile.py - drop duplicated module-level _default_causal_mask dict in softmax.py - del unused 'output' arg in the conversion setup_context helpers Signed-off-by: Pawel Gadzinski --- tests/pytorch/test_torch_compile.py | 2 -- .../dot_product_attention/backends.py | 1 + .../dot_product_attention/softmax.py | 21 +++++-------------- .../attention/dot_product_attention/utils.py | 15 ++++++------- 4 files changed, 12 insertions(+), 27 deletions(-) diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py index 0115370da1..a0b959dcde 100644 --- a/tests/pytorch/test_torch_compile.py +++ b/tests/pytorch/test_torch_compile.py @@ -3,7 +3,6 @@ # See LICENSE for license information. import abc -import os import pytest import torch @@ -393,7 +392,6 @@ def fn(inp): out.sum().backward() - _UNFUSED_DPA_CONFIG = dict( batch_size=2, num_heads=4, diff --git a/transformer_engine/pytorch/attention/dot_product_attention/backends.py b/transformer_engine/pytorch/attention/dot_product_attention/backends.py index 1cd875ab19..16a78ebef8 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/backends.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/backends.py @@ -3,6 +3,7 @@ # See LICENSE for license information. """Attention Backends.""" + from contextlib import nullcontext from importlib.metadata import version as get_pkg_version from importlib.metadata import PackageNotFoundError diff --git a/transformer_engine/pytorch/attention/dot_product_attention/softmax.py b/transformer_engine/pytorch/attention/dot_product_attention/softmax.py index 3b75df2f64..6e8a14402f 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/softmax.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/softmax.py @@ -3,6 +3,7 @@ # See LICENSE for license information. """Fused scaled masked softmax functions""" + import os from typing import Callable, Optional import torch @@ -10,14 +11,10 @@ import transformer_engine_torch as tex from transformer_engine.pytorch.export import is_in_onnx_export_mode - THREADS_PER_WARP = 32 THREADS_PER_BLOCK = 128 -_default_causal_mask = {} - - # ----------------------------- ScaledSoftmax ------------------------------- @@ -91,9 +88,7 @@ def scaled_masked_softmax_backward( output_grads: torch.Tensor, softmax_results: torch.Tensor, scale: float ) -> torch.Tensor: """Backward pass for ScaledMaskedSoftmax.""" - return tex.scaled_masked_softmax_backward( - output_grads, softmax_results, scale - ) + return tex.scaled_masked_softmax_backward(output_grads, softmax_results, scale) @scaled_masked_softmax_backward.register_fake @@ -128,9 +123,7 @@ def _scaled_masked_softmax_backward_wrapper(ctx, grad_output): @torch.library.custom_op("te_softmax::scaled_upper_triang_masked_softmax_fwd", mutates_args=()) -def scaled_upper_triang_masked_softmax_forward( - inputs: torch.Tensor, scale: float -) -> torch.Tensor: +def scaled_upper_triang_masked_softmax_forward(inputs: torch.Tensor, scale: float) -> torch.Tensor: """Forward pass for ScaledUpperTriangMaskedSoftmax.""" return tex.scaled_upper_triang_masked_softmax_forward(inputs, scale) @@ -148,9 +141,7 @@ def scaled_upper_triang_masked_softmax_backward( output_grads: torch.Tensor, softmax_results: torch.Tensor, scale: float ) -> torch.Tensor: """Backward pass for ScaledUpperTriangMaskedSoftmax.""" - return tex.scaled_upper_triang_masked_softmax_backward( - output_grads, softmax_results, scale - ) + return tex.scaled_upper_triang_masked_softmax_backward(output_grads, softmax_results, scale) @scaled_upper_triang_masked_softmax_backward.register_fake @@ -205,9 +196,7 @@ def scaled_aligned_causal_masked_softmax_backward( output_grads: torch.Tensor, softmax_results: torch.Tensor, scale: float ) -> torch.Tensor: """Backward pass for ScaledAlignedCausalMaskedSoftmax.""" - return tex.scaled_aligned_causal_masked_softmax_backward( - output_grads, softmax_results, scale - ) + return tex.scaled_aligned_causal_masked_softmax_backward(output_grads, softmax_results, scale) @scaled_aligned_causal_masked_softmax_backward.register_fake diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py index 883af16d11..f28270cf13 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py @@ -5,6 +5,7 @@ """ Utils/Helper classes and methods for attention """ + import math import os from typing import Any, Callable, Dict, List, Optional, Tuple, Union @@ -2197,12 +2198,11 @@ def _convert_bshd_to_thd_fake( ) -> torch.Tensor: del cu_seqlens h, d = bshd_tensor.shape[2], bshd_tensor.shape[3] - return torch.empty( - (num_tokens, h, d), dtype=bshd_tensor.dtype, device=bshd_tensor.device - ) + return torch.empty((num_tokens, h, d), dtype=bshd_tensor.dtype, device=bshd_tensor.device) def _convert_thd_to_bshd_setup_context(ctx, inputs, output): + del output thd_tensor, cu_seqlens, _batch_size, _max_seqlen = inputs ctx.save_for_backward(cu_seqlens) ctx.num_tokens = thd_tensor.size(0) @@ -2210,9 +2210,7 @@ def _convert_thd_to_bshd_setup_context(ctx, inputs, output): def _convert_thd_to_bshd_backward_wrapper(ctx, grad_bshd): (cu_seqlens,) = ctx.saved_tensors - grad_thd = torch.ops.te_attention.convert_bshd_to_thd( - grad_bshd, cu_seqlens, ctx.num_tokens - ) + grad_thd = torch.ops.te_attention.convert_bshd_to_thd(grad_bshd, cu_seqlens, ctx.num_tokens) return grad_thd, None, None, None @@ -2223,6 +2221,7 @@ def _convert_thd_to_bshd_backward_wrapper(ctx, grad_bshd): def _convert_bshd_to_thd_setup_context(ctx, inputs, output): + del output bshd_tensor, cu_seqlens, _num_tokens = inputs ctx.save_for_backward(cu_seqlens) ctx.batch_size = bshd_tensor.size(0) @@ -2271,9 +2270,7 @@ class ConvertBSHDtoTHD: @staticmethod def apply(bshd_tensor, cu_seqlens, num_tokens): # pylint: disable=missing-function-docstring - return torch.ops.te_attention.convert_bshd_to_thd( - bshd_tensor, cu_seqlens, num_tokens - ) + return torch.ops.te_attention.convert_bshd_to_thd(bshd_tensor, cu_seqlens, num_tokens) def get_qkv_format( From 8a7c2e8c3b0bf772b4ff49b75c01c2ea6880850a Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 10 Jul 2026 13:22:00 +0000 Subject: [PATCH 11/11] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- transformer_engine/pytorch/csrc/extensions/softmax.cpp | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/transformer_engine/pytorch/csrc/extensions/softmax.cpp b/transformer_engine/pytorch/csrc/extensions/softmax.cpp index 3c771ae5b6..be976d2cc4 100644 --- a/transformer_engine/pytorch/csrc/extensions/softmax.cpp +++ b/transformer_engine/pytorch/csrc/extensions/softmax.cpp @@ -186,9 +186,9 @@ at::Tensor scaled_upper_triang_masked_softmax_backward(at::Tensor output_grads_, auto softmax_results_cu = makeTransformerEngineTensor(softmax_results); auto input_grads_cu = makeTransformerEngineTensor(input_grads); - nvte_scaled_upper_triang_masked_softmax_backward( - output_grads_cu.data(), softmax_results_cu.data(), input_grads_cu.data(), scale_factor, - at::cuda::getCurrentCUDAStream()); + nvte_scaled_upper_triang_masked_softmax_backward(output_grads_cu.data(), + softmax_results_cu.data(), input_grads_cu.data(), + scale_factor, at::cuda::getCurrentCUDAStream()); return input_grads; }