From ab65772dea64fcbe082a263330c4021ae01a03d3 Mon Sep 17 00:00:00 2001 From: Xiaodong Ye Date: Fri, 4 Sep 2026 21:18:20 +0800 Subject: [PATCH 1/3] perf: add MAGI2 S5000 fused MoE configs --- .../E=768,N=1280,device_name=MTT_S5000.json | 18 +++++++++++++ ...768,N=1280,device_name=MTT_S5000_down.json | 18 +++++++++++++ tests/test_magi2_s5000_moe_configs.py | 25 +++++++++++++++++++ 3 files changed, 61 insertions(+) create mode 100644 src/torchada/triton/autotune/fused_moe/configs/triton_3_2_0/E=768,N=1280,device_name=MTT_S5000.json create mode 100644 src/torchada/triton/autotune/fused_moe/configs/triton_3_2_0/E=768,N=1280,device_name=MTT_S5000_down.json create mode 100644 tests/test_magi2_s5000_moe_configs.py 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..cdbda8a --- /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,18 @@ +{ + "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 + } +} 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..cdbda8a --- /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,18 @@ +{ + "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 + } +} diff --git a/tests/test_magi2_s5000_moe_configs.py b/tests/test_magi2_s5000_moe_configs.py new file mode 100644 index 0000000..634dc6a --- /dev/null +++ b/tests/test_magi2_s5000_moe_configs.py @@ -0,0 +1,25 @@ +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"} + for config in data.values(): + assert config == { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 16, + "num_warps": 16, + "num_stages": 1, + } From 245817299a8f43dad6717ee60679dda9fc321e61 Mon Sep 17 00:00:00 2001 From: Xiaodong Ye Date: Sat, 5 Sep 2026 08:34:54 +0800 Subject: [PATCH 2/3] perf: tune MAGI2 CFG2 fused MoE shape --- .../E=768,N=1280,device_name=MTT_S5000.json | 8 ++++++ ...768,N=1280,device_name=MTT_S5000_down.json | 8 ++++++ tests/test_magi2_s5000_moe_configs.py | 27 ++++++++++++------- 3 files changed, 33 insertions(+), 10 deletions(-) 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 index cdbda8a..32c7856 100644 --- 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 @@ -14,5 +14,13 @@ "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 index cdbda8a..99f9bbc 100644 --- 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 @@ -14,5 +14,13 @@ "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/tests/test_magi2_s5000_moe_configs.py b/tests/test_magi2_s5000_moe_configs.py index 634dc6a..45a5d9e 100644 --- a/tests/test_magi2_s5000_moe_configs.py +++ b/tests/test_magi2_s5000_moe_configs.py @@ -13,13 +13,20 @@ 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"} - for config in data.values(): - assert config == { - "BLOCK_SIZE_M": 128, - "BLOCK_SIZE_N": 128, - "BLOCK_SIZE_K": 32, - "GROUP_SIZE_M": 16, - "num_warps": 16, - "num_stages": 1, - } + 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, + } From a53691c2097be6f9302a6cf7da0a27f63d81f092 Mon Sep 17 00:00:00 2001 From: Xiaodong Ye Date: Sat, 12 Sep 2026 08:56:48 +0800 Subject: [PATCH 3/3] perf(triton): add S5000 BF16 MoE config for Triton 3.6 Signed-off-by: Xiaodong Ye --- .../E=768,N=1280,device_name=MTT_S5000,dtype=bf16.json | 10 ++++++++++ src/torchada/triton/runtime/fused_moe/config.py | 2 +- 2 files changed, 11 insertions(+), 1 deletion(-) create mode 100644 src/torchada/triton/autotune/fused_moe/configs/triton_3_6_0/E=768,N=1280,device_name=MTT_S5000,dtype=bf16.json 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