From 355b7734702020f3f9f169d025d64f2485ae511e Mon Sep 17 00:00:00 2001 From: agent Date: Sat, 19 Sep 2026 17:28:34 +0800 Subject: [PATCH 1/2] MUSA: one pipeline stage for the tiny-M MoE default (S5000: -17.8% TPOT, byte-identical) --- tests/test_moe_default_config.py | 37 +++++++++++++++++++ ..._executor.layers.fused_moe.fused_moe.patch | 35 ++++++++++++++++++ 2 files changed, 72 insertions(+) create mode 100644 tests/test_moe_default_config.py create mode 100644 vllm_musa/patches/series/0166-MUSA-vllm.model_executor.layers.fused_moe.fused_moe.patch diff --git a/tests/test_moe_default_config.py b/tests/test_moe_default_config.py new file mode 100644 index 000000000000..7093e3bf4d7e --- /dev/null +++ b/tests/test_moe_default_config.py @@ -0,0 +1,37 @@ +"""MUSA-100038: the tiny-M MoE default must not collapse occupancy on S5000. + +At the tiny-M decode shape (7 tokens x topk 6 = 42 useful rows over 13-41 of 128 +experts, one tile per expert) the upstream default pipelines the K loop 4 stages +deep at BLOCK_K=128, which reserves ~4x shared memory per block and costs resident +blocks. Measured on S5000: -17.80% TPOT at 1 stage, byte-identical output. +""" + +from vllm.model_executor.layers.fused_moe.fused_moe import get_default_config +from vllm.platforms import current_platform + + +def _cfg(M: int = 7) -> dict[str, int]: + return get_default_config(M=M, E=128, N=1856, K=2688, topk=6, dtype="bfloat16") + + +def test_tiny_m_num_stages_follows_platform() -> None: + assert _cfg()["num_stages"] == (1 if current_platform.is_musa() else 4) + + +def test_tiny_m_other_axes_unchanged() -> None: + config = _cfg() + assert config["BLOCK_SIZE_M"] == 16 + assert config["BLOCK_SIZE_N"] == 64 + assert config["BLOCK_SIZE_K"] == 128 + assert config["GROUP_SIZE_M"] == 1 + assert config["SPLIT_K"] == 1 + assert config["num_warps"] == 4 + + +def test_all_buckets_up_to_32_take_one_stage_on_musa() -> None: + for M in (1, 4, 7, 8, 16, 32): + assert _cfg(M)["num_stages"] == (1 if current_platform.is_musa() else 4) + + +def test_larger_m_keeps_upstream_staging() -> None: + assert _cfg(512)["num_stages"] == 3 diff --git a/vllm_musa/patches/series/0166-MUSA-vllm.model_executor.layers.fused_moe.fused_moe.patch b/vllm_musa/patches/series/0166-MUSA-vllm.model_executor.layers.fused_moe.fused_moe.patch new file mode 100644 index 000000000000..28cc8de94cc7 --- /dev/null +++ b/vllm_musa/patches/series/0166-MUSA-vllm.model_executor.layers.fused_moe.fused_moe.patch @@ -0,0 +1,35 @@ +From 0000000000000000000000000000000000000000 Mon Sep 17 00:00:00 2001 +From: musa +Date: Sat, 20 Sep 2026 08:30:00 +0800 +Subject: [PATCH] MUSA: vllm.model_executor.layers.fused_moe.fused_moe + +--- + vllm/model_executor/layers/fused_moe/fused_moe.py | 15 ++++++++++++++ + 1 file changed, 14 insertions(+), 1 deletions(-) + +diff --git a/vllm/model_executor/layers/fused_moe/fused_moe.py b/vllm/model_executor/layers/fused_moe/fused_moe.py +index be4930052a..03d1d9b8a1 100644 +--- a/vllm/model_executor/layers/fused_moe/fused_moe.py ++++ b/vllm/model_executor/layers/fused_moe/fused_moe.py +@@ -1397,7 +1397,20 @@ def get_default_config( + if current_platform.is_rocm(): + num_stages = num_stages_rocm + elif M <= 32: +- num_stages = 4 ++ # MUSA: 4 pipelining stages collapse occupancy for tiny-M MoE on S5000. ++ # At the Nemotron-3.5 MTP6 decode shape (7 tokens x topk 6 = 42 useful ++ # rows over 13-41 of 128 experts, one 16-row tile each) the tiny-M default ++ # is 4 stages / BLOCK_K=128, which reserves ~4x the shared memory per ++ # block and costs most of the resident blocks. Measured on S5000, same ++ # container and image, 4096-in/128-out, one request: 11.635 ms/token at ++ # 4 stages vs 9.563 ms/token at 1 stage (-17.80%), byte-identical output ++ # (same SHA over all 9 measured requests) and unchanged TTFT (1648 vs ++ # 1678 ms). BLOCK_SIZE_N=128 alone regresses (+8.46%), so staging is the ++ # whole effect. CUDA keeps the upstream 4 stages. ++ num_stages = 1 if getattr(current_platform, "is_musa", lambda: False)() else 4 ++ logger.info_once( ++ "MoE default config for tiny M: M=%d num_stages=%d", M, num_stages ++ ) + else: + num_stages = 3 + From 1527bc0b6f017d0959736e3d99fa9eb6bd9c3149 Mon Sep 17 00:00:00 2001 From: Xiaodong Ye Date: Sun, 20 Sep 2026 21:57:50 +0800 Subject: [PATCH 2/2] chore(patches): renumber the tiny-M stage patch 0166 -> 0168 takes 0167, so this one moves to 0168. The series README comes along: the file is taken from the current base (which has since added 0164/0165) and only the count line moves, 165 -> 168 - 167 once 0167 is in, 168 with this patch. --- ...8-MUSA-vllm.model_executor.layers.fused_moe.fused_moe.patch} | 0 vllm_musa/patches/series/README.md | 2 +- 2 files changed, 1 insertion(+), 1 deletion(-) rename vllm_musa/patches/series/{0166-MUSA-vllm.model_executor.layers.fused_moe.fused_moe.patch => 0168-MUSA-vllm.model_executor.layers.fused_moe.fused_moe.patch} (100%) diff --git a/vllm_musa/patches/series/0166-MUSA-vllm.model_executor.layers.fused_moe.fused_moe.patch b/vllm_musa/patches/series/0168-MUSA-vllm.model_executor.layers.fused_moe.fused_moe.patch similarity index 100% rename from vllm_musa/patches/series/0166-MUSA-vllm.model_executor.layers.fused_moe.fused_moe.patch rename to vllm_musa/patches/series/0168-MUSA-vllm.model_executor.layers.fused_moe.fused_moe.patch diff --git a/vllm_musa/patches/series/README.md b/vllm_musa/patches/series/README.md index ac3a1a7306db..a2ad4e16cfbe 100644 --- a/vllm_musa/patches/series/README.md +++ b/vllm_musa/patches/series/README.md @@ -17,7 +17,7 @@ is pre-patched. Author headers are normalized to the synthetic `musa ` identity. -Currently **167 patches**. This branch includes the Qwen3.6 patches for common +Currently **168 patches**. This branch includes the Qwen3.6 patches for common GDN decode metadata reuse, uniform-decode SSM slot-mapping removal, and the BF16 W1 tile specialization, plus the contract-bound DeepSeek-V4 MTP sparse-prefill headroom and mixed-prefill queue-fence patches. It additionally