Kvax: Flash Attention for JAX

Kvax supports context parallelism to balance token loads across GPUs and reduce memory consumption. To use it, follow these steps:

1. Permute Tokens: Use permute_tokens_context_parallelism to redistribute input tokens (embeddings, positions, and segment IDs) across the mesh.
2. Compute Attention: Use the standard flash_attention and create_attention_mask operations.
3. Unpermute Tokens: Use unpermute_tokens_context_parallelism to restore the original token order for the outputs.
4. Configure Sharding: Use attention_specs within a jax.sharding.Mesh context to define how query and KV tensors are sharded across the mesh axes.

from kvax.utils import (
    permute_tokens_context_parallelism,
    unpermute_tokens_context_parallelism,
    attention_specs
)
import jax

# 1. Permute input
if self._should_permute_input_tokens:
    embeddings, query_positions, query_segment_ids = permute_tokens_context_parallelism(
        (embeddings, positions, segment_ids),
    )

# ... perform attention operations ...

# 2. Unpermute output
if self._should_permute_input_tokens:
    embeddings = unpermute_tokens_context_parallelism(embeddings)

# 3. Setup sharding context
mesh = jax.sharding.Mesh(mesh_devices, mesh_names)
with mesh, attention_specs(
    query_specs=("data", "context", None, None),
    kv_specs=("data", None, None, None),
):
    # Run transformer call
    logits = Transformer(...)(embeddings, positions, segment_ids)

A context manager used to set the sharding specifications for query and key/value tensors. All other Kvax specifications are derived from these.

Arguments:

• query_specs: Sharding specs for the query tensor. Must have 4 dimensions: (batch, query_sequence, heads, attention_head_dim).
• kv_specs: Sharding specs for the key/value tensors. Must have 4 dimensions: (batch, kv_sequence, kv_heads, attention_head_dim).

Constraints:

• batch and attention_head_dim dimensions must have identical sharding specs in both query_specs and kv_specs.

Common Patterns:

• Tensor Parallelism (Mesh with data and model axes):
  • query_specs: ("data", None, "model", None)
  • kv_specs: ("data", None, "model", None)
• Context Parallelism (Mesh with data and context axes):
  • query_specs: ("data", "context", None, None)
  • kv_specs: ("data", None, None, None)

# Example for Tensor Parallelism
query_specs = ("data", None, "model", None)
kv_specs = ("data", None, "model", None)

with attention_specs(query_specs=query_specs, kv_specs=kv_specs):
    # flash_attention calls go here
    pass

The flash_attention function performs the attention operation using precomputed masks and input tensor sharding specifications. It must be called within an attention_specs context manager.

Arguments:

• query: Tensor of shape (batch_size, query_seq_length, num_heads, head_dim).
• key: Tensor of shape (batch_size, kv_seq_length, num_kv_heads, head_dim).
• value: Tensor of shape (batch_size, kv_seq_length, num_kv_heads, head_dim).
• query_positions: Tensor of shape (batch_size, query_seq_length). Use range(0, query_seq_length) for sequential tokens. Ignored if assume_sequential_positions=True.
• query_segment_ids: Tensor of shape (batch_size, query_seq_length). Tokens in the same sequence must share an ID. Use PADDING_SEGMENT_ID for padding.
• kv_positions: Tensor of shape (batch_size, kv_seq_length). Use range(0, kv_seq_length) for sequential tokens. Ignored if assume_sequential_positions=True.
• kv_segment_ids: Tensor of shape (batch_size, kv_seq_length). Tokens in the same sequence must share an ID. Use PADDING_SEGMENT_ID for padding.
• mask: Precomputed block-wise mask from create_attention_mask.
• scale: Scaling factor (default: 1.0).
• fwd_params: FlashAttentionParamsConfig for forward pass.
• bwd_params: FlashAttentionParamsConfig for backward pass.
• assume_sequential_positions: If True, behaves like is_causal=True in jax.nn.dot_product_attention and skips position loading.
• memory_optimized_gqa_backward: If True, enables memory-optimized gradient computation for GQA (slower but saves GPU memory).
• permute_tokens_for_load_balance: If True, permutes tokens for better load balancing in context parallelism (default: True).
• debug: If True, prints low-level IR of the kernel.
• mesh: Device mesh configuration. If None, uses the global context mesh.

# Note: Must be used within attention_specs context
with attention_specs(query_specs=..., kv_specs=...):
    output = flash_attention(
        query=query,
        key=key,
        value=value,
        query_positions=q_pos,
        query_segment_ids=q_seg,
        kv_positions=kv_pos,
        kv_segment_ids=kv_seg,
        mask=mask,
        assume_sequential_positions=True
    )

Calculates attention masks for forward and (optionally) backward Flash Attention operations using Triton kernels for block-wise computation.

Arguments:

• query_positions: Tensor of shape (batch_size, query_seq_length).
• query_segment_ids: Tensor of shape (batch_size, query_seq_length). Use PADDING_SEGMENT_ID for padding.
• kv_positions: Tensor of shape (batch_size, kv_seq_length).
• kv_segment_ids: Tensor of shape (batch_size, kv_seq_length). Use PADDING_SEGMENT_ID for padding.
• fwd_params: FlashAttentionParamsConfig for forward pass.
• bwd_params: FlashAttentionParamsConfig for backward pass.
• calc_bwd_mask: If True, calculates masks for the backward pass (default: False).
• skip_pad_tokens: If True, skips blocks containing only padding tokens (default: True).
• mesh: Device mesh configuration. If None, uses the global context mesh.

mask = create_attention_mask(
    query_positions=q_pos,
    query_segment_ids=q_seg,
    kv_positions=kv_pos,
    kv_segment_ids=kv_seg,
    calc_bwd_mask=True
)