diff --git a/build_tools/pytorch.py b/build_tools/pytorch.py index 98331ccbd8..918b5d941c 100644 --- a/build_tools/pytorch.py +++ b/build_tools/pytorch.py @@ -31,7 +31,7 @@ def install_requirements() -> List[str]: "packaging", "pydantic", "nvdlfw-inspect", - "nvidia-cudnn-frontend>=1.25.0", + "nvidia-cudnn-frontend>=1.28.0", ] diff --git a/pyproject.toml b/pyproject.toml index 2c9f224c14..efb6f40f5f 100755 --- a/pyproject.toml +++ b/pyproject.toml @@ -3,7 +3,7 @@ # See LICENSE for license information. [build-system] -requires = ["setuptools>=61.0", "cmake>=3.21", "wheel", "pybind11[global]", "ninja", "pip", "torch>=2.1", "jax>=0.5.0", "flax>=0.7.1", "nvidia-cudnn-frontend>=1.25.0"] +requires = ["setuptools>=61.0", "cmake>=3.21", "wheel", "pybind11[global]", "ninja", "pip", "torch>=2.1", "jax>=0.5.0", "flax>=0.7.1", "nvidia-cudnn-frontend>=1.28.0"] # Use legacy backend to import local packages in setup.py build-backend = "setuptools.build_meta:__legacy__" diff --git a/qa/L0_pytorch_unittest/test.sh b/qa/L0_pytorch_unittest/test.sh index 14a5f4fe3d..81acd07bee 100644 --- a/qa/L0_pytorch_unittest/test.sh +++ b/qa/L0_pytorch_unittest/test.sh @@ -63,6 +63,7 @@ python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_hybrid_quantizat python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_identity_quantizer.xml $TE_PATH/tests/pytorch/test_identity_quantizer.py || test_fail "test_identity_quantizer.py" 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_flex_attention.xml $TE_PATH/tests/pytorch/attention/test_flex_attention.py || test_fail "test_flex_attention.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_ALLOW_UNSAFE_PICKLE_EXTRA_STATE=1 NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_attention_deterministic.xml $TE_PATH/tests/pytorch/attention/test_attention.py || test_fail "NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 test_attention.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_linear_mxfp8_attention.xml $TE_PATH/tests/pytorch/attention/test_linear_mxfp8_attention.py || test_fail "test_linear_mxfp8_attention.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_fused_mla_q_uproj.xml $TE_PATH/tests/pytorch/attention/test_fused_mla_q_uproj.py || test_fail "test_fused_mla_q_uproj.py" diff --git a/tests/pytorch/attention/test_gdn_attention.py b/tests/pytorch/attention/test_gdn_attention.py new file mode 100644 index 0000000000..9db14364a7 --- /dev/null +++ b/tests/pytorch/attention/test_gdn_attention.py @@ -0,0 +1,499 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Tests for Gated DeltaNet through the DotProductAttention API.""" + +import importlib.util +import math +import os +from typing import Optional, Tuple + +import pytest +import torch +import torch.nn.functional as F + +from transformer_engine.pytorch import DotProductAttention, autocast, is_fp8_available + + +def _gdn_available() -> bool: + if not torch.cuda.is_available(): + return False + try: + from cudnn.linear_attention.ops import gated_delta_net # noqa: F401 + except (AttributeError, ImportError): + return False + try: + has_cutlass = importlib.util.find_spec("cutlass") is not None + has_cu_tile = importlib.util.find_spec("cuda.tile") is not None + except (ImportError, ModuleNotFoundError): + return False + return has_cutlass or has_cu_tile + + +_GDN_AVAILABLE = _gdn_available() +if os.getenv("NVTE_GDN_TEST_REQUIRED", "0") == "1" and not _GDN_AVAILABLE: + raise RuntimeError( + "NVTE_GDN_TEST_REQUIRED=1, but the cuDNN frontend GDN op or its " + "cutedsl runtime is unavailable." + ) + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +requires_gdn = pytest.mark.skipif( + not _GDN_AVAILABLE, reason="GDN requires a supported cuDNN frontend kernel runtime" +) + +_FWD_TOL = {torch.bfloat16: 2e-2, torch.float16: 1e-2} +_STATE_TOL = {torch.bfloat16: 2e-2, torch.float16: 1e-2} +_BWD_TOL = {torch.bfloat16: 4e-2, torch.float16: 2e-2} + + +def _rms_ratio(actual: torch.Tensor, expected: torch.Tensor) -> float: + """Return relative RMS error, which is appropriate for bfloat16 results.""" + actual = actual.detach().double() + expected = expected.detach().double() + return ( + (actual - expected).square().mean().sqrt() + / expected.square().mean().sqrt().clamp_min(1e-12) + ).item() + + +def _assert_rms_close( + actual: torch.Tensor, expected: torch.Tensor, tolerance: float, name: str +) -> None: + ratio = _rms_ratio(actual, expected) + assert ratio < tolerance, f"{name} RMS ratio {ratio:.4g} >= {tolerance}" + + +def _gdn_recurrence( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + alpha: torch.Tensor, + beta: torch.Tensor, + state: torch.Tensor, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Evaluate GDN for tensors in [batch, heads, sequence, ...] layout. + + ``state`` follows the cuDNN frontend's [..., v_head_dim, qk_head_dim] convention. + """ + outputs = [] + for token_idx in range(q.shape[2]): + q_t = q[:, :, token_idx] + k_t = k[:, :, token_idx] + v_t = v[:, :, token_idx] + alpha_t = alpha[:, :, token_idx] + beta_t = beta[:, :, token_idx] + + prediction = torch.matmul(state, k_t.unsqueeze(-1)).squeeze(-1) + residual = v_t - alpha_t.unsqueeze(-1) * prediction + state = alpha_t[..., None, None] * state + beta_t[..., None, None] * ( + residual.unsqueeze(-1) @ k_t.unsqueeze(-2) + ) + outputs.append(torch.matmul(state, q_t.unsqueeze(-1)).squeeze(-1)) + + return torch.stack(outputs, dim=2), state + + +def _gdn_reference( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + *, + scale: Optional[float] = None, + initial_state: Optional[torch.Tensor] = None, + cu_seqlens: Optional[torch.Tensor] = None, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Pure PyTorch FP64 implementation of the Gated DeltaNet recurrence. + + Inputs use [batch, sequence, heads, dimension] layout. Packed inputs use a + singleton batch dimension and provide sequence boundaries via cu_seqlens. + ``initial_state``/the returned final state use the cuDNN frontend's + [batch, heads, v_head_dim, qk_head_dim] convention throughout. + """ + qk_dim = q.shape[-1] + if scale is None: + scale = 1.0 / math.sqrt(qk_dim) + + output_heads = max(q.shape[2], v.shape[2]) + + def _expand_heads(tensor: torch.Tensor) -> torch.Tensor: + repeats = output_heads // tensor.shape[2] + return tensor.repeat_interleave(repeats, dim=2) if repeats > 1 else tensor + + q_ref = _expand_heads(q.double() * scale) + k_ref = _expand_heads(k.double()) + v_ref = _expand_heads(v.double()) + alpha_ref = _expand_heads(g.double().exp()) + beta_ref = _expand_heads(beta.double()) + + # [batch, sequence, heads, ...] -> [batch, heads, sequence, ...] + q_ref = q_ref.permute(0, 2, 1, 3) + k_ref = k_ref.permute(0, 2, 1, 3) + v_ref = v_ref.permute(0, 2, 1, 3) + alpha_ref = alpha_ref.permute(0, 2, 1) + beta_ref = beta_ref.permute(0, 2, 1) + + def _initial_state(sequence_idx: Optional[int] = None) -> torch.Tensor: + if initial_state is None: + return torch.zeros( + 1 if sequence_idx is not None else q.shape[0], + output_heads, + v.shape[-1], + qk_dim, + dtype=torch.float64, + device=q.device, + ) + if sequence_idx is None: + return initial_state.double() + return initial_state[sequence_idx : sequence_idx + 1].double() + + if cu_seqlens is None: + output, final_state = _gdn_recurrence( + q_ref, k_ref, v_ref, alpha_ref, beta_ref, _initial_state() + ) + return output.permute(0, 2, 1, 3), final_state + + assert q.shape[0] == 1 + bounds = cu_seqlens.tolist() + outputs = [] + final_states = [] + for sequence_idx, (start, end) in enumerate(zip(bounds[:-1], bounds[1:])): + state = _initial_state(sequence_idx) + if start == end: + final_states.append(state) + continue + output, state = _gdn_recurrence( + q_ref[:, :, start:end], + k_ref[:, :, start:end], + v_ref[:, :, start:end], + alpha_ref[:, :, start:end], + beta_ref[:, :, start:end], + state, + ) + outputs.append(output) + final_states.append(state) + + if outputs: + packed_output = torch.cat(outputs, dim=2).permute(0, 2, 1, 3) + else: + packed_output = q_ref.new_zeros(1, 0, output_heads, v.shape[-1]) + return packed_output, torch.cat(final_states, dim=0) + + +def _inputs( + batch, + sequence, + q_heads, + v_heads, + qk_dim=64, + v_dim=64, + dtype=torch.bfloat16, +): + torch.manual_seed(1234) + q = torch.randn(batch, sequence, q_heads, qk_dim, device="cuda", dtype=dtype) + k = F.normalize(torch.randn_like(q, dtype=torch.float32), dim=-1).to(dtype) + v = torch.randn(batch, sequence, v_heads, v_dim, device="cuda", dtype=dtype) + output_heads = max(q_heads, v_heads) + g = torch.rand(batch, sequence, output_heads, device="cuda", dtype=torch.float32).log() + beta = torch.rand(batch, sequence, output_heads, device="cuda", dtype=torch.float32) + return q, k, v, g, beta + + +@requires_gdn +@pytest.mark.parametrize("checkpoint_core_attention", [False, True], ids=["eager", "checkpoint"]) +@pytest.mark.parametrize( + ("qk_dim", "v_dim"), + [(64, 64), (128, 128), (64, 128)], + ids=["qk64_v64_cutile", "qk128_v128_frost", "qk64_v128"], +) +@pytest.mark.parametrize("use_qk_l2norm_in_kernel", [False, True], ids=["no_l2norm", "l2norm"]) +def test_gdn_thd_forward_final_state_and_backward( + checkpoint_core_attention, qk_dim, v_dim, use_qk_l2norm_in_kernel +): + """THD GDN matches a PyTorch recurrence in forward and backward. + + head_dim=128 exercises the FROST engine, while head_dim=64 falls back to cuTile. + """ + batch, sequence, heads = 2, 128, 2 + q, k, v, g, beta = ( + tensor.reshape(batch * sequence, *tensor.shape[2:]).requires_grad_(True) + for tensor in _inputs(batch, sequence, heads, heads, qk_dim, v_dim) + ) + cu_seqlens = torch.arange(batch + 1, device="cuda", dtype=torch.int32) * sequence + initial_state = ( + torch.randn(batch, heads, v_dim, qk_dim, device="cuda", dtype=torch.float32) * 0.05 + ).requires_grad_() + + attention = DotProductAttention( + num_attention_heads=heads, + kv_channels=(qk_dim, v_dim), + qkv_format="thd", + attn_mask_type="padding_causal", + ) + output, final_state = attention( + q, + k, + v, + cu_seqlens_q=cu_seqlens, + g=g, + beta=beta, + initial_state=initial_state, + output_final_state=True, + checkpoint_core_attention=checkpoint_core_attention, + use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel, + ) + reference_inputs = { + name: tensor.detach().double().reshape(1, -1, *tensor.shape[1:]).requires_grad_() + for name, tensor in (("q", q), ("k", k), ("v", v), ("g", g), ("beta", beta)) + } + initial_state_ref = initial_state.detach().double().requires_grad_() + reference_q = reference_inputs["q"] + reference_k = reference_inputs["k"] + if use_qk_l2norm_in_kernel: + reference_q = F.normalize(reference_q, dim=-1) + reference_k = F.normalize(reference_k, dim=-1) + output_ref, final_state_ref = _gdn_reference( + reference_q, + reference_k, + reference_inputs["v"], + reference_inputs["g"], + reference_inputs["beta"], + initial_state=initial_state_ref, + cu_seqlens=cu_seqlens, + ) + output_ref = output_ref.squeeze(0).flatten(-2) + + _assert_rms_close(output, output_ref, _FWD_TOL[q.dtype], "output") + _assert_rms_close(final_state, final_state_ref, _STATE_TOL[q.dtype], "final state") + + output_weight = torch.randn_like(output, dtype=torch.float32) + state_weight = torch.randn_like(final_state, dtype=torch.float32) + ((output.float() * output_weight).sum() + (final_state.float() * state_weight).sum()).backward() + ( + (output_ref * output_weight.double()).sum() + + (final_state_ref * state_weight.double()).sum() + ).backward() + + for name, tensor in (("q", q), ("k", k), ("v", v), ("g", g), ("beta", beta)): + assert tensor.grad is not None, f"no gradient for {name}" + assert torch.isfinite(tensor.grad).all(), f"non-finite gradient for {name}" + reference_grad = reference_inputs[name].grad.reshape_as(tensor) + _assert_rms_close(tensor.grad, reference_grad, _BWD_TOL[q.dtype], f"d{name}") + assert initial_state.grad is not None, "no gradient for initial_state" + _assert_rms_close( + initial_state.grad, + initial_state_ref.grad, + _BWD_TOL[q.dtype], + "dinitial_state", + ) + + +@pytest.mark.parametrize("qkv_format", ["bshd", "sbhd"]) +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16], ids=["bf16", "fp16"]) +@requires_gdn +def test_gdn_dense_layout(qkv_format, dtype): + """Dense TE layouts and supported input dtypes match a PyTorch reference.""" + batch, sequence, heads, dim = 1, 128, 2, 64 + q, k, v, g, beta = _inputs(batch, sequence, heads, heads, dim, dim, dtype) + with torch.no_grad(): + output_ref, _ = _gdn_reference( + F.normalize(q.float(), dim=-1), + F.normalize(k.float(), dim=-1), + v, + g, + beta, + ) + + if qkv_format == "sbhd": + q, k, v, g, beta = (tensor.transpose(0, 1).contiguous() for tensor in (q, k, v, g, beta)) + expected = output_ref.reshape(batch, sequence, -1).transpose(0, 1).contiguous() + else: + expected = output_ref.reshape(batch, sequence, -1) + + attention = DotProductAttention( + num_attention_heads=heads, + kv_channels=dim, + qkv_format=qkv_format, + attn_mask_type="causal", + ) + output = attention(q, k, v, g=g, beta=beta, use_qk_l2norm_in_kernel=True) + assert output.shape == expected.shape + _assert_rms_close(output, expected, _FWD_TOL[dtype], "output") + + +def test_gdn_rejects_value_head_count_that_changes_output_width(): + """DPA's configured output width cannot be changed by the runtime V tensor.""" + q, k, v, g, beta = _inputs(1, 128, 1, 2) + attention = DotProductAttention( + num_attention_heads=1, + kv_channels=64, + qkv_format="bshd", + attn_mask_type="causal", + ) + with pytest.raises(ValueError, match="GDN V must have 1 heads"): + attention(q, k, v, g=g, beta=beta) + + +@requires_gdn +def test_gdn_state_round_trip_matches_single_shot(): + """A final state can seed the next chunk without changing the result.""" + batch, sequence, heads, dim = 2, 128, 2, 64 + split = 64 + q, k, v, g, beta = _inputs(batch, sequence, heads, heads, dim, dim) + attention = DotProductAttention( + num_attention_heads=heads, + kv_channels=dim, + qkv_format="bshd", + attn_mask_type="causal", + ) + + with torch.no_grad(): + full_output, full_state = attention(q, k, v, g=g, beta=beta, output_final_state=True) + first_output, first_state = attention( + q[:, :split], + k[:, :split], + v[:, :split], + g=g[:, :split], + beta=beta[:, :split], + output_final_state=True, + ) + second_output, chunked_state = attention( + q[:, split:], + k[:, split:], + v[:, split:], + g=g[:, split:], + beta=beta[:, split:], + initial_state=first_state, + output_final_state=True, + ) + + chunked_output = torch.cat((first_output, second_output), dim=1) + _assert_rms_close(chunked_output, full_output, _FWD_TOL[q.dtype], "chunked output") + _assert_rms_close(chunked_state, full_state, _STATE_TOL[q.dtype], "chunked state") + + +@requires_gdn +@pytest.mark.parametrize( + "bounds", + [(0, 48, 160), (0, 0, 64, 160)], + ids=["unequal", "leading-empty"], +) +def test_gdn_thd_ragged_sequences(bounds): + """Packed GDN handles unequal lengths and empty sequences.""" + total_tokens, heads, dim = bounds[-1], 2, 64 + q, k, v, g, beta = ( + tensor.squeeze(0) for tensor in _inputs(1, total_tokens, heads, heads, dim, dim) + ) + cu_seqlens = torch.tensor(bounds, device="cuda", dtype=torch.int32) + initial_state = torch.randn( + len(bounds) - 1, + heads, + dim, + dim, + device="cuda", + dtype=torch.float32, + ) + attention = DotProductAttention( + num_attention_heads=heads, + kv_channels=dim, + qkv_format="thd", + attn_mask_type="padding_causal", + ) + + with torch.no_grad(): + output, final_state = attention( + q, + k, + v, + cu_seqlens_q=cu_seqlens, + g=g, + beta=beta, + initial_state=initial_state, + output_final_state=True, + ) + output_ref, final_state_ref = _gdn_reference( + q.unsqueeze(0), + k.unsqueeze(0), + v.unsqueeze(0), + g.unsqueeze(0), + beta.unsqueeze(0), + initial_state=initial_state, + cu_seqlens=cu_seqlens, + ) + + _assert_rms_close( + output, + output_ref.squeeze(0).flatten(-2), + _FWD_TOL[q.dtype], + "ragged output", + ) + _assert_rms_close( + final_state, + final_state_ref, + _STATE_TOL[q.dtype], + "ragged final state", + ) + + +def test_gdn_requires_both_gates(): + """A partial GDN invocation fails before entering a softmax-attention backend.""" + q, k, v, g, _ = _inputs(1, 128, 1, 1) + attention = DotProductAttention( + num_attention_heads=1, + kv_channels=64, + qkv_format="bshd", + attn_mask_type="causal", + ) + with pytest.raises(ValueError, match="require both g and beta"): + attention(q, k, v, g=g) + + +@pytest.mark.skipif(not is_fp8_available(), reason="FP8 is not available") +def test_gdn_rejects_fp8_autocast(): + """GDN must not silently run in high precision inside FP8 autocast.""" + q, k, v, g, beta = _inputs(1, 128, 1, 1) + attention = DotProductAttention( + num_attention_heads=1, + kv_channels=64, + qkv_format="bshd", + attn_mask_type="causal", + ) + with autocast(enabled=True), pytest.raises(ValueError, match="does not support FP8 autocast"): + attention(q, k, v, g=g, beta=beta) + + +def test_gdn_runs_te_forward_lifecycle(monkeypatch): + """GDN calls pair prepare_forward with end_forward even without the kernel runtime.""" + q, k, v, g, beta = _inputs(1, 128, 1, 1) + attention = DotProductAttention( + num_attention_heads=1, + kv_channels=64, + qkv_format="bshd", + attn_mask_type="causal", + ) + events = [] + prepare_forward = attention.prepare_forward + end_forward = attention.end_forward + + def traced_prepare_forward(*args, **kwargs): + events.append("prepare") + return prepare_forward(*args, **kwargs) + + def traced_end_forward(): + events.append("end") + return end_forward() + + def fake_gdn_forward(query, key, value, gate_g, gate_beta, initial_state, **kwargs): + del key, gate_g, gate_beta, initial_state, kwargs + return value.reshape(*query.shape[:-2], -1) + + monkeypatch.setattr(attention, "prepare_forward", traced_prepare_forward) + monkeypatch.setattr(attention, "end_forward", traced_end_forward) + monkeypatch.setattr(attention.gdn_attention, "forward", fake_gdn_forward) + + output = attention(q, k, v, g=g, beta=beta) + assert output.shape == (1, 128, 64) + assert events == ["prepare", "end"] 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 1e5f5552d4..d05bba9e79 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 @@ -66,6 +66,7 @@ FusedAttention, FlashAttention, ) +from transformer_engine.pytorch.attention.dot_product_attention.gdn import _GatedDeltaNetAttention # Setup Attention Logging @@ -824,6 +825,13 @@ def __init__( return_max_logit=self.return_max_logit, ) + self.gdn_attention = _GatedDeltaNetAttention( + softmax_scale, + num_attention_heads // self.tp_size, + self.hidden_size_per_attention_head_k, + self.hidden_size_per_attention_head_v, + ) + def remove_extra_states_check(self, incompatible_keys): # pylint: disable=unused-argument """ Temporarily remove core_attention._extra_state as a missing key @@ -863,7 +871,7 @@ def _checkpointed_attention_forward( attention_func: Callable, *forward_args: Tuple[torch.Tensor, ...], **forward_kwargs: Dict[str, Any], - ) -> torch.Tensor: + ) -> Union[torch.Tensor, Tuple[torch.Tensor, ...]]: """Forward method with activation checkpointing.""" def custom_forward(*input_args, **input_kwargs): @@ -1391,6 +1399,218 @@ def get_quantizer_roles( ] return base[:num_quantizers] + def _validate_gdn_request( + self, + g: Optional[torch.Tensor], + beta: Optional[torch.Tensor], + *, + attention_mask: Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]], + cu_seqlens_kv: Optional[torch.Tensor], + cu_seqlens_q_padded: Optional[torch.Tensor], + cu_seqlens_kv_padded: Optional[torch.Tensor], + max_seqlen_q: Optional[int], + max_seqlen_kv: Optional[int], + attn_mask_type: Optional[str], + window_size: Optional[Tuple[int, int]], + bottom_right_diagonal: Optional[bool], + core_attention_bias_type: str, + core_attention_bias: Optional[torch.Tensor], + alibi_slopes: Optional[torch.Tensor], + fast_zero_fill: bool, + inference_params: Optional[InferenceParams], + pad_between_seqs: Optional[bool], + fp8_output: Optional[bool], + bf16_backward: Optional[bool], + num_splits: Optional[int], + score_mod: Optional[Callable], + score_mod_bprop: Optional[Callable], + score_mod_tensors: Optional[Dict[str, torch.Tensor]], + score_mod_bprop_tensors: Optional[Dict[str, torch.Tensor]], + ) -> Tuple[str, Tuple[int, int], bool]: + """Reject DPA options and configurations that GDN does not support. + + Returns the normalized (attn_mask_type, window_size, bottom_right_diagonal), + since these may be defaulted from module attributes. + """ + if g is None or beta is None: + raise ValueError( + "GDN-specific arguments require both g and beta; " + f"got g={'set' if g is not None else 'None'} and " + f"beta={'set' if beta is not None else 'None'}." + ) + if FP8GlobalStateManager.is_fp8_enabled() or FP8GlobalStateManager.is_fp8_calibration(): + raise ValueError("GDN attention does not support FP8 autocast or FP8 calibration.") + if self.attention_type != "self": + raise ValueError("GDN attention only supports attention_type='self'.") + if self.attention_dropout != 0.0: + raise ValueError("GDN attention does not support attention dropout.") + if self.softmax_type != "vanilla": + raise ValueError("GDN attention does not support sink-attention softmax types.") + if self.return_max_logit: + raise ValueError("GDN attention does not support return_max_logit.") + if self.cp_group is not None: + raise ValueError("GDN attention does not support context parallelism.") + + if attn_mask_type is None: + attn_mask_type = self.attn_mask_type + else: + attn_mask_type = attn_mask_type.replace(",", "_") + if attn_mask_type == "causal_padding": + attn_mask_type = "padding_causal" + if attn_mask_type not in {"causal", "padding_causal"}: + raise ValueError( + "GDN is inherently causal and only supports attn_mask_type='causal' or " + f"'padding_causal', got {attn_mask_type!r}." + ) + if attention_mask is not None: + raise ValueError( + "GDN does not accept attention_mask; use cu_seqlens_q with qkv_format='thd' " + "for variable-length sequences." + ) + if cu_seqlens_kv is not None: + raise ValueError( + "GDN is self-attention over fully packed sequences and derives sequence " + "boundaries from cu_seqlens_q alone; pass only cu_seqlens_q, not " + "cu_seqlens_kv." + ) + if window_size is None: + window_size = self.window_size + window_size = dpa_utils.check_set_window_size(attn_mask_type, window_size) + if window_size != (-1, 0): + raise ValueError("GDN does not support sliding-window attention.") + if bottom_right_diagonal is None: + bottom_right_diagonal = self.bottom_right_diagonal + if bottom_right_diagonal: + raise ValueError("GDN does not support bottom-right causal alignment.") + + if cu_seqlens_q_padded is not None or cu_seqlens_kv_padded is not None: + raise ValueError("GDN does not support padded cumulative sequence lengths.") + if max_seqlen_q is not None or max_seqlen_kv is not None: + raise ValueError("GDN does not use max_seqlen_q or max_seqlen_kv.") + if core_attention_bias_type != "no_bias" or core_attention_bias is not None: + raise ValueError("GDN does not support core attention bias.") + if alibi_slopes is not None: + raise ValueError("GDN does not support ALiBi.") + if inference_params is not None: + raise ValueError( + "GDN does not use the KV cache. Pass initial_state and set " + "output_final_state=True for recurrent inference." + ) + if pad_between_seqs: + raise ValueError("GDN does not support pad_between_seqs=True.") + if not fast_zero_fill: + raise ValueError("GDN does not support fast_zero_fill=False.") + if fp8_output: + raise ValueError("GDN does not support FP8 output.") + if bf16_backward: + raise ValueError("GDN does not support bf16_backward.") + if num_splits not in {None, 1}: + raise ValueError("GDN does not support num_splits.") + if any( + value is not None + for value in (score_mod, score_mod_bprop, score_mod_tensors, score_mod_bprop_tensors) + ): + raise ValueError("GDN does not support Flex Attention score modifications.") + + return attn_mask_type, window_size, bottom_right_diagonal + + def _forward_gdn( + self, + query_layer: torch.Tensor, + key_layer: torch.Tensor, + value_layer: torch.Tensor, + g: Optional[torch.Tensor], + beta: Optional[torch.Tensor], + *, + qkv_format: str, + attention_mask: Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]], + cu_seqlens_q: Optional[torch.Tensor], + cu_seqlens_kv: Optional[torch.Tensor], + cu_seqlens_q_padded: Optional[torch.Tensor], + cu_seqlens_kv_padded: Optional[torch.Tensor], + max_seqlen_q: Optional[int], + max_seqlen_kv: Optional[int], + attn_mask_type: Optional[str], + window_size: Optional[Tuple[int, int]], + bottom_right_diagonal: Optional[bool], + checkpoint_core_attention: bool, + core_attention_bias_type: str, + core_attention_bias: Optional[torch.Tensor], + alibi_slopes: Optional[torch.Tensor], + fast_zero_fill: bool, + inference_params: Optional[InferenceParams], + pad_between_seqs: Optional[bool], + fp8_output: Optional[bool], + bf16_backward: Optional[bool], + num_splits: Optional[int], + score_mod: Optional[Callable], + score_mod_bprop: Optional[Callable], + score_mod_tensors: Optional[Dict[str, torch.Tensor]], + score_mod_bprop_tensors: Optional[Dict[str, torch.Tensor]], + initial_state: Optional[torch.Tensor], + output_final_state: bool, + use_qk_l2norm_in_kernel: bool, + ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]: + """Validate GDN-specific constraints and execute the cuDNN frontend op.""" + attn_mask_type, window_size, bottom_right_diagonal = self._validate_gdn_request( + g, + beta, + attention_mask=attention_mask, + cu_seqlens_kv=cu_seqlens_kv, + cu_seqlens_q_padded=cu_seqlens_q_padded, + cu_seqlens_kv_padded=cu_seqlens_kv_padded, + max_seqlen_q=max_seqlen_q, + max_seqlen_kv=max_seqlen_kv, + attn_mask_type=attn_mask_type, + window_size=window_size, + bottom_right_diagonal=bottom_right_diagonal, + core_attention_bias_type=core_attention_bias_type, + core_attention_bias=core_attention_bias, + alibi_slopes=alibi_slopes, + fast_zero_fill=fast_zero_fill, + inference_params=inference_params, + pad_between_seqs=pad_between_seqs, + fp8_output=fp8_output, + bf16_backward=bf16_backward, + num_splits=num_splits, + score_mod=score_mod, + score_mod_bprop=score_mod_bprop, + score_mod_tensors=score_mod_tensors, + score_mod_bprop_tensors=score_mod_bprop_tensors, + ) + + gdn_kwargs = { + "qkv_format": qkv_format, + "cu_seqlens_q": cu_seqlens_q, + "output_final_state": output_final_state, + "use_qk_l2norm_in_kernel": use_qk_l2norm_in_kernel, + } + with self.prepare_forward_ctx( + query_layer, + num_gemms=3, + allow_non_contiguous=True, + ) as query_layer: + if checkpoint_core_attention: + return self._checkpointed_attention_forward( + self.gdn_attention, + query_layer, + key_layer, + value_layer, + g, + beta, + initial_state, + **gdn_kwargs, + ) + return self.gdn_attention( + query_layer, + key_layer, + value_layer, + g, + beta, + initial_state, + **gdn_kwargs, + ) + @no_torch_dynamo(when=_needs_eager_dpa) def forward( self, @@ -1425,7 +1645,12 @@ def forward( qkv_layer: Optional[torch.Tensor] = None, kv_layer: Optional[torch.Tensor] = None, qkv_interleave_dim: int = -3, - ) -> torch.Tensor: + g: Optional[torch.Tensor] = None, + beta: Optional[torch.Tensor] = None, + initial_state: Optional[torch.Tensor] = None, + output_final_state: bool = False, + use_qk_l2norm_in_kernel: bool = False, + ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]: r""" Dot Product Attention Layer. @@ -1436,14 +1661,15 @@ def forward( .. note:: - DotProductAttention supports three backends: 1) FlashAttention which calls + Scaled-softmax attention supports three backends: 1) FlashAttention which calls HazyResearch/Dao-AILab's `flash-attn `_ PyTorch API, 2) FusedAttention which has multiple fused attention implementations based on `cuDNN Graph API `_ (see :attr:`FusedAttention` for more details on FusedAttention backends), and 3) UnfusedDotProductAttention which is the native PyTorch implementation - with fused scaled masked softmax. + with fused scaled masked softmax. GDN requests use the cuDNN frontend linear-attention + custom op directly. .. note:: @@ -1521,6 +1747,13 @@ def forward( :attr:`score_mod_bprop_tensors` are experimental cuDNN frontend Flex Attention APIs. Their callback signatures and supported configurations may change. + .. note:: + + Providing both :attr:`g` and :attr:`beta` selects cuDNN Gated DeltaNet (GDN) + linear attention instead of scaled-softmax attention. GDN is causal, uses + recurrent state instead of a KV cache, and does not support dropout, attention + bias, sliding windows, context parallelism, or FP8 attention. + Parameters ---------- query_layer : Optional[torch.Tensor], default = None @@ -1664,6 +1897,25 @@ def forward( 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. + g: Optional[torch.Tensor], default = None + GDN log-space scalar decay gate, with ``alpha = exp(g)``. Its shape is the + token dimensions followed by ``num_attention_heads`` and its dtype must be + ``torch.float32``. Expected to satisfy ``g <= 0`` so that ``alpha <= 1``. + Providing any GDN-specific argument selects GDN, and both :attr:`g` and + :attr:`beta` are required. + beta: Optional[torch.Tensor], default = None + GDN per-token write strength, expected in ``[0, 1]``. It has the same shape + and dtype requirements as :attr:`g`. + initial_state: Optional[torch.Tensor], default = None + Optional GDN recurrent state with shape + ``[batch_size, output_heads, v_head_dim, qk_head_dim]`` and dtype + ``torch.float32``. + output_final_state: bool, default = False + Return ``(output, final_state)`` for GDN when ``True``. The final state can be + passed as :attr:`initial_state` to a later invocation. + use_qk_l2norm_in_kernel: bool, default = False + L2-normalize GDN Q and K rows inside the kernel before applying the attention + scale to Q. Engines that do not support this option will decline the operation. """ query_layer, key_layer, value_layer, declared_qkv_layout = _unpack_packed_qkv( @@ -1677,6 +1929,48 @@ def forward( inference_params, ) + gdn_requested = ( + any(value is not None for value in (g, beta, initial_state)) + or output_final_state + or use_qk_l2norm_in_kernel + ) + if gdn_requested: + return self._forward_gdn( + query_layer, + key_layer, + value_layer, + g, + beta, + qkv_format=qkv_format if qkv_format is not None else self.qkv_format, + attention_mask=attention_mask, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_kv=cu_seqlens_kv, + cu_seqlens_q_padded=cu_seqlens_q_padded, + cu_seqlens_kv_padded=cu_seqlens_kv_padded, + max_seqlen_q=max_seqlen_q, + max_seqlen_kv=max_seqlen_kv, + attn_mask_type=attn_mask_type, + window_size=window_size, + bottom_right_diagonal=bottom_right_diagonal, + checkpoint_core_attention=checkpoint_core_attention, + core_attention_bias_type=core_attention_bias_type, + core_attention_bias=core_attention_bias, + alibi_slopes=alibi_slopes, + fast_zero_fill=fast_zero_fill, + inference_params=inference_params, + pad_between_seqs=pad_between_seqs, + fp8_output=fp8_output, + bf16_backward=bf16_backward, + num_splits=num_splits, + score_mod=score_mod, + score_mod_bprop=score_mod_bprop, + score_mod_tensors=score_mod_tensors, + score_mod_bprop_tensors=score_mod_bprop_tensors, + initial_state=initial_state, + output_final_state=output_final_state, + use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel, + ) + with self.prepare_forward_ctx( query_layer, num_gemms=3, diff --git a/transformer_engine/pytorch/attention/dot_product_attention/gdn.py b/transformer_engine/pytorch/attention/dot_product_attention/gdn.py new file mode 100644 index 0000000000..d0c10d76ff --- /dev/null +++ b/transformer_engine/pytorch/attention/dot_product_attention/gdn.py @@ -0,0 +1,251 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""cuDNN frontend Gated DeltaNet attention backend.""" + +import importlib +from functools import lru_cache +from typing import Callable, Optional, Tuple, Union + +import torch + + +@lru_cache(maxsize=1) +def _import_gated_delta_net() -> Callable: + """Import the cuDNN frontend GDN custom op lazily.""" + try: + ops = importlib.import_module("cudnn.linear_attention.ops") + return ops.gated_delta_net + except (AttributeError, ImportError) as exc: + raise ImportError( + "GDN attention requires a nvidia-cudnn-frontend installation that provides " + "cudnn.linear_attention.ops.gated_delta_net and its kernel runtime. Install " + "cuDNN frontend from the matching source revision with the 'cutedsl' extra." + ) from exc + + +def _to_thd(tensor: torch.Tensor, qkv_format: str) -> torch.Tensor: + """Convert a dense sequence tensor to the THD layout required by cuDNN GDN.""" + if qkv_format == "thd": + return tensor + if qkv_format == "sbhd": + tensor = tensor.transpose(0, 1) + return tensor.reshape(-1, *tensor.shape[2:]) + + +def _validate_cu_seqlens( + cu_seqlens: torch.Tensor, + *, + device: torch.device, + name: str, +) -> None: + """Validate a cumulative sequence-length tensor for GDN.""" + if not isinstance(cu_seqlens, torch.Tensor): + raise TypeError(f"{name} must be a torch.Tensor.") + if cu_seqlens.dim() != 1: + raise ValueError(f"{name} must have shape [batch_size + 1].") + if cu_seqlens.numel() < 2: + raise ValueError(f"{name} must contain at least a start and end offset.") + if cu_seqlens.dtype != torch.int32: + raise TypeError(f"{name} must have dtype torch.int32, got {cu_seqlens.dtype}.") + if cu_seqlens.device != device: + raise ValueError(f"{name} must be on {device}, got {cu_seqlens.device}.") + + +class _GatedDeltaNetAttention(torch.nn.Module): + """Adapter from TransformerEngine attention layouts to cuDNN frontend GDN. + + The cuDNN frontend GDN op is differentiated through ``torch.autograd`` (it registers + its own backward internally), so this module does not define an explicit backward. + """ + + def __init__( + self, + scale: float, + num_q_heads: int, + qk_head_dim: int, + v_head_dim: int, + ) -> None: + super().__init__() + self.scale = scale + self.num_q_heads = num_q_heads + self.qk_head_dim = qk_head_dim + self.v_head_dim = v_head_dim + self._dense_cu_seqlens_key: Optional[Tuple[torch.device, int, int]] = None + self._dense_cu_seqlens: Optional[torch.Tensor] = None + + def forward( + self, + query_layer: torch.Tensor, + key_layer: torch.Tensor, + value_layer: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + initial_state: Optional[torch.Tensor] = None, + *, + qkv_format: str, + cu_seqlens_q: Optional[torch.Tensor] = None, + output_final_state: bool = False, + use_qk_l2norm_in_kernel: bool = False, + ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]: + """Run GDN and return a TE-layout output, optionally with the final state.""" + if qkv_format not in {"thd", "bshd", "sbhd"}: + raise ValueError( + "GDN attention only supports qkv_format={'thd', 'bshd', 'sbhd'}, " + f"got {qkv_format!r}." + ) + + if not all( + isinstance(tensor, torch.Tensor) + for tensor in (query_layer, key_layer, value_layer, g, beta) + ): + raise TypeError("GDN Q, K, V, g, and beta must be torch.Tensor instances.") + if initial_state is not None and not isinstance(initial_state, torch.Tensor): + raise TypeError("GDN initial_state must be a torch.Tensor when provided.") + + expected_rank = 3 if qkv_format == "thd" else 4 + qkv = (query_layer, key_layer, value_layer) + if any(tensor.dim() != expected_rank for tensor in qkv): + raise ValueError( + f"Q, K, and V must be {expected_rank}D tensors for qkv_format={qkv_format!r}." + ) + if query_layer.shape != key_layer.shape: + raise ValueError( + "GDN requires Q and K to have the same shape; got " + f"{tuple(query_layer.shape)} and {tuple(key_layer.shape)}." + ) + if query_layer.shape[:-2] != value_layer.shape[:-2]: + raise ValueError("GDN requires Q, K, and V to have the same token dimensions.") + if query_layer.shape[-2] != self.num_q_heads: + raise ValueError( + f"GDN Q and K must have {self.num_q_heads} heads, got {query_layer.shape[-2]}." + ) + # The underlying op supports grouped value heads, but DPA's output contract is fixed + # at construction time. Integrations that use more V heads must expand Q/K before DPA. + # TODO: Support num_q_heads != num_v_heads once DPA's output-width contract allows it. + if value_layer.shape[-2] != self.num_q_heads: + raise ValueError( + f"GDN V must have {self.num_q_heads} heads, got {value_layer.shape[-2]}. " + "DotProductAttention requires its output width to match " + "num_attention_heads * v_head_dim." + ) + if query_layer.shape[-1] != self.qk_head_dim: + raise ValueError( + "GDN Q and K head dimension must match kv_channels; expected " + f"{self.qk_head_dim}, got {query_layer.shape[-1]}." + ) + if value_layer.shape[-1] != self.v_head_dim: + raise ValueError( + "GDN V head dimension must match kv_channels; expected " + f"{self.v_head_dim}, got {value_layer.shape[-1]}." + ) + + device = query_layer.device + if any(not tensor.is_cuda for tensor in qkv): + raise ValueError("GDN attention only supports CUDA tensors.") + if any(tensor.device != device for tensor in (*qkv, g, beta)): + raise ValueError("Q, K, V, g, and beta must be on the same CUDA device.") + if query_layer.dtype != key_layer.dtype or query_layer.dtype != value_layer.dtype: + raise TypeError("Q, K, and V must have the same dtype for GDN attention.") + if query_layer.dtype not in {torch.float16, torch.bfloat16}: + raise TypeError( + "GDN Q, K, and V must have dtype float16 or bfloat16 (no cuDNN frontend " + f"GDN engine supports float32 inputs), got {query_layer.dtype}." + ) + if g.dtype != torch.float32 or beta.dtype != torch.float32: + raise TypeError( + "GDN g and beta must have dtype torch.float32 (the kernel-native dtype)." + ) + + num_output_heads = self.num_q_heads + expected_gate_shape = (*query_layer.shape[:-2], num_output_heads) + if g.shape != expected_gate_shape or beta.shape != expected_gate_shape: + raise ValueError( + "GDN g and beta must both have shape " + f"{expected_gate_shape}; got {tuple(g.shape)} and {tuple(beta.shape)}." + ) + if qkv_format == "thd": + if cu_seqlens_q is None: + raise ValueError("cu_seqlens_q is required for GDN with qkv_format='thd'.") + _validate_cu_seqlens( + cu_seqlens_q, + device=device, + name="cu_seqlens_q", + ) + cu_seqlens = cu_seqlens_q + else: + if cu_seqlens_q is not None: + raise ValueError( + "Dense GDN inputs do not accept cu_seqlens_q. " + "Use qkv_format='thd' for packed or ragged batches." + ) + if qkv_format == "bshd": + batch_size, sequence_length = query_layer.shape[:2] + else: + sequence_length, batch_size = query_layer.shape[:2] + cache_key = (device, batch_size, sequence_length) + if self._dense_cu_seqlens_key != cache_key: + self._dense_cu_seqlens = ( + torch.arange(batch_size + 1, dtype=torch.int32, device=device) * sequence_length + ) + self._dense_cu_seqlens_key = cache_key + cu_seqlens = self._dense_cu_seqlens + + batch_size = cu_seqlens.shape[0] - 1 + expected_state_shape = ( + batch_size, + num_output_heads, + value_layer.shape[-1], + query_layer.shape[-1], + ) + if initial_state is not None: + if initial_state.device != device: + raise ValueError( + f"GDN initial_state must be on {device}, got {initial_state.device}." + ) + if initial_state.dtype != torch.float32: + raise TypeError( + f"GDN initial_state must have dtype torch.float32, got {initial_state.dtype}." + ) + if initial_state.shape != expected_state_shape: + raise ValueError( + f"GDN initial_state must have shape {expected_state_shape}, " + f"got {tuple(initial_state.shape)}." + ) + + q_thd = _to_thd(query_layer, qkv_format) + k_thd = _to_thd(key_layer, qkv_format) + v_thd = _to_thd(value_layer, qkv_format) + g_thd = _to_thd(g, qkv_format) + beta_thd = _to_thd(beta, qkv_format) + gated_delta_net = _import_gated_delta_net() + try: + output, final_state = gated_delta_net( + q_thd, + k_thd, + v_thd, + g_thd, + beta_thd, + cu_seqlens, + scale=self.scale, + initial_state=initial_state, + output_final_state=output_final_state, + use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel, + ) + except ImportError as exc: + raise ImportError( + "The cuDNN frontend GDN kernel runtime is unavailable. Install cuDNN " + "frontend from the matching source revision with the 'cutedsl' extra." + ) from exc + + if qkv_format == "thd": + output = output.reshape(output.shape[0], -1) + else: + output = output.reshape(batch_size, sequence_length, -1) + if qkv_format == "sbhd": + output = output.transpose(0, 1).contiguous() + + if output_final_state: + return output, final_state + return output