├── .clang-format ├── .coveragerc ├── .flake8 ├── .github ├── scripts │ └── install_requirements.sh └── workflows │ ├── build_test.yaml │ ├── lint.yaml │ └── sphinx.yaml ├── .gitignore ├── .gitmodules ├── .pre-commit-config.yaml ├── CONTRIBUTING.md ├── LICENSE ├── MANIFEST.in ├── README.md ├── assets ├── ffa_exp │ ├── attn with causal mask │ │ ├── .DS_Store │ │ ├── attn-bwd with causal mask │ │ │ ├── flops_report.pdf │ │ │ └── flops_report.png │ │ ├── attn-fwd with causal mask │ │ │ ├── flops_report.pdf │ │ │ └── flops_report.png │ │ ├── perf_report_all.pdf │ │ └── perf_report_all.png │ ├── attn with fulll mask │ │ ├── .DS_Store │ │ ├── attn-bwd with full mask │ │ │ ├── flops_report.csv │ │ │ ├── flops_report.pdf │ │ │ └── flops_report.png │ │ ├── attn-fwd with full mask │ │ │ ├── flops_report.csv │ │ │ ├── flops_report.pdf │ │ │ └── flops_report.png │ │ ├── perf_report_all.pdf │ │ └── perf_report_all.png │ ├── attn with sw causal mask │ │ ├── .DS_Store │ │ ├── attn-bwd with sliding_window_causal mask │ │ │ ├── flops_report.pdf │ │ │ └── flops_report.png │ │ ├── attn-fwd with sliding_window_causal mask │ │ │ ├── flops_report.pdf │ │ │ └── flops_report.png │ │ ├── perf_report_all.pdf │ │ └── perf_report_all.png │ ├── attn with varlen block causal mask │ │ ├── .DS_Store │ │ ├── attn-bwd with varlen_block_causal mask │ │ │ └── flops_report.png │ │ ├── attn-fwd with varlen_block_causal mask │ │ │ └── flops_report.png │ │ └── perf_report_all.png │ ├── attn with varlen causal mask │ │ ├── .DS_Store │ │ ├── attn-bwd with varlen_causal mask │ │ │ ├── .DS_Store │ │ │ ├── flops_report.csv │ │ │ ├── flops_report.pdf │ │ │ └── flops_report.png │ │ ├── attn-fwd with varlen_causal mask │ │ │ ├── .DS_Store │ │ │ ├── flops_report.csv │ │ │ ├── flops_report.pdf │ │ │ └── flops_report.png │ │ ├── perf_report_all.pdf │ │ └── perf_report_all.png │ └── attn with varlen full mask │ │ ├── .DS_Store │ │ ├── attn-bwd with varlen_full mask │ │ ├── flops_report.csv │ │ ├── flops_report.pdf │ │ └── flops_report.png │ │ ├── attn-fwd with varlen_full mask │ │ ├── flops_report.csv │ │ ├── flops_report.pdf │ │ └── flops_report.png │ │ ├── perf_report_all.pdf │ │ └── perf_report_all.png ├── magi_attention_exp │ ├── full_mask_bwd_per_gpu │ │ ├── flops_report.pdf │ │ └── flops_report.png │ ├── full_mask_fwd_per_gpu │ │ ├── flops_report.pdf │ │ └── flops_report.png │ ├── varlen_full_mask_bwd_per_gpu │ │ ├── flops_report.pdf │ │ └── flops_report.png │ └── varlen_full_mask_fwd_per_gpu │ │ ├── flops_report.pdf │ │ └── flops_report.png └── magiattn_overview.png ├── conftest.py ├── docs ├── Makefile ├── README.md ├── make.bat ├── requirements.txt └── source │ ├── _static │ ├── custom.css │ └── sand-logos │ │ ├── magi-black.png │ │ └── magi-icon.png │ ├── conf.py │ ├── env_variables.md │ ├── guide.md │ ├── index.md │ ├── install.md │ ├── magi_api.md │ └── quickstart.md ├── examples ├── megatron │ ├── README.md │ └── results.png ├── torch_native │ ├── README.md │ ├── configuration_llama.py │ ├── llama_pretrain_config.py │ ├── main.py │ ├── modeling_llama.py │ └── run.sh └── transformers │ ├── README.md │ ├── magi_attention_func.py │ ├── magi_trainer.py │ ├── result.png │ ├── run_magi_clm.py │ ├── run_magi_clm.sh │ ├── run_origin_clm.py │ └── run_origin_clm.sh ├── exps ├── .gitignore ├── README.md ├── __init__.py ├── attn │ ├── __init__.py │ ├── baselines │ │ ├── __init__.py │ │ ├── attn_impl.py │ │ ├── block_sparse_attn_triton.py │ │ └── utils.py │ ├── draw_benchmark.py │ ├── profile_ffa │ │ ├── README.md │ │ ├── __init__.py │ │ ├── compare_ffa_results.py │ │ ├── ffa_benchmark.py │ │ ├── profile_ffa.sh │ │ └── run_branch_profile.sh │ ├── requirements.txt │ ├── run_benchmark.py │ ├── run_benchmark.sh │ ├── run_block_sparse_bench.sh │ ├── run_block_sparse_benchmark.py │ ├── run_var_block_sparse_benchmark.py │ └── tests │ │ └── __init__.py ├── device_a2av │ ├── __init__.py │ ├── run.sh │ └── test.py ├── dist_attn │ ├── __init__.py │ ├── baselines │ │ ├── __init__.py │ │ ├── interface.py │ │ ├── loongtrain.py │ │ ├── nsa.py │ │ ├── ring_attn.py │ │ ├── shard.py │ │ ├── ulysess.py │ │ ├── usp.py │ │ ├── usp_nsa.py │ │ └── utils_cp.py │ ├── benchmark │ │ ├── __init__.py │ │ ├── datasets │ │ │ └── default │ │ │ │ └── doc_length_distribution.csv │ │ ├── enums.py │ │ ├── mask.py │ │ ├── metric.py │ │ └── utils.py │ ├── benchmark_conf.py │ ├── main.py │ ├── requirements.txt │ ├── run_benchmark.py │ ├── run_benchmark.sh │ ├── run_profile.py │ ├── run_profile.sh │ └── tests │ │ ├── __init__.py │ │ ├── attn_profile.py │ │ ├── attn_profile.sh │ │ ├── nsa_profile.py │ │ ├── test_baseline_attn.py │ │ ├── test_usp_nsa.py │ │ └── test_utils.py └── grpcoll │ ├── __init__.py │ ├── grpcoll_utils.py │ ├── run_grpcoll_test.sh │ ├── test_internode_grpcoll.py │ ├── test_intranode_grpcoll.py │ └── test_low_latency_grpcoll.py ├── extensions ├── README.md ├── __init__.py ├── magi_attn_extensions │ ├── __init__.py │ ├── fa2_interface_with_sink.py │ └── fa3_interface_with_sink.py ├── requirements.txt ├── setup.py └── tests │ ├── __init__.py │ └── test_fa_interface_with_sink.py ├── magi_attention ├── __init__.py ├── api │ ├── __init__.py │ ├── functools.py │ └── magi_attn_interface.py ├── benchmarking │ ├── __init__.py │ ├── bench.py │ ├── image_grid.py │ └── utils.py ├── comm │ ├── __init__.py │ ├── functional │ │ ├── __init__.py │ │ └── _gather_scatter_v.py │ ├── primitive │ │ ├── __init__.py │ │ ├── _all2all_v.py │ │ ├── _all_gather_v.py │ │ ├── _scatter_v.py │ │ ├── device_a2av │ │ │ ├── __init__.py │ │ │ ├── interface.py │ │ │ ├── triton_barrier.py │ │ │ ├── triton_impl.py │ │ │ └── triton_utils.py │ │ └── grpcoll │ │ │ ├── __init__.py │ │ │ ├── _a2av_grpcoll_impl.py │ │ │ ├── _buffer.py │ │ │ ├── _config.py │ │ │ ├── _event.py │ │ │ ├── _group_collective.py │ │ │ ├── _group_collective_hier.py │ │ │ ├── _handle.py │ │ │ ├── _mgr.py │ │ │ ├── _native_grpcoll_impl.py │ │ │ └── utils.py │ └── work.py ├── common │ ├── __init__.py │ ├── enum.py │ ├── jit │ │ ├── __init__.py │ │ ├── core.py │ │ ├── cpp_ext.py │ │ ├── env.py │ │ └── utils.py │ ├── mask.py │ ├── range.py │ ├── range_op │ │ ├── __init__.py │ │ ├── _range_fill.py │ │ ├── _range_gather.py │ │ ├── _range_reduce.py │ │ └── utils.py │ ├── ranges.py │ ├── rectangle.py │ └── rectangles.py ├── config.py ├── csrc │ ├── comm │ │ └── grpcoll │ │ │ ├── buffer.cpp │ │ │ ├── buffer.hpp │ │ │ ├── config.hpp │ │ │ ├── event.hpp │ │ │ ├── kernels │ │ │ ├── api.cuh │ │ │ ├── buffer.cuh │ │ │ ├── configs.cuh │ │ │ ├── exception.cuh │ │ │ ├── ibgda_device.cuh │ │ │ ├── internode.cu │ │ │ ├── internode_ll.cu │ │ │ ├── intranode.cu │ │ │ ├── launch.cuh │ │ │ ├── layout.cu │ │ │ ├── reduce_op.cuh │ │ │ ├── runtime.cu │ │ │ └── utils.cuh │ │ │ └── meta.hpp │ ├── common │ │ ├── cuda_check.h │ │ ├── static_switch.h │ │ └── utils.h │ ├── extensions │ │ ├── attn_ranges.hpp │ │ ├── magi_attn_ext.cpp │ │ └── magi_attn_ext.hpp │ ├── flexible_flash_attention │ │ ├── block.h │ │ ├── bwd_inst_template.jinja │ │ ├── copy_sm90_bulk_reduce.hpp │ │ ├── epilogue_bwd.hpp │ │ ├── epilogue_fwd.hpp │ │ ├── flash.h │ │ ├── flash_bwd_kernel_sm90.h │ │ ├── flash_bwd_launch_template.h │ │ ├── flash_bwd_preprocess_kernel.h │ │ ├── flash_fwd_kernel_sm90.h │ │ ├── flash_fwd_launch_template.h │ │ ├── flash_fwd_postprocess.cu │ │ ├── flash_fwd_postprocess_kernel.h │ │ ├── flash_fwd_postprocess_launch_template.h │ │ ├── flex_flash_bwd.hpp │ │ ├── flex_flash_common.cpp │ │ ├── flex_flash_common.hpp │ │ ├── flex_flash_fwd.hpp │ │ ├── fwd_inst_template.jinja │ │ ├── mainloop_bwd_sm90_tma_gmma_ws.hpp │ │ ├── mainloop_fwd_sm90_tma_gmma_ws.hpp │ │ ├── mask.h │ │ ├── named_barrier.hpp │ │ ├── seqlen.h │ │ ├── sink_layout.cuh │ │ ├── sm90_pipeline_no_cluster.hpp │ │ ├── softmax.h │ │ ├── tile_scheduler.hpp │ │ └── tile_size.h │ └── utils │ │ ├── bindings.cpp │ │ ├── profile_utils.cu │ │ ├── profile_utils.h │ │ └── unique_consecutive_pairs.cu ├── dist_attn_runtime_mgr.py ├── functional │ ├── __init__.py │ ├── _flex_flash_attn_jit.py │ ├── dispatch.py │ ├── dist_attn.py │ ├── flex_flash_attn.py │ ├── sdpa.py │ └── utils.py ├── meta │ ├── __init__.py │ ├── _make_attn_meta.py │ ├── _make_dispatch_meta.py │ ├── algorithms │ │ ├── __init__.py │ │ ├── base.py │ │ ├── grg.py │ │ └── ncq.py │ ├── collection │ │ ├── __init__.py │ │ ├── calc_meta.py │ │ ├── comm_meta.py │ │ └── dispatch_meta.py │ ├── container │ │ ├── __init__.py │ │ ├── bucket.py │ │ ├── chunk.py │ │ ├── rank_entry.py │ │ ├── slice.py │ │ └── transfer_table.py │ └── solver │ │ ├── __init__.py │ │ ├── dispatch_solver.py │ │ ├── dist_attn_solver.py │ │ ├── dynamic_attn_solver.py │ │ ├── overlap_solver.py │ │ └── slice_maker.py ├── testing │ ├── __init__.py │ ├── dist_common.py │ ├── flag_generator.py │ ├── gt_dispatcher.py │ ├── precision.py │ ├── ref_attn.py │ └── utils.py └── utils │ ├── __init__.py │ ├── _utils.py │ ├── debug.py │ ├── metaclass.py │ ├── nvtx.py │ └── sparse_utils.py ├── makefile ├── mypy.ini ├── pyproject.toml ├── requirements.txt ├── requirements_dev.txt ├── scripts ├── install_clang_format.sh └── run_csrc_code_formatter.sh ├── setup.py ├── tests ├── __init__.py ├── test_api │ ├── __init__.py │ ├── test_functools.py │ └── test_interface.py ├── test_attn │ ├── __init__.py │ ├── test_block_sparse_attn.py │ ├── test_dist_attn.py │ ├── test_flex_flash_attn.py │ ├── test_merge_range.py │ └── test_ref_attn.py ├── test_attn_solver │ ├── __init__.py │ ├── test_dist_attn_solver.py │ └── test_dynamic_attn_solver.py ├── test_comm │ ├── __init__.py │ ├── test_all_gather_v.py │ ├── test_group_collective.py │ ├── test_group_collective_utils.py │ └── test_scatter_v.py ├── test_common │ ├── __init__.py │ ├── test_attn_mask.py │ ├── test_attn_range.py │ ├── test_attn_ranges.py │ ├── test_range_op │ │ ├── __init__.py │ │ ├── test_range_fill.py │ │ ├── test_range_gather.py │ │ ├── test_range_op_utils.py │ │ └── test_range_reduce.py │ ├── test_rectangle.py │ └── test_rectangles.py ├── test_dispatch │ ├── __init__.py │ ├── test_calc_self_attn_areas.py │ ├── test_dispatch_solver.py │ ├── test_dispatcher.py │ └── test_gt_dispatcher.py ├── test_dist_runtime_mgr │ ├── __init__.py │ └── test_dist_runtime_mgr.py ├── test_functional │ └── __init__.py ├── test_pipeline.py ├── test_pipeline_sdpa.py └── test_utils │ ├── __init__.py │ ├── test_common_utils.py │ └── test_flag_generator.py └── tools ├── __init__.py ├── build_helper.py └── codestyle ├── check_for_chinese.py └── copyright.py /.clang-format: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/.clang-format -------------------------------------------------------------------------------- /.coveragerc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/.coveragerc -------------------------------------------------------------------------------- /.flake8: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/.flake8 -------------------------------------------------------------------------------- /.github/scripts/install_requirements.sh: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/.github/scripts/install_requirements.sh -------------------------------------------------------------------------------- /.github/workflows/build_test.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/.github/workflows/build_test.yaml -------------------------------------------------------------------------------- /.github/workflows/lint.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/.github/workflows/lint.yaml -------------------------------------------------------------------------------- /.github/workflows/sphinx.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/.github/workflows/sphinx.yaml -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/.gitignore -------------------------------------------------------------------------------- /.gitmodules: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/.gitmodules -------------------------------------------------------------------------------- /.pre-commit-config.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/.pre-commit-config.yaml -------------------------------------------------------------------------------- /CONTRIBUTING.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/CONTRIBUTING.md -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/LICENSE -------------------------------------------------------------------------------- /MANIFEST.in: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/MANIFEST.in -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/README.md -------------------------------------------------------------------------------- /assets/ffa_exp/attn with causal mask/.DS_Store: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/assets/ffa_exp/attn with causal mask/.DS_Store -------------------------------------------------------------------------------- /assets/ffa_exp/attn with causal mask/attn-bwd with causal mask/flops_report.pdf: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/assets/ffa_exp/attn with causal mask/attn-bwd with causal mask/flops_report.pdf -------------------------------------------------------------------------------- /assets/ffa_exp/attn with causal mask/attn-bwd with causal mask/flops_report.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/assets/ffa_exp/attn with causal mask/attn-bwd with causal mask/flops_report.png -------------------------------------------------------------------------------- /assets/ffa_exp/attn with causal mask/attn-fwd with causal mask/flops_report.pdf: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/assets/ffa_exp/attn with causal mask/attn-fwd with causal mask/flops_report.pdf -------------------------------------------------------------------------------- /assets/ffa_exp/attn with causal mask/attn-fwd with causal mask/flops_report.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/assets/ffa_exp/attn with causal mask/attn-fwd with causal mask/flops_report.png -------------------------------------------------------------------------------- /assets/ffa_exp/attn with causal mask/perf_report_all.pdf: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/assets/ffa_exp/attn with causal mask/perf_report_all.pdf -------------------------------------------------------------------------------- /assets/ffa_exp/attn with causal mask/perf_report_all.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/assets/ffa_exp/attn with causal mask/perf_report_all.png -------------------------------------------------------------------------------- /assets/ffa_exp/attn with fulll mask/.DS_Store: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/assets/ffa_exp/attn with fulll mask/.DS_Store -------------------------------------------------------------------------------- /assets/ffa_exp/attn with fulll mask/attn-bwd with full mask/flops_report.csv: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/assets/ffa_exp/attn with fulll mask/attn-bwd with full mask/flops_report.csv -------------------------------------------------------------------------------- /assets/ffa_exp/attn with fulll mask/attn-bwd with full mask/flops_report.pdf: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/assets/ffa_exp/attn with fulll mask/attn-bwd with full mask/flops_report.pdf -------------------------------------------------------------------------------- /assets/ffa_exp/attn with fulll mask/attn-bwd with full mask/flops_report.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/assets/ffa_exp/attn with fulll mask/attn-bwd with full mask/flops_report.png -------------------------------------------------------------------------------- /assets/ffa_exp/attn with fulll mask/attn-fwd with full mask/flops_report.csv: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/assets/ffa_exp/attn with fulll mask/attn-fwd with full mask/flops_report.csv -------------------------------------------------------------------------------- /assets/ffa_exp/attn with fulll mask/attn-fwd with full mask/flops_report.pdf: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/assets/ffa_exp/attn with fulll mask/attn-fwd with full mask/flops_report.pdf -------------------------------------------------------------------------------- /assets/ffa_exp/attn with fulll mask/attn-fwd with full mask/flops_report.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/assets/ffa_exp/attn with fulll mask/attn-fwd with full mask/flops_report.png -------------------------------------------------------------------------------- /assets/ffa_exp/attn with fulll mask/perf_report_all.pdf: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/assets/ffa_exp/attn with fulll mask/perf_report_all.pdf -------------------------------------------------------------------------------- /assets/ffa_exp/attn with fulll mask/perf_report_all.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/assets/ffa_exp/attn with fulll mask/perf_report_all.png -------------------------------------------------------------------------------- /assets/ffa_exp/attn with sw causal mask/.DS_Store: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/assets/ffa_exp/attn with sw causal mask/.DS_Store -------------------------------------------------------------------------------- /assets/ffa_exp/attn with sw causal mask/attn-bwd with sliding_window_causal mask/flops_report.pdf: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/assets/ffa_exp/attn with sw causal mask/attn-bwd with sliding_window_causal mask/flops_report.pdf -------------------------------------------------------------------------------- /assets/ffa_exp/attn with sw causal mask/attn-bwd with sliding_window_causal mask/flops_report.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/assets/ffa_exp/attn with sw causal mask/attn-bwd with sliding_window_causal mask/flops_report.png -------------------------------------------------------------------------------- /assets/ffa_exp/attn with sw causal mask/attn-fwd with sliding_window_causal mask/flops_report.pdf: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/assets/ffa_exp/attn with sw causal mask/attn-fwd with sliding_window_causal mask/flops_report.pdf -------------------------------------------------------------------------------- /assets/ffa_exp/attn with sw causal mask/attn-fwd with sliding_window_causal mask/flops_report.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/assets/ffa_exp/attn with sw causal mask/attn-fwd with sliding_window_causal mask/flops_report.png -------------------------------------------------------------------------------- /assets/ffa_exp/attn with sw causal mask/perf_report_all.pdf: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/assets/ffa_exp/attn with sw causal mask/perf_report_all.pdf -------------------------------------------------------------------------------- /assets/ffa_exp/attn with sw causal mask/perf_report_all.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/assets/ffa_exp/attn with sw causal mask/perf_report_all.png -------------------------------------------------------------------------------- /assets/ffa_exp/attn with varlen block causal mask/.DS_Store: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/assets/ffa_exp/attn with varlen block causal mask/.DS_Store -------------------------------------------------------------------------------- /assets/ffa_exp/attn with varlen block causal mask/attn-bwd with varlen_block_causal mask/flops_report.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/assets/ffa_exp/attn with varlen block causal mask/attn-bwd with varlen_block_causal mask/flops_report.png -------------------------------------------------------------------------------- /assets/ffa_exp/attn with varlen block causal mask/attn-fwd with varlen_block_causal mask/flops_report.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/assets/ffa_exp/attn with varlen block causal mask/attn-fwd with varlen_block_causal mask/flops_report.png -------------------------------------------------------------------------------- /assets/ffa_exp/attn with varlen block causal mask/perf_report_all.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/assets/ffa_exp/attn with varlen block causal mask/perf_report_all.png -------------------------------------------------------------------------------- /assets/ffa_exp/attn with varlen causal mask/.DS_Store: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/assets/ffa_exp/attn with varlen causal mask/.DS_Store -------------------------------------------------------------------------------- /assets/ffa_exp/attn with varlen causal mask/attn-bwd with varlen_causal mask/.DS_Store: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/assets/ffa_exp/attn with varlen causal mask/attn-bwd with varlen_causal mask/.DS_Store -------------------------------------------------------------------------------- /assets/ffa_exp/attn with varlen causal mask/attn-bwd with varlen_causal mask/flops_report.csv: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/assets/ffa_exp/attn with varlen causal mask/attn-bwd with varlen_causal mask/flops_report.csv -------------------------------------------------------------------------------- /assets/ffa_exp/attn with varlen causal mask/attn-bwd with varlen_causal mask/flops_report.pdf: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/assets/ffa_exp/attn with varlen causal mask/attn-bwd with varlen_causal mask/flops_report.pdf -------------------------------------------------------------------------------- /assets/ffa_exp/attn with varlen causal mask/attn-bwd with varlen_causal mask/flops_report.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/assets/ffa_exp/attn with varlen causal mask/attn-bwd with varlen_causal mask/flops_report.png -------------------------------------------------------------------------------- /assets/ffa_exp/attn with varlen causal mask/attn-fwd with varlen_causal mask/.DS_Store: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/assets/ffa_exp/attn with varlen causal mask/attn-fwd with varlen_causal mask/.DS_Store -------------------------------------------------------------------------------- /assets/ffa_exp/attn with varlen causal mask/attn-fwd with varlen_causal mask/flops_report.csv: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/assets/ffa_exp/attn with varlen causal mask/attn-fwd with varlen_causal mask/flops_report.csv -------------------------------------------------------------------------------- /assets/ffa_exp/attn with varlen causal mask/attn-fwd with varlen_causal mask/flops_report.pdf: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/assets/ffa_exp/attn with varlen causal mask/attn-fwd with varlen_causal mask/flops_report.pdf -------------------------------------------------------------------------------- /assets/ffa_exp/attn with varlen causal mask/attn-fwd with varlen_causal mask/flops_report.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/assets/ffa_exp/attn with varlen causal mask/attn-fwd with varlen_causal mask/flops_report.png -------------------------------------------------------------------------------- /assets/ffa_exp/attn with varlen causal mask/perf_report_all.pdf: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/assets/ffa_exp/attn with varlen causal mask/perf_report_all.pdf -------------------------------------------------------------------------------- /assets/ffa_exp/attn with varlen causal mask/perf_report_all.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/assets/ffa_exp/attn with varlen causal mask/perf_report_all.png -------------------------------------------------------------------------------- /assets/ffa_exp/attn with varlen full mask/.DS_Store: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/assets/ffa_exp/attn with varlen full mask/.DS_Store -------------------------------------------------------------------------------- /assets/ffa_exp/attn with varlen full mask/attn-bwd with varlen_full mask/flops_report.csv: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/assets/ffa_exp/attn with varlen full mask/attn-bwd with varlen_full mask/flops_report.csv -------------------------------------------------------------------------------- /assets/ffa_exp/attn with varlen full mask/attn-bwd with varlen_full mask/flops_report.pdf: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/assets/ffa_exp/attn with varlen full mask/attn-bwd with varlen_full mask/flops_report.pdf -------------------------------------------------------------------------------- /assets/ffa_exp/attn with varlen full mask/attn-bwd with varlen_full mask/flops_report.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/assets/ffa_exp/attn with varlen full mask/attn-bwd with varlen_full mask/flops_report.png -------------------------------------------------------------------------------- /assets/ffa_exp/attn with varlen full mask/attn-fwd with varlen_full mask/flops_report.csv: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/assets/ffa_exp/attn with varlen full mask/attn-fwd with varlen_full mask/flops_report.csv -------------------------------------------------------------------------------- /assets/ffa_exp/attn with varlen full mask/attn-fwd with varlen_full mask/flops_report.pdf: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/assets/ffa_exp/attn with varlen full mask/attn-fwd with varlen_full mask/flops_report.pdf -------------------------------------------------------------------------------- /assets/ffa_exp/attn with varlen full mask/attn-fwd with varlen_full mask/flops_report.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/assets/ffa_exp/attn with varlen full mask/attn-fwd with varlen_full mask/flops_report.png -------------------------------------------------------------------------------- /assets/ffa_exp/attn with varlen full mask/perf_report_all.pdf: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/assets/ffa_exp/attn with varlen full mask/perf_report_all.pdf -------------------------------------------------------------------------------- /assets/ffa_exp/attn with varlen full mask/perf_report_all.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/assets/ffa_exp/attn with varlen full mask/perf_report_all.png -------------------------------------------------------------------------------- /assets/magi_attention_exp/full_mask_bwd_per_gpu/flops_report.pdf: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/assets/magi_attention_exp/full_mask_bwd_per_gpu/flops_report.pdf -------------------------------------------------------------------------------- /assets/magi_attention_exp/full_mask_bwd_per_gpu/flops_report.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/assets/magi_attention_exp/full_mask_bwd_per_gpu/flops_report.png -------------------------------------------------------------------------------- /assets/magi_attention_exp/full_mask_fwd_per_gpu/flops_report.pdf: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/assets/magi_attention_exp/full_mask_fwd_per_gpu/flops_report.pdf -------------------------------------------------------------------------------- /assets/magi_attention_exp/full_mask_fwd_per_gpu/flops_report.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/assets/magi_attention_exp/full_mask_fwd_per_gpu/flops_report.png -------------------------------------------------------------------------------- /assets/magi_attention_exp/varlen_full_mask_bwd_per_gpu/flops_report.pdf: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/assets/magi_attention_exp/varlen_full_mask_bwd_per_gpu/flops_report.pdf -------------------------------------------------------------------------------- /assets/magi_attention_exp/varlen_full_mask_bwd_per_gpu/flops_report.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/assets/magi_attention_exp/varlen_full_mask_bwd_per_gpu/flops_report.png -------------------------------------------------------------------------------- /assets/magi_attention_exp/varlen_full_mask_fwd_per_gpu/flops_report.pdf: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/assets/magi_attention_exp/varlen_full_mask_fwd_per_gpu/flops_report.pdf -------------------------------------------------------------------------------- /assets/magi_attention_exp/varlen_full_mask_fwd_per_gpu/flops_report.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/assets/magi_attention_exp/varlen_full_mask_fwd_per_gpu/flops_report.png -------------------------------------------------------------------------------- /assets/magiattn_overview.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/assets/magiattn_overview.png -------------------------------------------------------------------------------- /conftest.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/conftest.py -------------------------------------------------------------------------------- /docs/Makefile: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/docs/Makefile -------------------------------------------------------------------------------- /docs/README.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/docs/README.md -------------------------------------------------------------------------------- /docs/make.bat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/docs/make.bat -------------------------------------------------------------------------------- /docs/requirements.txt: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/docs/requirements.txt -------------------------------------------------------------------------------- /docs/source/_static/custom.css: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/docs/source/_static/custom.css -------------------------------------------------------------------------------- /docs/source/_static/sand-logos/magi-black.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/docs/source/_static/sand-logos/magi-black.png -------------------------------------------------------------------------------- /docs/source/_static/sand-logos/magi-icon.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/docs/source/_static/sand-logos/magi-icon.png -------------------------------------------------------------------------------- /docs/source/conf.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/docs/source/conf.py -------------------------------------------------------------------------------- /docs/source/env_variables.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/docs/source/env_variables.md -------------------------------------------------------------------------------- /docs/source/guide.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/docs/source/guide.md -------------------------------------------------------------------------------- /docs/source/index.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/docs/source/index.md -------------------------------------------------------------------------------- /docs/source/install.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/docs/source/install.md -------------------------------------------------------------------------------- /docs/source/magi_api.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/docs/source/magi_api.md -------------------------------------------------------------------------------- /docs/source/quickstart.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/docs/source/quickstart.md -------------------------------------------------------------------------------- /examples/megatron/README.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/examples/megatron/README.md -------------------------------------------------------------------------------- /examples/megatron/results.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/examples/megatron/results.png -------------------------------------------------------------------------------- /examples/torch_native/README.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/examples/torch_native/README.md -------------------------------------------------------------------------------- /examples/torch_native/configuration_llama.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/examples/torch_native/configuration_llama.py -------------------------------------------------------------------------------- /examples/torch_native/llama_pretrain_config.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/examples/torch_native/llama_pretrain_config.py -------------------------------------------------------------------------------- /examples/torch_native/main.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/examples/torch_native/main.py -------------------------------------------------------------------------------- /examples/torch_native/modeling_llama.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/examples/torch_native/modeling_llama.py -------------------------------------------------------------------------------- /examples/torch_native/run.sh: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/examples/torch_native/run.sh -------------------------------------------------------------------------------- /examples/transformers/README.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/examples/transformers/README.md -------------------------------------------------------------------------------- /examples/transformers/magi_attention_func.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/examples/transformers/magi_attention_func.py -------------------------------------------------------------------------------- /examples/transformers/magi_trainer.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/examples/transformers/magi_trainer.py -------------------------------------------------------------------------------- /examples/transformers/result.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/examples/transformers/result.png -------------------------------------------------------------------------------- /examples/transformers/run_magi_clm.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/examples/transformers/run_magi_clm.py -------------------------------------------------------------------------------- /examples/transformers/run_magi_clm.sh: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/examples/transformers/run_magi_clm.sh -------------------------------------------------------------------------------- /examples/transformers/run_origin_clm.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/examples/transformers/run_origin_clm.py -------------------------------------------------------------------------------- /examples/transformers/run_origin_clm.sh: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/examples/transformers/run_origin_clm.sh -------------------------------------------------------------------------------- /exps/.gitignore: -------------------------------------------------------------------------------- 1 | **/outs/ 2 | -------------------------------------------------------------------------------- /exps/README.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/exps/README.md -------------------------------------------------------------------------------- /exps/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/exps/__init__.py -------------------------------------------------------------------------------- /exps/attn/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/exps/attn/__init__.py -------------------------------------------------------------------------------- /exps/attn/baselines/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/exps/attn/baselines/__init__.py -------------------------------------------------------------------------------- /exps/attn/baselines/attn_impl.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/exps/attn/baselines/attn_impl.py -------------------------------------------------------------------------------- /exps/attn/baselines/block_sparse_attn_triton.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/exps/attn/baselines/block_sparse_attn_triton.py -------------------------------------------------------------------------------- /exps/attn/baselines/utils.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/exps/attn/baselines/utils.py -------------------------------------------------------------------------------- /exps/attn/draw_benchmark.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/exps/attn/draw_benchmark.py -------------------------------------------------------------------------------- /exps/attn/profile_ffa/README.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/exps/attn/profile_ffa/README.md -------------------------------------------------------------------------------- /exps/attn/profile_ffa/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/exps/attn/profile_ffa/__init__.py -------------------------------------------------------------------------------- /exps/attn/profile_ffa/compare_ffa_results.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/exps/attn/profile_ffa/compare_ffa_results.py -------------------------------------------------------------------------------- /exps/attn/profile_ffa/ffa_benchmark.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/exps/attn/profile_ffa/ffa_benchmark.py -------------------------------------------------------------------------------- /exps/attn/profile_ffa/profile_ffa.sh: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/exps/attn/profile_ffa/profile_ffa.sh -------------------------------------------------------------------------------- /exps/attn/profile_ffa/run_branch_profile.sh: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/exps/attn/profile_ffa/run_branch_profile.sh -------------------------------------------------------------------------------- /exps/attn/requirements.txt: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/exps/attn/requirements.txt -------------------------------------------------------------------------------- /exps/attn/run_benchmark.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/exps/attn/run_benchmark.py -------------------------------------------------------------------------------- /exps/attn/run_benchmark.sh: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/exps/attn/run_benchmark.sh -------------------------------------------------------------------------------- /exps/attn/run_block_sparse_bench.sh: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/exps/attn/run_block_sparse_bench.sh -------------------------------------------------------------------------------- /exps/attn/run_block_sparse_benchmark.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/exps/attn/run_block_sparse_benchmark.py -------------------------------------------------------------------------------- /exps/attn/run_var_block_sparse_benchmark.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/exps/attn/run_var_block_sparse_benchmark.py -------------------------------------------------------------------------------- /exps/attn/tests/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/exps/attn/tests/__init__.py -------------------------------------------------------------------------------- /exps/device_a2av/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/exps/device_a2av/__init__.py -------------------------------------------------------------------------------- /exps/device_a2av/run.sh: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/exps/device_a2av/run.sh -------------------------------------------------------------------------------- /exps/device_a2av/test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/exps/device_a2av/test.py -------------------------------------------------------------------------------- /exps/dist_attn/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/exps/dist_attn/__init__.py -------------------------------------------------------------------------------- /exps/dist_attn/baselines/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/exps/dist_attn/baselines/__init__.py -------------------------------------------------------------------------------- /exps/dist_attn/baselines/interface.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/exps/dist_attn/baselines/interface.py -------------------------------------------------------------------------------- /exps/dist_attn/baselines/loongtrain.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/exps/dist_attn/baselines/loongtrain.py -------------------------------------------------------------------------------- /exps/dist_attn/baselines/nsa.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/exps/dist_attn/baselines/nsa.py -------------------------------------------------------------------------------- /exps/dist_attn/baselines/ring_attn.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/exps/dist_attn/baselines/ring_attn.py -------------------------------------------------------------------------------- /exps/dist_attn/baselines/shard.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/exps/dist_attn/baselines/shard.py -------------------------------------------------------------------------------- /exps/dist_attn/baselines/ulysess.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/exps/dist_attn/baselines/ulysess.py -------------------------------------------------------------------------------- /exps/dist_attn/baselines/usp.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/exps/dist_attn/baselines/usp.py -------------------------------------------------------------------------------- /exps/dist_attn/baselines/usp_nsa.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/exps/dist_attn/baselines/usp_nsa.py -------------------------------------------------------------------------------- /exps/dist_attn/baselines/utils_cp.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/exps/dist_attn/baselines/utils_cp.py -------------------------------------------------------------------------------- /exps/dist_attn/benchmark/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/exps/dist_attn/benchmark/__init__.py -------------------------------------------------------------------------------- /exps/dist_attn/benchmark/datasets/default/doc_length_distribution.csv: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/exps/dist_attn/benchmark/datasets/default/doc_length_distribution.csv -------------------------------------------------------------------------------- /exps/dist_attn/benchmark/enums.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/exps/dist_attn/benchmark/enums.py -------------------------------------------------------------------------------- /exps/dist_attn/benchmark/mask.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/exps/dist_attn/benchmark/mask.py -------------------------------------------------------------------------------- /exps/dist_attn/benchmark/metric.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/exps/dist_attn/benchmark/metric.py -------------------------------------------------------------------------------- /exps/dist_attn/benchmark/utils.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/exps/dist_attn/benchmark/utils.py -------------------------------------------------------------------------------- /exps/dist_attn/benchmark_conf.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/exps/dist_attn/benchmark_conf.py -------------------------------------------------------------------------------- /exps/dist_attn/main.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/exps/dist_attn/main.py -------------------------------------------------------------------------------- /exps/dist_attn/requirements.txt: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/exps/dist_attn/requirements.txt -------------------------------------------------------------------------------- /exps/dist_attn/run_benchmark.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/exps/dist_attn/run_benchmark.py -------------------------------------------------------------------------------- /exps/dist_attn/run_benchmark.sh: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/exps/dist_attn/run_benchmark.sh -------------------------------------------------------------------------------- /exps/dist_attn/run_profile.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/exps/dist_attn/run_profile.py -------------------------------------------------------------------------------- /exps/dist_attn/run_profile.sh: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/exps/dist_attn/run_profile.sh -------------------------------------------------------------------------------- /exps/dist_attn/tests/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/exps/dist_attn/tests/__init__.py -------------------------------------------------------------------------------- /exps/dist_attn/tests/attn_profile.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/exps/dist_attn/tests/attn_profile.py -------------------------------------------------------------------------------- /exps/dist_attn/tests/attn_profile.sh: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/exps/dist_attn/tests/attn_profile.sh -------------------------------------------------------------------------------- /exps/dist_attn/tests/nsa_profile.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/exps/dist_attn/tests/nsa_profile.py -------------------------------------------------------------------------------- /exps/dist_attn/tests/test_baseline_attn.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/exps/dist_attn/tests/test_baseline_attn.py -------------------------------------------------------------------------------- /exps/dist_attn/tests/test_usp_nsa.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/exps/dist_attn/tests/test_usp_nsa.py -------------------------------------------------------------------------------- /exps/dist_attn/tests/test_utils.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/exps/dist_attn/tests/test_utils.py -------------------------------------------------------------------------------- /exps/grpcoll/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/exps/grpcoll/__init__.py -------------------------------------------------------------------------------- /exps/grpcoll/grpcoll_utils.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/exps/grpcoll/grpcoll_utils.py -------------------------------------------------------------------------------- /exps/grpcoll/run_grpcoll_test.sh: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/exps/grpcoll/run_grpcoll_test.sh -------------------------------------------------------------------------------- /exps/grpcoll/test_internode_grpcoll.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/exps/grpcoll/test_internode_grpcoll.py -------------------------------------------------------------------------------- /exps/grpcoll/test_intranode_grpcoll.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/exps/grpcoll/test_intranode_grpcoll.py -------------------------------------------------------------------------------- /exps/grpcoll/test_low_latency_grpcoll.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/exps/grpcoll/test_low_latency_grpcoll.py -------------------------------------------------------------------------------- /extensions/README.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/extensions/README.md -------------------------------------------------------------------------------- /extensions/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/extensions/__init__.py -------------------------------------------------------------------------------- /extensions/magi_attn_extensions/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/extensions/magi_attn_extensions/__init__.py -------------------------------------------------------------------------------- /extensions/magi_attn_extensions/fa2_interface_with_sink.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/extensions/magi_attn_extensions/fa2_interface_with_sink.py -------------------------------------------------------------------------------- /extensions/magi_attn_extensions/fa3_interface_with_sink.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/extensions/magi_attn_extensions/fa3_interface_with_sink.py -------------------------------------------------------------------------------- /extensions/requirements.txt: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/extensions/requirements.txt -------------------------------------------------------------------------------- /extensions/setup.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/extensions/setup.py -------------------------------------------------------------------------------- /extensions/tests/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/extensions/tests/__init__.py -------------------------------------------------------------------------------- /extensions/tests/test_fa_interface_with_sink.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/extensions/tests/test_fa_interface_with_sink.py -------------------------------------------------------------------------------- /magi_attention/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/__init__.py -------------------------------------------------------------------------------- /magi_attention/api/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/api/__init__.py -------------------------------------------------------------------------------- /magi_attention/api/functools.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/api/functools.py -------------------------------------------------------------------------------- /magi_attention/api/magi_attn_interface.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/api/magi_attn_interface.py -------------------------------------------------------------------------------- /magi_attention/benchmarking/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/benchmarking/__init__.py -------------------------------------------------------------------------------- /magi_attention/benchmarking/bench.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/benchmarking/bench.py -------------------------------------------------------------------------------- /magi_attention/benchmarking/image_grid.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/benchmarking/image_grid.py -------------------------------------------------------------------------------- /magi_attention/benchmarking/utils.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/benchmarking/utils.py -------------------------------------------------------------------------------- /magi_attention/comm/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/comm/__init__.py -------------------------------------------------------------------------------- /magi_attention/comm/functional/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/comm/functional/__init__.py -------------------------------------------------------------------------------- /magi_attention/comm/functional/_gather_scatter_v.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/comm/functional/_gather_scatter_v.py -------------------------------------------------------------------------------- /magi_attention/comm/primitive/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/comm/primitive/__init__.py -------------------------------------------------------------------------------- /magi_attention/comm/primitive/_all2all_v.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/comm/primitive/_all2all_v.py -------------------------------------------------------------------------------- /magi_attention/comm/primitive/_all_gather_v.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/comm/primitive/_all_gather_v.py -------------------------------------------------------------------------------- /magi_attention/comm/primitive/_scatter_v.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/comm/primitive/_scatter_v.py -------------------------------------------------------------------------------- /magi_attention/comm/primitive/device_a2av/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/comm/primitive/device_a2av/__init__.py -------------------------------------------------------------------------------- /magi_attention/comm/primitive/device_a2av/interface.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/comm/primitive/device_a2av/interface.py -------------------------------------------------------------------------------- /magi_attention/comm/primitive/device_a2av/triton_barrier.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/comm/primitive/device_a2av/triton_barrier.py -------------------------------------------------------------------------------- /magi_attention/comm/primitive/device_a2av/triton_impl.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/comm/primitive/device_a2av/triton_impl.py -------------------------------------------------------------------------------- /magi_attention/comm/primitive/device_a2av/triton_utils.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/comm/primitive/device_a2av/triton_utils.py -------------------------------------------------------------------------------- /magi_attention/comm/primitive/grpcoll/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/comm/primitive/grpcoll/__init__.py -------------------------------------------------------------------------------- /magi_attention/comm/primitive/grpcoll/_a2av_grpcoll_impl.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/comm/primitive/grpcoll/_a2av_grpcoll_impl.py -------------------------------------------------------------------------------- /magi_attention/comm/primitive/grpcoll/_buffer.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/comm/primitive/grpcoll/_buffer.py -------------------------------------------------------------------------------- /magi_attention/comm/primitive/grpcoll/_config.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/comm/primitive/grpcoll/_config.py -------------------------------------------------------------------------------- /magi_attention/comm/primitive/grpcoll/_event.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/comm/primitive/grpcoll/_event.py -------------------------------------------------------------------------------- /magi_attention/comm/primitive/grpcoll/_group_collective.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/comm/primitive/grpcoll/_group_collective.py -------------------------------------------------------------------------------- /magi_attention/comm/primitive/grpcoll/_group_collective_hier.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/comm/primitive/grpcoll/_group_collective_hier.py -------------------------------------------------------------------------------- /magi_attention/comm/primitive/grpcoll/_handle.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/comm/primitive/grpcoll/_handle.py -------------------------------------------------------------------------------- /magi_attention/comm/primitive/grpcoll/_mgr.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/comm/primitive/grpcoll/_mgr.py -------------------------------------------------------------------------------- /magi_attention/comm/primitive/grpcoll/_native_grpcoll_impl.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/comm/primitive/grpcoll/_native_grpcoll_impl.py -------------------------------------------------------------------------------- /magi_attention/comm/primitive/grpcoll/utils.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/comm/primitive/grpcoll/utils.py -------------------------------------------------------------------------------- /magi_attention/comm/work.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/comm/work.py -------------------------------------------------------------------------------- /magi_attention/common/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/common/__init__.py -------------------------------------------------------------------------------- /magi_attention/common/enum.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/common/enum.py -------------------------------------------------------------------------------- /magi_attention/common/jit/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/common/jit/__init__.py -------------------------------------------------------------------------------- /magi_attention/common/jit/core.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/common/jit/core.py -------------------------------------------------------------------------------- /magi_attention/common/jit/cpp_ext.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/common/jit/cpp_ext.py -------------------------------------------------------------------------------- /magi_attention/common/jit/env.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/common/jit/env.py -------------------------------------------------------------------------------- /magi_attention/common/jit/utils.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/common/jit/utils.py -------------------------------------------------------------------------------- /magi_attention/common/mask.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/common/mask.py -------------------------------------------------------------------------------- /magi_attention/common/range.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/common/range.py -------------------------------------------------------------------------------- /magi_attention/common/range_op/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/common/range_op/__init__.py -------------------------------------------------------------------------------- /magi_attention/common/range_op/_range_fill.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/common/range_op/_range_fill.py -------------------------------------------------------------------------------- /magi_attention/common/range_op/_range_gather.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/common/range_op/_range_gather.py -------------------------------------------------------------------------------- /magi_attention/common/range_op/_range_reduce.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/common/range_op/_range_reduce.py -------------------------------------------------------------------------------- /magi_attention/common/range_op/utils.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/common/range_op/utils.py -------------------------------------------------------------------------------- /magi_attention/common/ranges.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/common/ranges.py -------------------------------------------------------------------------------- /magi_attention/common/rectangle.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/common/rectangle.py -------------------------------------------------------------------------------- /magi_attention/common/rectangles.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/common/rectangles.py -------------------------------------------------------------------------------- /magi_attention/config.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/config.py -------------------------------------------------------------------------------- /magi_attention/csrc/comm/grpcoll/buffer.cpp: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/csrc/comm/grpcoll/buffer.cpp -------------------------------------------------------------------------------- /magi_attention/csrc/comm/grpcoll/buffer.hpp: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/csrc/comm/grpcoll/buffer.hpp -------------------------------------------------------------------------------- /magi_attention/csrc/comm/grpcoll/config.hpp: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/csrc/comm/grpcoll/config.hpp -------------------------------------------------------------------------------- /magi_attention/csrc/comm/grpcoll/event.hpp: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/csrc/comm/grpcoll/event.hpp -------------------------------------------------------------------------------- /magi_attention/csrc/comm/grpcoll/kernels/api.cuh: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/csrc/comm/grpcoll/kernels/api.cuh -------------------------------------------------------------------------------- /magi_attention/csrc/comm/grpcoll/kernels/buffer.cuh: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/csrc/comm/grpcoll/kernels/buffer.cuh -------------------------------------------------------------------------------- /magi_attention/csrc/comm/grpcoll/kernels/configs.cuh: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/csrc/comm/grpcoll/kernels/configs.cuh -------------------------------------------------------------------------------- /magi_attention/csrc/comm/grpcoll/kernels/exception.cuh: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/csrc/comm/grpcoll/kernels/exception.cuh -------------------------------------------------------------------------------- /magi_attention/csrc/comm/grpcoll/kernels/ibgda_device.cuh: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/csrc/comm/grpcoll/kernels/ibgda_device.cuh -------------------------------------------------------------------------------- /magi_attention/csrc/comm/grpcoll/kernels/internode.cu: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/csrc/comm/grpcoll/kernels/internode.cu -------------------------------------------------------------------------------- /magi_attention/csrc/comm/grpcoll/kernels/internode_ll.cu: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/csrc/comm/grpcoll/kernels/internode_ll.cu -------------------------------------------------------------------------------- /magi_attention/csrc/comm/grpcoll/kernels/intranode.cu: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/csrc/comm/grpcoll/kernels/intranode.cu -------------------------------------------------------------------------------- /magi_attention/csrc/comm/grpcoll/kernels/launch.cuh: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/csrc/comm/grpcoll/kernels/launch.cuh -------------------------------------------------------------------------------- /magi_attention/csrc/comm/grpcoll/kernels/layout.cu: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/csrc/comm/grpcoll/kernels/layout.cu -------------------------------------------------------------------------------- /magi_attention/csrc/comm/grpcoll/kernels/reduce_op.cuh: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/csrc/comm/grpcoll/kernels/reduce_op.cuh -------------------------------------------------------------------------------- /magi_attention/csrc/comm/grpcoll/kernels/runtime.cu: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/csrc/comm/grpcoll/kernels/runtime.cu -------------------------------------------------------------------------------- /magi_attention/csrc/comm/grpcoll/kernels/utils.cuh: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/csrc/comm/grpcoll/kernels/utils.cuh -------------------------------------------------------------------------------- /magi_attention/csrc/comm/grpcoll/meta.hpp: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/csrc/comm/grpcoll/meta.hpp -------------------------------------------------------------------------------- /magi_attention/csrc/common/cuda_check.h: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/csrc/common/cuda_check.h -------------------------------------------------------------------------------- /magi_attention/csrc/common/static_switch.h: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/csrc/common/static_switch.h -------------------------------------------------------------------------------- /magi_attention/csrc/common/utils.h: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/csrc/common/utils.h -------------------------------------------------------------------------------- /magi_attention/csrc/extensions/attn_ranges.hpp: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/csrc/extensions/attn_ranges.hpp -------------------------------------------------------------------------------- /magi_attention/csrc/extensions/magi_attn_ext.cpp: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/csrc/extensions/magi_attn_ext.cpp -------------------------------------------------------------------------------- /magi_attention/csrc/extensions/magi_attn_ext.hpp: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/csrc/extensions/magi_attn_ext.hpp -------------------------------------------------------------------------------- /magi_attention/csrc/flexible_flash_attention/block.h: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/csrc/flexible_flash_attention/block.h -------------------------------------------------------------------------------- /magi_attention/csrc/flexible_flash_attention/bwd_inst_template.jinja: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/csrc/flexible_flash_attention/bwd_inst_template.jinja -------------------------------------------------------------------------------- /magi_attention/csrc/flexible_flash_attention/copy_sm90_bulk_reduce.hpp: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/csrc/flexible_flash_attention/copy_sm90_bulk_reduce.hpp -------------------------------------------------------------------------------- /magi_attention/csrc/flexible_flash_attention/epilogue_bwd.hpp: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/csrc/flexible_flash_attention/epilogue_bwd.hpp -------------------------------------------------------------------------------- /magi_attention/csrc/flexible_flash_attention/epilogue_fwd.hpp: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/csrc/flexible_flash_attention/epilogue_fwd.hpp -------------------------------------------------------------------------------- /magi_attention/csrc/flexible_flash_attention/flash.h: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/csrc/flexible_flash_attention/flash.h -------------------------------------------------------------------------------- /magi_attention/csrc/flexible_flash_attention/flash_bwd_kernel_sm90.h: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/csrc/flexible_flash_attention/flash_bwd_kernel_sm90.h -------------------------------------------------------------------------------- /magi_attention/csrc/flexible_flash_attention/flash_bwd_launch_template.h: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/csrc/flexible_flash_attention/flash_bwd_launch_template.h -------------------------------------------------------------------------------- /magi_attention/csrc/flexible_flash_attention/flash_bwd_preprocess_kernel.h: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/csrc/flexible_flash_attention/flash_bwd_preprocess_kernel.h -------------------------------------------------------------------------------- /magi_attention/csrc/flexible_flash_attention/flash_fwd_kernel_sm90.h: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/csrc/flexible_flash_attention/flash_fwd_kernel_sm90.h -------------------------------------------------------------------------------- /magi_attention/csrc/flexible_flash_attention/flash_fwd_launch_template.h: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/csrc/flexible_flash_attention/flash_fwd_launch_template.h -------------------------------------------------------------------------------- /magi_attention/csrc/flexible_flash_attention/flash_fwd_postprocess.cu: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/csrc/flexible_flash_attention/flash_fwd_postprocess.cu -------------------------------------------------------------------------------- /magi_attention/csrc/flexible_flash_attention/flash_fwd_postprocess_kernel.h: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/csrc/flexible_flash_attention/flash_fwd_postprocess_kernel.h -------------------------------------------------------------------------------- /magi_attention/csrc/flexible_flash_attention/flash_fwd_postprocess_launch_template.h: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/csrc/flexible_flash_attention/flash_fwd_postprocess_launch_template.h -------------------------------------------------------------------------------- /magi_attention/csrc/flexible_flash_attention/flex_flash_bwd.hpp: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/csrc/flexible_flash_attention/flex_flash_bwd.hpp -------------------------------------------------------------------------------- /magi_attention/csrc/flexible_flash_attention/flex_flash_common.cpp: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/csrc/flexible_flash_attention/flex_flash_common.cpp -------------------------------------------------------------------------------- /magi_attention/csrc/flexible_flash_attention/flex_flash_common.hpp: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/csrc/flexible_flash_attention/flex_flash_common.hpp -------------------------------------------------------------------------------- /magi_attention/csrc/flexible_flash_attention/flex_flash_fwd.hpp: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/csrc/flexible_flash_attention/flex_flash_fwd.hpp -------------------------------------------------------------------------------- /magi_attention/csrc/flexible_flash_attention/fwd_inst_template.jinja: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/csrc/flexible_flash_attention/fwd_inst_template.jinja -------------------------------------------------------------------------------- /magi_attention/csrc/flexible_flash_attention/mainloop_bwd_sm90_tma_gmma_ws.hpp: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/csrc/flexible_flash_attention/mainloop_bwd_sm90_tma_gmma_ws.hpp -------------------------------------------------------------------------------- /magi_attention/csrc/flexible_flash_attention/mainloop_fwd_sm90_tma_gmma_ws.hpp: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/csrc/flexible_flash_attention/mainloop_fwd_sm90_tma_gmma_ws.hpp -------------------------------------------------------------------------------- /magi_attention/csrc/flexible_flash_attention/mask.h: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/csrc/flexible_flash_attention/mask.h -------------------------------------------------------------------------------- /magi_attention/csrc/flexible_flash_attention/named_barrier.hpp: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/csrc/flexible_flash_attention/named_barrier.hpp -------------------------------------------------------------------------------- /magi_attention/csrc/flexible_flash_attention/seqlen.h: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/csrc/flexible_flash_attention/seqlen.h -------------------------------------------------------------------------------- /magi_attention/csrc/flexible_flash_attention/sink_layout.cuh: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/csrc/flexible_flash_attention/sink_layout.cuh -------------------------------------------------------------------------------- /magi_attention/csrc/flexible_flash_attention/sm90_pipeline_no_cluster.hpp: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/csrc/flexible_flash_attention/sm90_pipeline_no_cluster.hpp -------------------------------------------------------------------------------- /magi_attention/csrc/flexible_flash_attention/softmax.h: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/csrc/flexible_flash_attention/softmax.h -------------------------------------------------------------------------------- /magi_attention/csrc/flexible_flash_attention/tile_scheduler.hpp: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/csrc/flexible_flash_attention/tile_scheduler.hpp -------------------------------------------------------------------------------- /magi_attention/csrc/flexible_flash_attention/tile_size.h: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/csrc/flexible_flash_attention/tile_size.h -------------------------------------------------------------------------------- /magi_attention/csrc/utils/bindings.cpp: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/csrc/utils/bindings.cpp -------------------------------------------------------------------------------- /magi_attention/csrc/utils/profile_utils.cu: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/csrc/utils/profile_utils.cu -------------------------------------------------------------------------------- /magi_attention/csrc/utils/profile_utils.h: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/csrc/utils/profile_utils.h -------------------------------------------------------------------------------- /magi_attention/csrc/utils/unique_consecutive_pairs.cu: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/csrc/utils/unique_consecutive_pairs.cu -------------------------------------------------------------------------------- /magi_attention/dist_attn_runtime_mgr.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/dist_attn_runtime_mgr.py -------------------------------------------------------------------------------- /magi_attention/functional/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/functional/__init__.py -------------------------------------------------------------------------------- /magi_attention/functional/_flex_flash_attn_jit.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/functional/_flex_flash_attn_jit.py -------------------------------------------------------------------------------- /magi_attention/functional/dispatch.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/functional/dispatch.py -------------------------------------------------------------------------------- /magi_attention/functional/dist_attn.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/functional/dist_attn.py -------------------------------------------------------------------------------- /magi_attention/functional/flex_flash_attn.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/functional/flex_flash_attn.py -------------------------------------------------------------------------------- /magi_attention/functional/sdpa.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/functional/sdpa.py -------------------------------------------------------------------------------- /magi_attention/functional/utils.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/functional/utils.py -------------------------------------------------------------------------------- /magi_attention/meta/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/meta/__init__.py -------------------------------------------------------------------------------- /magi_attention/meta/_make_attn_meta.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/meta/_make_attn_meta.py -------------------------------------------------------------------------------- /magi_attention/meta/_make_dispatch_meta.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/meta/_make_dispatch_meta.py -------------------------------------------------------------------------------- /magi_attention/meta/algorithms/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/meta/algorithms/__init__.py -------------------------------------------------------------------------------- /magi_attention/meta/algorithms/base.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/meta/algorithms/base.py -------------------------------------------------------------------------------- /magi_attention/meta/algorithms/grg.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/meta/algorithms/grg.py -------------------------------------------------------------------------------- /magi_attention/meta/algorithms/ncq.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/meta/algorithms/ncq.py -------------------------------------------------------------------------------- /magi_attention/meta/collection/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/meta/collection/__init__.py -------------------------------------------------------------------------------- /magi_attention/meta/collection/calc_meta.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/meta/collection/calc_meta.py -------------------------------------------------------------------------------- /magi_attention/meta/collection/comm_meta.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/meta/collection/comm_meta.py -------------------------------------------------------------------------------- /magi_attention/meta/collection/dispatch_meta.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/meta/collection/dispatch_meta.py -------------------------------------------------------------------------------- /magi_attention/meta/container/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/meta/container/__init__.py -------------------------------------------------------------------------------- /magi_attention/meta/container/bucket.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/meta/container/bucket.py -------------------------------------------------------------------------------- /magi_attention/meta/container/chunk.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/meta/container/chunk.py -------------------------------------------------------------------------------- /magi_attention/meta/container/rank_entry.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/meta/container/rank_entry.py -------------------------------------------------------------------------------- /magi_attention/meta/container/slice.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/meta/container/slice.py -------------------------------------------------------------------------------- /magi_attention/meta/container/transfer_table.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/meta/container/transfer_table.py -------------------------------------------------------------------------------- /magi_attention/meta/solver/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/meta/solver/__init__.py -------------------------------------------------------------------------------- /magi_attention/meta/solver/dispatch_solver.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/meta/solver/dispatch_solver.py -------------------------------------------------------------------------------- /magi_attention/meta/solver/dist_attn_solver.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/meta/solver/dist_attn_solver.py -------------------------------------------------------------------------------- /magi_attention/meta/solver/dynamic_attn_solver.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/meta/solver/dynamic_attn_solver.py -------------------------------------------------------------------------------- /magi_attention/meta/solver/overlap_solver.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/meta/solver/overlap_solver.py -------------------------------------------------------------------------------- /magi_attention/meta/solver/slice_maker.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/meta/solver/slice_maker.py -------------------------------------------------------------------------------- /magi_attention/testing/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/testing/__init__.py -------------------------------------------------------------------------------- /magi_attention/testing/dist_common.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/testing/dist_common.py -------------------------------------------------------------------------------- /magi_attention/testing/flag_generator.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/testing/flag_generator.py -------------------------------------------------------------------------------- /magi_attention/testing/gt_dispatcher.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/testing/gt_dispatcher.py -------------------------------------------------------------------------------- /magi_attention/testing/precision.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/testing/precision.py -------------------------------------------------------------------------------- /magi_attention/testing/ref_attn.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/testing/ref_attn.py -------------------------------------------------------------------------------- /magi_attention/testing/utils.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/testing/utils.py -------------------------------------------------------------------------------- /magi_attention/utils/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/utils/__init__.py -------------------------------------------------------------------------------- /magi_attention/utils/_utils.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/utils/_utils.py -------------------------------------------------------------------------------- /magi_attention/utils/debug.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/utils/debug.py -------------------------------------------------------------------------------- /magi_attention/utils/metaclass.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/utils/metaclass.py -------------------------------------------------------------------------------- /magi_attention/utils/nvtx.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/utils/nvtx.py -------------------------------------------------------------------------------- /magi_attention/utils/sparse_utils.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/magi_attention/utils/sparse_utils.py -------------------------------------------------------------------------------- /makefile: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/makefile -------------------------------------------------------------------------------- /mypy.ini: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/mypy.ini -------------------------------------------------------------------------------- /pyproject.toml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/pyproject.toml -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/requirements.txt -------------------------------------------------------------------------------- /requirements_dev.txt: -------------------------------------------------------------------------------- 1 | pre-commit==4.0.1 2 | coverage 3 | pytest-sugar 4 | build 5 | -------------------------------------------------------------------------------- /scripts/install_clang_format.sh: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/scripts/install_clang_format.sh -------------------------------------------------------------------------------- /scripts/run_csrc_code_formatter.sh: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/scripts/run_csrc_code_formatter.sh -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/setup.py -------------------------------------------------------------------------------- /tests/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/tests/__init__.py -------------------------------------------------------------------------------- /tests/test_api/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/tests/test_api/__init__.py -------------------------------------------------------------------------------- /tests/test_api/test_functools.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/tests/test_api/test_functools.py -------------------------------------------------------------------------------- /tests/test_api/test_interface.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/tests/test_api/test_interface.py -------------------------------------------------------------------------------- /tests/test_attn/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/tests/test_attn/__init__.py -------------------------------------------------------------------------------- /tests/test_attn/test_block_sparse_attn.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/tests/test_attn/test_block_sparse_attn.py -------------------------------------------------------------------------------- /tests/test_attn/test_dist_attn.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/tests/test_attn/test_dist_attn.py -------------------------------------------------------------------------------- /tests/test_attn/test_flex_flash_attn.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/tests/test_attn/test_flex_flash_attn.py -------------------------------------------------------------------------------- /tests/test_attn/test_merge_range.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/tests/test_attn/test_merge_range.py -------------------------------------------------------------------------------- /tests/test_attn/test_ref_attn.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/tests/test_attn/test_ref_attn.py -------------------------------------------------------------------------------- /tests/test_attn_solver/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/tests/test_attn_solver/__init__.py -------------------------------------------------------------------------------- /tests/test_attn_solver/test_dist_attn_solver.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/tests/test_attn_solver/test_dist_attn_solver.py -------------------------------------------------------------------------------- /tests/test_attn_solver/test_dynamic_attn_solver.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/tests/test_attn_solver/test_dynamic_attn_solver.py -------------------------------------------------------------------------------- /tests/test_comm/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/tests/test_comm/__init__.py -------------------------------------------------------------------------------- /tests/test_comm/test_all_gather_v.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/tests/test_comm/test_all_gather_v.py -------------------------------------------------------------------------------- /tests/test_comm/test_group_collective.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/tests/test_comm/test_group_collective.py -------------------------------------------------------------------------------- /tests/test_comm/test_group_collective_utils.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/tests/test_comm/test_group_collective_utils.py -------------------------------------------------------------------------------- /tests/test_comm/test_scatter_v.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/tests/test_comm/test_scatter_v.py -------------------------------------------------------------------------------- /tests/test_common/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/tests/test_common/__init__.py -------------------------------------------------------------------------------- /tests/test_common/test_attn_mask.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/tests/test_common/test_attn_mask.py -------------------------------------------------------------------------------- /tests/test_common/test_attn_range.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/tests/test_common/test_attn_range.py -------------------------------------------------------------------------------- /tests/test_common/test_attn_ranges.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/tests/test_common/test_attn_ranges.py -------------------------------------------------------------------------------- /tests/test_common/test_range_op/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/tests/test_common/test_range_op/__init__.py -------------------------------------------------------------------------------- /tests/test_common/test_range_op/test_range_fill.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/tests/test_common/test_range_op/test_range_fill.py -------------------------------------------------------------------------------- /tests/test_common/test_range_op/test_range_gather.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/tests/test_common/test_range_op/test_range_gather.py -------------------------------------------------------------------------------- /tests/test_common/test_range_op/test_range_op_utils.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/tests/test_common/test_range_op/test_range_op_utils.py -------------------------------------------------------------------------------- /tests/test_common/test_range_op/test_range_reduce.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/tests/test_common/test_range_op/test_range_reduce.py -------------------------------------------------------------------------------- /tests/test_common/test_rectangle.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/tests/test_common/test_rectangle.py -------------------------------------------------------------------------------- /tests/test_common/test_rectangles.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/tests/test_common/test_rectangles.py -------------------------------------------------------------------------------- /tests/test_dispatch/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/tests/test_dispatch/__init__.py -------------------------------------------------------------------------------- /tests/test_dispatch/test_calc_self_attn_areas.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/tests/test_dispatch/test_calc_self_attn_areas.py -------------------------------------------------------------------------------- /tests/test_dispatch/test_dispatch_solver.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/tests/test_dispatch/test_dispatch_solver.py -------------------------------------------------------------------------------- /tests/test_dispatch/test_dispatcher.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/tests/test_dispatch/test_dispatcher.py -------------------------------------------------------------------------------- /tests/test_dispatch/test_gt_dispatcher.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/tests/test_dispatch/test_gt_dispatcher.py -------------------------------------------------------------------------------- /tests/test_dist_runtime_mgr/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/tests/test_dist_runtime_mgr/__init__.py -------------------------------------------------------------------------------- /tests/test_dist_runtime_mgr/test_dist_runtime_mgr.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/tests/test_dist_runtime_mgr/test_dist_runtime_mgr.py -------------------------------------------------------------------------------- /tests/test_functional/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/tests/test_functional/__init__.py -------------------------------------------------------------------------------- /tests/test_pipeline.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/tests/test_pipeline.py -------------------------------------------------------------------------------- /tests/test_pipeline_sdpa.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/tests/test_pipeline_sdpa.py -------------------------------------------------------------------------------- /tests/test_utils/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/tests/test_utils/__init__.py -------------------------------------------------------------------------------- /tests/test_utils/test_common_utils.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/tests/test_utils/test_common_utils.py -------------------------------------------------------------------------------- /tests/test_utils/test_flag_generator.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/tests/test_utils/test_flag_generator.py -------------------------------------------------------------------------------- /tools/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/tools/__init__.py -------------------------------------------------------------------------------- /tools/build_helper.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/tools/build_helper.py -------------------------------------------------------------------------------- /tools/codestyle/check_for_chinese.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/tools/codestyle/check_for_chinese.py -------------------------------------------------------------------------------- /tools/codestyle/copyright.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SandAI-org/MagiAttention/HEAD/tools/codestyle/copyright.py --------------------------------------------------------------------------------