[PyTorch] Fused GDN attention

[ksivaman]

Description

Add support for GDN attention kernels from cudnn-frontend.

Type of change

• Documentation change (change only to the documentation, either a fix or a new content)
• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Infra/Build change
• Code refactoring

Changes

• Kernel integration and testing.
• cudnnFE/cutlass version guards.

Checklist:

• I have read and followed the contributing guidelines
• The functionality is complete
• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes

[greptile-apps]

Greptile Summary

Adds a cuDNN-frontend Gated DeltaNet backend to PyTorch DotProductAttention, including recurrent-state support and dense/packed layout validation.

• Selects GDN when gates, recurrent state, final-state output, or in-kernel Q/K normalization are requested.
• Adapts BSHD, SBHD, and THD inputs to the cuDNN linear-attention operation.
• Adds forward, backward, final-state, grouped-head, and layout tests to the L0 PyTorch suite.

Confidence Score: 4/5

The PR appears safe to merge, with a non-blocking opportunity to reject malformed packed sequence boundaries before invoking the CUDA kernel.

The GDN dispatch and tested numerical paths are coherent, but the new THD adapter forwards cumulative sequence-length values without checking their ordering or bounds, reducing diagnostics and robustness for invalid public inputs.

Files Needing Attention: transformer_engine/pytorch/attention/dot_product_attention/gdn.py

Important Files Changed

Filename	Overview
transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py	Adds GDN construction, feature selection, unsupported-option validation, dispatch, and public recurrent-state arguments.
transformer_engine/pytorch/attention/dot_product_attention/gdn.py	Implements layout conversion and the cuDNN GDN adapter; packed offset values should be validated before kernel dispatch.
tests/pytorch/attention/test_gdn_attention.py	Adds numerical forward/backward, final-state, dense-layout, grouped-head, and partial-gate validation coverage.
qa/L0_pytorch_unittest/test.sh	Adds the GDN attention tests to the L0 PyTorch CI job.

Sequence Diagram

sequenceDiagram
    participant User
    participant DPA as DotProductAttention
    participant GDN as GatedDeltaNetAttention
    participant cuDNN as cuDNN frontend GDN
    User->>DPA: Q, K, V, g, beta, optional state
    DPA->>DPA: Validate unsupported attention options
    DPA->>GDN: Dispatch GDN request
    GDN->>GDN: Validate shapes and convert to THD
    GDN->>cuDNN: gated_delta_net(..., cu_seqlens)
    cuDNN-->>GDN: output, final_state
    GDN-->>DPA: TE-layout output, optional final_state
    DPA-->>User: result

Loading

_{Reviews (1): Last reviewed commit: "[pre-commit.ci] auto fixes from pre-comm..." | Re-trigger Greptile}

[greptile-apps]

Validate packed sequence offsets

The THD path checks only the rank, length, dtype, and device of cu_seqlens, so negative, non-monotonic, nonzero-start, or incorrect terminal offsets are forwarded to gated_delta_net. Validate that the offsets begin at zero, are nondecreasing, and terminate at the flattened token count to provide a deterministic argument error instead of an opaque kernel failure or incorrect sequence/state partitioning.

Knowledge Base Used: PyTorch Attention Stack

[cyanguwa]

I put my AI to work and I think some of the comments make sense. Could you please take a look at these first? Happy to get on a call and discuss the details. Thanks.

Summary
The adapter itself is careful and well-validated, but there are three things I'd want resolved before this leaves draft: the tests can never run in CI as configured, gradients for initial_state are silently dropped when checkpoint_core_attention=True, and the GDN path returns before prepare_forward_ctx so it silently ignores fp8_autocast. There's also a design question worth settling first, since it determines how much of the rest matters. CI currently shows PyTorch, JAX, and All failing with Core cancelled — JAX failing on a PyTorch-only PR suggests infra noise, but the PyTorch failure is worth triaging before review effort goes further.

Design question (worth settling first)
Before digging into details: is DotProductAttention.forward the right entry point for GDN? Gated DeltaNet isn't dot-product attention, and the current shape of this is that a single module now has two disjoint behaviors selected by sniffing five new kwargs, plus a ~60-line block in _forward_gdn whose only job is to reject the other ~25 kwargs. The reject-list is already incomplete — max_seqlen_q, max_seqlen_kv, fast_zero_fill, and is_first_microbatch are accepted and silently ignored — and it has to be manually re-audited every time someone adds an argument to forward(). A separate GatedDeltaNet module (or at minimum an explicit constructor-time linear_attention=True switch, so the unsupported-option check happens once at init rather than per call) would avoid all of that. If reuse of the DotProductAttention surface is a hard requirement from the model-integration side, it'd help to say so in the PR description, because it's the main thing shaping this diff.

Correctness
1. initial_state gradients are silently dropped under activation checkpointing.

In _forward_gdn, initial_state is passed inside gdn_kwargs rather than positionally:

if checkpoint_core_attention:
    return self._checkpointed_attention_forward(
        self.gdn_attention, query_layer, key_layer, value_layer, g, beta, **gdn_kwargs,
    )
That routes through _CheckpointFunction, which only treats positional args as differentiable inputs:


distributed.py
Lines 386-388
ctx.inputs = [arg if not torch.is_tensor(arg) else None for arg in args]
tensor_inputs = [arg if torch.is_tensor(arg) else None for arg in args]
ctx.save_for_backward(*tensor_inputs)
and returns grads only for those:


distributed.py
Lines 469-472
grads = tuple(
    inp.grad if isinstance(inp, torch.Tensor) else None for inp in detached_inputs
)
return (None, None, None, None, None, None) + grads
Since the forward runs under torch.no_grad() inside the Function, a tensor arriving via ctx.kwargs has no path back to autograd, so initial_state.grad comes back None with no error. g and beta are fine because they're positional. Please pass initial_state positionally too, and add a checkpoint_core_attention=True case to the tests — the current suite asserts initial_state.grad matches the reference, which is exactly the assertion that would have caught this, but only on the non-checkpointed path.

2. The GDN path returns before prepare_forward_ctx, so fp8_autocast is silently ignored.

The early return is inserted immediately above with self.prepare_forward_ctx(...), which means prepare_forward/end_forward never run for GDN calls. The new docstring note says GDN doesn't support FP8 attention, but nothing enforces it: calling GDN inside an fp8_autocast() block quietly computes in high precision instead of raising, which is inconsistent with how every other unsupported option here is handled. I'd add an explicit check against the FP8 global state manager alongside the other raise ValueError guards. Separately, could you confirm with the module owners that bypassing prepare_forward_ctx wholesale is intended? Skipping end_forward() for one code path in a TransformerEngineBaseModule is the kind of thing that tends to surface later as recipe-state weirdness rather than as a clean failure.

3. Head-count validation is asymmetric, and the output width can silently disagree with num_attention_heads.

GatedDeltaNetAttention.forward validates Q heads strictly against self.num_q_heads but never checks V heads against anything — output heads are inferred from the tensor via num_output_heads = max(query_layer.shape[-2], value_layer.shape[-2]). test_gdn_dense_layout_and_grouped_value_heads leans on this: it builds DotProductAttention(num_attention_heads=1, kv_channels=64) and then passes V with 2 heads, so the returned tensor is [b, s, 128] while the module was configured for a hidden size of 64. That breaks the invariant the rest of TE relies on, and it means GDN can't currently be dropped into MultiheadAttention, whose output projection is sized from num_attention_heads * kv_channels. Note also that the normal path does enforce this:


dot_product_attention.py
Lines 1554-1555
assert num_gqa_groups == self.num_gqa_groups_per_partition, (
    "Keys and values must have num_gqa_group ="
GDN skips it entirely. Since more V heads than QK heads is the normal configuration for Gated DeltaNet and TE's num_gqa_groups convention assumes the opposite, I don't think this is fixable by just adding an assert — it needs a decision on how V/output heads are expressed in the constructor, and then validation against that.

4. Two torch.equal calls force a device sync on every forward.

if not torch.equal(cu_seqlens_q, full_cu_seqlens):
...
if cu_seqlens.data_ptr() != cu_seqlens_kv.data_ptr() and not torch.equal(cu_seqlens, cu_seqlens_kv):
Both return a Python bool from CUDA tensors, so each is a host synchronization in the training inner loop, and both will break CUDA graph capture. The second one fires in the common THD case where a caller passes the same logical cu_seqlens as two separate tensors (the data_ptr fast path only helps when it's literally the same tensor). Consider dropping these to shape/dtype checks, or gating them behind a debug flag. The per-call torch.arange(batch_size + 1, ...) in the dense branch is a smaller version of the same concern and could be cached on the module.

Tests
5. The test file can never run in CI as configured — this is my biggest concern.

pytestmark = pytest.mark.skipif(not _gdn_available(), ...) skips the entire module unless the cuDNN frontend GDN op plus cutlass or cuda.tile are importable, and the new line in qa/L0_pytorch_unittest/test.sh will exit 0 with everything skipped. Combined with the unchecked "cudnnFE/cutlass version guards" box in your own PR description, that means this feature would merge with tests that pass by not running, and nobody would notice when it regresses. Could you confirm whether the L0 PyTorch image actually installs the frontend with the cutedsl extra? If it doesn't yet, I'd rather this PR either adds that to the image or makes the skip loud in CI (e.g. hard-fail when a NVTE_* CI env var is set and the runtime is missing) so the gap is visible rather than silent.

6. Coverage gaps I'd want filled, in rough priority order.

State round-trip. The docstring advertises "the final state can be passed as initial_state to a later invocation," which is the headline inference use case, and nothing tests it. Splitting a sequence into two chunks and checking that chunked execution matches single-shot would be the single most valuable test here.
Ragged sequences. test_gdn_thd_forward_final_state_and_backward uses cu_seqlens = arange(batch + 1) * sequence, i.e. uniform lengths, so the packed path is only ever exercised with equal-length sequences. The reference implementation even has a start == end empty-sequence branch that no test reaches.
checkpoint_core_attention=True (see item 1).
fp16, which the validation accepts but no test covers.
The unsupported-option guards: roughly fifteen raise ValueErrors were added and exactly one is tested (test_gdn_requires_both_gates). A single parametrized table over (kwarg, expected message) would cover the rest cheaply, and would be the thing that catches it when one of those checks gets dropped in a refactor.
Smaller items
_import_gated_delta_net() runs on every forward call. Worth a functools.lru_cache since it's pure.
Two different conditions (cu_seqlens_q_padded/cu_seqlens_kv_padded and pad_between_seqs) raise the identical string "GDN does not support padding between packed sequences.", so the message doesn't tell you which one you tripped.
gdn_requested is true if any of g, beta, initial_state, output_final_state, or use_qk_l2norm_in_kernel is set, but the docstring says GDN is selected by "providing both g and beta". A user who passes only use_qk_l2norm_in_kernel=True gets "GDN attention requires both g and beta," which won't be obvious. Either match the doc to the code or narrow the sniffing.
_checkpointed_attention_forward is still annotated -> torch.Tensor but can now return a tuple through the GDN path.
GDN accepts qkv_format='thd' with attn_mask_type='causal', whereas the softmax path asserts "padding" in attn_mask_type for THD. Minor, but worth matching for consistency.
gdn.py isn't referenced from any __init__.py or from docs/api/pytorch.rst — is GatedDeltaNetAttention meant to be public, or purely internal to DotProductAttention?
Worth confirming against the kernel: use_qk_l2norm_in_kernel combined with scale. The test's reference L2-normalizes Q/K and then applies scale, so the test encodes an assumption about ordering that isn't documented anywhere.