diff --git a/src/torchada/triton/autotune/fused_moe/configs/triton_3_2_0/E=768,N=1280,device_name=MTT_S5000.json b/src/torchada/triton/autotune/fused_moe/configs/triton_3_2_0/E=768,N=1280,device_name=MTT_S5000.json new file mode 100644 index 0000000..32c7856 --- /dev/null +++ b/src/torchada/triton/autotune/fused_moe/configs/triton_3_2_0/E=768,N=1280,device_name=MTT_S5000.json @@ -0,0 +1,26 @@ +{ + "21996": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 16, + "num_warps": 16, + "num_stages": 1 + }, + "24468": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 16, + "num_warps": 16, + "num_stages": 1 + }, + "45012": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 16, + "num_warps": 16, + "num_stages": 1 + } +} diff --git a/src/torchada/triton/autotune/fused_moe/configs/triton_3_2_0/E=768,N=1280,device_name=MTT_S5000_down.json b/src/torchada/triton/autotune/fused_moe/configs/triton_3_2_0/E=768,N=1280,device_name=MTT_S5000_down.json new file mode 100644 index 0000000..99f9bbc --- /dev/null +++ b/src/torchada/triton/autotune/fused_moe/configs/triton_3_2_0/E=768,N=1280,device_name=MTT_S5000_down.json @@ -0,0 +1,26 @@ +{ + "21996": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 16, + "num_warps": 16, + "num_stages": 1 + }, + "24468": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 16, + "num_warps": 16, + "num_stages": 1 + }, + "45012": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 16, + "num_stages": 1 + } +} diff --git a/src/torchada/triton/autotune/fused_moe/configs/triton_3_6_0/E=768,N=1280,device_name=MTT_S5000,dtype=bf16.json b/src/torchada/triton/autotune/fused_moe/configs/triton_3_6_0/E=768,N=1280,device_name=MTT_S5000,dtype=bf16.json new file mode 100644 index 0000000..ba6cb1c --- /dev/null +++ b/src/torchada/triton/autotune/fused_moe/configs/triton_3_6_0/E=768,N=1280,device_name=MTT_S5000,dtype=bf16.json @@ -0,0 +1,10 @@ +{ + "45012": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 16, + "num_warps": 16, + "num_stages": 1 + } +} diff --git a/src/torchada/triton/runtime/fused_moe/config.py b/src/torchada/triton/runtime/fused_moe/config.py index c34c3d6..19da803 100644 --- a/src/torchada/triton/runtime/fused_moe/config.py +++ b/src/torchada/triton/runtime/fused_moe/config.py @@ -71,7 +71,7 @@ def get_moe_configs( be picked and the associated configuration chosen to invoke the kernel. """ # Supported Triton versions, should be sorted from the newest to the oldest - supported_triton_versions = ["3.4.0", "3.3.1", "3.2.0", "3.1.0"] + supported_triton_versions = ["3.6.0", "3.4.0", "3.3.1", "3.2.0", "3.1.0"] # First look up if an optimized configuration is available in the configs # directory diff --git a/tests/test_magi2_s5000_moe_configs.py b/tests/test_magi2_s5000_moe_configs.py new file mode 100644 index 0000000..45a5d9e --- /dev/null +++ b/tests/test_magi2_s5000_moe_configs.py @@ -0,0 +1,32 @@ +import json +from pathlib import Path + +CONFIG_DIR = ( + Path(__file__).parents[1] / "src/torchada/triton/autotune/fused_moe/configs/triton_3_2_0" +) +NAMES = ( + "E=768,N=1280,device_name=MTT_S5000.json", + "E=768,N=1280,device_name=MTT_S5000_down.json", +) + + +def test_magi2_s5000_configs_have_tuned_runtime_shapes(): + for name in NAMES: + data = json.loads((CONFIG_DIR / name).read_text()) + assert set(data) == {"21996", "24468", "45012"} + for shape in ("21996", "24468"): + assert data[shape] == _expected_config(block_size_k=32) + + expected_k = 64 if name.endswith("_down.json") else 32 + assert data["45012"] == _expected_config(block_size_k=expected_k) + + +def _expected_config(*, block_size_k: int): + return { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": block_size_k, + "GROUP_SIZE_M": 16, + "num_warps": 16, + "num_stages": 1, + }