From 438ecfe559fcc3b400c67fb8b13533c46caa616e Mon Sep 17 00:00:00 2001 From: janEbert Date: Mon, 10 Aug 2026 20:45:17 +0200 Subject: [PATCH 1/6] 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 | 171 +++++++++++------- .../attention/dot_product_attention/utils.py | 46 ++++- 2 files changed, 150 insertions(+), 67 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 e6f50eb0da..42bf477460 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 @@ -2042,6 +2042,31 @@ def forward( has_score_mod=score_mod is not None, has_score_mod_bprop=score_mod_bprop is not None, ) + + # 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 + global _attention_backends if is_in_onnx_export_mode(): # We do not want to call get_attention_backend() in ONNX mode @@ -2159,8 +2184,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: @@ -2178,7 +2205,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, @@ -2213,51 +2240,57 @@ 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, - 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, + 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, @@ -2282,28 +2315,34 @@ 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, - 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, + 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 6eb3ce54f1..e5de6f7342 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 @@ -1648,6 +1648,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_padding_mask( batch_size: int, From 573edd800a417cda78e824322185f2426ff8f763 Mon Sep 17 00:00:00 2001 From: janEbert Date: Mon, 10 Aug 2026 21:21:57 +0200 Subject: [PATCH 2/6] Add padding-related tests Signed-off-by: janEbert --- .../test_dpa_mla_qkv_head_dim_pad.py | 140 ++++++++++++++++++ 1 file changed, 140 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 0000000000..febc0ddd64 --- /dev/null +++ b/tests/pytorch/attention/test_dpa_mla_qkv_head_dim_pad.py @@ -0,0 +1,140 @@ +# 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 50dbfef02d341680f1978fed50d3c3321687191f Mon Sep 17 00:00:00 2001 From: janEbert Date: Tue, 11 Aug 2026 01:09:48 +0200 Subject: [PATCH 3/6] Fix tests Signed-off-by: janEbert --- .../test_dpa_mla_qkv_head_dim_pad.py | 39 +++++++++++-------- 1 file changed, 23 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 febc0ddd64..e40ed5ec41 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 @@ -27,13 +28,15 @@ 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") @@ -115,26 +118,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 d24be6daa5aa16053e49f6978fce63861283c733 Mon Sep 17 00:00:00 2001 From: janEbert Date: Tue, 11 Aug 2026 01:21:09 +0200 Subject: [PATCH 4/6] Add memoization Signed-off-by: janEbert --- .../test_dpa_mla_qkv_head_dim_pad.py | 43 +++++++++++++++++ .../attention/dot_product_attention/utils.py | 48 +++++++++++++------ 2 files changed, 77 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 e40ed5ec41..fee2003932 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 @@ -89,6 +89,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 @@ -103,6 +108,44 @@ 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 e5de6f7342..771c8db9b3 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py @@ -72,6 +72,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: """ @@ -1650,7 +1656,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. @@ -1661,6 +1667,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 @@ -1673,23 +1684,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 bdc487ecdd8749e4376864c640d0d6dad8a92212 Mon Sep 17 00:00:00 2001 From: janEbert Date: Tue, 11 Aug 2026 01:28:59 +0200 Subject: [PATCH 5/6] 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 759432857a..5931c33607 100644 --- a/qa/L0_pytorch_unittest/test.sh +++ b/qa/L0_pytorch_unittest/test.sh @@ -63,6 +63,7 @@ python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_linear_mxfp8_att python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_fused_mla_q_uproj.xml $TE_PATH/tests/pytorch/attention/test_fused_mla_q_uproj.py || test_fail "test_fused_mla_q_uproj.py" 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_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 836f919570b8afae2f652f9f436c038cff1434b9 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 10 Aug 2026 23:31:44 +0000 Subject: [PATCH 6/6] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- .../test_dpa_mla_qkv_head_dim_pad.py | 86 +++++++++++++------ 1 file changed, 61 insertions(+), 25 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 fee2003932..2cca7e477c 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 @@ -20,7 +20,9 @@ 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 +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() @@ -28,8 +30,9 @@ 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): +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), @@ -50,26 +53,36 @@ def _thd_inputs(qk, v, t=32, h=4): 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, + 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 -]) +@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, + 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, @@ -77,10 +90,22 @@ def test_should_pad_qkv_head_dim(monkeypatch, native_unfused, padded_fused, expe # 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]) + 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). @@ -100,10 +125,16 @@ def fake_backend(p): 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, + 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 @@ -114,9 +145,14 @@ def test_should_pad_qkv_head_dim_is_memoized(monkeypatch): 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, + 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}