├── .gitignore ├── .gitmodules ├── Dockerfile ├── LICENSE ├── MANIFEST.in ├── Makefile ├── README.md ├── csrc ├── cuda_util.h ├── cumsum.h ├── histogram.h ├── indices.h ├── ops.cu ├── replicate.h └── sort.h ├── docker.sh ├── exp ├── dmoe │ ├── dmoe_125m_8gpu.sh │ ├── dmoe_356m_8gpu.sh │ ├── dmoe_46m_8gpu.sh │ └── dmoe_760m_8gpu.sh ├── gpt2 │ ├── gpt2_125m_1gpu.sh │ ├── gpt2_125m_8gpu.sh │ ├── gpt2_1315m_1gpu.sh │ ├── gpt2_1315m_8gpu.sh │ ├── gpt2_356m_1gpu.sh │ ├── gpt2_356m_8gpu.sh │ ├── gpt2_46m_1gpu.sh │ ├── gpt2_46m_8gpu.sh │ ├── gpt2_760m_1gpu.sh │ └── gpt2_760m_8gpu.sh └── moe │ ├── moe_125m_8gpu.sh │ ├── moe_356m_8gpu.sh │ └── moe_46m_8gpu.sh ├── media └── dropping_end_to_end.png ├── megablocks ├── __init__.py ├── backend │ ├── __init__.py │ └── kernels.py ├── benchmark_util.py ├── grouped_gemm_util.py ├── layers │ ├── __init__.py │ ├── activation_fn.py │ ├── all_to_all.py │ ├── arguments.py │ ├── common.py │ ├── dmlp_registry.py │ ├── dmoe.py │ ├── dmoe_test.py │ ├── gelu.py │ ├── glu.py │ ├── glu_test.py │ ├── memory_test.py │ ├── memory_test.sh │ ├── mlp.py │ ├── moa.py │ ├── moe.py │ ├── moe_test.py │ ├── mpu.py │ ├── parallelism_test.py │ ├── parallelism_test.sh │ ├── router.py │ ├── testing.py │ └── weight_parallel.py └── ops │ ├── __init__.py │ ├── all_to_all_benchmark.py │ ├── all_to_all_benchmark.sh │ ├── binned_gather.py │ ├── binned_gather_test.py │ ├── binned_scatter.py │ ├── binned_scatter_test.py │ ├── cumsum.py │ ├── cumsum_test.py │ ├── gather.py │ ├── histogram.py │ ├── histogram_benchmark.py │ ├── histogram_test.py │ ├── matmul_benchmark.py │ ├── padded_gather.py │ ├── padded_gather_test.py │ ├── padded_scatter.py │ ├── padded_scatter_benchmark.py │ ├── padded_scatter_test.py │ ├── permute_benchmark.py │ ├── repeat.py │ ├── replicate.py │ ├── replicate_test.py │ ├── round_up.py │ ├── scatter.py │ ├── sort.py │ ├── sort_benchmark.py │ ├── sort_test.py │ ├── sum.py │ ├── topology.py │ └── topology_test.py ├── requirements.txt ├── setup.py └── yamls ├── matmul_benchmark.yaml └── triton_benchmark.yaml /.gitignore: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yikangshen/megablocks/HEAD/.gitignore -------------------------------------------------------------------------------- /.gitmodules: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yikangshen/megablocks/HEAD/.gitmodules -------------------------------------------------------------------------------- /Dockerfile: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yikangshen/megablocks/HEAD/Dockerfile -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yikangshen/megablocks/HEAD/LICENSE -------------------------------------------------------------------------------- /MANIFEST.in: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yikangshen/megablocks/HEAD/MANIFEST.in -------------------------------------------------------------------------------- /Makefile: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yikangshen/megablocks/HEAD/Makefile -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yikangshen/megablocks/HEAD/README.md -------------------------------------------------------------------------------- /csrc/cuda_util.h: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yikangshen/megablocks/HEAD/csrc/cuda_util.h -------------------------------------------------------------------------------- /csrc/cumsum.h: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yikangshen/megablocks/HEAD/csrc/cumsum.h -------------------------------------------------------------------------------- /csrc/histogram.h: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yikangshen/megablocks/HEAD/csrc/histogram.h -------------------------------------------------------------------------------- /csrc/indices.h: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yikangshen/megablocks/HEAD/csrc/indices.h -------------------------------------------------------------------------------- /csrc/ops.cu: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yikangshen/megablocks/HEAD/csrc/ops.cu -------------------------------------------------------------------------------- /csrc/replicate.h: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yikangshen/megablocks/HEAD/csrc/replicate.h -------------------------------------------------------------------------------- /csrc/sort.h: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yikangshen/megablocks/HEAD/csrc/sort.h -------------------------------------------------------------------------------- /docker.sh: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yikangshen/megablocks/HEAD/docker.sh -------------------------------------------------------------------------------- /exp/dmoe/dmoe_125m_8gpu.sh: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yikangshen/megablocks/HEAD/exp/dmoe/dmoe_125m_8gpu.sh -------------------------------------------------------------------------------- /exp/dmoe/dmoe_356m_8gpu.sh: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yikangshen/megablocks/HEAD/exp/dmoe/dmoe_356m_8gpu.sh -------------------------------------------------------------------------------- /exp/dmoe/dmoe_46m_8gpu.sh: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yikangshen/megablocks/HEAD/exp/dmoe/dmoe_46m_8gpu.sh -------------------------------------------------------------------------------- /exp/dmoe/dmoe_760m_8gpu.sh: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yikangshen/megablocks/HEAD/exp/dmoe/dmoe_760m_8gpu.sh -------------------------------------------------------------------------------- /exp/gpt2/gpt2_125m_1gpu.sh: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yikangshen/megablocks/HEAD/exp/gpt2/gpt2_125m_1gpu.sh -------------------------------------------------------------------------------- /exp/gpt2/gpt2_125m_8gpu.sh: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yikangshen/megablocks/HEAD/exp/gpt2/gpt2_125m_8gpu.sh -------------------------------------------------------------------------------- /exp/gpt2/gpt2_1315m_1gpu.sh: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yikangshen/megablocks/HEAD/exp/gpt2/gpt2_1315m_1gpu.sh -------------------------------------------------------------------------------- /exp/gpt2/gpt2_1315m_8gpu.sh: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yikangshen/megablocks/HEAD/exp/gpt2/gpt2_1315m_8gpu.sh -------------------------------------------------------------------------------- /exp/gpt2/gpt2_356m_1gpu.sh: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yikangshen/megablocks/HEAD/exp/gpt2/gpt2_356m_1gpu.sh -------------------------------------------------------------------------------- /exp/gpt2/gpt2_356m_8gpu.sh: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yikangshen/megablocks/HEAD/exp/gpt2/gpt2_356m_8gpu.sh -------------------------------------------------------------------------------- /exp/gpt2/gpt2_46m_1gpu.sh: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yikangshen/megablocks/HEAD/exp/gpt2/gpt2_46m_1gpu.sh -------------------------------------------------------------------------------- /exp/gpt2/gpt2_46m_8gpu.sh: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yikangshen/megablocks/HEAD/exp/gpt2/gpt2_46m_8gpu.sh -------------------------------------------------------------------------------- /exp/gpt2/gpt2_760m_1gpu.sh: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yikangshen/megablocks/HEAD/exp/gpt2/gpt2_760m_1gpu.sh -------------------------------------------------------------------------------- /exp/gpt2/gpt2_760m_8gpu.sh: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yikangshen/megablocks/HEAD/exp/gpt2/gpt2_760m_8gpu.sh -------------------------------------------------------------------------------- /exp/moe/moe_125m_8gpu.sh: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yikangshen/megablocks/HEAD/exp/moe/moe_125m_8gpu.sh -------------------------------------------------------------------------------- /exp/moe/moe_356m_8gpu.sh: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yikangshen/megablocks/HEAD/exp/moe/moe_356m_8gpu.sh -------------------------------------------------------------------------------- /exp/moe/moe_46m_8gpu.sh: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yikangshen/megablocks/HEAD/exp/moe/moe_46m_8gpu.sh -------------------------------------------------------------------------------- /media/dropping_end_to_end.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yikangshen/megablocks/HEAD/media/dropping_end_to_end.png -------------------------------------------------------------------------------- /megablocks/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yikangshen/megablocks/HEAD/megablocks/__init__.py -------------------------------------------------------------------------------- /megablocks/backend/__init__.py: -------------------------------------------------------------------------------- 1 | 2 | -------------------------------------------------------------------------------- /megablocks/backend/kernels.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yikangshen/megablocks/HEAD/megablocks/backend/kernels.py -------------------------------------------------------------------------------- /megablocks/benchmark_util.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yikangshen/megablocks/HEAD/megablocks/benchmark_util.py -------------------------------------------------------------------------------- /megablocks/grouped_gemm_util.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yikangshen/megablocks/HEAD/megablocks/grouped_gemm_util.py -------------------------------------------------------------------------------- /megablocks/layers/__init__.py: -------------------------------------------------------------------------------- 1 | 2 | -------------------------------------------------------------------------------- /megablocks/layers/activation_fn.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yikangshen/megablocks/HEAD/megablocks/layers/activation_fn.py -------------------------------------------------------------------------------- /megablocks/layers/all_to_all.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yikangshen/megablocks/HEAD/megablocks/layers/all_to_all.py -------------------------------------------------------------------------------- /megablocks/layers/arguments.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yikangshen/megablocks/HEAD/megablocks/layers/arguments.py -------------------------------------------------------------------------------- /megablocks/layers/common.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yikangshen/megablocks/HEAD/megablocks/layers/common.py -------------------------------------------------------------------------------- /megablocks/layers/dmlp_registry.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yikangshen/megablocks/HEAD/megablocks/layers/dmlp_registry.py -------------------------------------------------------------------------------- /megablocks/layers/dmoe.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yikangshen/megablocks/HEAD/megablocks/layers/dmoe.py -------------------------------------------------------------------------------- /megablocks/layers/dmoe_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yikangshen/megablocks/HEAD/megablocks/layers/dmoe_test.py -------------------------------------------------------------------------------- /megablocks/layers/gelu.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yikangshen/megablocks/HEAD/megablocks/layers/gelu.py -------------------------------------------------------------------------------- /megablocks/layers/glu.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yikangshen/megablocks/HEAD/megablocks/layers/glu.py -------------------------------------------------------------------------------- /megablocks/layers/glu_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yikangshen/megablocks/HEAD/megablocks/layers/glu_test.py -------------------------------------------------------------------------------- /megablocks/layers/memory_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yikangshen/megablocks/HEAD/megablocks/layers/memory_test.py -------------------------------------------------------------------------------- /megablocks/layers/memory_test.sh: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yikangshen/megablocks/HEAD/megablocks/layers/memory_test.sh -------------------------------------------------------------------------------- /megablocks/layers/mlp.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yikangshen/megablocks/HEAD/megablocks/layers/mlp.py -------------------------------------------------------------------------------- /megablocks/layers/moa.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yikangshen/megablocks/HEAD/megablocks/layers/moa.py -------------------------------------------------------------------------------- /megablocks/layers/moe.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yikangshen/megablocks/HEAD/megablocks/layers/moe.py -------------------------------------------------------------------------------- /megablocks/layers/moe_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yikangshen/megablocks/HEAD/megablocks/layers/moe_test.py -------------------------------------------------------------------------------- /megablocks/layers/mpu.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yikangshen/megablocks/HEAD/megablocks/layers/mpu.py -------------------------------------------------------------------------------- /megablocks/layers/parallelism_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yikangshen/megablocks/HEAD/megablocks/layers/parallelism_test.py -------------------------------------------------------------------------------- /megablocks/layers/parallelism_test.sh: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yikangshen/megablocks/HEAD/megablocks/layers/parallelism_test.sh -------------------------------------------------------------------------------- /megablocks/layers/router.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yikangshen/megablocks/HEAD/megablocks/layers/router.py -------------------------------------------------------------------------------- /megablocks/layers/testing.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yikangshen/megablocks/HEAD/megablocks/layers/testing.py -------------------------------------------------------------------------------- /megablocks/layers/weight_parallel.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yikangshen/megablocks/HEAD/megablocks/layers/weight_parallel.py -------------------------------------------------------------------------------- /megablocks/ops/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yikangshen/megablocks/HEAD/megablocks/ops/__init__.py -------------------------------------------------------------------------------- /megablocks/ops/all_to_all_benchmark.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yikangshen/megablocks/HEAD/megablocks/ops/all_to_all_benchmark.py -------------------------------------------------------------------------------- /megablocks/ops/all_to_all_benchmark.sh: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yikangshen/megablocks/HEAD/megablocks/ops/all_to_all_benchmark.sh -------------------------------------------------------------------------------- /megablocks/ops/binned_gather.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yikangshen/megablocks/HEAD/megablocks/ops/binned_gather.py -------------------------------------------------------------------------------- /megablocks/ops/binned_gather_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yikangshen/megablocks/HEAD/megablocks/ops/binned_gather_test.py -------------------------------------------------------------------------------- /megablocks/ops/binned_scatter.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yikangshen/megablocks/HEAD/megablocks/ops/binned_scatter.py -------------------------------------------------------------------------------- /megablocks/ops/binned_scatter_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yikangshen/megablocks/HEAD/megablocks/ops/binned_scatter_test.py -------------------------------------------------------------------------------- /megablocks/ops/cumsum.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yikangshen/megablocks/HEAD/megablocks/ops/cumsum.py -------------------------------------------------------------------------------- /megablocks/ops/cumsum_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yikangshen/megablocks/HEAD/megablocks/ops/cumsum_test.py -------------------------------------------------------------------------------- /megablocks/ops/gather.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yikangshen/megablocks/HEAD/megablocks/ops/gather.py -------------------------------------------------------------------------------- /megablocks/ops/histogram.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yikangshen/megablocks/HEAD/megablocks/ops/histogram.py -------------------------------------------------------------------------------- /megablocks/ops/histogram_benchmark.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yikangshen/megablocks/HEAD/megablocks/ops/histogram_benchmark.py -------------------------------------------------------------------------------- /megablocks/ops/histogram_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yikangshen/megablocks/HEAD/megablocks/ops/histogram_test.py -------------------------------------------------------------------------------- /megablocks/ops/matmul_benchmark.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yikangshen/megablocks/HEAD/megablocks/ops/matmul_benchmark.py -------------------------------------------------------------------------------- /megablocks/ops/padded_gather.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yikangshen/megablocks/HEAD/megablocks/ops/padded_gather.py -------------------------------------------------------------------------------- /megablocks/ops/padded_gather_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yikangshen/megablocks/HEAD/megablocks/ops/padded_gather_test.py -------------------------------------------------------------------------------- /megablocks/ops/padded_scatter.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yikangshen/megablocks/HEAD/megablocks/ops/padded_scatter.py -------------------------------------------------------------------------------- /megablocks/ops/padded_scatter_benchmark.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yikangshen/megablocks/HEAD/megablocks/ops/padded_scatter_benchmark.py -------------------------------------------------------------------------------- /megablocks/ops/padded_scatter_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yikangshen/megablocks/HEAD/megablocks/ops/padded_scatter_test.py -------------------------------------------------------------------------------- /megablocks/ops/permute_benchmark.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yikangshen/megablocks/HEAD/megablocks/ops/permute_benchmark.py -------------------------------------------------------------------------------- /megablocks/ops/repeat.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yikangshen/megablocks/HEAD/megablocks/ops/repeat.py -------------------------------------------------------------------------------- /megablocks/ops/replicate.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yikangshen/megablocks/HEAD/megablocks/ops/replicate.py -------------------------------------------------------------------------------- /megablocks/ops/replicate_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yikangshen/megablocks/HEAD/megablocks/ops/replicate_test.py -------------------------------------------------------------------------------- /megablocks/ops/round_up.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yikangshen/megablocks/HEAD/megablocks/ops/round_up.py -------------------------------------------------------------------------------- /megablocks/ops/scatter.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yikangshen/megablocks/HEAD/megablocks/ops/scatter.py -------------------------------------------------------------------------------- /megablocks/ops/sort.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yikangshen/megablocks/HEAD/megablocks/ops/sort.py -------------------------------------------------------------------------------- /megablocks/ops/sort_benchmark.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yikangshen/megablocks/HEAD/megablocks/ops/sort_benchmark.py -------------------------------------------------------------------------------- /megablocks/ops/sort_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yikangshen/megablocks/HEAD/megablocks/ops/sort_test.py -------------------------------------------------------------------------------- /megablocks/ops/sum.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yikangshen/megablocks/HEAD/megablocks/ops/sum.py -------------------------------------------------------------------------------- /megablocks/ops/topology.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yikangshen/megablocks/HEAD/megablocks/ops/topology.py -------------------------------------------------------------------------------- /megablocks/ops/topology_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yikangshen/megablocks/HEAD/megablocks/ops/topology_test.py -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | stanford-stk>=0.0.6 2 | triton==2.1.0 3 | -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yikangshen/megablocks/HEAD/setup.py -------------------------------------------------------------------------------- /yamls/matmul_benchmark.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yikangshen/megablocks/HEAD/yamls/matmul_benchmark.yaml -------------------------------------------------------------------------------- /yamls/triton_benchmark.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yikangshen/megablocks/HEAD/yamls/triton_benchmark.yaml --------------------------------------------------------------------------------