├── .gitignore ├── LICENSE ├── README.md ├── README_RING_ATTN.md ├── bpt.py ├── llamabpt ├── __init__.py ├── data.py ├── llama.py └── train.py ├── scripts ├── gpu_requirements.yml ├── hf2jax.py ├── jax2hf.py ├── prepare_data.py └── tpu_requirements.sh └── tests └── test_fwd_bwd.py /.gitignore: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/exists-forall/striped_attention/HEAD/.gitignore -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/exists-forall/striped_attention/HEAD/LICENSE -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/exists-forall/striped_attention/HEAD/README.md -------------------------------------------------------------------------------- /README_RING_ATTN.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/exists-forall/striped_attention/HEAD/README_RING_ATTN.md -------------------------------------------------------------------------------- /bpt.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/exists-forall/striped_attention/HEAD/bpt.py -------------------------------------------------------------------------------- /llamabpt/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /llamabpt/data.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/exists-forall/striped_attention/HEAD/llamabpt/data.py -------------------------------------------------------------------------------- /llamabpt/llama.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/exists-forall/striped_attention/HEAD/llamabpt/llama.py -------------------------------------------------------------------------------- /llamabpt/train.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/exists-forall/striped_attention/HEAD/llamabpt/train.py -------------------------------------------------------------------------------- /scripts/gpu_requirements.yml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/exists-forall/striped_attention/HEAD/scripts/gpu_requirements.yml -------------------------------------------------------------------------------- /scripts/hf2jax.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/exists-forall/striped_attention/HEAD/scripts/hf2jax.py -------------------------------------------------------------------------------- /scripts/jax2hf.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/exists-forall/striped_attention/HEAD/scripts/jax2hf.py -------------------------------------------------------------------------------- /scripts/prepare_data.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/exists-forall/striped_attention/HEAD/scripts/prepare_data.py -------------------------------------------------------------------------------- /scripts/tpu_requirements.sh: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/exists-forall/striped_attention/HEAD/scripts/tpu_requirements.sh -------------------------------------------------------------------------------- /tests/test_fwd_bwd.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/exists-forall/striped_attention/HEAD/tests/test_fwd_bwd.py --------------------------------------------------------------------------------