From 15a34b4e93d7b9ff5181fda844026f2c3e0fa1ef Mon Sep 17 00:00:00 2001 From: janEbert Date: Mon, 10 Aug 2026 20:45:17 +0200 Subject: [PATCH 01/30] Automate MLA padding-related kernel selection While support for native QKV shapes has been added, some MLA kernels still benefit from having QKV padded to the same head dimension. Since it's not easy to automate this decision, we probe `get_attention_backend` twice with and without padding to figure out which version selects the better kernel. If padding results in a better kernel than not padding, we pad; otherwise, Q, K, and V stay native. Signed-off-by: janEbert --- .../dot_product_attention.py | 175 +++++++++++------- .../attention/dot_product_attention/utils.py | 46 ++++- 2 files changed, 152 insertions(+), 69 deletions(-) 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 e8e1396652d..0fd89bf5d34 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 @@ -2841,6 +2841,31 @@ def forward( bottom_right_diagonal=bottom_right_diagonal, pad_between_seqs=pad_between_seqs, ) + + # Optional MLA head-dim pad: padding Q/K/V to a common (the wider) head dim can upgrade + # the selected backend off the slow `UnfusedDotProductAttention` for certain setups. + # Probe both shapes and pad only when padding escapes the unfused path (and leaving the + # dims native would land on the unfused path). The pad-then-trim is an identity, so this + # never changes the result, only which kernel runs. + qkv_head_pad = False + orig_head_dim_v = head_dim_v + if ( + not is_in_onnx_export_mode() + and head_dim_qk != head_dim_v + and value_layer is not None + and not isinstance(value_layer, Float8TensorStorage) + and dpa_utils.should_pad_qkv_head_dim(attention_params) + ): + # Pad Q/K/V to the wider head dim so a fused backend can run. + 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 + qkv_head_pad = 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. @@ -2966,8 +2991,10 @@ 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) + if (orig_qk_dim is not None and orig_qk_dim > orig_v_dim) or qkv_head_pad: + attn_out = _trim_output( + attn_out, num_attention_heads, head_dim_qk, orig_head_dim_v + ) return attn_out if use_fused_attention: @@ -2985,7 +3012,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 +3048,58 @@ 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 qkv_head_pad: + attn_out = _trim_output( + attn_out, num_attention_heads, head_dim_qk, orig_head_dim_v + ) + return attn_out if 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 +3125,35 @@ 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, - ) + 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, + ) + if qkv_head_pad: + attn_out = _trim_output( + attn_out, num_attention_heads, head_dim_qk, orig_head_dim_v + ) + return attn_out return None diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py index a773091cd4e..fa79e6a6a97 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 @@ -1984,6 +1984,50 @@ 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", +) -> bool: + """Decide whether padding Q/K/V to a common head dim upgrades the attention backend. + + Padding Q, K, and V to the wider of the QK and V head dims is an identity -- the padded head + dims contribute zero to the Q*K dot product (so scores and softmax are unchanged) and the padded + V columns are sums of zeros (so trimming them back changes nothing). This returns `True` iff the + native mismatched shape selects the slow `UnfusedDotProductAttention` while the padded equal + shape selects a fused backend (`FlashAttention` or `FusedAttention`), i.e. the only case in + which the pad actually upgrades the kernel instead of being pure overhead. + + Parameters + ---------- + attention_params : AttentionParams + Attention parameters for the native (unpadded) QK and V head dimensions. + + Returns + ------- + bool + Whether Q/K/V should be padded to a common (the wider) head dim. + """ + if attention_params.head_dim_qk == attention_params.head_dim_v: + return False + native_backend = get_attention_backend(attention_params) + native_use_unfused_attention = native_backend[4] + # No point padding if native path is already fused. + if not native_use_unfused_attention: + return False + # Probe on a copy so the caller's `attention_params` are never mutated, even if + # `get_attention_backend` raises. + 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] + return bool(padded_use_flash_attention or padded_use_fused_attention) + + @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. From 6e5bf3641fc74fbf93ed7247389c473a2d7c727c Mon Sep 17 00:00:00 2001 From: janEbert Date: Mon, 10 Aug 2026 21:21:57 +0200 Subject: [PATCH 02/30] Add padding-related tests Signed-off-by: janEbert --- .../test_dpa_mla_qkv_head_dim_pad.py | 170 ++++++++++++++++++ 1 file changed, 170 insertions(+) create mode 100644 tests/pytorch/attention/test_dpa_mla_qkv_head_dim_pad.py diff --git a/tests/pytorch/attention/test_dpa_mla_qkv_head_dim_pad.py b/tests/pytorch/attention/test_dpa_mla_qkv_head_dim_pad.py new file mode 100644 index 00000000000..5bd97838b73 --- /dev/null +++ b/tests/pytorch/attention/test_dpa_mla_qkv_head_dim_pad.py @@ -0,0 +1,170 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Tests for the optional MLA head-dim pad in DotProductAttention. + +Covers: + * `should_pad_qkv_head_dim` decides correctly (native unfused vs padded fused). + * DPA with `head_dim_v > head_dim_qk` runs and produces a V-width output. + * The pad-then-trim is an identity for both `qk > v` and `v > qk`: padding Q/K/V + to a common width and trimming back yields the same result as running with the equal + (padded) shape directly. +""" + +import pathlib +import sys + +import pytest +import torch + +from transformer_engine.pytorch.attention.dot_product_attention import DotProductAttention +from transformer_engine.pytorch.attention.dot_product_attention import ( + dot_product_attention as dpa_module, +) +import transformer_engine.pytorch.attention.dot_product_attention.utils as dpa_utils + +_current_file = pathlib.Path(__file__).resolve() +sys.path = [str(_current_file.parent.parent)] + sys.path +from utils import reset_rng_states + + +def _build_dpa(qk, v, num_heads=4, qkv_format="thd", attn_mask_type="padding_causal"): + return DotProductAttention( + num_attention_heads=num_heads, + kv_channels=(qk, v), + attention_type="self", + attn_mask_type=attn_mask_type, + qkv_format=qkv_format, + ).to(dtype=torch.bfloat16, device="cuda") + + +def _thd_inputs(qk, v, t=32, h=4): + cu = torch.IntTensor([0, 6, 19, 22, t]).cuda() + q = torch.randn(t, h, qk, device="cuda", dtype=torch.bfloat16, requires_grad=True) + k = torch.randn(t, h, qk, device="cuda", dtype=torch.bfloat16) + v = torch.randn(t, h, v, device="cuda", dtype=torch.bfloat16) + return q, k, v, cu + + +def _run_dpa(dpa, q, k, v, cu, max_seqlen=13): + return dpa( + q, + k, + v, + cu_seqlens_q=cu, + cu_seqlens_kv=cu, + max_seqlen_q=max_seqlen, + max_seqlen_kv=max_seqlen, + attn_mask_type="padding_causal", + ) + + +# should_pad_qkv_head_dim +@pytest.mark.parametrize( + "native_unfused,padded_fused,expected", + [ + (False, False, False), # native already fused -> no pad + (True, False, False), # both unfused -> no upgrade -> no pad + (True, True, True), # native unfused, padded fused -> pad + ], +) +def test_should_pad_qkv_head_dim(monkeypatch, native_unfused, padded_fused, expected): + """`should_pad_qkv_head_dim` returns True iff native is unfused and padded is fused.""" + params = dpa_utils.AttentionParams( + qkv_layout="thd_thd_thd", + num_heads=4, + num_gqa_groups=4, + max_seqlen_q=13, + max_seqlen_kv=13, + head_dim_qk=96, + head_dim_v=128, + attn_mask_type="padding_causal", + is_training=True, + qkv_dtype=torch.bfloat16, + ) + + # get_attention_backend returns + # (use_flash, flash_backend, use_fused, fused_backend, use_unfused, available) + native = ( + False, + None, + not native_unfused, + None, + native_unfused, + [False, not native_unfused, native_unfused], + ) + padded = ( + False, + None, + padded_fused, + None, + not padded_fused, + [False, padded_fused, not padded_fused], + ) + + def fake_backend(p): + # native probe: real (mismatched) head_dim_qk/v; padded probe: both = max(qk, v). + # Distinguish by head_dim_qk (native=96, padded=max(96,128)=128). + is_padded = p.head_dim_qk != params.head_dim_qk + return padded if is_padded else native + + monkeypatch.setattr(dpa_utils, "get_attention_backend", fake_backend) + assert dpa_utils.should_pad_qkv_head_dim(params) is expected + + +def test_should_pad_qkv_head_dim_equal_dims(): + """No pad when head_dim_qk == head_dim_v.""" + params = dpa_utils.AttentionParams( + qkv_layout="thd_thd_thd", + num_heads=4, + num_gqa_groups=4, + max_seqlen_q=13, + max_seqlen_kv=13, + head_dim_qk=128, + head_dim_v=128, + attn_mask_type="padding_causal", + is_training=True, + qkv_dtype=torch.bfloat16, + ) + assert dpa_utils.should_pad_qkv_head_dim(params) is False + + +# v > qk end-to-end +@pytest.mark.parametrize("qk,v", [(64, 192), (96, 192)]) +def test_dpa_v_gt_qk_runs(qk, v): + """DPA with head_dim_v > head_dim_qk runs and produces a V-width output.""" + reset_rng_states() + dpa = _build_dpa(qk, v) + q, k, v_t, cu = _thd_inputs(qk, v) + out = _run_dpa(dpa, q, k, v_t, cu) + assert tuple(out.shape) == (32, 4 * v), out.shape # V-width + out.float().sum().backward() # backward must not crash + + +# pad-then-trim is an identity (both directions) +@pytest.mark.parametrize("qk,v", [(192, 128), (64, 192)]) +def test_dpa_mla_pad_is_identity(qk, v): + """Padding Q/K/V to a common width and trimming back equals running with the + equal (padded) shape directly -- for both qk > v and v > qk.""" + reset_rng_states() + m = max(qk, v) + # Reference: run with Q/K/V already at the common width (no pad needed). + dpa_ref = _build_dpa(m, m) + q_ref = torch.randn(32, 4, m, device="cuda", dtype=torch.bfloat16, requires_grad=True) + k_ref = torch.randn(32, 4, m, device="cuda", dtype=torch.bfloat16) + v_ref = torch.randn(32, 4, m, device="cuda", dtype=torch.bfloat16) + cu = torch.IntTensor([0, 6, 19, 22, 32]).cuda() + out_ref = _run_dpa(dpa_ref, q_ref, k_ref, v_ref, cu) + out_ref.float().sum().backward() + + # Test: run with mismatched QK/V, force the pad, then trim. + dpa = _build_dpa(qk, v) + q, k, v_t, _ = _thd_inputs(qk, v) + # Force the pad by calling the same path the forward uses. + q_p, k_p, v_p, _, _ = dpa_module._pad_qkv_head_dim(q, k, v_t) + assert q_p.shape[-1] == k_p.shape[-1] == v_p.shape[-1] == m + out = _run_dpa(dpa, q_p, k_p, v_p, cu) + # Trim back to the original V width. + out = dpa_module._trim_output(out, 4, m, v) + torch.testing.assert_close(out, out_ref, atol=1e-2, rtol=1e-2) From 748957eea539cc41b8670064cb91cc9ba9a6c906 Mon Sep 17 00:00:00 2001 From: janEbert Date: Tue, 11 Aug 2026 01:09:48 +0200 Subject: [PATCH 03/30] Fix tests Signed-off-by: janEbert --- .../test_dpa_mla_qkv_head_dim_pad.py | 40 +++++++++++-------- 1 file changed, 24 insertions(+), 16 deletions(-) diff --git a/tests/pytorch/attention/test_dpa_mla_qkv_head_dim_pad.py b/tests/pytorch/attention/test_dpa_mla_qkv_head_dim_pad.py index 5bd97838b73..420dac4ac71 100644 --- a/tests/pytorch/attention/test_dpa_mla_qkv_head_dim_pad.py +++ b/tests/pytorch/attention/test_dpa_mla_qkv_head_dim_pad.py @@ -7,11 +7,12 @@ Covers: * `should_pad_qkv_head_dim` decides correctly (native unfused vs padded fused). * DPA with `head_dim_v > head_dim_qk` runs and produces a V-width output. - * The pad-then-trim is an identity for both `qk > v` and `v > qk`: padding Q/K/V - to a common width and trimming back yields the same result as running with the equal - (padded) shape directly. + * The pad-then-trim is an identity for both `qk > v` and `v > qk`: padding Q/K/V to the + wider head dim, running with the equal (padded) shape, and trimming back equals the + native mismatched-dim run. """ +import math import pathlib import sys @@ -29,13 +30,16 @@ from utils import reset_rng_states -def _build_dpa(qk, v, num_heads=4, qkv_format="thd", attn_mask_type="padding_causal"): +def _build_dpa( + qk, v, num_heads=4, qkv_format="thd", attn_mask_type="padding_causal", softmax_scale=None +): return DotProductAttention( num_attention_heads=num_heads, kv_channels=(qk, v), attention_type="self", attn_mask_type=attn_mask_type, qkv_format=qkv_format, + softmax_scale=softmax_scale, ).to(dtype=torch.bfloat16, device="cuda") @@ -145,26 +149,30 @@ def test_dpa_v_gt_qk_runs(qk, v): # pad-then-trim is an identity (both directions) @pytest.mark.parametrize("qk,v", [(192, 128), (64, 192)]) def test_dpa_mla_pad_is_identity(qk, v): - """Padding Q/K/V to a common width and trimming back equals running with the - equal (padded) shape directly -- for both qk > v and v > qk.""" + """Pad-then-trim is an identity: padding Q/K/V to the wider head dim, running with the equal + (padded) shape, and trimming back equals the native mismatched-dim run -- for both qk > v and v + > qk. Both runs use the same `softmax_scale` (`1/sqrt(qk)`) that the production forward keeps + when padding. + """ reset_rng_states() m = max(qk, v) - # Reference: run with Q/K/V already at the common width (no pad needed). - dpa_ref = _build_dpa(m, m) - q_ref = torch.randn(32, 4, m, device="cuda", dtype=torch.bfloat16, requires_grad=True) - k_ref = torch.randn(32, 4, m, device="cuda", dtype=torch.bfloat16) - v_ref = torch.randn(32, 4, m, device="cuda", dtype=torch.bfloat16) + scale = 1.0 / math.sqrt(qk) cu = torch.IntTensor([0, 6, 19, 22, 32]).cuda() - out_ref = _run_dpa(dpa_ref, q_ref, k_ref, v_ref, cu) - out_ref.float().sum().backward() - # Test: run with mismatched QK/V, force the pad, then trim. - dpa = _build_dpa(qk, v) + # Reference: native mismatched-dim run (the production forward; it pads internally + # only when should_pad_qkv_head_dim upgrades the selected backend). + dpa_ref = _build_dpa(qk, v) # softmax_scale defaults to 1/sqrt(qk) q, k, v_t, _ = _thd_inputs(qk, v) - # Force the pad by calling the same path the forward uses. + out_ref = _run_dpa(dpa_ref, q, k, v_t, cu) + assert tuple(out_ref.shape) == (32, 4 * v), out_ref.shape + + # Test: manually pad to the common width, run with the equal (padded) shape, trim. + # Same softmax_scale as the reference so pad-then-trim is a true identity. + dpa = _build_dpa(m, m, softmax_scale=scale) q_p, k_p, v_p, _, _ = dpa_module._pad_qkv_head_dim(q, k, v_t) assert q_p.shape[-1] == k_p.shape[-1] == v_p.shape[-1] == m out = _run_dpa(dpa, q_p, k_p, v_p, cu) # Trim back to the original V width. out = dpa_module._trim_output(out, 4, m, v) torch.testing.assert_close(out, out_ref, atol=1e-2, rtol=1e-2) + out.float().sum().backward() # padded path backward must not crash From 64c347fc1601b4578b8d7c0343255a45461b1dda Mon Sep 17 00:00:00 2001 From: janEbert Date: Tue, 11 Aug 2026 01:21:09 +0200 Subject: [PATCH 04/30] Add memoization Signed-off-by: janEbert --- .../test_dpa_mla_qkv_head_dim_pad.py | 48 +++++++++++++++++++ .../attention/dot_product_attention/utils.py | 48 +++++++++++++------ 2 files changed, 82 insertions(+), 14 deletions(-) diff --git a/tests/pytorch/attention/test_dpa_mla_qkv_head_dim_pad.py b/tests/pytorch/attention/test_dpa_mla_qkv_head_dim_pad.py index 420dac4ac71..2cca7e477ce 100644 --- a/tests/pytorch/attention/test_dpa_mla_qkv_head_dim_pad.py +++ b/tests/pytorch/attention/test_dpa_mla_qkv_head_dim_pad.py @@ -114,6 +114,11 @@ def fake_backend(p): return padded if is_padded else native monkeypatch.setattr(dpa_utils, "get_attention_backend", fake_backend) + # The decision is memoized on `attention_params` (see + # `_should_pad_qkv_head_dim_cache`); reset the cache so each parametrization + # re-probes the freshly monkeypatched backend instead of returning a stale result. + dpa_utils._should_pad_qkv_head_dim_cache["attention_params"] = None + dpa_utils._should_pad_qkv_head_dim_cache["result"] = None assert dpa_utils.should_pad_qkv_head_dim(params) is expected @@ -134,6 +139,49 @@ def test_should_pad_qkv_head_dim_equal_dims(): assert dpa_utils.should_pad_qkv_head_dim(params) is False +def test_should_pad_qkv_head_dim_is_memoized(monkeypatch): + """`should_pad_qkv_head_dim` memoizes on the native (pre-pad) params: a second call + with an equal config skips the `get_attention_backend` probes entirely, even if the + caller mutated the first params object in place -- as the production forward does when + it pads `head_dim_qk`/`head_dim_v` after this returns.""" + base = dict( + qkv_layout="thd_thd_thd", + num_heads=4, + num_gqa_groups=4, + max_seqlen_q=13, + max_seqlen_kv=13, + attn_mask_type="padding_causal", + is_training=True, + qkv_dtype=torch.bfloat16, + ) + calls = {"n": 0} + + def fake_backend(p): + calls["n"] += 1 + # native probe has head_dim_qk=96; padded probe has head_dim_qk=128. + if p.head_dim_qk != 96: + return (False, None, True, None, False, [False, True, False]) # fused -> pad + return (False, None, False, None, True, [False, False, True]) # unfused native + + monkeypatch.setattr(dpa_utils, "get_attention_backend", fake_backend) + dpa_utils._should_pad_qkv_head_dim_cache["attention_params"] = None + dpa_utils._should_pad_qkv_head_dim_cache["result"] = None + + params = dpa_utils.AttentionParams(head_dim_qk=96, head_dim_v=128, **base) + assert dpa_utils.should_pad_qkv_head_dim(params) is True + assert calls["n"] == 2 # one native + one padded probe + + # Simulate the production forward mutating the live params in place after the call. + params.head_dim_qk = 128 + params.head_dim_v = 128 + + # A fresh native params with the same config must still hit the memo (the key is a + # copy, not the mutated live object) and must not re-probe. + params2 = dpa_utils.AttentionParams(head_dim_qk=96, head_dim_v=128, **base) + assert dpa_utils.should_pad_qkv_head_dim(params2) is True + assert calls["n"] == 2 # cache hit: no new probes + + # v > qk end-to-end @pytest.mark.parametrize("qk,v", [(64, 192), (96, 192)]) def test_dpa_v_gt_qk_runs(qk, v): diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py index fa79e6a6a97..7be331dd9a0 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py @@ -73,6 +73,12 @@ _cu_seqlens_cache = {} +# Global var for MLA padding cache. +_should_pad_qkv_head_dim_cache: Dict[str, Any] = { + "attention_params": None, + "result": None, +} + class AttentionLogging: """ @@ -1986,7 +1992,7 @@ 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", + attention_params: AttentionParams, ) -> bool: """Decide whether padding Q/K/V to a common head dim upgrades the attention backend. @@ -1997,6 +2003,11 @@ def should_pad_qkv_head_dim( shape selects a fused backend (`FlashAttention` or `FusedAttention`), i.e. the only case in which the pad actually upgrades the kernel instead of being pure overhead. + The decision depends only on `attention_params` (and the process-global backend availability, + which is fixed for the lifetime of the process), so it is memoized across calls: only the first + call with a given native config pays for the two `get_attention_backend` probes; subsequent + calls with an equal config return the cached result. + Parameters ---------- attention_params : AttentionParams @@ -2009,23 +2020,32 @@ def should_pad_qkv_head_dim( """ if attention_params.head_dim_qk == attention_params.head_dim_v: return False + cached_params = _should_pad_qkv_head_dim_cache["attention_params"] + if cached_params is not None and cached_params == attention_params: + return _should_pad_qkv_head_dim_cache["result"] native_backend = get_attention_backend(attention_params) native_use_unfused_attention = native_backend[4] # No point padding if native path is already fused. if not native_use_unfused_attention: - return False - # Probe on a copy so the caller's `attention_params` are never mutated, even if - # `get_attention_backend` raises. - 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] - return bool(padded_use_flash_attention or padded_use_fused_attention) + result = False + else: + # Probe on a copy so the caller's `attention_params` are never mutated, even if + # `get_attention_backend` raises. + 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] + result = bool(padded_use_flash_attention or padded_use_fused_attention) + # Store a shallow copy as the key: the caller pads `attention_params` in place after this + # returns, so storing the live object would make every subsequent call a miss. + _should_pad_qkv_head_dim_cache["attention_params"] = replace(attention_params) + _should_pad_qkv_head_dim_cache["result"] = result + return result @torch.no_grad() From 6bdac9b8b821865719cca248215b101dffa68bdb Mon Sep 17 00:00:00 2001 From: janEbert Date: Tue, 11 Aug 2026 01:28:59 +0200 Subject: [PATCH 05/30] Add test to CI Signed-off-by: janEbert --- qa/L0_pytorch_unittest/test.sh | 1 + 1 file changed, 1 insertion(+) diff --git a/qa/L0_pytorch_unittest/test.sh b/qa/L0_pytorch_unittest/test.sh index 2a486450a20..09056249a63 100644 --- a/qa/L0_pytorch_unittest/test.sh +++ b/qa/L0_pytorch_unittest/test.sh @@ -74,6 +74,7 @@ python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_fused_mla_q_upro python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_kv_cache.xml $TE_PATH/tests/pytorch/attention/test_kv_cache.py || test_fail "test_kv_cache.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_cu_seqlens_cache.xml $TE_PATH/tests/pytorch/attention/test_cu_seqlens_cache.py || test_fail "test_cu_seqlens_cache.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_mixed_thd_attention.xml $TE_PATH/tests/pytorch/attention/test_mixed_thd_attention.py || test_fail "test_mixed_thd_attention.py" +python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_dpa_mla_qkv_head_dim_pad.xml $TE_PATH/tests/pytorch/attention/test_dpa_mla_qkv_head_dim_pad.py || test_fail "test_dpa_mla_qkv_head_dim_pad.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_hf_integration.xml $TE_PATH/tests/pytorch/test_hf_integration.py || test_fail "test_hf_integration.py" export NVTE_TEST_CHECKPOINT_ARTIFACT_PATH=$TE_PATH/artifacts/tests/pytorch/test_checkpoint if [ ! -d "$NVTE_TEST_CHECKPOINT_ARTIFACT_PATH" ]; then From cf0e40122febdb795c8c5865637bdaf68a7e6ac6 Mon Sep 17 00:00:00 2001 From: janEbert Date: Tue, 11 Aug 2026 02:27:28 +0200 Subject: [PATCH 06/30] Skip padding when using fused+packed FP8 path Signed-off-by: janEbert --- .../test_dpa_mla_qkv_head_dim_pad.py | 68 +++++++++++++++++++ .../dot_product_attention.py | 5 ++ .../attention/dot_product_attention/utils.py | 35 ++++++++++ 3 files changed, 108 insertions(+) diff --git a/tests/pytorch/attention/test_dpa_mla_qkv_head_dim_pad.py b/tests/pytorch/attention/test_dpa_mla_qkv_head_dim_pad.py index 2cca7e477ce..88434a3791d 100644 --- a/tests/pytorch/attention/test_dpa_mla_qkv_head_dim_pad.py +++ b/tests/pytorch/attention/test_dpa_mla_qkv_head_dim_pad.py @@ -19,6 +19,8 @@ import pytest import torch +import transformer_engine.pytorch as te +from transformer_engine.common import recipe from transformer_engine.pytorch.attention.dot_product_attention import DotProductAttention from transformer_engine.pytorch.attention.dot_product_attention import ( dot_product_attention as dpa_module, @@ -224,3 +226,69 @@ def test_dpa_mla_pad_is_identity(qk, v): out = dpa_module._trim_output(out, 4, m, v) torch.testing.assert_close(out, out_ref, atol=1e-2, rtol=1e-2) out.float().sum().backward() # padded path backward must not crash + + +def test_fp8_packed_skips_qkv_head_dim_pad(): + """`fp8_packed_skips_qkv_head_dim_pad` returns True (skip the pad) only in the one + unsafe combination -- FP8 DPA active AND packed inputs -- and False (allow the pad) + in every safe combination (bf16 regardless of packing; FP8 DPA + unpacked).""" + fp8_dpa_on = {"recipe": recipe.DelayedScaling(fp8_dpa=True)} + fp8_dpa_off = {"recipe": recipe.DelayedScaling(fp8_dpa=False)} + + def params(fp8, fp8_meta): + return dpa_utils.AttentionParams( + qkv_layout="thd_thd_thd", + num_heads=4, + num_gqa_groups=4, + max_seqlen_q=13, + max_seqlen_kv=13, + head_dim_qk=96, + head_dim_v=128, + attn_mask_type="padding_causal", + is_training=True, + qkv_dtype=torch.bfloat16, + fp8=fp8, + fp8_meta=fp8_meta, + ) + + # bf16 (fp8=False): pad is safe & beneficial regardless of packing. + assert dpa_utils.fp8_packed_skips_qkv_head_dim_pad(params(False, None), False) is False + assert dpa_utils.fp8_packed_skips_qkv_head_dim_pad(params(False, None), True) is False + # FP8 autocast but fp8_dpa=False (attention itself runs in bf16): pad is safe. + assert dpa_utils.fp8_packed_skips_qkv_head_dim_pad(params(True, fp8_dpa_off), False) is False + assert dpa_utils.fp8_packed_skips_qkv_head_dim_pad(params(True, fp8_dpa_off), True) is False + # FP8 DPA + unpacked: pad is safe (combine_and_quantize re-derives from padded views). + assert dpa_utils.fp8_packed_skips_qkv_head_dim_pad(params(True, fp8_dpa_on), False) is False + # FP8 DPA + packed: the only unsafe combination -> skip the pad. + assert dpa_utils.fp8_packed_skips_qkv_head_dim_pad(params(True, fp8_dpa_on), True) is True + # Defensive: fp8=True but no recipe -> treat as "do not skip" (pad allowed). + assert dpa_utils.fp8_packed_skips_qkv_head_dim_pad(params(True, {}), True) is False + assert dpa_utils.fp8_packed_skips_qkv_head_dim_pad(params(True, None), True) is False + + +# FP8 DPA + unpacked MLA: the guard must not over-block the safe beneficial case. +@pytest.mark.parametrize("qk,v", [(192, 128), (64, 192)]) +def test_dpa_mla_pad_fp8_dpa_unpacked(qk, v, monkeypatch): + """Regression test for the FP8+packed guard: under FP8 DPA with *unpacked* mismatched MLA + inputs, the pad is safe and beneficial (`combine_and_quantize` re-derives from the padded + views, so the kernel sees the equal padded dims), so the guard must NOT skip it. Force + the conditions that make `should_pad_qkv_head_dim` return True (native unfused via + the `NVTE_UnfusedDPA_Emulate_FP8` env) and assert the forward runs and yields a + correct V-width output.""" + monkeypatch.setenv("NVTE_UnfusedDPA_Emulate_FP8", "1") + reset_rng_states() + fp8_recipe = recipe.DelayedScaling( + margin=0, + fp8_format=recipe.Format.HYBRID, + amax_history_len=1, + amax_compute_algo="most_recent", + fp8_dpa=True, + fp8_mha=False, + ) + dpa = _build_dpa(qk, v) + q, k, v_t, cu = _thd_inputs(qk, v) + with te.fp8_autocast(enabled=True, fp8_recipe=fp8_recipe): + out = _run_dpa(dpa, q, k, v_t, cu) + assert tuple(out.shape) == (32, 4 * v), out.shape + with te.fp8_autocast(enabled=True, fp8_recipe=fp8_recipe): + out.float().sum().backward() # FP8 padded path backward must not crash 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 0fd89bf5d34..317e2e7956d 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 @@ -2854,6 +2854,11 @@ def forward( and head_dim_qk != head_dim_v and value_layer is not None and not isinstance(value_layer, Float8TensorStorage) + # Skip padding when packed inputs are used with the FP8 fused path, which would lead + # to using unpadded data with a kernel that assumes padded data. + and not dpa_utils.fp8_packed_skips_qkv_head_dim_pad( + attention_params, qkv_layer is not None or kv_layer is not None + ) and dpa_utils.should_pad_qkv_head_dim(attention_params) ): # Pad Q/K/V to the wider head dim so a fused backend can run. diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py index 7be331dd9a0..9a689ad0056 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py @@ -2048,6 +2048,41 @@ def should_pad_qkv_head_dim( return result +def fp8_packed_skips_qkv_head_dim_pad( + attention_params: AttentionParams, + has_packed_qkv: bool, +) -> bool: + """Whether the optional QKV head-dim pad must be skipped because FP8 DPA is active and the + caller passed packed QKV/KV buffers. + + The pad in `DotProductAttention.forward` pads only the unpacked q/k/v views, not the packed + `qkv_layer`/`kv_layer` buffers. In the FP8 fused path, `combine_and_quantize` consumes those + packed buffers directly, bypassing the padded views -- so padding would feed the kernel the + original mismatched-dim packed data while the backend was selected for the equal padded dims. + This is only unsafe when FP8 DPA is active AND packed inputs are present; every other + combination is safe (FP8 DPA + unpacked re-derives from the padded views; bf16 + packed ignores + the packed buffers in `fused_attn_fwd`). + + 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. + + Returns + ------- + bool + True iff the pad should be skipped (not applied). + """ + if not attention_params.fp8: + return False + recipe = (attention_params.fp8_meta or {}).get("recipe", None) + if recipe is None or not getattr(recipe, "fp8_dpa", False): + return False + return has_packed_qkv + + @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. From 0546cdf45c9ccf5e8b3b41c04955787adc72e049 Mon Sep 17 00:00:00 2001 From: janEbert Date: Tue, 11 Aug 2026 23:27:56 +0200 Subject: [PATCH 07/30] Refactor optional pad if-guard Signed-off-by: janEbert --- .../dot_product_attention.py | 32 ++++++++++++------- 1 file changed, 21 insertions(+), 11 deletions(-) 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 317e2e7956d..2f411c5e854 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 @@ -423,6 +423,25 @@ def _delayed_scaling_recipe() -> Optional[DelayedScaling]: __all__ = ["DotProductAttention"] +def _should_pad_qkv_head_dim( + head_dim_qk, head_dim_v, value_layer, qkv_layer, kv_layer, attention_params +): + has_packed_inputs = qkv_layer is not None or kv_layer is not None + is_shape_pad_eligible = ( + not is_in_onnx_export_mode() + and head_dim_qk != head_dim_v + and value_layer is not None + and not isinstance(value_layer, Float8TensorStorage) + ) + # Skip padding when packed inputs are used with the FP8 fused path, which would lead to using + # unpadded data with a kernel that assumes padded data. + return ( + is_shape_pad_eligible + and not dpa_utils.fp8_packed_skips_qkv_head_dim_pad(attention_params, has_packed_inputs) + and dpa_utils.should_pad_qkv_head_dim(attention_params) + ) + + def _pad_qkv_head_dim(query_layer, key_layer, value_layer): """Pad Q/K/V to the same head dimension for FlashAttention 2 MLA.""" orig_head_dim_qk = query_layer.shape[-1] @@ -2849,17 +2868,8 @@ def forward( # never changes the result, only which kernel runs. qkv_head_pad = False orig_head_dim_v = head_dim_v - if ( - not is_in_onnx_export_mode() - and head_dim_qk != head_dim_v - and value_layer is not None - and not isinstance(value_layer, Float8TensorStorage) - # Skip padding when packed inputs are used with the FP8 fused path, which would lead - # to using unpadded data with a kernel that assumes padded data. - and not dpa_utils.fp8_packed_skips_qkv_head_dim_pad( - attention_params, qkv_layer is not None or kv_layer is not None - ) - and dpa_utils.should_pad_qkv_head_dim(attention_params) + if _should_pad_qkv_head_dim( + head_dim_qk, head_dim_v, value_layer, qkv_layer, kv_layer, attention_params ): # Pad Q/K/V to the wider head dim so a fused backend can run. query_layer, key_layer, value_layer, _, _ = _pad_qkv_head_dim( From 1416805bce0b300241828799dcee4108e37dd73b Mon Sep 17 00:00:00 2001 From: janEbert Date: Tue, 11 Aug 2026 23:28:23 +0200 Subject: [PATCH 08/30] Merge output trim code paths Signed-off-by: janEbert --- .../dot_product_attention.py | 30 ++++++++----------- 1 file changed, 12 insertions(+), 18 deletions(-) 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 2f411c5e854..3fc213b835a 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 @@ -2868,6 +2868,8 @@ def forward( # never changes the result, only which kernel runs. qkv_head_pad = False orig_head_dim_v = head_dim_v + orig_qk_dim = None + orig_v_dim = None if _should_pad_qkv_head_dim( head_dim_qk, head_dim_v, value_layer, qkv_layer, kv_layer, attention_params ): @@ -3006,13 +3008,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) or qkv_head_pad: - attn_out = _trim_output( - attn_out, num_attention_heads, head_dim_qk, orig_head_dim_v - ) - 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): @@ -3103,13 +3100,8 @@ def forward( packed_kv=kv_layer, bf16_backward=bf16_backward, ) - if qkv_head_pad: - attn_out = _trim_output( - attn_out, num_attention_heads, head_dim_qk, orig_head_dim_v - ) - return attn_out - if use_unfused_attention: + elif use_unfused_attention: allow_emulation = ( os.getenv("NVTE_UnfusedDPA_Emulate_FP8", "0") == "1" or is_in_onnx_export_mode() ) @@ -3166,9 +3158,11 @@ def forward( quantizers=self.quantizers, fp8_output=fp8_output, ) - if qkv_head_pad: - attn_out = _trim_output( - attn_out, num_attention_heads, head_dim_qk, orig_head_dim_v - ) - return attn_out - return None + else: + attn_out = None + + if attn_out is not None and ( + orig_qk_dim is not None and orig_qk_dim > orig_v_dim or qkv_head_pad + ): + attn_out = _trim_output(attn_out, num_attention_heads, head_dim_qk, orig_head_dim_v) + return attn_out From 78c1348c00e8f9b4cbcc377c0d03940b9e80ea61 Mon Sep 17 00:00:00 2001 From: janEbert Date: Tue, 11 Aug 2026 23:31:36 +0200 Subject: [PATCH 09/30] Reduce verbosity Signed-off-by: janEbert --- .../dot_product_attention/dot_product_attention.py | 13 +++---------- .../attention/dot_product_attention/utils.py | 6 ++---- 2 files changed, 5 insertions(+), 14 deletions(-) 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 3fc213b835a..d4de31bafdb 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 @@ -2861,11 +2861,7 @@ def forward( pad_between_seqs=pad_between_seqs, ) - # Optional MLA head-dim pad: padding Q/K/V to a common (the wider) head dim can upgrade - # the selected backend off the slow `UnfusedDotProductAttention` for certain setups. - # Probe both shapes and pad only when padding escapes the unfused path (and leaving the - # dims native would land on the unfused path). The pad-then-trim is an identity, so this - # never changes the result, only which kernel runs. + # Optionally pad inputs when it leads to faster kernel selection. qkv_head_pad = False orig_head_dim_v = head_dim_v orig_qk_dim = None @@ -2873,7 +2869,6 @@ def forward( if _should_pad_qkv_head_dim( head_dim_qk, head_dim_v, value_layer, qkv_layer, kv_layer, attention_params ): - # Pad Q/K/V to the wider head dim so a fused backend can run. query_layer, key_layer, value_layer, _, _ = _pad_qkv_head_dim( query_layer, key_layer, value_layer ) @@ -3159,10 +3154,8 @@ def forward( fp8_output=fp8_output, ) else: - attn_out = None + return None - if attn_out is not None and ( - orig_qk_dim is not None and orig_qk_dim > orig_v_dim or qkv_head_pad - ): + if orig_qk_dim is not None and orig_qk_dim > orig_v_dim or qkv_head_pad: attn_out = _trim_output(attn_out, num_attention_heads, head_dim_qk, orig_head_dim_v) 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 9a689ad0056..a103d94f4b9 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py @@ -2029,8 +2029,6 @@ def should_pad_qkv_head_dim( if not native_use_unfused_attention: result = False else: - # Probe on a copy so the caller's `attention_params` are never mutated, even if - # `get_attention_backend` raises. padded_head_dim = max(attention_params.head_dim_qk, attention_params.head_dim_v) padded_params = replace( attention_params, @@ -2041,8 +2039,8 @@ def should_pad_qkv_head_dim( padded_use_flash_attention = padded_backend[0] padded_use_fused_attention = padded_backend[2] result = bool(padded_use_flash_attention or padded_use_fused_attention) - # Store a shallow copy as the key: the caller pads `attention_params` in place after this - # returns, so storing the live object would make every subsequent call a miss. + # 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["result"] = result return result From 69f7044bfd8a4917eec51fa77b403af96e3bbe3c Mon Sep 17 00:00:00 2001 From: janEbert Date: Wed, 12 Aug 2026 13:38:04 +0200 Subject: [PATCH 10/30] Explicitly mention padding target and limits Signed-off-by: janEbert --- .../attention/dot_product_attention/dot_product_attention.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) 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 d4de31bafdb..c70c76cbb9e 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 @@ -2861,7 +2861,9 @@ def forward( pad_between_seqs=pad_between_seqs, ) - # Optionally pad inputs when it leads to faster kernel selection. + # Optionally pad inputs to `max(head_dim_qk, head_dim_v)` when it leads to faster kernel + # selection. Note that the padded tensors could still be unsupported inputs shapes by + # {Flash,Fused}Attention. qkv_head_pad = False orig_head_dim_v = head_dim_v orig_qk_dim = None From 60a69a2b0f85f3a3cd3b0abc2fdd7f43104a4862 Mon Sep 17 00:00:00 2001 From: janEbert Date: Wed, 12 Aug 2026 12:09:32 +0000 Subject: [PATCH 11/30] Merge input pad code paths Signed-off-by: janEbert --- .../dot_product_attention.py | 26 +++---------------- .../attention/dot_product_attention/utils.py | 24 ++++++++--------- 2 files changed, 15 insertions(+), 35 deletions(-) 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 c70c76cbb9e..320244a0921 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 ( @@ -2861,13 +2860,11 @@ def forward( pad_between_seqs=pad_between_seqs, ) - # Optionally pad inputs to `max(head_dim_qk, head_dim_v)` when it leads to faster kernel - # selection. Note that the padded tensors could still be unsupported inputs shapes by - # {Flash,Fused}Attention. + # 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 inputs shapes by `{Flash,Fused}Attention`. qkv_head_pad = False orig_head_dim_v = head_dim_v - orig_qk_dim = None - orig_v_dim = None if _should_pad_qkv_head_dim( head_dim_qk, head_dim_v, value_layer, qkv_layer, kv_layer, attention_params ): @@ -2952,21 +2949,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, @@ -3158,6 +3140,6 @@ def forward( else: return None - if orig_qk_dim is not None and orig_qk_dim > orig_v_dim or qkv_head_pad: + if qkv_head_pad: attn_out = _trim_output(attn_out, num_attention_heads, head_dim_qk, orig_head_dim_v) 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 a103d94f4b9..7c87f0f8baa 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py @@ -1994,19 +1994,14 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt def should_pad_qkv_head_dim( attention_params: AttentionParams, ) -> bool: - """Decide whether padding Q/K/V to a common head dim upgrades the attention backend. + """Decide whether to pad Q/K/V to a common head dim. - Padding Q, K, and V to the wider of the QK and V head dims is an identity -- the padded head - dims contribute zero to the Q*K dot product (so scores and softmax are unchanged) and the padded - V columns are sums of zeros (so trimming them back changes nothing). This returns `True` iff the - native mismatched shape selects the slow `UnfusedDotProductAttention` while the padded equal - shape selects a fused backend (`FlashAttention` or `FusedAttention`), i.e. the only case in - which the pad actually upgrades the kernel instead of being pure overhead. + This returns `True` in two cases: + 1. when the native mismatched shape would otherwise land on a backend that cannot take + mismatched head dims (FlashAttention-2). + 2. when padding results in a faster kernel - The decision depends only on `attention_params` (and the process-global backend availability, - which is fixed for the lifetime of the process), so it is memoized across calls: only the first - call with a given native config pays for the two `get_attention_backend` probes; subsequent - calls with an equal config return the cached result. + The selection is memoized for the given `attention_params`. Parameters ---------- @@ -2024,10 +2019,13 @@ def should_pad_qkv_head_dim( if cached_params is not None and cached_params == attention_params: return _should_pad_qkv_head_dim_cache["result"] native_backend = get_attention_backend(attention_params) + native_flash_attention_backend = native_backend[1] native_use_unfused_attention = native_backend[4] # No point padding if native path is already fused. if not native_use_unfused_attention: - result = False + # Native already lands on FlashAttention or FusedAttention, so no need to pad. However, + # FlashAttention-2 cannot take mismatched QK/V head dims, so we still pad in that case. + result = native_flash_attention_backend == FlashAttentionUtils.version else: padded_head_dim = max(attention_params.head_dim_qk, attention_params.head_dim_v) padded_params = replace( @@ -2038,7 +2036,7 @@ def should_pad_qkv_head_dim( padded_backend = get_attention_backend(padded_params) padded_use_flash_attention = padded_backend[0] padded_use_fused_attention = padded_backend[2] - result = bool(padded_use_flash_attention or padded_use_fused_attention) + result = padded_use_flash_attention or padded_use_fused_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) From 7559cada4b215d8a92dac6c4577aecc9031b8f7a Mon Sep 17 00:00:00 2001 From: janEbert Date: Thu, 13 Aug 2026 21:48:42 +0000 Subject: [PATCH 12/30] Remove standalone MLA pad tests Signed-off-by: janEbert --- .../test_dpa_mla_qkv_head_dim_pad.py | 294 ------------------ 1 file changed, 294 deletions(-) delete mode 100644 tests/pytorch/attention/test_dpa_mla_qkv_head_dim_pad.py diff --git a/tests/pytorch/attention/test_dpa_mla_qkv_head_dim_pad.py b/tests/pytorch/attention/test_dpa_mla_qkv_head_dim_pad.py deleted file mode 100644 index 88434a3791d..00000000000 --- a/tests/pytorch/attention/test_dpa_mla_qkv_head_dim_pad.py +++ /dev/null @@ -1,294 +0,0 @@ -# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# -# See LICENSE for license information. - -"""Tests for the optional MLA head-dim pad in DotProductAttention. - -Covers: - * `should_pad_qkv_head_dim` decides correctly (native unfused vs padded fused). - * DPA with `head_dim_v > head_dim_qk` runs and produces a V-width output. - * The pad-then-trim is an identity for both `qk > v` and `v > qk`: padding Q/K/V to the - wider head dim, running with the equal (padded) shape, and trimming back equals the - native mismatched-dim run. -""" - -import math -import pathlib -import sys - -import pytest -import torch - -import transformer_engine.pytorch as te -from transformer_engine.common import recipe -from transformer_engine.pytorch.attention.dot_product_attention import DotProductAttention -from transformer_engine.pytorch.attention.dot_product_attention import ( - dot_product_attention as dpa_module, -) -import transformer_engine.pytorch.attention.dot_product_attention.utils as dpa_utils - -_current_file = pathlib.Path(__file__).resolve() -sys.path = [str(_current_file.parent.parent)] + sys.path -from utils import reset_rng_states - - -def _build_dpa( - qk, v, num_heads=4, qkv_format="thd", attn_mask_type="padding_causal", softmax_scale=None -): - return DotProductAttention( - num_attention_heads=num_heads, - kv_channels=(qk, v), - attention_type="self", - attn_mask_type=attn_mask_type, - qkv_format=qkv_format, - softmax_scale=softmax_scale, - ).to(dtype=torch.bfloat16, device="cuda") - - -def _thd_inputs(qk, v, t=32, h=4): - cu = torch.IntTensor([0, 6, 19, 22, t]).cuda() - q = torch.randn(t, h, qk, device="cuda", dtype=torch.bfloat16, requires_grad=True) - k = torch.randn(t, h, qk, device="cuda", dtype=torch.bfloat16) - v = torch.randn(t, h, v, device="cuda", dtype=torch.bfloat16) - return q, k, v, cu - - -def _run_dpa(dpa, q, k, v, cu, max_seqlen=13): - return dpa( - q, - k, - v, - cu_seqlens_q=cu, - cu_seqlens_kv=cu, - max_seqlen_q=max_seqlen, - max_seqlen_kv=max_seqlen, - attn_mask_type="padding_causal", - ) - - -# should_pad_qkv_head_dim -@pytest.mark.parametrize( - "native_unfused,padded_fused,expected", - [ - (False, False, False), # native already fused -> no pad - (True, False, False), # both unfused -> no upgrade -> no pad - (True, True, True), # native unfused, padded fused -> pad - ], -) -def test_should_pad_qkv_head_dim(monkeypatch, native_unfused, padded_fused, expected): - """`should_pad_qkv_head_dim` returns True iff native is unfused and padded is fused.""" - params = dpa_utils.AttentionParams( - qkv_layout="thd_thd_thd", - num_heads=4, - num_gqa_groups=4, - max_seqlen_q=13, - max_seqlen_kv=13, - head_dim_qk=96, - head_dim_v=128, - attn_mask_type="padding_causal", - is_training=True, - qkv_dtype=torch.bfloat16, - ) - - # get_attention_backend returns - # (use_flash, flash_backend, use_fused, fused_backend, use_unfused, available) - native = ( - False, - None, - not native_unfused, - None, - native_unfused, - [False, not native_unfused, native_unfused], - ) - padded = ( - False, - None, - padded_fused, - None, - not padded_fused, - [False, padded_fused, not padded_fused], - ) - - def fake_backend(p): - # native probe: real (mismatched) head_dim_qk/v; padded probe: both = max(qk, v). - # Distinguish by head_dim_qk (native=96, padded=max(96,128)=128). - is_padded = p.head_dim_qk != params.head_dim_qk - return padded if is_padded else native - - monkeypatch.setattr(dpa_utils, "get_attention_backend", fake_backend) - # The decision is memoized on `attention_params` (see - # `_should_pad_qkv_head_dim_cache`); reset the cache so each parametrization - # re-probes the freshly monkeypatched backend instead of returning a stale result. - dpa_utils._should_pad_qkv_head_dim_cache["attention_params"] = None - dpa_utils._should_pad_qkv_head_dim_cache["result"] = None - assert dpa_utils.should_pad_qkv_head_dim(params) is expected - - -def test_should_pad_qkv_head_dim_equal_dims(): - """No pad when head_dim_qk == head_dim_v.""" - params = dpa_utils.AttentionParams( - qkv_layout="thd_thd_thd", - num_heads=4, - num_gqa_groups=4, - max_seqlen_q=13, - max_seqlen_kv=13, - head_dim_qk=128, - head_dim_v=128, - attn_mask_type="padding_causal", - is_training=True, - qkv_dtype=torch.bfloat16, - ) - assert dpa_utils.should_pad_qkv_head_dim(params) is False - - -def test_should_pad_qkv_head_dim_is_memoized(monkeypatch): - """`should_pad_qkv_head_dim` memoizes on the native (pre-pad) params: a second call - with an equal config skips the `get_attention_backend` probes entirely, even if the - caller mutated the first params object in place -- as the production forward does when - it pads `head_dim_qk`/`head_dim_v` after this returns.""" - base = dict( - qkv_layout="thd_thd_thd", - num_heads=4, - num_gqa_groups=4, - max_seqlen_q=13, - max_seqlen_kv=13, - attn_mask_type="padding_causal", - is_training=True, - qkv_dtype=torch.bfloat16, - ) - calls = {"n": 0} - - def fake_backend(p): - calls["n"] += 1 - # native probe has head_dim_qk=96; padded probe has head_dim_qk=128. - if p.head_dim_qk != 96: - return (False, None, True, None, False, [False, True, False]) # fused -> pad - return (False, None, False, None, True, [False, False, True]) # unfused native - - monkeypatch.setattr(dpa_utils, "get_attention_backend", fake_backend) - dpa_utils._should_pad_qkv_head_dim_cache["attention_params"] = None - dpa_utils._should_pad_qkv_head_dim_cache["result"] = None - - params = dpa_utils.AttentionParams(head_dim_qk=96, head_dim_v=128, **base) - assert dpa_utils.should_pad_qkv_head_dim(params) is True - assert calls["n"] == 2 # one native + one padded probe - - # Simulate the production forward mutating the live params in place after the call. - params.head_dim_qk = 128 - params.head_dim_v = 128 - - # A fresh native params with the same config must still hit the memo (the key is a - # copy, not the mutated live object) and must not re-probe. - params2 = dpa_utils.AttentionParams(head_dim_qk=96, head_dim_v=128, **base) - assert dpa_utils.should_pad_qkv_head_dim(params2) is True - assert calls["n"] == 2 # cache hit: no new probes - - -# v > qk end-to-end -@pytest.mark.parametrize("qk,v", [(64, 192), (96, 192)]) -def test_dpa_v_gt_qk_runs(qk, v): - """DPA with head_dim_v > head_dim_qk runs and produces a V-width output.""" - reset_rng_states() - dpa = _build_dpa(qk, v) - q, k, v_t, cu = _thd_inputs(qk, v) - out = _run_dpa(dpa, q, k, v_t, cu) - assert tuple(out.shape) == (32, 4 * v), out.shape # V-width - out.float().sum().backward() # backward must not crash - - -# pad-then-trim is an identity (both directions) -@pytest.mark.parametrize("qk,v", [(192, 128), (64, 192)]) -def test_dpa_mla_pad_is_identity(qk, v): - """Pad-then-trim is an identity: padding Q/K/V to the wider head dim, running with the equal - (padded) shape, and trimming back equals the native mismatched-dim run -- for both qk > v and v - > qk. Both runs use the same `softmax_scale` (`1/sqrt(qk)`) that the production forward keeps - when padding. - """ - reset_rng_states() - m = max(qk, v) - scale = 1.0 / math.sqrt(qk) - cu = torch.IntTensor([0, 6, 19, 22, 32]).cuda() - - # Reference: native mismatched-dim run (the production forward; it pads internally - # only when should_pad_qkv_head_dim upgrades the selected backend). - dpa_ref = _build_dpa(qk, v) # softmax_scale defaults to 1/sqrt(qk) - q, k, v_t, _ = _thd_inputs(qk, v) - out_ref = _run_dpa(dpa_ref, q, k, v_t, cu) - assert tuple(out_ref.shape) == (32, 4 * v), out_ref.shape - - # Test: manually pad to the common width, run with the equal (padded) shape, trim. - # Same softmax_scale as the reference so pad-then-trim is a true identity. - dpa = _build_dpa(m, m, softmax_scale=scale) - q_p, k_p, v_p, _, _ = dpa_module._pad_qkv_head_dim(q, k, v_t) - assert q_p.shape[-1] == k_p.shape[-1] == v_p.shape[-1] == m - out = _run_dpa(dpa, q_p, k_p, v_p, cu) - # Trim back to the original V width. - out = dpa_module._trim_output(out, 4, m, v) - torch.testing.assert_close(out, out_ref, atol=1e-2, rtol=1e-2) - out.float().sum().backward() # padded path backward must not crash - - -def test_fp8_packed_skips_qkv_head_dim_pad(): - """`fp8_packed_skips_qkv_head_dim_pad` returns True (skip the pad) only in the one - unsafe combination -- FP8 DPA active AND packed inputs -- and False (allow the pad) - in every safe combination (bf16 regardless of packing; FP8 DPA + unpacked).""" - fp8_dpa_on = {"recipe": recipe.DelayedScaling(fp8_dpa=True)} - fp8_dpa_off = {"recipe": recipe.DelayedScaling(fp8_dpa=False)} - - def params(fp8, fp8_meta): - return dpa_utils.AttentionParams( - qkv_layout="thd_thd_thd", - num_heads=4, - num_gqa_groups=4, - max_seqlen_q=13, - max_seqlen_kv=13, - head_dim_qk=96, - head_dim_v=128, - attn_mask_type="padding_causal", - is_training=True, - qkv_dtype=torch.bfloat16, - fp8=fp8, - fp8_meta=fp8_meta, - ) - - # bf16 (fp8=False): pad is safe & beneficial regardless of packing. - assert dpa_utils.fp8_packed_skips_qkv_head_dim_pad(params(False, None), False) is False - assert dpa_utils.fp8_packed_skips_qkv_head_dim_pad(params(False, None), True) is False - # FP8 autocast but fp8_dpa=False (attention itself runs in bf16): pad is safe. - assert dpa_utils.fp8_packed_skips_qkv_head_dim_pad(params(True, fp8_dpa_off), False) is False - assert dpa_utils.fp8_packed_skips_qkv_head_dim_pad(params(True, fp8_dpa_off), True) is False - # FP8 DPA + unpacked: pad is safe (combine_and_quantize re-derives from padded views). - assert dpa_utils.fp8_packed_skips_qkv_head_dim_pad(params(True, fp8_dpa_on), False) is False - # FP8 DPA + packed: the only unsafe combination -> skip the pad. - assert dpa_utils.fp8_packed_skips_qkv_head_dim_pad(params(True, fp8_dpa_on), True) is True - # Defensive: fp8=True but no recipe -> treat as "do not skip" (pad allowed). - assert dpa_utils.fp8_packed_skips_qkv_head_dim_pad(params(True, {}), True) is False - assert dpa_utils.fp8_packed_skips_qkv_head_dim_pad(params(True, None), True) is False - - -# FP8 DPA + unpacked MLA: the guard must not over-block the safe beneficial case. -@pytest.mark.parametrize("qk,v", [(192, 128), (64, 192)]) -def test_dpa_mla_pad_fp8_dpa_unpacked(qk, v, monkeypatch): - """Regression test for the FP8+packed guard: under FP8 DPA with *unpacked* mismatched MLA - inputs, the pad is safe and beneficial (`combine_and_quantize` re-derives from the padded - views, so the kernel sees the equal padded dims), so the guard must NOT skip it. Force - the conditions that make `should_pad_qkv_head_dim` return True (native unfused via - the `NVTE_UnfusedDPA_Emulate_FP8` env) and assert the forward runs and yields a - correct V-width output.""" - monkeypatch.setenv("NVTE_UnfusedDPA_Emulate_FP8", "1") - reset_rng_states() - fp8_recipe = recipe.DelayedScaling( - margin=0, - fp8_format=recipe.Format.HYBRID, - amax_history_len=1, - amax_compute_algo="most_recent", - fp8_dpa=True, - fp8_mha=False, - ) - dpa = _build_dpa(qk, v) - q, k, v_t, cu = _thd_inputs(qk, v) - with te.fp8_autocast(enabled=True, fp8_recipe=fp8_recipe): - out = _run_dpa(dpa, q, k, v_t, cu) - assert tuple(out.shape) == (32, 4 * v), out.shape - with te.fp8_autocast(enabled=True, fp8_recipe=fp8_recipe): - out.float().sum().backward() # FP8 padded path backward must not crash From ca7fd180c04a776b82b521774bdac37c92f1b87f Mon Sep 17 00:00:00 2001 From: janEbert Date: Thu, 13 Aug 2026 21:54:32 +0000 Subject: [PATCH 13/30] Assert MLA pad selects a faster backend Signed-off-by: janEbert --- tests/pytorch/attention/test_attention.py | 63 +++++++++++++++++++++++ 1 file changed, 63 insertions(+) diff --git a/tests/pytorch/attention/test_attention.py b/tests/pytorch/attention/test_attention.py index 6bcff5ce3f1..6f96e5bb443 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,67 @@ def test_dpa_mla(dtype, model_configs, model): test_dot_product_attention(dtype, model_configs, model, True, None, False, False) +@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. + """ + + def make_attn_params(head_dim_qk, head_dim_v): + return dpa_utils.AttentionParams( + qkv_dtype=dtype, + qkv_layout="thd_thd_thd", + 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, + ) + + config = model_configs[model] + native_attn_params = make_attn_params(config.head_dim_qk, config.head_dim_v) + if config.head_dim_qk != config.head_dim_v: + native_backend = dpa_utils.get_attention_backend(native_attn_params) + native_flash_attention_backend = native_backend[1] + native_use_unfused_attention = native_backend[4] + + padded_head_dim = max(config.head_dim_qk, config.head_dim_v) + padded_attn_params = make_attn_params(padded_head_dim, padded_head_dim) + padded_backend = dpa_utils.get_attention_backend(padded_attn_params) + padded_use_flash_attention = padded_backend[0] + padded_use_fused_attention = padded_backend[2] + + if native_flash_attention_backend == FlashAttentionUtils.version: + # FA2 needs padding. + assert should_pad_qkv_head_dim( + native_attn_params + ), "did not suggest padding when required for functioning kernel" + elif native_use_unfused_attention and ( + padded_use_flash_attention or padded_use_fused_attention + ): + # Native is unfused and a fused backend is available via padding. + assert should_pad_qkv_head_dim( + native_attn_params + ), "did not suggest padding when faster kernel was available" + else: + # Native is already fused, so we don't want to pad. + assert not should_pad_qkv_head_dim( + native_attn_params + ), "suggested padding when no faster kernel was available" + else: + assert not should_pad_qkv_head_dim( + native_attn_params + ), "suggested padding already equal-sized QK/V" + + model_configs_mask = { # test: ModelConfig(b, sq, hq, dqk) "mask_1_0": ModelConfig(2, 2048, 16, 64, attn_mask_type="causal"), From e576da3381a951539760653906c6f1028d7c1c40 Mon Sep 17 00:00:00 2001 From: janEbert Date: Fri, 14 Aug 2026 01:48:22 +0200 Subject: [PATCH 14/30] Remove removed test from test runner script Signed-off-by: janEbert --- qa/L0_pytorch_unittest/test.sh | 1 - 1 file changed, 1 deletion(-) diff --git a/qa/L0_pytorch_unittest/test.sh b/qa/L0_pytorch_unittest/test.sh index 09056249a63..2a486450a20 100644 --- a/qa/L0_pytorch_unittest/test.sh +++ b/qa/L0_pytorch_unittest/test.sh @@ -74,7 +74,6 @@ python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_fused_mla_q_upro python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_kv_cache.xml $TE_PATH/tests/pytorch/attention/test_kv_cache.py || test_fail "test_kv_cache.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_cu_seqlens_cache.xml $TE_PATH/tests/pytorch/attention/test_cu_seqlens_cache.py || test_fail "test_cu_seqlens_cache.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_mixed_thd_attention.xml $TE_PATH/tests/pytorch/attention/test_mixed_thd_attention.py || test_fail "test_mixed_thd_attention.py" -python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_dpa_mla_qkv_head_dim_pad.xml $TE_PATH/tests/pytorch/attention/test_dpa_mla_qkv_head_dim_pad.py || test_fail "test_dpa_mla_qkv_head_dim_pad.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_hf_integration.xml $TE_PATH/tests/pytorch/test_hf_integration.py || test_fail "test_hf_integration.py" export NVTE_TEST_CHECKPOINT_ARTIFACT_PATH=$TE_PATH/artifacts/tests/pytorch/test_checkpoint if [ ! -d "$NVTE_TEST_CHECKPOINT_ARTIFACT_PATH" ]; then From 37022203e8c3743a5bb8ab46041480adafb4343d Mon Sep 17 00:00:00 2001 From: janEbert Date: Fri, 14 Aug 2026 01:50:30 +0200 Subject: [PATCH 15/30] Reduce verbosity Signed-off-by: janEbert --- .../pytorch/attention/dot_product_attention/utils.py | 10 +--------- 1 file changed, 1 insertion(+), 9 deletions(-) diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py index 7c87f0f8baa..c379603f919 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py @@ -2051,14 +2051,6 @@ def fp8_packed_skips_qkv_head_dim_pad( """Whether the optional QKV head-dim pad must be skipped because FP8 DPA is active and the caller passed packed QKV/KV buffers. - The pad in `DotProductAttention.forward` pads only the unpacked q/k/v views, not the packed - `qkv_layer`/`kv_layer` buffers. In the FP8 fused path, `combine_and_quantize` consumes those - packed buffers directly, bypassing the padded views -- so padding would feed the kernel the - original mismatched-dim packed data while the backend was selected for the equal padded dims. - This is only unsafe when FP8 DPA is active AND packed inputs are present; every other - combination is safe (FP8 DPA + unpacked re-derives from the padded views; bf16 + packed ignores - the packed buffers in `fused_attn_fwd`). - Parameters ---------- attention_params : AttentionParams @@ -2069,7 +2061,7 @@ def fp8_packed_skips_qkv_head_dim_pad( Returns ------- bool - True iff the pad should be skipped (not applied). + True iff the pad should be skipped. """ if not attention_params.fp8: return False From aae27cc573734c4ff087cb2b46fa1636d84d20c1 Mon Sep 17 00:00:00 2001 From: janEbert Date: Mon, 12 Oct 2026 22:14:02 +0200 Subject: [PATCH 16/30] Update Attn fwd type hints The tuple is returned when `return_max_logit=True`. Signed-off-by: janEbert --- .../pytorch/attention/dot_product_attention/backends.py | 6 +++--- .../dot_product_attention/dot_product_attention.py | 2 +- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/transformer_engine/pytorch/attention/dot_product_attention/backends.py b/transformer_engine/pytorch/attention/dot_product_attention/backends.py index bccb4f1e26f..749988310d1 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 320244a0921..43a6d2548ca 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 @@ -1986,7 +1986,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. From 8f06b41741809e385d35617be5f95bcc6b6d6330 Mon Sep 17 00:00:00 2001 From: janEbert Date: Mon, 12 Oct 2026 22:15:30 +0200 Subject: [PATCH 17/30] Improve readability The two values are the same here, so it doesn't matter whether we use `head_dim_qk` or `head_dim_v`. Signed-off-by: janEbert --- .../attention/dot_product_attention/dot_product_attention.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 43a6d2548ca..d6bb37925ba 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 @@ -3141,5 +3141,5 @@ def forward( return None if qkv_head_pad: - attn_out = _trim_output(attn_out, num_attention_heads, head_dim_qk, orig_head_dim_v) + attn_out = _trim_output(attn_out, num_attention_heads, head_dim_v, orig_head_dim_v) return attn_out From f15da4f93a800e32511731cb01654f7cd58f68ac Mon Sep 17 00:00:00 2001 From: janEbert Date: Mon, 12 Oct 2026 22:16:32 +0200 Subject: [PATCH 18/30] Improve docstring No longer just related to FA2 and this encapsulates what we're actually doing by padding. Signed-off-by: janEbert --- .../attention/dot_product_attention/dot_product_attention.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 d6bb37925ba..43f5ce9a24a 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 @@ -442,7 +442,7 @@ def _should_pad_qkv_head_dim( 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) From 2da41778432c8e354e2136b96c3735711585d143 Mon Sep 17 00:00:00 2001 From: janEbert Date: Tue, 13 Oct 2026 00:19:06 +0200 Subject: [PATCH 19/30] Fix padding selection It could be that all `native_use_*_attention` values are False, so this needs to be handled. Signed-off-by: janEbert --- tests/pytorch/attention/test_attention.py | 6 +++++- .../attention/dot_product_attention/utils.py | 16 +++++++++++----- 2 files changed, 16 insertions(+), 6 deletions(-) diff --git a/tests/pytorch/attention/test_attention.py b/tests/pytorch/attention/test_attention.py index 6f96e5bb443..ebf8da2ba1b 100644 --- a/tests/pytorch/attention/test_attention.py +++ b/tests/pytorch/attention/test_attention.py @@ -1415,6 +1415,7 @@ def make_attn_params(head_dim_qk, head_dim_v): native_attn_params = make_attn_params(config.head_dim_qk, config.head_dim_v) if config.head_dim_qk != config.head_dim_v: native_backend = dpa_utils.get_attention_backend(native_attn_params) + native_use_flash_attention = native_backend[0] native_flash_attention_backend = native_backend[1] native_use_unfused_attention = native_backend[4] @@ -1424,7 +1425,10 @@ def make_attn_params(head_dim_qk, head_dim_v): padded_use_flash_attention = padded_backend[0] padded_use_fused_attention = padded_backend[2] - if native_flash_attention_backend == FlashAttentionUtils.version: + if ( + native_use_flash_attention + and native_flash_attention_backend == FlashAttentionUtils.version + ): # FA2 needs padding. assert should_pad_qkv_head_dim( native_attn_params diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py index c379603f919..4823f8f1e66 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py @@ -2019,14 +2019,18 @@ def should_pad_qkv_head_dim( if cached_params is not None and cached_params == attention_params: return _should_pad_qkv_head_dim_cache["result"] native_backend = get_attention_backend(attention_params) + native_use_flash_attention = native_backend[0] native_flash_attention_backend = native_backend[1] + native_use_fused_attention = native_backend[2] native_use_unfused_attention = native_backend[4] - # No point padding if native path is already fused. - if not native_use_unfused_attention: - # Native already lands on FlashAttention or FusedAttention, so no need to pad. However, - # FlashAttention-2 cannot take mismatched QK/V head dims, so we still pad in that case. + if native_use_flash_attention: + # Native already lands on FlashAttention, so no need to pad. However, FlashAttention-2 + # cannot take mismatched QK/V head dims, so we still pad in that case. result = native_flash_attention_backend == FlashAttentionUtils.version - else: + elif native_use_fused_attention: + # Native already lands on FusedAttention, so no need to pad. + result = False + elif native_use_unfused_attention: padded_head_dim = max(attention_params.head_dim_qk, attention_params.head_dim_v) padded_params = replace( attention_params, @@ -2037,6 +2041,8 @@ def should_pad_qkv_head_dim( padded_use_flash_attention = padded_backend[0] padded_use_fused_attention = padded_backend[2] result = padded_use_flash_attention or padded_use_fused_attention + else: + result = False # 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) From def8ad1d1e7a7b03c41a78a542defc57d5de5785 Mon Sep 17 00:00:00 2001 From: janEbert Date: Tue, 13 Oct 2026 00:24:20 +0200 Subject: [PATCH 20/30] Handle tuple Attn backend outputs Signed-off-by: janEbert --- .../dot_product_attention/dot_product_attention.py | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) 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 43f5ce9a24a..e8cb3a2d04c 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 @@ -3141,5 +3141,13 @@ def forward( return None if qkv_head_pad: - attn_out = _trim_output(attn_out, num_attention_heads, head_dim_v, orig_head_dim_v) + # Keep the rest of the result tuple as-is. + if not self.return_max_logit: + attn_out = (attn_out,) + attn_out = ( + # Note the comma at the end; we're creating a 1-tuple here. + _trim_output(attn_out[0], num_attention_heads, head_dim_v, orig_head_dim_v), + ) + attn_out[1:] + if not self.return_max_logit: + attn_out = attn_out[0] return attn_out From 4dcde582b0024e20c0cf64f4e4c7593e5cd8714a Mon Sep 17 00:00:00 2001 From: janEbert Date: Tue, 13 Oct 2026 00:45:41 +0200 Subject: [PATCH 21/30] Fold FP8 DSA+packed check into main padding check Improve variable names along the way. Signed-off-by: janEbert --- tests/pytorch/attention/test_attention.py | 9 +- .../dot_product_attention.py | 6 +- .../attention/dot_product_attention/utils.py | 113 +++++++++--------- 3 files changed, 61 insertions(+), 67 deletions(-) diff --git a/tests/pytorch/attention/test_attention.py b/tests/pytorch/attention/test_attention.py index ebf8da2ba1b..b7fa2147624 100644 --- a/tests/pytorch/attention/test_attention.py +++ b/tests/pytorch/attention/test_attention.py @@ -1396,6 +1396,7 @@ def test_dpa_mla_should_pad(dtype, model_configs, model): 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. """ + has_packed_qkv = False def make_attn_params(head_dim_qk, head_dim_v): return dpa_utils.AttentionParams( @@ -1431,23 +1432,23 @@ def make_attn_params(head_dim_qk, head_dim_v): ): # FA2 needs padding. assert should_pad_qkv_head_dim( - native_attn_params + native_attn_params, has_packed_qkv ), "did not suggest padding when required for functioning kernel" elif native_use_unfused_attention and ( padded_use_flash_attention or padded_use_fused_attention ): # Native is unfused and a fused backend is available via padding. assert should_pad_qkv_head_dim( - native_attn_params + native_attn_params, has_packed_qkv ), "did not suggest padding when faster kernel was available" else: # Native is already fused, so we don't want to pad. assert not should_pad_qkv_head_dim( - native_attn_params + native_attn_params, has_packed_qkv ), "suggested padding when no faster kernel was available" else: assert not should_pad_qkv_head_dim( - native_attn_params + native_attn_params, has_packed_qkv ), "suggested padding already equal-sized QK/V" 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 e8cb3a2d04c..0ab0da1b1b4 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 @@ -434,10 +434,8 @@ def _should_pad_qkv_head_dim( ) # Skip padding when packed inputs are used with the FP8 fused path, which would lead to using # unpadded data with a kernel that assumes padded data. - return ( - is_shape_pad_eligible - and not dpa_utils.fp8_packed_skips_qkv_head_dim_pad(attention_params, has_packed_inputs) - and dpa_utils.should_pad_qkv_head_dim(attention_params) + return is_shape_pad_eligible and dpa_utils.should_pad_qkv_head_dim( + attention_params, has_packed_inputs ) diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py index 4823f8f1e66..5a2ff10bc40 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py @@ -76,7 +76,8 @@ # Global var for MLA padding cache. _should_pad_qkv_head_dim_cache: Dict[str, Any] = { "attention_params": None, - "result": None, + "has_packed_qkv": None, + "should_pad": None, } @@ -1993,6 +1994,7 @@ 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, ) -> bool: """Decide whether to pad Q/K/V to a common head dim. @@ -2001,12 +2003,14 @@ def should_pad_qkv_head_dim( mismatched head dims (FlashAttention-2). 2. when padding results in a faster kernel - The selection is memoized for the given `attention_params`. + The selection is memoized for the given `attention_params` and `has_packed_qkv`. 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. Returns ------- @@ -2015,66 +2019,57 @@ def should_pad_qkv_head_dim( """ if attention_params.head_dim_qk == attention_params.head_dim_v: return False - cached_params = _should_pad_qkv_head_dim_cache["attention_params"] - if cached_params is not None and cached_params == attention_params: - return _should_pad_qkv_head_dim_cache["result"] - native_backend = get_attention_backend(attention_params) - native_use_flash_attention = native_backend[0] - native_flash_attention_backend = native_backend[1] - native_use_fused_attention = native_backend[2] - native_use_unfused_attention = native_backend[4] - if native_use_flash_attention: - # Native already lands on FlashAttention, so no need to pad. However, FlashAttention-2 - # cannot take mismatched QK/V head dims, so we still pad in that case. - result = native_flash_attention_backend == FlashAttentionUtils.version - elif native_use_fused_attention: - # Native already lands on FusedAttention, so no need to pad. - result = False - elif native_use_unfused_attention: - 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] - result = padded_use_flash_attention or padded_use_fused_attention - else: - result = False + cached_attention_params = _should_pad_qkv_head_dim_cache["attention_params"] + cached_has_packed_qkv = _should_pad_qkv_head_dim_cache["has_packed_qkv"] + if ( + cached_attention_params is not None + and cached_has_packed_qkv is has_packed_qkv + and cached_attention_params == attention_params + ): + 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_flash_attention_backend = native_backend[1] + native_use_fused_attention = native_backend[2] + native_use_unfused_attention = native_backend[4] + if native_use_flash_attention: + # Native already lands on FlashAttention, so no need to pad. However, FlashAttention-2 + # cannot take mismatched QK/V head dims, so we still pad in that case. + should_pad = native_flash_attention_backend == FlashAttentionUtils.version + elif native_use_fused_attention: + # Native already lands on FusedAttention, so no need to pad. + should_pad = False + elif native_use_unfused_attention: + 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] + should_pad = padded_use_flash_attention or padded_use_fused_attention + else: + should_pad = False # 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["result"] = result - return result - - -def fp8_packed_skips_qkv_head_dim_pad( - attention_params: AttentionParams, - has_packed_qkv: bool, -) -> bool: - """Whether the optional QKV head-dim pad must be skipped because FP8 DPA is active and the - caller passed packed QKV/KV buffers. - - 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. - - Returns - ------- - bool - True iff the pad should be skipped. - """ - if not attention_params.fp8: - return False - recipe = (attention_params.fp8_meta or {}).get("recipe", None) - if recipe is None or not getattr(recipe, "fp8_dpa", False): - return False - return has_packed_qkv + _should_pad_qkv_head_dim_cache["has_packed_qkv"] = has_packed_qkv + _should_pad_qkv_head_dim_cache["should_pad"] = should_pad + return should_pad @torch.no_grad() From 3f0079962660851c7b094eef8d79232220c5f615 Mon Sep 17 00:00:00 2001 From: janEbert Date: Tue, 13 Oct 2026 17:43:22 +0200 Subject: [PATCH 22/30] Fix backend selection cache invalidation Signed-off-by: janEbert --- .../dot_product_attention.py | 4 ++- .../attention/dot_product_attention/utils.py | 28 +++++++++++++++++-- 2 files changed, 29 insertions(+), 3 deletions(-) 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 0ab0da1b1b4..a6c35dfd765 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 @@ -435,7 +435,9 @@ def _should_pad_qkv_head_dim( # Skip padding when packed inputs are used with the FP8 fused path, which would lead to using # unpadded data with a kernel that assumes padded data. return is_shape_pad_eligible and dpa_utils.should_pad_qkv_head_dim( - attention_params, has_packed_inputs + attention_params, + has_packed_inputs, + _attention_backends["attention_params"], ) diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py index 5a2ff10bc40..300142d02eb 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py @@ -1995,6 +1995,7 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt def should_pad_qkv_head_dim( attention_params: AttentionParams, has_packed_qkv: bool, + backend_selection_params: Optional[AttentionParams], ) -> bool: """Decide whether to pad Q/K/V to a common head dim. @@ -2003,7 +2004,8 @@ def should_pad_qkv_head_dim( mismatched head dims (FlashAttention-2). 2. when padding results in a faster kernel - The selection is memoized for the given `attention_params` and `has_packed_qkv`. + The selection is memoized for the given `attention_params` and `has_packed_qkv`. A cached result + is only reused when the attention params used for backend selection match. Parameters ---------- @@ -2011,6 +2013,9 @@ def should_pad_qkv_head_dim( 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. + backend_selection_params : Optional[AttentionParams] + The attention params that were last used for the backend selection, or `None` if no backend + has been selected previously. Returns ------- @@ -2026,7 +2031,26 @@ def should_pad_qkv_head_dim( and cached_has_packed_qkv is has_packed_qkv and cached_attention_params == attention_params ): - return _should_pad_qkv_head_dim_cache["should_pad"] + cached_should_pad = _should_pad_qkv_head_dim_cache["should_pad"] + # Recreate the attention params that would've flown into backend selection with the cached + # results to check whether they match those currently used selection. + # + # This is required because the `dot_product_attention._attention_backends` cache can be + # invalidated from the outside. + if cached_should_pad: + padded_head_dim = max(attention_params.head_dim_qk, attention_params.head_dim_v) + expected_selection_params = replace( + attention_params, + head_dim_qk=padded_head_dim, + head_dim_v=padded_head_dim, + ) + else: + expected_selection_params = attention_params + if ( + backend_selection_params is not None + and backend_selection_params == expected_selection_params + ): + return cached_should_pad # Check whether padding must be skipped because FP8 DPA is active and the caller passed packed # QKV/KV buffers. From 0dc265d0be8dcedacf5d3f1760c5d704809a3705 Mon Sep 17 00:00:00 2001 From: janEbert Date: Tue, 13 Oct 2026 18:00:02 +0200 Subject: [PATCH 23/30] Fix tests Signed-off-by: janEbert --- tests/pytorch/attention/test_attention.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/tests/pytorch/attention/test_attention.py b/tests/pytorch/attention/test_attention.py index b7fa2147624..2fb4be426f7 100644 --- a/tests/pytorch/attention/test_attention.py +++ b/tests/pytorch/attention/test_attention.py @@ -1397,6 +1397,8 @@ def test_dpa_mla_should_pad(dtype, model_configs, model): backend. Otherwise padding should be avoided since it becomes redundant. """ has_packed_qkv = False + # `None` to computes each padding decision from scratch. + backend_selection_params = None def make_attn_params(head_dim_qk, head_dim_v): return dpa_utils.AttentionParams( @@ -1432,23 +1434,23 @@ def make_attn_params(head_dim_qk, head_dim_v): ): # FA2 needs padding. assert should_pad_qkv_head_dim( - native_attn_params, has_packed_qkv + native_attn_params, has_packed_qkv, backend_selection_params ), "did not suggest padding when required for functioning kernel" elif native_use_unfused_attention and ( padded_use_flash_attention or padded_use_fused_attention ): # Native is unfused and a fused backend is available via padding. assert should_pad_qkv_head_dim( - native_attn_params, has_packed_qkv + native_attn_params, has_packed_qkv, backend_selection_params ), "did not suggest padding when faster kernel was available" else: # Native is already fused, so we don't want to pad. assert not should_pad_qkv_head_dim( - native_attn_params, has_packed_qkv + native_attn_params, has_packed_qkv, backend_selection_params ), "suggested padding when no faster kernel was available" else: assert not should_pad_qkv_head_dim( - native_attn_params, has_packed_qkv + native_attn_params, has_packed_qkv, backend_selection_params ), "suggested padding already equal-sized QK/V" From 220c370779bc9f6b656e4cb4660e91359219a142 Mon Sep 17 00:00:00 2001 From: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> Date: Tue, 13 Oct 2026 16:54:21 -0700 Subject: [PATCH 24/30] To-be-reviewed: minor changes Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> --- tests/pytorch/attention/test_attention.py | 129 +++++++++++------- .../dot_product_attention.py | 63 +++------ .../attention/dot_product_attention/utils.py | 59 ++++---- 3 files changed, 125 insertions(+), 126 deletions(-) diff --git a/tests/pytorch/attention/test_attention.py b/tests/pytorch/attention/test_attention.py index 2fb4be426f7..8e9e7105ea6 100644 --- a/tests/pytorch/attention/test_attention.py +++ b/tests/pytorch/attention/test_attention.py @@ -1386,6 +1386,22 @@ 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="thd_thd_thd", + 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]) @@ -1396,62 +1412,69 @@ def test_dpa_mla_should_pad(dtype, model_configs, model): 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. """ - has_packed_qkv = False - # `None` to computes each padding decision from scratch. - backend_selection_params = None + 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) - def make_attn_params(head_dim_qk, head_dim_v): - return dpa_utils.AttentionParams( - qkv_dtype=dtype, - qkv_layout="thd_thd_thd", - 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, - ) + # 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_flash_backend, native_fused = native_backend[:3] + native_unfused = native_backend[4] + padded_backend = dpa_utils.get_attention_backend(padded_params) + padded_flash, padded_fused = padded_backend[0], padded_backend[2] + + # FlashAttention-2 cannot run mismatched head dims at all, so padding is required rather + # than an optimization. Otherwise it is only worthwhile when the native dims are stuck on + # the unfused backend and padding reaches a fused one. + pad_is_required = native_flash and native_flash_backend == FlashAttentionUtils.version + pad_unlocks_backend = ( + not native_flash and not native_fused and native_unfused and (padded_flash or padded_fused) + ) + expected = bool(pad_is_required or pad_unlocks_backend) + + assert should_pad_qkv_head_dim(native_params, False) == expected, ( + "did not suggest padding when required for a functioning kernel or a faster one" + if expected + else "suggested padding when no faster kernel was available" + ) + + +@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_attn_params = make_attn_params(config.head_dim_qk, config.head_dim_v) - if config.head_dim_qk != config.head_dim_v: - native_backend = dpa_utils.get_attention_backend(native_attn_params) - native_use_flash_attention = native_backend[0] - native_flash_attention_backend = native_backend[1] - native_use_unfused_attention = native_backend[4] - - padded_head_dim = max(config.head_dim_qk, config.head_dim_v) - padded_attn_params = make_attn_params(padded_head_dim, padded_head_dim) - padded_backend = dpa_utils.get_attention_backend(padded_attn_params) - padded_use_flash_attention = padded_backend[0] - padded_use_fused_attention = padded_backend[2] - - if ( - native_use_flash_attention - and native_flash_attention_backend == FlashAttentionUtils.version - ): - # FA2 needs padding. - assert should_pad_qkv_head_dim( - native_attn_params, has_packed_qkv, backend_selection_params - ), "did not suggest padding when required for functioning kernel" - elif native_use_unfused_attention and ( - padded_use_flash_attention or padded_use_fused_attention - ): - # Native is unfused and a fused backend is available via padding. - assert should_pad_qkv_head_dim( - native_attn_params, has_packed_qkv, backend_selection_params - ), "did not suggest padding when faster kernel was available" - else: - # Native is already fused, so we don't want to pad. - assert not should_pad_qkv_head_dim( - native_attn_params, has_packed_qkv, backend_selection_params - ), "suggested padding when no faster kernel was available" - else: - assert not should_pad_qkv_head_dim( - native_attn_params, has_packed_qkv, backend_selection_params - ), "suggested padding already equal-sized QK/V" + 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) + dpa_utils.clear_should_pad_qkv_head_dim_cache() + fresh = should_pad_qkv_head_dim(native_params, False) + 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 = { 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 a6c35dfd765..ab34eb993ea 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 @@ -85,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 = { @@ -422,25 +415,6 @@ def _delayed_scaling_recipe() -> Optional[DelayedScaling]: __all__ = ["DotProductAttention"] -def _should_pad_qkv_head_dim( - head_dim_qk, head_dim_v, value_layer, qkv_layer, kv_layer, attention_params -): - has_packed_inputs = qkv_layer is not None or kv_layer is not None - is_shape_pad_eligible = ( - not is_in_onnx_export_mode() - and head_dim_qk != head_dim_v - and value_layer is not None - and not isinstance(value_layer, Float8TensorStorage) - ) - # Skip padding when packed inputs are used with the FP8 fused path, which would lead to using - # unpadded data with a kernel that assumes padded data. - return is_shape_pad_eligible and dpa_utils.should_pad_qkv_head_dim( - attention_params, - has_packed_inputs, - _attention_backends["attention_params"], - ) - - def _pad_qkv_head_dim(query_layer, key_layer, value_layer): """Pad Q/K/V to the same head dimension for better backend selection result.""" orig_head_dim_qk = query_layer.shape[-1] @@ -2860,13 +2834,22 @@ def forward( pad_between_seqs=pad_between_seqs, ) + if _attention_backends["backend_selection_requires_update"]: + dpa_utils.clear_should_pad_qkv_head_dim_cache() + # 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 inputs shapes by `{Flash,Fused}Attention`. - qkv_head_pad = False + # 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 _should_pad_qkv_head_dim( - head_dim_qk, head_dim_v, value_layer, qkv_layer, kv_layer, attention_params + if can_pad_qkv_head_dim and dpa_utils.should_pad_qkv_head_dim( + attention_params, has_packed_qkv ): query_layer, key_layer, value_layer, _, _ = _pad_qkv_head_dim( query_layer, key_layer, value_layer @@ -2875,7 +2858,7 @@ def forward( 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 - qkv_head_pad = True + is_qkv_head_dim_padded = True if is_in_onnx_export_mode(): # We do not want to call get_attention_backend() in ONNX mode @@ -3140,14 +3123,10 @@ def forward( else: return None - if qkv_head_pad: - # Keep the rest of the result tuple as-is. - if not self.return_max_logit: - attn_out = (attn_out,) - attn_out = ( - # Note the comma at the end; we're creating a 1-tuple here. - _trim_output(attn_out[0], num_attention_heads, head_dim_v, orig_head_dim_v), - ) + attn_out[1:] - if not self.return_max_logit: - attn_out = attn_out[0] + 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 300142d02eb..6760b23f3ca 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py @@ -73,14 +73,32 @@ _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: """ Manage logging for attention module @@ -1995,7 +2013,6 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt def should_pad_qkv_head_dim( attention_params: AttentionParams, has_packed_qkv: bool, - backend_selection_params: Optional[AttentionParams], ) -> bool: """Decide whether to pad Q/K/V to a common head dim. @@ -2004,8 +2021,9 @@ def should_pad_qkv_head_dim( mismatched head dims (FlashAttention-2). 2. when padding results in a faster kernel - The selection is memoized for the given `attention_params` and `has_packed_qkv`. A cached result - is only reused when the attention params used for backend selection match. + 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 also call + `clear_should_pad_qkv_head_dim_cache`. Parameters ---------- @@ -2013,9 +2031,6 @@ def should_pad_qkv_head_dim( 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. - backend_selection_params : Optional[AttentionParams] - The attention params that were last used for the backend selection, or `None` if no backend - has been selected previously. Returns ------- @@ -2024,33 +2039,14 @@ def should_pad_qkv_head_dim( """ if attention_params.head_dim_qk == attention_params.head_dim_v: return False - cached_attention_params = _should_pad_qkv_head_dim_cache["attention_params"] - cached_has_packed_qkv = _should_pad_qkv_head_dim_cache["has_packed_qkv"] + + environment = tuple(os.environ.get(name) for name in BACKEND_SELECTION_ENV_VARS) if ( - cached_attention_params is not None - and cached_has_packed_qkv is has_packed_qkv - and cached_attention_params == attention_params + _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 ): - cached_should_pad = _should_pad_qkv_head_dim_cache["should_pad"] - # Recreate the attention params that would've flown into backend selection with the cached - # results to check whether they match those currently used selection. - # - # This is required because the `dot_product_attention._attention_backends` cache can be - # invalidated from the outside. - if cached_should_pad: - padded_head_dim = max(attention_params.head_dim_qk, attention_params.head_dim_v) - expected_selection_params = replace( - attention_params, - head_dim_qk=padded_head_dim, - head_dim_v=padded_head_dim, - ) - else: - expected_selection_params = attention_params - if ( - backend_selection_params is not None - and backend_selection_params == expected_selection_params - ): - return cached_should_pad + 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. @@ -2092,6 +2088,7 @@ def should_pad_qkv_head_dim( # 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 From 9739216631ecde4b853381b599977ef2b453083d Mon Sep 17 00:00:00 2001 From: janEbert Date: Tue, 13 Oct 2026 19:20:22 +0200 Subject: [PATCH 25/30] Fix QKV layout Also add batch size config. Signed-off-by: janEbert --- tests/pytorch/attention/test_attention.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/pytorch/attention/test_attention.py b/tests/pytorch/attention/test_attention.py index 8e9e7105ea6..f9ffda91922 100644 --- a/tests/pytorch/attention/test_attention.py +++ b/tests/pytorch/attention/test_attention.py @@ -1390,7 +1390,8 @@ 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="thd_thd_thd", + 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, From 4029149cbb93ca65279d26801972f578f1c7caf3 Mon Sep 17 00:00:00 2001 From: janEbert Date: Wed, 14 Oct 2026 13:44:47 +0200 Subject: [PATCH 26/30] Use varlength tuples with type hints Signed-off-by: janEbert --- .../pytorch/attention/dot_product_attention/backends.py | 6 +++--- .../dot_product_attention/dot_product_attention.py | 2 +- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/transformer_engine/pytorch/attention/dot_product_attention/backends.py b/transformer_engine/pytorch/attention/dot_product_attention/backends.py index 749988310d1..1cfc23ba01d 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, - ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]: + ) -> Union[torch.Tensor, Tuple[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]]: + ) -> Union[torch.Tensor, Tuple[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, - ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]: + ) -> Union[torch.Tensor, Tuple[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 ab34eb993ea..e72b6fac5d8 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 @@ -1960,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", - ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor], None]: + ) -> Union[torch.Tensor, Tuple[torch.Tensor, ...], None]: r""" Dot Product Attention Layer. From f9a2e77892a8417c2a6bfa0e83726b0273ddd9a4 Mon Sep 17 00:00:00 2001 From: janEbert Date: Wed, 14 Oct 2026 19:13:32 +0200 Subject: [PATCH 27/30] Revert "Use varlength tuples with type hints" This reverts commit 4029149cbb93ca65279d26801972f578f1c7caf3. Signed-off-by: janEbert --- .../pytorch/attention/dot_product_attention/backends.py | 6 +++--- .../dot_product_attention/dot_product_attention.py | 2 +- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/transformer_engine/pytorch/attention/dot_product_attention/backends.py b/transformer_engine/pytorch/attention/dot_product_attention/backends.py index 1cfc23ba01d..749988310d1 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, - ) -> Union[torch.Tensor, Tuple[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, ...]]: + ) -> 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, - ) -> Union[torch.Tensor, Tuple[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 e72b6fac5d8..ab34eb993ea 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 @@ -1960,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", - ) -> Union[torch.Tensor, Tuple[torch.Tensor, ...], None]: + ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor], None]: r""" Dot Product Attention Layer. From 0620e6d1fd5db1beab8b5ccd119ab42c7036f4fa Mon Sep 17 00:00:00 2001 From: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> Date: Wed, 14 Oct 2026 11:43:59 -0700 Subject: [PATCH 28/30] fix FA2 pad logic and should_pad overall logic Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> --- tests/pytorch/attention/test_attention.py | 24 +++----- .../attention/dot_product_attention/utils.py | 56 +++++++++++-------- 2 files changed, 41 insertions(+), 39 deletions(-) diff --git a/tests/pytorch/attention/test_attention.py b/tests/pytorch/attention/test_attention.py index f9ffda91922..5ca19d19eda 100644 --- a/tests/pytorch/attention/test_attention.py +++ b/tests/pytorch/attention/test_attention.py @@ -1427,25 +1427,17 @@ def test_dpa_mla_should_pad(dtype, model_configs, model): 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_flash_backend, native_fused = native_backend[:3] - native_unfused = native_backend[4] + 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] - # FlashAttention-2 cannot run mismatched head dims at all, so padding is required rather - # than an optimization. Otherwise it is only worthwhile when the native dims are stuck on - # the unfused backend and padding reaches a fused one. - pad_is_required = native_flash and native_flash_backend == FlashAttentionUtils.version - pad_unlocks_backend = ( - not native_flash and not native_fused and native_unfused and (padded_flash or padded_fused) - ) - expected = bool(pad_is_required or pad_unlocks_backend) - - assert should_pad_qkv_head_dim(native_params, False) == expected, ( - "did not suggest padding when required for a functioning kernel or a faster one" - if expected - else "suggested padding when no faster kernel was available" - ) + if should_pad_qkv_head_dim(native_params, False): + assert padded_flash or padded_fused, "suggested padding onto no faster backend" + assert padded_flash != native_flash, "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.") diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py index 6760b23f3ca..cebadd876ec 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py @@ -1037,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( @@ -1058,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 @@ -2017,9 +2023,10 @@ def should_pad_qkv_head_dim( """Decide whether to pad Q/K/V to a common head dim. This returns `True` in two cases: - 1. when the native mismatched shape would otherwise land on a backend that cannot take - mismatched head dims (FlashAttention-2). - 2. when padding results in a faster kernel + 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 also call @@ -2061,17 +2068,13 @@ def should_pad_qkv_head_dim( if should_pad: native_backend = get_attention_backend(attention_params) native_use_flash_attention = native_backend[0] - native_flash_attention_backend = native_backend[1] native_use_fused_attention = native_backend[2] - native_use_unfused_attention = native_backend[4] - if native_use_flash_attention: - # Native already lands on FlashAttention, so no need to pad. However, FlashAttention-2 - # cannot take mismatched QK/V head dims, so we still pad in that case. - should_pad = native_flash_attention_backend == FlashAttentionUtils.version - elif native_use_fused_attention: - # Native already lands on FusedAttention, so no need to pad. + 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 - elif native_use_unfused_attention: + else: padded_head_dim = max(attention_params.head_dim_qk, attention_params.head_dim_v) padded_params = replace( attention_params, @@ -2081,9 +2084,16 @@ def should_pad_qkv_head_dim( padded_backend = get_attention_backend(padded_params) padded_use_flash_attention = padded_backend[0] padded_use_fused_attention = padded_backend[2] - should_pad = padded_use_flash_attention or padded_use_fused_attention - else: - should_pad = False + 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) From d69ff67e82422de964256415a9c02e256d3a71bb Mon Sep 17 00:00:00 2001 From: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> Date: Wed, 14 Oct 2026 12:45:55 -0700 Subject: [PATCH 29/30] Update tests/pytorch/attention/test_attention.py Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> --- tests/pytorch/attention/test_attention.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/tests/pytorch/attention/test_attention.py b/tests/pytorch/attention/test_attention.py index 5ca19d19eda..5df85ffad39 100644 --- a/tests/pytorch/attention/test_attention.py +++ b/tests/pytorch/attention/test_attention.py @@ -1433,7 +1433,10 @@ def test_dpa_mla_should_pad(dtype, model_configs, model): if should_pad_qkv_head_dim(native_params, False): assert padded_flash or padded_fused, "suggested padding onto no faster backend" - assert padded_flash != native_flash, "suggested padding onto the same 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) From f8455f106ae8ce3a515b966c06e1ebe06f03758a Mon Sep 17 00:00:00 2001 From: janEbert Date: Wed, 21 Oct 2026 10:21:22 +0200 Subject: [PATCH 30/30] Improve cache clear usability Signed-off-by: janEbert --- tests/pytorch/attention/test_attention.py | 3 +-- .../dot_product_attention/dot_product_attention.py | 7 +++---- .../pytorch/attention/dot_product_attention/utils.py | 9 +++++++-- 3 files changed, 11 insertions(+), 8 deletions(-) diff --git a/tests/pytorch/attention/test_attention.py b/tests/pytorch/attention/test_attention.py index 5df85ffad39..96bd4a947c1 100644 --- a/tests/pytorch/attention/test_attention.py +++ b/tests/pytorch/attention/test_attention.py @@ -1461,8 +1461,7 @@ def test_dpa_mla_should_pad_cache(dtype, model_configs, model): os.environ["NVTE_FLASH_ATTN"] = "0" os.environ["NVTE_FUSED_ATTN"] = "0" cached = should_pad_qkv_head_dim(native_params, False) - dpa_utils.clear_should_pad_qkv_head_dim_cache() - fresh = 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: 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 ab34eb993ea..3731829c983 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 @@ -2834,9 +2834,6 @@ def forward( pad_between_seqs=pad_between_seqs, ) - if _attention_backends["backend_selection_requires_update"]: - dpa_utils.clear_should_pad_qkv_head_dim_cache() - # 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`. @@ -2849,7 +2846,9 @@ def forward( 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 + 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 diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py index ebded24a9db..b850aaf8be0 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py @@ -2019,6 +2019,7 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt 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. @@ -2029,8 +2030,8 @@ def should_pad_qkv_head_dim( 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 also call - `clear_should_pad_qkv_head_dim_cache`. + environment. Callers that invalidate backend selection for any other reason must pass + `invalidate_cache=True`. Parameters ---------- @@ -2038,12 +2039,16 @@ def should_pad_qkv_head_dim( 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