[PyTorch] GDN2 support and linear attention refactor by ksivaman · Pull Request #3521 · NVIDIA/TransformerEngine

ksivaman · 2026-09-15T18:52:17Z

Description

Adds support for GDN2 variant and redesign the linear attention API.

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

Integrate fused GDN2 kernel from cudnn-frontend (v1.29.0 required).
Linear attention base module.

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

Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

greptile-apps · 2026-09-15T18:56:17Z

Confidence Score: 4/5

The PR does not yet appear safe to merge because the previously reported required-runtime installation gap remains unresolved; the new architecture wording should also be corrected.

Findings

FSDP setup calls missing method ▶
Required runtime is not installed ▶
FP8 Test Fails Early ▶
Architecture Support Is Overstated ▶

Summary

Adds Gated DeltaNet v2 through the cuDNN frontend and refactors linear-attention variants around shared layout validation, lifecycle handling, recurrent-state support, checkpointing, and CUDA-graph integration.

Exposes and documents the new GatedDeltaNet2Attention API.
Adds dense and packed-layout GDN2 coverage, including state continuation and backward checks.
Raises the cuDNN frontend dependency floor to 1.29.0.
Extends TransformerEngine distributed and graph infrastructure to recognize parameter-free linear-attention modules.

Diagram

%%{init: {'theme': 'neutral'}}%%
flowchart LR
    U[Linear-attention caller] --> B[LinearAttentionBase]
    B --> V[Lifecycle and input validation]
    V --> A[Variant kernel adapter]
    A --> G1[cuDNN GDN op]
    A --> G2[cuDNN GDN2 op]
    G1 --> O[Output and optional final state]
    G2 --> O

_{Reviews (8) · Last reviewed commit: "docsting"}

greptile-apps · 2026-09-15T18:56:37Z

        LayerNorm,
        RMSNorm,
        TransformerEngineBaseModule,
+        LinearAttentionBase,


FSDP setup calls missing method

Adding LinearAttentionBase to the TE-class registry makes directly or individually FSDP-wrapped GDN modules enter prepare_te_modules_for_fsdp. That function unconditionally calls fast_setattr, but the new base is a plain torch.nn.Module and does not implement this method. FSDP preparation therefore raises AttributeError before the model can run.

greptile-apps · 2026-09-15T18:56:37Z

        "pydantic",
        "nvdlfw-inspect",
-        "nvidia-cudnn-frontend>=1.28.0",
+        "nvidia-cudnn-frontend>=1.29.0",


Required runtime is not installed

GDN2 explicitly requires the cuDNN frontend cutedsl runtime, but both installation declarations request only the base nvidia-cudnn-frontend package. If that base package does not include the optional cutlass runtime, the newly mandatory L0 GDN2 test fails during collection and the advertised feature cannot run.

for more information, see https://pre-commit.ci

Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

greptile-apps · 2026-09-15T20:59:30Z

+        attention(q, k, v, g=g, beta=beta)
+
+
+def test_gdn2_rejects_fp8_autocast():


FP8 Test Fails Early

On a CUDA device where FP8 is unavailable, autocast(enabled=True) raises RuntimeError while entering the context, before GDN2 can raise the expected ValueError. Because this test no longer checks FP8 availability, it fails on those supported test configurations instead of validating the module's FP8 rejection behavior.

Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

ksivaman · 2026-09-16T17:47:51Z

/te-ci pytorch

KshitijLakhani

Can we consider moving both gdn and gdn2 tests into the same linear attention file for testing and then we can just add the newer linear attention types to it in the future. I'm also assuming here that the tests could share some setup/boilerplate info (e.g. is gdn/gdn op available ?) .
The other (maybe preferred) option is to have a these separate as you currently have them and then have a shared test file/class in which the shared utilities could cover (for e.g.) :
i) RMS comparison helpers.
ii) THD input conversion.
iii) TP/lifecycle tests.
iv) Checkpointing tests.
v) Common output-shape/layout conformance.
nit: I understand this might be for simplicity in interest of time that even thought cuDNN GDN2 supports distinct Q, K and V head counts, the TE wrapper requires all head counts to be equal, but please one of the following - i) Add this support (ideal) or ii) Add a TODO and add either Charlene or my name for it for the future) . And on the same note, I learnt from my agent that cuDNN GDN API seems to be supporting a a lot of other args, but I'm guessing yoau are intentionally using a minimal GDN2 API ??
This is not supposed to be a blocker for this PR but @cyanguwa I think it would make sense to have a plan for integration into upper Transformer Layers ? (right now it is DPA/MHA hard-coded but we might have to change it with linear / spare right ?)
Something liek this:
TransformerLayer

    ├── MultiheadAttention
    │     QKV projection → DPA → output projection
    │
    └── GDN2 attention block
          QKV + gate projections → GDN2 core → output projection

Non-blocking: I understand the hesitation to refactor TransformerEngineBaseModule, especially since it is central and has been relatively stable. Keeping LinearAttentionBase separate is reasonable for this PR and avoids pulling unused FP8 state into parameter-free LA kernels.

My concern is less the duplicated implementation and more that this creates a second module type participating in TE-wide lifecycle behavior. This PR already updates get_te_classes(), CUDA-graph discovery, and implements the FSDP fast_setattr contract. Future framework-level code could recognize TransformerEngineBaseModule but unintentionally omit LinearAttentionBase.

For now, as discussed offline, could you please document the minimal TE lifecycle contract being implemented here ?

KshitijLakhani · 2026-09-16T22:40:45Z

+    return True
+
+
+def _gdn2_supported_arch() -> bool:


nit: Any reason this check does not exist for GDN and only GDN-2 tests ?

GDN-2 doesn't have a cu-tile backend, only FROST engine is supported here, so blackwell+

KshitijLakhani · 2026-09-16T22:53:03Z

+    return None
+
+
+class LinearAttentionKernelAdapter(torch.nn.Module):


Non-blocking follow-up: I like the private adapter as a thin shim, but _run() currently encodes several GDN/GDN2-specific assumptions: aligned Q/K/V token dimensions, identical Q/K shapes, equal Q/V head counts, and state heads derived from Q heads. These do not describe LA generally: GDP uses T Q/g rows but T ×num_householder K/V/beta rows, while KDA supports HK = H or HK = HV and uses HO = max(H, HV) for output/state heads.

Based on our offline discussion, I am okay deferring that generalization until KDA/GDP are added. At that point, I think these checks should move into the flavor-specific adapter rather than simply being removed or weakened globally, so GDN/GDN2 retain strict validation.
It would be helpful to leave a TODO or tracking issue recording that intended boundary.

KshitijLakhani · 2026-09-16T23:25:35Z

+
+    Gated DeltaNet v2 replaces Gated DeltaNet's scalar gates with channel-wise
+    ones: ``g`` and ``beta`` carry one value per query/key channel and ``w`` one
+    value per value channel. The cuDNN frontend serves GDN-2 on Blackwell


Incorrectly clubs 107 as Blackwell ?

KshitijLakhani · 2026-09-16T23:26:53Z

+    if not torch.cuda.is_available():
+        return False
+    major, minor = torch.cuda.get_device_capability()
+    return (major, minor) in {(10, 0), (10, 1), (10, 2), (10, 3), (10, 7)}


KshitijLakhani · 2026-09-16T23:27:56Z

+
+
+_GDN2_OP_AVAILABLE = _gdn2_op_available()
+if os.getenv("NVTE_GDN2_TEST_REQUIRED", "0") == "1" and not _GDN2_OP_AVAILABLE:


I'm unsure, but, could required mode also fail when the GPU architecture is unsupported? Currently it checks only whether the op/runtime imports; requires_gdn2 can then skip every numerical test on the wrong GPU, allowing the required CI command to pass without exercising GDN2. MAybe use a shared condition such as _GDN2_SUPPORTED = _GDN2_OP_AVAILABLE and _gdn2_supported_arch() for both the required-mode failure and the skip marker.

We run the CI with NVTE_GDN2_TEST_REQUIRED=1 and the _GDN2_OP_AVAILABLE checks GDN support. Later, _gdn2_supported_arch skips this test on unsupported architectures. So we will only run this test on supported hardware and skip otherwise.

Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

ksivaman · 2026-09-17T00:47:35Z

/te-ci pytorch L0

greptile-apps · 2026-09-17T00:49:27Z

+    value per value channel. The cuDNN frontend serves GDN-2 on Blackwell+
+    (SM100/SM103/SM107) only.


Architecture Support Is Overstated

Describing support as “Blackwell+” implies compatibility with architectures beyond Blackwell, but the same sentence limits the available engine to SM100, SM103, and SM107. This may lead users of later architectures to expect GDN2 to work when that support is not documented.

Suggested change

value per value channel. The cuDNN frontend serves GDN-2 on Blackwell+

(SM100/SM103/SM107) only.

value per value channel. The cuDNN frontend serves GDN-2 on Blackwell

(SM100/SM103/SM107) only.

Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!

ksivaman added 3 commits September 11, 2026 12:21

GDN2 attention

5c0e47e

Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

Update cudnnFE requirement

f9bb3ae

Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

Merge branch 'NVIDIA:main' into gdn2_attention

3566b27

ksivaman requested a review from KshitijLakhani September 15, 2026 18:52

ksivaman requested a review from cyanguwa as a code owner September 15, 2026 18:52

ksivaman added the 2.20 label Sep 15, 2026

greptile-apps Bot reviewed Sep 15, 2026

View reviewed changes

[pre-commit.ci] auto fixes from pre-commit.com hooks

de902ac

for more information, see https://pre-commit.ci

ksivaman marked this pull request as draft September 15, 2026 19:02

ksivaman added 2 commits September 15, 2026 13:52

Minor doc changes and review

7aa3b61

Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

Remove negative test cases

ab2cf2f

Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

ksivaman marked this pull request as ready for review September 15, 2026 20:56

greptile-apps Bot reviewed Sep 15, 2026

View reviewed changes

ksivaman added 2 commits September 16, 2026 10:29

Merge branch 'main' into gdn2_attention

5e1dc0f

Remove fp8 autocast checks

b506d72

Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

KshitijLakhani reviewed Sep 16, 2026

View reviewed changes

ksivaman added 3 commits September 16, 2026 16:47

Merge branch 'main' into gdn2_attention

2495c5f

review

ae7676a

Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

docsting

65e5133

Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

greptile-apps Bot reviewed Sep 17, 2026

View reviewed changes

		attention(q, k, v, g=g, beta=beta)


		def test_gdn2_rejects_fp8_autocast():

		return None


		class LinearAttentionKernelAdapter(torch.nn.Module):



		_GDN2_OP_AVAILABLE = _gdn2_op_available()
		if os.getenv("NVTE_GDN2_TEST_REQUIRED", "0") == "1" and not _GDN2_OP_AVAILABLE:

		value per value channel. The cuDNN frontend serves GDN-2 on Blackwell+
		(SM100/SM103/SM107) only.

Uh oh!

Conversation

ksivaman commented Sep 15, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Type of change

Changes

Checklist:

Uh oh!

greptile-apps Bot commented Sep 15, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Confidence Score: 4/5

Findings

Summary

Diagram

Uh oh!

greptile-apps Bot Sep 15, 2026

Choose a reason for hiding this comment

Uh oh!

greptile-apps Bot Sep 15, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

greptile-apps Bot Sep 15, 2026

Choose a reason for hiding this comment

Uh oh!

ksivaman commented Sep 16, 2026

Uh oh!

KshitijLakhani left a comment

Choose a reason for hiding this comment

Uh oh!

KshitijLakhani Sep 16, 2026

Choose a reason for hiding this comment

Uh oh!

ksivaman Sep 17, 2026

Choose a reason for hiding this comment

Uh oh!

KshitijLakhani Sep 16, 2026

Choose a reason for hiding this comment

Uh oh!

KshitijLakhani Sep 16, 2026

Choose a reason for hiding this comment

Uh oh!

KshitijLakhani Sep 16, 2026

Choose a reason for hiding this comment

Uh oh!

KshitijLakhani Sep 16, 2026

Choose a reason for hiding this comment

Uh oh!

ksivaman Sep 17, 2026

Choose a reason for hiding this comment

Uh oh!

ksivaman commented Sep 17, 2026

Uh oh!

greptile-apps Bot Sep 17, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

ksivaman commented Sep 15, 2026 •

edited

Loading

greptile-apps Bot commented Sep 15, 2026 •

edited

Loading