From c068937844106f35c0a5a4411e659d95553b37bf Mon Sep 17 00:00:00 2001 From: "Alex Y. Chan" Date: Sun, 1 Nov 2026 16:15:17 +0000 Subject: [PATCH 1/2] [JAX] Compare the attention tutorial to its baseline on a shorter sequence The native baseline in the attention tutorial materializes fp32 scores, so at the tutorial's 4K sequence its backward pass needs a ~49 GiB buffer. That does not fit when CI runs several test workers per GPU, and test_bshd_gqa_swa_matches_baseline fails with RESOURCE_EXHAUSTED. Compare TE to the baseline on a 1K prefix instead. Signed-off-by: Alex Y. Chan --- docs/examples/jax/test_attention.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/docs/examples/jax/test_attention.py b/docs/examples/jax/test_attention.py index 15321e85da3..c7f03c62518 100644 --- a/docs/examples/jax/test_attention.py +++ b/docs/examples/jax/test_attention.py @@ -14,6 +14,7 @@ AttnMaskType, AttnSoftmaxType, QKVLayout, + SequenceDescriptor, is_fused_attn_kernel_available, ) from transformer_engine_jax import get_device_compute_capability @@ -47,7 +48,13 @@ def test_bshd_gqa_swa_runs(): def test_bshd_gqa_swa_matches_baseline(): import attention - attention.compare_te_to_baseline() + # The native baseline materializes fp32 attention scores, so at the tutorial's 4K + # sequence its backward pass needs a ~49 GiB buffer, which does not fit when CI runs + # several test workers per GPU. Compare on a 1K prefix instead (1/16 of the memory). + seq = 1024 + qkv = tuple(t[:, :seq] for t in attention.qkv) + seq_desc = SequenceDescriptor.from_seqlens(jnp.full((attention.batch,), seq, dtype=jnp.int32)) + attention.compare_te_to_baseline(qkv, attention.dout[:, :seq], seq_desc) @requires_hopper_or_newer From df78bf7d553531ea6a1583e2bdbf3f80f10acbba Mon Sep 17 00:00:00 2001 From: "Alex Y. Chan" Date: Fri, 6 Nov 2026 20:04:53 +0000 Subject: [PATCH 2/2] [JAX] Generalize the attention tutorial test comments Describe the native baseline's backward buffer without a specific size, since the allocation XLA makes can change between versions, and note that the Hopper skip reason refers to the 4K baseline rather than the 1K comparison. Signed-off-by: Alex Y. Chan --- docs/examples/jax/test_attention.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/docs/examples/jax/test_attention.py b/docs/examples/jax/test_attention.py index c7f03c62518..c86631dbc2a 100644 --- a/docs/examples/jax/test_attention.py +++ b/docs/examples/jax/test_attention.py @@ -44,13 +44,14 @@ def test_bshd_gqa_swa_runs(): assert out.dtype == attention.dtype +# The skip reason refers to the 4K baseline; the 1K comparison keeps the gate unchanged. @requires_hopper_or_newer def test_bshd_gqa_swa_matches_baseline(): import attention # The native baseline materializes fp32 attention scores, so at the tutorial's 4K - # sequence its backward pass needs a ~49 GiB buffer, which does not fit when CI runs - # several test workers per GPU. Compare on a 1K prefix instead (1/16 of the memory). + # sequence its backward pass needs a much larger buffer, which does not fit when CI runs + # several test workers per GPU. Compare on a smaller prefix instead seq = 1024 qkv = tuple(t[:, :seq] for t in attention.qkv) seq_desc = SequenceDescriptor.from_seqlens(jnp.full((attention.batch,), seq, dtype=jnp.int32))