JAX: Context-Parallel Attention with TransformerEngine

Transformer Engine fused attention supports context parallelism (CP) for selected BSHD and packed THD Q/K/V layouts; see the JAX DotProductAttention API reference for the layout definitions and full interface. This tutorial focuses on a representative packed THD configuration with grouped-query attention (GQA), padded segments, causal sliding-window attention (SWA), and both Ring and AllGather strategies.

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.

Note

CP is most useful when attention does not fit on one GPU, or when long sequences and sufficiently wide attention windows provide enough computation to amortize communication. For instance, applications that use GQA may be good candidates for CP because GQA’s lower K/V head count reduces communication across devices. Conversely, workloads with narrow SWA windows may be better suited to single-GPU fused attention: CP still communicates K/V across devices while each query attends to relatively few tokens.

Prerequisite: this example requires four GPUs.

← Back to the Attention overview

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. CP supports separate Q/K/V (THD_THD_THD) and packed K/V (THD_T2HD) layouts, but not fully packed QKV (T3HD); this tutorial uses separate tensors. It uses a batch of two 64k sequences. Each sequence contains four padded, 16k-capacity segment slots with 12,288 valid tokens and 4,096 padding tokens per slot. It also uses GQA with 128 query heads and 8 K/V heads.

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

The tensor inputs and packed-sequence metadata are created as follows.

cp_size = 4
batch, seq, num_query_heads, num_kv_heads, head_dim = 2, 65536, 128, 8, 128
runtime_segments_per_seq = 4
max_segments_per_seq = runtime_segments_per_seq
window_size = (8192, 0)
dtype = jnp.bfloat16
timing_iters = 5
warmup_iters = 2
ring_stripe_size = 1
ag_stripe_size = 512


def create_qkv_inputs(seed: int = 2026):
    """Create separate THD GQA tensors and an output gradient."""

    q_key, k_key, v_key, dout_key = jax.random.split(jax.random.PRNGKey(seed), 4)
    q_shape = (batch, seq, num_query_heads, head_dim)
    kv_shape = (batch, seq, num_kv_heads, head_dim)
    q = jax.random.normal(q_key, q_shape).astype(dtype)
    k = jax.random.normal(k_key, kv_shape).astype(dtype)
    v = jax.random.normal(v_key, kv_shape).astype(dtype)
    dout = jax.random.normal(dout_key, q_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)

2. Context-parallel mesh

The JAX Mesh describes the physical devices. MeshResource tells TE which mesh axis is used for context parallelism.

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",))
    # Also set the corresponding MeshResource fields when other parallelisms
    # use additional mesh axis names in the surrounding model.
    mesh_resource = MeshResource(cp_resource="cp")
    return mesh, mesh_resource

3. Fused 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.

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,
    )

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=512 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.

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)

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),
    }

5. Ring and AllGather

The single-GPU baseline and both CP examples use the same packed THD GQA shape, causal masking, 8192-token SWA window, and dropout-free fused attention. The only strategy-specific difference between the two CP cases is the 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.

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_single_gpu_bench():
    grad_fn = _single_gpu_grad_fn()

    print("Single-GPU THD GQA + SWA:")
    for _ in range(warmup_iters):
        result = grad_fn((q, k, v), sequence_descriptor, dout)
    jax.block_until_ready(result)

    start = time.time()
    for _ in range(timing_iters):
        result = grad_fn((q, k, v), sequence_descriptor, dout)
    jax.block_until_ready(result)
    mean_ms = (time.time() - start) * 1000 / timing_iters
    print(f"Mean time: {mean_ms} ms")
    return mean_ms


def run_context_parallel_bench(
    strategy: CPStrategy,
    stripe_size: int,
    single_gpu_ms: float | None = None,
):
    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()

    mean_ms = (end - start) * 1000 / timing_iters
    print(f"Mean time: {mean_ms} ms")
    if single_gpu_ms is not None:
        print(f"Speedup vs single GPU: {single_gpu_ms / mean_ms:.2f}x")

Single-GPU output:

Single-GPU THD GQA + SWA:
Mean time: 126.68747901916504 ms

Ring output:

THD CP Ring stripe_size=1:
Mean time: 57.16729164123535 ms
Speedup vs single GPU: 2.22x

AllGather output:

THD CP AllGather stripe_size=512:
Mean time: 53.79219055175781 ms
Speedup vs single GPU: 2.36x

On four GB200s, Ring is roughly 2.22x faster and AllGather roughly 2.36x faster than the equivalent single-GPU fused-attention forward+backward pass. These results are specific to this workload and system. Applications with long segments and wide attention windows generally have more attention computation relative to communication and are stronger CP candidates; workloads with short windows may see less benefit. Performance also depends on the batch and segment lengths, head configuration, CP strategy, stripe size, and interconnect.

Next steps

• Single-GPU attention: BSHD GQA, SWA, and DeepSeek-style MLA head dimensions.

• ← Attention overview

• ← Hub