diff --git a/tests/pytorch/attention/test_attention.py b/tests/pytorch/attention/test_attention.py index 6bcff5ce3f..96bd4a947c 100644 --- a/tests/pytorch/attention/test_attention.py +++ b/tests/pytorch/attention/test_attention.py @@ -32,6 +32,7 @@ from transformer_engine.pytorch.attention.dot_product_attention import ( _attention_backends, ) +import transformer_engine.pytorch.attention.dot_product_attention.utils as dpa_utils from transformer_engine.pytorch.attention.dot_product_attention.utils import ( AttentionParams, FlashAttentionUtils, @@ -41,6 +42,7 @@ get_attention_backend, get_fused_attn_spec, get_thd_padding_mask, + should_pad_qkv_head_dim, ) from transformer_engine.pytorch.attention import RotaryPositionEmbedding import transformer_engine.pytorch.cpp_extensions as ext @@ -1384,6 +1386,92 @@ def test_dpa_mla(dtype, model_configs, model): test_dot_product_attention(dtype, model_configs, model, True, None, False, False) +def _make_mla_attn_params(dtype, config, head_dim_qk, head_dim_v): + """Build the attention params a THD MLA config would feed into backend selection.""" + return dpa_utils.AttentionParams( + qkv_dtype=dtype, + qkv_layout="sbhd_sbhd_sbhd" if config.attn_type == "self" else "bshd_bshd_bshd", + batch_size=config.batch_size, + num_heads=config.num_heads, + num_gqa_groups=config.num_gqa_groups, + max_seqlen_q=config.max_seqlen_q, + max_seqlen_kv=config.max_seqlen_kv, + head_dim_qk=head_dim_qk, + head_dim_v=head_dim_v, + attn_mask_type=config.attn_mask_type, + is_training=True, + ) + + +@pytest.mark.skipif(get_cudnn_version() < (8, 9, 1), reason="cuDNN 8.9.1+ is required.") +@pytest.mark.parametrize("dtype", param_types) +@pytest.mark.parametrize("model_configs", [model_configs_mla]) +@pytest.mark.parametrize("model", model_configs_mla.keys()) +def test_dpa_mla_should_pad(dtype, model_configs, model): + """Test whether optional padding works correctly. + + It either pads (1) in case it's required for a functioning kernel or (2) to select a faster + backend. Otherwise padding should be avoided since it becomes redundant. + """ + config = model_configs[model] + native_params = _make_mla_attn_params(dtype, config, config.head_dim_qk, config.head_dim_v) + padded_head_dim = max(config.head_dim_qk, config.head_dim_v) + padded_params = _make_mla_attn_params(dtype, config, padded_head_dim, padded_head_dim) + + # Equal head dims run on any backend as-is, so padding is never warranted. + assert not should_pad_qkv_head_dim(padded_params, False), "suggested padding equal-sized QK/V" + + # Packed FP8 buffers must never be padded; the kernel would read the unpadded data. + fp8_recipe = recipe.DelayedScaling(fp8_format=recipe.Format.HYBRID, fp8_dpa=True) + fp8_params = replace(native_params, fp8=True, fp8_meta={"recipe": fp8_recipe}) + assert not should_pad_qkv_head_dim(fp8_params, True), "suggested padding packed FP8 inputs" + + native_backend = dpa_utils.get_attention_backend(native_params) + native_flash, native_fused = native_backend[0], native_backend[2] + padded_backend = dpa_utils.get_attention_backend(padded_params) + padded_flash, padded_fused = padded_backend[0], padded_backend[2] + + if should_pad_qkv_head_dim(native_params, False): + assert padded_flash or padded_fused, "suggested padding onto no faster backend" + assert (padded_flash, padded_fused) != ( + native_flash, + native_fused, + ), "suggested padding onto the same backend" + else: + assert ( + native_flash or native_fused or not (padded_flash or padded_fused) + ), "did not suggest padding when it was the only way onto a faster backend" + + +@pytest.mark.skipif(get_cudnn_version() < (8, 9, 1), reason="cuDNN 8.9.1+ is required.") +@pytest.mark.parametrize("dtype", param_types) +@pytest.mark.parametrize("model_configs", [model_configs_mla]) +@pytest.mark.parametrize("model", model_configs_mla.keys()) +def test_dpa_mla_should_pad_cache(dtype, model_configs, model): + """A memoized padding decision must match a freshly computed one. + + `run_dot_product_attention` enables one backend at a time without changing the attention + params, so a decision cached for one backend must not leak into the next. + """ + config = model_configs[model] + native_params = _make_mla_attn_params(dtype, config, config.head_dim_qk, config.head_dim_v) + saved = {name: os.environ.get(name) for name in dpa_utils.BACKEND_SELECTION_ENV_VARS} + try: + should_pad_qkv_head_dim(native_params, False) + os.environ["NVTE_FLASH_ATTN"] = "0" + os.environ["NVTE_FUSED_ATTN"] = "0" + cached = should_pad_qkv_head_dim(native_params, False) + fresh = should_pad_qkv_head_dim(native_params, False, invalidate_cache=True) + finally: + for name, value in saved.items(): + if value is None: + os.environ.pop(name, None) + else: + os.environ[name] = value + dpa_utils.clear_should_pad_qkv_head_dim_cache() + assert cached == fresh, "reused a stale padding decision after backend availability changed" + + model_configs_mask = { # test: ModelConfig(b, sq, hq, dqk) "mask_1_0": ModelConfig(2, 2048, 16, 64, attn_mask_type="causal"), diff --git a/transformer_engine/pytorch/attention/dot_product_attention/backends.py b/transformer_engine/pytorch/attention/dot_product_attention/backends.py index bccb4f1e26..749988310d 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/backends.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/backends.py @@ -525,7 +525,7 @@ def _forward( fp8_meta: Optional[Dict[str, Any]] = None, quantizers=None, fp8_output: bool = False, - ) -> torch.Tensor: + ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]: """Unfused attention fprop""" assert ( qkv_layout in QKVLayouts @@ -2395,7 +2395,7 @@ def forward( attn_bias, softmax_offset, fwd_args: FusedAttnFwdArgs, - ): + ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]: # pylint: disable=missing-function-docstring # q/k/v/attn_bias/softmax_offset are positional so autograd tracks them; # they are re-attached to fwd_args for the single-argument helpers. @@ -2709,7 +2709,7 @@ def forward( packed_kv: Optional[torch.Tensor] = None, bf16_backward: bool = False, load_balancing_strategy=CPLoadBalancingStrategy.DUAL_CHUNK_SWAP, - ) -> torch.Tensor: + ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]: """fused attention fprop""" assert ( fused_attention_backend != FusedAttnBackend["No_Backend"] diff --git a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py index e8e1396652..3731829c98 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py @@ -58,7 +58,6 @@ import transformer_engine.pytorch.attention.dot_product_attention.utils as dpa_utils from transformer_engine.pytorch.attention.dot_product_attention.utils import ( AttentionLogging as attn_log, - FlashAttentionUtils, ) from transformer_engine.pytorch.attention.dot_product_attention.backends import ( @@ -86,14 +85,7 @@ # configurations. Cache their full backend-selection results to avoid re-running the # selector for every policy in every layer and forward pass. _THD_POLICY_BACKEND_CACHE_LIMIT = 64 -_THD_POLICY_BACKEND_ENV_VARS = ( - "NVTE_FLASH_ATTN", - "NVTE_FLASH_ATTN_V2", - "NVTE_FLASH_ATTN_V3", - "NVTE_FLASH_ATTN_V4", - "NVTE_FUSED_ATTN", - "NVTE_UNFUSED_ATTN", -) +_THD_POLICY_BACKEND_ENV_VARS = dpa_utils.BACKEND_SELECTION_ENV_VARS _thd_policy_backend_cache = [] _alibi_cache = { @@ -424,7 +416,7 @@ def _delayed_scaling_recipe() -> Optional[DelayedScaling]: def _pad_qkv_head_dim(query_layer, key_layer, value_layer): - """Pad Q/K/V to the same head dimension for FlashAttention 2 MLA.""" + """Pad Q/K/V to the same head dimension for better backend selection result.""" orig_head_dim_qk = query_layer.shape[-1] orig_head_dim_v = value_layer.shape[-1] padded_head_dim = max(orig_head_dim_qk, orig_head_dim_v) @@ -1968,7 +1960,7 @@ def forward( qkv_interleave_dim: int = -3, thd_attention_policies: Optional[List[Dict[str, Any]]] = None, thd_attention_policy_dispatch: str = "auto", - ) -> torch.Tensor: + ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor], None]: r""" Dot Product Attention Layer. @@ -2841,6 +2833,32 @@ def forward( bottom_right_diagonal=bottom_right_diagonal, pad_between_seqs=pad_between_seqs, ) + + # Optionally pad inputs to `max(head_dim_qk, head_dim_v)` when it's required by the + # kernel or when it leads to faster kernel selection. Note that the padded tensors could + # still be unsupported input shapes by `{Flash,Fused}Attention`. + can_pad_qkv_head_dim = ( + value_layer is not None + and not isinstance(value_layer, Float8TensorStorage) + and not is_in_onnx_export_mode() + ) + has_packed_qkv = qkv_layer is not None or kv_layer is not None + is_qkv_head_dim_padded = False + orig_head_dim_v = head_dim_v + if can_pad_qkv_head_dim and dpa_utils.should_pad_qkv_head_dim( + attention_params, + has_packed_qkv, + invalidate_cache=_attention_backends["backend_selection_requires_update"], + ): + query_layer, key_layer, value_layer, _, _ = _pad_qkv_head_dim( + query_layer, key_layer, value_layer + ) + padded_head_dim = max(head_dim_qk, head_dim_v) + head_dim_qk = head_dim_v = padded_head_dim + attention_params.head_dim_qk = padded_head_dim + attention_params.head_dim_v = padded_head_dim + is_qkv_head_dim_padded = True + if is_in_onnx_export_mode(): # We do not want to call get_attention_backend() in ONNX mode # and we want to avoid using any global variables like _attention_backends. @@ -2913,21 +2931,6 @@ def forward( ) if use_flash_attention: - orig_qk_dim = None - orig_v_dim = None - if ( - flash_attention_backend == FlashAttentionUtils.version - and not isinstance(value_layer, Float8TensorStorage) - and head_dim_qk != head_dim_v - ): - ( - query_layer, - key_layer, - value_layer, - orig_qk_dim, - orig_v_dim, - ) = _pad_qkv_head_dim(query_layer, key_layer, value_layer) - if core_attention_bias_type == "alibi": alibi_slopes, _ = dpa_utils.get_alibi( _alibi_cache, @@ -2966,11 +2969,8 @@ def forward( cu_seqlens_q_padded=cu_seqlens_q_padded, cu_seqlens_kv_padded=cu_seqlens_kv_padded, ) - if orig_qk_dim is not None and orig_qk_dim > orig_v_dim: - return _trim_output(attn_out, num_attention_heads, orig_qk_dim, orig_v_dim) - return attn_out - if use_fused_attention: + elif use_fused_attention: fu_core_attention_bias_type = core_attention_bias_type fu_core_attention_bias = core_attention_bias if core_attention_bias_type == "alibi" and (alibi_slopes is not None): @@ -2985,7 +2985,7 @@ def forward( bottom_right_alignment=bottom_right_diagonal, ) if checkpoint_core_attention: - return self._checkpointed_attention_forward( + attn_out = self._checkpointed_attention_forward( self.fused_attention, query_layer, key_layer, @@ -3021,52 +3021,53 @@ def forward( packed_kv=kv_layer, bf16_backward=bf16_backward, ) - return self.fused_attention( - query_layer, - key_layer, - value_layer, - qkv_layout=qkv_layout, - cu_seqlens_q=cu_seqlens_q, - cu_seqlens_kv=cu_seqlens_kv, - cu_seqlens_q_padded=cu_seqlens_q_padded, - cu_seqlens_kv_padded=cu_seqlens_kv_padded, - max_seqlen_q=max_seqlen_q, - max_seqlen_kv=max_seqlen_kv, - attn_mask_type=attn_mask_type, - attention_mask=attention_mask, - window_size=window_size, - bottom_right_diagonal=bottom_right_diagonal, - fused_attention_backend=fused_attention_backend, - core_attention_bias_type=fu_core_attention_bias_type, - core_attention_bias=fu_core_attention_bias, - fast_zero_fill=fast_zero_fill, - cp_group=self.cp_group, - cp_global_ranks=self.cp_global_ranks, - cp_stream=self.cp_stream, - cp_comm_type=self.cp_comm_type, - load_balancing_strategy=self.load_balancing_strategy, - fp8=self.fp8 and self.fp8_meta["recipe"].fp8_dpa, - fp8_meta=self.fp8_meta, - quantizers=self.quantizers, - pad_between_seqs=pad_between_seqs, - inference_params=inference_params, - softmax_offset=softmax_offset, - fp8_output=fp8_output, - score_mod=score_mod, - score_mod_bprop=score_mod_bprop, - score_mod_tensors=score_mod_tensors, - score_mod_bprop_tensors=score_mod_bprop_tensors, - packed_qkv=qkv_layer, - packed_kv=kv_layer, - bf16_backward=bf16_backward, - ) + else: + attn_out = self.fused_attention( + query_layer, + key_layer, + value_layer, + qkv_layout=qkv_layout, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_kv=cu_seqlens_kv, + cu_seqlens_q_padded=cu_seqlens_q_padded, + cu_seqlens_kv_padded=cu_seqlens_kv_padded, + max_seqlen_q=max_seqlen_q, + max_seqlen_kv=max_seqlen_kv, + attn_mask_type=attn_mask_type, + attention_mask=attention_mask, + window_size=window_size, + bottom_right_diagonal=bottom_right_diagonal, + fused_attention_backend=fused_attention_backend, + core_attention_bias_type=fu_core_attention_bias_type, + core_attention_bias=fu_core_attention_bias, + fast_zero_fill=fast_zero_fill, + cp_group=self.cp_group, + cp_global_ranks=self.cp_global_ranks, + cp_stream=self.cp_stream, + cp_comm_type=self.cp_comm_type, + load_balancing_strategy=self.load_balancing_strategy, + fp8=self.fp8 and self.fp8_meta["recipe"].fp8_dpa, + fp8_meta=self.fp8_meta, + quantizers=self.quantizers, + pad_between_seqs=pad_between_seqs, + inference_params=inference_params, + softmax_offset=softmax_offset, + fp8_output=fp8_output, + score_mod=score_mod, + score_mod_bprop=score_mod_bprop, + score_mod_tensors=score_mod_tensors, + score_mod_bprop_tensors=score_mod_bprop_tensors, + packed_qkv=qkv_layer, + packed_kv=kv_layer, + bf16_backward=bf16_backward, + ) - if use_unfused_attention: + elif use_unfused_attention: allow_emulation = ( os.getenv("NVTE_UnfusedDPA_Emulate_FP8", "0") == "1" or is_in_onnx_export_mode() ) if checkpoint_core_attention: - return self._checkpointed_attention_forward( + attn_out = self._checkpointed_attention_forward( self.unfused_attention, _alibi_cache, query_layer, @@ -3092,29 +3093,39 @@ def forward( quantizers=self.quantizers, fp8_output=fp8_output, ) - return self.unfused_attention( - _alibi_cache, - query_layer, - key_layer, - value_layer, - qkv_layout=qkv_layout, - cu_seqlens_q=cu_seqlens_q, - cu_seqlens_kv=cu_seqlens_kv, - max_seqlen_q=max_seqlen_q, - max_seqlen_kv=max_seqlen_kv, - attn_mask_type=attn_mask_type, - attention_mask=attention_mask, - window_size=window_size, - bottom_right_diagonal=bottom_right_diagonal, - softcap=softcap, - core_attention_bias_type=core_attention_bias_type, - core_attention_bias=core_attention_bias, - alibi_slopes=alibi_slopes, - inference_params=inference_params, - softmax_offset=softmax_offset, - fp8=self.fp8 and self.fp8_meta["recipe"].fp8_dpa and allow_emulation, - fp8_meta=self.fp8_meta, - quantizers=self.quantizers, - fp8_output=fp8_output, - ) - return None + else: + attn_out = self.unfused_attention( + _alibi_cache, + query_layer, + key_layer, + value_layer, + qkv_layout=qkv_layout, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_kv=cu_seqlens_kv, + max_seqlen_q=max_seqlen_q, + max_seqlen_kv=max_seqlen_kv, + attn_mask_type=attn_mask_type, + attention_mask=attention_mask, + window_size=window_size, + bottom_right_diagonal=bottom_right_diagonal, + softcap=softcap, + core_attention_bias_type=core_attention_bias_type, + core_attention_bias=core_attention_bias, + alibi_slopes=alibi_slopes, + inference_params=inference_params, + softmax_offset=softmax_offset, + fp8=self.fp8 and self.fp8_meta["recipe"].fp8_dpa and allow_emulation, + fp8_meta=self.fp8_meta, + quantizers=self.quantizers, + fp8_output=fp8_output, + ) + else: + return None + + if is_qkv_head_dim_padded: + trim_args = (num_attention_heads, head_dim_v, orig_head_dim_v) + if self.return_max_logit: + attn_out = (_trim_output(attn_out[0], *trim_args), *attn_out[1:]) + else: + attn_out = _trim_output(attn_out, *trim_args) + return attn_out diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py index e8cbd8a34c..b850aaf8be 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py @@ -13,7 +13,7 @@ import logging import functools -from dataclasses import dataclass, fields +from dataclasses import dataclass, fields, replace import numpy as np from packaging.version import Version as PkgVersion @@ -73,6 +73,31 @@ _cu_seqlens_cache = {} +# `get_attention_backend` reads these at call time, so they belong in the key of any cache +# that memoizes a backend selection or a decision derived from one. +BACKEND_SELECTION_ENV_VARS = ( + "NVTE_FLASH_ATTN", + "NVTE_FLASH_ATTN_V2", + "NVTE_FLASH_ATTN_V3", + "NVTE_FLASH_ATTN_V4", + "NVTE_FUSED_ATTN", + "NVTE_UNFUSED_ATTN", +) + +# Global var for MLA padding cache. +_should_pad_qkv_head_dim_cache: Dict[str, Any] = { + "attention_params": None, + "has_packed_qkv": None, + "environment": None, + "should_pad": None, +} + + +def clear_should_pad_qkv_head_dim_cache() -> None: + """Drop the memoized padding decision when backend selection is invalidated.""" + for key in _should_pad_qkv_head_dim_cache: + _should_pad_qkv_head_dim_cache[key] = None + class AttentionLogging: """ @@ -1012,6 +1037,14 @@ def _disable_all_flash_attention() -> None: # Filter: Head dimension if head_dim_qk != head_dim_v: + if use_flash_attention_2 and FlashAttentionUtils.is_installed: + logger.debug( + "Disabling FlashAttention 2 as it does not support head_dim_qk != head_dim_v. " + "Found: head_dim_qk = %s, head_dim_v = %s.", + head_dim_qk, + head_dim_v, + ) + use_flash_attention_2 = False qkv_layout_group = qkv_layout.replace("b", "").replace("s", "").replace("t", "") if use_fused_attention and qkv_layout_group != "hd_hd_hd": logger.debug( @@ -1033,27 +1066,25 @@ def _disable_all_flash_attention() -> None: ) use_fused_attention = False - fa2_padded_head_dim = max(head_dim_qk, head_dim_v) if ( # pylint: disable=too-many-boolean-expressions use_flash_attention_2 and FlashAttentionUtils.is_installed and ( - fa2_padded_head_dim > 256 - or fa2_padded_head_dim % 8 != 0 + max(head_dim_qk, head_dim_v) > 256 + or head_dim_qk % 8 != 0 + or head_dim_v % 8 != 0 or ( - fa2_padded_head_dim > 192 + max(head_dim_qk, head_dim_v) > 192 and device_compute_capability not in ((8, 0), (9, 0), (10, 0), (12, 0)) ) ) ): logger.debug( "Disabling FlashAttention 2 due to unsupported head_dim_qk and head_dim_v. " - "Supported after padding: padded head_dim %%8 = 0, padded head_dim <= 256 " - "(>192 requires sm80/90/100+). " - "Found: head_dim_qk = %s, head_dim_v = %s, padded head_dim = %s, on sm%s.", + "Supported: head_dim %%8 = 0, head_dim <= 256 (>192 requires sm80/90/100+). " + "Found: head_dim_qk = %s, head_dim_v = %s, on sm%s.", head_dim_qk, head_dim_v, - fa2_padded_head_dim, ".".join([str(i) for i in device_compute_capability]), ) use_flash_attention_2 = False @@ -1984,6 +2015,99 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt ) +@torch.no_grad() +def should_pad_qkv_head_dim( + attention_params: AttentionParams, + has_packed_qkv: bool, + invalidate_cache: bool = False, +) -> bool: + """Decide whether to pad Q/K/V to a common head dim. + + This returns `True` in two cases: + 1. when neither `FlashAttention` nor `FusedAttention` runs the native head dims, so padding + may enable one of them. + 2. when only one of them runs the native head dims, and padding reaches the other one, which + backend selection ranks higher for this configuration. + + The decision is memoized on `attention_params`, `has_packed_qkv` and the backend-selection + environment. Callers that invalidate backend selection for any other reason must pass + `invalidate_cache=True`. + + Parameters + ---------- + attention_params : AttentionParams + Attention parameters for the native (unpadded) QK and V head dimensions. + has_packed_qkv : bool + Whether the caller passed a packed `qkv_layer` or `kv_layer` buffer. + invalidate_cache : bool + Whether to invalidate the cache prior to checking it. + + Returns + ------- + bool + Whether Q/K/V should be padded to a common (the wider) head dim. + """ + if invalidate_cache: + clear_should_pad_qkv_head_dim_cache() + if attention_params.head_dim_qk == attention_params.head_dim_v: + return False + + environment = tuple(os.environ.get(name) for name in BACKEND_SELECTION_ENV_VARS) + if ( + _should_pad_qkv_head_dim_cache["attention_params"] == attention_params + and _should_pad_qkv_head_dim_cache["has_packed_qkv"] is has_packed_qkv + and _should_pad_qkv_head_dim_cache["environment"] == environment + ): + return _should_pad_qkv_head_dim_cache["should_pad"] + + # Check whether padding must be skipped because FP8 DPA is active and the caller passed packed + # QKV/KV buffers. + recipe = (attention_params.fp8_meta or {}).get("recipe", None) + should_pad = not ( + attention_params.fp8 and recipe is not None and getattr(recipe, "fp8_dpa", False) + ) + # `should_pad=False` indicates that FP8 DPA is active. + if not should_pad: + should_pad = not has_packed_qkv + + if should_pad: + native_backend = get_attention_backend(attention_params) + native_use_flash_attention = native_backend[0] + native_use_fused_attention = native_backend[2] + native_available_backends = native_backend[5] + if native_available_backends[0] and native_available_backends[1]: + # Both FlashAttention and FusedAttention run the native head dims, so backend + # selection has already picked the better of the two and padding cannot improve on it. + should_pad = False + else: + padded_head_dim = max(attention_params.head_dim_qk, attention_params.head_dim_v) + padded_params = replace( + attention_params, + head_dim_qk=padded_head_dim, + head_dim_v=padded_head_dim, + ) + padded_backend = get_attention_backend(padded_params) + padded_use_flash_attention = padded_backend[0] + padded_use_fused_attention = padded_backend[2] + if not (native_use_flash_attention or native_use_fused_attention): + # Neither FlashAttention nor FusedAttention runs the native head dims, so pad if + # either runs the padded ones. + should_pad = padded_use_flash_attention or padded_use_fused_attention + else: + # Exactly one of FlashAttention and FusedAttention runs the native head dims. Pad + # only if padding reaches the other one, which backend selection ranks higher. + should_pad = ( + padded_use_flash_attention or padded_use_fused_attention + ) and padded_use_flash_attention != native_use_flash_attention + # Store a shallow copy as the key to avoid later `attention_params` in-place modifications + # resulting in subsequent cache misses. + _should_pad_qkv_head_dim_cache["attention_params"] = replace(attention_params) + _should_pad_qkv_head_dim_cache["has_packed_qkv"] = has_packed_qkv + _should_pad_qkv_head_dim_cache["environment"] = environment + _should_pad_qkv_head_dim_cache["should_pad"] = should_pad + return should_pad + + @torch.no_grad() def get_thd_padding_mask(num_tokens, cu_seqlens, cu_seqlens_padded): """Return a boolean mask identifying padding token positions in a THD buffer.