perf(moe): one pipeline stage for the tiny-M MoE default on MUSA (S5000 4k/128: -18% TPOT on Nemotron 3.5 MTP6)

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

vllm_musa/patches/series/0168-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 <musa@local>
+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
+

vllm_musa/patches/series/README.md

@@ -17,7 +17,7 @@ is pre-patched.
   Author headers are normalized to the synthetic
   `musa <musa@local>` 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