Skip to content

Examples#

End-to-end examples grouped by the same high-level concepts as the library. Source files live at most one directory below examples/:

  • examples/flex_attention/ — FlexAttention masks and score mods, model recipes, and runtime integration.
  • examples/linear/ — KDA/GDN training, context parallelism, and decoding.
  • examples/sparse/ — compressed sparse attention, VSA with FlexAttention (examples/sparse/vsa_flex.py), and FastWan integration (examples/sparse/fastwan_vsa.py).

Run the commands below from the repository root.

Interactive Notebooks#

FlexAttention Walkthrough#

examples/flex_attention/flex_attn.ipynb — Comprehensive Jupyter notebook covering the full FlexAttention API: basic usage, score_mod vs mask_mod, causal masking, sliding window, prefix LM, document masking, NATTEN (with tiled layout), ALiBi, tanh soft-capping, nested jagged tensors, and Flamingo cross-attention. Each section includes performance benchmarks against F.scaled_dot_product_attention.

jupyter notebook examples/flex_attention/flex_attn.ipynb

Debugging Score Mods#

examples/flex_attention/debug_score_mod.py — Interactive marimo notebook that walks through debugging a broken score_mod step-by-step. Demonstrates the _FLEX_ATTENTION_DISABLE_COMPILE_DEBUG flag and how to use get_unwrapped to inspect tensor values inside vmap. See Concepts - Debugging for more details.

pip install -e ".[viz]"
marimo run examples/flex_attention/debug_score_mod.py

Benchmarking#

examples/flex_attention/benchmark.py — Compare FlexAttention against causal FA2 and F.sdpa with a dense mask across multiple attention patterns.

# Run all benchmarks
python examples/flex_attention/benchmark.py

# Run specific patterns
python examples/flex_attention/benchmark.py --examples causal sliding_window document

Available benchmark patterns: causal, alibi, sliding_window, prefix_lm, document, softcap, softcap_approx.

Each benchmark reports forward and backward time (ms) and TFLOPS for three implementations:

Implementation Description
causal FA2 F.scaled_dot_product_attention(is_causal=True)
F.sdpa + mask F.scaled_dot_product_attention(attn_mask=...)
flexattention flex_attention with BlockMask

The key takeaway: FlexAttention with block sparsity can match or beat dense SDPA while supporting arbitrary attention patterns.

Flash Backend Comparison#

examples/flex_attention/flex_flash_attention.py — Compare the Flash (CuTeDSL-based) and Triton backends for FlexAttention.

python examples/flex_attention/flex_flash_attention.py                    # Run all
python examples/flex_attention/flex_flash_attention.py --mode benchmark   # Just performance
python examples/flex_attention/flex_flash_attention.py --mode compare     # Just numerical accuracy

Reports both numerical accuracy (max absolute error vs FP32 reference) and performance (ms, TFLOPS) for forward and backward passes. Requires PyTorch >= 2.10 and flash-attn with CuTeDSL support.

Multi-Head Latent Attention (MLA)#

examples/flex_attention/mla.py — DeepSeek-V2 style MLA with weight absorption, showing how FlexAttention can express the absorbed RoPE score modification.

python examples/flex_attention/mla.py --mode acc     # Accuracy test (vanilla vs absorbed vs flex)
python examples/flex_attention/mla.py --mode perf    # Performance comparison

Not recommended for production use

This is a demonstration of what FlexAttention can express. The flex-based MLA path is not optimized for production workloads.

Determinism Testing#

examples/flex_attention/flex_determinism.py — Test bitwise determinism of FlexAttention across multiple compilation settings and tensor shapes.

python examples/flex_attention/flex_determinism.py

Tests forward and backward determinism across configurations (eager, Inductor default, forced reduction filtering) with shapes covering standard attention, decode, GQA, and long-context scenarios. See the FlexAttention determinism guide for recommended production settings and an explanation of the backward-pass reduction behavior.

Ring Attention#

examples/flex_attention/ring_attention.py — Build a custom distributed ring-attention op by directly invoking FlexAttention's forward and backward primitives.

Implements a generic single-node ring-attention example launched with torchrun. Each rank owns a contiguous sequence shard of q/k/v, rotates k/v with point-to-point communication, merges the local (out, lse) online, then routes dk/dv contributions back to the owning rank by circulating gradient accumulators with each shard in backward. The script validates local slices and gathered outputs and gradients against a single-process causal reference.

torchrun --standalone --nproc_per_node=4 examples/flex_attention/ring_attention.py --seq-len 131072

Context-Parallel Attention Styles#

examples/flex_attention/context_parallel_styles.py — Compare the common ways to shard softmax attention over the sequence, all starting from the same zigzag-sharded q/k/v:

  • allgather / allgather-overlap: gather k/v before one attention call, or attend to the local shard while the gather is in flight and merge by log-sum-exp.
  • ulysses / ulysses-overlap: all-to-all from sequence to head sharding and back, optionally pipelined over head groups.
  • ring: rotate k/v with a prefetched point-to-point shift and merge each step by log-sum-exp; a hand-written backward overlaps the k/v and dk/dv rotations with compute.

--transport auto|nccl|copy-engine picks how the all-gather and ring styles move k/v: NCCL, or copy-engine pushes into symmetric memory from attn_gym.distributed, which need every rank in one NVLink domain (auto picks copy engines when every rank shares one, otherwise NCCL). The copy-engine all-gather writes each shard straight into the layout attention reads. The NCCL path and Ulysses use torch.distributed._functional_collectives, whose async tensors make the overlap explicit; their backwards come from the collectives' autograd formulas. Pick styles with repeated --style. --cuda-graph replays each step from a CUDA graph to remove host launch overhead. --validate checks outputs and gradients against single-device attention, and --trace-dir writes per-rank Perfetto traces through transformer_nuggets.

torchrun --standalone --nproc_per_node=4 examples/flex_attention/context_parallel_styles.py --seq-len 16384 --validate --cuda-graph
torchrun --standalone --nproc_per_node=4 examples/flex_attention/context_parallel_styles.py --seq-len 131072 --backend FLASH --trace-dir traces/

Delta-Rule Context Parallelism#

examples/linear/delta_rule_context_parallel.py — Run the complete transformer-style module from delta_rule_training.py over packed context-parallel fragments using the reference recipe in attn_gym.linear.context_parallel (see Context Parallelism). This includes projections, Q/K/V short convolution, the KDA or GDN core (--variant), normalization, gating, and the output projection.

The input uses the same shared Zipf sampler as the base training example's --packed mode: --batch-size counts logical sequences and --tokens bounds their lengths. Use --sequence-lengths to override sampling with an explicit layout. --num-heads and --core-backend match the base example; --heads and --kda-backend remain aliases.

Each rank owns fragments (global token ranges) from partition_fragments directly in the CP example. Edit this function to define your own rank mappings: --partition contiguous gives each rank one near-equal block, while --partition zigzag gives it two mirrored blocks, matching the layout hybrid models inherit from ring softmax attention. Uneven totals are partitioned without dropping or padding tokens. The model, batch construction, loss/backward, and capture recipes are shown in examples/linear/delta_rule_training.py and imported by the CP example. Numerical assertions, profiling, and benchmark reporting stay in attn_gym.testing.delta_rule and attn_gym.testing.profiling. The example keeps the run-mode branches visible: validate first, then benchmark, capture/replay, or profile an eager step. The short convolution all-gathers one W - 1 token tail per fragment (its last subsequence's) and routes its initial-state gradient back to the owning ranks. The delta-rule recurrence computes forward and reverse [bias; transition] state summaries with portable Triton kernels on Hopper or native TMA/UMMA kernels on SM100, all-gathers them, and folds each fragment's predecessors or successors onto its first / last subsequence before running the ordinary local KDA or GDN kernel. The example validates local outputs, convolution and recurrent endpoint states, input gradients, and all-reduced parameter gradients against the complete unsharded module. It can also capture the full forward, NCCL communication, and backward in one CUDA Graph and validate a changed-input replay. Use --compute-dtype=float16 to exercise the FP16 route. Add --profile to use transformer-nuggets to export one merged multi-rank trace in native Perfetto .pftrace format. This requires a transformer-nuggets source revision with native Perfetto support; older releases lack the exporter:

uv pip install --python .venv/bin/python 'git+https://github.com/drisspg/transformer_nuggets.git'

torchrun --standalone --nproc_per_node=2 examples/linear/delta_rule_context_parallel.py --batch-size 4 --tokens 1024

torchrun --standalone --nproc_per_node=2 examples/linear/delta_rule_context_parallel.py --partition zigzag

torchrun --standalone --nproc_per_node=2 examples/linear/delta_rule_context_parallel.py --compute-dtype=float16

torchrun --standalone --nproc_per_node=2 examples/linear/delta_rule_context_parallel.py --cuda-graph

torchrun --standalone --nproc_per_node=2 examples/linear/delta_rule_context_parallel.py --profile

Kernel Tuning#

Autotune Replay#

examples/flex_attention/flex_autotune_replay.py — Capture kernel tuning decisions from max-autotune and replay them deterministically. Useful for locking in performance-critical configurations.

Grid Sweep#

examples/flex_attention/flex_grid_sweep.py — Exhaustive manual sweep over kernel configurations (block sizes, stages, warps) for a given attention pattern and problem size. Edit MASK_MOD and SCORE_MOD at the top of the file to test your specific pattern.

python examples/flex_attention/flex_grid_sweep.py --B 4 --H 32 --S 8192 --D 128

Paged Attention#

examples/flex_attention/paged_attention.py — Efficient inference examples using FlexAttention with batches of variable-length KV caches. KV tensors are split into fixed-size pages and stored compactly instead of padding to the maximum length.

From a repository checkout:

import torch
from examples.flex_attention.paged_attention import PagedAttention

paged_attn = PagedAttention(
    n_pages=256,
    page_size=128,
    max_batch_size=32,
    device="cuda",
)

# Reserve capacity for a batch element
paged_attn.reserve(batch_idx=torch.tensor(0), seq_len=torch.tensor(512))

# Assign KV values into the paged cache
paged_attn.assign(batch_idx, input_pos, k_val, v_val, k_cache, v_cache)

# Convert a logical block mask to physical page layout
physical_block_mask = paged_attn.convert_logical_block_mask(block_mask)

Key methods:

  • reserve(batch_idx, seq_len) — ensure capacity for at least seq_len tokens
  • erase(batch_idx) — free all pages for a batch element
  • assign(...) — write KV values into the paged cache
  • convert_logical_block_mask(block_mask) — remap logical block indices to physical page indices
  • get_mask_mod(mask_mod) / get_score_mod(score_mod) — wrap mods to operate in physical page space

See paged_attention.py for the full example implementation.

ROWS_GUARANTEED_SAFE Demo#

examples/flex_attention/rows_guaranteed_safe_demo.py — Demonstrates the subtle ROWS_GUARANTEED_SAFE contract in BlockMask metadata. Shows how block-level metadata can declare a block "safe" even when individual rows within it have no surviving attention entries, leading to NaN under torch.compile but not in eager mode.

python examples/flex_attention/rows_guaranteed_safe_demo.py

AOTInductor Integration#

Block Mask Export#

examples/flex_attention/aoti_create_block_mask.py — Export create_block_mask as an AOTInductor package with dynamic batch/head dimensions. Useful for deploying block mask creation outside of Python.

python examples/flex_attention/aoti_create_block_mask.py

Full FlexAttention Export#

examples/flex_attention/aoti_flex_attention.py — Export the full flex_attention forward (with LSE) as an AOTInductor artifact that can be serialized, saved, and reloaded.

python examples/flex_attention/aoti_flex_attention.py --mode aoti    # AOTInductor backend
python examples/flex_attention/aoti_flex_attention.py --mode python  # Python inductor backend