Implementation:NVIDIA TransformerEngine JAX Cpp Attention
| Field | Value |
|---|---|
| Sources | TransformerEngine |
| Domains | Deep_Learning, JAX, Attention |
| Last Updated | 2026-02-07 14:00 GMT |
Overview
Implements JAX custom primitives for fused multi-head attention (forward and backward), including context parallelism variants using AllGather and Ring P2P communication strategies.
Description
FusedAttnHelper queries the C++ backend for available cuDNN fused attention kernels based on configuration (dtype, layout, mask type, head dimensions). Forward/backward primitives (FusedAttnFwdPrimitive, FusedAttnBwdPrimitive) call into the TE C++ FFI bindings. Context parallelism is implemented through separate primitive classes for AllGather (FusedAttnCPWithAllGatherFwdPrimitive) and Ring P2P (FusedRingAttnFwdPrimitive) strategies, with both standard and striped variants. _FusedAttnConfig carries static attention configuration as a frozen dataclass registered with JAX's pytree system.
This is the largest and most complex cpp_extensions file, providing the critical bridge between JAX and cuDNN's optimized fused attention kernels, with full support for distributed attention across multiple GPUs via context parallelism.
Usage
Use this module when the high-level fused_attn function in transformer_engine.jax.attention dispatches to the underlying primitives. Direct usage is needed for implementing custom attention patterns with specific context parallelism strategies.
Code Reference
Source Location
- Repository
NVIDIA/TransformerEngine- File
transformer_engine/jax/cpp_extensions/attention.py- Lines
- 1--3642
Signature
class FusedAttnHelper:
"""Queries cuDNN for available fused attention kernel backends."""
...
class FusedAttnFwdPrimitive(BasePrimitive): ...
class FusedAttnBwdPrimitive(BasePrimitive): ...
class FusedAttnCPWithAllGatherFwdPrimitive(BasePrimitive): ...
class FusedAttnCPWithAllGatherBwdPrimitive(BasePrimitive): ...
class FusedAttnCPStripedWithAllGatherFwdPrimitive(BasePrimitive): ...
class FusedAttnCPStripedWithAllGatherBwdPrimitive(BasePrimitive): ...
class FusedRingAttnFwdPrimitive(BasePrimitive): ...
class FusedRingAttnBwdPrimitive(BasePrimitive): ...
class FusedRingAttnStripedFwdPrimitive(BasePrimitive): ...
class FusedRingAttnStripedBwdPrimitive(BasePrimitive): ...
def fused_attn_fwd(...) -> Tuple: ...
def fused_attn_bwd(...) -> Tuple: ...
Import
from transformer_engine.jax.cpp_extensions.attention import fused_attn_fwd, fused_attn_bwd, FusedAttnHelper
I/O Contract
Inputs
| Name | Type | Required | Description |
|---|---|---|---|
| qkv | Tuple[jnp.ndarray, ...] |
Yes | Query, key, value tensors in specified layout |
| bias | Optional[jnp.ndarray] |
No | Attention bias tensor |
| seqlen_q | jnp.ndarray |
Yes | Query sequence lengths |
| seqlen_kv | jnp.ndarray |
Yes | Key/value sequence lengths |
| seed | Optional[jnp.ndarray] |
No | RNG seed for dropout |
| attn_bias_type | NVTE_Bias_Type |
Yes | Attention bias type enum |
| attn_mask_type | NVTE_Mask_Type |
Yes | Attention mask type enum |
| qkv_layout | NVTE_QKV_Layout |
Yes | QKV memory layout enum |
| scaling_factor | float |
Yes | Attention scaling factor |
| dropout_probability | float |
Yes | Dropout rate |
| is_training | bool |
Yes | Whether model is in training mode |
Outputs
| Name | Type | Description |
|---|---|---|
| output | jnp.ndarray |
Attention output |
| softmax_aux | jnp.ndarray |
Softmax auxiliary data for backward pass |
| rng_state | jnp.ndarray |
RNG state for backward pass dropout |
Usage Examples
from transformer_engine.jax.cpp_extensions.attention import FusedAttnHelper
# Check if a fused attention kernel is available
helper = FusedAttnHelper(
q_dtype, kv_dtype, qkv_layout, attn_bias_type, attn_mask_type,
dropout_probability, q_num_heads, kv_num_heads, q_max_seqlen,
kv_max_seqlen, head_dim
)
backend = helper.get_fused_attn_backend()