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 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..2cca7e477c --- /dev/null +++ b/tests/pytorch/attention/test_dpa_mla_qkv_head_dim_pad.py @@ -0,0 +1,226 @@ +# 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 + +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 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..771c8db9b3 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 @@ -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: """ @@ -1648,6 +1654,64 @@ 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. + + 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 + 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 + 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: + 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() def get_padding_mask( batch_size: int,