[None][perf] update MegaMoE kernels for Blackwell and Rubin

.github/CODEOWNERS

@@ -350,7 +350,7 @@
 /cpp/tensorrt_llm/nanobind/runtime/moe* @NVIDIA/trt-llm-moe-devs
 /cpp/tensorrt_llm/runtime/moeLoadBalancer @NVIDIA/trt-llm-moe-devs
 /tensorrt_llm/_torch/cute_dsl_kernels/blackwell/moe_as_dense_gemm @NVIDIA/trt-llm-moe-devs
-/tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4 @NVIDIA/trt-llm-moe-devs
+/tensorrt_llm/_torch/cute_dsl_kernels/cutedsl_megamoe @NVIDIA/trt-llm-moe-devs
 /tensorrt_llm/_torch/moe @NVIDIA/trt-llm-moe-devs
 /tensorrt_llm/deep_ep @NVIDIA/trt-llm-moe-devs
 /tests/microbenchmarks/bench_moe @NVIDIA/trt-llm-moe-devs

pyproject.toml

@@ -19,7 +19,7 @@ column_limit = 80
 
 [tool.codespell]
 skip = ".git,3rdparty,triton_kernels,tests/integration/test_input_files**,**.jsonl,**.json"
-ignore-words-list = "rouge,inout,atleast,strat,nd,subtile,thrid,improbe,NotIn,te,iteract,anythin,tru,Tracin,vEw,dOut,indext,asend,medias,thw"
+ignore-words-list = "rouge,inout,atleast,strat,nd,subtile,thrid,improbe,NotIn,te,iteract,anythin,tru,Tracin,vEw,dOut,indext,asend,medias,thw,inh"
 
 [tool.autoflake]
 in-place = true

tensorrt_llm/_torch/autotuner.py

@@ -1,3 +1,17 @@
+# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
+#
+# 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.
+
 import ast
 import contextlib
 import copy
@@ -924,6 +938,7 @@ def __init__(self, warmup=2, repeat=10, stream_delay_micro_secs=1000):
         self.warmup = warmup
         self.stream_delay_micro_secs = stream_delay_micro_secs
         self.profiling_cache = AutoTunerProfilingCache()
+        self._primed_cached_tactics: set[tuple[str, str, str, str]] = set()
         self.is_tuning_mode = False
         self.skip_dynamic_tuning_buckets = False
 
@@ -1168,6 +1183,8 @@ def choose_one(
 
         # If it's tuning mode and cache hit, return the best runner and tactic to avoid redundant profiling.
         if self.is_tuning_mode and is_cache_hit:
+            self._prime_cached_tactics(custom_op, runners, tuning_config,
+                                       inputs, **kwargs)
             return (runners[best_runner_id], best_tactic)
 
         # PP rank does not have cache hit, so we try to receive the cache from the previous rank
@@ -1214,6 +1231,8 @@ def choose_one(
 
         self._maybe_sync_cache_data(tuning_config.distributed_tuning_strategy,
                                     custom_op)
+        self._prime_cached_tactics(custom_op, runners, tuning_config, inputs,
+                                   **kwargs)
 
         # If failed profiling tactics occurs, log the error.
         if new_tuning_failure_occurred:
@@ -1237,6 +1256,68 @@ def choose_one(
                 custom_op, 0) + tuning_end_time - tuning_start_time
         return (runners[runner_id], tactic)
 
+    def _prime_cached_tactics(
+        self,
+        custom_op: str,
+        runners: List[TunableRunner],
+        tuning_config: TuningConfig,
+        inputs: List[torch.Tensor],
+        **kwargs,
+    ) -> None:
+        """Compile cached winners in this process during tuning, not inference."""
+        if os.environ.get("TLLM_AUTOTUNER_PRIME_CACHED_TACTICS", "1") != "1":
+            return
+        # Cache hits may differ across ranks; collective runners must stay
+        # inside lockstep profiling, including profiling scratch allocation.
+        if tuning_config.distributed_tuning_strategy == DistributedTuningStrategy.MERGE:
+            return
+        jit_classes = {
+            type(r)
+            for r in runners
+            if isinstance(getattr(type(r), "kernel_cache", None), dict)
+        }
+        if not jit_classes:
+            return
+        for profile in self._optimization_profiles(tuning_config, inputs):
+            hit, runner_id, tactic, _ = self.profiling_cache.search_cache(
+                custom_op,
+                runners,
+                profile.get_opt_shapes(),
+                tuning_config,
+                apply_map_to_tuning_buckets=False,
+            )
+            if not hit:
+                continue
+            runner = runners[runner_id]
+            if type(runner) not in jit_classes:
+                continue
+            key = (custom_op, type(runner).__name__, str(runner.unique_id()),
+                   str(tactic))
+            if key in self._primed_cached_tactics:
+                continue
+            tensors = self._prepare_input_tensors(profile, inputs)
+            if tuning_config.inputs_pre_hook is not None:
+                tensors = tuning_config.inputs_pre_hook(tensors)
+            try:
+                with nvtx_range(f"{custom_op} prime tactic {tactic}"):
+                    if "do_preparation" in inspect.signature(
+                            runner.forward).parameters:
+                        runner(tensors,
+                               tactic=-1,
+                               do_preparation=True,
+                               **kwargs)
+                    runner(tensors, tactic=tactic, **kwargs)
+                self._primed_cached_tactics.add(key)
+            except Exception as e:
+                try:
+                    torch.cuda.synchronize()
+                except RuntimeError:
+                    pass
+                logger.warning(
+                    f"[Autotuner] Priming cached tactic failed for custom_op={custom_op}, "
+                    f"runner={type(runner).__name__}, tactic={tactic}, "
+                    f"shapes={profile.get_opt_shapes()}. Error: {e}")
+
     def _profile_runners(
         self,
         custom_op: str,

tensorrt_llm/_torch/cute_dsl_kernels/cutedsl_megamoe/__init__.py

@@ -0,0 +1,24 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Exported MegaMoE KernelClass entry points."""
+
+from .kernel_src.blackwell.inference.mega.block_scaled_swap_ab_mega_moe_kernel import (
+    BlockScaledSwapAbMegaMoeKernel as BlackwellInferenceMegaMoE,
+)
+from .kernel_src.rubin.inference.local_mega.block_scaled_swap_ab_local_mega_moe_kernel import (
+    BlockScaledSwapAbLocalMegaMoeKernel as RubinInferenceLocalMegaMoE,
+)
+from .kernel_src.rubin.inference.mega.block_scaled_swap_ab_mega_moe_kernel import (
+    BlockScaledSwapAbMegaMoeKernel as RubinInferenceMegaMoE,
+)
+from .kernel_src.rubin.inference.mega.block_scaled_swap_ab_mega_moe_kernel_gen_specialized import (
+    BlockScaledSwapAbGenphaseMoeKernel as RubinInferenceGenphaseMegaMoE,
+)
+
+__all__ = [
+    "BlackwellInferenceMegaMoE",
+    "RubinInferenceGenphaseMegaMoE",
+    "RubinInferenceLocalMegaMoE",
+    "RubinInferenceMegaMoE",
+]

tensorrt_llm/_torch/cute_dsl_kernels/cutedsl_megamoe/api.py

@@ -0,0 +1,167 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Stable construction API for composable kernel implementations."""
+
+import types
+from abc import ABC, ABCMeta, abstractmethod
+from collections.abc import Iterator, Mapping
+from types import MappingProxyType
+from typing import Dict, Union, get_args, get_origin
+
+from cutlass.cute.typing import SymInt
+
+RuntimeIntegerType = SymInt
+StaticIntegerType = int
+StaticOrRuntimeIntegerType = Union[int, SymInt]
+
+
+class OptionalRequirement:
+    """Descriptor field that a component may consume conditionally."""
+
+    __slots__ = ("expected_type",)
+
+    def __init__(self, expected_type) -> None:
+        self.expected_type = expected_type
+
+
+Requirement = Union[type, OptionalRequirement]
+
+
+def _required_type(requirement: Requirement):
+    return (
+        requirement.expected_type if isinstance(requirement, OptionalRequirement) else requirement
+    )
+
+
+def _matches_type(value, expected_type) -> bool:
+    origin = get_origin(expected_type)
+    if origin in (Union, types.UnionType):
+        return any(_matches_type(value, candidate) for candidate in get_args(expected_type))
+    return isinstance(value, expected_type)
+
+
+class Desc(Mapping[str, object]):
+    """Immutable descriptor mapping validated against component schemas."""
+
+    def __init__(self, values: Mapping[str, object]) -> None:
+        self._values = MappingProxyType(dict(values))
+
+    def __getitem__(self, key: str):
+        return self._values[key]
+
+    def __iter__(self) -> Iterator[str]:
+        return iter(self._values)
+
+    def __len__(self) -> int:
+        return len(self._values)
+
+    def validate(
+        self,
+        requirements: Dict[str, Requirement],
+        *,
+        component_name: str,
+    ) -> None:
+        for name, requirement in requirements.items():
+            if name not in self._values:
+                if isinstance(requirement, OptionalRequirement):
+                    continue
+                raise KeyError(f"{component_name} requires descriptor field {name!r}.")
+            expected_type = _required_type(requirement)
+            value = self._values[name]
+            if not _matches_type(value, expected_type):
+                raise TypeError(
+                    f"{component_name} requires {name!r} to have type "
+                    f"{expected_type}, got {type(value)}."
+                )
+
+
+class ProblemDesc(Desc):
+    """Problem semantics shared by every component in one kernel."""
+
+
+class ImplDesc(Desc):
+    """Fully static implementation choices shared by kernel components."""
+
+
+class _KernelComponentMeta(ABCMeta):
+    def __call__(cls, *args, **kwargs):
+        instance = super().__call__(*args, **kwargs)
+        instance._validate_component_init()
+        return instance
+
+
+class KernelComponent(ABC, metaclass=_KernelComponentMeta):
+    """Component that consumes descriptors only during construction."""
+
+    @classmethod
+    @abstractmethod
+    def problem_desc_require(cls) -> Dict[str, Requirement]: ...
+
+    @classmethod
+    @abstractmethod
+    def impl_desc_require(cls) -> Dict[str, Requirement]: ...
+
+    def _validate_desc_inputs(
+        self,
+        problem_desc: ProblemDesc,
+        impl_desc: ImplDesc,
+    ) -> None:
+        component_name = type(self).__name__
+        overlap = self.problem_desc_require().keys() & self.impl_desc_require().keys()
+        if overlap:
+            raise ValueError(
+                f"{component_name} requires fields from both descriptors: {sorted(overlap)}."
+            )
+        problem_desc.validate(
+            self.problem_desc_require(),
+            component_name=component_name,
+        )
+        impl_desc.validate(
+            self.impl_desc_require(),
+            component_name=component_name,
+        )
+
+    def _validate_component_init(self) -> None:
+        component_name = type(self).__name__
+        requirements = {
+            **self.problem_desc_require(),
+            **self.impl_desc_require(),
+        }
+        for name, requirement in requirements.items():
+            if not hasattr(self, name):
+                if isinstance(requirement, OptionalRequirement):
+                    continue
+                raise RuntimeError(f"{component_name} did not bind required field {name!r}.")
+            expected_type = _required_type(requirement)
+            value = getattr(self, name)
+            if not _matches_type(value, expected_type):
+                raise TypeError(
+                    f"{component_name}.{name} must have type {expected_type}, got {type(value)}."
+                )
+        for name, value in vars(self).items():
+            if isinstance(value, Desc):
+                raise RuntimeError(f"{component_name}.{name} retains a descriptor.")
+
+
+class KernelClass(KernelComponent):
+    """Top-level host wrapper for one composable kernel implementation."""
+
+    @abstractmethod
+    def name(self) -> str: ...
+
+    @abstractmethod
+    def aot_compile(self): ...
+
+
+__all__ = [
+    "Desc",
+    "ImplDesc",
+    "KernelClass",
+    "KernelComponent",
+    "OptionalRequirement",
+    "ProblemDesc",
+    "RuntimeIntegerType",
+    "StaticIntegerType",
+    "StaticOrRuntimeIntegerType",
+]

tensorrt_llm/_torch/cute_dsl_kernels/cutedsl_megamoe/communication/__init__.py

@@ -0,0 +1,2 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0

tensorrt_llm/_torch/cute_dsl_kernels/cutedsl_megamoe/communication/nvlink_domain/__init__.py

@@ -0,0 +1,2 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0