diff --git a/docs/api/pytorch.rst b/docs/api/pytorch.rst index 2d8cc87827..c60ea86fff 100644 --- a/docs/api/pytorch.rst +++ b/docs/api/pytorch.rst @@ -22,6 +22,10 @@ PyTorch .. autoapiclass:: transformer_engine.pytorch.LayerNormMLP(hidden_size, ffn_hidden_size, eps=1e-5, bias=True, **kwargs) :members: forward, set_tensor_parallel_group +.. autoapifunction:: transformer_engine.pytorch.attention.register_cu_seqlens + +.. autoapifunction:: transformer_engine.pytorch.attention.attention_backend_workspace + .. autoapiclass:: transformer_engine.pytorch.DotProductAttention(num_attention_heads, kv_channels, **kwargs) :members: forward, set_context_parallel_group diff --git a/docs/envvars.rst b/docs/envvars.rst index 8dd5043e38..9329fed1ec 100644 --- a/docs/envvars.rst +++ b/docs/envvars.rst @@ -223,6 +223,7 @@ backend-selection overview. :Type: ``int`` (0 or 1) :Default: ``1`` :Description: Enable or disable FusedAttention backend (cuDNN-based) for DotProductAttention. When set to ``0``, FusedAttention will not be used. + ``NVTE_CUDNN_COMPACT_GQA_BWD=1`` enables experimental SM107 compact GQA backward (default: ``0``); call ``attention.register_cu_seqlens(tensor, offsets)`` for each prefix tensor after its final write and before forward. .. envvar:: NVTE_UNFUSED_ATTN diff --git a/qa/L0_pytorch_unittest/test.sh b/qa/L0_pytorch_unittest/test.sh index 2aa0234001..1350005080 100644 --- a/qa/L0_pytorch_unittest/test.sh +++ b/qa/L0_pytorch_unittest/test.sh @@ -71,6 +71,7 @@ python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_identity_quantiz NVTE_ALLOW_UNSAFE_PICKLE_EXTRA_STATE=1 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_attention.xml $TE_PATH/tests/pytorch/attention/test_attention.py || test_fail "test_attention.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_attention_backend_selection.xml $TE_PATH/tests/pytorch/attention/test_attention_backend_selection.py || test_fail "test_attention_backend_selection.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_flex_attention.xml $TE_PATH/tests/pytorch/attention/test_flex_attention.py || test_fail "test_flex_attention.py" +python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_compact_gqa.xml $TE_PATH/tests/pytorch/attention/test_compact_gqa.py || test_fail "test_compact_gqa.py" NVTE_GDN_TEST_REQUIRED=1 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_gdn_attention.xml $TE_PATH/tests/pytorch/attention/test_gdn_attention.py || test_fail "test_gdn_attention.py" NVTE_GDN2_TEST_REQUIRED=1 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_gdn2_attention.xml $TE_PATH/tests/pytorch/attention/test_gdn2_attention.py || test_fail "test_gdn2_attention.py" NVTE_GDP_TEST_REQUIRED=1 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_gdp_attention.xml $TE_PATH/tests/pytorch/attention/test_gdp_attention.py || test_fail "test_gdp_attention.py" diff --git a/tests/pytorch/attention/test_compact_gqa.py b/tests/pytorch/attention/test_compact_gqa.py new file mode 100644 index 0000000000..89dad1b69f --- /dev/null +++ b/tests/pytorch/attention/test_compact_gqa.py @@ -0,0 +1,386 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Compact GQA metadata, workspace, dispatch, and GPU numerics.""" +import gc +import weakref +from dataclasses import replace +from itertools import accumulate +from types import ModuleType, SimpleNamespace +from unittest.mock import Mock, patch + +import pytest +import torch + +from transformer_engine.pytorch.attention import packed_sequence as metadata +from transformer_engine.pytorch.attention.dot_product_attention import compact_gqa + + +@pytest.mark.parametrize("offsets", [(0,), (0, 1), (0, 0, 257, 4097), (0, 16384, 65536)]) +def test_registered_metadata_never_reads_tensor(offsets, monkeypatch): + tensor = torch.tensor(offsets, dtype=torch.int32) + monkeypatch.setattr(torch.Tensor, "tolist", lambda *_: pytest.fail("Unexpected readback")) + metadata.register_cu_seqlens(tensor, offsets) + assert metadata._get_cu_seqlens(tensor) == offsets + + +def test_mutation_and_clone_require_registration(): + tensor = torch.tensor([0, 3, 8], dtype=torch.int32) + metadata.register_cu_seqlens(tensor, [0, 3, 8]) + assert metadata._get_cu_seqlens(tensor.clone()) is None + alias = tensor.view_as(tensor) + assert metadata._get_cu_seqlens(alias) == (0, 3, 8) + alias[-1] = 9 + assert metadata._get_cu_seqlens(tensor) is None + metadata.register_cu_seqlens(tensor, [0, 3, 9]) + assert metadata._get_cu_seqlens(tensor) == (0, 3, 9) + + +def test_registration_does_not_retain_tensor_or_mutable_host_list(): + tensor = torch.tensor([0, 3], dtype=torch.int32) + key = (tensor.device, tensor.data_ptr()) + host = [0, 3] + metadata.register_cu_seqlens(tensor, host) + host[-1] = 7 + assert metadata._get_cu_seqlens(tensor) == (0, 3) + ref = weakref.ref(tensor) + del tensor + gc.collect() + assert ref() is None + assert key not in metadata._PREFIXES + + +@pytest.mark.parametrize("offsets", [(1, 3), (0, 3, 2), (), (0, 2**31)]) +def test_invalid_offsets(offsets): + with pytest.raises(ValueError): + metadata.register_cu_seqlens(torch.zeros(len(offsets), dtype=torch.int32), offsets) + + +def test_tensor_metadata_is_not_implicitly_copied_to_host(): + tensor = torch.tensor([0, 3], dtype=torch.int32) + with pytest.raises(TypeError): + metadata.register_cu_seqlens(tensor, tensor) + + +def test_shape_and_dtype_mismatch(): + with pytest.raises(ValueError): + metadata.register_cu_seqlens(torch.tensor([0, 3]), [0, 3]) + with pytest.raises(ValueError): + metadata.register_cu_seqlens(torch.zeros(3, dtype=torch.int32), [0, 3]) + + +def test_inference_tensor_is_not_retained(): + with torch.inference_mode(): + tensor = torch.tensor([0, 3], dtype=torch.int32) + metadata.register_cu_seqlens(tensor, [0, 3]) + assert metadata._get_cu_seqlens(tensor) is None + + +def test_workspace_nested_and_exception_cleanup(): + assert metadata._get_attention_backend_workspace() is None + with pytest.raises(RuntimeError, match="training failed"): + with metadata.attention_backend_workspace(): + owner = metadata._get_attention_backend_workspace() + tensor = torch.empty(10) + ref = weakref.ref(tensor) + owner.buffers["test"] = tensor + del tensor + with metadata.attention_backend_workspace(): + assert metadata._get_attention_backend_workspace() is owner + assert owner.active and ref() is not None + raise RuntimeError("training failed") + assert not owner.active and not owner.buffers and ref() is None + assert metadata._get_attention_backend_workspace() is None + + +def inputs(lengths=(3, 5)): + offsets = [0] + for length in lengths: + offsets.append(offsets[-1] + length) + tokens = offsets[-1] + q = torch.zeros((tokens, 8, 256), dtype=torch.bfloat16) + k = torch.zeros((tokens, 1, 256), dtype=torch.bfloat16) + cu = torch.tensor(offsets, dtype=torch.int32) + metadata.register_cu_seqlens(cu, offsets) + return SimpleNamespace( + q=q, + k=k, + v=k.clone(), + out=q.clone(), + cu_seqlens_q=cu, + cu_seqlens_kv=cu, + cu_seqlens_q_padded=None, + cu_seqlens_kv_padded=None, + max_seqlen_q=max(lengths), + max_seqlen_kv=max(lengths), + backend_workspace=metadata._get_attention_backend_workspace(), + fp8=False, + is_input_fp8=False, + use_FAv2_bwd=False, + deterministic=False, + qkv_layout="thd_thd_thd", + dqkv_layout="thd_thd_thd", + o_format="thd", + nominal_dtype=torch.bfloat16, + dropout_p=0.0, + attn_bias_type="no_bias", + softmax_type="vanilla", + attn_mask_type="padding_causal", + attn_scale=1 / 16, + window_size=(-1, -1), + ) + + +def test_disabled_dispatch_does_not_inspect_or_import_backend(monkeypatch): + monkeypatch.setattr(compact_gqa, "_ENABLED", False) + assert compact_gqa.try_compact_gqa_backward(None, None, None) is None + + +def test_cpu_inputs_use_stock(): + args = inputs() + assert not compact_gqa._eligible(args, args.q, torch.zeros(8, 8)) + + +def test_unregistered_and_mutated_prefixes_use_stock(): + args = inputs() + assert compact_gqa._packing(args) == ((0, 3, 8), (3, 5)) + args.cu_seqlens_q[-1] = 7 + assert compact_gqa._packing(args) is None + + +def test_padding_and_mismatched_kv(): + args = inputs() + padded = torch.tensor([0, 4, 12], dtype=torch.int32) + metadata.register_cu_seqlens(padded, [0, 4, 12]) + args.q = torch.empty((12, 8, 256)) + args.cu_seqlens_q_padded = args.cu_seqlens_kv_padded = padded + assert compact_gqa._packing(args) == ((0, 4, 12), (3, 5)) + other = torch.tensor([0, 4, 8], dtype=torch.int32) + metadata.register_cu_seqlens(other, [0, 4, 8]) + args.cu_seqlens_kv = other + assert compact_gqa._packing(args) is None + + +def test_workspace_reuse_resize_and_release(monkeypatch): + import sys + + monkeypatch.setattr(compact_gqa, "_ENABLED", True) + monkeypatch.setattr(compact_gqa, "_eligible", lambda *_: True) + monkeypatch.setattr(compact_gqa, "_PLANS", {}) + monkeypatch.setattr(compact_gqa, "_BOUND_WORKSPACES", {}) + monkeypatch.setattr(torch.cuda, "current_stream", lambda *_: SimpleNamespace(cuda_stream=3)) + initialized_capacity = 0 + + def initialize_workspace(buffer, stream, *, max_seqlen): + nonlocal initialized_capacity + initialized_capacity = max_seqlen + + def backward(q, k, v, *args, sequence_lengths, **kwargs): + assert max(sequence_lengths) <= initialized_capacity + return {"dq": q, "dk": k, "dv": v} + + plan = SimpleNamespace( + compile=Mock(), + initialize_workspace=Mock(side_effect=initialize_workspace), + scratch_workspace_bytes=lambda n: 256 + n, + ) + api = ModuleType("cudnn.sdpa.bwd.compact_gqa") + api.CompactGqaBackward = Mock(return_value=plan) + api.compact_gqa_backward = Mock(side_effect=backward) + monkeypatch.setitem(sys.modules, api.__name__, api) + with metadata.attention_backend_workspace(): + owner = metadata._get_attention_backend_workspace() + for lengths in ((3, 5), (3, 5), (7, 9), (3, 5)): + args = inputs(lengths) + lse = torch.zeros(args.q.shape[0], 8) + result = compact_gqa.try_compact_gqa_backward(args, args.q, lse) + assert result[0] is args.q + assert plan.compile.call_count == 1 + assert plan.initialize_workspace.call_count == 2 + assert api.CompactGqaBackward.call_count == 1 + assert api.compact_gqa_backward.call_args.kwargs["sequence_lengths"] == (3, 5) + assert owner.buffers + assert not owner.buffers and not owner.active + assert compact_gqa.try_compact_gqa_backward(args, args.q, lse) is None + + +@pytest.mark.parametrize( + "field,value", + [ + ("fp8", True), + ("is_input_fp8", True), + ("use_FAv2_bwd", True), + ("deterministic", True), + ("dropout_p", 0.1), + ("attn_bias_type", "post_scale_bias"), + ("softmax_type", "off-by-one"), + ("attn_mask_type", "no_mask"), + ("qkv_layout", "bshd_bshd_bshd"), + ("window_size", (128, 0)), + ], +) +def test_unsupported_modes(field, value, monkeypatch): + class CudaMetadataTensor(torch.Tensor): + @property + def is_cuda(self): + return True + + args = inputs() + args.q = args.q.as_subclass(CudaMetadataTensor) + monkeypatch.setattr(torch.cuda, "get_device_capability", lambda *_: (10, 7)) + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False) + lse = torch.zeros(8, 8) + assert compact_gqa._eligible(args, args.q, lse) + setattr(args, field, value) + assert not compact_gqa._eligible(args, args.q, lse) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize( + "lengths,spans", + [ + ((3, 125), (16, 128)), + ((257, 1025), (384, 1152)), + ((4097, 8191), (4097, 8191)), + ], +) +def test_native_compact_dispatch_matches_stock(lengths, spans, monkeypatch): + if not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 7): + pytest.skip("Compact GQA requires SM107") + monkeypatch.setattr(compact_gqa, "_PLANS", {}) + monkeypatch.setattr(compact_gqa, "_BOUND_WORKSPACES", {}) + small = ((3, 125), (16, 128)) + cached_plan = None + allocated = 0 + try: + with metadata.attention_backend_workspace(): + owner = metadata._get_attention_backend_workspace() + for current_lengths, current_spans in (small, (lengths, spans), small): + _check_native_compact_dispatch(current_lengths, current_spans) + assert len(compact_gqa._PLANS) == 1 + plan = next(iter(compact_gqa._PLANS.values())) + if cached_plan is None: + cached_plan = plan + assert plan is cached_plan + assert len(owner.buffers) == 1 + _, capacity = next(iter(owner.buffers.values())) + allocated = max(allocated, *current_lengths) + assert capacity == allocated + finally: + torch.cuda.synchronize() + for plan in compact_gqa._PLANS.values(): + plan.close() + + +def _check_native_compact_dispatch(lengths, spans): + """Compare one pack with stock and an independent reference.""" + if torch.cuda.get_device_capability() != (10, 7): + pytest.skip("Compact GQA requires SM107") + api = pytest.importorskip("cudnn.sdpa.bwd.compact_gqa") + from transformer_engine.pytorch.attention import packed_sequence as metadata + from transformer_engine.pytorch.attention.dot_product_attention import compact_gqa + from transformer_engine.pytorch.attention.dot_product_attention.backends import ( + FusedAttnBwdArgs, + _fused_attn_backward_impl, + ) + from transformer_engine.pytorch.cpp_extensions.fused_attn import ( + FusedAttnBackend, + fused_attn_fwd, + ) + + torch.manual_seed(2300) + logical = tuple(accumulate(lengths, initial=0)) + physical = tuple(accumulate(spans, initial=0)) + cu = torch.tensor(logical, device="cuda", dtype=torch.int32) + padded = torch.tensor(physical, device="cuda", dtype=torch.int32) + metadata.register_cu_seqlens(cu, logical) + metadata.register_cu_seqlens(padded, physical) + q = torch.randn(sum(spans), 8, 256, device="cuda", dtype=torch.bfloat16) + k = torch.randn(sum(spans), 1, 256, device="cuda", dtype=torch.bfloat16) + v, d_out = torch.randn_like(k), torch.randn_like(q) + backend = FusedAttnBackend["F16_arbitrary_seqlen"] + out, aux = fused_attn_fwd( + True, + max(lengths), + max(lengths), + cu, + cu, + q, + k, + v, + torch.bfloat16, + backend, + cu_seqlens_q_padded=padded, + cu_seqlens_kv_padded=padded, + attn_scale=1 / 16, + dropout=0.0, + qkv_layout="thd_thd_thd", + o_format="thd", + attn_mask_type="padding_causal", + ) + args = FusedAttnBwdArgs( + q=q, + k=k, + v=v, + out=out, + grad_output=d_out, + cu_seqlens_q=cu, + cu_seqlens_kv=cu, + cu_seqlens_q_padded=padded, + cu_seqlens_kv_padded=padded, + softmax_stats=aux[0], + rng_state=aux[1], + max_seqlen_q=max(lengths), + max_seqlen_kv=max(lengths), + attn_scale=1 / 16, + dropout_p=0.0, + fast_zero_fill=True, + qkv_layout="thd_thd_thd", + dqkv_layout="thd_thd_thd", + o_format="thd", + attn_bias_type="no_bias", + attn_mask_type="padding_causal", + softmax_type="vanilla", + window_size=(-1, -1), + fused_attention_backend=backend, + deterministic=False, + use_FAv2_bwd=False, + nominal_dtype=torch.bfloat16, + fp8=False, + is_input_fp8=False, + ) + with patch.object(compact_gqa, "_ENABLED", False): + expected = _fused_attn_backward_impl(args) + with metadata.attention_backend_workspace(): + args = replace(args, backend_workspace=metadata._get_attention_backend_workspace()) + with patch.object(compact_gqa, "_ENABLED", True), patch.object( + api, "compact_gqa_backward", wraps=api.compact_gqa_backward + ) as call: + actual = _fused_attn_backward_impl(args) + assert call.call_count == 1 + # Unregistered prefixes must dispatch to stock. + unregistered = replace(args, cu_seqlens_q=cu.clone(), cu_seqlens_kv=cu.clone()) + fallback = _fused_attn_backward_impl(unregistered) + assert call.call_count == 1 + for result, reference, stock in zip(actual[:3], expected[:3], fallback[:3]): + torch.testing.assert_close(result, reference, atol=0.025, rtol=0.025) + torch.testing.assert_close(stock, reference, atol=0.025, rtol=0.025) + if max(lengths) <= 1025: + # Check valid tokens against independent FP32 causal attention. + for start, length in zip(physical, lengths): + q_ref, k_ref, v_ref = [ + tensor[start : start + length].float().detach().requires_grad_() + for tensor in (q, k, v) + ] + scores = torch.einsum("thd,shd->hts", q_ref, k_ref.expand(-1, 8, -1)) / 16 + mask = torch.ones(length, length, device=q.device, dtype=torch.bool).tril() + probability = scores.masked_fill(~mask, float("-inf")).softmax(dim=-1) + output = torch.einsum("hts,shd->thd", probability, v_ref.expand(-1, 8, -1)) + gradients = torch.autograd.grad( + output, (q_ref, k_ref, v_ref), d_out[start : start + length].float() + ) + for result, reference in zip(actual[:3], gradients): + torch.testing.assert_close( + result[start : start + length].float(), reference, atol=0.025, rtol=0.025 + ) diff --git a/transformer_engine/pytorch/attention/__init__.py b/transformer_engine/pytorch/attention/__init__.py index e0c5608770..a64aed0a7e 100644 --- a/transformer_engine/pytorch/attention/__init__.py +++ b/transformer_engine/pytorch/attention/__init__.py @@ -5,6 +5,7 @@ """Python interface for attention""" from .dot_product_attention import DotProductAttention +from .packed_sequence import register_cu_seqlens, attention_backend_workspace from .linear_attention import ( GatedDeltaNetAttention, GatedDeltaNet2Attention, @@ -17,6 +18,8 @@ from .rope import RotaryPositionEmbedding __all__ = [ + "register_cu_seqlens", + "attention_backend_workspace", "DotProductAttention", "GatedDeltaNetAttention", "GatedDeltaNet2Attention", diff --git a/transformer_engine/pytorch/attention/dot_product_attention/backends.py b/transformer_engine/pytorch/attention/dot_product_attention/backends.py index c3503d0f1b..c657c59f7d 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/backends.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/backends.py @@ -101,6 +101,12 @@ from transformer_engine.pytorch.export import is_in_onnx_export_mode from transformer_engine.pytorch.graph import is_graph_capturing +from transformer_engine.pytorch.attention.packed_sequence import ( + _AttentionBackendWorkspace, + _get_attention_backend_workspace, +) +from .compact_gqa import compact_gqa_enabled, try_compact_gqa_backward + # Global vars for flash attn v2 flash_attn_cuda_bwd = None flash_attn_func = None @@ -1613,6 +1619,8 @@ class FusedAttnFwdArgs: class FusedAttnBwdArgs: """Single-argument bag for the backward path of :class:`FusedAttnFunc`.""" + backend_workspace: Optional[_AttentionBackendWorkspace] = None + # --- Saved / restored tensors (populated at backward entry) --- grad_output: Optional[torch.Tensor] = None q_fp8: Optional[TensorOrQuantized] = None @@ -2055,6 +2063,8 @@ def _fused_attn_setup_ctx( out = fwd_outputs[0] fp8 = ctx_attrs["fp8"] + if compact_gqa_enabled() and not torch.compiler.is_compiling() and not is_graph_capturing(): + bwd_args.backend_workspace = _get_attention_backend_workspace() bwd_args.fp8 = fp8 bwd_args.is_input_fp8 = ctx_attrs["is_input_fp8"] # assume fwd and bwd always use the same high precision, i.e. torch.float16 or torch.bfloat16 @@ -2189,6 +2199,9 @@ def _fused_attn_backward_impl( if not aux_ctx_tensors[0].is_contiguous(): aux_ctx_tensors[0] = aux_ctx_tensors[0].contiguous() + compact_grads = try_compact_gqa_backward(args, d_out, aux_ctx_tensors[0]) + if compact_grads is not None: + return (*compact_grads, None, None) rest = [None] if args.use_FAv2_bwd: softmax_lse, rng_state = aux_ctx_tensors diff --git a/transformer_engine/pytorch/attention/dot_product_attention/compact_gqa.py b/transformer_engine/pytorch/attention/dot_product_attention/compact_gqa.py new file mode 100644 index 0000000000..29722df334 --- /dev/null +++ b/transformer_engine/pytorch/attention/dot_product_attention/compact_gqa.py @@ -0,0 +1,172 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Opt-in cuDNN Frontend compact GQA backward dispatch.""" +import os +import weakref +from threading import RLock +from typing import TYPE_CHECKING, Optional + +import torch + +from transformer_engine.pytorch.attention.packed_sequence import ( + _AttentionBackendWorkspace, + _get_cu_seqlens, +) + +if TYPE_CHECKING: + from .backends import FusedAttnBwdArgs + +_ENABLED = os.getenv("NVTE_CUDNN_COMPACT_GQA_BWD", "0") == "1" +_PLANS = {} +_BOUND_WORKSPACES = {} +_PLAN_LOCK = RLock() + + +def compact_gqa_enabled() -> bool: + """Whether the optional backend was enabled before TE was imported.""" + return _ENABLED + + +def _packing(args: "FusedAttnBwdArgs"): + logical = _get_cu_seqlens(args.cu_seqlens_q) + if logical is None or logical != _get_cu_seqlens(args.cu_seqlens_kv): + return None + physical = [ + logical if tensor is None else _get_cu_seqlens(tensor) + for tensor in (args.cu_seqlens_q_padded, args.cu_seqlens_kv_padded) + ] + offsets = physical[0] + if ( + offsets is None + or offsets != physical[1] + or len(offsets) != len(logical) + or offsets[-1] != args.q.shape[0] + ): + return None + lengths = tuple(b - a for a, b in zip(logical, logical[1:])) + if any(length > b - a for length, a, b in zip(lengths, offsets, offsets[1:])): + return None + if max(lengths) > min(args.max_seqlen_q, args.max_seqlen_kv): + return None + return offsets, lengths + + +def _eligible(args: "FusedAttnBwdArgs", d_out: torch.Tensor, lse: torch.Tensor) -> bool: + q = args.q + if args.fp8 or args.is_input_fp8 or args.use_FAv2_bwd or args.deterministic: + return False + if args.qkv_layout != "thd_thd_thd" or args.dqkv_layout != "thd_thd_thd": + return False + if args.o_format != "thd" or args.nominal_dtype != torch.bfloat16: + return False + if args.dropout_p != 0 or args.attn_bias_type != "no_bias" or args.softmax_type != "vanilla": + return False + if args.attn_mask_type not in ("causal", "padding_causal") or args.attn_scale not in ( + None, + 1 / 16, + ): + return False + if args.window_size not in (None, (-1, -1), (-1, 0)): + return False + if not isinstance(q, torch.Tensor) or not q.is_cuda: + return False + if torch.cuda.get_device_capability(q.device) != (10, 7): + return False + if torch.cuda.is_current_stream_capturing() or torch.compiler.is_compiling(): + return False + tokens = q.shape[0] + if not tokens: + return False + for tensor, heads in ((q, 8), (args.k, 1), (args.v, 1), (args.out, 8), (d_out, 8)): + if not isinstance(tensor, torch.Tensor): + return False + if ( + tensor.device != q.device + or tensor.dtype != torch.bfloat16 + or tuple(tensor.shape) != (tokens, heads, 256) + or not tensor.is_contiguous() + or tensor.data_ptr() % 16 + ): + return False + return ( + isinstance(lse, torch.Tensor) + and lse.device == q.device + and lse.dtype == torch.float32 + and lse.is_contiguous() + and tuple(lse.shape) in ((tokens, 8), (tokens, 8, 1)) + ) + + +def try_compact_gqa_backward( + args: "FusedAttnBwdArgs", d_out: torch.Tensor, lse: torch.Tensor +) -> Optional[tuple[torch.Tensor, torch.Tensor, torch.Tensor]]: + """Return gradients, or None when stock backward should handle this call.""" + if not _ENABLED or not _eligible(args, d_out, lse): + return None + packing = _packing(args) + if packing is None: + return None + # Experimental dependency is loaded only for an eligible opt-in call. + try: + from cudnn.sdpa.bwd.compact_gqa import CompactGqaBackward, compact_gqa_backward + except ImportError as error: + raise RuntimeError( + "NVTE_CUDNN_COMPACT_GQA_BWD requires the cuDNN Frontend compact GQA API." + ) from error + scope = args.backend_workspace + if scope is None: + scope = _AttentionBackendWorkspace() + offsets, lengths = packing + capacity = max(1, *lengths) + stream = torch.cuda.current_stream(args.q.device).cuda_stream + # cuDNN compiles symbolic lengths; initialize_workspace updates scratch capacity. + key = (args.q.device, stream, lse.ndim) + # Forward captures the owner so autograd threads do not depend on ContextVar propagation. + with scope.lock: + if not scope.active: + return None + with _PLAN_LOCK: + if key not in _PLANS: + plan = CompactGqaBackward( + args.q, + args.k, + args.v, + args.out, + d_out, + lse, + max_seqlen=capacity, + query_rows=32768, + groups=4, + fast_store=True, + ) + plan.compile(stream) + _PLANS[key] = plan + plan = _PLANS[key] + buffer_key = ("compact_gqa", *key) + buffer, allocated = scope.buffers.get(buffer_key, (None, 0)) + if allocated < capacity: + buffer = torch.empty( + plan.scratch_workspace_bytes(capacity), dtype=torch.uint8, device=args.q.device + ) + allocated = capacity + scope.buffers[buffer_key] = (buffer, allocated) + bound = _BOUND_WORKSPACES.get(key) + if bound is None or bound() is not buffer: + plan.initialize_workspace(buffer, stream, max_seqlen=allocated) + _BOUND_WORKSPACES[key] = weakref.ref(buffer) + result = compact_gqa_backward( + args.q, + args.k, + args.v, + args.out, + d_out, + lse, + plan=plan, + workspace=buffer, + current_stream=stream, + sequence_offsets=offsets, + sequence_lengths=lengths, + ) + return result["dq"], result["dk"], result["dv"] diff --git a/transformer_engine/pytorch/attention/packed_sequence.py b/transformer_engine/pytorch/attention/packed_sequence.py new file mode 100644 index 0000000000..6470f64a74 --- /dev/null +++ b/transformer_engine/pytorch/attention/packed_sequence.py @@ -0,0 +1,105 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Host metadata and scoped scratch for packed attention.""" +from contextlib import contextmanager +from contextvars import ContextVar +from dataclasses import dataclass, field +from numbers import Integral +from threading import RLock +from typing import Any, Iterator, Optional, Sequence +import weakref + +import torch + +_PREFIXES = {} + + +def register_cu_seqlens(tensor: torch.Tensor, offsets: Sequence[int]) -> None: + """Associate a prefix tensor with the host integers used to create it. + + This is a producer contract: values are not read back from the tensor. + Call after its final write. Mutation, resizing, and tensor destruction + invalidate the registration. No model, head count, or length is assumed. + """ + if torch.is_tensor(offsets) or any(not isinstance(n, Integral) for n in offsets): + raise TypeError("Offsets must be a sequence of host integers.") + values = tuple(int(n) for n in offsets) + if len(values) < 1 or values[0] != 0 or any(b < a for a, b in zip(values, values[1:])): + raise ValueError("Offsets must start at zero and be nondecreasing.") + if ( + tensor.dtype != torch.int32 + or tensor.ndim != 1 + or not tensor.is_contiguous() + or tensor.numel() != len(values) + or values[-1] > torch.iinfo(torch.int32).max + ): + raise ValueError("Expected a matching contiguous int32 prefix tensor.") + if tensor.is_inference(): + return + key = (tensor.device, tensor.data_ptr()) + + def expired(ref): + if _PREFIXES.get(key, (None,))[0] is ref: + _PREFIXES.pop(key, None) + + _PREFIXES[key] = (weakref.ref(tensor, expired), tensor._version, values) + + +def _get_cu_seqlens(tensor: Optional[torch.Tensor]) -> Optional[tuple[int, ...]]: + if tensor is None or tensor.dtype != torch.int32 or tensor.ndim != 1: + return None + if not tensor.is_contiguous() or tensor.is_inference(): + return None + entry = _PREFIXES.get((tensor.device, tensor.data_ptr())) + if entry is None: + return None + owner, version, values = entry + owner = owner() + if ( + owner is not None + and owner._version == version == tensor._version + and owner.data_ptr() == tensor.data_ptr() + and owner.numel() == tensor.numel() == len(values) + ): + return values + return None + + +@dataclass +class _AttentionBackendWorkspace: + buffers: dict[Any, Any] = field(default_factory=dict) + lock: Any = field(default_factory=RLock) + active: bool = True + + +_WORKSPACE: ContextVar[Optional[_AttentionBackendWorkspace]] = ContextVar( + "te_attention_backend_workspace", default=None +) + + +def _get_attention_backend_workspace() -> Optional[_AttentionBackendWorkspace]: + return _WORKSPACE.get() + + +@contextmanager +def attention_backend_workspace() -> Iterator[None]: + """Reuse optional attention scratch until this scope exits. + + Enclose forward and backward. Nested scopes share the outer owner. + Scratch references are released on success or exception; cached compiled + plans remain reusable. This does not enable any optional backend. + """ + if _WORKSPACE.get() is not None: + yield + return + owner = _AttentionBackendWorkspace() + token = _WORKSPACE.set(owner) + try: + yield + finally: + with owner.lock: + owner.active = False + owner.buffers.clear() + _WORKSPACE.reset(token)