From 7a55b3868373a1bf69c4f15a18630c8b83bd237a Mon Sep 17 00:00:00 2001 From: Kshitij Lakhani Date: Tue, 30 Jun 2026 18:50:45 -0700 Subject: [PATCH 1/2] Add JAX attention tutorials Signed-off-by: Kshitij Lakhani --- docs/examples/jax/attention.out | 13 + docs/examples/jax/attention.py | 290 ++++++++++++++ docs/examples/jax/attention.rst | 125 +++++- .../jax/attention_context_parallel.out | 13 + .../jax/attention_context_parallel.py | 378 ++++++++++++++++++ .../jax/attention_context_parallel.rst | 145 +++++++ docs/examples/jax/test_attention.py | 121 ++++++ docs/examples/te_jax_integration.rst | 8 +- 8 files changed, 1088 insertions(+), 5 deletions(-) create mode 100644 docs/examples/jax/attention.out create mode 100644 docs/examples/jax/attention.py create mode 100644 docs/examples/jax/attention_context_parallel.out create mode 100644 docs/examples/jax/attention_context_parallel.py create mode 100644 docs/examples/jax/attention_context_parallel.rst create mode 100644 docs/examples/jax/test_attention.py diff --git a/docs/examples/jax/attention.out b/docs/examples/jax/attention.out new file mode 100644 index 0000000000..058d39e7a4 --- /dev/null +++ b/docs/examples/jax/attention.out @@ -0,0 +1,13 @@ +# SINGLE_GPU_OUTPUT_START +Native JAX bf16 GQA + SWA: +Mean time: 5.109810829162598 ms + +TE DotProductAttention GQA + SWA: +Mean time: 0.09856224060058594 ms +# SINGLE_GPU_OUTPUT_END + +# MLA_OUTPUT_START +TE MLA-style BSHD: q/k head dim=128, v head dim=64 +Output shape=(2, 4096, 128, 64), dtype=bfloat16 +Grad shapes=[(2, 4096, 128, 128), (2, 4096, 8, 128), (2, 4096, 8, 64)] +# MLA_OUTPUT_END diff --git a/docs/examples/jax/attention.py b/docs/examples/jax/attention.py new file mode 100644 index 0000000000..86e446dbe9 --- /dev/null +++ b/docs/examples/jax/attention.py @@ -0,0 +1,290 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""JAX: BSHD attention with TransformerEngine. + +Companion source for ``attention.rst``. Code blocks between +``# ATTENTION_*_START`` / ``# ATTENTION_*_END`` markers are pulled into the RST +via ``literalinclude``. + +Run as a script to exercise the example end-to-end: + + python docs/examples/jax/attention.py +""" + +# ATTENTION_IMPORTS_START +from typing import Optional, Tuple + +import jax +import jax.numpy as jnp +import numpy as np +from flax import linen as nn + +import quickstart_jax_utils as utils + +from transformer_engine.jax.attention import SequenceDescriptor +from transformer_engine.jax.flax import DotProductAttention + +# ATTENTION_IMPORTS_END + + +# ATTENTION_INPUTS_START +batch, seq, num_query_heads, num_kv_heads, head_dim = 2, 4096, 128, 8, 128 +window_size = (128, 0) +dtype = jnp.bfloat16 +timing_iters = 20 +warmup_iters = 10 + + +def create_qkv_inputs( + *, + seed: int, + kv_heads: int = num_kv_heads, + qk_head_dim: int = head_dim, + v_head_dim: int = head_dim, +): + """Create separate BSHD query, key, value tensors and an output gradient.""" + + q_key, k_key, v_key, dout_key = jax.random.split(jax.random.PRNGKey(seed), 4) + q = jax.random.normal(q_key, (batch, seq, num_query_heads, qk_head_dim)).astype(dtype) + k = jax.random.normal(k_key, (batch, seq, kv_heads, qk_head_dim)).astype(dtype) + v = jax.random.normal(v_key, (batch, seq, kv_heads, v_head_dim)).astype(dtype) + dout = jax.random.normal(dout_key, (batch, seq, num_query_heads, v_head_dim)).astype(dtype) + return q, k, v, dout + + +def create_full_sequence_descriptor(): + """Describe a BSHD batch with no padding.""" + + seqlens = jnp.full((batch,), seq, dtype=jnp.int32) + return SequenceDescriptor.from_seqlens(seqlens) + + +q, k, v, dout = create_qkv_inputs(seed=2026) +qkv = (q, k, v) +sequence_descriptor = create_full_sequence_descriptor() +# ATTENTION_INPUTS_END + + +# ATTENTION_BASELINE_MODEL_START +def _repeat_kv_for_gqa(x, query_heads): + """Repeat each KV head across its group of query heads.""" + + repeats = query_heads // x.shape[2] + return jnp.repeat(x, repeats, axis=2) + + +def _make_causal_swa_mask(q_len, kv_len, window: Optional[Tuple[int, int]]): + """Create a boolean causal mask, optionally restricted to an SWA window.""" + + q_pos = jnp.arange(q_len)[:, None] + kv_pos = jnp.arange(kv_len)[None, :] + + if window is None: + return kv_pos <= q_pos + + left, right = window + allowed = kv_pos <= q_pos + right + if left >= 0: + allowed = allowed & (kv_pos >= q_pos - left) + return allowed + + +class FlaxNativeGQAAttention(nn.Module): + """Plain JAX/Flax GQA used as the bf16 baseline.""" + + window_size: Optional[Tuple[int, int]] = None + + @nn.compact + def __call__(self, qkv_tensors): + query, key, value = qkv_tensors + key = _repeat_kv_for_gqa(key, query.shape[2]) + value = _repeat_kv_for_gqa(value, query.shape[2]) + + scale = query.shape[-1] ** -0.5 + scores = jnp.einsum( + "bqhd,bkhd->bhqk", + query.astype(jnp.float32), + key.astype(jnp.float32), + ) + scores *= scale + + mask = _make_causal_swa_mask(query.shape[1], key.shape[1], self.window_size) + scores = jnp.where(mask[None, None, :, :], scores, jnp.finfo(jnp.float32).min) + probs = jax.nn.softmax(scores, axis=-1) + out = jnp.einsum("bhqk,bkhd->bqhd", probs, value.astype(jnp.float32)) + return out.astype(query.dtype) + + +baseline = FlaxNativeGQAAttention(window_size=window_size) +baseline_vars = baseline.init(jax.random.PRNGKey(2026), qkv) +# ATTENTION_BASELINE_MODEL_END + + +# ATTENTION_TE_MODEL_START +class TEDotProductAttention(nn.Module): + """Thin Flax wrapper around TE's DotProductAttention.""" + + num_kv_heads: int + qk_head_dim: int = head_dim + attn_mask_type: str = "causal" + qkv_layout: str = "bshd_bshd_bshd" + window_size: Optional[Tuple[int, int]] = None + + @nn.compact + def __call__( + self, + qkv_tensors, + sequence_descriptor: Optional[SequenceDescriptor] = None, + *, + deterministic: bool = False, + ): + query, key, value = qkv_tensors + return DotProductAttention( + head_dim=self.qk_head_dim, + num_attention_heads=num_query_heads, + num_gqa_groups=self.num_kv_heads, + attn_mask_type=self.attn_mask_type, + qkv_layout=self.qkv_layout, + attention_dropout=0.0, + transpose_batch_sequence=False, + window_size=self.window_size, + )( + query, + key, + value, + sequence_descriptor=sequence_descriptor, + deterministic=deterministic, + ) + + +te_model = TEDotProductAttention(num_kv_heads=num_kv_heads, window_size=window_size) +te_vars = te_model.init( + jax.random.PRNGKey(2026), + qkv, + sequence_descriptor=sequence_descriptor, + deterministic=False, +) +# ATTENTION_TE_MODEL_END + + +def run_forward_backward(model, variables, input_qkv, output_grad, seq_desc=None): + """Run one compiled forward+backward pass through an attention module.""" + + def loss_fn(qkv_arg): + if seq_desc is None: + out = model.apply(variables, qkv_arg) + else: + out = model.apply( + variables, + qkv_arg, + sequence_descriptor=seq_desc, + deterministic=False, + ) + return jnp.vdot(out.astype(jnp.float32), output_grad.astype(jnp.float32)) + + return jax.jit(jax.value_and_grad(loss_fn))(input_qkv) + + +def compare_te_to_baseline(input_qkv=qkv, output_grad=dout, seq_desc=sequence_descriptor): + """Compare the TE example to the native baseline.""" + + loss_ref, grads_ref = run_forward_backward( + baseline, baseline_vars, input_qkv, output_grad + ) + loss_te, grads_te = run_forward_backward(te_model, te_vars, input_qkv, output_grad, seq_desc) + out_ref = baseline.apply(baseline_vars, input_qkv) + out_te = te_model.apply(te_vars, input_qkv, sequence_descriptor=seq_desc, deterministic=False) + + jax.block_until_ready((loss_ref, grads_ref, loss_te, grads_te, out_ref, out_te)) + np.testing.assert_allclose(out_te, out_ref, rtol=5e-2, atol=5e-2) + for got, expected in zip(grads_te, grads_ref): + np.testing.assert_allclose(got, expected, rtol=8e-2, atol=8e-2) + + +# ATTENTION_SINGLE_GPU_BENCH_START +def run_single_gpu_bench(): + forward_kwargs = { + "sequence_descriptor": sequence_descriptor, + "deterministic": False, + } + + print("Native JAX bf16 GQA + SWA:") + utils.speedometer( + model_apply_fn=baseline.apply, + variables=baseline_vars, + input=qkv, + output_grad=dout, + timing_iters=timing_iters, + warmup_iters=warmup_iters, + ) + + print("\nTE DotProductAttention GQA + SWA:") + utils.speedometer( + model_apply_fn=te_model.apply, + variables=te_vars, + input=qkv, + output_grad=dout, + forward_kwargs=forward_kwargs, + timing_iters=timing_iters, + warmup_iters=warmup_iters, + ) + + +# ATTENTION_SINGLE_GPU_BENCH_END + + +# ATTENTION_MLA_START +mla_head_dim_qk, mla_head_dim_v = 128, 64 +mla_q, mla_k, mla_v, mla_dout = create_qkv_inputs( + seed=2027, + kv_heads=num_kv_heads, + qk_head_dim=mla_head_dim_qk, + v_head_dim=mla_head_dim_v, +) +mla_qkv = (mla_q, mla_k, mla_v) + +mla_model = TEDotProductAttention( + num_kv_heads=num_kv_heads, + qk_head_dim=mla_head_dim_qk, + window_size=None, +) +mla_vars = mla_model.init( + jax.random.PRNGKey(4), + mla_qkv, + sequence_descriptor=sequence_descriptor, + deterministic=False, +) + + +def run_mla_variant(): + out = mla_model.apply( + mla_vars, + mla_qkv, + sequence_descriptor=sequence_descriptor, + deterministic=False, + ) + loss, grads = run_forward_backward( + mla_model, mla_vars, mla_qkv, mla_dout, sequence_descriptor + ) + jax.block_until_ready((out, loss, grads)) + print( + "TE MLA-style BSHD: " + f"q/k head dim={mla_head_dim_qk}, v head dim={mla_head_dim_v}" + ) + print(f"Output shape={tuple(out.shape)}, dtype={out.dtype}") + print(f"Grad shapes={[tuple(grad.shape) for grad in grads]}") + + +# ATTENTION_MLA_END + + +if __name__ == "__main__": + print("# SINGLE_GPU_OUTPUT_START") + run_single_gpu_bench() + print("# SINGLE_GPU_OUTPUT_END") + + print("\n# MLA_OUTPUT_START") + run_mla_variant() + print("# MLA_OUTPUT_END") diff --git a/docs/examples/jax/attention.rst b/docs/examples/jax/attention.rst index c9f84da634..8d6de590a8 100644 --- a/docs/examples/jax/attention.rst +++ b/docs/examples/jax/attention.rst @@ -3,9 +3,128 @@ See LICENSE for license information. -JAX: Attention with TransformerEngine -===================================== +JAX: BSHD Attention with TransformerEngine +========================================== -**TODO — Coming soon.** +This document walks through replacing a plain JAX implementation of BSHD +attention with TransformerEngine's fused ``DotProductAttention``. The example +uses `grouped-query attention (GQA) `_ and +sliding-window attention (SWA). `← Back to the JAX integration overview <../te_jax_integration.html>`_ + +1. Baseline: native BSHD GQA + SWA +---------------------------------- + +The baseline keeps query, key, and value as separate BSHD tensors. GQA is modeled +by repeating each KV head across a group of query heads, then applying a causal +sliding-window mask before softmax. + +.. literalinclude:: attention.py + :language: python + :start-after: # ATTENTION_IMPORTS_START + :end-before: # ATTENTION_IMPORTS_END + +.. literalinclude:: attention.py + :language: python + :start-after: # ATTENTION_INPUTS_START + :end-before: # ATTENTION_INPUTS_END + +.. literalinclude:: attention.py + :language: python + :start-after: # ATTENTION_BASELINE_MODEL_START + :end-before: # ATTENTION_BASELINE_MODEL_END + + +2. Transformer Engine ``DotProductAttention`` +---------------------------------------------- + +The Transformer Engine version keeps the same separate BSHD inputs. The important arguments are +``num_gqa_groups`` for GQA, ``attn_mask_type="causal"`` for autoregressive +attention, and ``window_size`` for SWA. + +.. literalinclude:: attention.py + :language: python + :start-after: # ATTENTION_TE_MODEL_START + :end-before: # ATTENTION_TE_MODEL_END + + +3. Single-GPU performance +------------------------- + +``speedometer`` runs a JIT-compiled forward+backward loop with warmup for both +implementations. + +.. literalinclude:: attention.py + :language: python + :start-after: # ATTENTION_SINGLE_GPU_BENCH_START + :end-before: # ATTENTION_SINGLE_GPU_BENCH_END + +.. raw:: html + +
+ Output: +
+ +.. container:: program-output + + .. literalinclude:: attention.out + :language: text + :start-after: # SINGLE_GPU_OUTPUT_START + :end-before: # SINGLE_GPU_OUTPUT_END + +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(..., +implementation="cudnn")``. + + +4. MLA-style head dimensions +---------------------------- + +In TE/JAX, the simple MLA-style attention case is represented by separate Q, K, +and V tensors where Q/K and V use different per-head dimensions. Keep +``qkv_layout="bshd_bshd_bshd"`` so TE can see the Q/K head dimension and the V +head dimension separately. + +.. literalinclude:: attention.py + :language: python + :start-after: # ATTENTION_MLA_START + :end-before: # ATTENTION_MLA_END + +.. raw:: html + +
+ Output: +
+ +.. container:: program-output + + .. literalinclude:: attention.out + :language: text + :start-after: # MLA_OUTPUT_START + :end-before: # MLA_OUTPUT_END + + +Other attention knobs +--------------------- + +The examples above intentionally stay focused. Other ``DotProductAttention`` +features are enabled through the same module arguments: + +* Dropout: set ``attention_dropout > 0``, call with ``deterministic=False``, and + pass a Flax ``dropout`` RNG to ``apply``. +* Bias: pass ``bias`` and set ``attn_bias_type`` when the selected fused kernel + supports that bias mode. +* Sink attention: use ``softmax_type="off_by_one"`` or ``"learnable"``. +* Determinism: set ``NVTE_ALLOW_NONDETERMINISTIC_ALGO=0`` before launching the + process if deterministic fused kernels are required. + + +Next steps +---------- + +* `Context-parallel attention `_: packed THD + attention over a context-parallel mesh. +* `← Hub <../te_jax_integration.html>`_ diff --git a/docs/examples/jax/attention_context_parallel.out b/docs/examples/jax/attention_context_parallel.out new file mode 100644 index 0000000000..7c6bd7af38 --- /dev/null +++ b/docs/examples/jax/attention_context_parallel.out @@ -0,0 +1,13 @@ +# RING_OUTPUT_START +THD CP Ring stripe_size=1: +Mean time: 71.53654098510742 ms +Ring output shape=(2, 65536, 128, 128), dtype=bfloat16 +Ring grad shapes=[(2, 65536, 128, 128), (2, 65536, 128, 128), (2, 65536, 128, 128)] +# RING_OUTPUT_END + +# AG_OUTPUT_START +THD CP AllGather stripe_size=4096: +Mean time: 50.125694274902344 ms +AllGather output shape=(2, 65536, 128, 128), dtype=bfloat16 +AllGather grad shapes=[(2, 65536, 128, 128), (2, 65536, 128, 128), (2, 65536, 128, 128)] +# AG_OUTPUT_END diff --git a/docs/examples/jax/attention_context_parallel.py b/docs/examples/jax/attention_context_parallel.py new file mode 100644 index 0000000000..45b22b7327 --- /dev/null +++ b/docs/examples/jax/attention_context_parallel.py @@ -0,0 +1,378 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""JAX: context-parallel THD attention with Transformer Engine. + +Companion source for ``attention_context_parallel.rst``. Code blocks between +``# ATTENTION_CP_*_START`` / ``# ATTENTION_CP_*_END`` markers are pulled into +the RST via ``literalinclude``. + +Run as a script to exercise the example end-to-end: + + python docs/examples/jax/attention_context_parallel.py +""" + +# ATTENTION_CP_IMPORTS_START +import os +import time +from typing import Tuple + +# Ring + SWA uses the non-scan Ring implementation. Set this before JAX compiles +# the first fused attention call so the example follows the distributed tests. +os.environ.setdefault("NVTE_FUSED_RING_ATTENTION_USE_SCAN", "0") + +import jax +import jax.numpy as jnp +import numpy as np +from jax.sharding import Mesh, NamedSharding, PartitionSpec as P + +import transformer_engine.jax as te +from transformer_engine.jax.attention import ( + AttnBiasType, + AttnMaskType, + AttnSoftmaxType, + CPStrategy, + QKVLayout, + ReorderStrategy, + SequenceDescriptor, + fused_attn, + inverse_reorder_causal_load_balancing, + is_fused_attn_kernel_available, + reorder_causal_load_balancing, +) +from transformer_engine.jax.sharding import MeshResource + +# ATTENTION_CP_IMPORTS_END + + +# ATTENTION_CP_INPUTS_START +cp_size = 4 +batch, seq, num_attention_heads, head_dim = 2, 65536, 128, 128 +runtime_segments_per_seq = 16 +max_segments_per_seq = runtime_segments_per_seq +window_size = (128, 0) +dtype = jnp.bfloat16 +timing_iters = 5 +warmup_iters = 2 +ring_stripe_size = 1 +ag_stripe_size = 4096 + + +def create_qkv_inputs(seed: int = 2026): + q_key, k_key, v_key, dout_key = jax.random.split(jax.random.PRNGKey(seed), 4) + shape = (batch, seq, num_attention_heads, head_dim) + q = jax.random.normal(q_key, shape).astype(dtype) + k = jax.random.normal(k_key, shape).astype(dtype) + v = jax.random.normal(v_key, shape).astype(dtype) + dout = jax.random.normal(dout_key, shape).astype(dtype) + return q, k, v, dout + + +def create_packed_segment_ids_and_pos(): + """Pack padded causal segments into each THD batch row.""" + + segment_slot_len = seq // runtime_segments_per_seq + valid_segment_len = 3 * segment_slot_len // 4 + segment_ids_per_row = [] + segment_pos_per_row = [] + + for segment_id in range(1, runtime_segments_per_seq + 1): + valid_ids = jnp.full((valid_segment_len,), segment_id, dtype=jnp.int32) + padded_ids = jnp.zeros((segment_slot_len - valid_segment_len,), dtype=jnp.int32) + segment_ids_per_row.append(jnp.concatenate([valid_ids, padded_ids])) + segment_pos_per_row.append(jnp.arange(segment_slot_len, dtype=jnp.int32)) + + segment_ids = jnp.concatenate(segment_ids_per_row) + segment_pos = jnp.concatenate(segment_pos_per_row) + segment_ids = jnp.tile(segment_ids[None, :], (batch, 1)) + segment_pos = jnp.tile(segment_pos[None, :], (batch, 1)) + return segment_ids, segment_pos + + +def create_sequence_descriptor(segment_ids_arg, segment_pos_arg): + """Create the THD sequence descriptor from segment IDs and positions.""" + + return SequenceDescriptor.from_segment_ids_and_pos(segment_ids_arg, segment_pos_arg) + + +q, k, v, dout = create_qkv_inputs() +segment_ids, segment_pos = create_packed_segment_ids_and_pos() +sequence_descriptor = create_sequence_descriptor(segment_ids, segment_pos) +# ATTENTION_CP_INPUTS_END + + +# ATTENTION_CP_MESH_START +def build_cp_mesh(): + """Use one JAX mesh axis for context parallelism over sequence.""" + + devices = np.asarray(jax.devices()[:cp_size]) + mesh = Mesh(devices, axis_names=("cp",)) + mesh_resource = MeshResource(cp_resource="cp") + return mesh, mesh_resource + + +# ATTENTION_CP_MESH_END + + +# ATTENTION_CP_FUSED_ATTENTION_START +def fused_thd_attention( + qkv_tensors, + seq_desc, + *, + context_parallel_axis: str = "", + context_parallel_strategy: CPStrategy = CPStrategy.DEFAULT, + context_parallel_causal_load_balanced: bool = False, + stripe_size: int | None = None, +): + """Call TE fused attention on separate THD Q, K, V tensors.""" + + return fused_attn( + qkv_tensors, + None, + seq_desc, + None, + attn_bias_type=AttnBiasType.NO_BIAS, + attn_mask_type=AttnMaskType.PADDING_CAUSAL_MASK, + qkv_layout=QKVLayout.THD_THD_THD, + softmax_type=AttnSoftmaxType.VANILLA_SOFTMAX, + scaling_factor=head_dim**-0.5, + dropout_probability=0.0, + is_training=True, + max_segments_per_seq=max_segments_per_seq, + window_size=window_size, + context_parallel_strategy=context_parallel_strategy, + context_parallel_causal_load_balanced=context_parallel_causal_load_balanced, + context_parallel_axis=context_parallel_axis, + stripe_size=stripe_size, + ) + + +def apply_context_parallel_attention( + _variables, + qkv_tensors, + *, + seq_desc, + context_parallel_strategy: CPStrategy, + stripe_size: int, + rngs=None, +): + del rngs + return fused_thd_attention( + qkv_tensors, + seq_desc, + context_parallel_axis="cp", + context_parallel_strategy=context_parallel_strategy, + context_parallel_causal_load_balanced=True, + stripe_size=stripe_size, + ) + + +# ATTENTION_CP_FUSED_ATTENTION_END + + +# ATTENTION_CP_REORDER_START +def reorder_for_context_parallel(x, stripe_size: int): + return reorder_causal_load_balancing( + x, + strategy=ReorderStrategy.Striped, + cp_size=cp_size, + seq_dim=1, + stripe_size=stripe_size, + ) + + +def inverse_reorder_from_context_parallel(x, stripe_size: int): + return inverse_reorder_causal_load_balancing( + x, + strategy=ReorderStrategy.Striped, + cp_size=cp_size, + seq_dim=1, + stripe_size=stripe_size, + ) + + +def create_reordered_sequence_descriptor(stripe_size: int): + reordered_ids = reorder_for_context_parallel(segment_ids, stripe_size) + reordered_pos = reorder_for_context_parallel(segment_pos, stripe_size) + return create_sequence_descriptor(reordered_ids, reordered_pos) + + +# ATTENTION_CP_REORDER_END + + +# ATTENTION_CP_SHARD_START +def shard_sequence_descriptor(mesh, seq_desc): + def put_leaf(x): + if x.ndim == 1: + sharding = NamedSharding(mesh, P(None)) + else: + sharding = NamedSharding(mesh, P(None, "cp")) + return jax.device_put(x, sharding) + + return jax.tree.map(put_leaf, seq_desc) + + +def shard_for_context_parallel(mesh, stripe_size: int): + qkv_sharding = NamedSharding(mesh, P(None, "cp", None, None)) + dout_sharding = NamedSharding(mesh, P(None, "cp", None, None)) + reordered_seq_desc = create_reordered_sequence_descriptor(stripe_size) + + return { + "qkv": tuple( + jax.device_put(reorder_for_context_parallel(x, stripe_size), qkv_sharding) + for x in (q, k, v) + ), + "dout": jax.device_put(reorder_for_context_parallel(dout, stripe_size), dout_sharding), + "sequence_descriptor": shard_sequence_descriptor(mesh, reordered_seq_desc), + } + + +# ATTENTION_CP_SHARD_END + + +def _strategy_name(strategy: CPStrategy): + return "Ring" if strategy == CPStrategy.RING else "AllGather" + + +def context_parallel_supported() -> Tuple[bool, str]: + if len(jax.devices()) < cp_size: + return False, f"needs {cp_size} GPUs" + + has_kernel = is_fused_attn_kernel_available( + True, + dtype, + dtype, + QKVLayout.THD_THD_THD, + AttnBiasType.NO_BIAS, + AttnMaskType.PADDING_CAUSAL_MASK, + AttnSoftmaxType.VANILLA_SOFTMAX, + 0.0, + num_attention_heads, + num_attention_heads, + seq, + seq, + head_dim, + head_dim, + window_size, + ) + if not has_kernel: + return False, "no fused attention kernel for the THD SWA shape" + return True, "" + + +def run_reference_attention(): + out = fused_thd_attention((q, k, v), sequence_descriptor) + return jax.block_until_ready(out) + + +# ATTENTION_CP_RUN_START +def _context_parallel_jit_fns(strategy: CPStrategy, stripe_size: int, sharded): + qkv_shardings = tuple(x.sharding for x in sharded["qkv"]) + seq_desc_shardings = jax.tree.map(lambda x: x.sharding, sharded["sequence_descriptor"]) + dout_sharding = sharded["dout"].sharding + + def loss_fn(qkv_arg, seq_desc_arg, dout_arg): + out = apply_context_parallel_attention( + {}, + qkv_arg, + seq_desc=seq_desc_arg, + context_parallel_strategy=strategy, + stripe_size=stripe_size, + ) + return jnp.vdot(out.astype(jnp.float32), dout_arg.astype(jnp.float32)) + + def forward_fn(qkv_arg, seq_desc_arg): + out = apply_context_parallel_attention( + {}, + qkv_arg, + seq_desc=seq_desc_arg, + context_parallel_strategy=strategy, + stripe_size=stripe_size, + ) + return inverse_reorder_from_context_parallel(out, stripe_size) + + grad_fn = jax.jit( + jax.value_and_grad(loss_fn), + in_shardings=(qkv_shardings, seq_desc_shardings, dout_sharding), + out_shardings=(None, qkv_shardings), + ) + forward_jit = jax.jit( + forward_fn, + in_shardings=(qkv_shardings, seq_desc_shardings), + ) + return grad_fn, forward_jit + + +def run_context_parallel_case(strategy: CPStrategy, stripe_size: int): + mesh, mesh_resource = build_cp_mesh() + sharded = shard_for_context_parallel(mesh, stripe_size) + grad_fn, forward_jit = _context_parallel_jit_fns(strategy, stripe_size, sharded) + + with jax.set_mesh(mesh), te.autocast(mesh_resource=mesh_resource): + loss, grads = grad_fn( + sharded["qkv"], + sharded["sequence_descriptor"], + sharded["dout"], + ) + out = forward_jit(sharded["qkv"], sharded["sequence_descriptor"]) + + jax.block_until_ready((loss, grads, out)) + return {"loss": loss, "grads": grads, "output": out} + + +def run_context_parallel_bench(strategy: CPStrategy, stripe_size: int): + mesh, mesh_resource = build_cp_mesh() + sharded = shard_for_context_parallel(mesh, stripe_size) + grad_fn, _ = _context_parallel_jit_fns(strategy, stripe_size, sharded) + + print(f"THD CP {_strategy_name(strategy)} stripe_size={stripe_size}:") + with jax.set_mesh(mesh), te.autocast(mesh_resource=mesh_resource): + for _ in range(warmup_iters): + result = grad_fn( + sharded["qkv"], + sharded["sequence_descriptor"], + sharded["dout"], + ) + jax.block_until_ready(result) + + start = time.time() + for _ in range(timing_iters): + result = grad_fn( + sharded["qkv"], + sharded["sequence_descriptor"], + sharded["dout"], + ) + jax.block_until_ready(result) + end = time.time() + + print(f"Mean time: {(end - start) * 1000 / timing_iters} ms") + + +# ATTENTION_CP_RUN_END + + +if __name__ == "__main__": + supported, reason = context_parallel_supported() + if not supported: + print(f"skipped context-parallel example: {reason}") + else: + print("# RING_OUTPUT_START") + run_context_parallel_bench(CPStrategy.RING, ring_stripe_size) + ring_result = run_context_parallel_case(CPStrategy.RING, ring_stripe_size) + print( + "Ring output shape=" + f"{tuple(ring_result['output'].shape)}, dtype={ring_result['output'].dtype}" + ) + print(f"Ring grad shapes={[tuple(grad.shape) for grad in ring_result['grads']]}") + print("# RING_OUTPUT_END") + + print("\n# AG_OUTPUT_START") + run_context_parallel_bench(CPStrategy.ALL_GATHER, ag_stripe_size) + ag_result = run_context_parallel_case(CPStrategy.ALL_GATHER, ag_stripe_size) + print( + "AllGather output shape=" + f"{tuple(ag_result['output'].shape)}, dtype={ag_result['output'].dtype}" + ) + print(f"AllGather grad shapes={[tuple(grad.shape) for grad in ag_result['grads']]}") + print("# AG_OUTPUT_END") diff --git a/docs/examples/jax/attention_context_parallel.rst b/docs/examples/jax/attention_context_parallel.rst new file mode 100644 index 0000000000..c865e2c62b --- /dev/null +++ b/docs/examples/jax/attention_context_parallel.rst @@ -0,0 +1,145 @@ +.. + Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + + See LICENSE for license information. + +JAX: Context-Parallel THD Attention with Transformer Engine +=========================================================== + +This document demonstrates context parallelism (CP) with packed THD attention. +CP shards the sequence dimension over a JAX mesh axis so long-context attention +can split activation memory and attention work across devices while +Transformer Engine (TE) runs the required collectives inside the fused attention call. + +CP is most useful when sequence length is large enough that single-device +attention becomes memory- or latency-limited. It is usually not worth adding for +short sequences or small local batches where communication overhead can dominate +the attention work. + +**Prerequisite:** this example requires four GPUs. + +`← Back to the JAX integration overview <../te_jax_integration.html>`_ + +1. Packed THD inputs +-------------------- + +In the separate-QKV THD layout used here, Q/K/V are shaped +``[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. + +.. literalinclude:: attention_context_parallel.py + :language: python + :start-after: # ATTENTION_CP_IMPORTS_START + :end-before: # ATTENTION_CP_IMPORTS_END + +.. literalinclude:: attention_context_parallel.py + :language: python + :start-after: # ATTENTION_CP_INPUTS_START + :end-before: # ATTENTION_CP_INPUTS_END + + +2. Context-parallel mesh +------------------------ + +The JAX ``Mesh`` describes the physical devices. ``MeshResource`` tells TE which +mesh axis is used for context parallelism. + +.. literalinclude:: attention_context_parallel.py + :language: python + :start-after: # ATTENTION_CP_MESH_START + :end-before: # ATTENTION_CP_MESH_END + + +3. Fused THD attention call +--------------------------- + +This example calls ``transformer_engine.jax.attention.fused_attn`` directly. The +Flax ``DotProductAttention`` wrapper covers the common path, but the lower-level +function exposes ``stripe_size``. This tutorial uses ``fused_attn`` directly +because stripe-size tuning matters for CP + THD striped load balancing. + +.. literalinclude:: attention_context_parallel.py + :language: python + :start-after: # ATTENTION_CP_FUSED_ATTENTION_START + :end-before: # ATTENTION_CP_FUSED_ATTENTION_END + + +4. Striped load balancing and sharding +-------------------------------------- + +For THD causal CP, TE uses striped load balancing. Ring attention requires +``stripe_size=1``. AllGather can use a larger stripe size; this tutorial uses +``stripe_size=4096`` for the 64k sequence shape. Ring + SWA uses the non-scan +Ring path, set in the example before the first fused attention call is compiled. + +.. literalinclude:: attention_context_parallel.py + :language: python + :start-after: # ATTENTION_CP_REORDER_START + :end-before: # ATTENTION_CP_REORDER_END + +.. literalinclude:: attention_context_parallel.py + :language: python + :start-after: # ATTENTION_CP_SHARD_START + :end-before: # ATTENTION_CP_SHARD_END + + +5. Ring and AllGather +--------------------- + +Both examples use packed THD, causal masking, SWA, and dropout-free fused +attention. The only strategy-specific difference is the CP strategy and stripe +size. CP collectives depend on the compiler seeing the intended sharding, so the +forward and forward+backward functions are compiled with explicit +``in_shardings``; the forward+backward path also pins the gradient sharding. The +timing loop follows the same forward+backward pattern as ``speedometer`` while +keeping those sharding controls visible. + +.. literalinclude:: attention_context_parallel.py + :language: python + :start-after: # ATTENTION_CP_RUN_START + :end-before: # ATTENTION_CP_RUN_END + +.. raw:: html + +
+ Ring output: +
+ +.. container:: program-output + + .. literalinclude:: attention_context_parallel.out + :language: text + :start-after: # RING_OUTPUT_START + :end-before: # RING_OUTPUT_END + +.. raw:: html + +
+ AllGather output: +
+ +.. container:: program-output + + .. literalinclude:: attention_context_parallel.out + :language: text + :start-after: # AG_OUTPUT_START + :end-before: # AG_OUTPUT_END + + +Other attention knobs +--------------------- + +CP supports a narrower set of feature combinations than non-CP attention. These +examples intentionally use no bias, no dropout, and vanilla softmax. Enable other +features one at a time and keep the CP strategy, layout, mask, and sequence +descriptor choices explicit when debugging unsupported combinations. + + +Next steps +---------- + +* `BSHD attention `_: single-GPU BSHD GQA, SWA, and MLA-style + head dimensions. +* `← Hub <../te_jax_integration.html>`_ diff --git a/docs/examples/jax/test_attention.py b/docs/examples/jax/test_attention.py new file mode 100644 index 0000000000..0c19cf1492 --- /dev/null +++ b/docs/examples/jax/test_attention.py @@ -0,0 +1,121 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Pytest entry points for the JAX attention tutorials.""" + +import jax +import numpy as np +import pytest + +import attention +import attention_context_parallel as cp_attention + + +def test_bshd_gqa_swa_runs(): + out = attention.te_model.apply( + attention.te_vars, + attention.qkv, + sequence_descriptor=attention.sequence_descriptor, + deterministic=False, + ) + + assert out.shape == attention.dout.shape + assert out.dtype == attention.dtype + + +def test_bshd_gqa_swa_matches_baseline(): + attention.compare_te_to_baseline() + + +def test_single_gpu_benchmark(): + attention.run_single_gpu_bench() + + +def test_mla_variant_runs(): + out = attention.mla_model.apply( + attention.mla_vars, + attention.mla_qkv, + sequence_descriptor=attention.sequence_descriptor, + deterministic=False, + ) + loss, grads = attention.run_forward_backward( + attention.mla_model, + attention.mla_vars, + attention.mla_qkv, + attention.mla_dout, + attention.sequence_descriptor, + ) + jax.block_until_ready((out, loss, grads)) + + assert out.shape == attention.mla_dout.shape + assert out.dtype == attention.dtype + assert loss.shape == () + assert [grad.shape for grad in grads] == [x.shape for x in attention.mla_qkv] + + +_cp_supported, _cp_reason = cp_attention.context_parallel_supported() +requires_cp = pytest.mark.skipif( + not _cp_supported, + reason=f"context-parallel attention tutorial skipped: {_cp_reason}", +) + + +def _assert_cp_result(strategy, stripe_size): + result = cp_attention.run_context_parallel_case(strategy, stripe_size) + reference = cp_attention.run_reference_attention() + + assert result["output"].shape == ( + cp_attention.batch, + cp_attention.seq, + cp_attention.num_attention_heads, + cp_attention.head_dim, + ) + assert result["output"].dtype == cp_attention.dtype + assert result["loss"].shape == () + assert [grad.shape for grad in result["grads"]] == [ + x.shape for x in cp_attention.create_qkv_inputs()[:3] + ] + + valid_tokens = cp_attention.segment_ids.astype(bool)[..., None, None] + valid_diff = jax.numpy.max( + jax.numpy.where( + valid_tokens, + jax.numpy.abs( + result["output"].astype(jax.numpy.float32) + - reference.astype(jax.numpy.float32) + ), + 0.0, + ) + ) + padded_max = jax.numpy.max( + jax.numpy.where( + valid_tokens, + 0.0, + jax.numpy.abs(result["output"].astype(jax.numpy.float32)), + ) + ) + np.testing.assert_allclose(valid_diff, 0, rtol=5e-2, atol=5e-2) + np.testing.assert_allclose(padded_max, 0, rtol=5e-2, atol=5e-2) + + +@requires_cp +def test_multi_gpu_context_parallel_ring_case(): + _assert_cp_result(cp_attention.CPStrategy.RING, cp_attention.ring_stripe_size) + + +@requires_cp +def test_multi_gpu_context_parallel_allgather_case(): + _assert_cp_result(cp_attention.CPStrategy.ALL_GATHER, cp_attention.ag_stripe_size) + + +@requires_cp +def test_multi_gpu_context_parallel_benchmarks(): + cp_attention.run_context_parallel_bench( + cp_attention.CPStrategy.RING, + cp_attention.ring_stripe_size, + ) + cp_attention.run_context_parallel_bench( + cp_attention.CPStrategy.ALL_GATHER, + cp_attention.ag_stripe_size, + ) diff --git a/docs/examples/te_jax_integration.rst b/docs/examples/te_jax_integration.rst index a15a10e0b3..82bb1f910e 100644 --- a/docs/examples/te_jax_integration.rst +++ b/docs/examples/te_jax_integration.rst @@ -28,8 +28,11 @@ Pick a topic - *Coming soon* - * - `Attention `_ - - *Coming soon* - - + - **Available** + - BSHD GQA + SWA; single-GPU perf comparison; MLA-style head dimensions + * - `Context-Parallel Attention `_ + - **Available** + - Packed THD attention; Ring and AllGather CP; striped load balancing; SWA * - `Expert Parallelism `_ - *Coming soon* - @@ -92,4 +95,5 @@ Conventions used across these documents jax/dense jax/collective_gemm jax/attention + jax/attention_context_parallel jax/expert_parallelism From eeaab44e17d11427ec8eb16064ef57827e312023 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Wed, 1 Jul 2026 16:30:59 +0000 Subject: [PATCH 2/2] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- docs/examples/jax/attention.py | 13 +++---------- docs/examples/jax/test_attention.py | 3 +-- 2 files changed, 4 insertions(+), 12 deletions(-) diff --git a/docs/examples/jax/attention.py b/docs/examples/jax/attention.py index 86e446dbe9..038d439277 100644 --- a/docs/examples/jax/attention.py +++ b/docs/examples/jax/attention.py @@ -190,9 +190,7 @@ def loss_fn(qkv_arg): def compare_te_to_baseline(input_qkv=qkv, output_grad=dout, seq_desc=sequence_descriptor): """Compare the TE example to the native baseline.""" - loss_ref, grads_ref = run_forward_backward( - baseline, baseline_vars, input_qkv, output_grad - ) + loss_ref, grads_ref = run_forward_backward(baseline, baseline_vars, input_qkv, output_grad) loss_te, grads_te = run_forward_backward(te_model, te_vars, input_qkv, output_grad, seq_desc) out_ref = baseline.apply(baseline_vars, input_qkv) out_te = te_model.apply(te_vars, input_qkv, sequence_descriptor=seq_desc, deterministic=False) @@ -265,14 +263,9 @@ def run_mla_variant(): sequence_descriptor=sequence_descriptor, deterministic=False, ) - loss, grads = run_forward_backward( - mla_model, mla_vars, mla_qkv, mla_dout, sequence_descriptor - ) + loss, grads = run_forward_backward(mla_model, mla_vars, mla_qkv, mla_dout, sequence_descriptor) jax.block_until_ready((out, loss, grads)) - print( - "TE MLA-style BSHD: " - f"q/k head dim={mla_head_dim_qk}, v head dim={mla_head_dim_v}" - ) + print(f"TE MLA-style BSHD: q/k head dim={mla_head_dim_qk}, v head dim={mla_head_dim_v}") print(f"Output shape={tuple(out.shape)}, dtype={out.dtype}") print(f"Grad shapes={[tuple(grad.shape) for grad in grads]}") diff --git a/docs/examples/jax/test_attention.py b/docs/examples/jax/test_attention.py index 0c19cf1492..c299f95dad 100644 --- a/docs/examples/jax/test_attention.py +++ b/docs/examples/jax/test_attention.py @@ -82,8 +82,7 @@ def _assert_cp_result(strategy, stripe_size): jax.numpy.where( valid_tokens, jax.numpy.abs( - result["output"].astype(jax.numpy.float32) - - reference.astype(jax.numpy.float32) + result["output"].astype(jax.numpy.float32) - reference.astype(jax.numpy.float32) ), 0.0, )