From f73f2031aa65c974617e4fea99fce07020704625 Mon Sep 17 00:00:00 2001
From: kikig <29116997+guqiqi@users.noreply.github.com>
Date: Mon, 13 Jul 2026 19:40:04 -0700
Subject: [PATCH] [None][feat] Add kimi_k2/glm_5 grouped routing and fused
 router to bench_moe

- specs.py: set kimi_k2 n_group=1/topk_group=1 (real Kimi-K2 uses a single
  expert group; avoids the TRTLLM routing kernel topk_group<=4 failure); add
  glm_5 model spec (256 experts, top-8, DeepSeek-V3-style single-group routing).
- mapping.py: use fused DeepSeek-V3 routing (is_fused=True).
- search.py: prune MEGAMOE_DEEPGEMM + TEP as skipped instead of a hard build
  failure (kernel assumes partitioned MoE input; TEP replicates TP-wide).

Signed-off-by: kikig <29116997+guqiqi@users.noreply.github.com>
---
 tests/microbenchmarks/bench_moe/mapping.py |  2 +-
 tests/microbenchmarks/bench_moe/search.py  | 13 +++++++++++++
 tests/microbenchmarks/bench_moe/specs.py   | 17 +++++++++++++++++
 3 files changed, 31 insertions(+), 1 deletion(-)

diff --git a/tests/microbenchmarks/bench_moe/mapping.py b/tests/microbenchmarks/bench_moe/mapping.py
index 545b9109edcd..456cc1003c79 100644
--- a/tests/microbenchmarks/bench_moe/mapping.py
+++ b/tests/microbenchmarks/bench_moe/mapping.py
@@ -129,7 +129,7 @@ def _create_routing_method(
             topk_group=topk_group,
             routed_scaling_factor=1.0,
             callable_e_score_correction_bias=lambda: e_score_correction_bias,
-            is_fused=False,
+            is_fused=True,
         )
 
     if routing_method_cls is MiniMaxM2MoeRoutingMethod:
diff --git a/tests/microbenchmarks/bench_moe/search.py b/tests/microbenchmarks/bench_moe/search.py
index f692911e7ff2..4c6e45c23226 100644
--- a/tests/microbenchmarks/bench_moe/search.py
+++ b/tests/microbenchmarks/bench_moe/search.py
@@ -177,6 +177,19 @@ def is_candidate_valid(
             "use DEP/TEP modes only"
         )
 
+    # MegaMoEDeepGemm does not support TEP: its DeepGEMM fp8_fp4_mega_moe kernel
+    # assumes the MoE input is partitioned across ranks (single rank, or DEP with
+    # ep_size == parallel_size). TEP replicates the input TP-wide, so
+    # MegaMoEDeepGemm.__init__ raises NotImplementedError (enable_attention_dp=False,
+    # parallel_size>1). Prune here so the sweep records status="skipped" instead of a
+    # hard build failure.
+    if config.backend.upper() == "MEGAMOE_DEEPGEMM" and not enable_dp and world_size > 1:
+        return False, (
+            f"MEGAMOE_DEEPGEMM does not support TEP (enable_attention_dp=False, "
+            f"parallel_size={world_size}>1); use DEP with ep_size==parallel_size "
+            "or enable attention-DP"
+        )
+
     # DENSEGEMM DTP: FC2 kernel requires (intermediate_size / moe_tp_size) % 256 == 0.
     # DENSEGEMM __init__ only checks the full intermediate_size, so a model like
     # DeepSeek V3 (intermediate_size=2048, 2048%256=0) passes __init__ but fails
diff --git a/tests/microbenchmarks/bench_moe/specs.py b/tests/microbenchmarks/bench_moe/specs.py
index 36134c9b429a..dff98ad5d06d 100644
--- a/tests/microbenchmarks/bench_moe/specs.py
+++ b/tests/microbenchmarks/bench_moe/specs.py
@@ -308,6 +308,23 @@ class RunResult:
         intermediate_size=2048,
         quant_algo="FP8_BLOCK_SCALES",
         routing_method="DEEPSEEK_V3",
+        n_group=1,
+        topk_group=1,
+    ),
+    # GLM-5 (zai-org/GLM-5): 256 routed experts, top-8, sigmoid/noaux_tc
+    # (DeepSeek-V3-style) routing with a single expert group (n_group=1,
+    # topk_group=1). intermediate_size is the per-expert moe_intermediate_size.
+    # quant_algo left None: pass --quant on the CLI (the glm_5 sweep uses NVFP4).
+    "glm_5": ModelSpec(
+        name="glm_5",
+        num_experts=256,
+        top_k=8,
+        hidden_size=6144,
+        intermediate_size=2048,
+        quant_algo=None,
+        routing_method="DEEPSEEK_V3",
+        n_group=1,
+        topk_group=1,
     ),
     # DeepSeek-V4-Pro: 1.6T total / 49B activated. quant_algo intentionally
     # left None: pass --quant on the CLI to pin the mode (the released