[PyTorch] Fused GDN attention

[ksivaman]

Description

Add support for GDN attention kernels from cudnn-frontend.

New tests B100 wall clock time: 1m 7s

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

The PR integrates cuDNN frontend Gated DeltaNet attention into DotProductAttention, including dense and packed layouts, recurrent state, checkpointing, dependency guards, and GPU tests.

• Adds GDN request validation and dispatch through the public PyTorch attention API.
• Adds a cuDNN frontend adapter for layout conversion, state handling, and kernel invocation.
• Raises the cuDNN frontend dependency floor and adds GDN coverage to L0 PyTorch QA.

Confidence Score: 3/5

The PR is not yet safe to merge because valid THD self-attention metadata is rejected and malformed packed offsets can still reach the GDN kernel.

Matching Q/KV sequence metadata forwarded by the documented THD calling path now raises before GDN dispatch, while the packed-offset validator still forwards offsets without checking their boundary and ordering invariants.

Files Needing Attention: transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py; 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 public GDN dispatch and capability validation, but the KV-offset fix rejects valid conventional THD self-attention metadata.
transformer_engine/pytorch/attention/dot_product_attention/gdn.py	Adds the cuDNN GDN adapter and tensor validation; previously reported malformed packed offsets remain insufficiently validated.
tests/pytorch/attention/test_gdn_attention.py	Adds broad numerical, backward, state, layout, lifecycle, and checkpoint coverage for GDN.
build_tools/pytorch.py	Raises the PyTorch build requirement for nvidia-cudnn-frontend to 1.28.0.
pyproject.toml	Aligns the isolated build-system cuDNN frontend requirement with the GDN integration.
qa/L0_pytorch_unittest/test.sh	Adds the GDN attention suite as a required L0 GPU test.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart LR
  A[DotProductAttention request] --> B{GDN arguments present?}
  B -- No --> C[Scaled-softmax backend selection]
  B -- Yes --> D[Validate GDN options and metadata]
  D --> E[Convert Q K V gates to THD]
  E --> F[cuDNN frontend gated_delta_net]
  F --> G[Restore requested output layout]
  G --> H[Output]
  G --> I[Optional final recurrent state]

Loading

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

[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.

[ksivaman]

/te-ci pytorch L0

[ksivaman]

/te-ci pytorch

[ksivaman]

/te-ci pytorch

[KshitijLakhani]

Three general important things I noticed that would be good to fix is :

1. Different Q and KV sequence boundaries can be silently accepted (via cu_seqlens_q and cu_seqlens_kv) while only the Q boundaries (cus_seqlens) are passed down to cuDNNN FE. So maybe we should have some checks to ensure that these are exactly the same. IIRC, currently checks exist only for their shape but they could have different values, which could be misleading for the user.
   So, for e.g. this would be acceptable right now, but I do not think we should allow it as it conveys incorrect usage to a TE user

  cu_seqlens_q  = [0, 48, 160]
  cu_seqlens_kv = [0, 64, 160]

2. FP32 Q/K/V are advertised as supported even though neither cuDNN FE 1.27 GDN engine supports FP32 (FROST or cuTile IIUC) so we might need to gate against this - please correct me if I'm wrong.
3. cuDNN FE version/state-layout compatibility
   • v1.27 exposes state as [N, H, K, V].
   • Current develop uses [N, H, V, K] -> I'd imagine this would become 2.18+
     So this only affects the shape of initial_state . We might want to choose the shape conditional on the cuDNN FE version within TE so that this remains opaque to the TE user

[KshitijLakhani]

I think I mentioned this in my general comments in the PR (please look at that for an example), but should we not just have a check to ensure that the contents of cu_seqlens_q and cu_seqlens_kv are the same ? Because GDN expects that for the cuDNN FE API IIUC, so we don't want to mislead users

[KshitijLakhani]

Not in this PR but for the future:

This is a question for the cuDNN FE team: I'm wondering if the chunk_size can be exposed as a user configurable arg )I understand it is heuristic-based selection inside cuDNN FE right now) ?
I'm wondering if we can repurpose the chunk to be context and choose a very large chunk_size, split it across ranks, perform linear attention on the sub-chunk per rank and then AG the sub-chunks.
This would enable large context lengths
cc: @cyanguwa

[cyanguwa]

Not in this PR but for the future:

This is a question for the cuDNN FE team: I'm wondering if the chunk_size can be exposed as a user configurable arg )I understand it is heuristic-based selection inside cuDNN FE right now) ? I'm wondering if we can repurpose the chunk to be context and choose a very large chunk_size, split it across ranks, perform linear attention on the sub-chunk per rank and then AG the sub-chunks. This would enable large context lengths cc: @cyanguwa

I think that makes sense. We should follow up with the cuDNN team. Thanks!

[greptile-apps]

Matching KV offsets are rejected

When a THD self-attention caller supplies the documented equal cu_seqlens_q and cu_seqlens_kv tensors, this unconditional check rejects the request before GDN dispatch, causing an otherwise valid packed self-attention call to raise ValueError.

Knowledge Base Used: PyTorch attention execution