diff --git a/flashinfer/autotuner/autotuner.py b/flashinfer/autotuner/autotuner.py index 2bbd75891f0..749d9be2075 100644 --- a/flashinfer/autotuner/autotuner.py +++ b/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() diff --git a/flashinfer/fused_moe/runners.py b/flashinfer/fused_moe/runners.py index 9e0ef645f0b..234dba33217 100644 --- a/flashinfer/fused_moe/runners.py +++ b/flashinfer/fused_moe/runners.py @@ -181,7 +181,12 @@ def _validate_logits_inputs( class MoERunner(TunableRunner): - """Base class for unified MoE backend runners.""" + """Unified MoE runner lifecycle: validate, build once, then execute. + + Concrete runners implement ``_check_support()`` and ``_build()``. Keeping + the public methods here ensures a failed support check cannot authorize a + build and execution cannot silently initialize backend resources. + """ backend_key: ClassVar[str] = "" supported_routing_modes: tuple[RoutingInputMode, ...] = () @@ -189,7 +194,16 @@ class MoERunner(TunableRunner): config: MoEConfig + def __init__(self) -> None: + self._support_checked = False + self._built = False + def check_support(self) -> None: + self._support_checked = False + self._check_support() + self._support_checked = True + + def _check_support(self) -> None: """Raise if the initialized runner cannot execute its configuration.""" variant = self.config.quant.variant if variant not in self.supported_quant_variants: @@ -198,12 +212,23 @@ def check_support(self) -> None: ) def build(self) -> None: - """Prepare backend resources after support has been validated. + if getattr(self, "_built", False): + return + if not getattr(self, "_support_checked", False): + raise RuntimeError( + f"{type(self).__name__}.check_support() must succeed before build()." + ) + self._build() + self._built = True - Existing runners that finish initialization in ``__init__`` inherit - this no-op. Backends can migrate expensive module loading here - incrementally. - """ + def _build(self) -> None: + """Prepare shape-independent resources for a supported runner.""" + + def _require_built(self) -> None: + if not getattr(self, "_built", False): + raise RuntimeError( + f"{type(self).__name__}.build() must be called before execution." + ) # --------------------------------------------------------------------------- @@ -228,9 +253,13 @@ class _CutlassRunnerBase(MoERunner): _use_w4_group_scaling: ClassVar[bool] _required_weight_keys: ClassVar[tuple[str, ...]] _expected_num_inputs: ClassVar[int] + # Keep top-k tactics per GEMM stage, then return their Cartesian product as + # compound candidates for the outer end-to-end autotuner. k=1 preserves the + # legacy independent-winner behavior. + _stage_tactic_top_k: ClassVar[int] = 2 - def check_support(self) -> None: - super().check_support() + def _check_support(self) -> None: + super()._check_support() if self.config.activation.type is not ActivationType.Swiglu: raise NotImplementedError( f"{type(self).__name__} supports only the Swiglu activation." @@ -255,6 +284,7 @@ def check_support(self) -> None: ) def __init__(self, config: MoEConfig, device: torch.device): + super().__init__() from ..utils import device_support_pdl, get_compute_capability self.config = config @@ -283,11 +313,8 @@ def __init__(self, config: MoEConfig, device: torch.device): self._workspace_num_tokens = 0 self._workspace_hidden_size: int | None = None - def build(self) -> None: - """Load the CUTLASS module and create the inner runner once.""" - if self._inner is not None: - return - + def _build(self) -> None: + """Load the CUTLASS module and create the inner runner.""" from .core import get_cutlass_fused_moe_module with torch.cuda.device(self.device): @@ -315,14 +342,6 @@ def build(self) -> None: use_wfp4afp8_humming=False, ) - def _require_built(self) -> None: - """Reject execution before the explicit runner lifecycle completes.""" - if self._inner is None: - raise RuntimeError( - f"{type(self).__name__} must be initialized with " - "check_support() followed by build()." - ) - def _prepare_tuning_inputs(self, inputs: List[torch.Tensor]) -> List[torch.Tensor]: """Populate synthesized routing inputs with a valid balanced pattern.""" num_tokens = inputs[1].shape[0] @@ -343,37 +362,41 @@ def get_valid_tactics(self, inputs: List[torch.Tensor], _profile: Any) -> List[A self._require_built() self._validate_input_count(inputs) # The two GEMMs have independent tactic spaces. Preserve the legacy - # O(n1+n2) tuning flow, then let the outer unified tuner profile only - # the selected pair as one full-MoE candidate. - # FIXME: get_valid_tactics() is supposed to list candidates, but the - # autotuner cannot express multi-stage / factorized tactics yet, so we - # nest choose_one() per GEMM and return one compound winner. Refine the - # autotuner to own staged tuning instead of selecting winners here. + # O(n1+n2) stage search, keep the top-k tactics per stage, then let the + # outer unified tuner profile only the k² compound pairs end-to-end. + # FIXME: Prefer a first-class factorized/multi-stage autotuner API so + # runners do not need to nest stage ranking inside get_valid_tactics(). tuner = AutoTuner.get() profile_inputs = [inputs[1], inputs[4], None, inputs[5], None] stage_tuning_config = TuningConfig() + top_k = self._stage_tactic_top_k try: self._inner.gemm_idx_for_tuning = 1 - _, gemm1 = tuner.choose_one( + gemm1_tactics = tuner.rank_tactics( f"moe_{self.backend_key}_sm{self._device_arch}_gemm1", [self._inner], stage_tuning_config, profile_inputs, + k=top_k, gemm_idx=1, ) self._inner.gemm_idx_for_tuning = 2 - _, gemm2 = tuner.choose_one( + gemm2_tactics = tuner.rank_tactics( f"moe_{self.backend_key}_sm{self._device_arch}_gemm2", [self._inner], stage_tuning_config, profile_inputs, + k=top_k, gemm_idx=2, ) finally: self._inner.gemm_idx_for_tuning = None - if gemm1 is None or gemm2 is None: - return [-1] - return [(int(gemm1), int(gemm2))] + pairs: List[Any] = [ + (int(gemm1), int(gemm2)) + for gemm1 in gemm1_tactics + for gemm2 in gemm2_tactics + ] + return pairs if pairs else [-1] def _ensure_workspace(self, num_tokens: int, hidden_size: int) -> None: max_num_tokens = self.config.execution.tune_max_num_tokens @@ -695,8 +718,8 @@ class CuteDslNvfp4Runner(MoERunner): supported_routing_modes = (RoutingInputMode.PackedPrecomputed,) supported_quant_variants = (QuantVariant.NVFP4, QuantVariant.W4A16) - def check_support(self) -> None: - super().check_support() + def _check_support(self) -> None: + super()._check_support() if self.config.activation.type is not ActivationType.Swiglu: raise NotImplementedError( f"{type(self).__name__} supports only the Swiglu activation." @@ -711,45 +734,53 @@ def check_support(self) -> None: ) def __init__(self, config: MoEConfig, device: torch.device): + super().__init__() + self.config = config + self.device = torch.device(device) + self._inner: Any = None + self.tuning_config = TuningConfig() + + def _build(self) -> None: + """Create the shape-independent CuTe DSL tuning runner.""" from .cute_dsl.fused_moe import _cute_dsl_fused_moe_nvfp4_impl from .cute_dsl.tuner import ( CuteDslFusedMoENvfp4Runner, CuteDslFusedMoEW4A16Runner, ) - self.config = config - experts = config.experts - routing = config.routing + experts = self.config.experts + routing = self.config.routing num_local_experts = experts.local_num_experts or routing.num_experts enable_pdl = ( - True if config.execution.enable_pdl is None else config.execution.enable_pdl + True + if self.config.execution.enable_pdl is None + else self.config.execution.enable_pdl ) - self._inner: CuteDslFusedMoENvfp4Runner | CuteDslFusedMoEW4A16Runner - if config.quant.variant is QuantVariant.NVFP4: + if self.config.quant.variant is QuantVariant.NVFP4: self._inner = CuteDslFusedMoENvfp4Runner( forward_impl=_cute_dsl_fused_moe_nvfp4_impl, num_experts=routing.num_experts, top_k=routing.top_k, num_local_experts=num_local_experts, local_expert_offset=experts.local_expert_offset, - use_fused_finalize=config.execution.use_fused_finalize, + use_fused_finalize=self.config.execution.use_fused_finalize, enable_pdl=enable_pdl, - activation_type=int(config.activation.type), - use_per_token_activation=bool(config.quant.per_token_scale), + activation_type=int(self.config.activation.type), + use_per_token_activation=bool(self.config.quant.per_token_scale), ) - elif config.quant.variant is QuantVariant.W4A16: + elif self.config.quant.variant is QuantVariant.W4A16: self._inner = CuteDslFusedMoEW4A16Runner( num_experts=routing.num_experts, top_k=routing.top_k, num_local_experts=num_local_experts, local_expert_offset=experts.local_expert_offset, - use_fused_finalize=config.execution.use_fused_finalize, + use_fused_finalize=self.config.execution.use_fused_finalize, enable_pdl=enable_pdl, - activation_type=int(config.activation.type), + activation_type=int(self.config.activation.type), ) else: raise NotImplementedError( - f"CuteDslNvfp4Runner does not support {config.quant.variant}." + f"CuteDslNvfp4Runner does not support {self.config.quant.variant}." ) # tuning_config is an instance attribute on the inner runner (its # dummy expert-id span depends on num_experts/offset), so read it from @@ -757,6 +788,7 @@ def __init__(self, config: MoEConfig, device: torch.device): self.tuning_config = self._inner.tuning_config def get_valid_tactics(self, inputs: List[torch.Tensor], profile: Any) -> List[Any]: + self._require_built() return self._inner.get_valid_tactics(inputs, profile) def get_cache_key_extras(self, inputs: List[torch.Tensor]) -> tuple: @@ -769,6 +801,7 @@ def forward( do_preparation: bool = False, **kwargs: Any, ) -> torch.Tensor: + self._require_built() return self._inner.forward( inputs, tactic=tactic, do_preparation=do_preparation, **kwargs ) @@ -785,6 +818,7 @@ def pack_inputs( tuning configurations include the output buffer so profiling can replace it for each token bucket. """ + self._require_built() # MoELayer already filters by supported_routing_modes; this guards the # direct-runner path (tests/benchmarks) against silently forwarding a # logits pack's None topk tensors into the kernel launch. @@ -879,15 +913,25 @@ def pack_inputs( ) def __hash__(self): + self._require_built() return hash(("cute_dsl_nvfp4", hash(self._inner))) # --------------------------------------------------------------------------- -# TRTLLM FP4 routed runner — delegates to the canonical trtllm-gen MoERunner +# TRTLLM runners — shared module lifecycle, shape-specific inner runners # --------------------------------------------------------------------------- -class TrtllmFp4RoutedRunner(MoERunner): +class _TrtllmRunnerBase(MoERunner): + """Load the shared TRTLLM-gen module after support validation.""" + + def _build(self) -> None: + from .core import get_trtllm_moe_sm100_module + + self._module = get_trtllm_moe_sm100_module() + + +class TrtllmFp4RoutedRunner(_TrtllmRunnerBase): """FP4 adapter over the canonical trtllm-gen ``MoERunner``. Translates (MoEActivationPack, MoEWeightPack) into the ``MoeRunnerInputs`` list @@ -928,8 +972,8 @@ class TrtllmFp4RoutedRunner(MoERunner): QuantVariant.W4A16, ) - def check_support(self) -> None: - super().check_support() + def _check_support(self) -> None: + super()._check_support() if self.config.activation.type is not ActivationType.Swiglu: raise NotImplementedError( f"{type(self).__name__} supports only the Swiglu activation." @@ -953,13 +997,13 @@ def check_support(self) -> None: ) def __init__(self, config: MoEConfig, device: torch.device): + super().__init__() from ..tllm_enums import DtypeTrtllmGen, Fp8QuantizationType from ..utils import device_support_pdl - from .core import get_trtllm_moe_sm100_module self.config = config self.device = device - self._module = get_trtllm_moe_sm100_module() + self._module: Any = None routing = config.routing experts = config.experts @@ -1001,6 +1045,7 @@ def __init__(self, config: MoEConfig, device: torch.device): self.tuning_config: Any = None def _ensure_inner(self, hidden_size: int) -> None: + self._require_built() if self._inner is not None: return from ..tllm_enums import WeightLayout @@ -1023,6 +1068,7 @@ def _ensure_inner(self, hidden_size: int) -> None: def get_valid_tactics( # type: ignore[override] self, inputs: List[torch.Tensor], profile: Any ) -> List[Any]: + self._require_built() # The inner runner reads num_tokens from inputs + its own instance key; # no static kwargs are needed for tactic enumeration. return self._inner.get_valid_tactics(inputs, profile) @@ -1034,6 +1080,7 @@ def forward( do_preparation: bool = False, **kwargs: Any, ) -> torch.Tensor: + self._require_built() # MoELayer's autotuner call passes no kwargs, so the static weight/config # kwargs are injected here. The inner runner writes the result in-place # into inputs[0] (the output buffer of the MoeRunnerInputs list). @@ -1173,6 +1220,7 @@ def pack_inputs( (passed via the static kwargs) and dropping ids outside ``[offset, offset + local_num_experts)``. """ + self._require_built() from .core import MoeRunnerInputs, RoutingInputMode v = weights.get_view(self.backend_key) @@ -1329,7 +1377,7 @@ def __hash__(self): # --------------------------------------------------------------------------- -class TrtllmFp8BlockRunner(MoERunner): +class TrtllmFp8BlockRunner(_TrtllmRunnerBase): """Block-FP8 adapter over the canonical trtllm-gen ``MoERunner``. DeepSeek FP8 and MXFP8 share the kernel family but not scale contracts: @@ -1347,8 +1395,8 @@ class TrtllmFp8BlockRunner(MoERunner): QuantVariant.MxFp8, ) - def check_support(self) -> None: - super().check_support() + def _check_support(self) -> None: + super()._check_support() if self.config.activation.type is not ActivationType.Swiglu: raise NotImplementedError( f"{type(self).__name__} supports only the Swiglu activation." @@ -1369,10 +1417,10 @@ def check_support(self) -> None: ) def __init__(self, config: MoEConfig, device: torch.device): + super().__init__() from ..tllm_enums import DtypeTrtllmGen, Fp8QuantizationType from ..utils import device_support_pdl from .api import QuantVariant - from .core import get_trtllm_moe_sm100_module if config.quant.variant is QuantVariant.MxFp8: dtype = DtypeTrtllmGen.MxE4m3 @@ -1385,7 +1433,7 @@ def __init__(self, config: MoEConfig, device: torch.device): self.config = config self.device = device - self._module = get_trtllm_moe_sm100_module() + self._module: Any = None self._variant = config.quant.variant self._dtype_act = dtype self._dtype_weights = dtype @@ -1411,6 +1459,7 @@ def __init__(self, config: MoEConfig, device: torch.device): self.tuning_config: Any = None def _ensure_inner(self, hidden_size: int) -> None: + self._require_built() if self._inner is not None: return from ..tllm_enums import WeightLayout @@ -1433,6 +1482,7 @@ def _ensure_inner(self, hidden_size: int) -> None: def get_valid_tactics( # type: ignore[override] self, inputs: List[torch.Tensor], profile: Any ) -> List[Any]: + self._require_built() return self._inner.get_valid_tactics(inputs, profile) def forward( @@ -1442,6 +1492,7 @@ def forward( do_preparation: bool = False, **kwargs: Any, ) -> torch.Tensor: + self._require_built() self._inner.forward( inputs, tactic=tactic, @@ -1538,6 +1589,7 @@ def _validate_fp8_tensors( def pack_inputs( self, act: MoEActivationPack, weights: MoEWeightPack ) -> List[torch.Tensor]: + self._require_built() from ..tllm_enums import WeightLayout from .core import MoeRunnerInputs, RoutingInputMode @@ -1633,7 +1685,7 @@ def __hash__(self): # --------------------------------------------------------------------------- -class TrtllmFp8PerTensorRunner(MoERunner): +class TrtllmFp8PerTensorRunner(_TrtllmRunnerBase): """Per-tensor-FP8 adapter over the canonical trtllm-gen ``MoERunner``. The kernel consumes prequantized E4M3 activations and weights. Its calibrated @@ -1650,8 +1702,8 @@ class TrtllmFp8PerTensorRunner(MoERunner): ) supported_quant_variants = (QuantVariant.FP8PerTensor,) - def check_support(self) -> None: - super().check_support() + def _check_support(self) -> None: + super()._check_support() from ..tllm_enums import RoutingMethodType from ..utils import get_compute_capability from .api import TrtllmFp8PerTensorConfig @@ -1680,13 +1732,13 @@ def check_support(self) -> None: ) def __init__(self, config: MoEConfig, device: torch.device): + super().__init__() from ..tllm_enums import DtypeTrtllmGen, Fp8QuantizationType from ..utils import device_support_pdl - from .core import get_trtllm_moe_sm100_module self.config = config self.device = device - self._module = get_trtllm_moe_sm100_module() + self._module: Any = None self._dtype_act = DtypeTrtllmGen.E4m3 self._dtype_weights = DtypeTrtllmGen.E4m3 self._fp8_quantization_type = Fp8QuantizationType.NoneFp8 @@ -1710,6 +1762,7 @@ def __init__(self, config: MoEConfig, device: torch.device): self.tuning_config: Any = None def _ensure_inner(self, hidden_size: int) -> None: + self._require_built() if self._inner is not None: return from ..tllm_enums import WeightLayout @@ -1732,6 +1785,7 @@ def _ensure_inner(self, hidden_size: int) -> None: def get_valid_tactics( # type: ignore[override] self, inputs: List[torch.Tensor], profile: Any ) -> List[Any]: + self._require_built() return self._inner.get_valid_tactics(inputs, profile) def forward( @@ -1741,6 +1795,7 @@ def forward( do_preparation: bool = False, **kwargs: Any, ) -> torch.Tensor: + self._require_built() self._inner.forward( inputs, tactic=tactic, @@ -1806,6 +1861,7 @@ def _validate_tensors( def pack_inputs( self, act: MoEActivationPack, weights: MoEWeightPack ) -> List[torch.Tensor]: + self._require_built() from ..tllm_enums import RoutingMethodType from .core import MoeRunnerInputs @@ -1897,7 +1953,7 @@ def __hash__(self): # --------------------------------------------------------------------------- -class TrtllmBf16RoutedRunner(MoERunner): +class TrtllmBf16RoutedRunner(_TrtllmRunnerBase): """BF16 adapter over the canonical trtllm-gen ``MoERunner``. Mirrors :class:`TrtllmFp4RoutedRunner` but with ``Bfloat16`` activation + @@ -1920,8 +1976,8 @@ class TrtllmBf16RoutedRunner(MoERunner): ) supported_quant_variants = (QuantVariant.BF16,) - def check_support(self) -> None: - super().check_support() + def _check_support(self) -> None: + super()._check_support() if self.config.activation.type is not ActivationType.Swiglu: raise NotImplementedError( f"{type(self).__name__} supports only the Swiglu activation." @@ -1941,13 +1997,13 @@ def check_support(self) -> None: ) def __init__(self, config: MoEConfig, device: torch.device): + super().__init__() from ..tllm_enums import DtypeTrtllmGen, Fp8QuantizationType from ..utils import device_support_pdl - from .core import get_trtllm_moe_sm100_module self.config = config self.device = device - self._module = get_trtllm_moe_sm100_module() + self._module: Any = None routing = config.routing experts = config.experts @@ -1972,6 +2028,7 @@ def __init__(self, config: MoEConfig, device: torch.device): self.tuning_config: Any = None def _ensure_inner(self, hidden_size: int) -> None: + self._require_built() if self._inner is not None: return from ..tllm_enums import WeightLayout @@ -1994,6 +2051,7 @@ def _ensure_inner(self, hidden_size: int) -> None: def get_valid_tactics( # type: ignore[override] self, inputs: List[torch.Tensor], profile: Any ) -> List[Any]: + self._require_built() return self._inner.get_valid_tactics(inputs, profile) def forward( @@ -2003,6 +2061,7 @@ def forward( do_preparation: bool = False, **kwargs: Any, ) -> torch.Tensor: + self._require_built() self._inner.forward( inputs, tactic=tactic, @@ -2021,6 +2080,7 @@ def pack_inputs( (the EP bridge does not quantize on the bf16 path); ``act.hidden_states_scale`` is unused. """ + self._require_built() from .core import MoeRunnerInputs, RoutingInputMode v = weights.get_view(self.backend_key) @@ -2119,7 +2179,7 @@ def __hash__(self): # --------------------------------------------------------------------------- -class TrtllmMxInt4RoutedRunner(MoERunner): +class TrtllmMxInt4RoutedRunner(_TrtllmRunnerBase): """MxInt4 adapter over the canonical TRTLLM MoE runner.""" backend_key = "trtllm_mxint4_routed" @@ -2129,8 +2189,8 @@ class TrtllmMxInt4RoutedRunner(MoERunner): ) supported_quant_variants = (QuantVariant.MxInt4,) - def check_support(self) -> None: - super().check_support() + def _check_support(self) -> None: + super()._check_support() if self.config.activation.type is not ActivationType.Swiglu: raise NotImplementedError( f"{type(self).__name__} supports only the Swiglu activation." @@ -2151,13 +2211,13 @@ def check_support(self) -> None: ) def __init__(self, config: MoEConfig, device: torch.device): + super().__init__() from ..tllm_enums import DtypeTrtllmGen, Fp8QuantizationType from ..utils import device_support_pdl - from .core import get_trtllm_moe_sm100_module self.config = config self.device = device - self._module = get_trtllm_moe_sm100_module() + self._module: Any = None routing = config.routing experts = config.experts @@ -2182,6 +2242,7 @@ def __init__(self, config: MoEConfig, device: torch.device): self.tuning_config: Any = None def _ensure_inner(self, hidden_size: int) -> None: + self._require_built() if self._inner is not None: return from ..tllm_enums import WeightLayout @@ -2204,6 +2265,7 @@ def _ensure_inner(self, hidden_size: int) -> None: def get_valid_tactics( # type: ignore[override] self, inputs: List[torch.Tensor], profile: Any ) -> List[Any]: + self._require_built() return self._inner.get_valid_tactics(inputs, profile) def forward( @@ -2213,6 +2275,7 @@ def forward( do_preparation: bool = False, **kwargs: Any, ) -> torch.Tensor: + self._require_built() self._inner.forward( inputs, tactic=tactic, @@ -2224,6 +2287,7 @@ def forward( def pack_inputs( self, act: MoEActivationPack, weights: MoEWeightPack ) -> List[torch.Tensor]: + self._require_built() from .core import MoeRunnerInputs view = weights.get_view(self.backend_key) @@ -2419,8 +2483,8 @@ class _B12xRunner(MoERunner): backend_key: ClassVar[str] = "" required_weight_keys: ClassVar[tuple[str, ...]] = () - def check_support(self) -> None: - super().check_support() + def _check_support(self) -> None: + super()._check_support() from ..cute_dsl import is_cute_dsl_available from ..jit.cpp_ext import get_cuda_version @@ -2448,6 +2512,7 @@ def check_support(self) -> None: raise NotImplementedError("b12x unified MoE requires do_finalize=True.") def __init__(self, config: MoEConfig, device: torch.device): + super().__init__() from .utils import get_b12x_activation_name self.config = config @@ -2458,8 +2523,16 @@ def __init__(self, config: MoEConfig, device: torch.device): self.tuning_config = TuningConfig() self._prepared_weights: dict[str, torch.Tensor] | None = None self._inner: Any = None + self._wrapper_cls: Any = None + + def _build(self) -> None: + """Load the b12x wrapper factory; shapes remain per-call.""" + from .cute_dsl import B12xMoEWrapper + + self._wrapper_cls = B12xMoEWrapper def get_valid_tactics(self, inputs: List[torch.Tensor], profile: Any) -> List[Any]: + self._require_built() return [-1] def _get_quant_mode_name(self) -> str: @@ -2491,15 +2564,14 @@ def _validate_prepared_weights( raise TypeError(f"{self.backend_key} prepared weights must be tensors.") def _ensure_inner(self, hidden_size: int, num_tokens: int) -> None: + self._require_built() if ( self._inner is not None and hidden_size == self._inner.hidden_size and num_tokens <= self._inner.max_num_tokens ): return - from .cute_dsl import B12xMoEWrapper - - self._inner = B12xMoEWrapper( + self._inner = self._wrapper_cls( num_experts=self.config.routing.num_experts, top_k=self.config.routing.top_k, hidden_size=hidden_size, @@ -2515,6 +2587,7 @@ def _ensure_inner(self, hidden_size: int, num_tokens: int) -> None: def pack_inputs( self, act: MoEActivationPack, weights: MoEWeightPack ) -> List[torch.Tensor]: + self._require_built() v = weights.get_view(self.backend_key) self._validate_prepared_weights(v) first_weight = v[self.required_weight_keys[0]] @@ -2578,6 +2651,7 @@ def forward( do_preparation: bool = False, **kwargs: Any, ) -> torch.Tensor: + self._require_built() if tactic != -1: raise ValueError(f"{self.backend_key} supports only tactic -1.") if self._prepared_weights is None: @@ -2635,8 +2709,8 @@ class B12xW4A16Runner(_B12xRunner): "w2_alpha", ) - def check_support(self) -> None: - super().check_support() + def _check_support(self) -> None: + super()._check_support() if self.config.activation.type not in ( ActivationType.Swiglu, ActivationType.Relu2, diff --git a/tests/autotuner/test_autotuner_core.py b/tests/autotuner/test_autotuner_core.py index 0d5c92285a5..bddb23ccc1e 100644 --- a/tests/autotuner/test_autotuner_core.py +++ b/tests/autotuner/test_autotuner_core.py @@ -417,6 +417,128 @@ def fake_profile( assert tuner.stats.tuned_op_successful_configs["dummy_tune"] >= 1 +def test_rank_tactics_returns_top_k_and_caches_winner(monkeypatch): + """rank_tactics should return best-first shortlist and cache the winner.""" + tuner = reset_autotuner() + runner = DummyRunner(valid_tactics=(0, 1, 2)) + inputs = [torch.empty((16, 32), dtype=torch.float32)] + config = TuningConfig() + profile_calls = [] + + def fake_profile( + self, runner_obj, prof_inputs, tactic, tuning_config=None, **kwargs + ): + profile_calls.append(tactic) + return {0: 5.0, 1: 1.0, 2: 3.0}[tactic] + + monkeypatch.setattr(AutoTuner, "_profile_single_kernel", fake_profile) + with autotune(tune_mode=True): + ranked = tuner.rank_tactics("dummy_rank", [runner], config, inputs, k=2) + cached_ranking = tuner.rank_tactics("dummy_rank", [runner], config, inputs, k=3) + + assert ranked == [1, 2] + assert cached_ranking == [1, 2, 0] + assert profile_calls == [0, 1, 2] + _, tactic = tuner.choose_one("dummy_rank", [runner], config, inputs) + assert tactic == 1 + + +def test_rank_tactics_rebuilds_shortlist_from_winner_only_cache(monkeypatch): + tuner = reset_autotuner() + runner = DummyRunner(valid_tactics=(0, 1, 2)) + inputs = [torch.empty((16, 32), dtype=torch.float32)] + config = TuningConfig() + key = AutoTuner._get_cache_key( + "dummy_rank_winner_only", runner, (inputs[0].shape,), config + ) + tuner.profiling_cache[key] = (0, None) + + def fake_profile( + self, runner_obj, prof_inputs, tactic, tuning_config=None, **kwargs + ): + return {0: 5.0, 1: 1.0, 2: 3.0}[tactic] + + monkeypatch.setattr(AutoTuner, "_profile_single_kernel", fake_profile) + with autotune(tune_mode=True): + ranked = tuner.rank_tactics( + "dummy_rank_winner_only", [runner], config, inputs, k=2 + ) + + assert ranked == [1, 2] + + +def test_rank_tactics_rebuilds_shortlist_from_file_cache(monkeypatch): + tuner = reset_autotuner() + runner = DummyRunner(valid_tactics=(0, 1, 2)) + inputs = [torch.empty((16, 32), dtype=torch.float32)] + config = TuningConfig() + key = AutoTuner._get_cache_key( + "dummy_rank_file_cache", runner, (inputs[0].shape,), config + ) + tuner._file_configs[key.file_key] = ("DummyRunner", 0) + profile_calls = [] + + def fake_profile( + self, runner_obj, prof_inputs, tactic, tuning_config=None, **kwargs + ): + profile_calls.append(tactic) + return {0: 5.0, 1: 1.0, 2: 3.0}[tactic] + + monkeypatch.setattr(AutoTuner, "_profile_single_kernel", fake_profile) + with autotune(tune_mode=True): + ranked = tuner.rank_tactics( + "dummy_rank_file_cache", [runner], config, inputs, k=2 + ) + + assert ranked == [1, 2] + assert profile_calls == [0, 1, 2] + + +def test_rank_tactics_uses_mapped_dynamic_profile(monkeypatch): + tuner = reset_autotuner() + runner = DummyRunner(valid_tactics=(0, 1)) + inputs = [torch.empty((12, 4), dtype=torch.float32)] + config = TuningConfig( + dynamic_tensor_specs=( + DynamicTensorSpec( + input_idx=(0,), + dim_idx=(0,), + gen_tuning_buckets=(8, 16), + map_to_tuning_buckets=lambda x: 8 if x <= 8 else 16, + ), + ) + ) + profiled_shapes = [] + + def fake_profile( + self, runner_obj, prof_inputs, tactic, tuning_config=None, **kwargs + ): + profiled_shapes.append(tuple(prof_inputs[0].shape)) + return float(tactic) + + monkeypatch.setattr(AutoTuner, "_profile_single_kernel", fake_profile) + with autotune(tune_mode=True): + ranked = tuner.rank_tactics("dummy_rank_dynamic", [runner], config, inputs, k=2) + + assert ranked == [0, 1] + assert profiled_shapes == [(16, 4), (16, 4)] + + +def test_rank_tactics_outside_tuning_returns_single_cached_or_fallback(): + tuner = reset_autotuner() + runner = DummyRunner(valid_tactics=(0, 1, 2)) + inputs = [torch.empty((4, 8), dtype=torch.float32)] + config = TuningConfig() + + assert tuner.rank_tactics("dummy_rank_infer", [runner], config, inputs, k=3) == [-1] + + key = AutoTuner._get_cache_key( + "dummy_rank_infer", runner, (inputs[0].shape,), config + ) + tuner.profiling_cache[key] = (2, None) + assert tuner.rank_tactics("dummy_rank_infer", [runner], config, inputs, k=3) == [2] + + def test_prepare_input_tensors_reuses_static_and_recreates_dynamic(): """Profiles apply constraints, dynamic inputs are recreated, static inputs are reused.""" tuner = reset_autotuner() diff --git a/tests/moe/test_unified_moe.py b/tests/moe/test_unified_moe.py index c514e08e543..f59006f3dd6 100644 --- a/tests/moe/test_unified_moe.py +++ b/tests/moe/test_unified_moe.py @@ -51,6 +51,7 @@ TrtllmBf16RoutedRunner, TrtllmFp8BlockRunner, TrtllmFp8PerTensorRunner, + TrtllmMxInt4RoutedRunner, ) from flashinfer.fused_moe.api import ( ActivationConfig, @@ -74,6 +75,14 @@ ) from flashinfer.utils import get_compute_capability + +def _build_direct_runner(runner_type, config, device): + runner = runner_type(config, device=device) + runner.check_support() + runner.build() + return runner + + # Reuse the canonical reference implementation + accuracy helpers from the # existing CuteDSL test — keeps tolerance bounds consistent across tests. from tests.moe.test_cute_dsl_fused_moe import ( # noqa: E402 @@ -756,6 +765,88 @@ def forward(self, inputs, **kwargs): assert runner.check_support() is None +class TestBuiltInRunnerLifecycle: + @staticmethod + def _config(variant): + return MoEConfig( + routing=RoutingConfig(num_experts=32, top_k=2), + quant=QuantConfig(variant=variant), + experts=ExpertConfig(intermediate_size=512), + activation=ActivationConfig.swiglu, + execution=ExecutionConfig(enable_pdl=False), + ) + + @pytest.mark.parametrize( + "runner_type,variant", + ( + (TrtllmFp4RoutedRunner, QuantVariant.NVFP4), + (TrtllmFp8BlockRunner, QuantVariant.DeepSeekFp8), + (TrtllmFp8PerTensorRunner, QuantVariant.FP8PerTensor), + (TrtllmBf16RoutedRunner, QuantVariant.BF16), + (TrtllmMxInt4RoutedRunner, QuantVariant.MxInt4), + ), + ) + def test_trtllm_constructor_defers_idempotent_module_build( + self, monkeypatch, runner_type, variant + ): + from flashinfer.fused_moe import core + + module = object() + loads = [] + + def load_module(): + loads.append("module") + return module + + monkeypatch.setattr(core, "get_trtllm_moe_sm100_module", load_module) + runner = runner_type(self._config(variant), torch.device("cuda:0")) + runner._check_support = lambda: None + + assert runner._module is None + assert loads == [] + + runner.check_support() + runner.build() + runner.build() + + assert runner._module is module + assert runner._built + assert loads == ["module"] + + def test_cute_dsl_constructor_defers_idempotent_inner_build(self, monkeypatch): + from flashinfer.fused_moe.cute_dsl import fused_moe, tuner + + events = [] + tuning_config = object() + + class Inner: + def __init__(self, **kwargs): + events.append(("build", kwargs)) + self.tuning_config = tuning_config + + def __hash__(self): + return 0 + + monkeypatch.setattr(tuner, "CuteDslFusedMoENvfp4Runner", Inner) + monkeypatch.setattr(fused_moe, "_cute_dsl_fused_moe_nvfp4_impl", object()) + runner = CuteDslNvfp4Runner( + self._config(QuantVariant.NVFP4), torch.device("cuda:0") + ) + runner._check_support = lambda: None + + assert runner._inner is None + assert events == [] + + runner.check_support() + runner.build() + runner.build() + + assert len(events) == 1 + assert runner._built + assert runner.tuning_config is tuning_config + assert hash(runner) == hash(("cute_dsl_nvfp4", 0)) + + # --------------------------------------------------------------------------- # MoEActivationPack construction + runner-boundary validation (CPU-only) # --------------------------------------------------------------------------- @@ -1702,7 +1793,7 @@ def test_pack_inputs_keeps_global_ids(self, spec, local_expert_offset): local_num_experts=local_num_experts, ), ) - runner = spec.runner_cls(config, device=device) + runner = _build_direct_runner(spec.runner_cls, config, device) # Global expert ids drawn from this rank's local shard. selected_experts = ( @@ -1776,7 +1867,7 @@ def test_pack_inputs_forwards_separate_routing_tensors(self, weights_dtype): local_num_experts=32, ), ) - runner = TrtllmFp4RoutedRunner(config, device=device) + runner = _build_direct_runner(TrtllmFp4RoutedRunner, config, device) ids = torch.randint( 32, 64, (num_tokens, top_k), dtype=torch.int32, device=device ) @@ -1835,7 +1926,7 @@ def test_cuda_graph_replay_matches_eager(self, weights_dtype): local_num_experts=num_experts, ), ) - runner = TrtllmFp4RoutedRunner(config, device=device) + runner = _build_direct_runner(TrtllmFp4RoutedRunner, config, device) act_pack = MoEActivationPack( hidden_states_q=tensors["x"], hidden_states_scale=tensors["x_sf"].squeeze(-1), @@ -1947,7 +2038,7 @@ def run(offset: int) -> torch.Tensor: ), routing_input_mode=routing_input_mode, ) - runner = TrtllmFp4RoutedRunner(config, device=device) + runner = _build_direct_runner(TrtllmFp4RoutedRunner, config, device) inputs = runner.pack_inputs(act_pack, weight_pack) return runner.forward(inputs, tactic=-1).clone() @@ -1994,7 +2085,7 @@ def test_expert_weights_buffer_is_bf16(self, logits_dtype): quant=QuantConfig(variant=QuantVariant.NVFP4), experts=ExpertConfig(intermediate_size=512), ) - runner = TrtllmFp4RoutedRunner(config, device=device) + runner = _build_direct_runner(TrtllmFp4RoutedRunner, config, device) routing_logits = torch.randn( num_tokens, num_experts, dtype=logits_dtype, device=device @@ -2058,7 +2149,7 @@ def _make_bf16_from_logits_inputs(self, logits_dtype): routing_input_mode=RoutingInputMode.FromLogits, routing_logits=logits, ) - runner = TrtllmBf16RoutedRunner(config, device=logits.device) + runner = _build_direct_runner(TrtllmBf16RoutedRunner, config, logits.device) inputs = runner.pack_inputs(logits_act, weights) return runner, inputs, MoeRunnerInputs.from_list(inputs), logits diff --git a/tests/moe/test_unified_moe_b12x.py b/tests/moe/test_unified_moe_b12x.py index b344930b046..153fee1455f 100644 --- a/tests/moe/test_unified_moe_b12x.py +++ b/tests/moe/test_unified_moe_b12x.py @@ -130,6 +130,27 @@ def test_b12x_supports_precomputed_routing(self, runner_type): RoutingInputMode.PackedPrecomputed, ) + def test_b12x_constructor_defers_idempotent_wrapper_build(self, monkeypatch): + import flashinfer.fused_moe.cute_dsl as cute_dsl + + config = self._config(B12xNvfp4Config(), QuantVariant.NVFP4) + + class Wrapper: + pass + + monkeypatch.setattr(cute_dsl, "B12xMoEWrapper", Wrapper) + runner = B12xNvfp4Runner(config, torch.device("cuda:0")) + runner._check_support = lambda: None + + assert runner._wrapper_cls is None + + runner.check_support() + runner.build() + runner.build() + + assert runner._wrapper_cls is Wrapper + assert runner._built + @pytest.mark.parametrize( "runner_type,expected", ((B12xNvfp4Runner, "nvfp4"), (B12xW4A16Runner, "w4a16")), @@ -325,6 +346,7 @@ def run(self, **kwargs): prepared["fc2_input_scale"] = weight runner = object.__new__(runner_type) + runner._built = True runner._prepared_weights = prepared runner._inner = Wrapper() hidden = torch.empty(1, 16) diff --git a/tests/moe/test_unified_moe_cutlass.py b/tests/moe/test_unified_moe_cutlass.py index 353662289b2..d0af7c2a69a 100644 --- a/tests/moe/test_unified_moe_cutlass.py +++ b/tests/moe/test_unified_moe_cutlass.py @@ -27,6 +27,7 @@ RoutingInputMode, ) from flashinfer.fused_moe.layer import _BACKEND_RUNNERS +from flashinfer.fused_moe.runners import MoERunner from flashinfer.fused_moe.prepare import _quantize_mxfp4_linear from flashinfer.fused_moe.utils import map_to_hybrid_bucket from flashinfer.tllm_enums import ActivationType @@ -71,6 +72,86 @@ def test_cutlass_bf16_config_architectures_and_registration(): assert _BACKEND_RUNNERS[CutlassW4A16Config] is CutlassW4A16Runner +def test_all_registered_runners_use_enforced_lifecycle(): + for runner_type in _BACKEND_RUNNERS.values(): + assert issubclass(runner_type, MoERunner) + assert runner_type.check_support is MoERunner.check_support + assert runner_type.build is MoERunner.build + + +@pytest.mark.parametrize("runner_type", tuple(_BACKEND_RUNNERS.values())) +@pytest.mark.parametrize( + "method,args", + ( + ("pack_inputs", (None, None)), + ("get_valid_tactics", ([], None)), + ("forward", ([],)), + ), +) +def test_registered_runner_execution_requires_build(runner_type, method, args): + runner = runner_type.__new__(runner_type) + + with pytest.raises(RuntimeError, match=r"build\(\).*before execution"): + getattr(runner, method)(*args) + + +def test_moe_runner_enforces_lifecycle_order(): + events = [] + + class Runner(MoERunner): + supported_quant_variants = (QuantVariant.BF16,) + + def _check_support(self): + events.append("check_support") + super()._check_support() + + def _build(self): + events.append("build") + + def get_valid_tactics(self, inputs, profile): + self._require_built() + return [-1] + + def forward(self, inputs, **kwargs): + self._require_built() + events.append("execution") + + runner = Runner() + runner.config = _config() + + with pytest.raises(RuntimeError, match=r"check_support\(\).*build\(\)"): + runner.build() + with pytest.raises(RuntimeError, match=r"build\(\).*before execution"): + runner.forward([]) + + runner.check_support() + runner.build() + runner.build() + runner.forward([]) + + assert events == ["check_support", "build", "execution"] + + +def test_failed_support_check_does_not_authorize_build(): + class Runner(MoERunner): + supported_quant_variants = (QuantVariant.NVFP4,) + + def get_valid_tactics(self, inputs, profile): + return [-1] + + def forward(self, inputs, **kwargs): + return None + + runner = Runner() + runner.config = _config() + runner._support_checked = True + + with pytest.raises(NotImplementedError, match="QuantVariant.BF16"): + runner.check_support() + with pytest.raises(RuntimeError, match=r"check_support\(\).*build\(\)"): + runner.build() + + def test_legacy_cutlass_config_is_deprecated(): with pytest.warns(DeprecationWarning, match="CutlassConfig is deprecated"): CutlassConfig() @@ -393,7 +474,7 @@ def test_cutlass_direct_execution_requires_explicit_build(monkeypatch, execute): lambda *args, **kwargs: backend_calls.append("workspace"), ) - with pytest.raises(RuntimeError, match=r"check_support\(\).*build\(\)"): + with pytest.raises(RuntimeError, match=r"build\(\).*before execution"): execute(runner) assert backend_calls == [] @@ -419,6 +500,7 @@ def forward(self, inputs, **kwargs): runner = CutlassBf16Runner.__new__(CutlassBf16Runner) runner._inner = RecordingInner() + runner._built = True runner._workspace = torch.empty(1, dtype=torch.uint8) inputs = [torch.empty(1) for _ in range(6)] @@ -437,9 +519,13 @@ class RecordingTuner: def __init__(self): self.calls = [] - def choose_one(self, custom_op, runners, tuning_config, inputs, **kwargs): - self.calls.append((custom_op, kwargs["gemm_idx"])) - return runners[0], 3 if kwargs["gemm_idx"] == 1 else 9 + def rank_tactics( + self, custom_op, runners, tuning_config, inputs, k=1, **kwargs + ): + self.calls.append((custom_op, kwargs["gemm_idx"], k)) + if kwargs["gemm_idx"] == 1: + return [3, 5][:k] + return [9, 7][:k] class Inner: gemm_idx_for_tuning = None @@ -448,19 +534,39 @@ class Inner: monkeypatch.setattr(AutoTuner, "get", classmethod(lambda cls: tuner)) runner = CutlassBf16Runner.__new__(CutlassBf16Runner) runner._inner = Inner() + runner._built = True runner._device_arch = 100 + runner._stage_tactic_top_k = 2 inputs = [torch.empty(1) for _ in range(6)] tactics = runner.get_valid_tactics(inputs, None) - assert tactics == [(3, 9)] + assert tactics == [(3, 9), (3, 7), (5, 9), (5, 7)] assert tuner.calls == [ - ("moe_cutlass_bf16_sm100_gemm1", 1), - ("moe_cutlass_bf16_sm100_gemm2", 2), + ("moe_cutlass_bf16_sm100_gemm1", 1, 2), + ("moe_cutlass_bf16_sm100_gemm2", 2, 2), ] assert runner._inner.gemm_idx_for_tuning is None +def test_cutlass_stage_top_k_one_preserves_single_compound_pair(monkeypatch): + class RecordingTuner: + def rank_tactics( + self, custom_op, runners, tuning_config, inputs, k=1, **kwargs + ): + return [3] if kwargs["gemm_idx"] == 1 else [9] + + monkeypatch.setattr(AutoTuner, "get", classmethod(lambda cls: RecordingTuner())) + runner = CutlassBf16Runner.__new__(CutlassBf16Runner) + runner._inner = type("Inner", (), {"gemm_idx_for_tuning": None})() + runner._built = True + runner._device_arch = 100 + runner._stage_tactic_top_k = 1 + inputs = [torch.empty(1) for _ in range(6)] + + assert runner.get_valid_tactics(inputs, None) == [(3, 9)] + + def test_cutlass_outer_cache_key_includes_enable_pdl(): runner = CutlassBf16Runner.__new__(CutlassBf16Runner) runner._device_arch = 90 @@ -479,6 +585,7 @@ def test_cutlass_direct_runner_rejects_tokens_above_tuning_ceiling(): runner.config = _config() runner.device = torch.device("cpu") runner._inner = object() + runner._built = True num_tokens, hidden_size, top_k = 65, 128, 2 act = MoEActivationPack( torch.empty(num_tokens, hidden_size, dtype=torch.bfloat16), @@ -510,8 +617,8 @@ def build(): def ensure_workspace(num_tokens, hidden_size): events.append("workspace") - runner.check_support = check_support - runner.build = build + runner._check_support = check_support + runner._build = build runner._ensure_workspace = ensure_workspace runner._pack_weight_inputs = lambda view, hidden_size: [ torch.empty(1), diff --git a/tests/moe/test_unified_moe_fp8.py b/tests/moe/test_unified_moe_fp8.py index 70c1a373f38..a0d0c5bc25c 100644 --- a/tests/moe/test_unified_moe_fp8.py +++ b/tests/moe/test_unified_moe_fp8.py @@ -34,6 +34,13 @@ from tests.moe.trtllm_gen_fused_moe_utils import check_accuracy +def _build_per_tensor_fp8_runner(config): + runner = TrtllmFp8PerTensorRunner(config, torch.device("cuda")) + runner.check_support() + runner.build() + return runner + + def _is_trtllm_fp8_arch() -> bool: return torch.cuda.is_available() and get_compute_capability( torch.device("cuda") @@ -668,7 +675,7 @@ def test_fp8_per_tensor_layer_and_direct_runner_match_reference(routing_input_mo layer_out = MoELayer(config)(act, weights) _assert_per_tensor_fp8_close(layer_out, ref) - runner = TrtllmFp8PerTensorRunner(config, torch.device("cuda")) + runner = _build_per_tensor_fp8_runner(config) inputs = runner.pack_inputs(act, weights) direct_out = runner.forward(inputs) _assert_per_tensor_fp8_close(direct_out, ref) @@ -687,7 +694,7 @@ def test_fp8_per_tensor_llama4_routes_scale_on_input(routing_input_mode): ) _assert_per_tensor_fp8_close(MoELayer(config)(act, weights), ref) - runner = TrtllmFp8PerTensorRunner(config, torch.device("cuda")) + runner = _build_per_tensor_fp8_runner(config) _assert_per_tensor_fp8_close(runner.forward(runner.pack_inputs(act, weights)), ref) invalid_config = dataclasses.replace( @@ -715,7 +722,7 @@ def test_fp8_per_tensor_nonzero_expert_offset(routing_input_mode): assert torch.count_nonzero(ref) _assert_per_tensor_fp8_close(MoELayer(config)(act, weights), ref) - runner = TrtllmFp8PerTensorRunner(config, torch.device("cuda")) + runner = _build_per_tensor_fp8_runner(config) _assert_per_tensor_fp8_close(runner.forward(runner.pack_inputs(act, weights)), ref) @@ -734,7 +741,7 @@ def test_fp8_per_tensor_packed_ids_keep_global_ids_and_weight_bits(): act.topk_weights.to(torch.bfloat16).view(torch.int16).to(torch.int32) & 0xFFFF ) - runner = TrtllmFp8PerTensorRunner(config, torch.device("cuda")) + runner = _build_per_tensor_fp8_runner(config) moe_inputs = MoeRunnerInputs.from_list(runner.pack_inputs(act, weights)) packed = moe_inputs.topk_ids assert moe_inputs.expert_weights is None @@ -753,7 +760,7 @@ def test_fp8_per_tensor_noncontiguous_packed_routing_matches_reference(): assert not act.topk_ids.is_contiguous() assert not act.topk_weights.is_contiguous() - runner = TrtllmFp8PerTensorRunner(config, torch.device("cuda")) + runner = _build_per_tensor_fp8_runner(config) inputs = runner.pack_inputs(act, weights) assert MoeRunnerInputs.from_list(inputs).topk_ids.is_contiguous() _assert_per_tensor_fp8_close(runner.forward(inputs), ref) @@ -761,7 +768,7 @@ def test_fp8_per_tensor_noncontiguous_packed_routing_matches_reference(): def test_fp8_per_tensor_routing_replay_matches_reference(): act, weights, config, _, selected_experts = _make_per_tensor_fp8_case() - runner = TrtllmFp8PerTensorRunner(config, torch.device("cuda")) + runner = _build_per_tensor_fp8_runner(config) inputs = runner.pack_inputs(act, weights) replay = torch.full( (TOKENS, TOP_K), -1, dtype=torch.int16, device=torch.device("cuda") @@ -785,7 +792,7 @@ def test_fp8_per_tensor_cuda_graph_replay(routing_input_mode): act, weights, config, ref, _ = _make_per_tensor_fp8_case( routing_input_mode=routing_input_mode ) - runner = TrtllmFp8PerTensorRunner(config, torch.device("cuda")) + runner = _build_per_tensor_fp8_runner(config) inputs = runner.pack_inputs(act, weights) runner.forward(inputs) graph = torch.cuda.CUDAGraph() diff --git a/tests/moe/test_unified_moe_mxint4.py b/tests/moe/test_unified_moe_mxint4.py index 6e8e02c9e91..c4e7c5f8207 100644 --- a/tests/moe/test_unified_moe_mxint4.py +++ b/tests/moe/test_unified_moe_mxint4.py @@ -31,6 +31,13 @@ from flashinfer.utils import get_compute_capability +def _build_mxint4_runner(config): + runner = TrtllmMxInt4RoutedRunner(config, torch.device("cuda")) + runner.check_support() + runner.build() + return runner + + def _is_mxint4_arch() -> bool: return torch.cuda.is_available() and get_compute_capability( torch.device("cuda") @@ -349,7 +356,7 @@ def test_mxint4_layer_and_direct_runner_match_reference(routing_input_mode): layer_output = MoELayer(config)(act, weights) _assert_mxint4_close(layer_output, reference) - runner = TrtllmMxInt4RoutedRunner(config, torch.device("cuda")) + runner = _build_mxint4_runner(config) direct_output = runner.forward(runner.pack_inputs(act, weights)) _assert_mxint4_close(direct_output, reference) @@ -387,7 +394,7 @@ def test_mxint4_from_logits_rejects_fp32_until_validated(): routing_input_mode=RoutingInputMode.FromLogits ) act.routing_logits = act.routing_logits.float() - runner = TrtllmMxInt4RoutedRunner(config, torch.device("cuda")) + runner = _build_mxint4_runner(config) with pytest.raises(TypeError, match="requires bfloat16 routing_logits"): runner.pack_inputs(act, weights) @@ -399,7 +406,7 @@ def test_mxint4_from_logits_rejects_fp32_bias(): routing_method=RoutingMethodType.DeepSeekV3, ) act.routing_bias = act.routing_bias.float() - runner = TrtllmMxInt4RoutedRunner(config, torch.device("cuda")) + runner = _build_mxint4_runner(config) with pytest.raises(TypeError, match="routing_bias must be bfloat16"): runner.pack_inputs(act, weights) @@ -419,7 +426,7 @@ def test_mxint4_runner_rejects_noncontiguous_runtime_inputs(field): assert not tensor.is_contiguous() setattr(act, field, tensor) - runner = TrtllmMxInt4RoutedRunner(config, torch.device("cuda")) + runner = _build_mxint4_runner(config) with pytest.raises(ValueError, match=rf"{field} must be contiguous"): runner.pack_inputs(act, weights) @@ -438,7 +445,7 @@ def test_mxint4_runner_rejects_malformed_prepared_view(key): act, weights, config, _, _ = _make_case() view = weights.get_view("trtllm_mxint4_routed") view[key] = view[key][..., :-1].contiguous() - runner = TrtllmMxInt4RoutedRunner(config, torch.device("cuda")) + runner = _build_mxint4_runner(config) with pytest.raises(ValueError, match=rf"{key} shape"): runner.pack_inputs(act, weights) @@ -459,7 +466,7 @@ def test_mxint4_runner_rejects_unaligned_runtime_geometry(dimension): config, experts=dataclasses.replace(config.experts, intermediate_size=384), ) - runner = TrtllmMxInt4RoutedRunner(config, torch.device("cuda")) + runner = _build_mxint4_runner(config) with pytest.raises(ValueError, match="divisible by 256"): runner.pack_inputs(act, weights) @@ -497,7 +504,7 @@ def test_mxint4_runner_validates_optional_gemm1_params(mutation, error_type, mat assert not value.is_contiguous() view["gemm1_alpha"] = value - runner = TrtllmMxInt4RoutedRunner(config, torch.device("cuda")) + runner = _build_mxint4_runner(config) with pytest.raises(error_type, match=match): runner.pack_inputs(act, weights) @@ -512,7 +519,7 @@ def test_mxint4_cuda_graph_replay(routing_input_mode): act, weights, config, reference, _ = _make_case( routing_input_mode=routing_input_mode ) - runner = TrtllmMxInt4RoutedRunner(config, torch.device("cuda")) + runner = _build_mxint4_runner(config) inputs = runner.pack_inputs(act, weights) eager = runner.forward(inputs).clone() graph = torch.cuda.CUDAGraph()