[JAX] Add attention tutorials by KshitijLakhani · Pull Request #3162 · NVIDIA/TransformerEngine

KshitijLakhani · 2026-07-01T16:30:05Z

Description

Add attention tutorial for integrating TE into an existing framework

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

Add a page for non-CP attn tutorial
Add a page for CP attn tutorial

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: Kshitij Lakhani <klakhani@nvidia.com>

greptile-apps · 2026-07-01T23:01:43Z

Greptile Summary

This PR adds two JAX attention tutorials to the TransformerEngine documentation: a single-GPU BSHD tutorial covering GQA+SWA and DeepSeek-style MLA head dimensions, and a multi-GPU context-parallel tutorial covering packed THD GQA+SWA with Ring and AllGather strategies.

attention.py + attention_single_gpu.rst: introduces TEDotProductAttention, benchmarks it against a native JAX baseline, and demonstrates asymmetric Q/K vs V head dims for MLA.
attention_context_parallel.py + attention_context_parallel.rst: demonstrates striped load-balancing, JAX mesh sharding, and explicit in_shardings/out_shardings for Ring and AllGather CP strategies on a 4-GPU packed-THD workload.
test_attention.py: pytest entry points with correctly deferred imports and hardware-capability skip guards for both tutorials.

Confidence Score: 5/5

Safe to merge; the change is purely additive tutorial documentation with no modifications to the core TE library.

All changes are new documentation files and a test helper. The tutorial Python code is self-contained and only exercised when the appropriate hardware is present. The test file correctly defers module imports into test bodies so skip guards fire before any large tensor allocations happen for the CP case. The single finding is a wasteful tensor allocation in an assertion helper that runs only on 4-GPU Hopper hardware.

Files Needing Attention: docs/examples/jax/test_attention.py line 131 — unnecessary large-tensor allocation inside _assert_cp_result.

Important Files Changed

Filename	Overview
docs/examples/jax/attention.py	New single-GPU tutorial: implements a native JAX GQA+SWA baseline and a TE DotProductAttention wrapper with MLA head-dim support; module-scope tensor creation and model init are intentional for the tutorial pattern.
docs/examples/jax/attention_context_parallel.py	New context-parallel tutorial: packed THD GQA+SWA with Ring and AllGather strategies, correct striped load-balancing, sharding, and forward+backward timing loop.
docs/examples/jax/test_attention.py	New pytest file; imports correctly deferred to test bodies; CP tests gated by requires_cp; _assert_cp_result allocates large tensors unnecessarily to retrieve shapes already available as module attributes.
docs/examples/jax/attention_single_gpu.rst	New RST page for single-GPU tutorial; correct literalinclude markers, back-links, and API references.
docs/examples/jax/attention_context_parallel.rst	New RST page for CP tutorial; correctly embeds all code sections via literalinclude markers and captured output blocks.
docs/examples/jax/attention.rst	Replaced TODO placeholder with a real index page linking to the two sub-tutorials via list-table and toctree.
docs/examples/te_jax_integration.rst	Marks the Attention row as Available with a brief description; straightforward two-line update.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A[te_jax_integration.rst] --> B[attention.rst\nindex page]
    B --> C[attention_single_gpu.rst]
    B --> D[attention_context_parallel.rst]

    C -->|literalinclude| E[attention.py\nBSHD tutorial code]
    C -->|literalinclude| F[attention.out\ncaptured output]

    D -->|literalinclude| G[attention_context_parallel.py\nTHD CP tutorial code]
    D -->|literalinclude| H[attention_context_parallel.out\ncaptured output]

    E --> E1[FlaxNativeGQAAttention\nnative JAX baseline]
    E --> E2[TEDotProductAttention\nTE wrapper]
    E --> E3[MLA variant\nasymmetric Q/K vs V dim]

    G --> G1[create_packed_segment_ids_and_pos\npacked THD inputs]
    G --> G2[build_cp_mesh\nJAX Mesh + MeshResource]
    G --> G3[shard_for_context_parallel\nstripe reorder + device_put]
    G --> G4[fused_thd_attention\nRing / AllGather CP]

    I[test_attention.py] -->|deferred import| E
    I -->|deferred import under @requires_cp| G

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

jberchtold-nvidia · 2026-07-06T22:05:04Z


-**TODO — Coming soon.**
+This document walks through replacing a plain JAX implementation of BSHD
+attention with TransformerEngine's fused ``DotProductAttention``. The example


Thoughts on another sentence in here to mention TE/JAX attn has more features than shown in this tutorial?

Something like

TransformerEngine's DotProductAttention supports a variety of features. The full documentation can be found here "link". The tutorial example uses grouped-query attention (GQA) <https://arxiv.org/abs/2305.13245>_ and
sliding-window attention (SWA).

We can use the table we had in recent presentations (or just turn that into words to use here).

I am thinking of adding a sentence and API reference as @jberchtold-nvidia suggested
The table might change the scope of this documentation from samples to support matrix i feel. So that could be a separate effort that we can then link here in the samples once it is ready. Is that okay @cyanguwa ?

jberchtold-nvidia · 2026-07-06T22:09:21Z

+On a single GB200, this run is roughly **52x faster** for the fwd+bwd of this
+BSHD GQA + SWA example. This compares TE ``DotProductAttention`` against the
+native JAX baseline above, which materializes attention scores with XLA ops; it
+is not a comparison against ``jax.nn.dot_product_attention(...,


it is not a comparison against jax.nn.dot_product_attention(..., implementation="cudnn")

Does native JAX jax.nn.dot_product_attention with backend cudnn support SWA? My understanding is native JAX did not support SWA via cuDNN. If so, can you mention that as the reasoning for why we aren't comparing against jax.nn.dot_production_attention with cuDNN backend

The API does support SWA: https://docs.jax.dev/en/latest/_autosummary/jax.nn.dot_product_attention.html
I tried examples with the cuDNN backend separately to confirm that SWA is indeed supported

The motivation I had in mind for this sample/tutorial was to compare against native XLA backend and not the cuDNN backend (the comment here just helps clarify the same)

jberchtold-nvidia · 2026-07-06T22:35:47Z

+timing loop follows the same forward+backward pattern as ``speedometer`` while
+keeping those sharding controls visible.
+
+.. literalinclude:: attention_context_parallel.py


Should we add in the following entries to show where CP excels?

JAX attention jax.nn.dpa: OOM

TE DPA without CP: OOM

If you picked small/medium shapes for the example and they don't actually OOM without CP in these isolated examples, I think that's fine to keep as is without this suggestion

Can we add some speedup numbers as well, CP vs non-CP?

Okay let me add speed up numbers and tweak the config numbers to pick an appropriate config for the same

Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>

cyanguwa · 2026-07-22T17:14:07Z

-JAX: Attention with TransformerEngine
-=====================================
+JAX: BSHD Attention with TransformerEngine
+==========================================


How about keeping "JAX: Attention with TransformerEngine" and then use "JAX: Context-Parallel Attention with TransformerEngine" for the CP file? We can walk through some of the features non-CP/CP supports before getting into specific examples. I think jberchtold suggested this, but also, the last section of this file seems to cover "other features" in attention as well.

How about keeping "JAX: Attention with TransformerEngine" and then use "JAX: Context-Parallel Attention with TransformerEngine" for the CP file?

Sounds good

cyanguwa · 2026-07-22T17:15:06Z


-**TODO — Coming soon.**
+This document walks through replacing a plain JAX implementation of BSHD
+attention with TransformerEngine's fused ``DotProductAttention``. The example


We can use the table we had in recent presentations (or just turn that into words to use here).

cyanguwa · 2026-07-22T17:15:34Z

+This document walks through replacing a plain JAX implementation of BSHD
+attention with TransformerEngine's fused ``DotProductAttention``. The example
+uses `grouped-query attention (GQA) <https://arxiv.org/abs/2305.13245>`_ and
+sliding-window attention (SWA).


Maybe add a link to SWA as well, or remove the one for GQA for consistency?

cyanguwa · 2026-07-22T17:17:37Z


 `← Back to the JAX integration overview <../te_jax_integration.html>`_
+
+1. Baseline: native BSHD GQA + SWA


How about "1. Baseline: native-Jax BSHD GQA + SWA"? Also, how is the "unfused attn" path compare to this baseline? Should we use that path instead of reimplementing FlaxNativeGQAAttention here?

Sure I can do Baseline: native-xla BSHD GQA + SWA

I did think about this but the general flow of other sections seems to be: start with a native jax (non TE API) implementation and then show the user how to "port" it to TE so keeping that spirit in mind, I stayed away from using the TE unfused API

cyanguwa · 2026-07-22T17:20:30Z

+   :language: python
+   :start-after: # ATTENTION_IMPORTS_START
+   :end-before: # ATTENTION_IMPORTS_END
+


Maybe add a single line before the next code block starts? Same for the next block.

cyanguwa · 2026-07-22T17:25:47Z

+
+    See LICENSE for license information.
+
+JAX: Context-Parallel THD Attention with Transformer Engine


How about "JAX: Context-Parallel Attention with Transformer Engine" and

"CP supports BSHD, THD with blah (link to our docs), and we are going to demonstrate a few popular user configs..."?

cyanguwa · 2026-07-22T17:30:06Z

+``[batch, seq, heads, dim]``, and the sequence dimension can pack several
+shorter segments. The ``SequenceDescriptor`` tells TE which tokens belong to
+which packed segment and which token slots are padding. This tutorial uses
+sixteen padded segments per sequence and a 64k sequence length.


"sixteen padded segments per sequence, a 64k sequence length, and 2 sequences per batch."?

cyanguwa · 2026-07-22T17:30:28Z

+   :language: python
+   :start-after: # ATTENTION_CP_IMPORTS_START
+   :end-before: # ATTENTION_CP_IMPORTS_END
+


"The tensor inputs are created as follows."?

cyanguwa · 2026-07-22T17:34:06Z

+timing loop follows the same forward+backward pattern as ``speedometer`` while
+keeping those sharding controls visible.
+
+.. literalinclude:: attention_context_parallel.py


Can we add some speedup numbers as well, CP vs non-CP?

cyanguwa · 2026-07-22T17:34:42Z

+   :end-before: # ATTENTION_CP_MESH_END
+
+
+3. Fused THD attention call


"Fused attention call"?

Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>

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

Add JAX attention tutorials

7a55b38

Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>

KshitijLakhani self-assigned this Jul 1, 2026

KshitijLakhani added the 2.18 label Jul 1, 2026

KshitijLakhani marked this pull request as ready for review July 1, 2026 22:56

KshitijLakhani requested review from cyanguwa and jberchtold-nvidia July 1, 2026 22:56

KshitijLakhani changed the title ~~Add JAX attention tutorials~~ [JAX] Add attention tutorials Jul 1, 2026

greptile-apps Bot reviewed Jul 1, 2026

View reviewed changes

Comment thread docs/examples/jax/test_attention.py

Comment thread docs/examples/jax/attention.py

KshitijLakhani added the documentation Improvements or additions to documentation label Jul 2, 2026

jberchtold-nvidia reviewed Jul 6, 2026

View reviewed changes

KshitijLakhani added 3 commits July 7, 2026 18:07

Defer attention tutorial imports

03dc340

Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>

Pass query heads through attention wrapper

d67a84f

Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>

Clean up attention tutorial docs

ccea51f

Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>

cyanguwa reviewed Jul 22, 2026

View reviewed changes

KshitijLakhani added 2 commits July 27, 2026 09:40

Add JAX attention tutorial landing page

496b314

Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>

Refine context-parallel attention tutorial

82872d6

Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>

KshitijLakhani force-pushed the klakhani/feat/jax-attention-tutorials branch from eeaab44 to 82872d6 Compare July 27, 2026 16:49

Refine attention tutorial references

c8b9cc4

Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>

KshitijLakhani force-pushed the klakhani/feat/jax-attention-tutorials branch from 9908e58 to c8b9cc4 Compare July 27, 2026 17:33

cyanguwa previously approved these changes Jul 27, 2026

View reviewed changes

Skip large attention tests on older GPUs

a137c0e

Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>

KshitijLakhani dismissed cyanguwa’s stale review via a137c0e July 27, 2026 23:04

KshitijLakhani force-pushed the klakhani/feat/jax-attention-tutorials branch from 3d1fdbd to a137c0e Compare July 27, 2026 23:04

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

03c213b

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

KshitijLakhani requested review from cyanguwa and jberchtold-nvidia July 28, 2026 15:24

cyanguwa approved these changes Jul 28, 2026

View reviewed changes

jberchtold-nvidia approved these changes Jul 28, 2026

View reviewed changes

KshitijLakhani merged commit 6d3280c into NVIDIA:main Jul 28, 2026
10 of 14 checks passed


		`← Back to the JAX integration overview <../te_jax_integration.html>`_

		1. Baseline: native BSHD GQA + SWA


		See LICENSE for license information.

		JAX: Context-Parallel THD Attention with Transformer Engine

		:end-before: # ATTENTION_CP_MESH_END


		3. Fused THD attention call

Uh oh!

Conversation

KshitijLakhani commented Jul 1, 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 Jul 1, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Greptile Summary

Confidence Score: 5/5

Important Files Changed

Flowchart

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

KshitijLakhani Jul 23, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

KshitijLakhani commented Jul 1, 2026 •

edited

Loading

greptile-apps Bot commented Jul 1, 2026 •

edited

Loading

KshitijLakhani Jul 23, 2026 •

edited

Loading