[PyTorch] torch.compile support for FusedAttention (diff vs #3480 only)

tests/pytorch/test_torch_compile.py

@@ -898,11 +898,8 @@ def _skip_unsupported(
 ) -> None:
     """Skip what the backend under test cannot run, or -- for a test that
     compiles it -- cannot be compiled."""
-    if compiled and backend == "fused":
-        # FusedAttention's forward carries @no_torch_dynamo, so there is nothing
-        # to compile: it runs as an eager island. Drop this skip once it traces,
-        # and the tests below cover it as they do the others.
-        pytest.skip("FusedAttention is an eager island and does not compile")
+    if compiled and backend == "fused" and inference_params is not None:
+        pytest.skip("KV caching with FusedAttention falls back to eager under torch.compile")
     available, _, _ = get_available_attention_backends(
         spec["model_config"],
         dtype,
@@ -1046,32 +1043,39 @@ def test_dpa_torch_compile(monkeypatch, backend, config):
 
 
 def test_dpa_torch_compile_around_fused(monkeypatch):
-    """FusedAttention itself is an eager island, but everything around it is
-    compiled: DotProductAttention traces up to the backend call, breaks the
-    graph there and resumes afterwards. What crosses that break has to survive
-    it -- the sub-backend enum did not, and reached cuDNN as the function that
-    produced it."""
+    """Without the fused attention custom op, FusedAttention is an eager island
+    and everything around it is compiled: DotProductAttention traces up to the
+    backend call, breaks the graph there and resumes afterwards. What crosses
+    that break has to survive it -- the sub-backend enum did not, and reached
+    cuDNN as the function that produced it."""
+    from transformer_engine.pytorch.attention.dot_product_attention import backends
+
     dtype = torch.bfloat16
     spec = _DPA_COMPILE_CONFIGS["self_bshd_causal"]
     _skip_unsupported(spec, "fused", dtype, compiled=False)
     _force_dpa_backend(monkeypatch, "fused")
+    monkeypatch.setattr(backends, "_fused_attn_op", None)
 
     module = _make_dpa(spec, dtype)
     args, kwargs, grads = _make_dpa_inputs(spec, dtype)
     # No fullgraph: the graph break at the eager island is the point here.
-    _compare_compiled_to_eager(module, args, kwargs, grads, monkeypatch, "fused", dtype)
+    with pytest.warns(UserWarning, match="Falling back to eager execution"):
+        _compare_compiled_to_eager(module, args, kwargs, grads, monkeypatch, "fused", dtype)
 
 
-@pytest.mark.parametrize("backend", ["flash", "unfused"])
+@pytest.mark.parametrize("backend", ["flash", "fused", "unfused"])
 @pytest.mark.parametrize("config", ["self_bshd_causal", "kv_cache_bshd"])
 def test_dpa_torch_compile_cudagraphs(monkeypatch, backend, config):
     """`mode="reduce-overhead"`: forward and backward of DotProductAttention
     are captured into CUDA graphs and replayed on subsequent iterations."""
     dtype = torch.bfloat16
     spec = _DPA_COMPILE_CONFIGS[config]
-    _force_dpa_backend(monkeypatch, backend)
-
     module = _make_dpa(spec, dtype)
+    _, kwargs, _ = _make_dpa_inputs(spec, dtype)
+    # Before forcing the backend: probing the available backends re-runs the
+    # selection and would otherwise be cached over the forced one.
+    _skip_unsupported(spec, backend, dtype, inference_params=kwargs.get("inference_params"))
+    _force_dpa_backend(monkeypatch, backend)
 
     torch._dynamo.reset()
     counters.clear()
@@ -1230,7 +1234,7 @@ def _thd_without_max_seqlen_inputs(spec, dtype):
 }
 
 
-@pytest.mark.parametrize("backend", ["flash", "unfused"])
+@pytest.mark.parametrize("backend", ["flash", "fused", "unfused"])
 @pytest.mark.parametrize("case", _EAGER_FALLBACK_CASES.keys())
 def test_dpa_torch_compile_eager_fallback(monkeypatch, backend, case):
     """Calls that cannot be traced run as an eager island instead, with a
@@ -1241,6 +1245,7 @@ def test_dpa_torch_compile_eager_fallback(monkeypatch, backend, case):
     dtype = torch.bfloat16
     config_name, make_inputs = _EAGER_FALLBACK_CASES[case]
     spec = _DPA_COMPILE_CONFIGS[config_name]
+    _skip_unsupported(spec, backend, dtype, compiled=False)
     _force_dpa_backend(monkeypatch, backend)
 
     module = _make_dpa(spec, dtype)

transformer_engine/pytorch/attention/dot_product_attention/backends.py

@@ -17,6 +17,7 @@
 import torch
 import torch.nn.functional as F
 from transformer_engine.pytorch.utils import (
+    get_cudnn_version,
     get_device_compute_capability,
     split_tensor_along_dim,
 )
@@ -58,7 +59,8 @@
     fa_prepare_fwd,
 )
 from transformer_engine.pytorch.jit import no_torch_dynamo
-from transformer_engine.pytorch.dynamo.custom_op import TensorOrQuantized
+from transformer_engine.pytorch.dynamo.custom_op import TensorOrQuantized, register_custom_op
+from transformer_engine.pytorch.dynamo.tensor_spec import TensorSpec
 from transformer_engine.pytorch.attention.dot_product_attention.context_parallel import (
     attn_forward_func_with_cp,
 )
@@ -2250,6 +2252,134 @@ def backward(ctx, d_out, *_args):
         return (*_fused_attn_backward_impl(bwd_args), None)
 
 
+def _fused_attn_stats_shape(args: FusedAttnFwdArgs, q_format: str) -> Tuple[int, ...]:
+    """Shape of the softmax stats auxiliary tensor cuDNN returns."""
+    q_shape = args.q.shape
+    if q_format == "thd":
+        num_heads = q_shape[1]
+        major, minor, _ = get_cudnn_version()
+        if (major, minor) >= (9, 6) and get_device_compute_capability() != (12, 0):
+            return (q_shape[0], num_heads, 1)
+        batch_size = args.cu_seqlens_q.shape[0] - 1
+        return (batch_size, num_heads, args.max_seqlen_q, 1)
+    if q_format == "bshd":
+        batch_size, seqlen, num_heads = q_shape[0], q_shape[1], q_shape[2]
+    else:
+        seqlen, batch_size, num_heads = q_shape[0], q_shape[1], q_shape[2]
+    return (batch_size, num_heads, seqlen, 1)
+
+
+def _fused_attn_forward_fake(
+    args: FusedAttnFwdArgs,
+) -> Tuple[TensorSpec, Optional[TensorSpec], Tuple[Any, ...], Dict[str, Any]]:
+    """Data-free twin of :func:`_fused_attn_forward_impl` (non-FP8 only)."""
+    q, v = args.q, args.v
+    _, o_format, _ = dpa_utils.get_qkv_format(args.qkv_layout)
+    out = TensorSpec(shape=(*q.shape[:-1], v.shape[-1]), dtype=q.dtype, device=q.device)
+    max_logit = None
+    if args.return_max_logit:
+        max_logit = TensorSpec(shape=(q.shape[-2],), dtype=q.dtype, device=q.device)
+    softmax_stats = TensorSpec(
+        shape=_fused_attn_stats_shape(args, o_format), dtype=torch.float32, device=q.device
+    )
+    rng_state = TensorSpec(shape=(2,), dtype=torch.int64, device=q.device)
+    has_bias = args.attn_bias_type not in ["no_bias", "alibi"] and args.attn_bias is not None
+    has_softmax_offset = args.softmax_type != "vanilla" and args.softmax_offset is not None
+    tensors_to_save = (*(None,) * 8, softmax_stats, rng_state, None, None)
+    saved_from = (
+        *(None,) * 4,
+        "q",
+        "k",
+        "v",
+        "out",
+        None,
+        None,
+        "attn_bias" if has_bias else None,
+        "softmax_offset" if has_softmax_offset else None,
+    )
+    ctx_attrs = {
+        "fp8": False,
+        "is_input_fp8": False,
+        "saved_from": saved_from,
+        "fused_attention_backend": args.fused_attention_backend,
+        "nominal_dtype": q.dtype,
+        "qkv_type": None,
+        "qkv_layout": args.qkv_layout,
+        "original_qkv_layout": args.qkv_layout,
+        "o_format": o_format,
+        "qkv_scale_inv_format": None,
+        "QKV_quantizer": None,
+        "O_quantizer": None,
+        "S_quantizer": None,
+        "dQKV_quantizer": None,
+        "dO_quantizer": None,
+        "dP_quantizer": None,
+    }
+    return out, max_logit, tensors_to_save, ctx_attrs
+
+
+def _fused_attn_backward_op_impl(
+    args: FusedAttnBwdArgs,
+) -> Tuple[Optional[torch.Tensor], ...]:
+    dq, dk, dv, d_bias, d_softmax_offset = _fused_attn_backward_impl(args)
+    # Packed layouts return dq/dk/dv as views of one buffer; op outputs may not alias.
+    return dq.contiguous(), dk.contiguous(), dv.contiguous(), d_bias, d_softmax_offset
+
+
+def _fused_attn_backward_fake(
+    args: FusedAttnBwdArgs,
+) -> Tuple[Optional[TensorSpec], ...]:
+    """Data-free twin of :func:`_fused_attn_backward_op_impl`."""
+    q, k, v = args.q, args.k, args.v
+    dq = TensorSpec(shape=tuple(q.shape), dtype=q.dtype, device=q.device)
+    dk = TensorSpec(shape=tuple(k.shape), dtype=k.dtype, device=k.device)
+    dv = TensorSpec(shape=tuple(v.shape), dtype=v.dtype, device=v.device)
+    d_bias = None
+    if args.aux_bias is not None:
+        d_bias = TensorSpec(shape=tuple(args.aux_bias.shape), dtype=q.dtype, device=q.device)
+    d_softmax_offset = None
+    if args.aux_softmax_offset is not None:
+        d_softmax_offset = TensorSpec(
+            shape=(1, q.shape[-2], 1, 1), dtype=torch.float32, device=q.device
+        )
+    return dq, dk, dv, d_bias, d_softmax_offset
+
+
+# Custom op used under ``torch.compile``.
+_fused_attn_op = register_custom_op(
+    op_name="fused_attn",
+    input_tensors_for_grad=["q", "k", "v", "attn_bias", "softmax_offset"],
+    fwd_arg_type=FusedAttnFwdArgs,
+    fwd_impl=_fused_attn_forward_impl,
+    fwd_fake_impl=_fused_attn_forward_fake,
+    setup_context=_fused_attn_setup_ctx,
+    bwd_arg_type=FusedAttnBwdArgs,
+    bwd_impl=_fused_attn_backward_op_impl,
+    bwd_fake_impl=_fused_attn_backward_fake,
+)
+
+
+def _needs_eager_fused_attention(call: Dict[str, Any]) -> Optional[str]:
+    """Why this FusedAttention call has to run outside the graph, or None.
+
+    `call` maps `FusedAttention.forward`'s parameter names to the arguments
+    this call passed, including `self`.
+    """
+    if _fused_attn_op is None:
+        return "the fused attention custom op (unavailable on this PyTorch build)"
+    if call.get("fp8", False):
+        return "FP8 attention"
+    if call.get("cp_group") is not None:
+        return "context parallelism"
+    if call.get("score_mod") is not None:
+        return "score_mod"
+    if call["self"].use_FAv2_bwd:
+        return "NVTE_FUSED_ATTN_USE_FAv2_BWD"
+    if is_cpu_offload_enabled():
+        return "CPU activation offloading"
+    return None
+
+
 class FusedAttention(torch.nn.Module):
     """Dot product attention using cuDNN attention:
 
@@ -2305,7 +2435,7 @@ def remove_extra_states_check(self, incompatible_keys):  # pylint: disable=unuse
 
         self.register_load_state_dict_post_hook(remove_extra_states_check)
 
-    @no_torch_dynamo()
+    @no_torch_dynamo(when=_needs_eager_fused_attention)
     def forward(
         self,
         query_layer: torch.Tensor,
@@ -2560,21 +2690,26 @@ def forward(
                 return_max_logit=self.return_max_logit,
                 rng_gen=None,
                 fp8=fp8,
-                fp8_meta=fp8_meta,
-                quantizers=quantizers,
+                fp8_meta=fp8_meta if fp8 else None,
+                quantizers=quantizers if fp8 else None,
                 fp8_output=fp8_output,
                 bf16_backward=bf16_backward,
                 layer_number=self.layer_number,
             )
             with self.attention_dropout_ctx():
-                output = FusedAttnFunc.apply(
-                    query_layer,
-                    key_layer,
-                    value_layer,
-                    core_attention_bias,
-                    softmax_offset,
-                    fwd_args,
-                )
+                if torch.compiler.is_compiling() and _fused_attn_op is not None:
+                    output, max_logit = _fused_attn_op(fwd_args)
+                    if self.return_max_logit:
+                        output = (output, max_logit)
+                else:
+                    output = FusedAttnFunc.apply(
+                        query_layer,
+                        key_layer,
+                        value_layer,
+                        core_attention_bias,
+                        softmax_offset,
+                        fwd_args,
+                    )
 
         if self.return_max_logit:
             # ...hd -> ...(hd)