attention-gym
Helpful tools and examples for working with flex-attention
File Explorer
Download Latest Version (.zip)- copy_reads_example.py
- SKILL.md
- SKILL.md
- SKILL.md
- delete_old_pr_branches.py
- generate_ci_health.py
- test_delete_old_pr_branches.py
- test_generate_ci_health.py
- claude-code.yml
- delete-old-pr-branches.yml
- docs.yml
- modal-tests.yml
- prek.yml
- publish-to-pypi.yml
- test.yml
- zizmor.yml
- __init__.py
- _compile_driver.py
- _key.py
- cache.py
- compile.py
- device.py
- target.py
- tune.py
- utils.py
- __init__.py
- utils.py
- __init__.py
- __init__.py
- reference.py
- __init__.py
- api.py
- validation.py
- __init__.py
- chunk_delta_h_bwd_v1.py
- chunk_delta_h_bwd_v1_dispatch.py
- chunk_kda_bwd.py
- chunk_kda_bwd_intra.py
- chunk_kda_bwd_wy_dqkg_fused.py
- gate_bwd_fused.py
- __init__.py
- chunk_kda_bwd_dav.py
- gate_bwd.py
- l2norm_bwd.py
- __init__.py
- __init__.py
- chunk_kda_fwd.py
- chunk_kda_fwd_inter_solve.py
- chunk_kda_fwd_intra.py
- chunk_kda_k3b_offdiag_cutedsl.py
- chunk_kda_k4b_inverse_cutedsl.py
- chunk_schedule.py
- chunk_scheduler_cute.py
- recompute_w_u_fwd.py
- __init__.py
- chunk_delta_h.py
- chunk_gla_fwd_o.py
- chunk_kda_fwd_intra_sub_chunk_forloop.py
- gate_fwd.py
- l2norm_fwd.py
- masking.py
- recompute_w_u.py
- recurrent.py
- __init__.py
- __init__.py
- fused.py
- reference.py
- __init__.py
- activations.py
- cute.py
- __init__.py
- api.py
- chunk_schedule.py
- chunk_scheduler.py
- masking.py
- naive.py
- ops.py
- utils.py
- validation.py
- __init__.py
- types.py
- __init__.py
- batchify.py
- block_diffusion.py
- causal.py
- dilated_sliding_window.py
- document_mask.py
- flamingo.py
- global_sliding_window.py
- jetspec.py
- natten.py
- prefix_lm.py
- shared_prefix.py
- sliding_window.py
- sta.py
- svg.py
- vsa.py
- __init__.py
- activation.py
- alibi.py
- graphormer.py
- latent_attention.py
- sandwich.py
- softcapping.py
- __init__.py
- backward.py
- forward.py
- primitives.py
- shared_backward.py
- __init__.py
- cute.py
- reference.py
- __init__.py
- api.py
- __init__.py
- __init__.py
- kda.py
- __init__.py
- utils.py
- flash_b200.yaml
- flash_h200.yaml
- triton.yaml
- selected_attention_benchmark.py
- bench_block_mask.py
- flex_perf.py
- deterministic-mode.png
- fixed-reduction-config.png
- nondeterministic-backward.png
- ci-health.js
- ci-health.json
- hero.png
- ci-health.md
- extra.css
- concepts.md
- determinism.md
- examples.md
- getting-started.md
- index.md
- linear.md
- masks.md
- mods.md
- utilities.md
- latency.py
- model.py
- paged_attention.py
- throughput.py
- utils.py
- __init__.py
- aoti_create_block_mask.py
- aoti_flex_attention.py
- benchmark.py
- compressed_sparse_attention.py
- debug_score_mod.py
- fastwan_vsa.py
- flex_attn.ipynb
- flex_autotune_replay.py
- flex_determinism.py
- flex_flash_attention.py
- flex_grid_sweep.py
- kda_training.py
- kda_zero_drift_decode.py
- mla.py
- ring_attention.py
- rows_guaranteed_safe_demo.py
- shared_prefix_split_attention.py
- vsa.py
- __init__.py
- test_gdn.py
- AGENTS.md
- conftest.py
- test_block_diffusion.py
- test_causal_block_mask.py
- test_cute_cache.py
- test_cute_cache_gpu.py
- test_document_mask.py
- test_global_sliding_window.py
- test_graphormer.py
- test_jetspec.py
- test_kda.py
- test_kda_chunk_scheduler.py
- test_kda_compile_dynamic_shapes.py
- test_kda_cute_forward.py
- test_kda_cute_inter_solve.py
- test_kda_cute_recompute.py
- test_kda_fused_gate_bwd.py
- test_kda_impl_dispatch.py
- test_kda_imports.py
- test_kda_int64_offsets.py
- test_kda_kernels.py
- test_kda_masking.py
- test_kda_ragged_autograd_ops.py
- test_kda_ragged_bwd_dav.py
- test_kda_ragged_bwd_intra.py
- test_kda_ragged_bwd_wy.py
- test_kda_ragged_delta_h_bwd.py
- test_kda_ragged_forward_stages.py
- test_kda_ragged_gate.py
- test_kda_ragged_intra_diagonal.py
- test_kda_ragged_k3.py
- test_kda_ragged_k4.py
- test_kda_ragged_output.py
- test_kda_ragged_public_backward.py
- test_kda_ragged_recurrence.py
- test_kda_recurrent.py
- test_kda_short_conv_cute.py
- test_kda_training_example.py
- test_merge_attention.py
- test_mods.py
- test_namespaces.py
- test_natten.py
- test_sandwich.py
- test_selected_attention_cute.py
- test_selected_attention_matches_csa_reference.py
- test_selected_attention_reference.py
- test_selected_attention_semantics.py
- test_selected_attention_triton.py
- test_shared_prefix.py
- test_shared_prefix_split_attention.py
- test_triton_utils.py
- test_vsa.py
- .gitignore
- .pre-commit-config.yaml
- AGENTS.md
- CLAUDE.md
- CODE_OF_CONDUCT.md
- CONTRIBUTING.md
- DESIGN.md
- LICENSE
- mkdocs.yml
- modal_tests.py
- pyproject.toml
- README.md
๐ Installation Guide
1. Get the code
git clone https://github.com/meta-pytorch/attention-gym
Downloads the entire project code from GitHub to your computer.
cd attention-gym
Moves into the project folder you just downloaded.
2. Official Install Script
Easy RecommendedPrerequisites
- Python 3 Python is required to use pip.
pip install .
Installs the package published on PyPI directly โ no need to clone the source.
After installing, open a new terminal and run the program's version command (e.g. --version) to confirm it worked.
Pulled directly from this repo's README.
3. Python
EasyPrerequisites
pip install .
Installs the package published on PyPI directly โ no need to clone the source.
pip install -e ".[dev]"
Installs the Python libraries listed in requirements.txt (or similar).
If it runs without errors and prints output in the terminal, it worked.
Pulled directly from this repo's README.
// repository documentation
Was this content helpful?
(0 ratings)