From a9cd697fa2fb1da2df187f49e706e3e698ae6dc5 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Tue, 7 Jul 2026 17:35:28 +0200 Subject: [PATCH 01/15] [PyTorch] DotProductAttention: declarative packed qkv/kv inputs Fused QKV projections naturally produce one packed buffer, but DotProductAttention forces callers to slice it into q/k/v views that TE then reverse-engineers with pointer-based layout detection (get_qkv_layout inspects data_ptr/storage_offset on every forward, which graph-breaks under torch.compile and adds CPU overhead). Let callers declare the packing instead (JAX-style): * DotProductAttention.forward gains optional qkv_layer (fully packed QKV: [b,s,3,h,d]/[s,b,3,h,d]/[b,s,h,3,d]/[s,b,h,3,d] dense, [t,3,h,d]/ [t,h,3,d] thd), kv_layer (packed KV used with query_layer), and qkv_interleave_dim (-3 or -2; explicit knob rather than shape inference since h==3 or hg==2 would be ambiguous). * Q/K/V are derived as zero-copy select() views and the exact layout enum (bs3hd, bsh3d, sb3hd, bshd_bs2hd, t3hd, ...) is constructed declaratively -- it is truthful by construction, so get_qkv_layout is never called on this path, including for thd and FP8 DPA. * combine_and_quantize no longer re-combines what is already combined: a new optional combined= argument carries the caller's original packed buffer, which is quantized directly instead of rebuilding the packed buffer from q/k/v views via combine_tensors (a raw set_ with a silent adjacency/interleave assumption). The packed original is threaded from DPA.forward through FusedAttention to FusedAttnFunc.forward; all legacy call sites are untouched (combined=None preserves exact behavior), and backward combine calls are unchanged (gradients have no pre-packed original). Tests: dense fwd+grad bit-exactness vs separate contiguous q/k/v for bs3hd/bsh3d/sb3hd/kv-packed/GQA (fused + flash), validation errors, torch.compile (no data_ptr/UntypedStorage graph breaks), FP8 combined-vs-views bit equivalence, and detection-free declared t3hd. Co-Authored-By: Claude Opus 4.8 Signed-off-by: Pawel Gadzinski --- .../attention/test_dpa_packed_inputs.py | 355 ++++++++++++++++++ .../dot_product_attention/backends.py | 12 +- .../dot_product_attention.py | 126 ++++++- .../attention/dot_product_attention/utils.py | 30 +- 4 files changed, 515 insertions(+), 8 deletions(-) create mode 100644 tests/pytorch/attention/test_dpa_packed_inputs.py diff --git a/tests/pytorch/attention/test_dpa_packed_inputs.py b/tests/pytorch/attention/test_dpa_packed_inputs.py new file mode 100644 index 0000000000..7e63630014 --- /dev/null +++ b/tests/pytorch/attention/test_dpa_packed_inputs.py @@ -0,0 +1,355 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Tests for declarative packed QKV/KV inputs to DotProductAttention. + +Instead of slicing a fused-projection buffer into q/k/v views (which TE then +reverse-engineers via pointer-based layout detection), callers can pass the +packed tensor directly (``qkv_layer``/``kv_layer`` + ``qkv_interleave_dim``). +Q/K/V are derived as zero-copy views and the exact layout string (e.g. +``bs3hd``) is declared, not detected -- including for thd and FP8 DPA. +""" + +import pytest +import torch + +import transformer_engine.pytorch # noqa: F401 (loads libtransformer_engine.so) +import transformer_engine_torch as tex +from transformer_engine.pytorch 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 +from transformer_engine.pytorch.attention.dot_product_attention.utils import ( + combine_and_quantize, +) +from transformer_engine.pytorch.cpp_extensions.fused_attn import ( + fused_attn_fwd, +) +from transformer_engine.pytorch.tensor.float8_tensor import Float8Quantizer + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is not available") + +_B, _S, _H, _D = 2, 128, 8, 64 +_DTYPE = torch.bfloat16 + + +def _cu_seqlens(): + return torch.arange(0, (_B + 1) * _S, _S, dtype=torch.int32, device="cuda") + + +def _fused_backend_supported(): + try: + q = torch.randn(_B, _S, _H, _D, dtype=_DTYPE, device="cuda") + fused_attn_fwd( + True, _S, _S, _cu_seqlens(), _cu_seqlens(), q, q.clone(), q.clone(), _DTYPE, + tex.NVTE_Fused_Attn_Backend.NVTE_F16_arbitrary_seqlen, + dropout=0.0, qkv_layout="bshd_bshd_bshd", o_format="bshd", + attn_bias_type="no_bias", attn_mask_type="no_mask", + ) + return True + except Exception: + return False + + +requires_fused = pytest.mark.skipif( + not (torch.cuda.is_available() and _fused_backend_supported()), + reason="F16_arbitrary_seqlen fused attention backend is not supported on this device", +) + + +def _force_backend(monkeypatch, backend): + """Force a single attention backend via env and invalidate the selection cache.""" + flash, fused = {"flash": ("1", "0"), "fused": ("0", "1")}[backend] + monkeypatch.setenv("NVTE_FLASH_ATTN", flash) + monkeypatch.setenv("NVTE_FUSED_ATTN", fused) + monkeypatch.setenv("NVTE_UNFUSED_ATTN", "0") + if backend == "flash": + # flash-attn bwd uses atomics unless deterministic + monkeypatch.setenv("NVTE_ALLOW_NONDETERMINISTIC_ALGO", "0") + dpa_module._attention_backends["backend_selection_requires_update"] = True + + +def _make_dpa(qkv_format, num_gqa_groups=None): + return DotProductAttention( + _H, + _D, + num_gqa_groups=num_gqa_groups, + attention_dropout=0.0, + qkv_format=qkv_format, + attn_mask_type="no_mask", + ) + + +def _assert_bit_exact(result, reference, names=("out", "dq", "dk", "dv")): + for name, x, y in zip(names, result, reference): + assert torch.equal(x.contiguous(), y.contiguous()), f"{name} differs" + + +def _fresh_parts(qkv_format, num_heads, seed=0): + torch.manual_seed(seed) + shape = (_B, _S, num_heads, _D) if qkv_format == "bshd" else (_S, _B, num_heads, _D) + return [torch.randn(*shape, dtype=_DTYPE, device="cuda") for _ in range(3)] + + +def _dpa_separate_baseline(qkv_format, num_gqa_groups=None): + """Fwd+bwd on contiguous separate q/k/v leaves; returns (out, dq, dk, dv).""" + hg = num_gqa_groups or _H + torch.manual_seed(0) + q_shape = (_B, _S, _H, _D) if qkv_format == "bshd" else (_S, _B, _H, _D) + kv_shape = (_B, _S, hg, _D) if qkv_format == "bshd" else (_S, _B, hg, _D) + q = torch.randn(*q_shape, dtype=_DTYPE, device="cuda") + k = torch.randn(*kv_shape, dtype=_DTYPE, device="cuda") + v = torch.randn(*kv_shape, dtype=_DTYPE, device="cuda") + q, k, v = [x.clone().requires_grad_() for x in (q, k, v)] + dpa_module._attention_backends["backend_selection_requires_update"] = True + out = _make_dpa(qkv_format, num_gqa_groups)(q, k, v) + out.backward(torch.ones_like(out)) + return out, q.grad, k.grad, v.grad + + +# --------------------------------------------------------------------------- +# 1. Dense eager equivalence (fused backend), fwd + input grads, bit-exact +# --------------------------------------------------------------------------- + + +@requires_fused +@pytest.mark.parametrize( + "qkv_format, interleave_dim", + [ + pytest.param("bshd", -3, id="bshd_qkv_dim-3"), # (a) bs3hd + pytest.param("bshd", -2, id="bshd_qkv_dim-2"), # (b) bsh3d (Megatron-style) + pytest.param("sbhd", -3, id="sbhd_qkv_dim-3"), # (c) sb3hd + ], +) +def test_dpa_fused_qkv_layer_dense(monkeypatch, qkv_format, interleave_dim): + """qkv_layer packed input is bit-exact vs separate contiguous q/k/v, and grads + flow back into the packed tensor itself.""" + _force_backend(monkeypatch, "fused") + reference = _dpa_separate_baseline(qkv_format) + + torch.manual_seed(0) + q_shape = (_B, _S, _H, _D) if qkv_format == "bshd" else (_S, _B, _H, _D) + parts = [torch.randn(*q_shape, dtype=_DTYPE, device="cuda") for _ in range(3)] + stack_dim = len(q_shape) + interleave_dim + 1 # -3 -> before h, -2 -> before d + qkv = torch.stack(parts, dim=stack_dim).requires_grad_() + + dpa_module._attention_backends["backend_selection_requires_update"] = True + out = _make_dpa(qkv_format)(qkv_layer=qkv, qkv_interleave_dim=interleave_dim) + out.backward(torch.ones_like(out)) + + assert qkv.grad is not None and qkv.grad.shape == qkv.shape + grads = [qkv.grad.select(stack_dim, i) for i in range(3)] + _assert_bit_exact((out, *grads), reference) + + +@requires_fused +@pytest.mark.parametrize( + "num_gqa_groups", + [pytest.param(None, id="mha_kv"), pytest.param(2, id="gqa_kv")], # (d) and (e) +) +def test_dpa_fused_kv_layer_dense(monkeypatch, num_gqa_groups): + """kv_layer packed input (with separate query) is bit-exact vs separate + contiguous q/k/v; grads flow into the packed kv tensor.""" + _force_backend(monkeypatch, "fused") + reference = _dpa_separate_baseline("bshd", num_gqa_groups) + + hg = num_gqa_groups or _H + torch.manual_seed(0) + q = torch.randn(_B, _S, _H, _D, dtype=_DTYPE, device="cuda") + k = torch.randn(_B, _S, hg, _D, dtype=_DTYPE, device="cuda") + v = torch.randn(_B, _S, hg, _D, dtype=_DTYPE, device="cuda") + q = q.clone().requires_grad_() + kv = torch.stack([k, v], dim=2).requires_grad_() # [b,s,2,hg,d] + + dpa_module._attention_backends["backend_selection_requires_update"] = True + out = _make_dpa("bshd", num_gqa_groups)(query_layer=q, kv_layer=kv) + out.backward(torch.ones_like(out)) + + assert kv.grad is not None and kv.grad.shape == kv.shape + _assert_bit_exact((out, q.grad, kv.grad[:, :, 0], kv.grad[:, :, 1]), reference) + + +# --------------------------------------------------------------------------- +# 2. Flash backend smoke +# --------------------------------------------------------------------------- + + +def test_dpa_flash_qkv_layer(monkeypatch): + """Flash backend: packed qkv_layer [b,s,3,h,d] is bit-exact vs separate.""" + _force_backend(monkeypatch, "flash") + try: + reference = _dpa_separate_baseline("bshd") + except Exception as exc: + pytest.skip(f"flash attention backend not available: {exc}") + + torch.manual_seed(0) + parts = [torch.randn(_B, _S, _H, _D, dtype=_DTYPE, device="cuda") for _ in range(3)] + qkv = torch.stack(parts, dim=2).requires_grad_() + dpa_module._attention_backends["backend_selection_requires_update"] = True + out = _make_dpa("bshd")(qkv_layer=qkv) + out.backward(torch.ones_like(out)) + grads = [qkv.grad[:, :, i] for i in range(3)] + _assert_bit_exact((out, *grads), reference) + + +# --------------------------------------------------------------------------- +# 3. Validation errors +# --------------------------------------------------------------------------- + + +def test_dpa_packed_input_validation(): + dpa = _make_dpa("bshd") + qkv = torch.randn(_B, _S, 3, _H, _D, dtype=_DTYPE, device="cuda") + kv = torch.randn(_B, _S, 2, _H, _D, dtype=_DTYPE, device="cuda") + k = torch.randn(_B, _S, _H, _D, dtype=_DTYPE, device="cuda") + + with pytest.raises(ValueError, match="must be None when qkv_layer is provided"): + dpa(qkv_layer=qkv, key_layer=k) + with pytest.raises(ValueError, match="query_layer is required when kv_layer"): + dpa(kv_layer=kv) + with pytest.raises(ValueError, match="qkv_interleave_dim must be -3"): + dpa(qkv_layer=qkv, qkv_interleave_dim=-1) + with pytest.raises(ValueError, match="mutually exclusive"): + dpa(qkv_layer=qkv, kv_layer=kv) + with pytest.raises(ValueError, match="must have size 3 at dim"): + dpa(qkv_layer=kv) # 2 at the interleave dim, not 3 + with pytest.raises(ValueError, match="required unless packed"): + dpa() + + +# --------------------------------------------------------------------------- +# 4. torch.compile: no data_ptr/UntypedStorage graph breaks with qkv_layer +# --------------------------------------------------------------------------- + + +@requires_fused +def test_dpa_torch_compile_qkv_layer_no_pointer_graph_breaks(monkeypatch): + _force_backend(monkeypatch, "fused") + torch._dynamo.reset() + torch._dynamo.utils.counters.clear() + + torch.manual_seed(0) + parts = [torch.randn(_B, _S, _H, _D, dtype=_DTYPE, device="cuda") for _ in range(3)] + qkv = torch.stack(parts, dim=2).requires_grad_() + dpa = _make_dpa("bshd") + + def fn(x): + return dpa(qkv_layer=x) + + dpa_module._attention_backends["backend_selection_requires_update"] = True + eager_out = fn(qkv) # eager warm-up: backend selection happens outside dynamo + compiled_out = torch.compile(fn)(qkv) + compiled_out.backward(torch.ones_like(compiled_out)) + + breaks = dict(torch._dynamo.utils.counters["graph_break"]) + torch._dynamo.reset() + pointer_breaks = { + reason: count + for reason, count in breaks.items() + if "data_ptr" in reason or "UntypedStorage" in reason + } + assert not pointer_breaks, f"pointer-based graph breaks with qkv_layer: {pointer_breaks}" + assert torch.equal(compiled_out, eager_out), "compiled output differs from eager" + + +# --------------------------------------------------------------------------- +# 5. FP8 combine refactor: combined= path is bit-identical to combine_tensors path +# --------------------------------------------------------------------------- + + +def _fp8_quantizer(): + return Float8Quantizer( + scale=torch.ones(1, dtype=torch.float32, device="cuda"), + amax=torch.zeros(1, dtype=torch.float32, device="cuda"), + fp8_dtype=tex.DType.kFloat8E4M3, + ) + + +def test_combine_and_quantize_combined_matches_views(): + """Quantizing the caller's packed buffer directly (combined=) produces the same + _data bits and scale_inv as rebuilding the packed buffer from q/k/v views via + combine_tensors (the old set_-based path).""" + torch.manual_seed(0) + qkv = torch.randn(_B, _S, 3, _H, _D, dtype=_DTYPE, device="cuda") + q, k, v = qkv[:, :, 0], qkv[:, :, 1], qkv[:, :, 2] + + old = combine_and_quantize("bs3hd", q, k, v, _fp8_quantizer()) + new = combine_and_quantize("bs3hd", q, k, v, _fp8_quantizer(), combined=qkv) + + assert old[3] == new[3] == "bs3hd" + for name, x, y in zip(("q", "k", "v"), old[:3], new[:3]): + assert torch.equal(x._data, y._data), f"{name} fp8 bits differ" + assert torch.equal(x._scale_inv, y._scale_inv), f"{name} scale_inv differs" + + +def test_combine_and_quantize_combined_kv_matches_views(): + """Same for the kv-packed (group 2) layout.""" + torch.manual_seed(0) + q = torch.randn(_B, _S, _H, _D, dtype=_DTYPE, device="cuda") + kv = torch.randn(_B, _S, 2, _H, _D, dtype=_DTYPE, device="cuda") + k, v = kv[:, :, 0], kv[:, :, 1] + + old = combine_and_quantize("bshd_bs2hd", q, k, v, _fp8_quantizer()) + new = combine_and_quantize("bshd_bs2hd", q, k, v, _fp8_quantizer(), combined=kv) + + for name, x, y in zip(("q", "k", "v"), old[:3], new[:3]): + assert torch.equal(x._data, y._data), f"{name} fp8 bits differ" + assert torch.equal(x._scale_inv, y._scale_inv), f"{name} scale_inv differs" + + +# --------------------------------------------------------------------------- +# 6. thd declarative: t3hd is declared, get_qkv_layout is never called +# --------------------------------------------------------------------------- + + +def test_dpa_thd_qkv_layer_declared_no_detection(monkeypatch): + """Packed thd input (qkv_layer [t,3,h,d]) declares 't3hd' without calling + get_qkv_layout; full forward+backward runs if a thd backend is available.""" + calls = [] + orig_get_qkv_layout = dpa_utils.get_qkv_layout + + def counting(*args, **kwargs): + calls.append(kwargs.get("qkv_format")) + return orig_get_qkv_layout(*args, **kwargs) + + monkeypatch.setattr(dpa_utils, "get_qkv_layout", counting) + + seen_layouts = [] + orig_get_backend = dpa_utils.get_attention_backend + + def recording(params): + seen_layouts.append(params.qkv_layout) + return orig_get_backend(params) + + monkeypatch.setattr(dpa_utils, "get_attention_backend", recording) + + torch.manual_seed(0) + t = _B * _S + qkv = torch.randn(t, 3, _H, _D, dtype=_DTYPE, device="cuda", requires_grad=True) + cu = _cu_seqlens() + dpa = DotProductAttention( + _H, _D, attention_dropout=0.0, qkv_format="thd", attn_mask_type="padding" + ) + dpa_module._attention_backends["backend_selection_requires_update"] = True + ran_full = True + try: + out = dpa( + qkv_layer=qkv, + cu_seqlens_q=cu, + cu_seqlens_kv=cu, + max_seqlen_q=_S, + max_seqlen_kv=_S, + ) + out.backward(torch.ones_like(out)) + assert qkv.grad is not None and qkv.grad.shape == qkv.shape + except ValueError: + # No thd-capable backend on this device; the layout step (the subject of + # this test) runs before backend dispatch, so the assertions still hold. + ran_full = False + + assert not calls, "get_qkv_layout must not be called for declared packed thd input" + assert seen_layouts == ["t3hd"], f"expected declared layout 't3hd', got {seen_layouts}" + if not ran_full: + pytest.skip("layout declaration verified; no thd backend available for full run") diff --git a/transformer_engine/pytorch/attention/dot_product_attention/backends.py b/transformer_engine/pytorch/attention/dot_product_attention/backends.py index 8f42983553..50565bebf3 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/backends.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/backends.py @@ -1328,6 +1328,7 @@ def forward( fp8_output, layer_number, return_max_logit, + packed_qkv=None, ): # pylint: disable=missing-function-docstring @@ -1390,7 +1391,13 @@ def forward( q_fp8, k_fp8, v_fp8 = q, k, v else: q_fp8, k_fp8, v_fp8, qkv_layout, qkv_scale_inv_format = combine_and_quantize( - qkv_layout, q, k, v, QKV_quantizer, used_in_backward=is_training + qkv_layout, + q, + k, + v, + QKV_quantizer, + used_in_backward=is_training, + combined=packed_qkv, ) # print quantizers @@ -1897,6 +1904,7 @@ def backward(ctx, d_out, *_args): None, None, None, + None, # packed_qkv ) @@ -1991,6 +1999,7 @@ def forward( score_mod_bprop: Optional[Callable] = None, score_mod_tensors: Optional[Dict[str, torch.Tensor]] = None, score_mod_bprop_tensors: Optional[Dict[str, torch.Tensor]] = None, + packed_qkv: Optional[torch.Tensor] = None, ) -> torch.Tensor: """fused attention fprop""" assert ( @@ -2210,6 +2219,7 @@ def forward( fp8_output, self.layer_number, self.return_max_logit, + packed_qkv, ) if self.return_max_logit: 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 03008bb2d7..e968817b97 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 @@ -1008,9 +1008,9 @@ def get_quantizer_roles( @no_torch_dynamo(recursive=False) def forward( self, - query_layer: torch.Tensor, - key_layer: torch.Tensor, - value_layer: torch.Tensor, + query_layer: Optional[torch.Tensor] = None, + key_layer: Optional[torch.Tensor] = None, + value_layer: Optional[torch.Tensor] = None, attention_mask: Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]] = None, qkv_format: str = None, cu_seqlens_q: torch.Tensor = None, @@ -1035,6 +1035,9 @@ def forward( score_mod_bprop: Optional[Callable] = None, score_mod_tensors: Optional[Dict[str, torch.Tensor]] = None, score_mod_bprop_tensors: Optional[Dict[str, torch.Tensor]] = None, + qkv_layer: Optional[torch.Tensor] = None, + kv_layer: Optional[torch.Tensor] = None, + qkv_interleave_dim: int = -3, ) -> torch.Tensor: r""" Dot Product Attention Layer. @@ -1245,8 +1248,113 @@ def forward( Runtime tensors exposed to score_mod_bprop as cuDNN graph tensors. Keys are user-defined string names consumed by the callback through ``tensors[name]``; there is no predefined set of accepted keys. + qkv_layer: Optional[torch.Tensor], default = None + Fully packed QKV tensor. When the QKV projection produces one packed buffer + (e.g. a fused QKV GEMM), it can be passed here directly instead of slicing + it into :attr:`query_layer`/:attr:`key_layer`/:attr:`value_layer` views. + For :attr:`qkv_format` = {"bshd", "sbhd"}, it must be a 5D tensor of shape + ``[b, s, 3, h, d]``/``[s, b, 3, h, d]`` (:attr:`qkv_interleave_dim` = -3) or + ``[b, s, h, 3, d]``/``[s, b, h, 3, d]`` (:attr:`qkv_interleave_dim` = -2); + for :attr:`qkv_format` = "thd", a 4D tensor of shape ``[t, 3, h, d]`` or + ``[t, h, 3, d]``. Q/K/V are derived as zero-copy views and the memory layout + (e.g. ``bs3hd``) is declared from the packing itself, so no pointer-based + layout detection runs on this path -- including for "thd" and FP8 attention. + Mutually exclusive with :attr:`query_layer`, :attr:`key_layer`, + :attr:`value_layer` and :attr:`kv_layer`. + kv_layer: Optional[torch.Tensor], default = None + Packed KV tensor, used together with :attr:`query_layer` + (e.g. ``[b, s, 2, hg, d]`` for :attr:`qkv_interleave_dim` = -3, or + ``[b, s, hg, 2, d]`` for :attr:`qkv_interleave_dim` = -2). K/V are derived + as zero-copy views and the layout (e.g. ``bshd_bs2hd``) is declared, not + detected. Mutually exclusive with :attr:`key_layer`, :attr:`value_layer` + and :attr:`qkv_layer`. + qkv_interleave_dim: int, default = -3 + Dimension of :attr:`qkv_layer`/:attr:`kv_layer` where the 3 (QKV) or 2 (KV) + interleave sits; must be -3 (e.g. ``bs3hd``) or -2 (e.g. ``bsh3d``, + Megatron-style). This is an explicit knob rather than shape inference, + since e.g. ``h == 3`` would make the shapes ambiguous. """ + # Declarative packed inputs: derive q/k/v as zero-copy views of the packed + # buffer and construct the exact layout string from the declaration. The + # layout enum is truthful by construction, so no pointer-based detection + # is needed downstream (this also covers thd and FP8 DPA). + packed_tensor = None + declared_qkv_layout = None + if qkv_layer is not None or kv_layer is not None: + if qkv_layer is not None and kv_layer is not None: + raise ValueError("qkv_layer and kv_layer are mutually exclusive.") + if inference_params is not None: + raise ValueError( + "Packed inputs (qkv_layer/kv_layer) are not supported with KV caching" + " (inference_params); pass separate query/key/value tensors instead." + ) + if qkv_interleave_dim not in (-3, -2): + raise ValueError( + "qkv_interleave_dim must be -3 (e.g. bs3hd) or -2 (e.g. bsh3d), got" + f" {qkv_interleave_dim}." + ) + packed_format = qkv_format if qkv_format is not None else self.qkv_format + + def _packed_layout(fmt: str, num: int) -> str: + # bshd + 3 @ -3 -> bs3hd; bshd + 3 @ -2 -> bsh3d; thd + 2 @ -2 -> th2d + pos = len(fmt) + qkv_interleave_dim + 1 + return fmt[:pos] + str(num) + fmt[pos:] + + if qkv_layer is not None: + if any(x is not None for x in (query_layer, key_layer, value_layer)): + raise ValueError( + "qkv_layer already packs Q, K and V: query_layer, key_layer and" + " value_layer must be None when qkv_layer is provided." + ) + expected_rank = 4 if packed_format == "thd" else 5 + if qkv_layer.dim() != expected_rank: + raise ValueError( + f"qkv_layer must be a {expected_rank}D tensor for" + f" qkv_format={packed_format!r}, got {qkv_layer.dim()}D." + ) + if qkv_layer.shape[qkv_interleave_dim] != 3: + raise ValueError( + f"qkv_layer must have size 3 at dim {qkv_interleave_dim}" + f" (qkv_interleave_dim), got shape {tuple(qkv_layer.shape)}." + ) + query_layer, key_layer, value_layer = ( + qkv_layer.select(qkv_interleave_dim, i) for i in range(3) + ) + packed_tensor = qkv_layer + declared_qkv_layout = _packed_layout(packed_format, 3) + else: + if query_layer is None: + raise ValueError( + "kv_layer packs only K and V: query_layer is required when kv_layer" + " is provided." + ) + if key_layer is not None or value_layer is not None: + raise ValueError( + "kv_layer already packs K and V: key_layer and value_layer must be" + " None when kv_layer is provided." + ) + if kv_layer.dim() != query_layer.dim() + 1: + raise ValueError( + f"kv_layer must have one more dimension than query_layer, got" + f" {kv_layer.dim()}D kv_layer and {query_layer.dim()}D query_layer." + ) + if kv_layer.shape[qkv_interleave_dim] != 2: + raise ValueError( + f"kv_layer must have size 2 at dim {qkv_interleave_dim}" + f" (qkv_interleave_dim), got shape {tuple(kv_layer.shape)}." + ) + key_layer, value_layer = ( + kv_layer.select(qkv_interleave_dim, i) for i in range(2) + ) + packed_tensor = kv_layer + declared_qkv_layout = f"{packed_format}_{_packed_layout(packed_format, 2)}" + elif query_layer is None or key_layer is None or value_layer is None: + raise ValueError( + "query_layer, key_layer and value_layer are required unless packed" + " inputs (qkv_layer or query_layer + kv_layer) are provided." + ) + with self.prepare_forward_ctx( query_layer, num_gemms=3, @@ -1432,7 +1540,15 @@ def forward( cu_seqlens_kv_padded = None # get qkv's memory layout - if all( + if declared_qkv_layout is not None: + # Packed inputs (qkv_layer/kv_layer) declare the layout: the enum is + # truthful by construction, so the pointer-based detection in + # get_qkv_layout is skipped entirely -- for dense, thd (t3hd/th3d) + # and FP8 DPA alike. + qkv_layout = declared_qkv_layout + q_format = qkv_format + kv_format = qkv_format + elif all( isinstance(x, Float8TensorStorage) for x in [query_layer, key_layer, value_layer] ): ( @@ -1812,6 +1928,7 @@ def forward( inference_params=inference_params, softmax_offset=softmax_offset, fp8_output=fp8_output, + packed_qkv=packed_tensor, ) return self.fused_attention( query_layer, @@ -1847,6 +1964,7 @@ def forward( score_mod_bprop=score_mod_bprop, score_mod_tensors=score_mod_tensors, score_mod_bprop_tensors=score_mod_bprop_tensors, + packed_qkv=packed_tensor, ) if use_unfused_attention: diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py index 9913b78dfc..3176200ffa 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py @@ -2783,8 +2783,18 @@ def combine_and_quantize( used_in_forward=True, used_in_backward=False, keep_same_data_and_scale_inv_format=False, + combined: Optional[torch.Tensor] = None, ): - """Combine Q, K, V tensors based on qkv_layout and quantize them together.""" + """Combine Q, K, V tensors based on qkv_layout and quantize them together. + + When ``combined`` is provided, it must be the caller's original packed buffer + matching the packed group in ``qkv_layout`` (the full QKV tensor for ``3`` + layouts such as ``bs3hd``, or the KV tensor for ``2`` layouts such as + ``bshd_bs2hd``). It is then quantized directly instead of re-deriving the + packed buffer from the q/k/v views via ``combine_tensors`` (which rebuilds it + with a raw ``set_`` under a silent adjacency/interleave assumption). Ignored + for MXFP8 quantization. + """ if isinstance(qkv_quantizer, MXFP8Quantizer): qkv_format, q_format, kv_format = get_qkv_format(qkv_layout) assert qkv_format in ("bshd", "sbhd"), ( @@ -2884,12 +2894,26 @@ def combine_and_quantize( match qkv_group: case 1: dim = qkv_layout.find("3") - qkv = combine_tensors([q, k, v], dim) + if combined is not None: + assert combined.shape[dim] == 3, ( + f"combined QKV tensor does not match qkv_layout {qkv_layout}: expected" + f" size 3 at dim {dim}, got shape {tuple(combined.shape)}." + ) + qkv = combined + else: + qkv = combine_tensors([q, k, v], dim) qkv_fp8 = qkv_quantizer(qkv) q_data, k_data, v_data = SplitAlongDim.apply(qkv_fp8._data, dim, [1, 1, 1], True) case 2: dim = qkv_layout.split("_")[1].find("2") - kv = combine_tensors([k, v], dim) + if combined is not None: + assert combined.shape[dim] == 2, ( + f"combined KV tensor does not match qkv_layout {qkv_layout}: expected" + f" size 2 at dim {dim}, got shape {tuple(combined.shape)}." + ) + kv = combined + else: + kv = combine_tensors([k, v], dim) tensors = [q, kv] num_tensors = len(tensors) shapes = [x.shape for x in tensors] From 5f70d121989747878eb6baa5015b647aee98cd93 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Wed, 8 Jul 2026 12:55:58 +0200 Subject: [PATCH 02/15] [PyTorch] MultiheadAttention: pass packed projection output to DPA declaratively Adopt the new DotProductAttention packed API inside MultiheadAttention: * self-attention (np == ng): the fused QKV projection output, already viewed as [.., h, 3, d] (qkv_weight_interleaved) or [.., 3, h, d], is handed to DPA directly as qkv_layer with the matching qkv_interleave_dim (-2 / -3) -- no SplitAlongDim slicing in MHA and no pointer-based layout detection in DPA. * cross-attention: the packed KV projection output is exposed as [.., hg, 2, d] / [.., 2, hg, d] and passed as kv_layer. * The pass-through only engages when no per-tensor operation needs the individual q/k/v slices: it is skipped for RoPE, QK normalization, KV caching (inference_params), CPU offloading, GQA (np != ng, not a uniform 3-interleave) and quantized (FP8) projection outputs; those keep the legacy sliced-views path unchanged. Tests: MHA self (interleaved + non-interleaved) and cross (both interleaves) are bit-exact vs the same MHA with packed inputs converted back to separate contiguous q/k/v (output, input grad, weight grads); spy asserts the packed argument and interleave dim actually reach DPA; GQA and RoPE fall back to the views path. TransformerLayer regression suite unchanged; test_kv_cache failures on this device are pre-existing on origin/main (verified). Co-Authored-By: Claude Opus 4.8 Signed-off-by: Pawel Gadzinski --- .../attention/test_dpa_packed_inputs.py | 201 +++++++++++++++++- .../pytorch/attention/multi_head_attention.py | 143 +++++++++---- 2 files changed, 300 insertions(+), 44 deletions(-) diff --git a/tests/pytorch/attention/test_dpa_packed_inputs.py b/tests/pytorch/attention/test_dpa_packed_inputs.py index 7e63630014..b87e20a344 100644 --- a/tests/pytorch/attention/test_dpa_packed_inputs.py +++ b/tests/pytorch/attention/test_dpa_packed_inputs.py @@ -16,7 +16,7 @@ import transformer_engine.pytorch # noqa: F401 (loads libtransformer_engine.so) import transformer_engine_torch as tex -from transformer_engine.pytorch import DotProductAttention +from transformer_engine.pytorch import DotProductAttention, MultiheadAttention from transformer_engine.pytorch.attention.dot_product_attention import ( dot_product_attention as dpa_module, ) @@ -353,3 +353,202 @@ def recording(params): assert seen_layouts == ["t3hd"], f"expected declared layout 't3hd', got {seen_layouts}" if not ran_full: pytest.skip("layout declaration verified; no thd backend available for full run") + + +# --------------------------------------------------------------------------- +# 7. MultiheadAttention adoption: the fused QKV/KV projection output is passed +# packed (qkv_layer/kv_layer) to DotProductAttention +# --------------------------------------------------------------------------- + + +def _spy_packed_dpa(monkeypatch, record): + """Record (qkv_layer given, kv_layer given, qkv_interleave_dim) per DPA call.""" + orig = DotProductAttention.forward + + def spy(self, *args, **kwargs): + record.append( + ( + kwargs.get("qkv_layer") is not None, + kwargs.get("kv_layer") is not None, + kwargs.get("qkv_interleave_dim", None), + ) + ) + return orig(self, *args, **kwargs) + + monkeypatch.setattr(DotProductAttention, "forward", spy) + + +def _strip_packed_dpa(monkeypatch): + """Reference path: convert packed DPA inputs back to separate contiguous q/k/v.""" + orig = DotProductAttention.forward + + def stripped(self, query_layer=None, key_layer=None, value_layer=None, *args, **kwargs): + qkv = kwargs.pop("qkv_layer", None) + kv = kwargs.pop("kv_layer", None) + dim = kwargs.pop("qkv_interleave_dim", -3) + if qkv is not None: + query_layer, key_layer, value_layer = ( + qkv.select(dim, i).contiguous() for i in range(3) + ) + elif kv is not None: + key_layer, value_layer = (kv.select(dim, i).contiguous() for i in range(2)) + return orig(self, query_layer, key_layer, value_layer, *args, **kwargs) + + monkeypatch.setattr(DotProductAttention, "forward", stripped) + + +def _run_mha(mha, x, encoder_output=None): + dpa_module._attention_backends["backend_selection_requires_update"] = True + if encoder_output is not None: + out = mha(x, encoder_output=encoder_output) + else: + out = mha(x) + out.backward(torch.ones_like(out)) + wgrads = [p.grad.clone() for p in mha.parameters() if p.grad is not None] + xgrad = x.grad.clone() + x.grad = None + mha.zero_grad(set_to_none=True) + return out, xgrad, wgrads + + +def _assert_mha_equal(result, reference): + out, xgrad, wgrads = result + out_ref, xgrad_ref, wgrads_ref = reference + assert torch.equal(out, out_ref), "output differs" + assert torch.equal(xgrad, xgrad_ref), "input grad differs" + assert len(wgrads) == len(wgrads_ref) + for i, (w, w_ref) in enumerate(zip(wgrads, wgrads_ref)): + assert torch.equal(w, w_ref), f"weight grad {i} differs" + + +@requires_fused +@pytest.mark.parametrize( + "interleaved", [pytest.param(True, id="interleaved"), pytest.param(False, id="non_interleaved")] +) +def test_mha_self_attention_packed_pass_through(monkeypatch, interleaved): + """MHA self-attention passes its packed projection output straight to DPA as + qkv_layer (with the matching interleave dim), bit-exact vs the same MHA with + packed inputs converted back to separate contiguous q/k/v.""" + _force_backend(monkeypatch, "fused") + hidden = _H * _D + torch.manual_seed(0) + mha = MultiheadAttention( + hidden, + _H, + attention_dropout=0.0, + attn_mask_type="no_mask", + qkv_format="sbhd", + fuse_qkv_params=True, + qkv_weight_interleaved=interleaved, + params_dtype=_DTYPE, + device="cuda", + ) + torch.manual_seed(1) + x = torch.randn(_S, _B, hidden, dtype=_DTYPE, device="cuda", requires_grad=True) + + record = [] + _spy_packed_dpa(monkeypatch, record) + result = _run_mha(mha, x) + assert record == [(True, False, -2 if interleaved else -3)], f"unexpected DPA call: {record}" + monkeypatch.undo() + + _force_backend(monkeypatch, "fused") + _strip_packed_dpa(monkeypatch) + reference = _run_mha(mha, x) + _assert_mha_equal(result, reference) + + +@requires_fused +@pytest.mark.parametrize( + "interleaved", [pytest.param(True, id="interleaved"), pytest.param(False, id="non_interleaved")] +) +def test_mha_cross_attention_packed_kv_pass_through(monkeypatch, interleaved): + """MHA cross-attention passes its packed KV projection output to DPA as + kv_layer, bit-exact vs the separate contiguous reference.""" + _force_backend(monkeypatch, "fused") + hidden = _H * _D + torch.manual_seed(0) + mha = MultiheadAttention( + hidden, + _H, + attention_dropout=0.0, + attn_mask_type="no_mask", + qkv_format="sbhd", + attention_type="cross", + fuse_qkv_params=True, + qkv_weight_interleaved=interleaved, + params_dtype=_DTYPE, + device="cuda", + ) + torch.manual_seed(1) + x = torch.randn(_S, _B, hidden, dtype=_DTYPE, device="cuda", requires_grad=True) + enc = torch.randn(_S, _B, hidden, dtype=_DTYPE, device="cuda") + + record = [] + _spy_packed_dpa(monkeypatch, record) + result = _run_mha(mha, x, encoder_output=enc) + assert record == [(False, True, -2 if interleaved else -3)], f"unexpected DPA call: {record}" + monkeypatch.undo() + + _force_backend(monkeypatch, "fused") + _strip_packed_dpa(monkeypatch) + reference = _run_mha(mha, x, encoder_output=enc) + _assert_mha_equal(result, reference) + + +@requires_fused +def test_mha_gqa_falls_back_to_views(monkeypatch): + """GQA (np != ng) is not a uniform 3-interleave: MHA must keep the legacy + sliced-views path and still work.""" + _force_backend(monkeypatch, "fused") + hidden = _H * _D + torch.manual_seed(0) + mha = MultiheadAttention( + hidden, + _H, + num_gqa_groups=2, + attention_dropout=0.0, + attn_mask_type="no_mask", + qkv_format="sbhd", + fuse_qkv_params=True, + params_dtype=_DTYPE, + device="cuda", + ) + torch.manual_seed(1) + x = torch.randn(_S, _B, hidden, dtype=_DTYPE, device="cuda", requires_grad=True) + + record = [] + _spy_packed_dpa(monkeypatch, record) + out, _, _ = _run_mha(mha, x) + assert record == [(False, False, -3)], f"GQA must not use the packed path: {record}" + assert out.shape == (_S, _B, hidden) + + +@requires_fused +def test_mha_rope_falls_back_to_views(monkeypatch): + """RoPE needs the individual q/k slices: MHA must keep the legacy path.""" + _force_backend(monkeypatch, "fused") + from transformer_engine.pytorch.attention.rope import RotaryPositionEmbedding + + hidden = _H * _D + torch.manual_seed(0) + mha = MultiheadAttention( + hidden, + _H, + attention_dropout=0.0, + attn_mask_type="no_mask", + qkv_format="sbhd", + fuse_qkv_params=True, + params_dtype=_DTYPE, + device="cuda", + ) + rope = RotaryPositionEmbedding(_D)(max_seq_len=_S).to("cuda") + torch.manual_seed(1) + x = torch.randn(_S, _B, hidden, dtype=_DTYPE, device="cuda", requires_grad=True) + + record = [] + _spy_packed_dpa(monkeypatch, record) + dpa_module._attention_backends["backend_selection_requires_update"] = True + out = mha(x, rotary_pos_emb=rope) + assert record == [(False, False, -3)], f"RoPE must not use the packed path: {record}" + assert out.shape == (_S, _B, hidden) diff --git a/transformer_engine/pytorch/attention/multi_head_attention.py b/transformer_engine/pytorch/attention/multi_head_attention.py index 70ae9dfc21..c3b32e7c03 100644 --- a/transformer_engine/pytorch/attention/multi_head_attention.py +++ b/transformer_engine/pytorch/attention/multi_head_attention.py @@ -9,6 +9,7 @@ import torch from transformer_engine.pytorch.quantization import FP8GlobalStateManager, QuantizerRole +from transformer_engine.pytorch.quantized_tensor import QuantizedTensorStorage from transformer_engine.pytorch.tensor.float8_tensor import Float8Tensor from transformer_engine.pytorch.module.base import TransformerEngineBaseModule from transformer_engine.pytorch.module import LayerNormLinear, Linear, RMSNorm, LayerNorm @@ -903,6 +904,22 @@ def forward( self._update_output_quantizer_roles(qkv_fp8_output, proj_fp8_grad, dpa_fp8_output) + # Packed pass-through to DotProductAttention: the fused QKV/KV projection + # already produces one packed buffer, which DPA accepts directly via its + # declarative qkv_layer/kv_layer arguments (deriving q/k/v as zero-copy + # views and skipping pointer-based layout detection). Only possible when + # no per-tensor operation (RoPE, QK normalization, KV caching, CPU + # offloading) needs the individual q/k/v slices. + packed_dpa_eligible = ( + rotary_pos_emb is None + and self.q_norm is None + and inference_params is None + and not is_cpu_offload_enabled() + ) + packed_qkv_layer = None + packed_kv_layer = None + packed_interleave_dim = -3 + layernorm_output = None if self.attention_type == "self": # Attention heads [sq, b, h] --> [sq, b, ng * (np/ng + 2) * hn] @@ -947,28 +964,43 @@ def forward( mixed_x_layer = mixed_x_layer.view(*new_tensor_shape) - # qkv_weight_interleaved: - # [sq, b, ng, (np/ng + 2), hn] - # --> [sq, b, ng, np/ng, hn], [sq, b, ng, 1, hn], [sq, b, ng, 1, hn] - # not qkv_weight_interleaved: - # [sq, b, (np/ng + 2), ng, hn] - # --> [sq, b, np/ng, np, hn], [sq, b, 1, ng, hn], [sq, b, 1, ng, hn] - query_layer, key_layer, value_layer = SplitAlongDim.apply( - mixed_x_layer, split_dim, (num_queries_per_key_value, 1, 1) - ) - - if self.qkv_format == "thd": - query_layer, key_layer, value_layer = ( - x.reshape(x.size(0), -1, self.hidden_size_per_attention_head) - for x in (query_layer, key_layer, value_layer) - ) + if ( + num_queries_per_key_value == 1 + and packed_dpa_eligible + and not isinstance(mixed_x_layer, QuantizedTensorStorage) + ): + # np == ng: the projection output is a uniform 3-interleave + # ([.., h, 3, d] interleaved / [.., 3, h, d] otherwise), which + # DotProductAttention accepts directly as a declared packed + # qkv_layer -- no slicing here, no layout detection there. + packed_qkv_layer = mixed_x_layer + packed_interleave_dim = split_dim + query_layer = None + key_layer = None + value_layer = None else: - # query: -> [sq, b, np, hn] - # key, value: -> [sq, b, ng, hn] - query_layer, key_layer, value_layer = ( - x.reshape(x.size(0), x.size(1), -1, self.hidden_size_per_attention_head) - for x in (query_layer, key_layer, value_layer) + # qkv_weight_interleaved: + # [sq, b, ng, (np/ng + 2), hn] + # --> [sq, b, ng, np/ng, hn], [sq, b, ng, 1, hn], [sq, b, ng, 1, hn] + # not qkv_weight_interleaved: + # [sq, b, (np/ng + 2), ng, hn] + # --> [sq, b, np/ng, np, hn], [sq, b, 1, ng, hn], [sq, b, 1, ng, hn] + query_layer, key_layer, value_layer = SplitAlongDim.apply( + mixed_x_layer, split_dim, (num_queries_per_key_value, 1, 1) ) + + if self.qkv_format == "thd": + query_layer, key_layer, value_layer = ( + x.reshape(x.size(0), -1, self.hidden_size_per_attention_head) + for x in (query_layer, key_layer, value_layer) + ) + else: + # query: -> [sq, b, np, hn] + # key, value: -> [sq, b, ng, hn] + query_layer, key_layer, value_layer = ( + x.reshape(x.size(0), x.size(1), -1, self.hidden_size_per_attention_head) + for x in (query_layer, key_layer, value_layer) + ) elif self.attention_type == "cross": # Attention heads [sk, b, h] --> [sk, b, (ng * 2 * hn)] mixed_kv_layer = self.key_value( @@ -996,34 +1028,56 @@ def forward( mixed_kv_layer = mixed_kv_layer.view(*new_tensor_shape) - # mixed_kv_layer --> 2 [sk, b, ng, hn] - key_layer, value_layer = SplitAlongDim.apply( - mixed_kv_layer, - split_dim, - mixed_kv_layer.shape[split_dim] // 2, - ) - key_layer, value_layer = ( - x.reshape( - x.size(0), - x.size(1), - -1, - self.hidden_size_per_attention_head, - ) - for x in (key_layer, value_layer) - ) - - if self.qkv_format == "thd": - key_layer, value_layer = ( - x.reshape(x.size(0), -1, self.hidden_size_per_attention_head) - for x in (key_layer, value_layer) - ) + if packed_dpa_eligible and not isinstance(mixed_kv_layer, QuantizedTensorStorage): + # Declare the packed KV to DotProductAttention instead of + # slicing it: expose the 2-interleave as its own dimension. + if self.qkv_weight_interleaved: + # [.., ng, 2 * hn] --> [.., ng, 2, hn] + packed_kv_shape = mixed_kv_layer.size()[:-1] + ( + 2, + self.hidden_size_per_attention_head, + ) + packed_interleave_dim = -2 + else: + # [.., 2 * ng, hn] --> [.., 2, ng, hn] + packed_kv_shape = mixed_kv_layer.size()[:-2] + ( + 2, + self.num_gqa_groups_per_partition, + self.hidden_size_per_attention_head, + ) + packed_interleave_dim = -3 + packed_kv_layer = mixed_kv_layer.view(*packed_kv_shape) + key_layer = None + value_layer = None else: - # key, value: -> [sq, b, ng, hn] + # mixed_kv_layer --> 2 [sk, b, ng, hn] + key_layer, value_layer = SplitAlongDim.apply( + mixed_kv_layer, + split_dim, + mixed_kv_layer.shape[split_dim] // 2, + ) key_layer, value_layer = ( - x.reshape(x.size(0), x.size(1), -1, self.hidden_size_per_attention_head) + x.reshape( + x.size(0), + x.size(1), + -1, + self.hidden_size_per_attention_head, + ) for x in (key_layer, value_layer) ) + if self.qkv_format == "thd": + key_layer, value_layer = ( + x.reshape(x.size(0), -1, self.hidden_size_per_attention_head) + for x in (key_layer, value_layer) + ) + else: + # key, value: -> [sq, b, ng, hn] + key_layer, value_layer = ( + x.reshape(x.size(0), x.size(1), -1, self.hidden_size_per_attention_head) + for x in (key_layer, value_layer) + ) + # Attention head [sq, b, h] --> [sq, b, hp] if self.input_layernorm: layernorm_query_outputs = self.layernorm_query( @@ -1143,6 +1197,9 @@ def forward( inference_params=inference_params, pad_between_seqs=pad_between_seqs, fp8_output=dpa_fp8_output, + qkv_layer=packed_qkv_layer, + kv_layer=packed_kv_layer, + qkv_interleave_dim=packed_interleave_dim, ) # =================== From 5713048326289dceaae168d940ec267844f330c9 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Thu, 9 Jul 2026 14:52:12 +0200 Subject: [PATCH 03/15] [PyTorch] DotProductAttention: factor packed-input handling into _unpack_packed_qkv Signed-off-by: Pawel Gadzinski --- .../dot_product_attention.py | 187 ++++++++++-------- 1 file changed, 109 insertions(+), 78 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 e968817b97..8feca1b630 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 @@ -197,6 +197,104 @@ def _trim_output(attn_out, num_attention_heads, padded_head_dim_v, orig_head_dim return attn_out[..., :orig_head_dim_v].reshape(*out_shape, -1) +def _unpack_packed_qkv( + qkv_layer: Optional[torch.Tensor], + kv_layer: Optional[torch.Tensor], + query_layer: Optional[torch.Tensor], + key_layer: Optional[torch.Tensor], + value_layer: Optional[torch.Tensor], + qkv_format: str, + qkv_interleave_dim: int, + inference_params: Optional[InferenceParams], +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, Optional[torch.Tensor], Optional[str]]: + """Resolve declarative packed inputs into q/k/v. + + Derives q/k/v as zero-copy views of the packed buffer (``qkv_layer`` or + ``kv_layer``) and constructs the exact layout string from the declaration. + The layout enum is truthful by construction, so no pointer-based detection + is needed downstream (this also covers thd and FP8 DPA). + + Returns ``(query_layer, key_layer, value_layer, packed_tensor, + declared_qkv_layout)``; the last two are ``None`` when no packed input is + given. + """ + if qkv_layer is None and kv_layer is None: + if query_layer is None or key_layer is None or value_layer is None: + raise ValueError( + "query_layer, key_layer and value_layer are required unless packed" + " inputs (qkv_layer or query_layer + kv_layer) are provided." + ) + return query_layer, key_layer, value_layer, None, None + + if qkv_layer is not None and kv_layer is not None: + raise ValueError("qkv_layer and kv_layer are mutually exclusive.") + if inference_params is not None: + raise ValueError( + "Packed inputs (qkv_layer/kv_layer) are not supported with KV caching" + " (inference_params); pass separate query/key/value tensors instead." + ) + if qkv_interleave_dim not in (-3, -2): + raise ValueError( + "qkv_interleave_dim must be -3 (e.g. bs3hd) or -2 (e.g. bsh3d), got" + f" {qkv_interleave_dim}." + ) + + def _packed_layout(fmt: str, num: int) -> str: + # bshd + 3 @ -3 -> bs3hd; bshd + 3 @ -2 -> bsh3d; thd + 2 @ -2 -> th2d + pos = len(fmt) + qkv_interleave_dim + 1 + return fmt[:pos] + str(num) + fmt[pos:] + + if qkv_layer is not None: + if any(x is not None for x in (query_layer, key_layer, value_layer)): + raise ValueError( + "qkv_layer already packs Q, K and V: query_layer, key_layer and" + " value_layer must be None when qkv_layer is provided." + ) + expected_rank = 4 if qkv_format == "thd" else 5 + if qkv_layer.dim() != expected_rank: + raise ValueError( + f"qkv_layer must be a {expected_rank}D tensor for" + f" qkv_format={qkv_format!r}, got {qkv_layer.dim()}D." + ) + if qkv_layer.shape[qkv_interleave_dim] != 3: + raise ValueError( + f"qkv_layer must have size 3 at dim {qkv_interleave_dim}" + f" (qkv_interleave_dim), got shape {tuple(qkv_layer.shape)}." + ) + query_layer, key_layer, value_layer = ( + qkv_layer.select(qkv_interleave_dim, i) for i in range(3) + ) + return query_layer, key_layer, value_layer, qkv_layer, _packed_layout(qkv_format, 3) + + if query_layer is None: + raise ValueError( + "kv_layer packs only K and V: query_layer is required when kv_layer" " is provided." + ) + if key_layer is not None or value_layer is not None: + raise ValueError( + "kv_layer already packs K and V: key_layer and value_layer must be" + " None when kv_layer is provided." + ) + if kv_layer.dim() != query_layer.dim() + 1: + raise ValueError( + f"kv_layer must have one more dimension than query_layer, got" + f" {kv_layer.dim()}D kv_layer and {query_layer.dim()}D query_layer." + ) + if kv_layer.shape[qkv_interleave_dim] != 2: + raise ValueError( + f"kv_layer must have size 2 at dim {qkv_interleave_dim}" + f" (qkv_interleave_dim), got shape {tuple(kv_layer.shape)}." + ) + key_layer, value_layer = (kv_layer.select(qkv_interleave_dim, i) for i in range(2)) + return ( + query_layer, + key_layer, + value_layer, + kv_layer, + f"{qkv_format}_{_packed_layout(qkv_format, 2)}", + ) + + class DotProductAttention(TransformerEngineBaseModule): r"""Allows the model to jointly attend to information from different representation subspaces as described in the paper: @@ -1275,85 +1373,18 @@ def forward( since e.g. ``h == 3`` would make the shapes ambiguous. """ - # Declarative packed inputs: derive q/k/v as zero-copy views of the packed - # buffer and construct the exact layout string from the declaration. The - # layout enum is truthful by construction, so no pointer-based detection - # is needed downstream (this also covers thd and FP8 DPA). - packed_tensor = None - declared_qkv_layout = None - if qkv_layer is not None or kv_layer is not None: - if qkv_layer is not None and kv_layer is not None: - raise ValueError("qkv_layer and kv_layer are mutually exclusive.") - if inference_params is not None: - raise ValueError( - "Packed inputs (qkv_layer/kv_layer) are not supported with KV caching" - " (inference_params); pass separate query/key/value tensors instead." - ) - if qkv_interleave_dim not in (-3, -2): - raise ValueError( - "qkv_interleave_dim must be -3 (e.g. bs3hd) or -2 (e.g. bsh3d), got" - f" {qkv_interleave_dim}." - ) - packed_format = qkv_format if qkv_format is not None else self.qkv_format - - def _packed_layout(fmt: str, num: int) -> str: - # bshd + 3 @ -3 -> bs3hd; bshd + 3 @ -2 -> bsh3d; thd + 2 @ -2 -> th2d - pos = len(fmt) + qkv_interleave_dim + 1 - return fmt[:pos] + str(num) + fmt[pos:] - - if qkv_layer is not None: - if any(x is not None for x in (query_layer, key_layer, value_layer)): - raise ValueError( - "qkv_layer already packs Q, K and V: query_layer, key_layer and" - " value_layer must be None when qkv_layer is provided." - ) - expected_rank = 4 if packed_format == "thd" else 5 - if qkv_layer.dim() != expected_rank: - raise ValueError( - f"qkv_layer must be a {expected_rank}D tensor for" - f" qkv_format={packed_format!r}, got {qkv_layer.dim()}D." - ) - if qkv_layer.shape[qkv_interleave_dim] != 3: - raise ValueError( - f"qkv_layer must have size 3 at dim {qkv_interleave_dim}" - f" (qkv_interleave_dim), got shape {tuple(qkv_layer.shape)}." - ) - query_layer, key_layer, value_layer = ( - qkv_layer.select(qkv_interleave_dim, i) for i in range(3) - ) - packed_tensor = qkv_layer - declared_qkv_layout = _packed_layout(packed_format, 3) - else: - if query_layer is None: - raise ValueError( - "kv_layer packs only K and V: query_layer is required when kv_layer" - " is provided." - ) - if key_layer is not None or value_layer is not None: - raise ValueError( - "kv_layer already packs K and V: key_layer and value_layer must be" - " None when kv_layer is provided." - ) - if kv_layer.dim() != query_layer.dim() + 1: - raise ValueError( - f"kv_layer must have one more dimension than query_layer, got" - f" {kv_layer.dim()}D kv_layer and {query_layer.dim()}D query_layer." - ) - if kv_layer.shape[qkv_interleave_dim] != 2: - raise ValueError( - f"kv_layer must have size 2 at dim {qkv_interleave_dim}" - f" (qkv_interleave_dim), got shape {tuple(kv_layer.shape)}." - ) - key_layer, value_layer = ( - kv_layer.select(qkv_interleave_dim, i) for i in range(2) - ) - packed_tensor = kv_layer - declared_qkv_layout = f"{packed_format}_{_packed_layout(packed_format, 2)}" - elif query_layer is None or key_layer is None or value_layer is None: - raise ValueError( - "query_layer, key_layer and value_layer are required unless packed" - " inputs (qkv_layer or query_layer + kv_layer) are provided." + query_layer, key_layer, value_layer, packed_tensor, declared_qkv_layout = ( + _unpack_packed_qkv( + qkv_layer, + kv_layer, + query_layer, + key_layer, + value_layer, + qkv_format if qkv_format is not None else self.qkv_format, + qkv_interleave_dim, + inference_params, ) + ) with self.prepare_forward_ctx( query_layer, From 2255bcc8aa0ea17e4a5e412c57b6ff140eec4589 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Thu, 9 Jul 2026 15:06:23 +0200 Subject: [PATCH 04/15] [PyTorch] DotProductAttention: validate packed-input last-dim stride Signed-off-by: Pawel Gadzinski --- .../attention/test_dpa_packed_inputs.py | 19 ++++++++++++++++--- .../dot_product_attention.py | 9 +++++++++ 2 files changed, 25 insertions(+), 3 deletions(-) diff --git a/tests/pytorch/attention/test_dpa_packed_inputs.py b/tests/pytorch/attention/test_dpa_packed_inputs.py index b87e20a344..7de1d938eb 100644 --- a/tests/pytorch/attention/test_dpa_packed_inputs.py +++ b/tests/pytorch/attention/test_dpa_packed_inputs.py @@ -43,10 +43,21 @@ def _fused_backend_supported(): try: q = torch.randn(_B, _S, _H, _D, dtype=_DTYPE, device="cuda") fused_attn_fwd( - True, _S, _S, _cu_seqlens(), _cu_seqlens(), q, q.clone(), q.clone(), _DTYPE, + True, + _S, + _S, + _cu_seqlens(), + _cu_seqlens(), + q, + q.clone(), + q.clone(), + _DTYPE, tex.NVTE_Fused_Attn_Backend.NVTE_F16_arbitrary_seqlen, - dropout=0.0, qkv_layout="bshd_bshd_bshd", o_format="bshd", - attn_bias_type="no_bias", attn_mask_type="no_mask", + dropout=0.0, + qkv_layout="bshd_bshd_bshd", + o_format="bshd", + attn_bias_type="no_bias", + attn_mask_type="no_mask", ) return True except Exception: @@ -215,6 +226,8 @@ def test_dpa_packed_input_validation(): dpa(qkv_layer=qkv, kv_layer=kv) with pytest.raises(ValueError, match="must have size 3 at dim"): dpa(qkv_layer=kv) # 2 at the interleave dim, not 3 + with pytest.raises(ValueError, match="stride 1 in its last"): + dpa(qkv_layer=qkv.transpose(-2, -1)) # declared layout would lie about memory with pytest.raises(ValueError, match="required unless packed"): dpa() 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 8feca1b630..88e496fa28 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 @@ -238,6 +238,15 @@ def _unpack_packed_qkv( "qkv_interleave_dim must be -3 (e.g. bs3hd) or -2 (e.g. bsh3d), got" f" {qkv_interleave_dim}." ) + packed = qkv_layer if qkv_layer is not None else kv_layer + # The declared layout describes the packed buffer's memory, so it must have + # stride 1 in its last dimension (the check get_qkv_layout would otherwise + # perform on the derived q/k/v views). + if packed.stride(-1) != 1: + raise ValueError( + "The packed tensor (qkv_layer/kv_layer) must have stride 1 in its last" + f" dimension, got strides {tuple(packed.stride())}." + ) def _packed_layout(fmt: str, num: int) -> str: # bshd + 3 @ -3 -> bs3hd; bshd + 3 @ -2 -> bsh3d; thd + 2 @ -2 -> th2d From dc4469bf88f30f6afc2316e512dcb1954a99fc59 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Thu, 9 Jul 2026 15:28:06 +0200 Subject: [PATCH 05/15] [PyTorch] DotProductAttention: pass packed qkv/kv buffers as separate arguments Rename the single layout-dependent packed_qkv plumbing argument (which held the full QKV buffer for *3* layouts but the KV buffer for *_2* layouts) into explicit packed_qkv/packed_kv, mirroring the public qkv_layer/kv_layer API. combine_and_quantize's combined= is split into combined_qkv=/combined_kv= accordingly, and _unpack_packed_qkv no longer returns the packed tensor since the callers already hold qkv_layer/kv_layer. Signed-off-by: Pawel Gadzinski --- .../attention/test_dpa_packed_inputs.py | 13 ++++--- .../dot_product_attention/backends.py | 7 +++- .../dot_product_attention.py | 38 +++++++++---------- .../attention/dot_product_attention/utils.py | 36 +++++++++--------- 4 files changed, 49 insertions(+), 45 deletions(-) diff --git a/tests/pytorch/attention/test_dpa_packed_inputs.py b/tests/pytorch/attention/test_dpa_packed_inputs.py index 7de1d938eb..9fc71cee18 100644 --- a/tests/pytorch/attention/test_dpa_packed_inputs.py +++ b/tests/pytorch/attention/test_dpa_packed_inputs.py @@ -268,7 +268,8 @@ def fn(x): # --------------------------------------------------------------------------- -# 5. FP8 combine refactor: combined= path is bit-identical to combine_tensors path +# 5. FP8 combine refactor: combined_qkv/combined_kv path is bit-identical to +# the combine_tensors path # --------------------------------------------------------------------------- @@ -281,15 +282,15 @@ def _fp8_quantizer(): def test_combine_and_quantize_combined_matches_views(): - """Quantizing the caller's packed buffer directly (combined=) produces the same - _data bits and scale_inv as rebuilding the packed buffer from q/k/v views via - combine_tensors (the old set_-based path).""" + """Quantizing the caller's packed buffer directly (combined_qkv=) produces the + same _data bits and scale_inv as rebuilding the packed buffer from q/k/v views + via combine_tensors (the old set_-based path).""" torch.manual_seed(0) qkv = torch.randn(_B, _S, 3, _H, _D, dtype=_DTYPE, device="cuda") q, k, v = qkv[:, :, 0], qkv[:, :, 1], qkv[:, :, 2] old = combine_and_quantize("bs3hd", q, k, v, _fp8_quantizer()) - new = combine_and_quantize("bs3hd", q, k, v, _fp8_quantizer(), combined=qkv) + new = combine_and_quantize("bs3hd", q, k, v, _fp8_quantizer(), combined_qkv=qkv) assert old[3] == new[3] == "bs3hd" for name, x, y in zip(("q", "k", "v"), old[:3], new[:3]): @@ -305,7 +306,7 @@ def test_combine_and_quantize_combined_kv_matches_views(): k, v = kv[:, :, 0], kv[:, :, 1] old = combine_and_quantize("bshd_bs2hd", q, k, v, _fp8_quantizer()) - new = combine_and_quantize("bshd_bs2hd", q, k, v, _fp8_quantizer(), combined=kv) + new = combine_and_quantize("bshd_bs2hd", q, k, v, _fp8_quantizer(), combined_kv=kv) for name, x, y in zip(("q", "k", "v"), old[:3], new[:3]): assert torch.equal(x._data, y._data), f"{name} fp8 bits differ" diff --git a/transformer_engine/pytorch/attention/dot_product_attention/backends.py b/transformer_engine/pytorch/attention/dot_product_attention/backends.py index 50565bebf3..79fe695db2 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/backends.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/backends.py @@ -1329,6 +1329,7 @@ def forward( layer_number, return_max_logit, packed_qkv=None, + packed_kv=None, ): # pylint: disable=missing-function-docstring @@ -1397,7 +1398,8 @@ def forward( v, QKV_quantizer, used_in_backward=is_training, - combined=packed_qkv, + combined_qkv=packed_qkv, + combined_kv=packed_kv, ) # print quantizers @@ -1905,6 +1907,7 @@ def backward(ctx, d_out, *_args): None, None, None, # packed_qkv + None, # packed_kv ) @@ -2000,6 +2003,7 @@ def forward( score_mod_tensors: Optional[Dict[str, torch.Tensor]] = None, score_mod_bprop_tensors: Optional[Dict[str, torch.Tensor]] = None, packed_qkv: Optional[torch.Tensor] = None, + packed_kv: Optional[torch.Tensor] = None, ) -> torch.Tensor: """fused attention fprop""" assert ( @@ -2220,6 +2224,7 @@ def forward( self.layer_number, self.return_max_logit, packed_qkv, + packed_kv, ) if self.return_max_logit: 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 88e496fa28..990239e388 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 @@ -206,7 +206,7 @@ def _unpack_packed_qkv( qkv_format: str, qkv_interleave_dim: int, inference_params: Optional[InferenceParams], -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, Optional[torch.Tensor], Optional[str]]: +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, Optional[str]]: """Resolve declarative packed inputs into q/k/v. Derives q/k/v as zero-copy views of the packed buffer (``qkv_layer`` or @@ -214,9 +214,8 @@ def _unpack_packed_qkv( The layout enum is truthful by construction, so no pointer-based detection is needed downstream (this also covers thd and FP8 DPA). - Returns ``(query_layer, key_layer, value_layer, packed_tensor, - declared_qkv_layout)``; the last two are ``None`` when no packed input is - given. + Returns ``(query_layer, key_layer, value_layer, declared_qkv_layout)``; + the layout is ``None`` when no packed input is given. """ if qkv_layer is None and kv_layer is None: if query_layer is None or key_layer is None or value_layer is None: @@ -224,7 +223,7 @@ def _unpack_packed_qkv( "query_layer, key_layer and value_layer are required unless packed" " inputs (qkv_layer or query_layer + kv_layer) are provided." ) - return query_layer, key_layer, value_layer, None, None + return query_layer, key_layer, value_layer, None if qkv_layer is not None and kv_layer is not None: raise ValueError("qkv_layer and kv_layer are mutually exclusive.") @@ -273,7 +272,7 @@ def _packed_layout(fmt: str, num: int) -> str: query_layer, key_layer, value_layer = ( qkv_layer.select(qkv_interleave_dim, i) for i in range(3) ) - return query_layer, key_layer, value_layer, qkv_layer, _packed_layout(qkv_format, 3) + return query_layer, key_layer, value_layer, _packed_layout(qkv_format, 3) if query_layer is None: raise ValueError( @@ -299,7 +298,6 @@ def _packed_layout(fmt: str, num: int) -> str: query_layer, key_layer, value_layer, - kv_layer, f"{qkv_format}_{_packed_layout(qkv_format, 2)}", ) @@ -1382,17 +1380,15 @@ def forward( since e.g. ``h == 3`` would make the shapes ambiguous. """ - query_layer, key_layer, value_layer, packed_tensor, declared_qkv_layout = ( - _unpack_packed_qkv( - qkv_layer, - kv_layer, - query_layer, - key_layer, - value_layer, - qkv_format if qkv_format is not None else self.qkv_format, - qkv_interleave_dim, - inference_params, - ) + query_layer, key_layer, value_layer, declared_qkv_layout = _unpack_packed_qkv( + qkv_layer, + kv_layer, + query_layer, + key_layer, + value_layer, + qkv_format if qkv_format is not None else self.qkv_format, + qkv_interleave_dim, + inference_params, ) with self.prepare_forward_ctx( @@ -1968,7 +1964,8 @@ def forward( inference_params=inference_params, softmax_offset=softmax_offset, fp8_output=fp8_output, - packed_qkv=packed_tensor, + packed_qkv=qkv_layer, + packed_kv=kv_layer, ) return self.fused_attention( query_layer, @@ -2004,7 +2001,8 @@ def forward( score_mod_bprop=score_mod_bprop, score_mod_tensors=score_mod_tensors, score_mod_bprop_tensors=score_mod_bprop_tensors, - packed_qkv=packed_tensor, + packed_qkv=qkv_layer, + packed_kv=kv_layer, ) if use_unfused_attention: diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py index 3176200ffa..e1643283d2 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py @@ -2783,17 +2783,17 @@ def combine_and_quantize( used_in_forward=True, used_in_backward=False, keep_same_data_and_scale_inv_format=False, - combined: Optional[torch.Tensor] = None, + combined_qkv: Optional[torch.Tensor] = None, + combined_kv: Optional[torch.Tensor] = None, ): """Combine Q, K, V tensors based on qkv_layout and quantize them together. - When ``combined`` is provided, it must be the caller's original packed buffer - matching the packed group in ``qkv_layout`` (the full QKV tensor for ``3`` - layouts such as ``bs3hd``, or the KV tensor for ``2`` layouts such as - ``bshd_bs2hd``). It is then quantized directly instead of re-deriving the - packed buffer from the q/k/v views via ``combine_tensors`` (which rebuilds it - with a raw ``set_`` under a silent adjacency/interleave assumption). Ignored - for MXFP8 quantization. + When ``combined_qkv`` (for ``3`` layouts such as ``bs3hd``) or ``combined_kv`` + (for ``2`` layouts such as ``bshd_bs2hd``) is provided, it must be the + caller's original packed buffer that q/k/v are views of. It is then quantized + directly instead of re-deriving the packed buffer from the q/k/v views via + ``combine_tensors`` (which rebuilds it with a raw ``set_`` under a silent + adjacency/interleave assumption). Ignored for MXFP8 quantization. """ if isinstance(qkv_quantizer, MXFP8Quantizer): qkv_format, q_format, kv_format = get_qkv_format(qkv_layout) @@ -2894,24 +2894,24 @@ def combine_and_quantize( match qkv_group: case 1: dim = qkv_layout.find("3") - if combined is not None: - assert combined.shape[dim] == 3, ( - f"combined QKV tensor does not match qkv_layout {qkv_layout}: expected" - f" size 3 at dim {dim}, got shape {tuple(combined.shape)}." + if combined_qkv is not None: + assert combined_qkv.shape[dim] == 3, ( + f"combined_qkv does not match qkv_layout {qkv_layout}: expected" + f" size 3 at dim {dim}, got shape {tuple(combined_qkv.shape)}." ) - qkv = combined + qkv = combined_qkv else: qkv = combine_tensors([q, k, v], dim) qkv_fp8 = qkv_quantizer(qkv) q_data, k_data, v_data = SplitAlongDim.apply(qkv_fp8._data, dim, [1, 1, 1], True) case 2: dim = qkv_layout.split("_")[1].find("2") - if combined is not None: - assert combined.shape[dim] == 2, ( - f"combined KV tensor does not match qkv_layout {qkv_layout}: expected" - f" size 2 at dim {dim}, got shape {tuple(combined.shape)}." + if combined_kv is not None: + assert combined_kv.shape[dim] == 2, ( + f"combined_kv does not match qkv_layout {qkv_layout}: expected" + f" size 2 at dim {dim}, got shape {tuple(combined_kv.shape)}." ) - kv = combined + kv = combined_kv else: kv = combine_tensors([k, v], dim) tensors = [q, kv] From f82ba7033712e300e22048be00f10e3c0cfff1de Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Thu, 9 Jul 2026 15:37:59 +0200 Subject: [PATCH 06/15] [PyTorch] Move packed qkv/kv input tests into test_attention.py Fold the tests from the new test_dpa_packed_inputs.py file into the existing attention test suite as a dedicated section, reusing its imports. No new test file; test logic unchanged apart from renaming the module-level constants (_B/_S/_H/_D/_DTYPE -> _PACKED_*) to avoid collisions. Signed-off-by: Pawel Gadzinski --- tests/pytorch/attention/test_attention.py | 580 ++++++++++++++++++ .../attention/test_dpa_packed_inputs.py | 568 ----------------- 2 files changed, 580 insertions(+), 568 deletions(-) delete mode 100644 tests/pytorch/attention/test_dpa_packed_inputs.py diff --git a/tests/pytorch/attention/test_attention.py b/tests/pytorch/attention/test_attention.py index 2dbf94fc20..aa70f52565 100644 --- a/tests/pytorch/attention/test_attention.py +++ b/tests/pytorch/attention/test_attention.py @@ -26,10 +26,13 @@ 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 ( FlashAttentionUtils, check_set_window_size, + combine_and_quantize, ) +from transformer_engine.pytorch.tensor.float8_tensor import Float8Quantizer from transformer_engine.pytorch.attention import RotaryPositionEmbedding import transformer_engine.pytorch.cpp_extensions as ext from transformer_engine.pytorch.cpp_extensions.fused_attn import ( @@ -3002,3 +3005,580 @@ def forward( self.quantizers, ) return out + + +# --------------------------------------------------------------------------- +# Declarative packed QKV/KV inputs (qkv_layer/kv_layer + qkv_interleave_dim) +# +# Instead of slicing a fused-projection buffer into q/k/v views (which TE then +# reverse-engineers via pointer-based layout detection), callers can pass the +# packed tensor directly. Q/K/V are derived as zero-copy views and the exact +# layout string (e.g. bs3hd) is declared, not detected -- including for thd +# and FP8 DPA. +# --------------------------------------------------------------------------- + +_PACKED_B, _PACKED_S, _PACKED_H, _PACKED_D = 2, 128, 8, 64 +_PACKED_DTYPE = torch.bfloat16 + + +def _cu_seqlens(): + return torch.arange(0, (_PACKED_B + 1) * _PACKED_S, _PACKED_S, dtype=torch.int32, device="cuda") + + +def _fused_backend_supported(): + try: + q = torch.randn( + _PACKED_B, _PACKED_S, _PACKED_H, _PACKED_D, dtype=_PACKED_DTYPE, device="cuda" + ) + fused_attn_fwd( + True, + _PACKED_S, + _PACKED_S, + _cu_seqlens(), + _cu_seqlens(), + q, + q.clone(), + q.clone(), + _PACKED_DTYPE, + tex.NVTE_Fused_Attn_Backend.NVTE_F16_arbitrary_seqlen, + dropout=0.0, + qkv_layout="bshd_bshd_bshd", + o_format="bshd", + attn_bias_type="no_bias", + attn_mask_type="no_mask", + ) + return True + except Exception: + return False + + +requires_fused = pytest.mark.skipif( + not (torch.cuda.is_available() and _fused_backend_supported()), + reason="F16_arbitrary_seqlen fused attention backend is not supported on this device", +) + + +def _force_backend(monkeypatch, backend): + """Force a single attention backend via env and invalidate the selection cache.""" + flash, fused = {"flash": ("1", "0"), "fused": ("0", "1")}[backend] + monkeypatch.setenv("NVTE_FLASH_ATTN", flash) + monkeypatch.setenv("NVTE_FUSED_ATTN", fused) + monkeypatch.setenv("NVTE_UNFUSED_ATTN", "0") + if backend == "flash": + # flash-attn bwd uses atomics unless deterministic + monkeypatch.setenv("NVTE_ALLOW_NONDETERMINISTIC_ALGO", "0") + _attention_backends["backend_selection_requires_update"] = True + + +def _make_dpa(qkv_format, num_gqa_groups=None): + return DotProductAttention( + _PACKED_H, + _PACKED_D, + num_gqa_groups=num_gqa_groups, + attention_dropout=0.0, + qkv_format=qkv_format, + attn_mask_type="no_mask", + ) + + +def _assert_bit_exact(result, reference, names=("out", "dq", "dk", "dv")): + for name, x, y in zip(names, result, reference): + assert torch.equal(x.contiguous(), y.contiguous()), f"{name} differs" + + +def _dpa_separate_baseline(qkv_format, num_gqa_groups=None): + """Fwd+bwd on contiguous separate q/k/v leaves; returns (out, dq, dk, dv).""" + hg = num_gqa_groups or _PACKED_H + torch.manual_seed(0) + q_shape = ( + (_PACKED_B, _PACKED_S, _PACKED_H, _PACKED_D) + if qkv_format == "bshd" + else (_PACKED_S, _PACKED_B, _PACKED_H, _PACKED_D) + ) + kv_shape = ( + (_PACKED_B, _PACKED_S, hg, _PACKED_D) + if qkv_format == "bshd" + else (_PACKED_S, _PACKED_B, hg, _PACKED_D) + ) + q = torch.randn(*q_shape, dtype=_PACKED_DTYPE, device="cuda") + k = torch.randn(*kv_shape, dtype=_PACKED_DTYPE, device="cuda") + v = torch.randn(*kv_shape, dtype=_PACKED_DTYPE, device="cuda") + q, k, v = [x.clone().requires_grad_() for x in (q, k, v)] + _attention_backends["backend_selection_requires_update"] = True + out = _make_dpa(qkv_format, num_gqa_groups)(q, k, v) + out.backward(torch.ones_like(out)) + return out, q.grad, k.grad, v.grad + + +# --------------------------------------------------------------------------- +# 1. Dense eager equivalence (fused backend), fwd + input grads, bit-exact +# --------------------------------------------------------------------------- + + +@requires_fused +@pytest.mark.parametrize( + "qkv_format, interleave_dim", + [ + pytest.param("bshd", -3, id="bshd_qkv_dim-3"), # (a) bs3hd + pytest.param("bshd", -2, id="bshd_qkv_dim-2"), # (b) bsh3d (Megatron-style) + pytest.param("sbhd", -3, id="sbhd_qkv_dim-3"), # (c) sb3hd + ], +) +def test_dpa_fused_qkv_layer_dense(monkeypatch, qkv_format, interleave_dim): + """qkv_layer packed input is bit-exact vs separate contiguous q/k/v, and grads + flow back into the packed tensor itself.""" + _force_backend(monkeypatch, "fused") + reference = _dpa_separate_baseline(qkv_format) + + torch.manual_seed(0) + q_shape = ( + (_PACKED_B, _PACKED_S, _PACKED_H, _PACKED_D) + if qkv_format == "bshd" + else (_PACKED_S, _PACKED_B, _PACKED_H, _PACKED_D) + ) + parts = [torch.randn(*q_shape, dtype=_PACKED_DTYPE, device="cuda") for _ in range(3)] + stack_dim = len(q_shape) + interleave_dim + 1 # -3 -> before h, -2 -> before d + qkv = torch.stack(parts, dim=stack_dim).requires_grad_() + + _attention_backends["backend_selection_requires_update"] = True + out = _make_dpa(qkv_format)(qkv_layer=qkv, qkv_interleave_dim=interleave_dim) + out.backward(torch.ones_like(out)) + + assert qkv.grad is not None and qkv.grad.shape == qkv.shape + grads = [qkv.grad.select(stack_dim, i) for i in range(3)] + _assert_bit_exact((out, *grads), reference) + + +@requires_fused +@pytest.mark.parametrize( + "num_gqa_groups", + [pytest.param(None, id="mha_kv"), pytest.param(2, id="gqa_kv")], # (d) and (e) +) +def test_dpa_fused_kv_layer_dense(monkeypatch, num_gqa_groups): + """kv_layer packed input (with separate query) is bit-exact vs separate + contiguous q/k/v; grads flow into the packed kv tensor.""" + _force_backend(monkeypatch, "fused") + reference = _dpa_separate_baseline("bshd", num_gqa_groups) + + hg = num_gqa_groups or _PACKED_H + torch.manual_seed(0) + q = torch.randn(_PACKED_B, _PACKED_S, _PACKED_H, _PACKED_D, dtype=_PACKED_DTYPE, device="cuda") + k = torch.randn(_PACKED_B, _PACKED_S, hg, _PACKED_D, dtype=_PACKED_DTYPE, device="cuda") + v = torch.randn(_PACKED_B, _PACKED_S, hg, _PACKED_D, dtype=_PACKED_DTYPE, device="cuda") + q = q.clone().requires_grad_() + kv = torch.stack([k, v], dim=2).requires_grad_() # [b,s,2,hg,d] + + _attention_backends["backend_selection_requires_update"] = True + out = _make_dpa("bshd", num_gqa_groups)(query_layer=q, kv_layer=kv) + out.backward(torch.ones_like(out)) + + assert kv.grad is not None and kv.grad.shape == kv.shape + _assert_bit_exact((out, q.grad, kv.grad[:, :, 0], kv.grad[:, :, 1]), reference) + + +# --------------------------------------------------------------------------- +# 2. Flash backend smoke +# --------------------------------------------------------------------------- + + +def test_dpa_flash_qkv_layer(monkeypatch): + """Flash backend: packed qkv_layer [b,s,3,h,d] is bit-exact vs separate.""" + _force_backend(monkeypatch, "flash") + try: + reference = _dpa_separate_baseline("bshd") + except Exception as exc: + pytest.skip(f"flash attention backend not available: {exc}") + + torch.manual_seed(0) + parts = [ + torch.randn(_PACKED_B, _PACKED_S, _PACKED_H, _PACKED_D, dtype=_PACKED_DTYPE, device="cuda") + for _ in range(3) + ] + qkv = torch.stack(parts, dim=2).requires_grad_() + _attention_backends["backend_selection_requires_update"] = True + out = _make_dpa("bshd")(qkv_layer=qkv) + out.backward(torch.ones_like(out)) + grads = [qkv.grad[:, :, i] for i in range(3)] + _assert_bit_exact((out, *grads), reference) + + +# --------------------------------------------------------------------------- +# 3. Validation errors +# --------------------------------------------------------------------------- + + +def test_dpa_packed_input_validation(): + dpa = _make_dpa("bshd") + qkv = torch.randn( + _PACKED_B, _PACKED_S, 3, _PACKED_H, _PACKED_D, dtype=_PACKED_DTYPE, device="cuda" + ) + kv = torch.randn( + _PACKED_B, _PACKED_S, 2, _PACKED_H, _PACKED_D, dtype=_PACKED_DTYPE, device="cuda" + ) + k = torch.randn(_PACKED_B, _PACKED_S, _PACKED_H, _PACKED_D, dtype=_PACKED_DTYPE, device="cuda") + + with pytest.raises(ValueError, match="must be None when qkv_layer is provided"): + dpa(qkv_layer=qkv, key_layer=k) + with pytest.raises(ValueError, match="query_layer is required when kv_layer"): + dpa(kv_layer=kv) + with pytest.raises(ValueError, match="qkv_interleave_dim must be -3"): + dpa(qkv_layer=qkv, qkv_interleave_dim=-1) + with pytest.raises(ValueError, match="mutually exclusive"): + dpa(qkv_layer=qkv, kv_layer=kv) + with pytest.raises(ValueError, match="must have size 3 at dim"): + dpa(qkv_layer=kv) # 2 at the interleave dim, not 3 + with pytest.raises(ValueError, match="stride 1 in its last"): + dpa(qkv_layer=qkv.transpose(-2, -1)) # declared layout would lie about memory + with pytest.raises(ValueError, match="required unless packed"): + dpa() + + +# --------------------------------------------------------------------------- +# 4. torch.compile: no data_ptr/UntypedStorage graph breaks with qkv_layer +# --------------------------------------------------------------------------- + + +@requires_fused +def test_dpa_torch_compile_qkv_layer_no_pointer_graph_breaks(monkeypatch): + _force_backend(monkeypatch, "fused") + torch._dynamo.reset() + torch._dynamo.utils.counters.clear() + + torch.manual_seed(0) + parts = [ + torch.randn(_PACKED_B, _PACKED_S, _PACKED_H, _PACKED_D, dtype=_PACKED_DTYPE, device="cuda") + for _ in range(3) + ] + qkv = torch.stack(parts, dim=2).requires_grad_() + dpa = _make_dpa("bshd") + + def fn(x): + return dpa(qkv_layer=x) + + _attention_backends["backend_selection_requires_update"] = True + eager_out = fn(qkv) # eager warm-up: backend selection happens outside dynamo + compiled_out = torch.compile(fn)(qkv) + compiled_out.backward(torch.ones_like(compiled_out)) + + breaks = dict(torch._dynamo.utils.counters["graph_break"]) + torch._dynamo.reset() + pointer_breaks = { + reason: count + for reason, count in breaks.items() + if "data_ptr" in reason or "UntypedStorage" in reason + } + assert not pointer_breaks, f"pointer-based graph breaks with qkv_layer: {pointer_breaks}" + assert torch.equal(compiled_out, eager_out), "compiled output differs from eager" + + +# --------------------------------------------------------------------------- +# 5. FP8 combine refactor: combined_qkv/combined_kv path is bit-identical to +# the combine_tensors path +# --------------------------------------------------------------------------- + + +def _fp8_quantizer(): + return Float8Quantizer( + scale=torch.ones(1, dtype=torch.float32, device="cuda"), + amax=torch.zeros(1, dtype=torch.float32, device="cuda"), + fp8_dtype=tex.DType.kFloat8E4M3, + ) + + +def test_combine_and_quantize_combined_matches_views(): + """Quantizing the caller's packed buffer directly (combined_qkv=) produces the + same _data bits and scale_inv as rebuilding the packed buffer from q/k/v views + via combine_tensors (the old set_-based path).""" + torch.manual_seed(0) + qkv = torch.randn( + _PACKED_B, _PACKED_S, 3, _PACKED_H, _PACKED_D, dtype=_PACKED_DTYPE, device="cuda" + ) + q, k, v = qkv[:, :, 0], qkv[:, :, 1], qkv[:, :, 2] + + old = combine_and_quantize("bs3hd", q, k, v, _fp8_quantizer()) + new = combine_and_quantize("bs3hd", q, k, v, _fp8_quantizer(), combined_qkv=qkv) + + assert old[3] == new[3] == "bs3hd" + for name, x, y in zip(("q", "k", "v"), old[:3], new[:3]): + assert torch.equal(x._data, y._data), f"{name} fp8 bits differ" + assert torch.equal(x._scale_inv, y._scale_inv), f"{name} scale_inv differs" + + +def test_combine_and_quantize_combined_kv_matches_views(): + """Same for the kv-packed (group 2) layout.""" + torch.manual_seed(0) + q = torch.randn(_PACKED_B, _PACKED_S, _PACKED_H, _PACKED_D, dtype=_PACKED_DTYPE, device="cuda") + kv = torch.randn( + _PACKED_B, _PACKED_S, 2, _PACKED_H, _PACKED_D, dtype=_PACKED_DTYPE, device="cuda" + ) + k, v = kv[:, :, 0], kv[:, :, 1] + + old = combine_and_quantize("bshd_bs2hd", q, k, v, _fp8_quantizer()) + new = combine_and_quantize("bshd_bs2hd", q, k, v, _fp8_quantizer(), combined_kv=kv) + + for name, x, y in zip(("q", "k", "v"), old[:3], new[:3]): + assert torch.equal(x._data, y._data), f"{name} fp8 bits differ" + assert torch.equal(x._scale_inv, y._scale_inv), f"{name} scale_inv differs" + + +# --------------------------------------------------------------------------- +# 6. thd declarative: t3hd is declared, get_qkv_layout is never called +# --------------------------------------------------------------------------- + + +def test_dpa_thd_qkv_layer_declared_no_detection(monkeypatch): + """Packed thd input (qkv_layer [t,3,h,d]) declares 't3hd' without calling + get_qkv_layout; full forward+backward runs if a thd backend is available.""" + calls = [] + orig_get_qkv_layout = dpa_utils.get_qkv_layout + + def counting(*args, **kwargs): + calls.append(kwargs.get("qkv_format")) + return orig_get_qkv_layout(*args, **kwargs) + + monkeypatch.setattr(dpa_utils, "get_qkv_layout", counting) + + seen_layouts = [] + orig_get_backend = dpa_utils.get_attention_backend + + def recording(params): + seen_layouts.append(params.qkv_layout) + return orig_get_backend(params) + + monkeypatch.setattr(dpa_utils, "get_attention_backend", recording) + + torch.manual_seed(0) + t = _PACKED_B * _PACKED_S + qkv = torch.randn( + t, 3, _PACKED_H, _PACKED_D, dtype=_PACKED_DTYPE, device="cuda", requires_grad=True + ) + cu = _cu_seqlens() + dpa = DotProductAttention( + _PACKED_H, _PACKED_D, attention_dropout=0.0, qkv_format="thd", attn_mask_type="padding" + ) + _attention_backends["backend_selection_requires_update"] = True + ran_full = True + try: + out = dpa( + qkv_layer=qkv, + cu_seqlens_q=cu, + cu_seqlens_kv=cu, + max_seqlen_q=_PACKED_S, + max_seqlen_kv=_PACKED_S, + ) + out.backward(torch.ones_like(out)) + assert qkv.grad is not None and qkv.grad.shape == qkv.shape + except ValueError: + # No thd-capable backend on this device; the layout step (the subject of + # this test) runs before backend dispatch, so the assertions still hold. + ran_full = False + + assert not calls, "get_qkv_layout must not be called for declared packed thd input" + assert seen_layouts == ["t3hd"], f"expected declared layout 't3hd', got {seen_layouts}" + if not ran_full: + pytest.skip("layout declaration verified; no thd backend available for full run") + + +# --------------------------------------------------------------------------- +# 7. MultiheadAttention adoption: the fused QKV/KV projection output is passed +# packed (qkv_layer/kv_layer) to DotProductAttention +# --------------------------------------------------------------------------- + + +def _spy_packed_dpa(monkeypatch, record): + """Record (qkv_layer given, kv_layer given, qkv_interleave_dim) per DPA call.""" + orig = DotProductAttention.forward + + def spy(self, *args, **kwargs): + record.append( + ( + kwargs.get("qkv_layer") is not None, + kwargs.get("kv_layer") is not None, + kwargs.get("qkv_interleave_dim", None), + ) + ) + return orig(self, *args, **kwargs) + + monkeypatch.setattr(DotProductAttention, "forward", spy) + + +def _strip_packed_dpa(monkeypatch): + """Reference path: convert packed DPA inputs back to separate contiguous q/k/v.""" + orig = DotProductAttention.forward + + def stripped(self, query_layer=None, key_layer=None, value_layer=None, *args, **kwargs): + qkv = kwargs.pop("qkv_layer", None) + kv = kwargs.pop("kv_layer", None) + dim = kwargs.pop("qkv_interleave_dim", -3) + if qkv is not None: + query_layer, key_layer, value_layer = ( + qkv.select(dim, i).contiguous() for i in range(3) + ) + elif kv is not None: + key_layer, value_layer = (kv.select(dim, i).contiguous() for i in range(2)) + return orig(self, query_layer, key_layer, value_layer, *args, **kwargs) + + monkeypatch.setattr(DotProductAttention, "forward", stripped) + + +def _run_mha(mha, x, encoder_output=None): + _attention_backends["backend_selection_requires_update"] = True + if encoder_output is not None: + out = mha(x, encoder_output=encoder_output) + else: + out = mha(x) + out.backward(torch.ones_like(out)) + wgrads = [p.grad.clone() for p in mha.parameters() if p.grad is not None] + xgrad = x.grad.clone() + x.grad = None + mha.zero_grad(set_to_none=True) + return out, xgrad, wgrads + + +def _assert_mha_equal(result, reference): + out, xgrad, wgrads = result + out_ref, xgrad_ref, wgrads_ref = reference + assert torch.equal(out, out_ref), "output differs" + assert torch.equal(xgrad, xgrad_ref), "input grad differs" + assert len(wgrads) == len(wgrads_ref) + for i, (w, w_ref) in enumerate(zip(wgrads, wgrads_ref)): + assert torch.equal(w, w_ref), f"weight grad {i} differs" + + +@requires_fused +@pytest.mark.parametrize( + "interleaved", [pytest.param(True, id="interleaved"), pytest.param(False, id="non_interleaved")] +) +def test_mha_self_attention_packed_pass_through(monkeypatch, interleaved): + """MHA self-attention passes its packed projection output straight to DPA as + qkv_layer (with the matching interleave dim), bit-exact vs the same MHA with + packed inputs converted back to separate contiguous q/k/v.""" + _force_backend(monkeypatch, "fused") + hidden = _PACKED_H * _PACKED_D + torch.manual_seed(0) + mha = MultiheadAttention( + hidden, + _PACKED_H, + attention_dropout=0.0, + attn_mask_type="no_mask", + qkv_format="sbhd", + fuse_qkv_params=True, + qkv_weight_interleaved=interleaved, + params_dtype=_PACKED_DTYPE, + device="cuda", + ) + torch.manual_seed(1) + x = torch.randn( + _PACKED_S, _PACKED_B, hidden, dtype=_PACKED_DTYPE, device="cuda", requires_grad=True + ) + + record = [] + _spy_packed_dpa(monkeypatch, record) + result = _run_mha(mha, x) + assert record == [(True, False, -2 if interleaved else -3)], f"unexpected DPA call: {record}" + monkeypatch.undo() + + _force_backend(monkeypatch, "fused") + _strip_packed_dpa(monkeypatch) + reference = _run_mha(mha, x) + _assert_mha_equal(result, reference) + + +@requires_fused +@pytest.mark.parametrize( + "interleaved", [pytest.param(True, id="interleaved"), pytest.param(False, id="non_interleaved")] +) +def test_mha_cross_attention_packed_kv_pass_through(monkeypatch, interleaved): + """MHA cross-attention passes its packed KV projection output to DPA as + kv_layer, bit-exact vs the separate contiguous reference.""" + _force_backend(monkeypatch, "fused") + hidden = _PACKED_H * _PACKED_D + torch.manual_seed(0) + mha = MultiheadAttention( + hidden, + _PACKED_H, + attention_dropout=0.0, + attn_mask_type="no_mask", + qkv_format="sbhd", + attention_type="cross", + fuse_qkv_params=True, + qkv_weight_interleaved=interleaved, + params_dtype=_PACKED_DTYPE, + device="cuda", + ) + torch.manual_seed(1) + x = torch.randn( + _PACKED_S, _PACKED_B, hidden, dtype=_PACKED_DTYPE, device="cuda", requires_grad=True + ) + enc = torch.randn(_PACKED_S, _PACKED_B, hidden, dtype=_PACKED_DTYPE, device="cuda") + + record = [] + _spy_packed_dpa(monkeypatch, record) + result = _run_mha(mha, x, encoder_output=enc) + assert record == [(False, True, -2 if interleaved else -3)], f"unexpected DPA call: {record}" + monkeypatch.undo() + + _force_backend(monkeypatch, "fused") + _strip_packed_dpa(monkeypatch) + reference = _run_mha(mha, x, encoder_output=enc) + _assert_mha_equal(result, reference) + + +@requires_fused +def test_mha_gqa_falls_back_to_views(monkeypatch): + """GQA (np != ng) is not a uniform 3-interleave: MHA must keep the legacy + sliced-views path and still work.""" + _force_backend(monkeypatch, "fused") + hidden = _PACKED_H * _PACKED_D + torch.manual_seed(0) + mha = MultiheadAttention( + hidden, + _PACKED_H, + num_gqa_groups=2, + attention_dropout=0.0, + attn_mask_type="no_mask", + qkv_format="sbhd", + fuse_qkv_params=True, + params_dtype=_PACKED_DTYPE, + device="cuda", + ) + torch.manual_seed(1) + x = torch.randn( + _PACKED_S, _PACKED_B, hidden, dtype=_PACKED_DTYPE, device="cuda", requires_grad=True + ) + + record = [] + _spy_packed_dpa(monkeypatch, record) + out, _, _ = _run_mha(mha, x) + assert record == [(False, False, -3)], f"GQA must not use the packed path: {record}" + assert out.shape == (_PACKED_S, _PACKED_B, hidden) + + +@requires_fused +def test_mha_rope_falls_back_to_views(monkeypatch): + """RoPE needs the individual q/k slices: MHA must keep the legacy path.""" + _force_backend(monkeypatch, "fused") + hidden = _PACKED_H * _PACKED_D + torch.manual_seed(0) + mha = MultiheadAttention( + hidden, + _PACKED_H, + attention_dropout=0.0, + attn_mask_type="no_mask", + qkv_format="sbhd", + fuse_qkv_params=True, + params_dtype=_PACKED_DTYPE, + device="cuda", + ) + rope = RotaryPositionEmbedding(_PACKED_D)(max_seq_len=_PACKED_S).to("cuda") + torch.manual_seed(1) + x = torch.randn( + _PACKED_S, _PACKED_B, hidden, dtype=_PACKED_DTYPE, device="cuda", requires_grad=True + ) + + record = [] + _spy_packed_dpa(monkeypatch, record) + _attention_backends["backend_selection_requires_update"] = True + out = mha(x, rotary_pos_emb=rope) + assert record == [(False, False, -3)], f"RoPE must not use the packed path: {record}" + assert out.shape == (_PACKED_S, _PACKED_B, hidden) diff --git a/tests/pytorch/attention/test_dpa_packed_inputs.py b/tests/pytorch/attention/test_dpa_packed_inputs.py deleted file mode 100644 index 9fc71cee18..0000000000 --- a/tests/pytorch/attention/test_dpa_packed_inputs.py +++ /dev/null @@ -1,568 +0,0 @@ -# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# -# See LICENSE for license information. - -"""Tests for declarative packed QKV/KV inputs to DotProductAttention. - -Instead of slicing a fused-projection buffer into q/k/v views (which TE then -reverse-engineers via pointer-based layout detection), callers can pass the -packed tensor directly (``qkv_layer``/``kv_layer`` + ``qkv_interleave_dim``). -Q/K/V are derived as zero-copy views and the exact layout string (e.g. -``bs3hd``) is declared, not detected -- including for thd and FP8 DPA. -""" - -import pytest -import torch - -import transformer_engine.pytorch # noqa: F401 (loads libtransformer_engine.so) -import transformer_engine_torch as tex -from transformer_engine.pytorch import DotProductAttention, MultiheadAttention -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 -from transformer_engine.pytorch.attention.dot_product_attention.utils import ( - combine_and_quantize, -) -from transformer_engine.pytorch.cpp_extensions.fused_attn import ( - fused_attn_fwd, -) -from transformer_engine.pytorch.tensor.float8_tensor import Float8Quantizer - -pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is not available") - -_B, _S, _H, _D = 2, 128, 8, 64 -_DTYPE = torch.bfloat16 - - -def _cu_seqlens(): - return torch.arange(0, (_B + 1) * _S, _S, dtype=torch.int32, device="cuda") - - -def _fused_backend_supported(): - try: - q = torch.randn(_B, _S, _H, _D, dtype=_DTYPE, device="cuda") - fused_attn_fwd( - True, - _S, - _S, - _cu_seqlens(), - _cu_seqlens(), - q, - q.clone(), - q.clone(), - _DTYPE, - tex.NVTE_Fused_Attn_Backend.NVTE_F16_arbitrary_seqlen, - dropout=0.0, - qkv_layout="bshd_bshd_bshd", - o_format="bshd", - attn_bias_type="no_bias", - attn_mask_type="no_mask", - ) - return True - except Exception: - return False - - -requires_fused = pytest.mark.skipif( - not (torch.cuda.is_available() and _fused_backend_supported()), - reason="F16_arbitrary_seqlen fused attention backend is not supported on this device", -) - - -def _force_backend(monkeypatch, backend): - """Force a single attention backend via env and invalidate the selection cache.""" - flash, fused = {"flash": ("1", "0"), "fused": ("0", "1")}[backend] - monkeypatch.setenv("NVTE_FLASH_ATTN", flash) - monkeypatch.setenv("NVTE_FUSED_ATTN", fused) - monkeypatch.setenv("NVTE_UNFUSED_ATTN", "0") - if backend == "flash": - # flash-attn bwd uses atomics unless deterministic - monkeypatch.setenv("NVTE_ALLOW_NONDETERMINISTIC_ALGO", "0") - dpa_module._attention_backends["backend_selection_requires_update"] = True - - -def _make_dpa(qkv_format, num_gqa_groups=None): - return DotProductAttention( - _H, - _D, - num_gqa_groups=num_gqa_groups, - attention_dropout=0.0, - qkv_format=qkv_format, - attn_mask_type="no_mask", - ) - - -def _assert_bit_exact(result, reference, names=("out", "dq", "dk", "dv")): - for name, x, y in zip(names, result, reference): - assert torch.equal(x.contiguous(), y.contiguous()), f"{name} differs" - - -def _fresh_parts(qkv_format, num_heads, seed=0): - torch.manual_seed(seed) - shape = (_B, _S, num_heads, _D) if qkv_format == "bshd" else (_S, _B, num_heads, _D) - return [torch.randn(*shape, dtype=_DTYPE, device="cuda") for _ in range(3)] - - -def _dpa_separate_baseline(qkv_format, num_gqa_groups=None): - """Fwd+bwd on contiguous separate q/k/v leaves; returns (out, dq, dk, dv).""" - hg = num_gqa_groups or _H - torch.manual_seed(0) - q_shape = (_B, _S, _H, _D) if qkv_format == "bshd" else (_S, _B, _H, _D) - kv_shape = (_B, _S, hg, _D) if qkv_format == "bshd" else (_S, _B, hg, _D) - q = torch.randn(*q_shape, dtype=_DTYPE, device="cuda") - k = torch.randn(*kv_shape, dtype=_DTYPE, device="cuda") - v = torch.randn(*kv_shape, dtype=_DTYPE, device="cuda") - q, k, v = [x.clone().requires_grad_() for x in (q, k, v)] - dpa_module._attention_backends["backend_selection_requires_update"] = True - out = _make_dpa(qkv_format, num_gqa_groups)(q, k, v) - out.backward(torch.ones_like(out)) - return out, q.grad, k.grad, v.grad - - -# --------------------------------------------------------------------------- -# 1. Dense eager equivalence (fused backend), fwd + input grads, bit-exact -# --------------------------------------------------------------------------- - - -@requires_fused -@pytest.mark.parametrize( - "qkv_format, interleave_dim", - [ - pytest.param("bshd", -3, id="bshd_qkv_dim-3"), # (a) bs3hd - pytest.param("bshd", -2, id="bshd_qkv_dim-2"), # (b) bsh3d (Megatron-style) - pytest.param("sbhd", -3, id="sbhd_qkv_dim-3"), # (c) sb3hd - ], -) -def test_dpa_fused_qkv_layer_dense(monkeypatch, qkv_format, interleave_dim): - """qkv_layer packed input is bit-exact vs separate contiguous q/k/v, and grads - flow back into the packed tensor itself.""" - _force_backend(monkeypatch, "fused") - reference = _dpa_separate_baseline(qkv_format) - - torch.manual_seed(0) - q_shape = (_B, _S, _H, _D) if qkv_format == "bshd" else (_S, _B, _H, _D) - parts = [torch.randn(*q_shape, dtype=_DTYPE, device="cuda") for _ in range(3)] - stack_dim = len(q_shape) + interleave_dim + 1 # -3 -> before h, -2 -> before d - qkv = torch.stack(parts, dim=stack_dim).requires_grad_() - - dpa_module._attention_backends["backend_selection_requires_update"] = True - out = _make_dpa(qkv_format)(qkv_layer=qkv, qkv_interleave_dim=interleave_dim) - out.backward(torch.ones_like(out)) - - assert qkv.grad is not None and qkv.grad.shape == qkv.shape - grads = [qkv.grad.select(stack_dim, i) for i in range(3)] - _assert_bit_exact((out, *grads), reference) - - -@requires_fused -@pytest.mark.parametrize( - "num_gqa_groups", - [pytest.param(None, id="mha_kv"), pytest.param(2, id="gqa_kv")], # (d) and (e) -) -def test_dpa_fused_kv_layer_dense(monkeypatch, num_gqa_groups): - """kv_layer packed input (with separate query) is bit-exact vs separate - contiguous q/k/v; grads flow into the packed kv tensor.""" - _force_backend(monkeypatch, "fused") - reference = _dpa_separate_baseline("bshd", num_gqa_groups) - - hg = num_gqa_groups or _H - torch.manual_seed(0) - q = torch.randn(_B, _S, _H, _D, dtype=_DTYPE, device="cuda") - k = torch.randn(_B, _S, hg, _D, dtype=_DTYPE, device="cuda") - v = torch.randn(_B, _S, hg, _D, dtype=_DTYPE, device="cuda") - q = q.clone().requires_grad_() - kv = torch.stack([k, v], dim=2).requires_grad_() # [b,s,2,hg,d] - - dpa_module._attention_backends["backend_selection_requires_update"] = True - out = _make_dpa("bshd", num_gqa_groups)(query_layer=q, kv_layer=kv) - out.backward(torch.ones_like(out)) - - assert kv.grad is not None and kv.grad.shape == kv.shape - _assert_bit_exact((out, q.grad, kv.grad[:, :, 0], kv.grad[:, :, 1]), reference) - - -# --------------------------------------------------------------------------- -# 2. Flash backend smoke -# --------------------------------------------------------------------------- - - -def test_dpa_flash_qkv_layer(monkeypatch): - """Flash backend: packed qkv_layer [b,s,3,h,d] is bit-exact vs separate.""" - _force_backend(monkeypatch, "flash") - try: - reference = _dpa_separate_baseline("bshd") - except Exception as exc: - pytest.skip(f"flash attention backend not available: {exc}") - - torch.manual_seed(0) - parts = [torch.randn(_B, _S, _H, _D, dtype=_DTYPE, device="cuda") for _ in range(3)] - qkv = torch.stack(parts, dim=2).requires_grad_() - dpa_module._attention_backends["backend_selection_requires_update"] = True - out = _make_dpa("bshd")(qkv_layer=qkv) - out.backward(torch.ones_like(out)) - grads = [qkv.grad[:, :, i] for i in range(3)] - _assert_bit_exact((out, *grads), reference) - - -# --------------------------------------------------------------------------- -# 3. Validation errors -# --------------------------------------------------------------------------- - - -def test_dpa_packed_input_validation(): - dpa = _make_dpa("bshd") - qkv = torch.randn(_B, _S, 3, _H, _D, dtype=_DTYPE, device="cuda") - kv = torch.randn(_B, _S, 2, _H, _D, dtype=_DTYPE, device="cuda") - k = torch.randn(_B, _S, _H, _D, dtype=_DTYPE, device="cuda") - - with pytest.raises(ValueError, match="must be None when qkv_layer is provided"): - dpa(qkv_layer=qkv, key_layer=k) - with pytest.raises(ValueError, match="query_layer is required when kv_layer"): - dpa(kv_layer=kv) - with pytest.raises(ValueError, match="qkv_interleave_dim must be -3"): - dpa(qkv_layer=qkv, qkv_interleave_dim=-1) - with pytest.raises(ValueError, match="mutually exclusive"): - dpa(qkv_layer=qkv, kv_layer=kv) - with pytest.raises(ValueError, match="must have size 3 at dim"): - dpa(qkv_layer=kv) # 2 at the interleave dim, not 3 - with pytest.raises(ValueError, match="stride 1 in its last"): - dpa(qkv_layer=qkv.transpose(-2, -1)) # declared layout would lie about memory - with pytest.raises(ValueError, match="required unless packed"): - dpa() - - -# --------------------------------------------------------------------------- -# 4. torch.compile: no data_ptr/UntypedStorage graph breaks with qkv_layer -# --------------------------------------------------------------------------- - - -@requires_fused -def test_dpa_torch_compile_qkv_layer_no_pointer_graph_breaks(monkeypatch): - _force_backend(monkeypatch, "fused") - torch._dynamo.reset() - torch._dynamo.utils.counters.clear() - - torch.manual_seed(0) - parts = [torch.randn(_B, _S, _H, _D, dtype=_DTYPE, device="cuda") for _ in range(3)] - qkv = torch.stack(parts, dim=2).requires_grad_() - dpa = _make_dpa("bshd") - - def fn(x): - return dpa(qkv_layer=x) - - dpa_module._attention_backends["backend_selection_requires_update"] = True - eager_out = fn(qkv) # eager warm-up: backend selection happens outside dynamo - compiled_out = torch.compile(fn)(qkv) - compiled_out.backward(torch.ones_like(compiled_out)) - - breaks = dict(torch._dynamo.utils.counters["graph_break"]) - torch._dynamo.reset() - pointer_breaks = { - reason: count - for reason, count in breaks.items() - if "data_ptr" in reason or "UntypedStorage" in reason - } - assert not pointer_breaks, f"pointer-based graph breaks with qkv_layer: {pointer_breaks}" - assert torch.equal(compiled_out, eager_out), "compiled output differs from eager" - - -# --------------------------------------------------------------------------- -# 5. FP8 combine refactor: combined_qkv/combined_kv path is bit-identical to -# the combine_tensors path -# --------------------------------------------------------------------------- - - -def _fp8_quantizer(): - return Float8Quantizer( - scale=torch.ones(1, dtype=torch.float32, device="cuda"), - amax=torch.zeros(1, dtype=torch.float32, device="cuda"), - fp8_dtype=tex.DType.kFloat8E4M3, - ) - - -def test_combine_and_quantize_combined_matches_views(): - """Quantizing the caller's packed buffer directly (combined_qkv=) produces the - same _data bits and scale_inv as rebuilding the packed buffer from q/k/v views - via combine_tensors (the old set_-based path).""" - torch.manual_seed(0) - qkv = torch.randn(_B, _S, 3, _H, _D, dtype=_DTYPE, device="cuda") - q, k, v = qkv[:, :, 0], qkv[:, :, 1], qkv[:, :, 2] - - old = combine_and_quantize("bs3hd", q, k, v, _fp8_quantizer()) - new = combine_and_quantize("bs3hd", q, k, v, _fp8_quantizer(), combined_qkv=qkv) - - assert old[3] == new[3] == "bs3hd" - for name, x, y in zip(("q", "k", "v"), old[:3], new[:3]): - assert torch.equal(x._data, y._data), f"{name} fp8 bits differ" - assert torch.equal(x._scale_inv, y._scale_inv), f"{name} scale_inv differs" - - -def test_combine_and_quantize_combined_kv_matches_views(): - """Same for the kv-packed (group 2) layout.""" - torch.manual_seed(0) - q = torch.randn(_B, _S, _H, _D, dtype=_DTYPE, device="cuda") - kv = torch.randn(_B, _S, 2, _H, _D, dtype=_DTYPE, device="cuda") - k, v = kv[:, :, 0], kv[:, :, 1] - - old = combine_and_quantize("bshd_bs2hd", q, k, v, _fp8_quantizer()) - new = combine_and_quantize("bshd_bs2hd", q, k, v, _fp8_quantizer(), combined_kv=kv) - - for name, x, y in zip(("q", "k", "v"), old[:3], new[:3]): - assert torch.equal(x._data, y._data), f"{name} fp8 bits differ" - assert torch.equal(x._scale_inv, y._scale_inv), f"{name} scale_inv differs" - - -# --------------------------------------------------------------------------- -# 6. thd declarative: t3hd is declared, get_qkv_layout is never called -# --------------------------------------------------------------------------- - - -def test_dpa_thd_qkv_layer_declared_no_detection(monkeypatch): - """Packed thd input (qkv_layer [t,3,h,d]) declares 't3hd' without calling - get_qkv_layout; full forward+backward runs if a thd backend is available.""" - calls = [] - orig_get_qkv_layout = dpa_utils.get_qkv_layout - - def counting(*args, **kwargs): - calls.append(kwargs.get("qkv_format")) - return orig_get_qkv_layout(*args, **kwargs) - - monkeypatch.setattr(dpa_utils, "get_qkv_layout", counting) - - seen_layouts = [] - orig_get_backend = dpa_utils.get_attention_backend - - def recording(params): - seen_layouts.append(params.qkv_layout) - return orig_get_backend(params) - - monkeypatch.setattr(dpa_utils, "get_attention_backend", recording) - - torch.manual_seed(0) - t = _B * _S - qkv = torch.randn(t, 3, _H, _D, dtype=_DTYPE, device="cuda", requires_grad=True) - cu = _cu_seqlens() - dpa = DotProductAttention( - _H, _D, attention_dropout=0.0, qkv_format="thd", attn_mask_type="padding" - ) - dpa_module._attention_backends["backend_selection_requires_update"] = True - ran_full = True - try: - out = dpa( - qkv_layer=qkv, - cu_seqlens_q=cu, - cu_seqlens_kv=cu, - max_seqlen_q=_S, - max_seqlen_kv=_S, - ) - out.backward(torch.ones_like(out)) - assert qkv.grad is not None and qkv.grad.shape == qkv.shape - except ValueError: - # No thd-capable backend on this device; the layout step (the subject of - # this test) runs before backend dispatch, so the assertions still hold. - ran_full = False - - assert not calls, "get_qkv_layout must not be called for declared packed thd input" - assert seen_layouts == ["t3hd"], f"expected declared layout 't3hd', got {seen_layouts}" - if not ran_full: - pytest.skip("layout declaration verified; no thd backend available for full run") - - -# --------------------------------------------------------------------------- -# 7. MultiheadAttention adoption: the fused QKV/KV projection output is passed -# packed (qkv_layer/kv_layer) to DotProductAttention -# --------------------------------------------------------------------------- - - -def _spy_packed_dpa(monkeypatch, record): - """Record (qkv_layer given, kv_layer given, qkv_interleave_dim) per DPA call.""" - orig = DotProductAttention.forward - - def spy(self, *args, **kwargs): - record.append( - ( - kwargs.get("qkv_layer") is not None, - kwargs.get("kv_layer") is not None, - kwargs.get("qkv_interleave_dim", None), - ) - ) - return orig(self, *args, **kwargs) - - monkeypatch.setattr(DotProductAttention, "forward", spy) - - -def _strip_packed_dpa(monkeypatch): - """Reference path: convert packed DPA inputs back to separate contiguous q/k/v.""" - orig = DotProductAttention.forward - - def stripped(self, query_layer=None, key_layer=None, value_layer=None, *args, **kwargs): - qkv = kwargs.pop("qkv_layer", None) - kv = kwargs.pop("kv_layer", None) - dim = kwargs.pop("qkv_interleave_dim", -3) - if qkv is not None: - query_layer, key_layer, value_layer = ( - qkv.select(dim, i).contiguous() for i in range(3) - ) - elif kv is not None: - key_layer, value_layer = (kv.select(dim, i).contiguous() for i in range(2)) - return orig(self, query_layer, key_layer, value_layer, *args, **kwargs) - - monkeypatch.setattr(DotProductAttention, "forward", stripped) - - -def _run_mha(mha, x, encoder_output=None): - dpa_module._attention_backends["backend_selection_requires_update"] = True - if encoder_output is not None: - out = mha(x, encoder_output=encoder_output) - else: - out = mha(x) - out.backward(torch.ones_like(out)) - wgrads = [p.grad.clone() for p in mha.parameters() if p.grad is not None] - xgrad = x.grad.clone() - x.grad = None - mha.zero_grad(set_to_none=True) - return out, xgrad, wgrads - - -def _assert_mha_equal(result, reference): - out, xgrad, wgrads = result - out_ref, xgrad_ref, wgrads_ref = reference - assert torch.equal(out, out_ref), "output differs" - assert torch.equal(xgrad, xgrad_ref), "input grad differs" - assert len(wgrads) == len(wgrads_ref) - for i, (w, w_ref) in enumerate(zip(wgrads, wgrads_ref)): - assert torch.equal(w, w_ref), f"weight grad {i} differs" - - -@requires_fused -@pytest.mark.parametrize( - "interleaved", [pytest.param(True, id="interleaved"), pytest.param(False, id="non_interleaved")] -) -def test_mha_self_attention_packed_pass_through(monkeypatch, interleaved): - """MHA self-attention passes its packed projection output straight to DPA as - qkv_layer (with the matching interleave dim), bit-exact vs the same MHA with - packed inputs converted back to separate contiguous q/k/v.""" - _force_backend(monkeypatch, "fused") - hidden = _H * _D - torch.manual_seed(0) - mha = MultiheadAttention( - hidden, - _H, - attention_dropout=0.0, - attn_mask_type="no_mask", - qkv_format="sbhd", - fuse_qkv_params=True, - qkv_weight_interleaved=interleaved, - params_dtype=_DTYPE, - device="cuda", - ) - torch.manual_seed(1) - x = torch.randn(_S, _B, hidden, dtype=_DTYPE, device="cuda", requires_grad=True) - - record = [] - _spy_packed_dpa(monkeypatch, record) - result = _run_mha(mha, x) - assert record == [(True, False, -2 if interleaved else -3)], f"unexpected DPA call: {record}" - monkeypatch.undo() - - _force_backend(monkeypatch, "fused") - _strip_packed_dpa(monkeypatch) - reference = _run_mha(mha, x) - _assert_mha_equal(result, reference) - - -@requires_fused -@pytest.mark.parametrize( - "interleaved", [pytest.param(True, id="interleaved"), pytest.param(False, id="non_interleaved")] -) -def test_mha_cross_attention_packed_kv_pass_through(monkeypatch, interleaved): - """MHA cross-attention passes its packed KV projection output to DPA as - kv_layer, bit-exact vs the separate contiguous reference.""" - _force_backend(monkeypatch, "fused") - hidden = _H * _D - torch.manual_seed(0) - mha = MultiheadAttention( - hidden, - _H, - attention_dropout=0.0, - attn_mask_type="no_mask", - qkv_format="sbhd", - attention_type="cross", - fuse_qkv_params=True, - qkv_weight_interleaved=interleaved, - params_dtype=_DTYPE, - device="cuda", - ) - torch.manual_seed(1) - x = torch.randn(_S, _B, hidden, dtype=_DTYPE, device="cuda", requires_grad=True) - enc = torch.randn(_S, _B, hidden, dtype=_DTYPE, device="cuda") - - record = [] - _spy_packed_dpa(monkeypatch, record) - result = _run_mha(mha, x, encoder_output=enc) - assert record == [(False, True, -2 if interleaved else -3)], f"unexpected DPA call: {record}" - monkeypatch.undo() - - _force_backend(monkeypatch, "fused") - _strip_packed_dpa(monkeypatch) - reference = _run_mha(mha, x, encoder_output=enc) - _assert_mha_equal(result, reference) - - -@requires_fused -def test_mha_gqa_falls_back_to_views(monkeypatch): - """GQA (np != ng) is not a uniform 3-interleave: MHA must keep the legacy - sliced-views path and still work.""" - _force_backend(monkeypatch, "fused") - hidden = _H * _D - torch.manual_seed(0) - mha = MultiheadAttention( - hidden, - _H, - num_gqa_groups=2, - attention_dropout=0.0, - attn_mask_type="no_mask", - qkv_format="sbhd", - fuse_qkv_params=True, - params_dtype=_DTYPE, - device="cuda", - ) - torch.manual_seed(1) - x = torch.randn(_S, _B, hidden, dtype=_DTYPE, device="cuda", requires_grad=True) - - record = [] - _spy_packed_dpa(monkeypatch, record) - out, _, _ = _run_mha(mha, x) - assert record == [(False, False, -3)], f"GQA must not use the packed path: {record}" - assert out.shape == (_S, _B, hidden) - - -@requires_fused -def test_mha_rope_falls_back_to_views(monkeypatch): - """RoPE needs the individual q/k slices: MHA must keep the legacy path.""" - _force_backend(monkeypatch, "fused") - from transformer_engine.pytorch.attention.rope import RotaryPositionEmbedding - - hidden = _H * _D - torch.manual_seed(0) - mha = MultiheadAttention( - hidden, - _H, - attention_dropout=0.0, - attn_mask_type="no_mask", - qkv_format="sbhd", - fuse_qkv_params=True, - params_dtype=_DTYPE, - device="cuda", - ) - rope = RotaryPositionEmbedding(_D)(max_seq_len=_S).to("cuda") - torch.manual_seed(1) - x = torch.randn(_S, _B, hidden, dtype=_DTYPE, device="cuda", requires_grad=True) - - record = [] - _spy_packed_dpa(monkeypatch, record) - dpa_module._attention_backends["backend_selection_requires_update"] = True - out = mha(x, rotary_pos_emb=rope) - assert record == [(False, False, -3)], f"RoPE must not use the packed path: {record}" - assert out.shape == (_S, _B, hidden) From 5cd0351649d811421bff6f9d748f45d58732609c Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Thu, 9 Jul 2026 15:40:50 +0200 Subject: [PATCH 07/15] [PyTorch] Deprecate pointer-based detection of packed qkv layouts get_qkv_layout now emits a DeprecationWarning when it recognizes q/k/v as views of a packed buffer purely from data pointers/strides/offsets (detected *3*/*_2* layouts), pointing callers at the declarative qkv_layer/kv_layer API. Separate q/k/v tensors (hd_hd_hd layouts) never warn since there is nothing to declare. Signed-off-by: Pawel Gadzinski --- .../pytorch/attention/dot_product_attention/utils.py | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py index e1643283d2..bf709753f4 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py @@ -2428,6 +2428,17 @@ def run_iteratively(q, k, v): if qkv_layout == "not_supported": raise RuntimeError("The provided qkv memory layout is not supported!") + if len(qkv_layout.split("_")) < 3: + # q/k/v were recognized as views of a packed buffer only by inspecting + # their data pointers, strides and storage offsets. + warnings.warn( + "Relying on pointer-based detection of packed q/k/v layouts" + f" (detected {qkv_layout!r}) is deprecated: pass the packed buffer" + " explicitly via qkv_layer/kv_layer (with qkv_interleave_dim) to" + " DotProductAttention instead.", + DeprecationWarning, + ) + if inference_params is not None and inference_params.is_paged: qkv_layout = "paged_kv_" + qkv_layout From 361a2a0f37c32a5e24ec0f512504608a65c46e81 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Thu, 9 Jul 2026 15:55:24 +0200 Subject: [PATCH 08/15] [PyTorch] Fix implicit string concatenation lint warning Signed-off-by: Pawel Gadzinski --- .../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 990239e388..cab47964da 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 @@ -276,7 +276,7 @@ def _packed_layout(fmt: str, num: int) -> str: if query_layer is None: raise ValueError( - "kv_layer packs only K and V: query_layer is required when kv_layer" " is provided." + "kv_layer packs only K and V: query_layer is required when kv_layer is provided." ) if key_layer is not None or value_layer is not None: raise ValueError( From 627dd3fe69c179bf6b9c56b320f738b4d3cd9c01 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Thu, 9 Jul 2026 16:39:51 +0200 Subject: [PATCH 09/15] [PyTorch] Trim packed-input test section to equivalence and MHA tests Drop the API-validation, torch.compile graph-break, combine_and_quantize equivalence and thd no-detection spy tests; keep the dense/flash bit-exact equivalence tests and the MHA packed pass-through/fallback tests. Signed-off-by: Pawel Gadzinski --- tests/pytorch/attention/test_attention.py | 180 ---------------------- 1 file changed, 180 deletions(-) diff --git a/tests/pytorch/attention/test_attention.py b/tests/pytorch/attention/test_attention.py index aa70f52565..845968c137 100644 --- a/tests/pytorch/attention/test_attention.py +++ b/tests/pytorch/attention/test_attention.py @@ -26,13 +26,10 @@ 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 ( FlashAttentionUtils, check_set_window_size, - combine_and_quantize, ) -from transformer_engine.pytorch.tensor.float8_tensor import Float8Quantizer from transformer_engine.pytorch.attention import RotaryPositionEmbedding import transformer_engine.pytorch.cpp_extensions as ext from transformer_engine.pytorch.cpp_extensions.fused_attn import ( @@ -3202,183 +3199,6 @@ def test_dpa_flash_qkv_layer(monkeypatch): _assert_bit_exact((out, *grads), reference) -# --------------------------------------------------------------------------- -# 3. Validation errors -# --------------------------------------------------------------------------- - - -def test_dpa_packed_input_validation(): - dpa = _make_dpa("bshd") - qkv = torch.randn( - _PACKED_B, _PACKED_S, 3, _PACKED_H, _PACKED_D, dtype=_PACKED_DTYPE, device="cuda" - ) - kv = torch.randn( - _PACKED_B, _PACKED_S, 2, _PACKED_H, _PACKED_D, dtype=_PACKED_DTYPE, device="cuda" - ) - k = torch.randn(_PACKED_B, _PACKED_S, _PACKED_H, _PACKED_D, dtype=_PACKED_DTYPE, device="cuda") - - with pytest.raises(ValueError, match="must be None when qkv_layer is provided"): - dpa(qkv_layer=qkv, key_layer=k) - with pytest.raises(ValueError, match="query_layer is required when kv_layer"): - dpa(kv_layer=kv) - with pytest.raises(ValueError, match="qkv_interleave_dim must be -3"): - dpa(qkv_layer=qkv, qkv_interleave_dim=-1) - with pytest.raises(ValueError, match="mutually exclusive"): - dpa(qkv_layer=qkv, kv_layer=kv) - with pytest.raises(ValueError, match="must have size 3 at dim"): - dpa(qkv_layer=kv) # 2 at the interleave dim, not 3 - with pytest.raises(ValueError, match="stride 1 in its last"): - dpa(qkv_layer=qkv.transpose(-2, -1)) # declared layout would lie about memory - with pytest.raises(ValueError, match="required unless packed"): - dpa() - - -# --------------------------------------------------------------------------- -# 4. torch.compile: no data_ptr/UntypedStorage graph breaks with qkv_layer -# --------------------------------------------------------------------------- - - -@requires_fused -def test_dpa_torch_compile_qkv_layer_no_pointer_graph_breaks(monkeypatch): - _force_backend(monkeypatch, "fused") - torch._dynamo.reset() - torch._dynamo.utils.counters.clear() - - torch.manual_seed(0) - parts = [ - torch.randn(_PACKED_B, _PACKED_S, _PACKED_H, _PACKED_D, dtype=_PACKED_DTYPE, device="cuda") - for _ in range(3) - ] - qkv = torch.stack(parts, dim=2).requires_grad_() - dpa = _make_dpa("bshd") - - def fn(x): - return dpa(qkv_layer=x) - - _attention_backends["backend_selection_requires_update"] = True - eager_out = fn(qkv) # eager warm-up: backend selection happens outside dynamo - compiled_out = torch.compile(fn)(qkv) - compiled_out.backward(torch.ones_like(compiled_out)) - - breaks = dict(torch._dynamo.utils.counters["graph_break"]) - torch._dynamo.reset() - pointer_breaks = { - reason: count - for reason, count in breaks.items() - if "data_ptr" in reason or "UntypedStorage" in reason - } - assert not pointer_breaks, f"pointer-based graph breaks with qkv_layer: {pointer_breaks}" - assert torch.equal(compiled_out, eager_out), "compiled output differs from eager" - - -# --------------------------------------------------------------------------- -# 5. FP8 combine refactor: combined_qkv/combined_kv path is bit-identical to -# the combine_tensors path -# --------------------------------------------------------------------------- - - -def _fp8_quantizer(): - return Float8Quantizer( - scale=torch.ones(1, dtype=torch.float32, device="cuda"), - amax=torch.zeros(1, dtype=torch.float32, device="cuda"), - fp8_dtype=tex.DType.kFloat8E4M3, - ) - - -def test_combine_and_quantize_combined_matches_views(): - """Quantizing the caller's packed buffer directly (combined_qkv=) produces the - same _data bits and scale_inv as rebuilding the packed buffer from q/k/v views - via combine_tensors (the old set_-based path).""" - torch.manual_seed(0) - qkv = torch.randn( - _PACKED_B, _PACKED_S, 3, _PACKED_H, _PACKED_D, dtype=_PACKED_DTYPE, device="cuda" - ) - q, k, v = qkv[:, :, 0], qkv[:, :, 1], qkv[:, :, 2] - - old = combine_and_quantize("bs3hd", q, k, v, _fp8_quantizer()) - new = combine_and_quantize("bs3hd", q, k, v, _fp8_quantizer(), combined_qkv=qkv) - - assert old[3] == new[3] == "bs3hd" - for name, x, y in zip(("q", "k", "v"), old[:3], new[:3]): - assert torch.equal(x._data, y._data), f"{name} fp8 bits differ" - assert torch.equal(x._scale_inv, y._scale_inv), f"{name} scale_inv differs" - - -def test_combine_and_quantize_combined_kv_matches_views(): - """Same for the kv-packed (group 2) layout.""" - torch.manual_seed(0) - q = torch.randn(_PACKED_B, _PACKED_S, _PACKED_H, _PACKED_D, dtype=_PACKED_DTYPE, device="cuda") - kv = torch.randn( - _PACKED_B, _PACKED_S, 2, _PACKED_H, _PACKED_D, dtype=_PACKED_DTYPE, device="cuda" - ) - k, v = kv[:, :, 0], kv[:, :, 1] - - old = combine_and_quantize("bshd_bs2hd", q, k, v, _fp8_quantizer()) - new = combine_and_quantize("bshd_bs2hd", q, k, v, _fp8_quantizer(), combined_kv=kv) - - for name, x, y in zip(("q", "k", "v"), old[:3], new[:3]): - assert torch.equal(x._data, y._data), f"{name} fp8 bits differ" - assert torch.equal(x._scale_inv, y._scale_inv), f"{name} scale_inv differs" - - -# --------------------------------------------------------------------------- -# 6. thd declarative: t3hd is declared, get_qkv_layout is never called -# --------------------------------------------------------------------------- - - -def test_dpa_thd_qkv_layer_declared_no_detection(monkeypatch): - """Packed thd input (qkv_layer [t,3,h,d]) declares 't3hd' without calling - get_qkv_layout; full forward+backward runs if a thd backend is available.""" - calls = [] - orig_get_qkv_layout = dpa_utils.get_qkv_layout - - def counting(*args, **kwargs): - calls.append(kwargs.get("qkv_format")) - return orig_get_qkv_layout(*args, **kwargs) - - monkeypatch.setattr(dpa_utils, "get_qkv_layout", counting) - - seen_layouts = [] - orig_get_backend = dpa_utils.get_attention_backend - - def recording(params): - seen_layouts.append(params.qkv_layout) - return orig_get_backend(params) - - monkeypatch.setattr(dpa_utils, "get_attention_backend", recording) - - torch.manual_seed(0) - t = _PACKED_B * _PACKED_S - qkv = torch.randn( - t, 3, _PACKED_H, _PACKED_D, dtype=_PACKED_DTYPE, device="cuda", requires_grad=True - ) - cu = _cu_seqlens() - dpa = DotProductAttention( - _PACKED_H, _PACKED_D, attention_dropout=0.0, qkv_format="thd", attn_mask_type="padding" - ) - _attention_backends["backend_selection_requires_update"] = True - ran_full = True - try: - out = dpa( - qkv_layer=qkv, - cu_seqlens_q=cu, - cu_seqlens_kv=cu, - max_seqlen_q=_PACKED_S, - max_seqlen_kv=_PACKED_S, - ) - out.backward(torch.ones_like(out)) - assert qkv.grad is not None and qkv.grad.shape == qkv.shape - except ValueError: - # No thd-capable backend on this device; the layout step (the subject of - # this test) runs before backend dispatch, so the assertions still hold. - ran_full = False - - assert not calls, "get_qkv_layout must not be called for declared packed thd input" - assert seen_layouts == ["t3hd"], f"expected declared layout 't3hd', got {seen_layouts}" - if not ran_full: - pytest.skip("layout declaration verified; no thd backend available for full run") - - # --------------------------------------------------------------------------- # 7. MultiheadAttention adoption: the fused QKV/KV projection output is passed # packed (qkv_layer/kv_layer) to DotProductAttention From a6c454a43275ae44ee10a6f741b02b033a6c68c8 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Fri, 10 Jul 2026 12:22:15 +0200 Subject: [PATCH 10/15] [PyTorch] Drop MHA packed pass-through tests Signed-off-by: Pawel Gadzinski --- tests/pytorch/attention/test_attention.py | 205 ---------------------- 1 file changed, 205 deletions(-) diff --git a/tests/pytorch/attention/test_attention.py b/tests/pytorch/attention/test_attention.py index 845968c137..64548cdd19 100644 --- a/tests/pytorch/attention/test_attention.py +++ b/tests/pytorch/attention/test_attention.py @@ -3197,208 +3197,3 @@ def test_dpa_flash_qkv_layer(monkeypatch): out.backward(torch.ones_like(out)) grads = [qkv.grad[:, :, i] for i in range(3)] _assert_bit_exact((out, *grads), reference) - - -# --------------------------------------------------------------------------- -# 7. MultiheadAttention adoption: the fused QKV/KV projection output is passed -# packed (qkv_layer/kv_layer) to DotProductAttention -# --------------------------------------------------------------------------- - - -def _spy_packed_dpa(monkeypatch, record): - """Record (qkv_layer given, kv_layer given, qkv_interleave_dim) per DPA call.""" - orig = DotProductAttention.forward - - def spy(self, *args, **kwargs): - record.append( - ( - kwargs.get("qkv_layer") is not None, - kwargs.get("kv_layer") is not None, - kwargs.get("qkv_interleave_dim", None), - ) - ) - return orig(self, *args, **kwargs) - - monkeypatch.setattr(DotProductAttention, "forward", spy) - - -def _strip_packed_dpa(monkeypatch): - """Reference path: convert packed DPA inputs back to separate contiguous q/k/v.""" - orig = DotProductAttention.forward - - def stripped(self, query_layer=None, key_layer=None, value_layer=None, *args, **kwargs): - qkv = kwargs.pop("qkv_layer", None) - kv = kwargs.pop("kv_layer", None) - dim = kwargs.pop("qkv_interleave_dim", -3) - if qkv is not None: - query_layer, key_layer, value_layer = ( - qkv.select(dim, i).contiguous() for i in range(3) - ) - elif kv is not None: - key_layer, value_layer = (kv.select(dim, i).contiguous() for i in range(2)) - return orig(self, query_layer, key_layer, value_layer, *args, **kwargs) - - monkeypatch.setattr(DotProductAttention, "forward", stripped) - - -def _run_mha(mha, x, encoder_output=None): - _attention_backends["backend_selection_requires_update"] = True - if encoder_output is not None: - out = mha(x, encoder_output=encoder_output) - else: - out = mha(x) - out.backward(torch.ones_like(out)) - wgrads = [p.grad.clone() for p in mha.parameters() if p.grad is not None] - xgrad = x.grad.clone() - x.grad = None - mha.zero_grad(set_to_none=True) - return out, xgrad, wgrads - - -def _assert_mha_equal(result, reference): - out, xgrad, wgrads = result - out_ref, xgrad_ref, wgrads_ref = reference - assert torch.equal(out, out_ref), "output differs" - assert torch.equal(xgrad, xgrad_ref), "input grad differs" - assert len(wgrads) == len(wgrads_ref) - for i, (w, w_ref) in enumerate(zip(wgrads, wgrads_ref)): - assert torch.equal(w, w_ref), f"weight grad {i} differs" - - -@requires_fused -@pytest.mark.parametrize( - "interleaved", [pytest.param(True, id="interleaved"), pytest.param(False, id="non_interleaved")] -) -def test_mha_self_attention_packed_pass_through(monkeypatch, interleaved): - """MHA self-attention passes its packed projection output straight to DPA as - qkv_layer (with the matching interleave dim), bit-exact vs the same MHA with - packed inputs converted back to separate contiguous q/k/v.""" - _force_backend(monkeypatch, "fused") - hidden = _PACKED_H * _PACKED_D - torch.manual_seed(0) - mha = MultiheadAttention( - hidden, - _PACKED_H, - attention_dropout=0.0, - attn_mask_type="no_mask", - qkv_format="sbhd", - fuse_qkv_params=True, - qkv_weight_interleaved=interleaved, - params_dtype=_PACKED_DTYPE, - device="cuda", - ) - torch.manual_seed(1) - x = torch.randn( - _PACKED_S, _PACKED_B, hidden, dtype=_PACKED_DTYPE, device="cuda", requires_grad=True - ) - - record = [] - _spy_packed_dpa(monkeypatch, record) - result = _run_mha(mha, x) - assert record == [(True, False, -2 if interleaved else -3)], f"unexpected DPA call: {record}" - monkeypatch.undo() - - _force_backend(monkeypatch, "fused") - _strip_packed_dpa(monkeypatch) - reference = _run_mha(mha, x) - _assert_mha_equal(result, reference) - - -@requires_fused -@pytest.mark.parametrize( - "interleaved", [pytest.param(True, id="interleaved"), pytest.param(False, id="non_interleaved")] -) -def test_mha_cross_attention_packed_kv_pass_through(monkeypatch, interleaved): - """MHA cross-attention passes its packed KV projection output to DPA as - kv_layer, bit-exact vs the separate contiguous reference.""" - _force_backend(monkeypatch, "fused") - hidden = _PACKED_H * _PACKED_D - torch.manual_seed(0) - mha = MultiheadAttention( - hidden, - _PACKED_H, - attention_dropout=0.0, - attn_mask_type="no_mask", - qkv_format="sbhd", - attention_type="cross", - fuse_qkv_params=True, - qkv_weight_interleaved=interleaved, - params_dtype=_PACKED_DTYPE, - device="cuda", - ) - torch.manual_seed(1) - x = torch.randn( - _PACKED_S, _PACKED_B, hidden, dtype=_PACKED_DTYPE, device="cuda", requires_grad=True - ) - enc = torch.randn(_PACKED_S, _PACKED_B, hidden, dtype=_PACKED_DTYPE, device="cuda") - - record = [] - _spy_packed_dpa(monkeypatch, record) - result = _run_mha(mha, x, encoder_output=enc) - assert record == [(False, True, -2 if interleaved else -3)], f"unexpected DPA call: {record}" - monkeypatch.undo() - - _force_backend(monkeypatch, "fused") - _strip_packed_dpa(monkeypatch) - reference = _run_mha(mha, x, encoder_output=enc) - _assert_mha_equal(result, reference) - - -@requires_fused -def test_mha_gqa_falls_back_to_views(monkeypatch): - """GQA (np != ng) is not a uniform 3-interleave: MHA must keep the legacy - sliced-views path and still work.""" - _force_backend(monkeypatch, "fused") - hidden = _PACKED_H * _PACKED_D - torch.manual_seed(0) - mha = MultiheadAttention( - hidden, - _PACKED_H, - num_gqa_groups=2, - attention_dropout=0.0, - attn_mask_type="no_mask", - qkv_format="sbhd", - fuse_qkv_params=True, - params_dtype=_PACKED_DTYPE, - device="cuda", - ) - torch.manual_seed(1) - x = torch.randn( - _PACKED_S, _PACKED_B, hidden, dtype=_PACKED_DTYPE, device="cuda", requires_grad=True - ) - - record = [] - _spy_packed_dpa(monkeypatch, record) - out, _, _ = _run_mha(mha, x) - assert record == [(False, False, -3)], f"GQA must not use the packed path: {record}" - assert out.shape == (_PACKED_S, _PACKED_B, hidden) - - -@requires_fused -def test_mha_rope_falls_back_to_views(monkeypatch): - """RoPE needs the individual q/k slices: MHA must keep the legacy path.""" - _force_backend(monkeypatch, "fused") - hidden = _PACKED_H * _PACKED_D - torch.manual_seed(0) - mha = MultiheadAttention( - hidden, - _PACKED_H, - attention_dropout=0.0, - attn_mask_type="no_mask", - qkv_format="sbhd", - fuse_qkv_params=True, - params_dtype=_PACKED_DTYPE, - device="cuda", - ) - rope = RotaryPositionEmbedding(_PACKED_D)(max_seq_len=_PACKED_S).to("cuda") - torch.manual_seed(1) - x = torch.randn( - _PACKED_S, _PACKED_B, hidden, dtype=_PACKED_DTYPE, device="cuda", requires_grad=True - ) - - record = [] - _spy_packed_dpa(monkeypatch, record) - _attention_backends["backend_selection_requires_update"] = True - out = mha(x, rotary_pos_emb=rope) - assert record == [(False, False, -3)], f"RoPE must not use the packed path: {record}" - assert out.shape == (_PACKED_S, _PACKED_B, hidden) From f45338968c84e923ee87675f49d709d7cdf44505 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Fri, 10 Jul 2026 12:30:00 +0200 Subject: [PATCH 11/15] [PyTorch] Fold packed-input tests into test_dpa_qkv_layout via a declarative param Parametrize test_dpa_qkv_layout and test_dpa_qkv_layout_thd with declarative={views,declarative}: the declarative mode passes the packed buffer to DotProductAttention via qkv_layer/kv_layer (declared layout, gradients read off the packed buffer) instead of slicing it into q/k/v views for pointer-based detection. This reuses the whole existing config matrix (masks, bias, SWA, cross-attention, thd, all backends) for the declarative API, replacing the dedicated packed-input test section. Signed-off-by: Pawel Gadzinski --- tests/pytorch/attention/test_attention.py | 298 +++++++--------------- 1 file changed, 94 insertions(+), 204 deletions(-) diff --git a/tests/pytorch/attention/test_attention.py b/tests/pytorch/attention/test_attention.py index 64548cdd19..d08c98f619 100644 --- a/tests/pytorch/attention/test_attention.py +++ b/tests/pytorch/attention/test_attention.py @@ -134,6 +134,7 @@ def test_dot_product_attention( qkv_layout, swa, pad_between_seqs, + declarative_packed=False, ): """Test DotProductAttention module""" @@ -151,6 +152,8 @@ def test_dot_product_attention( qkv_layout = "bshd_bs2hd" if not is_mla and not is_mqa_gqa else "bshd_bshd_bshd" if "3" in qkv_layout and config.attn_type == "cross": pytest.skip("No need to test this layout for cross attention") + if declarative_packed and not any(c.isdigit() for c in qkv_layout): + pytest.skip("Declarative packed inputs only apply to packed qkv layouts.") if config.window_size == (-1, -1) and swa: config.window_size = [2, 2] @@ -225,6 +228,7 @@ def test_dot_product_attention( workspace_opt, pad_between_seqs, is_training, + declarative_packed=declarative_packed, ) # FlashAttention backend @@ -238,6 +242,7 @@ def test_dot_product_attention( workspace_opt, pad_between_seqs, is_training, + declarative_packed=declarative_packed, ) # Compare results @@ -929,9 +934,26 @@ def test_dpa_alibi_slopes(dtype, model_configs, model): @pytest.mark.parametrize("model_configs", [model_configs_layout]) @pytest.mark.parametrize("model", model_configs_layout.keys()) @pytest.mark.parametrize("qkv_layout", qkv_layouts) -def test_dpa_qkv_layout(dtype, model_configs, model, qkv_layout): - """Test DotProductAttention module with different QKV layouts""" - test_dot_product_attention(dtype, model_configs, model, False, True, qkv_layout, False, False) +@pytest.mark.parametrize( + "declarative", [pytest.param(False, id="views"), pytest.param(True, id="declarative")] +) +def test_dpa_qkv_layout(dtype, model_configs, model, qkv_layout, declarative): + """Test DotProductAttention module with different QKV layouts. + + declarative=False passes q/k/v as views of the packed buffer (pointer-based + layout detection); declarative=True passes the packed buffer itself via + qkv_layer/kv_layer (declared layout, gradients on the packed buffer).""" + test_dot_product_attention( + dtype, + model_configs, + model, + False, + True, + qkv_layout, + False, + False, + declarative_packed=declarative, + ) qkv_layouts_thd = ["t3hd", "th3d", "thd_t2hd", "thd_th2d", "thd_thd_thd"] @@ -998,7 +1020,10 @@ def test_dpa_qkv_layout(dtype, model_configs, model, qkv_layout): @pytest.mark.parametrize("model_configs", [model_configs_layout_thd]) @pytest.mark.parametrize("model", model_configs_layout_thd.keys()) @pytest.mark.parametrize("qkv_layout", qkv_layouts_thd) -def test_dpa_qkv_layout_thd(dtype, model_configs, model, qkv_layout): +@pytest.mark.parametrize( + "declarative", [pytest.param(False, id="views"), pytest.param(True, id="declarative")] +) +def test_dpa_qkv_layout_thd(dtype, model_configs, model, qkv_layout, declarative): """Test DotProductAttention module with different QKV layouts""" config = model_configs[model] if config.num_heads != config.num_gqa_groups and "3" in qkv_layout: @@ -1006,14 +1031,30 @@ def test_dpa_qkv_layout_thd(dtype, model_configs, model, qkv_layout): logging.info("[test_dpa_qkv_layout_thd]: pad_between_seqs = True") pad_between_seqs = True test_dot_product_attention( - dtype, model_configs, model, False, True, qkv_layout, False, pad_between_seqs + dtype, + model_configs, + model, + False, + True, + qkv_layout, + False, + pad_between_seqs, + declarative_packed=declarative, ) if get_cudnn_version() >= (9, 3, 0): logging.info("[test_dpa_qkv_layout_thd]: pad_between_seqs = False") # cuDNN 9.3.0+ is required to run pad_between_seqs = False/True in the same run pad_between_seqs = False test_dot_product_attention( - dtype, model_configs, model, False, True, qkv_layout, False, pad_between_seqs + dtype, + model_configs, + model, + False, + True, + qkv_layout, + False, + pad_between_seqs, + declarative_packed=declarative, ) @@ -1026,8 +1067,14 @@ def _run_dot_product_attention( workspace_opt: bool, pad_between_seqs: bool, is_training: bool, + declarative_packed: bool = False, ) -> Tuple[torch.Tensor, Tuple[torch.Tensor, torch.Tensor, torch.Tensor]]: - """Run DotProductAttention module with one forward pass and one backward pass""" + """Run DotProductAttention module with one forward pass and one backward pass. + + With declarative_packed=True (packed qkv_layout only), the packed buffer is + passed to DotProductAttention directly via qkv_layer/kv_layer instead of + slicing it into q/k/v views, and input gradients are read off the packed + buffer itself.""" # Set RNG and environment varables reset_rng_states() os.environ["NVTE_FLASH_ATTN"] = "0" @@ -1225,6 +1272,12 @@ def _run_dot_product_attention( tensor_count = int(l) split_dim = dim break + if declarative_packed and split_dim != 0: + # The packed buffer is the autograd leaf; q/k/v below are non-leaf + # views of it, and DPA receives the buffer via qkv_layer/kv_layer. + tensor.requires_grad_() + packed_tensor = tensor + packed_interleave_dim = split_dim - tensor.dim() tensors = torch.split(tensor, 1, dim=split_dim) if split_dim != 0 else [tensor] tensors_orig = ( torch.split(tensor_orig, 1, dim=split_dim) if split_dim != 0 else [tensor_orig] @@ -1237,8 +1290,10 @@ def _run_dot_product_attention( inp.append(tensors[j]) inp_orig.append(tensors_orig[j]) for i in range(3): - inp[i].requires_grad = True - inp_orig[i].requires_grad = True + if inp[i].is_leaf: + inp[i].requires_grad = True + if inp_orig[i].is_leaf: + inp_orig[i].requires_grad = True # Create output gradient qkv_format_kv = "_".join(qkv_format) @@ -1320,10 +1375,21 @@ def get_dummy_cuda_rng_tracker() -> CudaRNGStatesTracker: k = inp[1] v = inp[2] d_out = out_grad + packed_kwargs = {} + if declarative_packed: + assert backend in ["FusedAttention", "FlashAttention"] + packed_kwargs["qkv_interleave_dim"] = packed_interleave_dim + if len(qkv_layout.split("_")) == 1: + packed_kwargs["qkv_layer"] = packed_tensor + q, k, v = None, None, None + else: + packed_kwargs["kv_layer"] = packed_tensor + k, v = None, None out = block( q, k, v, + **packed_kwargs, window_size=config.window_size, attention_mask=attention_mask, qkv_format=qkv_format, @@ -1353,6 +1419,25 @@ def get_dummy_cuda_rng_tracker() -> CudaRNGStatesTracker: if is_training: out.backward(d_out) + if is_training and declarative_packed: + # Input gradients live on the packed buffer; slice them back out (and + # wrap them in .grad holders) so the cross-backend comparisons below + # stay uniform with the separate-q/k/v path. + assert packed_tensor.grad is not None and packed_tensor.grad.shape == packed_tensor.shape + + class _PackedGrad: + def __init__(self, grad): + self.grad = grad + + packed_grads = [ + _PackedGrad(packed_tensor.grad.select(packed_interleave_dim, j)) + for j in range(packed_tensor.shape[packed_interleave_dim]) + ] + if len(qkv_layout.split("_")) == 1: + q, k, v = packed_grads + else: + k, v = packed_grads + d_softmax_offset = None if is_training and config.softmax_type != "vanilla": d_softmax_offset = block.softmax_offset.grad @@ -3002,198 +3087,3 @@ def forward( self.quantizers, ) return out - - -# --------------------------------------------------------------------------- -# Declarative packed QKV/KV inputs (qkv_layer/kv_layer + qkv_interleave_dim) -# -# Instead of slicing a fused-projection buffer into q/k/v views (which TE then -# reverse-engineers via pointer-based layout detection), callers can pass the -# packed tensor directly. Q/K/V are derived as zero-copy views and the exact -# layout string (e.g. bs3hd) is declared, not detected -- including for thd -# and FP8 DPA. -# --------------------------------------------------------------------------- - -_PACKED_B, _PACKED_S, _PACKED_H, _PACKED_D = 2, 128, 8, 64 -_PACKED_DTYPE = torch.bfloat16 - - -def _cu_seqlens(): - return torch.arange(0, (_PACKED_B + 1) * _PACKED_S, _PACKED_S, dtype=torch.int32, device="cuda") - - -def _fused_backend_supported(): - try: - q = torch.randn( - _PACKED_B, _PACKED_S, _PACKED_H, _PACKED_D, dtype=_PACKED_DTYPE, device="cuda" - ) - fused_attn_fwd( - True, - _PACKED_S, - _PACKED_S, - _cu_seqlens(), - _cu_seqlens(), - q, - q.clone(), - q.clone(), - _PACKED_DTYPE, - tex.NVTE_Fused_Attn_Backend.NVTE_F16_arbitrary_seqlen, - dropout=0.0, - qkv_layout="bshd_bshd_bshd", - o_format="bshd", - attn_bias_type="no_bias", - attn_mask_type="no_mask", - ) - return True - except Exception: - return False - - -requires_fused = pytest.mark.skipif( - not (torch.cuda.is_available() and _fused_backend_supported()), - reason="F16_arbitrary_seqlen fused attention backend is not supported on this device", -) - - -def _force_backend(monkeypatch, backend): - """Force a single attention backend via env and invalidate the selection cache.""" - flash, fused = {"flash": ("1", "0"), "fused": ("0", "1")}[backend] - monkeypatch.setenv("NVTE_FLASH_ATTN", flash) - monkeypatch.setenv("NVTE_FUSED_ATTN", fused) - monkeypatch.setenv("NVTE_UNFUSED_ATTN", "0") - if backend == "flash": - # flash-attn bwd uses atomics unless deterministic - monkeypatch.setenv("NVTE_ALLOW_NONDETERMINISTIC_ALGO", "0") - _attention_backends["backend_selection_requires_update"] = True - - -def _make_dpa(qkv_format, num_gqa_groups=None): - return DotProductAttention( - _PACKED_H, - _PACKED_D, - num_gqa_groups=num_gqa_groups, - attention_dropout=0.0, - qkv_format=qkv_format, - attn_mask_type="no_mask", - ) - - -def _assert_bit_exact(result, reference, names=("out", "dq", "dk", "dv")): - for name, x, y in zip(names, result, reference): - assert torch.equal(x.contiguous(), y.contiguous()), f"{name} differs" - - -def _dpa_separate_baseline(qkv_format, num_gqa_groups=None): - """Fwd+bwd on contiguous separate q/k/v leaves; returns (out, dq, dk, dv).""" - hg = num_gqa_groups or _PACKED_H - torch.manual_seed(0) - q_shape = ( - (_PACKED_B, _PACKED_S, _PACKED_H, _PACKED_D) - if qkv_format == "bshd" - else (_PACKED_S, _PACKED_B, _PACKED_H, _PACKED_D) - ) - kv_shape = ( - (_PACKED_B, _PACKED_S, hg, _PACKED_D) - if qkv_format == "bshd" - else (_PACKED_S, _PACKED_B, hg, _PACKED_D) - ) - q = torch.randn(*q_shape, dtype=_PACKED_DTYPE, device="cuda") - k = torch.randn(*kv_shape, dtype=_PACKED_DTYPE, device="cuda") - v = torch.randn(*kv_shape, dtype=_PACKED_DTYPE, device="cuda") - q, k, v = [x.clone().requires_grad_() for x in (q, k, v)] - _attention_backends["backend_selection_requires_update"] = True - out = _make_dpa(qkv_format, num_gqa_groups)(q, k, v) - out.backward(torch.ones_like(out)) - return out, q.grad, k.grad, v.grad - - -# --------------------------------------------------------------------------- -# 1. Dense eager equivalence (fused backend), fwd + input grads, bit-exact -# --------------------------------------------------------------------------- - - -@requires_fused -@pytest.mark.parametrize( - "qkv_format, interleave_dim", - [ - pytest.param("bshd", -3, id="bshd_qkv_dim-3"), # (a) bs3hd - pytest.param("bshd", -2, id="bshd_qkv_dim-2"), # (b) bsh3d (Megatron-style) - pytest.param("sbhd", -3, id="sbhd_qkv_dim-3"), # (c) sb3hd - ], -) -def test_dpa_fused_qkv_layer_dense(monkeypatch, qkv_format, interleave_dim): - """qkv_layer packed input is bit-exact vs separate contiguous q/k/v, and grads - flow back into the packed tensor itself.""" - _force_backend(monkeypatch, "fused") - reference = _dpa_separate_baseline(qkv_format) - - torch.manual_seed(0) - q_shape = ( - (_PACKED_B, _PACKED_S, _PACKED_H, _PACKED_D) - if qkv_format == "bshd" - else (_PACKED_S, _PACKED_B, _PACKED_H, _PACKED_D) - ) - parts = [torch.randn(*q_shape, dtype=_PACKED_DTYPE, device="cuda") for _ in range(3)] - stack_dim = len(q_shape) + interleave_dim + 1 # -3 -> before h, -2 -> before d - qkv = torch.stack(parts, dim=stack_dim).requires_grad_() - - _attention_backends["backend_selection_requires_update"] = True - out = _make_dpa(qkv_format)(qkv_layer=qkv, qkv_interleave_dim=interleave_dim) - out.backward(torch.ones_like(out)) - - assert qkv.grad is not None and qkv.grad.shape == qkv.shape - grads = [qkv.grad.select(stack_dim, i) for i in range(3)] - _assert_bit_exact((out, *grads), reference) - - -@requires_fused -@pytest.mark.parametrize( - "num_gqa_groups", - [pytest.param(None, id="mha_kv"), pytest.param(2, id="gqa_kv")], # (d) and (e) -) -def test_dpa_fused_kv_layer_dense(monkeypatch, num_gqa_groups): - """kv_layer packed input (with separate query) is bit-exact vs separate - contiguous q/k/v; grads flow into the packed kv tensor.""" - _force_backend(monkeypatch, "fused") - reference = _dpa_separate_baseline("bshd", num_gqa_groups) - - hg = num_gqa_groups or _PACKED_H - torch.manual_seed(0) - q = torch.randn(_PACKED_B, _PACKED_S, _PACKED_H, _PACKED_D, dtype=_PACKED_DTYPE, device="cuda") - k = torch.randn(_PACKED_B, _PACKED_S, hg, _PACKED_D, dtype=_PACKED_DTYPE, device="cuda") - v = torch.randn(_PACKED_B, _PACKED_S, hg, _PACKED_D, dtype=_PACKED_DTYPE, device="cuda") - q = q.clone().requires_grad_() - kv = torch.stack([k, v], dim=2).requires_grad_() # [b,s,2,hg,d] - - _attention_backends["backend_selection_requires_update"] = True - out = _make_dpa("bshd", num_gqa_groups)(query_layer=q, kv_layer=kv) - out.backward(torch.ones_like(out)) - - assert kv.grad is not None and kv.grad.shape == kv.shape - _assert_bit_exact((out, q.grad, kv.grad[:, :, 0], kv.grad[:, :, 1]), reference) - - -# --------------------------------------------------------------------------- -# 2. Flash backend smoke -# --------------------------------------------------------------------------- - - -def test_dpa_flash_qkv_layer(monkeypatch): - """Flash backend: packed qkv_layer [b,s,3,h,d] is bit-exact vs separate.""" - _force_backend(monkeypatch, "flash") - try: - reference = _dpa_separate_baseline("bshd") - except Exception as exc: - pytest.skip(f"flash attention backend not available: {exc}") - - torch.manual_seed(0) - parts = [ - torch.randn(_PACKED_B, _PACKED_S, _PACKED_H, _PACKED_D, dtype=_PACKED_DTYPE, device="cuda") - for _ in range(3) - ] - qkv = torch.stack(parts, dim=2).requires_grad_() - _attention_backends["backend_selection_requires_update"] = True - out = _make_dpa("bshd")(qkv_layer=qkv) - out.backward(torch.ones_like(out)) - grads = [qkv.grad[:, :, i] for i in range(3)] - _assert_bit_exact((out, *grads), reference) From 2e0f32de143e8ed45bb1329518b73b9947ff47da Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Fri, 10 Jul 2026 12:41:52 +0200 Subject: [PATCH 12/15] [PyTorch] Shrink declarative packed-input tests to a dedicated small matrix Revert the declarative parametrization of test_dpa_qkv_layout(_thd) (which doubled their whole config x layout product) and instead add test_dpa_qkv_layout(_thd)_declarative covering all packed layouts on a trimmed config dimension: one self-attention and one cross-attention config (kv_layer path) for dense, one config for thd. Past the input handling the backend code is identical to the views mode, so the full config matrix added no coverage. Signed-off-by: Pawel Gadzinski --- tests/pytorch/attention/test_attention.py | 62 ++++++++++++++--------- 1 file changed, 39 insertions(+), 23 deletions(-) diff --git a/tests/pytorch/attention/test_attention.py b/tests/pytorch/attention/test_attention.py index d08c98f619..3569ca3cdf 100644 --- a/tests/pytorch/attention/test_attention.py +++ b/tests/pytorch/attention/test_attention.py @@ -934,25 +934,28 @@ def test_dpa_alibi_slopes(dtype, model_configs, model): @pytest.mark.parametrize("model_configs", [model_configs_layout]) @pytest.mark.parametrize("model", model_configs_layout.keys()) @pytest.mark.parametrize("qkv_layout", qkv_layouts) -@pytest.mark.parametrize( - "declarative", [pytest.param(False, id="views"), pytest.param(True, id="declarative")] -) -def test_dpa_qkv_layout(dtype, model_configs, model, qkv_layout, declarative): - """Test DotProductAttention module with different QKV layouts. +def test_dpa_qkv_layout(dtype, model_configs, model, qkv_layout): + """Test DotProductAttention module with different QKV layouts""" + test_dot_product_attention(dtype, model_configs, model, False, True, qkv_layout, False, False) + - declarative=False passes q/k/v as views of the packed buffer (pointer-based - layout detection); declarative=True passes the packed buffer itself via - qkv_layer/kv_layer (declared layout, gradients on the packed buffer).""" +qkv_layouts_packed = [l for l in qkv_layouts if any(c.isdigit() for c in l)] + + +@pytest.mark.skipif(get_cudnn_version() < (8, 9, 5), reason="cuDNN 8.9.5+ is required.") +@pytest.mark.parametrize("dtype", param_types_lean) +@pytest.mark.parametrize("model_configs", [model_configs_layout]) +@pytest.mark.parametrize("model", ["layout_1_1", "layout_1_2"]) +@pytest.mark.parametrize("qkv_layout", qkv_layouts_packed) +def test_dpa_qkv_layout_declarative(dtype, model_configs, model, qkv_layout): + """Declarative packed inputs: the packed buffer is passed to + DotProductAttention via qkv_layer/kv_layer (declared layout, gradients read + off the packed buffer) instead of q/k/v views + pointer-based detection. + Layout coverage is complete; the model-config dimension is trimmed to one + self-attention and one cross-attention config, since past the input + handling the backend code is identical to test_dpa_qkv_layout.""" test_dot_product_attention( - dtype, - model_configs, - model, - False, - True, - qkv_layout, - False, - False, - declarative_packed=declarative, + dtype, model_configs, model, False, True, qkv_layout, False, False, declarative_packed=True ) @@ -1020,10 +1023,7 @@ def test_dpa_qkv_layout(dtype, model_configs, model, qkv_layout, declarative): @pytest.mark.parametrize("model_configs", [model_configs_layout_thd]) @pytest.mark.parametrize("model", model_configs_layout_thd.keys()) @pytest.mark.parametrize("qkv_layout", qkv_layouts_thd) -@pytest.mark.parametrize( - "declarative", [pytest.param(False, id="views"), pytest.param(True, id="declarative")] -) -def test_dpa_qkv_layout_thd(dtype, model_configs, model, qkv_layout, declarative): +def test_dpa_qkv_layout_thd(dtype, model_configs, model, qkv_layout, declarative_packed=False): """Test DotProductAttention module with different QKV layouts""" config = model_configs[model] if config.num_heads != config.num_gqa_groups and "3" in qkv_layout: @@ -1039,7 +1039,7 @@ def test_dpa_qkv_layout_thd(dtype, model_configs, model, qkv_layout, declarative qkv_layout, False, pad_between_seqs, - declarative_packed=declarative, + declarative_packed=declarative_packed, ) if get_cudnn_version() >= (9, 3, 0): logging.info("[test_dpa_qkv_layout_thd]: pad_between_seqs = False") @@ -1054,10 +1054,26 @@ def test_dpa_qkv_layout_thd(dtype, model_configs, model, qkv_layout, declarative qkv_layout, False, pad_between_seqs, - declarative_packed=declarative, + declarative_packed=declarative_packed, ) +qkv_layouts_thd_packed = [l for l in qkv_layouts_thd if any(c.isdigit() for c in l)] + + +@pytest.mark.skipif(get_cudnn_version() < (9, 0, 0), reason="cuDNN 9.0.0+ is required.") +@pytest.mark.skipif( + get_device_compute_capability() < (9, 0), reason="THD is only supported on Hopper+." +) +@pytest.mark.parametrize("dtype", param_types_lean) +@pytest.mark.parametrize("model_configs", [model_configs_layout_thd]) +@pytest.mark.parametrize("model", ["layout_0_0"]) +@pytest.mark.parametrize("qkv_layout", qkv_layouts_thd_packed) +def test_dpa_qkv_layout_thd_declarative(dtype, model_configs, model, qkv_layout): + """Declarative packed thd inputs, see test_dpa_qkv_layout_declarative.""" + test_dpa_qkv_layout_thd(dtype, model_configs, model, qkv_layout, declarative_packed=True) + + def _run_dot_product_attention( dtype: torch.dtype, config: ModelConfig, From 36cf09e5f6bf0a879ae1fbb98845a2d8afad16dc Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Fri, 10 Jul 2026 12:52:13 +0200 Subject: [PATCH 13/15] [PyTorch] Use explicit q/k/v grad variables in the DPA test harness Replace the .grad-holder objects substituted for q/k/v in declarative packed mode with q_grad/k_grad/v_grad variables computed right after backward, used uniformly by all return paths. Signed-off-by: Pawel Gadzinski --- tests/pytorch/attention/test_attention.py | 52 +++++++++++------------ 1 file changed, 24 insertions(+), 28 deletions(-) diff --git a/tests/pytorch/attention/test_attention.py b/tests/pytorch/attention/test_attention.py index 3569ca3cdf..e0b2758440 100644 --- a/tests/pytorch/attention/test_attention.py +++ b/tests/pytorch/attention/test_attention.py @@ -1435,34 +1435,33 @@ def get_dummy_cuda_rng_tracker() -> CudaRNGStatesTracker: if is_training: out.backward(d_out) - if is_training and declarative_packed: - # Input gradients live on the packed buffer; slice them back out (and - # wrap them in .grad holders) so the cross-backend comparisons below - # stay uniform with the separate-q/k/v path. - assert packed_tensor.grad is not None and packed_tensor.grad.shape == packed_tensor.shape - - class _PackedGrad: - def __init__(self, grad): - self.grad = grad - - packed_grads = [ - _PackedGrad(packed_tensor.grad.select(packed_interleave_dim, j)) - for j in range(packed_tensor.shape[packed_interleave_dim]) - ] - if len(qkv_layout.split("_")) == 1: - q, k, v = packed_grads + q_grad, k_grad, v_grad = None, None, None + if is_training: + if declarative_packed: + # Input gradients live on the packed buffer; slice them back out so + # the cross-backend comparisons below stay uniform with the + # separate-q/k/v path. + assert ( + packed_tensor.grad is not None and packed_tensor.grad.shape == packed_tensor.shape + ) + packed_grads = [ + packed_tensor.grad.select(packed_interleave_dim, j) + for j in range(packed_tensor.shape[packed_interleave_dim]) + ] + if len(qkv_layout.split("_")) == 1: + q_grad, k_grad, v_grad = packed_grads + else: + q_grad = q.grad + k_grad, v_grad = packed_grads else: - k, v = packed_grads + q_grad, k_grad, v_grad = q.grad, k.grad, v.grad d_softmax_offset = None if is_training and config.softmax_type != "vanilla": d_softmax_offset = block.softmax_offset.grad if backend in ["UnfusedDotProductAttention"]: - if is_training: - return out, max_logit, (q.grad, k.grad, v.grad, d_softmax_offset) - else: - return out, max_logit, (None, None, None, d_softmax_offset) + return out, max_logit, (q_grad, k_grad, v_grad, d_softmax_offset) if backend in ["FusedAttention", "FlashAttention"]: if qkv_format == "thd" and pad_between_seqs: out_orig = torch.Tensor([]).to(device="cuda", dtype=dtype) @@ -1482,13 +1481,13 @@ def __init__(self, grad): out_orig = torch.cat([out_orig, out[valid_range_q[0] : valid_range_q[1]]], dim=0) if is_training: q_grad_orig = torch.cat( - [q_grad_orig, q.grad[valid_range_q[0] : valid_range_q[1]]], dim=0 + [q_grad_orig, q_grad[valid_range_q[0] : valid_range_q[1]]], dim=0 ) k_grad_orig = torch.cat( - [k_grad_orig, k.grad[valid_range_kv[0] : valid_range_kv[1]]], dim=0 + [k_grad_orig, k_grad[valid_range_kv[0] : valid_range_kv[1]]], dim=0 ) v_grad_orig = torch.cat( - [v_grad_orig, v.grad[valid_range_kv[0] : valid_range_kv[1]]], dim=0 + [v_grad_orig, v_grad[valid_range_kv[0] : valid_range_kv[1]]], dim=0 ) if is_training: return ( @@ -1499,10 +1498,7 @@ def __init__(self, grad): else: return out_orig, max_logit, (None, None, None, d_softmax_offset) else: - if is_training: - return out, max_logit, (q.grad, k.grad, v.grad, d_softmax_offset) - else: - return out, max_logit, (None, None, None, d_softmax_offset) + return out, max_logit, (q_grad, k_grad, v_grad, d_softmax_offset) model_configs_te_layer = { From 560e24649c7cf75931c9ba3769e30c9c4d807c02 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 10 Jul 2026 11:09:52 +0000 Subject: [PATCH 14/15] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- .../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 55ee89c701..82a23ff021 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 @@ -284,7 +284,7 @@ def _packed_layout(fmt: str, num: int) -> str: ) if kv_layer.dim() != query_layer.dim() + 1: raise ValueError( - f"kv_layer must have one more dimension than query_layer, got" + "kv_layer must have one more dimension than query_layer, got" f" {kv_layer.dim()}D kv_layer and {query_layer.dim()}D query_layer." ) if kv_layer.shape[qkv_interleave_dim] != 2: From 86bb5bc354e6058b43d41dbd23fc052d2716a418 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Fri, 10 Jul 2026 13:56:38 +0200 Subject: [PATCH 15/15] [PyTorch] Address review: warning stacklevel, offload suppression, explicit k_norm - get_qkv_layout deprecation warning: add stacklevel=2 and skip it while CPU offloading is enabled (offloading forces MultiheadAttention onto the sliced-views fallback, so the caller has no migration option there). - MultiheadAttention: gate the packed pass-through on k_norm explicitly instead of relying on q_norm/k_norm being created together. Signed-off-by: Pawel Gadzinski --- .../attention/dot_product_attention/utils.py | 22 ++++++++++++------- .../pytorch/attention/multi_head_attention.py | 1 + 2 files changed, 15 insertions(+), 8 deletions(-) diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py index 54ce11950e..4949641810 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py @@ -35,6 +35,7 @@ META_DP, ) from transformer_engine.pytorch.attention.inference import InferenceParams +from transformer_engine.pytorch.cpu_offload import is_cpu_offload_enabled from transformer_engine.pytorch.quantized_tensor import QuantizedTensorStorage from transformer_engine.pytorch.tensor.float8_tensor import ( Float8Tensor, @@ -2434,14 +2435,19 @@ def run_iteratively(q, k, v): if len(qkv_layout.split("_")) < 3: # q/k/v were recognized as views of a packed buffer only by inspecting - # their data pointers, strides and storage offsets. - warnings.warn( - "Relying on pointer-based detection of packed q/k/v layouts" - f" (detected {qkv_layout!r}) is deprecated: pass the packed buffer" - " explicitly via qkv_layer/kv_layer (with qkv_interleave_dim) to" - " DotProductAttention instead.", - DeprecationWarning, - ) + # their data pointers, strides and storage offsets. Skip the nudge while + # CPU offloading is enabled: offloading forces MultiheadAttention onto + # its sliced-views fallback, so packed views reaching detection are + # expected there and the caller has no migration option. + if not is_cpu_offload_enabled(): + warnings.warn( + "Relying on pointer-based detection of packed q/k/v layouts" + f" (detected {qkv_layout!r}) is deprecated: pass the packed buffer" + " explicitly via qkv_layer/kv_layer (with qkv_interleave_dim) to" + " DotProductAttention instead.", + DeprecationWarning, + stacklevel=2, + ) if inference_params is not None and inference_params.is_paged: qkv_layout = "paged_kv_" + qkv_layout diff --git a/transformer_engine/pytorch/attention/multi_head_attention.py b/transformer_engine/pytorch/attention/multi_head_attention.py index c3b32e7c03..0b73daf667 100644 --- a/transformer_engine/pytorch/attention/multi_head_attention.py +++ b/transformer_engine/pytorch/attention/multi_head_attention.py @@ -913,6 +913,7 @@ def forward( packed_dpa_eligible = ( rotary_pos_emb is None and self.q_norm is None + and self.k_norm is None and inference_params is None and not is_cpu_offload_enabled() )