diff --git a/qa/L0_pytorch_unittest/test.sh b/qa/L0_pytorch_unittest/test.sh index 759432857a..926b55f0c4 100644 --- a/qa/L0_pytorch_unittest/test.sh +++ b/qa/L0_pytorch_unittest/test.sh @@ -58,6 +58,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" +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..5bdff28edb --- /dev/null +++ b/tests/pytorch/attention/test_gdn_attention.py @@ -0,0 +1,288 @@ +# 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 +from typing import Optional, Tuple + +import pytest +import torch +import torch.nn.functional as F + +from transformer_engine.pytorch import DotProductAttention + + +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 + + +pytestmark = pytest.mark.skipif( + not _gdn_available(), reason="GDN requires a supported cuDNN frontend kernel runtime" +) + +_FWD_TOL = 2e-2 +_STATE_TOL = 2e-2 +_BWD_TOL = 4e-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.""" + 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(k_t.unsqueeze(-2), state).squeeze(-2) + residual = v_t - alpha_t.unsqueeze(-1) * prediction + state = alpha_t[..., None, None] * state + beta_t[..., None, None] * ( + k_t.unsqueeze(-1) @ residual.unsqueeze(-2) + ) + outputs.append(torch.matmul(q_t.unsqueeze(-2), state).squeeze(-2)) + + 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. + """ + 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, + qk_dim, + v.shape[-1], + 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): + torch.manual_seed(1234) + q = torch.randn(batch, sequence, q_heads, qk_dim, device="cuda", dtype=torch.bfloat16) + k = F.normalize(torch.randn_like(q, dtype=torch.float32), dim=-1).to(torch.bfloat16) + v = torch.randn(batch, sequence, v_heads, v_dim, device="cuda", dtype=torch.bfloat16) + 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 + + +def test_gdn_thd_forward_final_state_and_backward(): + """THD GDN matches a PyTorch recurrence in forward and backward.""" + batch, sequence, heads, dim = 2, 128, 2, 64 + q, k, v, g, beta = ( + tensor.reshape(batch * sequence, *tensor.shape[2:]).requires_grad_(True) + for tensor in _inputs(batch, sequence, heads, heads, dim, dim) + ) + cu_seqlens = torch.arange(batch + 1, device="cuda", dtype=torch.int32) * sequence + initial_state = ( + torch.randn(batch, heads, dim, dim, device="cuda", dtype=torch.float32) * 0.05 + ).requires_grad_() + + attention = DotProductAttention( + num_attention_heads=heads, + kv_channels=dim, + qkv_format="thd", + attn_mask_type="padding_causal", + ) + output, final_state = attention( + q, + k, + v, + cu_seqlens_q=cu_seqlens, + cu_seqlens_kv=cu_seqlens, + g=g, + beta=beta, + initial_state=initial_state, + output_final_state=True, + ) + 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_() + output_ref, final_state_ref = _gdn_reference( + reference_inputs["q"], + reference_inputs["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, "output") + _assert_rms_close(final_state, final_state_ref, _STATE_TOL, "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, f"d{name}") + _assert_rms_close(initial_state.grad, initial_state_ref.grad, _BWD_TOL, "dinitial_state") + + +@pytest.mark.parametrize("qkv_format", ["bshd", "sbhd"]) +def test_gdn_dense_layout_and_grouped_value_heads(qkv_format): + """Dense TE layouts and grouped heads match a pure PyTorch reference.""" + batch, sequence, q_heads, v_heads, dim = 1, 128, 1, 2, 64 + q, k, v, g, beta = _inputs(batch, sequence, q_heads, v_heads, dim, dim) + 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=q_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, "output") + + +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="requires both g and beta"): + attention(q, k, v, g=g) diff --git a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py index e6f50eb0da..2ce14e5337 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 @@ -59,6 +59,7 @@ FusedAttention, FlashAttention, ) +from transformer_engine.pytorch.attention.dot_product_attention.gdn import GatedDeltaNetAttention # Setup Attention Logging @@ -773,6 +774,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 @@ -1323,6 +1331,129 @@ def get_quantizer_roles( ] return base[:num_quantizers] + 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], + 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], + inference_params: Optional[InferenceParams], + pad_between_seqs: Optional[bool], + fp8_output: 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.""" + if g is None or beta is None: + raise ValueError("GDN attention requires both g and beta.") + 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 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 padding between packed sequences.") + 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 padding between packed sequences.") + if fp8_output: + raise ValueError("GDN does not support FP8 output.") + 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.") + + gdn_kwargs = { + "qkv_format": qkv_format, + "cu_seqlens_q": cu_seqlens_q, + "cu_seqlens_kv": cu_seqlens_kv, + "initial_state": initial_state, + "output_final_state": output_final_state, + "use_qk_l2norm_in_kernel": use_qk_l2norm_in_kernel, + } + if checkpoint_core_attention: + return self._checkpointed_attention_forward( + self.gdn_attention, + query_layer, + key_layer, + value_layer, + g, + beta, + **gdn_kwargs, + ) + return self.gdn_attention( + query_layer, + key_layer, + value_layer, + g, + beta, + **gdn_kwargs, + ) + @no_torch_dynamo(recursive=False) def forward( self, @@ -1357,7 +1488,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. @@ -1368,14 +1504,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:: @@ -1453,6 +1590,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 @@ -1596,6 +1740,23 @@ 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 ``max(num_q_heads, num_v_heads)`` and its dtype + must be ``torch.float32``. Providing :attr:`g` and :attr:`beta` selects GDN. + beta: Optional[torch.Tensor], default = None + GDN per-token write strength. 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, qk_head_dim, v_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. Engines that do not support + this option will decline the operation. """ query_layer, key_layer, value_layer, declared_qkv_layout = _unpack_packed_qkv( @@ -1609,6 +1770,44 @@ 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, + 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, + inference_params=inference_params, + pad_between_seqs=pad_between_seqs, + fp8_output=fp8_output, + 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..cd7b70a7fc --- /dev/null +++ b/transformer_engine/pytorch/attention/dot_product_attention/gdn.py @@ -0,0 +1,238 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""cuDNN frontend Gated DeltaNet attention backend.""" + +import importlib +from typing import Callable, Optional, Tuple, Union + +import torch + + +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.""" + + 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 + + def forward( + self, + query_layer: torch.Tensor, + key_layer: torch.Tensor, + value_layer: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + *, + qkv_format: str, + cu_seqlens_q: Optional[torch.Tensor] = None, + cu_seqlens_kv: 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]]: + """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]}." + ) + 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, torch.float32}: + raise TypeError( + "GDN Q, K, and V must have dtype float16, bfloat16, or float32, " + f"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 = max(query_layer.shape[-2], value_layer.shape[-2]) + 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 qkv_format == "bshd": + batch_size, sequence_length = query_layer.shape[:2] + else: + sequence_length, batch_size = query_layer.shape[:2] + full_cu_seqlens = ( + torch.arange(batch_size + 1, dtype=torch.int32, device=device) * sequence_length + ) + if cu_seqlens_q is not None: + _validate_cu_seqlens(cu_seqlens_q, device=device, name="cu_seqlens_q") + if not torch.equal(cu_seqlens_q, full_cu_seqlens): + raise ValueError( + "Dense GDN inputs do not support padding; cu_seqlens_q must describe " + "full, equal-length sequences. Use qkv_format='thd' for ragged batches." + ) + cu_seqlens = full_cu_seqlens + + if cu_seqlens_kv is not None: + _validate_cu_seqlens(cu_seqlens_kv, device=device, name="cu_seqlens_kv") + if cu_seqlens.data_ptr() != cu_seqlens_kv.data_ptr() and not torch.equal( + cu_seqlens, cu_seqlens_kv + ): + raise ValueError("GDN requires cu_seqlens_q and cu_seqlens_kv to be identical.") + + batch_size = cu_seqlens.shape[0] - 1 + expected_state_shape = ( + batch_size, + num_output_heads, + query_layer.shape[-1], + value_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