diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py index f8e09d5ce1..22f0ffc3bd 100644 --- a/tests/pytorch/test_torch_compile.py +++ b/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) diff --git a/transformer_engine/pytorch/attention/dot_product_attention/backends.py b/transformer_engine/pytorch/attention/dot_product_attention/backends.py index c3ee569c4b..aa17d859bd 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/backends.py +++ b/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)