[TRTLLM-12635][feat] add bench_moe microbenchmark

tests/microbenchmarks/bench_moe/__init__.py

@@ -0,0 +1,67 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""MoE microbenchmark package entrypoint.
+
+The implementation lives in focused submodules. This package mirrors the
+historical ``bench_moe.*`` import surface by re-exporting every non-dunder
+attribute, including underscore-prefixed helpers used by tests and scripts.
+"""
+
+from __future__ import annotations
+
+import importlib
+import sys
+from pathlib import Path
+
+_MICROBENCH_DIR = Path(__file__).resolve().parent.parent
+_TESTS_UNITTEST_DIR = _MICROBENCH_DIR.parent / "unittest"
+if str(_TESTS_UNITTEST_DIR) not in sys.path:
+    sys.path.insert(0, str(_TESTS_UNITTEST_DIR))
+
+_REPO_ROOT = _MICROBENCH_DIR.parent.parent
+if (_REPO_ROOT / "tensorrt_llm" / "bindings").is_dir() and str(_REPO_ROOT) not in sys.path:
+    sys.path.insert(0, str(_REPO_ROOT))
+
+_PACKAGE_PARENT = Path(__file__).resolve().parent.parent
+if str(_PACKAGE_PARENT) not in sys.path:
+    sys.path.insert(0, str(_PACKAGE_PARENT))
+
+
+def _mirror_module(module_name: str) -> None:
+    module = importlib.import_module(module_name)
+    for name, value in module.__dict__.items():
+        if not (name.startswith("__") and name.endswith("__")):
+            globals()[name] = value
+
+
+for _module_name in (
+    "bench_moe.backend",
+    "bench_moe.quantize",
+    "bench_moe.specs",
+    "bench_moe.search",
+    "bench_moe.utils",
+    "bench_moe.mapping",
+    "bench_moe.build",
+    "bench_moe.routing",
+    "bench_moe.timing",
+    "bench_moe.results",
+    "bench_moe.cli",
+    "bench_moe.case_runner",
+    "bench_moe.worker",
+):
+    _mirror_module(_module_name)
+
+del _module_name

tests/microbenchmarks/bench_moe/__main__.py

@@ -0,0 +1,65 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Script entry point for the MoE microbenchmark package.
+
+Two invocation styles are supported and tested:
+
+    # As a module (preferred, requires the parent of bench_moe/ on sys.path):
+    python -m bench_moe --help
+
+    # As a file path (matches CI and legacy script-style invocation):
+    python tests/microbenchmarks/bench_moe/__main__.py --help
+
+When invoked by file path, Python sets ``sys.path[0]`` to the directory
+containing the script (the ``bench_moe`` package itself), which would shadow
+the package's relative imports. We detect that case and prepend the package's
+parent directory so ``import bench_moe`` resolves to the package rather than
+this single file.
+"""
+
+from __future__ import annotations
+
+import sys
+from pathlib import Path
+
+
+def _ensure_package_importable() -> None:
+    """Make ``import bench_moe`` resolve to the parent package.
+
+    Only needed when this file is executed directly (``python __main__.py``),
+    not when run via ``python -m bench_moe`` (which sets ``__package__`` for
+    us).
+    """
+    if __package__:
+        return
+
+    pkg_dir = Path(__file__).resolve().parent
+    parent_dir = pkg_dir.parent
+
+    if str(parent_dir) not in sys.path:
+        sys.path.insert(0, str(parent_dir))
+
+    pkg_dir_str = str(pkg_dir)
+    while pkg_dir_str in sys.path:
+        sys.path.remove(pkg_dir_str)
+
+
+_ensure_package_importable()
+
+from bench_moe.worker import main  # noqa: E402
+
+if __name__ == "__main__":
+    main()

tests/microbenchmarks/bench_moe/backend.py

@@ -0,0 +1,125 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Local MoE backend registry for ``bench_moe``.
+
+The benchmark should not import unittest-wide MoE helpers at module import time:
+those helpers import every backend, including CUTEDSL, even when a run only asks
+for TRTLLM/CUTLASS/MegaMoE. Keep the small registry needed by the benchmark here
+and import concrete backend classes lazily.
+"""
+
+from __future__ import annotations
+
+import importlib
+import sys
+import types
+from dataclasses import dataclass
+from enum import Enum
+from typing import Optional
+
+
+class MoeBackendType(str, Enum):
+    """MoE backend identifiers accepted by bench_moe."""
+
+    CUTLASS = "CUTLASS"
+    TRTLLM = "TRTLLM"
+    CUTEDSL = "CUTEDSL"
+    DEEPGEMM = "DEEPGEMM"
+    DENSEGEMM = "DENSEGEMM"
+    MEGAMOE = "MEGAMOE_DEEPGEMM"
+
+
+@dataclass
+class MoeModelConfig:
+    """MoE model shape used by routing-method construction."""
+
+    num_experts: int
+    top_k: int
+    hidden_size: int
+    intermediate_size: int
+    n_group: Optional[int] = None
+    topk_group: Optional[int] = None
+
+    def __str__(self) -> str:
+        return f"e{self.num_experts}_k{self.top_k}_h{self.hidden_size}_i{self.intermediate_size}"
+
+
+def resolve_deepseek_group_config(model_config: MoeModelConfig) -> tuple[int, int]:
+    """Resolve DeepSeek-V3 routing group settings for built-in or custom shapes."""
+    if model_config.n_group is not None and model_config.topk_group is not None:
+        return model_config.n_group, model_config.topk_group
+    n_group = max(1, model_config.num_experts // 2)
+    topk_group = min(n_group, max(1, n_group // 2))
+    return n_group, topk_group
+
+
+def ensure_cute_dsl_importable_for_benchmark() -> None:
+    """Install a local sentinel module when optional CUTEDSL imports are absent.
+
+    Production MoE modules import ``fused_moe_cute_dsl`` at module load time.
+    ``bench_moe`` can still benchmark non-CUTEDSL backends in environments that
+    do not package CUTLASS DSL; keep that fallback local to the benchmark rather
+    than weakening the production module.
+    """
+    module_name = "tensorrt_llm._torch.modules.fused_moe.fused_moe_cute_dsl"
+    if module_name in sys.modules:
+        return
+    try:
+        importlib.import_module(module_name)
+        return
+    except ImportError as exc:
+        import_error = exc
+
+    class CuteDslFusedMoE:
+        @classmethod
+        def can_implement(cls, *_args, **_kwargs):
+            return False, f"CUTLASS DSL is unavailable: {import_error}"
+
+        def __init__(self, *_args, **_kwargs):
+            raise RuntimeError(f"CUTLASS DSL is unavailable: {import_error}")
+
+    module = types.ModuleType(module_name)
+    module.CuteDslFusedMoE = CuteDslFusedMoE
+    sys.modules[module_name] = module
+
+
+def get_backend_class(backend_type: MoeBackendType):
+    """Import and return the concrete backend class for ``backend_type`` lazily."""
+    if backend_type == MoeBackendType.CUTLASS:
+        from tensorrt_llm._torch.modules.fused_moe.fused_moe_cutlass import CutlassFusedMoE
+
+        return CutlassFusedMoE
+    if backend_type == MoeBackendType.TRTLLM:
+        from tensorrt_llm._torch.modules.fused_moe.fused_moe_trtllm_gen import TRTLLMGenFusedMoE
+
+        return TRTLLMGenFusedMoE
+    if backend_type == MoeBackendType.CUTEDSL:
+        from tensorrt_llm._torch.modules.fused_moe.fused_moe_cute_dsl import CuteDslFusedMoE
+
+        return CuteDslFusedMoE
+    if backend_type == MoeBackendType.DEEPGEMM:
+        from tensorrt_llm._torch.modules.fused_moe.fused_moe_deepgemm import DeepGemmFusedMoE
+
+        return DeepGemmFusedMoE
+    if backend_type == MoeBackendType.DENSEGEMM:
+        from tensorrt_llm._torch.modules.fused_moe.fused_moe_densegemm import DenseGEMMFusedMoE
+
+        return DenseGEMMFusedMoE
+    if backend_type == MoeBackendType.MEGAMOE:
+        from tensorrt_llm._torch.modules.fused_moe.mega_moe import MegaMoEDeepGemm
+
+        return MegaMoEDeepGemm
+    raise ValueError(f"unknown MoE backend {backend_type!r}")