refactor(moe): enforce build() for all runners and refine CUTLASS autotuning in unified MoE API

flashinfer/autotuner/autotuner.py

@@ -1074,6 +1074,10 @@ def __init__(
         self.profiling_cache: dict[
             ProfilingCacheKey, tuple[Any, OptimizationProfile | None]
         ] = {}
+        # Ranked shortlists are process-local. Persisted configs retain the
+        # selected winner; a later tuning session rebuilds the shortlist when
+        # compound refinement needs more than one candidate.
+        self._ranked_tactics_cache: dict[ProfilingCacheKey, tuple[Any, ...]] = {}
         self.is_tuning_mode = False
         self._active_tuning_contexts = 0
 
@@ -1721,6 +1725,130 @@ def choose_one(
 
             return runners[runner_id], tactic
 
+    def rank_tactics(
+        self,
+        custom_op: str,
+        runners: list[TunableRunner],
+        tuning_config: TuningConfig,
+        inputs: list[torch.Tensor],
+        k: int = 1,
+        **kwargs,
+    ) -> list[Any]:
+        """Return up to ``k`` tactics for ``runners[0]``, ordered best-first.
+
+        Outside tuning mode (or when ``k == 1``), this matches ``choose_one``
+        and returns a single cached / fallback tactic. During tuning with
+        ``k > 1``, every valid tactic is profiled once; the winner is cached
+        under the same key as ``choose_one``, and the top ``k`` by measured
+        time are returned for callers that need a shortlist (e.g. multi-stage
+        compound-tactic refinement).
+        """
+        if k < 1:
+            raise ValueError(f"k must be >= 1, got {k}")
+        if len(runners) != 1:
+            raise ValueError(
+                f"rank_tactics requires exactly one runner, got {len(runners)} "
+                f"for op '{custom_op}'"
+            )
+
+        if k == 1 or not self.is_tuning_mode:
+            _, tactic = self.choose_one(
+                custom_op, runners, tuning_config, inputs, **kwargs
+            )
+            return [-1 if tactic is None else tactic]
+
+        if custom_op in self._effective_skip_ops:
+            logger.debug(
+                f"[AutoTuner]: Skipping ranking for '{custom_op}' "
+                f"(in skip_ops). Using fallback tactic."
+            )
+            return [-1]
+
+        with self._lock:
+            if self._override_tuning_buckets is not None or self._override_round_up:
+                tuning_config = self._apply_tuning_overrides(tuning_config)
+
+            input_shapes = tuple(self._get_input_sizes(inputs))
+            profiles = self._generate_optimization_profiles(tuning_config, inputs)
+            runner = runners[0]
+            runner_arg_names = {
+                param.name
+                for param in inspect.signature(runner.forward).parameters.values()
+            }
+
+            nearest_profile = self._find_nearest_profile(input_shapes, tuning_config)
+            try:
+                profile = next(
+                    candidate
+                    for candidate in profiles
+                    if self._find_nearest_profile(
+                        candidate.get_opt_shapes(), tuning_config
+                    )
+                    == nearest_profile
+                )
+            except StopIteration as e:
+                raise RuntimeError(
+                    f"No optimization profile for mapped shapes {nearest_profile} "
+                    f"while ranking '{custom_op}'"
+                ) from e
+
+            cache_key = AutoTuner._get_cache_key(
+                custom_op,
+                runner,
+                profile.get_opt_shapes(),
+                tuning_config,
+                runner.get_cache_key_extras(inputs),
+            )
+            cached_ranking = self._ranked_tactics_cache.get(cache_key)
+            if cached_ranking is not None:
+                return list(cached_ranking[:k])
+
+            tensors = self._prepare_input_tensors(profile, inputs)
+            if tuning_config.inputs_pre_hook is not None:
+                tensors = list(tuning_config.inputs_pre_hook(tensors))
+
+            valid_tactics = runner.get_valid_tactics(tensors, profile)
+            valid_tactics = self._blocklist.filter(custom_op, runner, valid_tactics)
+            if not valid_tactics:
+                return [-1]
+
+            if "do_preparation" in runner_arg_names:
+                runner(tensors, tactic=-1, do_preparation=True, **kwargs)
+
+            scored: list[tuple[float, Any]] = []
+            for tac in valid_tactics:
+                try:
+                    time_measured = self._profile_single_kernel(
+                        runner, tensors, tac, tuning_config, **kwargs
+                    )
+                except Exception as e:
+                    logger.debug(
+                        f"[Autotuner]: Skipping tactic {runner} {tac} while "
+                        f"ranking {custom_op}: {e}"
+                    )
+                    with contextlib.suppress(Exception):
+                        torch.cuda.synchronize()
+                    with contextlib.suppress(Exception):
+                        torch.cuda.cudart().cudaGetLastError()
+                    time_measured = float("inf")
+                scored.append((time_measured, tac))
+
+            scored.sort(key=lambda item: item[0])
+            ranked = [
+                tac for time_measured, tac in scored if time_measured < float("inf")
+            ]
+            if not ranked:
+                return [-1]
+
+            # Populate the choose_one cache with the winner so stage lookups
+            # remain consistent between rank_tactics and choose_one.
+            self.profiling_cache[cache_key] = (ranked[0], profile)
+            self._ranked_tactics_cache[cache_key] = tuple(ranked)
+            self._dirty = True
+            self._dirty_seq += 1
+
+            return ranked[:k]
+
     def _get_input_sizes(self, inputs: list[Any]) -> tuple[tuple[int, ...], ...]:
         """Return ``torch.Size`` for each input, using ``(0,)`` for non-Tensor values."""
         return tuple(
@@ -2434,6 +2562,7 @@ def clear_cache(self) -> None:
         """Clear the profiling cache and user-loaded file configs."""
         with self._lock:
             self.profiling_cache.clear()
+            self._ranked_tactics_cache.clear()
             self._file_configs.clear()
             self._logged_file_hits.clear()
             self._logged_cache_miss_oor.clear()