JAX: Single-GPU Attention with TransformerEngine

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 Attention overview

1. Baseline: native JAX BSHD GQA + SWA

Start with the imports shared by the native JAX and Transformer Engine implementations.

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

Next, create reproducible BSHD Q/K/V tensors and the sequence descriptor. The SequenceDescriptor supplies TE with sequence lengths and, for packed inputs, segment boundaries and padding metadata.

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

The native JAX baseline repeats K/V heads for GQA and applies the causal sliding-window mask explicitly.

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)

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.

class TEDotProductAttention(nn.Module):
    """Thin Flax wrapper around TE's DotProductAttention."""

    num_query_heads: int
    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=self.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_query_heads=num_query_heads,
    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,
)

3. Single-GPU performance

speedometer runs a JIT-compiled forward+backward loop with warmup for both implementations.

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,
    )
Output: 
Native JAX bf16 GQA + SWA:
Mean time: 5.109810829162598 ms

TE DotProductAttention GQA + SWA:
Mean time: 0.09856224060058594 ms

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. DeepSeek-style MLA head dimensions

This example covers the attention-kernel interface used after DeepSeek-style MLA projections, not the latent projection layers themselves. At this point, separate Q, K, and V tensors can use different per-head dimensions for Q/K and V. Keep qkv_layout="bshd_bshd_bshd" so TE can see the Q/K head dimension and the V head dimension separately.

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_query_heads=num_query_heads,
    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 DeepSeek-style MLA head dimensions: "
        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]}")
Output: 
TE DeepSeek-style MLA head dimensions: 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)]

Other attention knobs

The examples above represent a subset of attention features. Other DotProductAttention features can be enabled through the same module arguments as below:

  • 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".

  • Score scaling: set scale_factor to override the default 1 / sqrt(head_dim) scaling.

  • Determinism: set NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 before launching the process if deterministic fused kernels are required.

  • Score modification (experimental): use score_mod for a FlexAttention-style cuDNN frontend callback, with runtime operands in score_mod_tensors and optional custom backward logic in score_mod_bprop. This path requires fused attention and currently cannot be combined with masks, bias, dropout, SWA, CP, or packed/ragged sequence metadata.

Next steps