diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py index 8ebce563ce..a0b959dcde 100644 --- a/tests/pytorch/test_torch_compile.py +++ b/tests/pytorch/test_torch_compile.py @@ -37,6 +37,9 @@ NVFP4Quantizer, ) 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) @@ -389,6 +392,164 @@ def fn(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 + + # --------------------------------------------------------------------------- # 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 785e438cda..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 @@ -365,7 +366,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: bool = False, fp8_output: bool = False, **kwargs) -> torch.Tensor: + """Unfused attention fprop; see `_forward` for the argument list. + + FP8 (emulation and/or Float8Tensor output) is not supported under + torch.compile -- run the backend as an eager island in that case. + """ + 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: + """Eager-only (dynamo-disabled) wrapper around `_forward`.""" + return self._forward(*args, **kwargs) + + def _forward( self, _alibi_cache: Dict[str, Any], query_layer: torch.Tensor, @@ -425,6 +441,10 @@ 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 + # 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, cu_seqlens_q, @@ -708,6 +728,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..6e8a14402f 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/softmax.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/softmax.py @@ -3,142 +3,248 @@ # See LICENSE for license information. """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 from transformer_engine.pytorch.export import is_in_onnx_export_mode - THREADS_PER_WARP = 32 THREADS_PER_BLOCK = 128 -_default_causal_mask = {} +# ----------------------------- ScaledSoftmax ------------------------------- -def _get_default_causal_mask(mask_type: str, sq: int, sk: int) -> torch.Tensor: - """Return the causal upper triangular mask for softmax input""" +@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) - 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] +@scaled_softmax_forward.register_fake +def _scaled_softmax_forward_fake(inputs: torch.Tensor, scale: float) -> torch.Tensor: + del scale + return torch.empty_like(inputs) -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_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) - @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_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 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 +def _scaled_softmax_setup_context(ctx, inputs, output): + _inp, scale = inputs + ctx.scale = scale + ctx.save_for_backward(output) -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. - """ +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 - @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] - ) +scaled_softmax_forward.register_autograd( + _scaled_softmax_backward_wrapper, + setup_context=_scaled_softmax_setup_context, +) - return input_grads, None +# --------------------------- ScaledMaskedSoftmax --------------------------- -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]) +@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) - 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 +@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) - 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_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) -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_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) - 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 +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) + + +@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) + + +@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) + + +@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) + + +@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 +340,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 7be94a6fa1..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 @@ -2142,76 +2143,134 @@ 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): + del 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): + del 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, + return torch.ops.te_attention.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 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, - ) - return bshd_tensor, None + return torch.ops.te_attention.convert_bshd_to_thd(bshd_tensor, cu_seqlens, num_tokens) def get_qkv_format( diff --git a/transformer_engine/pytorch/csrc/extensions/softmax.cpp b/transformer_engine/pytorch/csrc/extensions/softmax.cpp index 3bb6a5e7b3..be976d2cc4 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, - 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 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