diff --git a/.claude/skills/trtllm-moe-develop/references/moe-canonical-code-examples.md b/.claude/skills/trtllm-moe-develop/references/moe-canonical-code-examples.md index 2ed20d941465..fe5dee17e0f8 100644 --- a/.claude/skills/trtllm-moe-develop/references/moe-canonical-code-examples.md +++ b/.claude/skills/trtllm-moe-develop/references/moe-canonical-code-examples.md @@ -90,7 +90,7 @@ only names the code that demonstrates them. Paths are relative to `find_cutlass_leaf` / `cutlass_leaf`; gates composed in `eligibility.py`. - `trtllm_gen/`: per-kernel-ABI layers (`fp4_block_scale.py`, `fp8_block_scale.py`) between `TrtllmGenFusedMoEBase` and the - `trtllm_.py` / `flashinfer_.py` leaves. + `trtllm_.py` leaves and the unquantized `flashinfer_bf16.py` leaf. - `marlin/`: one kernel ABI, so `MarlinFusedMoEBase` implements three of the four abstract methods and a leaf is an identity plus `can_implement`. - `cutedsl/`: laid out like the id `trtllm.cutedsl..`. diff --git a/tensorrt_llm/_torch/moe/fused_moe/MOE_DEVELOPER_GUIDE.md b/tensorrt_llm/_torch/moe/fused_moe/MOE_DEVELOPER_GUIDE.md index a02033e6794f..fa528fe735ea 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/MOE_DEVELOPER_GUIDE.md +++ b/tensorrt_llm/_torch/moe/fused_moe/MOE_DEVELOPER_GUIDE.md @@ -221,7 +221,7 @@ Still on old path (standalone, with embedded communication): | `routing.py` | Routing methods (`TopKRouting`, etc.) | | `moe_comm_timeout_guard.py` | Device-side timeouts of communication kernels, relaxed during warmup — budget resolution and `MoECommTimeoutProxy` registration (see [Communication timeouts](#communication-timeouts)) | | `moe_load_balancer.py` | EPLB implementation | -| `moe_op_backend.py` | Op backend registry for TRTLLMGen (flashinfer/trtllm ops) | +| `moe_op_backend.py` | Op backend registry for TRTLLMGen: `trtllm` ops for every quantized leaf, the FlashInfer bf16 runner for the unquantized one | ### Backends (`fused_moe/`) @@ -308,48 +308,39 @@ holds no code: it defines `TRTLLMGenFusedMoE` (the abstract |------|------| | `identity.py` | TRTLLM-Gen's values of the types `impl_identity.py` defines: the identity constants (`PROVIDER_*`, `TECHNIQUE_TRTLLM_GEN`, `KERNEL_FUSED_MOE`), the two published contracts (`TRTLLM_GEN_CAPABILITIES`, `TRTLLM_GEN_INPUT_REQUIREMENT`), and the `trtllm_gen_descriptor` factory that stamps those segments into what the registry reads | | `base.py` | `TrtllmGenFusedMoEBase` — the abstract root every leaf shares: construction, the weight-creation skeleton, the two routing predicates the framework reads off a module, the fake-output shapes. Reads both axes as leaf-declared attributes and overridable hooks, never as a branch on `quant_config` or the provider string | -| `fp4_block_scale.py` | `run_fp4_block_scale_moe` and the formats that reach it — NVFP4, W4A16_MXFP4, W4A8_MXFP4_MXFP8 — each supplying only its weight/input preparation, its SiTu alignment, and its scale group size. Two leaves per format, one per provider | -| `fp8_block_scale.py` | DeepSeek-style FP8 with 1x128 block scales: `run_fp8_block_scale_moe`, the narrowed activation ABI (scalar clamp), and shared-expert fusion. Two leaves. Shares nothing with the fp4 module, which is why the two sit side by side rather than under a common quant layer | +| `fp4_block_scale.py` | `run_fp4_block_scale_moe` and the formats that reach it — NVFP4, W4A16_MXFP4, W4A8_MXFP4_MXFP8 — each supplying only its weight/input preparation, its SiTu alignment, and its scale group size. One leaf per format | +| `fp8_block_scale.py` | DeepSeek-style FP8 with 1x128 block scales: `run_fp8_block_scale_moe`, the narrowed activation ABI (scalar clamp), and shared-expert fusion. One leaf. Shares nothing with the fp4 module, which is why the two sit side by side rather than under a common quant layer | | `eligibility.py` | the eligibility checks each leaf's `can_implement` composes, as free functions — every caller is a leaf | | `kernel_inputs.py` | the `run_moe` prologue the kernel bodies share, as free functions taking the impl | | `trtllm_.py` | one registered native leaf each: `nvfp4`, `fp8_block_scales`, `w4a16_mxfp4`, `w4a8_mxfp4_mxfp8`, `w4a8_nvfp4_fp8`, `w4a8_mxfp4_fp8` | -| `flashinfer_.py` | one registered FlashInfer leaf each: `nvfp4`, `fp8_block_scales`, `w4a16_mxfp4`, `w4a8_mxfp4_mxfp8`, `bf16` | +| `flashinfer_bf16.py` | the one registered FlashInfer leaf: the unquantized format, which TRT-LLM's own cubins do not cover | #### The leaves Module names carry the two identity segments that differ: `provider` — the in-tree `trtllm` ops or the FlashInfer package — and `quant`. Technique and kernel are `trtllm_gen` / `fused_moe` for every leaf, which the package name -says once. The grid is not full: FlashInfer has no `w4a8_nvfp4_fp8` or -`w4a8_mxfp4_fp8` runner and alone serves the unquantized format, so some -cells of the provider x format grid have no leaf. - -| `quant` segment | `trtllm.` leaf | `flashinfer.` leaf | -|---|---|---| -| `nvfp4` | `TrtllmTrtllmGenNvfp4Impl` | `FlashinferTrtllmGenNvfp4Impl` | -| `fp8_block_scales` | `TrtllmTrtllmGenFp8BlockScalesImpl` | `FlashinferTrtllmGenFp8BlockScalesImpl` | -| `w4a16_mxfp4` | `TrtllmTrtllmGenW4a16Mxfp4Impl` | `FlashinferTrtllmGenW4a16Mxfp4Impl` | -| `w4a8_mxfp4_mxfp8` | `TrtllmTrtllmGenW4a8Mxfp4Mxfp8Impl` | `FlashinferTrtllmGenW4a8Mxfp4Mxfp8Impl` | -| `w4a8_nvfp4_fp8` | `TrtllmTrtllmGenW4a8Nvfp4Fp8Impl` | — | -| `w4a8_mxfp4_fp8` | `TrtllmTrtllmGenW4a8Mxfp4Fp8Impl` | — | -| `none` (BF16) | — | `FlashinferTrtllmGenBf16Impl` | +says once. Each format has exactly one leaf: every quantized format runs on the +native provider, and FlashInfer serves only the unquantized format, for which +TRT-LLM ships no bf16 x bf16 cubins. + +| `quant` segment | Leaf | +|---|---| +| `nvfp4` | `TrtllmTrtllmGenNvfp4Impl` | +| `fp8_block_scales` | `TrtllmTrtllmGenFp8BlockScalesImpl` | +| `w4a16_mxfp4` | `TrtllmTrtllmGenW4a16Mxfp4Impl` | +| `w4a8_mxfp4_mxfp8` | `TrtllmTrtllmGenW4a8Mxfp4Mxfp8Impl` | +| `w4a8_nvfp4_fp8` | `TrtllmTrtllmGenW4a8Nvfp4Fp8Impl` | +| `w4a8_mxfp4_fp8` | `TrtllmTrtllmGenW4a8Mxfp4Fp8Impl` | +| `none` (BF16) | `FlashinferTrtllmGenBf16Impl` | Ids follow the class names (`trtllm.trtllm_gen.fused_moe.nvfp4`, `flashinfer.trtllm_gen.fused_moe.none`), keeping the technique and kernel segments so a future non-TRTLLM-Gen kernel from the same wheel stays -distinguishable. The two native-only formats subclass the family base directly: -a per-quant parent implementing three methods for one subclass buys nothing. - -The provider split is a policy call, not a capability one. With the FlashInfer -wheel installed a `TRTLLM` literal still picks the native leaf, and only -`TRTLLM_GEN_FUSED_MOE_USE_FLASHINFER=1` moves it — read through -`MoEEnvironment.env_flag`, so `can_implement` never touches `os.environ`. A -FlashInfer leaf asked for without the flag declines with `PATH_NOT_ENABLED`, -which distinguishes "not chosen" from "not installed" (`DEP_MISSING`). -`IMPL_PRIORITY` lists each FlashInfer leaf ahead of the native leaf of the same -format, so the flag alone decides the order the two are asked in. -`FlashinferTrtllmGenBf16Impl` is outside that policy — see the `§` footnote -under [Quantization Support](#quantization-support). +distinguishable. The two native formats on runner ops of their own subclass the +family base directly: a per-quant parent implementing three methods for one +subclass buys nothing. Because the `quant` segments are disjoint, the order of +the TRTLLM-Gen leaves within `IMPL_PRIORITY` does not matter. Only a registered class may be resolved to, and importing the package is what registers, which `moe_resolution` guarantees. `trtllm_gen_leaf` sits at the @@ -660,8 +651,7 @@ Each backend's `can_implement(p, d)` classmethod declares what it supports. Sour § The unquantized TRTLLM-Gen path is `FlashinferTrtllmGenBf16Impl`, and it is not a TRTLLM-Gen kernel at all: it calls FlashInfer's `trtllm_bf16_moe` / `trtllm_bf16_routed_moe`, which is why it is gated on -`MoEDep.FLASHINFER_BF16_MOE` rather than on a quant algo or on the -`TRTLLM_GEN_FUSED_MOE_USE_FLASHINFER` opt-in, and why `TRTLLMOpBackend` raises +`MoEDep.FLASHINFER_BF16_MOE` and why `TRTLLMOpBackend` raises `NotImplementedError` for it. It is the only leaf reachable on a host with no quantization configured at all. The row reads `Y` because `can_implement` really can select it; without the FlashInfer symbols the layer degrades to diff --git a/tensorrt_llm/_torch/moe/fused_moe/communication/communication_factory.py b/tensorrt_llm/_torch/moe/fused_moe/communication/communication_factory.py index 7d452b358f72..dc5783999172 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/communication/communication_factory.py +++ b/tensorrt_llm/_torch/moe/fused_moe/communication/communication_factory.py @@ -42,7 +42,6 @@ from .nccl_ep import NcclEP from .nvlink_one_sided import NVLinkOneSided from .nvlink_two_sided import NVLinkTwoSided -from .nvlink_two_sided_flashinfer import NVLinkTwoSidedFlashinfer # Temporary NCCL-EP v0.1.0 limitation. The LL combine kernel derives its # warp-group count and dynamic-SMEM requirement here: @@ -92,7 +91,6 @@ def create_strategy( expert_size_per_partition: int, payload_in_workspace: bool = False, alltoall_result_do_sum: bool = True, - use_flashinfer: bool = False, hidden_size: Optional[int] = None, communication_method: Optional[str] = None, allow_alltoall: bool = True, @@ -196,7 +194,6 @@ def create_strategy( expert_size_per_partition, payload_in_workspace, alltoall_result_do_sum, - use_flashinfer, hidden_size=hidden_size, ) logger.info( @@ -241,24 +238,14 @@ def create_strategy( return AllGatherReduceScatter(mapping) try: - if use_flashinfer: - strategy = NVLinkTwoSidedFlashinfer( - mapping, - num_experts, - num_slots, - top_k, - use_low_precision_combine, - alltoall_result_do_sum=alltoall_result_do_sum, - ) - else: - strategy = NVLinkTwoSided( - mapping, - num_experts, - num_slots, - top_k, - use_low_precision_combine, - alltoall_result_do_sum=alltoall_result_do_sum, - ) + strategy = NVLinkTwoSided( + mapping, + num_experts, + num_slots, + top_k, + use_low_precision_combine, + alltoall_result_do_sum=alltoall_result_do_sum, + ) logger.info("Selected communication strategy: NVLinkTwoSided") return strategy except Exception as e: @@ -351,7 +338,6 @@ def _create_forced_method( expert_size_per_partition: int, payload_in_workspace: bool, alltoall_result_do_sum: bool, - use_flashinfer: bool, hidden_size: Optional[int] = None, ) -> Communication: """ @@ -385,24 +371,14 @@ def _create_forced_method( # Create strategy - will raise RuntimeError if platform not supported if method in ["NVLINK_TWO_SIDED"]: - if use_flashinfer: - return NVLinkTwoSidedFlashinfer( - mapping, - num_experts, - num_slots, - top_k, - use_low_precision_combine, - alltoall_result_do_sum=alltoall_result_do_sum, - ) - else: - return NVLinkTwoSided( - mapping, - num_experts, - num_slots, - top_k, - use_low_precision_combine, - alltoall_result_do_sum=alltoall_result_do_sum, - ) + return NVLinkTwoSided( + mapping, + num_experts, + num_slots, + top_k, + use_low_precision_combine, + alltoall_result_do_sum=alltoall_result_do_sum, + ) elif method in ["NVLINK_ONE_SIDED"]: enable_eplb = model_config.moe_load_balancer is not None ep_group_health, watchdog_timeout_s, watchdog_poll_interval_s = get_wide_ep_ft_options( diff --git a/tensorrt_llm/_torch/moe/fused_moe/communication/nvlink_two_sided_flashinfer.py b/tensorrt_llm/_torch/moe/fused_moe/communication/nvlink_two_sided_flashinfer.py deleted file mode 100644 index b3aaa2f8ad0e..000000000000 --- a/tensorrt_llm/_torch/moe/fused_moe/communication/nvlink_two_sided_flashinfer.py +++ /dev/null @@ -1,283 +0,0 @@ -# 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. - -""" -NVLINK Two-Sided AllToAll Communication Strategy - -This module implements the NVLINK two-sided comm AllToAll communication method for MoE. - -NVLINK Two-Sided supports post-quant dispatch for all quantization modes. -""" - -import functools -import os -from typing import List, Optional, Tuple - -import pynvml -import torch - -from tensorrt_llm.logger import logger -from tensorrt_llm.mapping import Mapping - -from .base import Communication - - -@functools.lru_cache(maxsize=1) -def _flashinfer_mnnvl(): - """Resolve FlashInfer's MNNVL symbols on first use. - - Deliberately not a module-level import: ``flashinfer.comm`` does not import - on every interpreter this package supports. FlashInfer >= 0.6.16 annotates - ``array.array[int]`` at module scope in ``comm/fd_exchange.py``, and - ``array.array`` only became subscriptable in Python 3.12, so that annotation - raises ``TypeError`` on the 3.10/3.11 interpreters covered by our - ``python_requires`` (FlashInfer declares ``requires_python >= 3.10`` too, so - this is an upstream bug). At module scope that would break ``import - tensorrt_llm`` outright; resolving it lazily confines the damage to this one - alltoall strategy. - - Returns ``(MnnvlMemory, MnnvlMoe)``, or ``None`` when unavailable. - """ - try: - from flashinfer.comm.mnnvl import MnnvlMemory - from flashinfer.comm.trtllm_alltoall import MnnvlMoe - except (ImportError, TypeError) as e: - # TypeError covers the annotation-evaluation failure described above; - # it is not a generic catch-all. - logger.warning( - f"flashinfer.comm is unavailable ({e}); the NVLinkTwoSidedFlashinfer " - f"MoE alltoall strategy is disabled." - ) - return None - return MnnvlMemory, MnnvlMoe - - -class NVLinkTwoSidedFlashinfer(Communication): - """ - NVLINK two-sided comm AllToAll strategy. - This implementation utilizes symmetric memory to enable peer-to-peer access between GPUs over NVLink. - The kernel takes the role as both sender and receiver: as the sender, it puts the data into a FIFO - quene in peer ranks' symmetric memory; as the receiver, it gets the data from the FIFO quene to the - local buffer. This communication model is akin to NCCL's collective operations. - The required symmetric memory size is proportional to the communication channels opened. - """ - - def __init__( - self, - mapping: Mapping, - num_experts: int, - num_slots: int, - top_k: int = 1, - use_low_precision_combine: bool = False, - alltoall_result_do_sum: bool = False, - ): - super().__init__(mapping) - if mapping.has_cp_helix(): - raise ValueError( - "NVLinkTwoSidedFlashinfer does not support Helix context parallelism " - "because its MNNVL communicator covers only the tensor-parallel group" - ) - - # Store needed parameters - self.num_experts = num_experts - self.num_slots = num_slots - self.top_k = top_k - - assert alltoall_result_do_sum, "flashinfer do_reduce must be True" - assert use_low_precision_combine is False, ( - "flashinfer use_low_precision_combine must be False" - ) - - self.use_low_precision_combine = use_low_precision_combine - self.alltoall_result_do_sum = alltoall_result_do_sum - - # Read from environment variable, same as wideEP - self.enable_postquant_alltoall = ( - os.environ.get("TRTLLM_MOE_POST_QUANT_ALLTOALLV", "1") == "1" - ) - - # Invalid token expert ID (default to -1), the kernels in TRTLLM-gen is hard-coded to support -1 only. - # CutlassFusedMoE kernels support any invalid value. - self.invalid_token_expert_id: int = -1 - - mnnvl = _flashinfer_mnnvl() - if mnnvl is None: - raise RuntimeError( - "NVLinkTwoSidedFlashinfer requires flashinfer.comm, which failed " - "to import (see the preceding warning); select a different MoE " - "alltoall strategy or run on Python >= 3.12." - ) - flashinfer_MnnvlMemory, self._mnnvl_moe = mnnvl - - # Initialize NVLINK workspaces - flashinfer_MnnvlMemory.initialize() - self.alltoall_workspace = self._mnnvl_moe.get_moe_workspaces(mapping) - self.alltoall_prepare_workspace = self._mnnvl_moe.get_moe_prepare_workspace(mapping) - - # Initialize dispatch state - self._dispatch_state = {} - - @staticmethod - def is_platform_supported() -> bool: - """ - Check if NVLINK two-sided comm is supported on current hardware. - """ - mnnvl = _flashinfer_mnnvl() - if mnnvl is None: - return False - flashinfer_MnnvlMemory, _ = mnnvl - - # flashinfer's MnnvlMemory.supports_mnnvl() queries NVML without - # initializing it (only its initialize() guards), so satisfy that - # precondition on our side of the boundary. - try: - pynvml.nvmlDeviceGetCount() - except pynvml.NVMLError_Uninitialized: - pynvml.nvmlInit() - return flashinfer_MnnvlMemory.supports_mnnvl() - - def supports_post_quant_dispatch(self) -> bool: - """ - NVLINK two-sided comm supports post-quant for all modes. - """ - return self.enable_postquant_alltoall - - def is_workload_feasible(self, all_rank_num_tokens: List[int], num_chunks: int) -> bool: - """ - Check if NVLINK two-sided comm is feasible for the given workload at runtime. - - This method performs runtime checks based on workload characteristics such as - token counts, number of chunks, and other runtime parameters. - """ - return True - - def prepare_dispatch( - self, - token_selected_slots: torch.Tensor, - all_rank_num_tokens: List[int], - local_statistic_tensor: Optional[torch.Tensor] = None, - ) -> Optional[torch.Tensor]: - """ - NVLINK two-sided comm prepare dispatch: gather EPLB statistics and prepare alltoall_info. - """ - all_rank_max_num_tokens = max(all_rank_num_tokens) - top_k = token_selected_slots.shape[1] - - # Call NVLINK prepare to get alltoall_info and gather EPLB statistics - alltoall_info, _, __, gathered_local_statistic_tensor = ( - self._mnnvl_moe.mnnvl_moe_alltoallv_prepare_without_allgather( - token_selected_slots, - None, - local_statistic_tensor, - self.alltoall_prepare_workspace, - all_rank_max_num_tokens, - self.ep_rank, - self.ep_size, - self.num_experts, - self.num_slots, - top_k, - ) - ) - - # Store alltoall_info in dispatch_state for use in dispatch() - self._dispatch_state["alltoall_info"] = alltoall_info - - return gathered_local_statistic_tensor - - def dispatch( - self, - hidden_states: torch.Tensor, - hidden_states_sf: Optional[torch.Tensor], - token_selected_slots: torch.Tensor, - token_final_scales: Optional[torch.Tensor], - all_rank_num_tokens: List[int], - use_dp_padding: Optional[bool] = None, - **kwargs, - ) -> Tuple[torch.Tensor, Optional[torch.Tensor], torch.Tensor, Optional[torch.Tensor]]: - """ - NVLINK two-sided comm dispatch (post-quant, uses alltoall_info from prepare_dispatch). - """ - - def mnnvl_moe_alltoallv_packed(x, alltoall_info, workspace, ep_rank, ep_size): - results = [] - for tensor in x: - if tensor is not None: - result = self._mnnvl_moe.mnnvl_moe_alltoallv( - tensor, alltoall_info, workspace, ep_rank, ep_size - ) - results.append(result) - else: - results.append(None) - return results - - # Read alltoall_info from dispatch_state (set by prepare_dispatch) - alltoall_info = self._dispatch_state.get("alltoall_info") - if alltoall_info is None: - raise ValueError( - "NVLinkTwoSidedFlashinfer dispatch requires prepare_dispatch() to be called first" - ) - - all_rank_max_num_tokens = max(all_rank_num_tokens) - original_token_count = hidden_states.shape[0] # Store for combine - top_k = token_selected_slots.shape[1] - - # Dispatch quantized data using AllToAll - hidden_states, hidden_states_sf, token_selected_slots, token_final_scales = ( - mnnvl_moe_alltoallv_packed( - [hidden_states, hidden_states_sf, token_selected_slots, token_final_scales], - alltoall_info, - self.alltoall_workspace, - self.ep_rank, - self.ep_size, - ) - ) - - # Set expert IDs after alltoall - torch.ops.trtllm.memset_expert_ids( - token_selected_slots, - alltoall_info.recv_rank_count_cumsum, - all_rank_max_num_tokens, - top_k, - self.invalid_token_expert_id, - self.ep_size, - ) - - # Store original_token_count for combine (alltoall_info already stored in prepare_dispatch) - self._dispatch_state["original_token_count"] = original_token_count - - return hidden_states, hidden_states_sf, token_selected_slots, token_final_scales - - def combine( - self, - final_hidden_states: torch.Tensor, - **kwargs, - ) -> torch.Tensor: - """ - NVLINK two-sided comm combine - reads from self._dispatch_state. - """ - if isinstance(final_hidden_states, list): - final_hidden_states = final_hidden_states[0] - - final_hidden_states = self._mnnvl_moe.mnnvl_moe_alltoallv_combine( - final_hidden_states, - self._dispatch_state["alltoall_info"], - self.alltoall_workspace, - ep_rank=self.ep_rank, - ep_size=self.ep_size, - top_k=self.top_k, - token_count=self._dispatch_state["original_token_count"], - ) - - return final_hidden_states diff --git a/tensorrt_llm/_torch/moe/fused_moe/configurable_moe.py b/tensorrt_llm/_torch/moe/fused_moe/configurable_moe.py index 1e213d342e69..d91ebd8b1acd 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/configurable_moe.py +++ b/tensorrt_llm/_torch/moe/fused_moe/configurable_moe.py @@ -388,7 +388,6 @@ def _create_and_sync_backend( # enforce ``moe.comm is None`` without ``getattr`` guards. self.backend = backend self._reject_non_divisible_ep_backend() - self.use_flashinfer = getattr(self.backend, "use_flashinfer", False) # Mirror wrapper-owned EPLB / layer-id state onto the backend so any # backend code path that reads e.g. ``self.layer_load_balancer`` or @@ -666,7 +665,6 @@ def _create_comm_strategy_auto(self) -> Optional[Communication]: # Currently the TRTLLMGEN reduce sum internally. # Keep updated with more supported backends. alltoall_result_do_sum=True, - use_flashinfer=self.use_flashinfer, hidden_size=self.hidden_size, communication_method=self.communication_method, allow_alltoall=self.backend.supports_alltoall, diff --git a/tensorrt_llm/_torch/moe/fused_moe/impl_environment.py b/tensorrt_llm/_torch/moe/fused_moe/impl_environment.py index c87c4107289a..3a289ce31196 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/impl_environment.py +++ b/tensorrt_llm/_torch/moe/fused_moe/impl_environment.py @@ -52,9 +52,6 @@ class MoEDep(str, Enum): class MoEEnvFlag(str, Enum): """Environment variables that MoE selection is allowed to read.""" - #: Opt-in to the FlashInfer provider for quantized TRTLLM-Gen. Changes the - #: routing split, which is why load-balancer eligibility depends on it. - TRTLLM_GEN_USE_FLASHINFER = "TRTLLM_GEN_FUSED_MOE_USE_FLASHINFER" #: Opt-in to folding the shared experts into the routed grouped GEMM. #: Collected rather than read from ``os.environ`` because it changes how #: many expert slots the GEMM builds: two ranks disagreeing on it would @@ -143,7 +140,6 @@ def _probe_megamoe_cutedsl_op() -> Tuple[bool, str]: # Preserve prior defaults when environment variables are unset. _ENV_FLAG_DEFAULTS: Dict[MoEEnvFlag, str] = { - MoEEnvFlag.TRTLLM_GEN_USE_FLASHINFER: "0", MoEEnvFlag.SHARED_EXPERT_FUSION: "0", } diff --git a/tensorrt_llm/_torch/moe/fused_moe/moe_op_backend.py b/tensorrt_llm/_torch/moe/fused_moe/moe_op_backend.py index 8040c523e9a3..e7df5844a674 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/moe_op_backend.py +++ b/tensorrt_llm/_torch/moe/fused_moe/moe_op_backend.py @@ -15,12 +15,12 @@ """Per-provider kernel entry points for the TRTLLM-Gen MoE leaves. -The two implementations differ in real ways -- FlashInfer translates activation -and routing enums and probes which module a given release keeps them in, the -native one validates SiTu inputs and picks between the NVFP4 and MXFP4 ops -- -so the argument marshalling in ``trtllm_gen/fp4_block_scale.py`` and -``trtllm_gen/fp8_block_scale.py`` can be written once per weight format and -shared by that format's native and FlashInfer leaves. +Every quantized leaf runs on :class:`TRTLLMOpBackend`, the in-tree ``trtllm`` +ops, which validates SiTu inputs and picks between the NVFP4 and MXFP4 runners. +The unquantized leaf runs on :class:`FlashinferOpBackend`, whose only runner is +``run_bf16_moe``: TRT-LLM's own cubins have no bf16 x bf16 kernel. A method a +provider does not implement raises ``NotImplementedError`` from +:class:`MoEOpBackend`. Follow-up (TRTLLM-16387): this is a vtable, not a factory. Provider choice now happens in ``can_implement`` plus ``IMPL_PRIORITY``, so @@ -552,20 +552,18 @@ def run_bf16_moe( tune_max_num_tokens=8192, ): raise NotImplementedError( - "TRTLLM native op backend does not support unquantized BF16 TRTLLM-Gen fused MoE. " - "Enable FlashInfer fused MoE for TRTLLM backend." + "TRTLLM native op backend does not support unquantized BF16 TRTLLM-Gen fused MoE; " + "the unquantized format runs on the FlashInfer provider." ) # ==================== Flashinfer Backend ==================== @register_op_backend("flashinfer") class FlashinferOpBackend(MoEOpBackend): - """Flashinfer op backend implementation.""" + """Flashinfer op backend implementation, for the unquantized BF16 runner.""" def __init__(self): import flashinfer.fused_moe as _flashinfer_fused_moe - from flashinfer.fp4_quantization import fp4_quantize as _flashinfer_fp4_quantize - from flashinfer.fp8_quantization import mxfp8_quantize as _flashinfer_mxfp8_quantize # Flashinfer reorganised these enums across releases: # - 0.6.9 release (main): RoutingMethodType in flashinfer.fused_moe, @@ -587,8 +585,6 @@ def __init__(self): self._activation_type = _flashinfer_activation_type self._routing_method_type = _flashinfer_routing_method_type self._fused_moe = _flashinfer_fused_moe - self._fp4_quantize = _flashinfer_fp4_quantize - self._mxfp8_quantize = _flashinfer_mxfp8_quantize # need to add this to the flashinfer side os.environ["FLASHINFER_EXTRA_LDFLAGS"] = "-Wl,-Bsymbolic-functions" @@ -623,255 +619,7 @@ def cvt_routing_method_type(self, routing_method_type) -> int: raise ValueError(f"Unsupported routing method type: {routing_method_type}") return int(_mapping[routing_method_type]) - # Quantization - def fp4_quantize( - self, - input: torch.Tensor, - global_scale: Optional[torch.Tensor] = None, - sf_vec_size: int = 16, - sf_use_ue8m0: bool = False, - is_sf_swizzled_layout: bool = True, - is_sf_8x4_layout: bool = False, - enable_pdl: Optional[bool] = None, - ): - # Keyword binding: the provider inserts new options (e.g. - # ``is_global_scale_inversed``) ahead of ``enable_pdl`` in the - # positional signature, so positional calls misbind across releases. - return self._fp4_quantize( - input, - global_scale=global_scale, - sf_vec_size=sf_vec_size, - sf_use_ue8m0=sf_use_ue8m0, - is_sf_swizzled_layout=is_sf_swizzled_layout, - is_sf_8x4_layout=is_sf_8x4_layout, - enable_pdl=enable_pdl, - ) - - def mxfp8_quantize( - self, - input: torch.Tensor, - is_sf_swizzled_layout: bool = True, - alignment: int = 32, - enable_pdl: Optional[bool] = None, - ): - return self._mxfp8_quantize( - input, is_sf_swizzled_layout, alignment=alignment, enable_pdl=enable_pdl - ) - # MoE Runners - def run_fp8_block_scale_moe( - self, - router_logits, - routing_bias, - hidden_states, - hidden_states_scale, - gemm1_weights, - gemm1_weights_scale, - gemm2_weights, - gemm2_weights_scale, - num_experts, - top_k, - num_fused_shared_experts, - n_group, - topk_group, - intermediate_size, - local_expert_offset, - local_num_experts, - routed_scaling_factor, - routing_method_type, - topk_weights=None, - topk_ids=None, - gated_act_type=0, - gemm1_clamp_limit=None, - output=None, - use_shuffled_weight=False, - weight_layout=0, - enable_pdl=None, - tune_max_num_tokens=8192, - use_dp=False, - ): - assert not num_fused_shared_experts, ( - "Flashinfer backend does not support fusing shared experts" - ) - if gemm1_clamp_limit is not None: - raise NotImplementedError( - "FlashinferOpBackend.run_fp8_block_scale_moe does not yet " - "forward gemm1_clamp_limit; use the trtllm op backend." - ) - if router_logits is not None: - return self._fused_moe.trtllm_fp8_block_scale_moe( - router_logits, - routing_bias, - hidden_states, - hidden_states_scale, - gemm1_weights, - gemm1_weights_scale, - gemm2_weights, - gemm2_weights_scale, - num_experts, - top_k, - n_group, - topk_group, - intermediate_size, - local_expert_offset, - local_num_experts, - routed_scaling_factor, - self.cvt_routing_method_type(routing_method_type), - use_shuffled_weight=use_shuffled_weight, - weight_layout=weight_layout, - enable_pdl=enable_pdl, - tune_max_num_tokens=tune_max_num_tokens, - ) - else: - packed_topk_ids = (topk_ids << 16) | topk_weights.view(torch.int16).to(torch.int32) - # Run with pre-computed routing (packed format) - return self._fused_moe.trtllm_fp8_block_scale_routed_moe( - topk_ids=packed_topk_ids, - routing_bias=routing_bias, - hidden_states=hidden_states, - hidden_states_scale=hidden_states_scale, - gemm1_weights=gemm1_weights, - gemm1_weights_scale=gemm1_weights_scale, - gemm2_weights=gemm2_weights, - gemm2_weights_scale=gemm2_weights_scale, - num_experts=num_experts, - top_k=top_k, - n_group=n_group, - topk_group=topk_group, - intermediate_size=intermediate_size, - local_expert_offset=local_expert_offset, - local_num_experts=num_experts, - routed_scaling_factor=routed_scaling_factor, - routing_method_type=self.cvt_routing_method_type(routing_method_type), - use_shuffled_weight=use_shuffled_weight, - weight_layout=weight_layout, - enable_pdl=enable_pdl, - output=output, - tune_max_num_tokens=tune_max_num_tokens, - ) - - def run_fp4_block_scale_moe( - self, - router_logits, - routing_bias, - hidden_states, - hidden_states_scale, - gemm1_weights, - gemm1_weights_scale, - gemm1_bias, - gemm1_alpha, - gemm1_beta, - gemm1_clamp_limit, - gemm2_weights, - gemm2_weights_scale, - gemm2_bias, - output1_scale_scalar, - output1_scale_gate_scalar, - output2_scale_scalar, - num_experts, - top_k, - n_group, - topk_group, - intermediate_size, - local_expert_offset, - local_num_experts, - routed_scaling_factor, - routing_method_type, - do_finalize=True, - topk_weights=None, - topk_ids=None, - valid_hidden_size=None, - valid_intermediate_size=None, - enable_pdl=None, - gated_act_type=0, - output=None, - tune_max_num_tokens=8192, - use_dp=False, - ): - # The framework hands the scales over token-major 2-D as - # ``quantize_input`` produced them, which is exactly what the - # FlashInfer API asserts (``shape[0] == num_tokens``). No reshape. - if router_logits is not None: - outputs = self._fused_moe.trtllm_fp4_block_scale_moe( - router_logits, - routing_bias, - hidden_states, - hidden_states_scale.view(torch.float8_e4m3fn) - if hidden_states_scale is not None - else None, - gemm1_weights, - gemm1_weights_scale.view(torch.float8_e4m3fn), - gemm1_bias, - gemm1_alpha, - gemm1_beta, - gemm1_clamp_limit, - gemm2_weights, - gemm2_weights_scale.view(torch.float8_e4m3fn), - gemm2_bias, - output1_scale_scalar, - output1_scale_gate_scalar, - output2_scale_scalar, - num_experts, - top_k, - n_group, - topk_group, - intermediate_size, - local_expert_offset, - local_num_experts, - routed_scaling_factor, - self.cvt_routing_method_type(routing_method_type), - do_finalize=do_finalize, - enable_pdl=enable_pdl, - activation_type=self.cvt_activation_type(gated_act_type), - output=output, - tune_max_num_tokens=tune_max_num_tokens, - ) - else: - packed_tensor = (topk_ids.to(torch.int32) << 16) | topk_weights.to(torch.bfloat16).view( - torch.int16 - ) - outputs = self._fused_moe.trtllm_fp4_block_scale_routed_moe( - packed_tensor, - routing_bias, - hidden_states, - hidden_states_scale.view(torch.float8_e4m3fn) - if hidden_states_scale is not None - else None, - gemm1_weights, - gemm1_weights_scale.view(torch.float8_e4m3fn), - gemm1_bias, - gemm1_alpha, - gemm1_beta, - gemm1_clamp_limit, - gemm2_weights, - gemm2_weights_scale.view(torch.float8_e4m3fn), - gemm2_bias, - output1_scale_scalar, - output1_scale_gate_scalar, - output2_scale_scalar, - num_experts, - top_k, - n_group, - topk_group, - intermediate_size, - local_expert_offset, - local_num_experts, - routed_scaling_factor, - self.cvt_routing_method_type(routing_method_type), - do_finalize=do_finalize, - enable_pdl=enable_pdl, - activation_type=self.cvt_activation_type(gated_act_type), - output=output, - tune_max_num_tokens=tune_max_num_tokens, - ) - if not do_finalize: - if outputs[2].dim() != 2: - outputs[2] = outputs[2].view(-1, top_k) - return outputs - else: - final_hidden_states = outputs[0] - return final_hidden_states - def run_bf16_moe( self, router_logits, diff --git a/tensorrt_llm/_torch/moe/fused_moe/moe_resolution.py b/tensorrt_llm/_torch/moe/fused_moe/moe_resolution.py index 932bb57cb130..da4ff4225ddd 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/moe_resolution.py +++ b/tensorrt_llm/_torch/moe/fused_moe/moe_resolution.py @@ -61,10 +61,6 @@ from .moe_load_balancer import get_moe_load_balancer from .trtllm_gen import ( FlashinferTrtllmGenBf16Impl, - FlashinferTrtllmGenFp8BlockScalesImpl, - FlashinferTrtllmGenNvfp4Impl, - FlashinferTrtllmGenW4a8Mxfp4Mxfp8Impl, - FlashinferTrtllmGenW4a16Mxfp4Impl, TrtllmTrtllmGenFp8BlockScalesImpl, TrtllmTrtllmGenNvfp4Impl, TrtllmTrtllmGenW4a8Mxfp4Fp8Impl, @@ -107,18 +103,9 @@ TrtllmCutedslGroupedGemmNvfp4Impl, TrtllmCutedslGroupedGemmBf16Impl, TrtllmCutedslFusedFc12Nvfp4Impl, - # The TRTLLM-Gen leaves. FlashInfer sits ahead of the native leaf - # for the same format because the opt-in flag is what selects it: with the - # flag unset every FlashInfer leaf rejects in ``check_flashinfer_provider`` - # and resolution walks on to the native one below. - # - # Order within a provider does not matter: the ``quant`` segments are - # disjoint, so at most one leaf per provider can admit a given problem. - FlashinferTrtllmGenNvfp4Impl, - FlashinferTrtllmGenFp8BlockScalesImpl, - FlashinferTrtllmGenW4a16Mxfp4Impl, - FlashinferTrtllmGenW4a8Mxfp4Mxfp8Impl, - FlashinferTrtllmGenBf16Impl, # no native counterpart; flag-independent + # The TRTLLM-Gen leaves. Order among them does not matter: the ``quant`` + # segments are disjoint, so at most one leaf can admit a given problem. + FlashinferTrtllmGenBf16Impl, # the unquantized format; no native counterpart TrtllmTrtllmGenNvfp4Impl, TrtllmTrtllmGenFp8BlockScalesImpl, TrtllmTrtllmGenW4a16Mxfp4Impl, @@ -159,10 +146,6 @@ # which. A pinned ``impl_id`` names exactly one of them. "TRTLLM": frozenset( { - FlashinferTrtllmGenNvfp4Impl, - FlashinferTrtllmGenFp8BlockScalesImpl, - FlashinferTrtllmGenW4a16Mxfp4Impl, - FlashinferTrtllmGenW4a8Mxfp4Mxfp8Impl, FlashinferTrtllmGenBf16Impl, TrtllmTrtllmGenNvfp4Impl, TrtllmTrtllmGenFp8BlockScalesImpl, diff --git a/tensorrt_llm/_torch/moe/fused_moe/moe_scheduler.py b/tensorrt_llm/_torch/moe/fused_moe/moe_scheduler.py index 33b9aaf14488..c531c4d8449d 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/moe_scheduler.py +++ b/tensorrt_llm/_torch/moe/fused_moe/moe_scheduler.py @@ -54,7 +54,6 @@ from tensorrt_llm.tools.layer_wise_benchmarks import get_calibrator from .communication import DeepEP, DeepEPLowLatency, NcclEP, NVLinkOneSided, NVLinkTwoSided -from .communication.nvlink_two_sided_flashinfer import NVLinkTwoSidedFlashinfer from .fused_moe_cutlass import raise_moe_lora_multichunk_unsupported from .impl_contract import MoECommPlan, MoERunContext from .interface import FORCE_SEPARATED_ROUTING, MoESchedulerKind @@ -252,7 +251,7 @@ def forward( # NVLink-specific EPLB stat-gather paths) # ------------------------------------------------------------------ def _is_using_nvlink_two_sided(self) -> bool: - return isinstance(self.moe.comm, (NVLinkTwoSided, NVLinkTwoSidedFlashinfer)) + return isinstance(self.moe.comm, NVLinkTwoSided) def _is_using_nvlink_one_sided(self) -> bool: return isinstance(self.moe.comm, NVLinkOneSided) diff --git a/tensorrt_llm/_torch/moe/fused_moe/trtllm_gen/__init__.py b/tensorrt_llm/_torch/moe/fused_moe/trtllm_gen/__init__.py index b13232c52222..5d59ae95a5af 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/trtllm_gen/__init__.py +++ b/tensorrt_llm/_torch/moe/fused_moe/trtllm_gen/__init__.py @@ -32,9 +32,8 @@ what registers them, which is why ``trtllm_gen_leaf`` is defined at the bottom of this file: the lookup is only correct once every leaf module has loaded. -NVFP4, FP8_BLOCK_SCALES, W4A16_MXFP4 and W4A8_MXFP4_MXFP8 exist on both -providers; W4A8_NVFP4_FP8 and W4A8_MXFP4_FP8 are native-only, and the -unquantized bf16 path is FlashInfer-only. +Each format has exactly one leaf: every quantized format runs on the native +``trtllm`` provider, and the unquantized bf16 path on FlashInfer. """ from tensorrt_llm.models.modeling_utils import QuantAlgo @@ -43,20 +42,13 @@ from ..impl_identity import MOE_IMPL_REGISTRY, MoEImplId from .base import TrtllmGenFusedMoEBase from .eligibility import ( - check_flashinfer_provider, check_flashinfer_shard_alignment, - check_mxfp4_flashinfer_shape, check_no_activation_constants, check_no_expert_bias, check_trtllm_gen_capabilities, check_trtllm_gen_leaf, - nvfp4_needs_padded_method, ) from .flashinfer_bf16 import FlashinferTrtllmGenBf16Impl -from .flashinfer_fp8_block_scales import FlashinferTrtllmGenFp8BlockScalesImpl -from .flashinfer_nvfp4 import FlashinferTrtllmGenNvfp4Impl -from .flashinfer_w4a8_mxfp4_mxfp8 import FlashinferTrtllmGenW4a8Mxfp4Mxfp8Impl -from .flashinfer_w4a16_mxfp4 import FlashinferTrtllmGenW4a16Mxfp4Impl from .fp4_block_scale import ( TRTLLMGenFp4BlockScaleBase, TRTLLMGenNvfp4Base, @@ -81,21 +73,18 @@ from .trtllm_w4a16_mxfp4 import TrtllmTrtllmGenW4a16Mxfp4Impl -def trtllm_gen_leaf(quant_algo: QuantAlgo | None, *, provider: str | None = None) -> type: - """The leaf implementing ``quant_algo``, on ``provider`` when given. +def trtllm_gen_leaf(quant_algo: QuantAlgo | None) -> type: + """The leaf implementing ``quant_algo``. For callers that know the format and want the class rather than a verdict. Resolution does not use this -- it walks ``IMPL_PRIORITY`` and asks ``can_implement`` -- so this stays a lookup and never becomes a second selection path. - With ``provider`` left open the native leaf wins where there is one, which - matches a deployment without the FlashInfer opt-in flag. - Raises when no leaf publishes the format. A caller enumerating formats wants ``find_trtllm_gen_leaf``, where absence is an ordinary answer. """ - cls = find_trtllm_gen_leaf(quant_algo, provider=provider) + cls = find_trtllm_gen_leaf(quant_algo) if cls is not None: return cls registered = sorted( @@ -105,60 +94,30 @@ def trtllm_gen_leaf(quant_algo: QuantAlgo | None, *, provider: str | None = None ) raise ValueError( f"no TRTLLM-Gen implementation for quant=" - f"{normalize_quant(canonical_quant(quant_algo))} on " - f"provider={'|'.join(_providers_to_try(provider))}; registered: {registered}" + f"{normalize_quant(canonical_quant(quant_algo))}; registered: {registered}" ) -def _providers_to_try(provider: str | None) -> tuple: - """The providers a lookup walks, in the order it walks them.""" - return (provider,) if provider is not None else (PROVIDER_TRTLLM, PROVIDER_FLASHINFER) - - -def find_trtllm_gen_leaf( - quant_algo: QuantAlgo | None, *, provider: str | None = None -) -> type | None: +def find_trtllm_gen_leaf(quant_algo: QuantAlgo | None) -> type | None: """The leaf implementing ``quant_algo``, or ``None`` if none publishes it. Goes through the registry rather than a table of its own, so a leaf that is - renamed or unregistered disappears from here too. + renamed or unregistered disappears from here too. The providers publish + disjoint formats, so at most one of them answers. """ # The same two steps resolution takes (``moe_resolution`` canonicalizes, # ``MoEImplId`` folds case), so a caller naming NVFP4_AWQ or MIXED_PRECISION # here reaches the leaf resolution would pick rather than a lookup miss. quant = normalize_quant(canonical_quant(quant_algo)) - for candidate in _providers_to_try(provider): + for provider in (PROVIDER_TRTLLM, PROVIDER_FLASHINFER): cls = MOE_IMPL_REGISTRY.lookup( - MoEImplId(candidate, TECHNIQUE_TRTLLM_GEN, KERNEL_FUSED_MOE, quant) + MoEImplId(provider, TECHNIQUE_TRTLLM_GEN, KERNEL_FUSED_MOE, quant) ) if cls is not None: return cls return None -def trtllm_gen_leaves_in_resolution_order(quant_algo: QuantAlgo | None) -> tuple[type, ...]: - """The leaves a resolution walk would try for a format, in its order. - - For a caller asking "would a run serve this?" rather than "which class - implements this format?". ``IMPL_PRIORITY`` ranks each FlashInfer leaf - ahead of its native sibling and the opt-in flag is what makes the - FlashInfer ones reject, so resolution falls through to the native leaf on a - *rejection* and not only on an absent leaf. A caller that picks one class - and asks it cannot reproduce that, whichever provider it picks first. - - Shorter than two entries where only one provider publishes the format. - """ - quant = normalize_quant(canonical_quant(quant_algo)) - leaves = [] - for provider in (PROVIDER_FLASHINFER, PROVIDER_TRTLLM): - cls = MOE_IMPL_REGISTRY.lookup( - MoEImplId(provider, TECHNIQUE_TRTLLM_GEN, KERNEL_FUSED_MOE, quant) - ) - if cls is not None: - leaves.append(cls) - return tuple(leaves) - - __all__ = [ # trtllm provider "TrtllmTrtllmGenNvfp4Impl", @@ -168,10 +127,6 @@ def trtllm_gen_leaves_in_resolution_order(quant_algo: QuantAlgo | None) -> tuple "TrtllmTrtllmGenW4a8Nvfp4Fp8Impl", "TrtllmTrtllmGenW4a8Mxfp4Fp8Impl", # flashinfer provider - "FlashinferTrtllmGenNvfp4Impl", - "FlashinferTrtllmGenFp8BlockScalesImpl", - "FlashinferTrtllmGenW4a16Mxfp4Impl", - "FlashinferTrtllmGenW4a8Mxfp4Mxfp8Impl", "FlashinferTrtllmGenBf16Impl", # shared layers "TrtllmGenFusedMoEBase", @@ -189,15 +144,11 @@ def trtllm_gen_leaves_in_resolution_order(quant_algo: QuantAlgo | None) -> tuple "TRTLLM_GEN_INPUT_REQUIREMENT", "trtllm_gen_leaf", "find_trtllm_gen_leaf", - "trtllm_gen_leaves_in_resolution_order", "trtllm_gen_descriptor", "check_trtllm_gen_leaf", "check_trtllm_gen_capabilities", "check_quant_matches_identity", - "check_flashinfer_provider", "check_flashinfer_shard_alignment", - "check_mxfp4_flashinfer_shape", "check_no_activation_constants", "check_no_expert_bias", - "nvfp4_needs_padded_method", ] diff --git a/tensorrt_llm/_torch/moe/fused_moe/trtllm_gen/base.py b/tensorrt_llm/_torch/moe/fused_moe/trtllm_gen/base.py index 2a5a63f695a5..94db2af5351d 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/trtllm_gen/base.py +++ b/tensorrt_llm/_torch/moe/fused_moe/trtllm_gen/base.py @@ -67,10 +67,11 @@ class TrtllmGenFusedMoEBase(MoEImplBase): - ``quant`` decides which kernel ``run_moe`` calls, how weights and inputs are prepared, and which activation ABI applies. Held by a per-quant class in the module named for that kernel ABI -- :mod:`.fp4_block_scale` or - :mod:`.fp8_block_scale` -- one per format that two providers share. A - format only one provider serves has no such class. - - ``provider`` decides which op backend executes, which configurations are - eligible, and whether the kernel writes into the caller's workspace. Read + :mod:`.fp8_block_scale` -- for the formats that share a runner op. A + format on a runner op of its own has no such class. + - ``provider`` decides which op backend executes and whether the kernel + writes into the caller's workspace. Every quantized format runs on the + native ``trtllm`` provider and the unquantized one on FlashInfer. Read off the leaf's own descriptor by ``__init_subclass__`` below, so a leaf is just an identity and a ``can_implement``. @@ -98,16 +99,15 @@ class TrtllmGenFusedMoEBase(MoEImplBase): #: ``__init_subclass__``, not declared per leaf. provider: ClassVar[str] #: ``provider == PROVIDER_FLASHINFER``, read from outside as a plain - #: attribute by ``ConfigurableMoE`` and the backend tests. Derived with - #: ``provider``. + #: attribute by the backend tests. Derived with ``provider``. use_flashinfer: ClassVar[bool] #: Whether this format has a fused cubin for the gpt-oss SwiGLU package #: (expert bias plus alpha/beta). Read off the *class* at resolution time, #: where there is no instance to ask. supports_gptoss_style: ClassVar[bool] = False #: Whether this leaf can run SiTu. Per-leaf and not per-format: the format - #: needs a fused SiTu FC1 cubin *and* the provider has to call it, so a - #: FlashInfer sibling of a supported format may still be False. + #: needs a fused SiTu FC1 cubin *and* the provider has to call it, and only + #: the native op backend does. supports_situ: ClassVar[bool] = False #: Whether this format's cubins always read an FC bias, so a model without #: one still needs a zeroed buffer. Declarative rather than an override diff --git a/tensorrt_llm/_torch/moe/fused_moe/trtllm_gen/eligibility.py b/tensorrt_llm/_torch/moe/fused_moe/trtllm_gen/eligibility.py index ba13e3e32b37..8552dda05541 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/trtllm_gen/eligibility.py +++ b/tensorrt_llm/_torch/moe/fused_moe/trtllm_gen/eligibility.py @@ -37,7 +37,6 @@ MoERejectReason, check_quant_matches_identity, ) -from ..impl_environment import MoEDep, MoEEnvFlag from ..interface import _reject @@ -97,66 +96,17 @@ def check_trtllm_gen_capabilities( return None -def check_flashinfer_provider(cls: type, p: MoEProblem, d: MoEDeployment) -> MoEEligibility | None: - """Gates that separate the FlashInfer provider from the ``trtllm`` one. - - Only the quantized leaves come through here. The unquantized one is - FlashInfer-exclusive and reached without the opt-in flag, so it states its - own dependency gate instead. - """ - if not d.env.has_dep(MoEDep.FLASHINFER): - return _reject(MoERejectReason.DEP_MISSING, f"{cls.__name__} requires the FlashInfer wheel") - - # Opt-in, not a capability: the trtllm provider serves these formats too, - # so routing traffic here without being asked would change which kernel a - # previously-working deployment runs. - if d.env.env_flag(MoEEnvFlag.TRTLLM_GEN_USE_FLASHINFER) != "1": - return _reject( - MoERejectReason.PATH_NOT_ENABLED, - f"{cls.__name__} is opt-in; set " - f"{MoEEnvFlag.TRTLLM_GEN_USE_FLASHINFER.value}=1 to select the " - f"FlashInfer provider for quantized TRTLLM-Gen", - ) - - # SiTu is a native TRTLLM-Gen cubin and is absent from FlashInfer's - # activation enum; Relu2 has no FlashInfer path either. - if p.activation_type in (ActivationType.SiTu, ActivationType.Relu2): - return _reject( - MoERejectReason.ACTIVATION_UNSUPPORTED, - f"{cls.__name__} does not implement {p.activation} " - f"(FlashInfer's activation enum has no such kernel)", - ) - - # These two fuse routing in a form only the native runner accepts. Matched - # on the algorithm the problem carries, not on the value handed to the C++ - # kernels: a method no kernel implements borrows a recognized encoding - # there, and matching on the borrowed one would reject a method this gate - # has nothing against. - if p.routing_method_type is not None and p.routing_method_type.name in ( - "DeepSeekV3", - "Default", - ): - return _reject( - MoERejectReason.ROUTING_UNSUPPORTED, - f"{cls.__name__} does not implement {p.routing} routing", - ) - - return None - - def check_flashinfer_shard_alignment( cls: type, p: MoEProblem, d: MoEDeployment, *, weight_alignment: int, - input_hidden_alignment: int | None = None, ) -> MoEEligibility | None: - """Per-rank shard alignment the FlashInfer kernels need, or ``None``. + """Per-rank intermediate-shard alignment the FlashInfer bf16 kernels need. - Two distinct spellings of "does not apply": an absent shape on ``p`` means - the caller did not say, so abstain; an ``input_hidden_alignment`` of - ``None`` means the kernels constrain the hidden dimension not at all. + ``None`` admits. An absent ``intermediate_size`` on ``p`` means the caller + did not say, so the gate abstains rather than rejects. """ if p.intermediate_size is not None: inter = p.intermediate_size @@ -175,17 +125,6 @@ def check_flashinfer_shard_alignment( f"({inter}) to be a multiple of {weight_alignment}", ) - if ( - input_hidden_alignment is not None - and p.hidden_size is not None - and p.hidden_size % input_hidden_alignment != 0 - ): - return _reject( - MoERejectReason.SHAPE_UNALIGNED, - f"{cls.__name__} requires hidden_size ({p.hidden_size}) to be a " - f"multiple of {input_hidden_alignment}", - ) - return None @@ -198,18 +137,12 @@ def check_no_expert_bias(cls: type, p: MoEProblem) -> MoEEligibility | None: ``capabilities.supports_expert_bias`` is declared once for the whole family, and ``create_moe`` reads it against the class resolution already picked -- so a leaf that answered only there would be chosen and then - refuse, instead of standing aside for a sibling that can serve. + refuse, instead of standing aside for a backend that can serve. A gate rather than only the ``_check_configs`` assert for the same reason: - some leaves reaching here have a sibling on the other provider that does - take a bias, and standing aside during selection is what lets it be - tried; for the rest it is the difference between a structured verdict - and an assert after resolution has committed. - - Its own function rather than a line inside - :func:`check_mxfp4_flashinfer_shape`, because a leaf whose shape gate is - conditional still has to answer about bias unconditionally, and one copy - of the rule keeps every leaf giving the same reason. + standing aside during selection is what lets the fallback be tried, and it + is the difference between a structured verdict and an assert after + resolution has committed. """ if p.bias: return _reject( @@ -241,61 +174,20 @@ def check_no_activation_constants(cls: type, p: MoEProblem) -> MoEEligibility | return None -def check_mxfp4_flashinfer_shape( - cls: type, - p: MoEProblem, - d: MoEDeployment, - *, - weight_alignment: int, - input_hidden_alignment: int, -) -> MoEEligibility | None: - """The quantized FlashInfer leaves' shape gate: no bias, then alignment. - - Bundled rather than folded into - :func:`check_flashinfer_shard_alignment`, because the unquantized leaf - shares the alignment rule but answers about bias through - ``swiglu_gptoss_style``. - """ - bias_verdict = check_no_expert_bias(cls, p) - if bias_verdict is not None: - return bias_verdict - return check_flashinfer_shard_alignment( - cls, - p, - d, - weight_alignment=weight_alignment, - input_hidden_alignment=input_hidden_alignment, - ) - - -def nvfp4_needs_padded_method(activation_type: ActivationType, has_alpha_constant: bool) -> bool: - """Whether NVFP4 needs the padded quant method rather than the base one. - - One definition read from two sides -- ``can_implement`` asks it of the - problem, ``_get_quant_method`` of the instance -- so that a configuration - cannot be admitted by one and laid out by the other. - """ - return has_alpha_constant or activation_type in ( - ActivationType.SiTu, - ActivationType.Relu2, - ActivationType.Silu, - ) - - def check_trtllm_gen_leaf( - cls: type, p: MoEProblem, d: MoEDeployment, *provider_gates: MoEEligibility | None + cls: type, p: MoEProblem, d: MoEDeployment, *leaf_gates: MoEEligibility | None ) -> MoEEligibility: - """Compose one leaf's verdict: identity, shared capability, then provider. + """Compose one leaf's verdict: identity, shared capability, then its own. - First rejection wins, in a fixed order so that two leaves differing only in - provider give the same reason for a problem neither can serve. - ``provider_gates`` are already-evaluated verdicts, not callables: the gates + First rejection wins, in a fixed order so that every leaf gives the + family-wide reason before its own for a problem neither can serve. + ``leaf_gates`` are already-evaluated verdicts, not callables: the gates are pure, so there is nothing to defer. """ for verdict in ( check_quant_matches_identity(cls, p), check_trtllm_gen_capabilities(cls, p, d), - *provider_gates, + *leaf_gates, ): if verdict is not None: return verdict diff --git a/tensorrt_llm/_torch/moe/fused_moe/trtllm_gen/flashinfer_bf16.py b/tensorrt_llm/_torch/moe/fused_moe/trtllm_gen/flashinfer_bf16.py index 585c2113ff51..cf5f76a4c40c 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/trtllm_gen/flashinfer_bf16.py +++ b/tensorrt_llm/_torch/moe/fused_moe/trtllm_gen/flashinfer_bf16.py @@ -45,9 +45,8 @@ class FlashinferTrtllmGenBf16Impl(TrtllmGenFusedMoEBase): """``flashinfer.trtllm_gen.fused_moe.none`` -- the unquantized path. - FlashInfer-exclusive, and the one leaf reached without the opt-in flag: - ``trtllm_bf16_moe`` has no counterpart in TRT-LLM's own cubins, so there is - no native leaf to prefer and nothing for a flag to switch between. It is + The only FlashInfer leaf: ``trtllm_bf16_moe`` has no counterpart in + TRT-LLM's own cubins, so there is no native leaf for this format. It is also the only leaf whose kernel does not route internally, which is what ``_requires_separated_routing`` below says. """ @@ -132,9 +131,8 @@ def _check_bf16_path(cls, p: MoEProblem, d: MoEDeployment) -> MoEEligibility | N MoERejectReason.DEP_MISSING, f"{cls.__name__} requires FlashInfer fused MoE with trtllm_bf16_moe support.", ) - # The same per-rank shard rule the quantized FlashInfer leaves run. - # These kernels constrain only the intermediate dimension, which is - # what passing no ``input_hidden_alignment`` says. + # ``trtllm_bf16_moe`` constrains only the intermediate dimension: each + # rank's shard has to be a multiple of 128. return check_flashinfer_shard_alignment(cls, p, d, weight_alignment=128) def _get_quant_method(self) -> object: diff --git a/tensorrt_llm/_torch/moe/fused_moe/trtllm_gen/flashinfer_fp8_block_scales.py b/tensorrt_llm/_torch/moe/fused_moe/trtllm_gen/flashinfer_fp8_block_scales.py deleted file mode 100644 index 6ac43b02f152..000000000000 --- a/tensorrt_llm/_torch/moe/fused_moe/trtllm_gen/flashinfer_fp8_block_scales.py +++ /dev/null @@ -1,42 +0,0 @@ -# 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. -"""``flashinfer.trtllm_gen.fused_moe.fp8_block_scales``.""" - -from ..impl_contract import MoEDeployment, MoEEligibility, MoEProblem -from ..impl_identity import register_moe_impl -from .eligibility import check_flashinfer_provider, check_no_expert_bias, check_trtllm_gen_leaf -from .fp8_block_scale import TRTLLMGenFp8BlockScalesBase -from .identity import PROVIDER_FLASHINFER, trtllm_gen_descriptor - - -@register_moe_impl -class FlashinferTrtllmGenFp8BlockScalesImpl(TRTLLMGenFp8BlockScalesBase): - """``flashinfer.trtllm_gen.fused_moe.fp8_block_scales``. - - No alignment gate: the padding rules belong to the FP4 weight layouts, and - this format has none of them. - """ - - descriptor = trtllm_gen_descriptor( - PROVIDER_FLASHINFER, - "fp8_block_scales", - "FlashInfer's TRTLLM-Gen FP8 block-scale fused MoE, SM100 family.", - ) - - @classmethod - def can_implement(cls, p: MoEProblem, d: MoEDeployment) -> MoEEligibility: - return check_trtllm_gen_leaf( - cls, p, d, check_flashinfer_provider(cls, p, d), check_no_expert_bias(cls, p) - ) diff --git a/tensorrt_llm/_torch/moe/fused_moe/trtllm_gen/flashinfer_nvfp4.py b/tensorrt_llm/_torch/moe/fused_moe/trtllm_gen/flashinfer_nvfp4.py deleted file mode 100644 index d2f6b75e8617..000000000000 --- a/tensorrt_llm/_torch/moe/fused_moe/trtllm_gen/flashinfer_nvfp4.py +++ /dev/null @@ -1,72 +0,0 @@ -# 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. -"""``flashinfer.trtllm_gen.fused_moe.nvfp4``.""" - -from ..impl_contract import MoEDeployment, MoEEligibility, MoEProblem -from ..impl_identity import register_moe_impl -from ..quantization import NVFP4TRTLLMGenFusedMoEMethod -from .eligibility import ( - check_flashinfer_provider, - check_mxfp4_flashinfer_shape, - check_no_expert_bias, - check_trtllm_gen_leaf, - nvfp4_needs_padded_method, -) -from .fp4_block_scale import TRTLLMGenNvfp4Base -from .identity import PROVIDER_FLASHINFER, trtllm_gen_descriptor - - -@register_moe_impl -class FlashinferTrtllmGenNvfp4Impl(TRTLLMGenNvfp4Base): - """``flashinfer.trtllm_gen.fused_moe.nvfp4``.""" - - descriptor = trtllm_gen_descriptor( - PROVIDER_FLASHINFER, "nvfp4", "FlashInfer's TRTLLM-Gen NVFP4 fused MoE, SM100 family." - ) - - @classmethod - def can_implement(cls, p: MoEProblem, d: MoEDeployment) -> MoEEligibility: - return check_trtllm_gen_leaf( - cls, - p, - d, - check_flashinfer_provider(cls, p, d), - # Alongside the shape gate rather than inside it: this leaf's shape - # gate abstains for the unpadded method, and bias is rejected for - # both methods. - check_no_expert_bias(cls, p), - cls._check_shape(p, d), - ) - - @classmethod - def _check_shape(cls, p: MoEProblem, d: MoEDeployment) -> MoEEligibility | None: - """Alignment gate, and only for the padded quant method. - - The unpadded ``NVFP4TRTLLMGenFusedMoEBaseMethod`` lays out any shape, - so it is admitted unconditionally. - - The alignments are the padded method's *class* attributes, not what - ``resolve_alignments`` would pick for this shape. Reading the resolved - values would reject shapes this provider currently serves. - """ - if not nvfp4_needs_padded_method(p.activation_type, "alpha" in p.activation_constants): - return None - return check_mxfp4_flashinfer_shape( - cls, - p, - d, - weight_alignment=NVFP4TRTLLMGenFusedMoEMethod.weight_alignment, - input_hidden_alignment=NVFP4TRTLLMGenFusedMoEMethod.input_hidden_alignment, - ) diff --git a/tensorrt_llm/_torch/moe/fused_moe/trtllm_gen/flashinfer_w4a16_mxfp4.py b/tensorrt_llm/_torch/moe/fused_moe/trtllm_gen/flashinfer_w4a16_mxfp4.py deleted file mode 100644 index fe981b0cf657..000000000000 --- a/tensorrt_llm/_torch/moe/fused_moe/trtllm_gen/flashinfer_w4a16_mxfp4.py +++ /dev/null @@ -1,53 +0,0 @@ -# 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. -"""``flashinfer.trtllm_gen.fused_moe.w4a16_mxfp4``.""" - -from ..impl_contract import MoEDeployment, MoEEligibility, MoEProblem -from ..impl_identity import register_moe_impl -from ..quantization import W4A16MXFP4TRTLLMGenFusedMoEMethod -from .eligibility import ( - check_flashinfer_provider, - check_mxfp4_flashinfer_shape, - check_trtllm_gen_leaf, -) -from .fp4_block_scale import TRTLLMGenW4a16Mxfp4Base -from .identity import PROVIDER_FLASHINFER, trtllm_gen_descriptor - - -@register_moe_impl -class FlashinferTrtllmGenW4a16Mxfp4Impl(TRTLLMGenW4a16Mxfp4Base): - """``flashinfer.trtllm_gen.fused_moe.w4a16_mxfp4``.""" - - descriptor = trtllm_gen_descriptor( - PROVIDER_FLASHINFER, - "w4a16_mxfp4", - "FlashInfer's TRTLLM-Gen MXFP4-weight fused MoE with bf16 activations.", - ) - - @classmethod - def can_implement(cls, p: MoEProblem, d: MoEDeployment) -> MoEEligibility: - return check_trtllm_gen_leaf( - cls, - p, - d, - check_flashinfer_provider(cls, p, d), - check_mxfp4_flashinfer_shape( - cls, - p, - d, - weight_alignment=W4A16MXFP4TRTLLMGenFusedMoEMethod.weight_alignment, - input_hidden_alignment=W4A16MXFP4TRTLLMGenFusedMoEMethod.input_hidden_alignment, - ), - ) diff --git a/tensorrt_llm/_torch/moe/fused_moe/trtllm_gen/flashinfer_w4a8_mxfp4_mxfp8.py b/tensorrt_llm/_torch/moe/fused_moe/trtllm_gen/flashinfer_w4a8_mxfp4_mxfp8.py deleted file mode 100644 index 1b1593edabe6..000000000000 --- a/tensorrt_llm/_torch/moe/fused_moe/trtllm_gen/flashinfer_w4a8_mxfp4_mxfp8.py +++ /dev/null @@ -1,53 +0,0 @@ -# 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. -"""``flashinfer.trtllm_gen.fused_moe.w4a8_mxfp4_mxfp8``.""" - -from ..impl_contract import MoEDeployment, MoEEligibility, MoEProblem -from ..impl_identity import register_moe_impl -from ..quantization import W4A8MXFP4MXFP8TRTLLMGenFusedMoEMethod -from .eligibility import ( - check_flashinfer_provider, - check_mxfp4_flashinfer_shape, - check_trtllm_gen_leaf, -) -from .fp4_block_scale import TRTLLMGenW4a8Mxfp4Mxfp8Base -from .identity import PROVIDER_FLASHINFER, trtllm_gen_descriptor - - -@register_moe_impl -class FlashinferTrtllmGenW4a8Mxfp4Mxfp8Impl(TRTLLMGenW4a8Mxfp4Mxfp8Base): - """``flashinfer.trtllm_gen.fused_moe.w4a8_mxfp4_mxfp8``.""" - - descriptor = trtllm_gen_descriptor( - PROVIDER_FLASHINFER, - "w4a8_mxfp4_mxfp8", - "FlashInfer's TRTLLM-Gen MXFP4-weight fused MoE with MXFP8 activations.", - ) - - @classmethod - def can_implement(cls, p: MoEProblem, d: MoEDeployment) -> MoEEligibility: - return check_trtllm_gen_leaf( - cls, - p, - d, - check_flashinfer_provider(cls, p, d), - check_mxfp4_flashinfer_shape( - cls, - p, - d, - weight_alignment=W4A8MXFP4MXFP8TRTLLMGenFusedMoEMethod.weight_alignment, - input_hidden_alignment=W4A8MXFP4MXFP8TRTLLMGenFusedMoEMethod.input_hidden_alignment, - ), - ) diff --git a/tensorrt_llm/_torch/moe/fused_moe/trtllm_gen/fp4_block_scale.py b/tensorrt_llm/_torch/moe/fused_moe/trtllm_gen/fp4_block_scale.py index 7be6441fcf27..a4d12e422966 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/trtllm_gen/fp4_block_scale.py +++ b/tensorrt_llm/_torch/moe/fused_moe/trtllm_gen/fp4_block_scale.py @@ -17,7 +17,7 @@ :class:`.TRTLLMGenFp4BlockScaleBase` owns the kernel call; NVFP4, W4A16_MXFP4 and W4A8_MXFP4_MXFP8 differ in how weights and inputs are prepared and not at all in how the kernel is called, so each subclass supplies only the -preparation. Two leaves land here per format, one per provider. +preparation. One leaf lands here per format. A format's class also answers whatever the family base would otherwise switch on ``quant_config`` for: the SiTu weight alignment and scale group size are @@ -27,7 +27,12 @@ import torch from torch import nn -from tensorrt_llm._torch.utils import ActType_TrtllmGen, Fp4QuantizedTensor, MxFp8QuantizedTensor +from tensorrt_llm._torch.utils import ( + ActivationType, + ActType_TrtllmGen, + Fp4QuantizedTensor, + MxFp8QuantizedTensor, +) from ..activation import materialize_activation_params, resolve_activation_support from ..impl_contract import MoERunContext @@ -38,7 +43,6 @@ W4A16MXFP4TRTLLMGenFusedMoEMethod, ) from .base import TrtllmGenFusedMoEBase -from .eligibility import nvfp4_needs_padded_method from .kernel_inputs import get_data_or_none, prepare_kernel_inputs, to_trtllm_gen_act_type @@ -198,7 +202,7 @@ class TRTLLMGenNvfp4Base(TRTLLMGenFp4BlockScaleBase): supports_gptoss_style = True def supports_unfinalized_output(self) -> bool: - """Both providers' NVFP4 runners honor ``do_finalize=False``. + """The NVFP4 runner honors ``do_finalize=False``. Declared here rather than on the family base: the MXFP4 siblings call runners without a ``do_finalize`` parameter, which finalize regardless. @@ -206,8 +210,8 @@ def supports_unfinalized_output(self) -> bool: return True # Group size of the ``Bmm_E2m1_E2m1E2m1_..._siTuGlu_*`` cubins. Whether a - # leaf may reach them is ``supports_situ``, declared per leaf because only - # the native op backend calls them. + # leaf may reach them is ``supports_situ``, declared on the leaf because it + # also depends on the op backend calling them. situ_scaling_vector_size = 16 def _situ_tp_weight_alignment(self) -> int: @@ -223,8 +227,10 @@ def _get_quant_method(self) -> object: # SiTu fills the act_alpha/act_beta slots from create_weights, which # runs after this, so keying off the tensor would make the method # depend on *when* it is asked for. - needs_padded_method = nvfp4_needs_padded_method( - self.activation_type, self.act_alpha is not None + needs_padded_method = self.act_alpha is not None or self.activation_type in ( + ActivationType.SiTu, + ActivationType.Relu2, + ActivationType.Silu, ) return ( NVFP4TRTLLMGenFusedMoEMethod() diff --git a/tensorrt_llm/_torch/moe/fused_moe/trtllm_gen/fp8_block_scale.py b/tensorrt_llm/_torch/moe/fused_moe/trtllm_gen/fp8_block_scale.py index fb44b0c0c6c2..557d88835b2e 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/trtllm_gen/fp8_block_scale.py +++ b/tensorrt_llm/_torch/moe/fused_moe/trtllm_gen/fp8_block_scale.py @@ -12,7 +12,7 @@ # 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. -"""DeepSeek-style FP8 with 1x128 block scales, served on both providers. +"""DeepSeek-style FP8 with 1x128 block scales. ``run_fp8_block_scale_moe`` is its own kernel ABI, and this format is also the only one that narrows the family's activation support and the only one whose @@ -81,10 +81,6 @@ def fused_shared_expert_count(cls, model_config: ModelConfig) -> int: """ if collect_moe_environment().env_flag(MoEEnvFlag.SHARED_EXPERT_FUSION) != "1": return 0 - # Only the trtllm op backend implements fused shared experts, so the - # FlashInfer leaf on this same format stays unfused. - if cls.use_flashinfer: - return 0 # Expert parallelism (moe_ep_size > 1) is not supported by the fused # path yet (the routing kernel's shared-expert append assumes the full # expert set is local); gate it out here so EP configs fall back to the diff --git a/tensorrt_llm/_torch/moe/fused_moe/trtllm_gen/trtllm_nvfp4.py b/tensorrt_llm/_torch/moe/fused_moe/trtllm_gen/trtllm_nvfp4.py index cfd15a714bdd..3ada648246f8 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/trtllm_gen/trtllm_nvfp4.py +++ b/tensorrt_llm/_torch/moe/fused_moe/trtllm_gen/trtllm_nvfp4.py @@ -29,9 +29,9 @@ class TrtllmTrtllmGenNvfp4Impl(TRTLLMGenNvfp4Base): PROVIDER_TRTLLM, "nvfp4", "TRTLLM-Gen batched-GEMM cubins over NVFP4, SM100 family." ) - #: The group-16 fused SiTu cubins exist for this format, but only the - #: native op backend calls them, so eligibility is per leaf and not on - #: ``TRTLLMGenNvfp4Base`` (which the FlashInfer leaf shares). + #: The group-16 fused SiTu cubins. Declared on the leaf rather than on + #: ``TRTLLMGenNvfp4Base``: reaching them also takes the native op backend, + #: which is a property of this identity, not of the format. supports_situ = True @classmethod diff --git a/tensorrt_llm/_torch/moe/fused_moe/trtllm_gen/trtllm_w4a8_mxfp4_fp8.py b/tensorrt_llm/_torch/moe/fused_moe/trtllm_gen/trtllm_w4a8_mxfp4_fp8.py index f2d04a5dc3ba..463e88f2ecd0 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/trtllm_gen/trtllm_w4a8_mxfp4_fp8.py +++ b/tensorrt_llm/_torch/moe/fused_moe/trtllm_gen/trtllm_w4a8_mxfp4_fp8.py @@ -29,8 +29,8 @@ class TrtllmTrtllmGenW4a8Mxfp4Fp8Impl(TrtllmGenFusedMoEBase): """``trtllm.trtllm_gen.fused_moe.w4a8_mxfp4_fp8``. - Single-provider like the NVFP4/FP8 leaf, and likewise on its own runner op - rather than ``self.op_backend``. + On its own runner op like the NVFP4/FP8 leaf, and likewise called directly + rather than through ``self.op_backend``. """ descriptor = trtllm_gen_descriptor( diff --git a/tensorrt_llm/_torch/moe/fused_moe/trtllm_gen/trtllm_w4a8_mxfp4_mxfp8.py b/tensorrt_llm/_torch/moe/fused_moe/trtllm_gen/trtllm_w4a8_mxfp4_mxfp8.py index 5870093b3ba6..c0bf844b756b 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/trtllm_gen/trtllm_w4a8_mxfp4_mxfp8.py +++ b/tensorrt_llm/_torch/moe/fused_moe/trtllm_gen/trtllm_w4a8_mxfp4_mxfp8.py @@ -55,7 +55,7 @@ def try_fused_route_quant( """Fuse Kimi K3 no-aux routing and MXFP8 input quantization. The op is MXFP8-activation and native-cubin only, which is exactly this - identity, so no format or provider test is needed here. + identity, so no format test is needed here. Beyond that the op is specialized to the K3 decode shape: it hardcodes 896 experts, top-16, hidden 3584 and at most 64 tokens. The checks diff --git a/tensorrt_llm/_torch/moe/fused_moe/trtllm_gen/trtllm_w4a8_nvfp4_fp8.py b/tensorrt_llm/_torch/moe/fused_moe/trtllm_gen/trtllm_w4a8_nvfp4_fp8.py index 88262536dd06..182cd66b917a 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/trtllm_gen/trtllm_w4a8_nvfp4_fp8.py +++ b/tensorrt_llm/_torch/moe/fused_moe/trtllm_gen/trtllm_w4a8_nvfp4_fp8.py @@ -36,9 +36,9 @@ class TrtllmTrtllmGenW4a8Nvfp4Fp8Impl(TrtllmGenFusedMoEBase): """``trtllm.trtllm_gen.fused_moe.w4a8_nvfp4_fp8``. - Single-provider, so there is no per-quant parent, and it calls its runner - op directly instead of going through ``self.op_backend``: the FlashInfer - wheel has no equivalent. + No other format shares its runner op, so there is no per-quant parent, + and it calls that op directly instead of going through + ``self.op_backend``. """ descriptor = trtllm_gen_descriptor( diff --git a/tests/microbenchmarks/bench_moe/backend.py b/tests/microbenchmarks/bench_moe/backend.py index d16071f5675b..0736a8d2e8a0 100644 --- a/tests/microbenchmarks/bench_moe/backend.py +++ b/tests/microbenchmarks/bench_moe/backend.py @@ -128,8 +128,8 @@ def find_backend_class(backend_type: MoeBackendType, quant_algo=None): rather than restated. Names one class. A caller deciding whether a configuration is *runnable* - wants :func:`find_backend_classes` instead, because for ``TRTLLM`` that - question spans a provider pair. + wants :func:`find_backend_classes` instead, because for CUTLASS + ``fp8_block_scales`` that question spans two leaves. """ if backend_type == MoeBackendType.TRTLLM: from tensorrt_llm._torch.moe.fused_moe.fused_moe_trtllm_gen import find_trtllm_gen_leaf @@ -149,19 +149,11 @@ def find_backend_class(backend_type: MoeBackendType, quant_algo=None): def find_backend_classes(backend_type: MoeBackendType, quant_algo=None) -> tuple: """The classes a resolution walk would try, in its order; empty if none. - A tuple and not one class because ``TRTLLM`` is a provider pair, and - resolution falls through to the native leaf when the FlashInfer one - rejects -- a missing wheel, an unmet shape -- not only when it is absent. - Asking a single pre-picked class prunes candidates a run would have served. - - For CUTLASS ``fp8_block_scales`` these are the SM90 and SM120 leaves. + A tuple and not one class because CUTLASS publishes ``fp8_block_scales`` + twice, as the SM90 and SM120 leaves, and resolution falls through to the + second when the first rejects. Asking a single pre-picked class prunes + candidates a run would have served. """ - if backend_type == MoeBackendType.TRTLLM: - from tensorrt_llm._torch.moe.fused_moe.trtllm_gen import ( - trtllm_gen_leaves_in_resolution_order, - ) - - return trtllm_gen_leaves_in_resolution_order(quant_algo) if backend_type == MoeBackendType.CUTLASS: from tensorrt_llm._torch.moe.fused_moe.cutlass import cutlass_leaves_in_resolution_order @@ -183,10 +175,9 @@ def get_backend_class(backend_type: MoeBackendType, quant_algo=None): return cutlass_leaf(quant_algo) if backend_type == MoeBackendType.TRTLLM: - # Leaves are keyed by (provider, quant), so the format has to be - # named. ``trtllm_gen_leaf`` returns the native leaf where one exists - # and the FlashInfer one for the unquantized format, which is what a - # run without the opt-in flag would select. + # Leaves are keyed by quant, so the format has to be named. + # ``trtllm_gen_leaf`` returns the native leaf for a quantized format + # and the FlashInfer one for the unquantized format. from tensorrt_llm._torch.moe.fused_moe.fused_moe_trtllm_gen import trtllm_gen_leaf return trtllm_gen_leaf(quant_algo) diff --git a/tests/microbenchmarks/bench_moe/search.py b/tests/microbenchmarks/bench_moe/search.py index 770cea4c2617..ff26d08d0c47 100644 --- a/tests/microbenchmarks/bench_moe/search.py +++ b/tests/microbenchmarks/bench_moe/search.py @@ -154,9 +154,6 @@ def _check_backend_can_implement( ) if verdict.eligible: return True, None - # The last candidate's reason, not the first: without the opt-in flag - # every FlashInfer leaf rejects for the flag alone, which says nothing - # about why the configuration is unrunnable. reason = f"{verdict.reject_reason.value}: {verdict.detail}" return False, reason diff --git a/tests/microbenchmarks/bench_moe_comm.py b/tests/microbenchmarks/bench_moe_comm.py index 8d47c4cf8143..68f8a3b28b7e 100644 --- a/tests/microbenchmarks/bench_moe_comm.py +++ b/tests/microbenchmarks/bench_moe_comm.py @@ -986,7 +986,6 @@ def _run_benchmark_worker_under_current_mpi( experts_per_rank, payload_in_workspace=not bool(args.no_payload_in_workspace), alltoall_result_do_sum=True, - use_flashinfer=False, ) if backend is None: diff --git a/tests/unittest/_torch/moe/moe_test_utils.py b/tests/unittest/_torch/moe/moe_test_utils.py index 91244581e9ab..9d19bb282165 100644 --- a/tests/unittest/_torch/moe/moe_test_utils.py +++ b/tests/unittest/_torch/moe/moe_test_utils.py @@ -68,7 +68,6 @@ from tensorrt_llm._torch.moe.fused_moe.mega_moe.mega_moe_cute_dsl import ( is_megamoe_cute_dsl_runtime_available, ) -from tensorrt_llm._torch.moe.fused_moe.trtllm_gen import trtllm_gen_leaves_in_resolution_order from tensorrt_llm._torch.utils import ActivationType, is_gated_activation from tensorrt_llm.models.modeling_utils import QuantAlgo @@ -126,15 +125,12 @@ def find_backend_class( """The MoE backend class, or ``None`` if the family has no leaf for the format. ``CUTLASS``, ``TRTLLM``, ``MARLIN``, ``CUTEDSL`` and ``CUTE_DSL_B12X`` are - not one class each but sets of leaves keyed by quant (and, for TRTLLM-Gen, - provider), so each needs ``quant_algo``. CUTLASS publishes - ``fp8_block_scales`` twice, once per SM family, so for that format the leaf - is the one this device runs. - TRTLLM-Gen prefers the native leaf for that format and falls back to the - FlashInfer sibling, which is what "the TRTLLM backend" has to mean for the - unquantized format: bf16 has no native leaf, only - ``FlashinferTrtllmGenBf16Impl``. For every other format the native leaf - exists and wins, so only bf16 tests are exercising a FlashInfer class. + not one class each but sets of leaves keyed by quant, so each needs + ``quant_algo``. CUTLASS publishes ``fp8_block_scales`` twice, once per SM + family, so for that format the leaf is the one this device runs. + TRTLLM-Gen has one leaf per format: native for every quantized one, and + ``FlashinferTrtllmGenBf16Impl`` for the unquantized one, so only bf16 tests + are exercising a FlashInfer class. Absence is returned rather than raised so that a parameter generator can tell it apart from a real failure: catching the exception instead would let @@ -178,19 +174,13 @@ def find_backend_classes( ) -> tuple[type[MoE], ...]: """The classes a resolution walk would try, in its order; empty if none. - A tuple and not one class because ``TRTLLM`` is a provider pair, and - ``IMPL_PRIORITY`` ranks each FlashInfer leaf ahead of its native sibling: - resolution falls through to the native one when the FlashInfer leaf - rejects -- a missing wheel, an unmet shape -- and not only when it is - absent. A caller asking whether a case is *supported* has to ask the same - way, or it skips cases the layer under test goes on to serve. - - For CUTLASS ``fp8_block_scales`` these are the SM90 and SM120 leaves. + A tuple and not one class because CUTLASS publishes ``fp8_block_scales`` + twice, as the SM90 and SM120 leaves, and resolution asks each. A caller + asking whether a case is *supported* has to ask the same way, or it skips + cases the layer under test goes on to serve. :func:`find_backend_class` stays for callers that want to name one class. """ - if backend_type is MoeBackendType.TRTLLM: - return trtllm_gen_leaves_in_resolution_order(quant_algo) if backend_type is MoeBackendType.CUTLASS: return cutlass_leaves_in_resolution_order(quant_algo) cls = find_backend_class(backend_type, quant_algo) @@ -416,11 +406,6 @@ def should_skip_trtllm( "TRTLLMGenFusedMoE BF16 path requires hidden_size and intermediate_size " f"to be multiples of 128 (got h={hidden_size}, i={intermediate_size})." ) - if comm_method == "NVLINK_TWO_SIDED" and top_k % 8 != 0: - return ( - "TRTLLMGenFusedMoE BF16 path with NVLinkTwoSidedFlashinfer requires " - f"top_k to be a multiple of 8 for 16-byte row alignment (got top_k={top_k})." - ) if moe_tp_size > 1: if intermediate_size % moe_tp_size != 0: return ( @@ -1199,9 +1184,9 @@ def get_quick_skip_reason( try: # Every candidate a run would try, in resolution's order, because - # resolution falls through to the native leaf when the FlashInfer one - # rejects. Asking a single pre-picked class would skip cases the layer - # under test goes on to serve. + # resolution falls through to the next leaf publishing the format when + # one rejects. Asking a single pre-picked class would skip cases the + # layer under test goes on to serve. candidates = find_backend_classes(backend_type, quant_algo) if not candidates: # No leaf of this family publishes the format. Asked as a query, @@ -1268,9 +1253,6 @@ def get_quick_skip_reason( if verdict.eligible: reason = None break - # The last candidate's reason, not the first: without the opt-in - # flag every FlashInfer leaf rejects for the flag alone, which says - # nothing about why the case is unsupported. reason = f"{verdict.reject_reason.value}: {verdict.detail}" if reason is not None: return reason diff --git a/tests/unittest/_torch/moe/test_communication_factory.py b/tests/unittest/_torch/moe/test_communication_factory.py index 59d5cc321ed2..ef39a283174b 100644 --- a/tests/unittest/_torch/moe/test_communication_factory.py +++ b/tests/unittest/_torch/moe/test_communication_factory.py @@ -27,9 +27,6 @@ from tensorrt_llm._torch.moe.fused_moe.communication import communication_factory from tensorrt_llm._torch.moe.fused_moe.communication import nvlink_one_sided as one_sided_module from tensorrt_llm._torch.moe.fused_moe.communication import nvlink_two_sided as two_sided_module -from tensorrt_llm._torch.moe.fused_moe.communication import ( - nvlink_two_sided_flashinfer as flashinfer_module, -) from tensorrt_llm._torch.moe.fused_moe.communication.allgather_reducescatter import ( AllGatherReduceScatter, ) @@ -97,10 +94,8 @@ def test_create_strategy_rejects_invalid_comm_timeout_before_selection( one_sided.assert_not_called() -@pytest.mark.parametrize("use_flashinfer", [False, True]) def test_forced_two_sided_rejects_helix_before_workspace_allocation( monkeypatch: pytest.MonkeyPatch, - use_flashinfer: bool, ) -> None: model_config = _make_model_config() model_config.mapping.has_cp_helix = Mock(return_value=True) @@ -109,15 +104,8 @@ def test_forced_two_sided_rejects_helix_before_workspace_allocation( "is_platform_supported", Mock(return_value=True), ) - monkeypatch.setattr( - communication_factory.NVLinkTwoSidedFlashinfer, - "is_platform_supported", - Mock(return_value=True), - ) native_initialize = Mock(side_effect=AssertionError("native allocation reached")) - flashinfer_symbols = Mock(side_effect=AssertionError("FlashInfer allocation reached")) monkeypatch.setattr(two_sided_module.MnnvlMemory, "initialize", native_initialize) - monkeypatch.setattr(flashinfer_module, "_flashinfer_mnnvl", flashinfer_symbols) with pytest.raises(ValueError, match="does not support Helix context parallelism"): communication_factory.CommunicationFactory._create_forced_method( @@ -129,12 +117,10 @@ def test_forced_two_sided_rejects_helix_before_workspace_allocation( expert_size_per_partition=16, payload_in_workspace=False, alltoall_result_do_sum=True, - use_flashinfer=use_flashinfer, hidden_size=4096, ) native_initialize.assert_not_called() - flashinfer_symbols.assert_not_called() def test_forced_one_sided_rejects_helix_before_workspace_allocation( @@ -161,7 +147,6 @@ def test_forced_one_sided_rejects_helix_before_workspace_allocation( expert_size_per_partition=16, payload_in_workspace=False, alltoall_result_do_sum=True, - use_flashinfer=False, hidden_size=4096, ) @@ -184,15 +169,8 @@ def test_auto_selection_rejects_unrepurposed_helix( "is_platform_supported", Mock(return_value=True), ) - monkeypatch.setattr( - communication_factory.NVLinkTwoSidedFlashinfer, - "is_platform_supported", - Mock(return_value=True), - ) native_initialize = Mock(side_effect=AssertionError("native allocation reached")) - flashinfer_symbols = Mock(side_effect=AssertionError("FlashInfer allocation reached")) monkeypatch.setattr(one_sided_module.MnnvlMemory, "initialize", native_initialize) - monkeypatch.setattr(flashinfer_module, "_flashinfer_mnnvl", flashinfer_symbols) with pytest.raises(ValueError, match="repurpose_helix_cp_to_tp"): communication_factory.CommunicationFactory.create_strategy( model_config, @@ -204,7 +182,6 @@ def test_auto_selection_rejects_unrepurposed_helix( ) native_initialize.assert_not_called() - flashinfer_symbols.assert_not_called() def _install_failing_nccl_module(monkeypatch: pytest.MonkeyPatch, error: BaseException): @@ -330,7 +307,6 @@ def test_forced_nccl_ep_validates_preconditions( expert_size_per_partition=16, payload_in_workspace=False, alltoall_result_do_sum=True, - use_flashinfer=False, hidden_size=4096, ) @@ -358,7 +334,6 @@ def test_forced_nccl_ep_does_not_apply_host_smem_gate( expert_size_per_partition=16, payload_in_workspace=False, alltoall_result_do_sum=True, - use_flashinfer=False, hidden_size=4096, ) @@ -381,7 +356,6 @@ def test_forced_nccl_ep_allows_missing_moe_max_num_tokens( expert_size_per_partition=16, payload_in_workspace=False, alltoall_result_do_sum=True, - use_flashinfer=False, hidden_size=4096, ) @@ -493,7 +467,6 @@ def test_forced_nccl_ep_forwards_low_precision_combine( expert_size_per_partition=16, payload_in_workspace=False, alltoall_result_do_sum=True, - use_flashinfer=False, hidden_size=4096, ) diff --git a/tests/unittest/_torch/moe/test_moe_backend.py b/tests/unittest/_torch/moe/test_moe_backend.py index 3285b5caa1d7..d63a898a0f82 100644 --- a/tests/unittest/_torch/moe/test_moe_backend.py +++ b/tests/unittest/_torch/moe/test_moe_backend.py @@ -146,7 +146,6 @@ W4A16NVFP4CutlassFusedMoEMethod, ) from tensorrt_llm._torch.moe.fused_moe.trtllm_gen import ( - FlashinferTrtllmGenNvfp4Impl, TrtllmTrtllmGenNvfp4Impl, TrtllmTrtllmGenW4a8Mxfp4Mxfp8Impl, ) @@ -1226,15 +1225,9 @@ def create_test_backend( activation_type: ActivationType = ActivationType.Swiglu, locality_domain_policy: Optional[LocalityDomainPolicy] = None, n_shared_experts: int = 0, - moe_cls: Optional[type] = None, ) -> MoE: - """Create a MoE backend for testing. - - ``moe_cls`` pins the leaf class directly, for the providers the default - lookup never selects (resolution prefers the native TRTLLM-Gen leaf - whenever one exists for the format). - """ - backend_cls = moe_cls or get_backend_class( + """Create a MoE backend for testing.""" + backend_cls = get_backend_class( backend_type, None if quant_config is None else quant_config.quant_algo ) if locality_domain_policy is None: @@ -1289,109 +1282,6 @@ def create_test_backend( return backend -# The quantized FlashInfer provider is otherwise never run: resolution prefers -# the native TRTLLM-Gen leaf whenever one exists for the format (only the bf16 -# leaf is FlashInfer-exclusive), so without pinning the class the FlashInfer -# NVFP4 kernels have zero coverage. The real provider is exercised here, which -# needs Blackwell: FlashInfer's FP4 path has no pre-Blackwell backend. -@skip_pre_blackwell -def test_flashinfer_provider_nvfp4_matches_reference(tmp_path: Path): - pytest.importorskip("flashinfer") - - model_config = MoeModelConfig(8, 2, 512, 512) - seq_len = 8 - dtype = torch.bfloat16 - mapping = Mapping() - mapping.rank = mpi_rank() - - with torch.device(f"cuda:{mapping.rank}"): - torch.manual_seed(0) - torch.cuda.manual_seed(0) - AutoTuner.get().setup_distributed_state(mapping) - - routing_method = RenormalizeMoeRoutingMethod(top_k=model_config.top_k) - x = torch.randn((seq_len, model_config.hidden_size), dtype=dtype, device="cuda") - router_logits = torch.randn((seq_len, model_config.num_experts), dtype=dtype, device="cuda") - - quantize_util_cls, quant_config, quant_kwargs = get_test_quant_params( - QuantAlgo.NVFP4, x, MoeBackendType.TRTLLM - ) - quantize_util = quantize_util_cls( - num_experts=model_config.num_experts, - dtype=dtype, - intermediate_size=model_config.intermediate_size, - hidden_size=model_config.hidden_size, - quant_config=quant_config, - bias=False, - swiglu_gptoss_style=False, - swiglu_alpha=None, - swiglu_beta=None, - swiglu_limit=None, - activation_type=ActivationType.Swiglu, - ) - NVFP4TRTLLMGenFusedMoEBaseMethod._cache_permute_indices.clear() - - backend = create_test_backend( - backend_type=MoeBackendType.TRTLLM, - routing_method=routing_method, - num_experts=model_config.num_experts, - hidden_size=model_config.hidden_size, - intermediate_size=model_config.intermediate_size, - dtype=dtype, - quant_config=quant_config, - mapping=mapping, - weight_loading_mode=getattr( - quantize_util, "weight_loading_mode", MoEWeightLoadingMode.VANILLA - ), - moe_cls=FlashinferTrtllmGenNvfp4Impl, - ) - assert backend.use_flashinfer - - ref_cls = quant_kwargs.pop("ref_cls", None) - weights = quantize_util.create_weights(**quant_kwargs) - backend.load_weights([weights]) - backend.post_load_weights() - backend.cuda() - - if ref_cls is not None: - ref_fused_moe = quantize_util.create_ref_module(routing_method, ref_cls=ref_cls) - else: - ref_fused_moe = quantize_util.create_ref_module(routing_method) - ref_fused_moe.load_weights([weights]) - ref_fused_moe.cuda() - - with torch.inference_mode(): - ref_output = ref_fused_moe.forward(x, router_logits) - - def run_moe(): - token_selected_experts, token_final_scales = routing_method.apply(router_logits) - x_quantized, x_sf = backend.quantize_input(x, post_quant_comm=False) - return run_backend_moe( - backend, - MoeBackendType.TRTLLM, - x_quantized, - x_sf, - token_selected_experts, - token_final_scales, - dtype, - router_logits, - ) - - autotuner = AutoTuner.get() - autotuner.warmup = 0 - autotuner.repeat = 1 - autotuner.stream_delay_micro_secs = 10 - autotuner.clear_cache() - with ( - torch.inference_mode(), - autotune(cache_path=str(tmp_path / "moe_autotuner_cache.json")), - ): - _ = run_moe() - with torch.inference_mode(): - output = run_moe() - ref_fused_moe.check_accuracy(output, ref_output) - - # ===================================================================== # Staged post-load hook lifecycle tests # ===================================================================== diff --git a/tests/unittest/_torch/moe/test_moe_comm.py b/tests/unittest/_torch/moe/test_moe_comm.py index 58778f9b0bf3..650d92a05344 100644 --- a/tests/unittest/_torch/moe/test_moe_comm.py +++ b/tests/unittest/_torch/moe/test_moe_comm.py @@ -16,7 +16,7 @@ Unified Communication Strategy Tests for MoE Tests all Communication subclasses (AllGatherReduceScatter, DeepEP, -DeepEPLowLatency, NVLinkOneSided, NVLinkTwoSided, NVLinkTwoSidedFlashinfer) via the public +DeepEPLowLatency, NVLinkOneSided, NVLinkTwoSided, NcclEP) via the public dispatch() and combine() interfaces defined in Communication base class. Each test runs the full pipeline: dispatch -> verify dispatch -> simple_moe @@ -80,9 +80,6 @@ MnnvlMoe, NVLinkTwoSided, ) -from tensorrt_llm._torch.moe.fused_moe.communication.nvlink_two_sided_flashinfer import ( - NVLinkTwoSidedFlashinfer, -) from tensorrt_llm._torch.moe.fused_moe.deep_ep_utils import deep_ep_installed from tensorrt_llm._torch.moe.fused_moe.ep_group_health import EPGroupHealth from tensorrt_llm._torch.moe.fused_moe.nccl_ep_utils import ( @@ -109,7 +106,6 @@ COMM_DEEP_EP_LL = "DeepEPLowLatency" COMM_NVLINK_ONE_SIDED = "NVLinkOneSided" COMM_NVLINK_TWO_SIDED = "NVLinkTwoSided" -COMM_NVLINK_TWO_SIDED_FLASHINFER = "NVLinkTwoSidedFlashinfer" COMM_NCCL_EP = "NcclEP" ALL_COMM_TYPES = [ @@ -118,7 +114,6 @@ COMM_DEEP_EP_LL, COMM_NVLINK_ONE_SIDED, COMM_NVLINK_TWO_SIDED, - COMM_NVLINK_TWO_SIDED_FLASHINFER, COMM_NCCL_EP, ] @@ -611,17 +606,6 @@ def create_comm_object( alltoall_result_do_sum=True, ) - elif comm_type == COMM_NVLINK_TWO_SIDED_FLASHINFER: - # FlashInfer requires reduced combine output and does not support the - # low-precision combine path. - return NVLinkTwoSidedFlashinfer( - mapping=mapping, - num_experts=num_experts, - num_slots=num_slots, - top_k=config.top_k, - alltoall_result_do_sum=True, - ) - elif comm_type == COMM_NCCL_EP: return NcclEP( mapping=mapping, @@ -699,14 +683,6 @@ def _comm_reuse_key(config: CommTestConfig) -> Tuple: config.use_low_precision_combine, ) - if config.comm_type == COMM_NVLINK_TWO_SIDED_FLASHINFER: - return ( - config.comm_type, - config.ep_size, - config.num_experts, - config.top_k, - ) - return (config.comm_type,) @@ -753,17 +729,6 @@ def _check_mnnvl_support() -> Optional[str]: return None -@lru_cache(maxsize=None) -def _check_flashinfer_mnnvl_support() -> Optional[str]: - """Return skip reason if FlashInfer MNNVL is not supported, else None.""" - try: - if not NVLinkTwoSidedFlashinfer.is_platform_supported(): - return "FlashInfer MNNVL not supported" - except Exception: - return "FlashInfer MNNVL initialization failed" - return None - - @lru_cache(maxsize=None) def check_platform_support(comm_type: str) -> Optional[str]: """Return skip reason string if comm type is unsupported, else None.""" @@ -781,9 +746,6 @@ def check_platform_support(comm_type: str) -> Optional[str]: if comm_type == COMM_NVLINK_TWO_SIDED: return _check_mnnvl_support() - if comm_type == COMM_NVLINK_TWO_SIDED_FLASHINFER: - return _check_flashinfer_mnnvl_support() - if comm_type == COMM_NCCL_EP: if not is_nccl_ep_installed(): return "NCCL EP not available (install the nccl4py wheel)" @@ -850,16 +812,6 @@ def check_feasibility(comm_type: str, config: CommTestConfig) -> Optional[str]: if config.top_k > NCCL_EP_MAX_TOP_K: return f"NcclEP MAX_TOP_K={NCCL_EP_MAX_TOP_K}, got top_k={config.top_k}" - if comm_type == COMM_NVLINK_TWO_SIDED_FLASHINFER: - # FlashInfer alltoallv requires every 2D payload row to be 16-byte aligned. - # This test dispatches both int32 slots [N, top_k] and bf16 scales - # [N, top_k], so top_k must satisfy both alignments. - if config.top_k % 8 != 0: - return ( - "NVLinkTwoSidedFlashinfer requires top_k to be a multiple of 8 " - f"for 16-byte row alignment, got top_k={config.top_k}" - ) - # W4AFP8: encode_source_info writes token_idx into fp8 low byte. # token_idx >= 127 can create fp8 NaN which may not survive bf16 roundtrip. if config.quant_mode == "w4afp8": @@ -1064,7 +1016,7 @@ def _run_worker_dispatch( config: CommTestConfig, ) -> DispatchOutputs: """Prepare dispatch metadata if needed, then run dispatch().""" - if config.comm_type in (COMM_NVLINK_TWO_SIDED, COMM_NVLINK_TWO_SIDED_FLASHINFER): + if config.comm_type == COMM_NVLINK_TWO_SIDED: comm.prepare_dispatch(worker_inputs.slots, config.all_num_tokens) recv_hs, recv_sf, recv_slots, recv_scales = comm.dispatch( @@ -1739,7 +1691,7 @@ def verify_dispatch_alltoall( """Verify AllToAll dispatch: slot validity, content integrity, completeness. Slot semantics differ by communication backend: - - NVLinkOneSided / NVLinkTwoSided / NVLinkTwoSidedFlashinfer: + - NVLinkOneSided / NVLinkTwoSided: valid tokens keep global expert IDs [0, num_experts), padding tokens are set to -1 by sanitize kernel. - DeepEP: valid tokens keep global expert IDs [0, num_experts), empty- @@ -2090,7 +2042,6 @@ def verify_combine_results( COMM_NCCL_EP, COMM_NVLINK_ONE_SIDED, COMM_NVLINK_TWO_SIDED, - COMM_NVLINK_TWO_SIDED_FLASHINFER, COMM_DEEP_EP_LL, COMM_ALLGATHER_RS, ], @@ -2098,7 +2049,6 @@ def verify_combine_results( COMM_NCCL_EP, COMM_NVLINK_ONE_SIDED, COMM_NVLINK_TWO_SIDED, - COMM_NVLINK_TWO_SIDED_FLASHINFER, COMM_DEEP_EP, COMM_DEEP_EP_LL, COMM_ALLGATHER_RS, @@ -2107,7 +2057,7 @@ def verify_combine_results( } """Only valid (quant_mode, comm_type) combinations for post-quant tests. - - fp8: NVLink (payload-agnostic, including FlashInfer variant) + DeepEP LL + - fp8: NVLink (payload-agnostic) + DeepEP LL (fp8 branch) + AllGatherRS. DeepEP normal only supports nvfp4 post-quant, NOT fp8. - nvfp4: All tested COMM types except w4afp8-only branches support nvfp4 @@ -2205,7 +2155,6 @@ def _make_boundary_test_params(): configs = [] for target_ep_size in [2, 4]: for comm_type in ALL_COMM_TYPES: - aligned_top_k = 8 if comm_type == COMM_NVLINK_TWO_SIDED_FLASHINFER else 2 boundary_cases = [] if comm_type == COMM_NVLINK_ONE_SIDED: # DSpark routes to three experts; include an idle ADP peer. @@ -2223,21 +2172,20 @@ def _make_boundary_test_params(): f"{comm_type}_topk3", ) ) - if comm_type != COMM_NVLINK_TWO_SIDED_FLASHINFER: - boundary_cases.append( - ( - 2, - CommTestConfig( - comm_type=comm_type, - ep_size=2, - num_experts=FIXED_NUM_EXPERTS, - top_k=1, - hidden_size=DEFAULT_HIDDEN_SIZE, - all_num_tokens=[8, 8], - ), - f"{comm_type}_topk1", - ) + boundary_cases.append( + ( + 2, + CommTestConfig( + comm_type=comm_type, + ep_size=2, + num_experts=FIXED_NUM_EXPERTS, + top_k=1, + hidden_size=DEFAULT_HIDDEN_SIZE, + all_num_tokens=[8, 8], + ), + f"{comm_type}_topk1", ) + ) boundary_cases.extend( [ ( @@ -2246,7 +2194,7 @@ def _make_boundary_test_params(): comm_type=comm_type, ep_size=2, num_experts=FIXED_NUM_EXPERTS, - top_k=aligned_top_k, + top_k=2, hidden_size=2048, all_num_tokens=[16, 16], ), @@ -2258,7 +2206,7 @@ def _make_boundary_test_params(): comm_type=comm_type, ep_size=4, num_experts=FIXED_NUM_EXPERTS, - top_k=aligned_top_k, + top_k=2, hidden_size=DEFAULT_HIDDEN_SIZE, all_num_tokens=[1, 1, 1, 1], ), @@ -2276,7 +2224,7 @@ def _make_boundary_test_params(): comm_type=comm_type, ep_size=4, num_experts=FIXED_NUM_EXPERTS, - top_k=aligned_top_k, + top_k=2, hidden_size=DEFAULT_HIDDEN_SIZE, all_num_tokens=[32, 0, 16, 0], ), @@ -2352,8 +2300,7 @@ def _make_postquant_test_params(): configs = [] for quant_mode, comm_types in POSTQUANT_COMM_MAP.items(): for comm_type in comm_types: - top_ks = [8] if comm_type == COMM_NVLINK_TWO_SIDED_FLASHINFER else [2, 16] - for top_k in top_ks: + for top_k in [2, 16]: for use_low_precision_combine in [False, True]: config = CommTestConfig( comm_type=comm_type, diff --git a/tests/unittest/_torch/moe/test_moe_impl.py b/tests/unittest/_torch/moe/test_moe_impl.py index d41abdc7ac4e..cc07124f53d5 100644 --- a/tests/unittest/_torch/moe/test_moe_impl.py +++ b/tests/unittest/_torch/moe/test_moe_impl.py @@ -95,11 +95,7 @@ identity_quant_of, normalize_quant, ) -from tensorrt_llm._torch.moe.fused_moe.impl_environment import ( - MoEDep, - MoEEnvFlag, - override_moe_environment, -) +from tensorrt_llm._torch.moe.fused_moe.impl_environment import MoEDep, override_moe_environment from tensorrt_llm._torch.moe.fused_moe.impl_identity import ( MOE_IMPL_REGISTRY, MoEImplDescriptor, @@ -124,12 +120,10 @@ from tensorrt_llm._torch.moe.fused_moe.routing import DeepSeekV4MoeRoutingMethod, RoutingMethodType from tensorrt_llm._torch.moe.fused_moe.trtllm_gen import ( FlashinferTrtllmGenBf16Impl, - FlashinferTrtllmGenNvfp4Impl, TrtllmTrtllmGenNvfp4Impl, TrtllmTrtllmGenW4a8Nvfp4Fp8Impl, - trtllm_gen_leaves_in_resolution_order, + find_trtllm_gen_leaf, ) -from tensorrt_llm._torch.moe.fused_moe.trtllm_gen.eligibility import check_flashinfer_provider from tensorrt_llm._torch.utils import ActivationType from tensorrt_llm.models.modeling_utils import QuantAlgo, QuantConfig @@ -511,16 +505,15 @@ def test_registering_megamoe_cutedsl_leaves_the_backend_literal_path_unchanged() # ===================================================================== # TRTLLM-Gen implementation identity # ===================================================================== -# Eleven identities over one abstract parent that holds the shared -# ``__init__``, ``create_weights`` and input preparation. Most formats have a -# native and a FlashInfer leaf that differ only in provider, so these tests -# cover which provider a request lands on, and which leaves decline a routing, -# bias, activation constant or epilogue their cubins do not run. +# Seven identities over one abstract parent that holds the shared +# ``__init__``, ``create_weights`` and input preparation, one per format. These +# tests cover which leaf a request lands on, and which leaves decline a bias, +# activation constant or epilogue their cubins do not run. # The full grid, written out rather than generated, so a leaf that silently # stops registering fails here instead of shrinking a computed expectation. -# Six native leaves; the FlashInfer wheel serves five of the seven formats -- -# it has no fp8/fp4-activation runner, and it alone has the unquantized one. +# Six native leaves for the quantized formats; the FlashInfer wheel serves only +# the unquantized one, which TRT-LLM's own cubins do not cover. _TRTLLM_GEN_IDS = ( "trtllm.trtllm_gen.fused_moe.nvfp4", "trtllm.trtllm_gen.fused_moe.fp8_block_scales", @@ -528,28 +521,16 @@ def test_registering_megamoe_cutedsl_leaves_the_backend_literal_path_unchanged() "trtllm.trtllm_gen.fused_moe.w4a8_mxfp4_mxfp8", "trtllm.trtllm_gen.fused_moe.w4a8_nvfp4_fp8", "trtllm.trtllm_gen.fused_moe.w4a8_mxfp4_fp8", - "flashinfer.trtllm_gen.fused_moe.nvfp4", - "flashinfer.trtllm_gen.fused_moe.fp8_block_scales", - "flashinfer.trtllm_gen.fused_moe.w4a16_mxfp4", - "flashinfer.trtllm_gen.fused_moe.w4a8_mxfp4_mxfp8", "flashinfer.trtllm_gen.fused_moe.none", ) -def _trtllm_gen_environment( - *, sm: int = 100, flashinfer: bool = False, use_flashinfer_flag: bool = False -) -> MoEEnvironment: - """An SM100 host, with the FlashInfer wheel and its opt-in flag optional.""" +def _trtllm_gen_environment(*, sm: int = 100, flashinfer: bool = False) -> MoEEnvironment: + """An SM100 host, with the FlashInfer wheel optional.""" deps = [] if flashinfer: deps = [MoEDep.FLASHINFER.value, MoEDep.FLASHINFER_BF16_MOE.value] - return MoEEnvironment( - sm=sm, - available_deps=tuple(sorted(deps)), - env_flags=( - (MoEEnvFlag.TRTLLM_GEN_USE_FLASHINFER.value, "1" if use_flashinfer_flag else "0"), - ), - ) + return MoEEnvironment(sm=sm, available_deps=tuple(sorted(deps))) def _single_rank_deployment(env: MoEEnvironment) -> MoEDeployment: @@ -571,8 +552,8 @@ def _trtllm_gen_model_config(quant_algo=QuantAlgo.NVFP4): return cfg -def test_trtllm_gen_registers_exactly_eleven_identities(): - """The grid is the deliverable: eleven addressable ids, no more, no fewer.""" +def test_trtllm_gen_registers_exactly_seven_identities(): + """The grid is the deliverable: seven addressable ids, no more, no fewer.""" registered = sorted( identity.canonical() for identity in MOE_IMPL_REGISTRY.identities() @@ -586,7 +567,7 @@ def test_trtllm_gen_identity_round_trips_through_registry(impl_id): impl = MOE_IMPL_REGISTRY.lookup(MoEImplId.parse(impl_id)) assert impl is not None # Declared on the leaf itself, not inherited from the parent: an inherited - # descriptor would give eleven classes one identity. + # descriptor would give seven classes one identity. assert "descriptor" in vars(impl) assert impl.descriptor.identity.canonical() == impl_id assert not impl.__abstractmethods__ @@ -616,7 +597,7 @@ def test_trtllm_gen_leaf_provider_agrees_with_its_identity(impl_id): @pytest.mark.parametrize("impl_id", _TRTLLM_GEN_IDS) def test_trtllm_gen_leaf_admits_only_its_own_format(impl_id): - """Two leaves differing only in provider still admit one format each.""" + """Each leaf admits its own format and no other.""" impl = MOE_IMPL_REGISTRY.lookup(MoEImplId.parse(impl_id)) quant = impl.descriptor.identity.quant other = "nvfp4" if quant != "nvfp4" else "w4a16_mxfp4" @@ -649,22 +630,13 @@ def test_every_trtllm_gen_leaf_descends_from_the_surviving_name(): assert issubclass(impl, TRTLLMGenFusedMoE) -def test_trtllm_literal_picks_the_native_leaf_without_the_opt_in_flag(): - """The default provider must not move just because FlashInfer is installed.""" +def test_trtllm_literal_picks_the_native_leaf_with_flashinfer_installed(): + """A quantized format stays on the native leaf when FlashInfer is present.""" with override_moe_environment(_trtllm_gen_environment(flashinfer=True)): report = resolve_moe_impl(_trtllm_gen_model_config()) assert impl_class_for(report) is TrtllmTrtllmGenNvfp4Impl -def test_trtllm_literal_moves_to_flashinfer_under_the_opt_in_flag(): - """The old ``_check_flashinfer_backend_support`` switch, now in selection.""" - with override_moe_environment( - _trtllm_gen_environment(flashinfer=True, use_flashinfer_flag=True) - ): - report = resolve_moe_impl(_trtllm_gen_model_config()) - assert impl_class_for(report).provider == "flashinfer" - - def test_deepseek_v4_routing_is_its_own_problem_but_keeps_v3s_kernel_encoding(): """The two answers a routing method gives, and why they have to differ. @@ -686,50 +658,9 @@ def test_deepseek_v4_routing_is_its_own_problem_but_keeps_v3s_kernel_encoding(): assert canonical_routing(routing) == RoutingMethodType.DeepSeekV4.name -@pytest.mark.parametrize( - "routing,fused_by_the_native_runner", - [ - ("DeepSeekV3", True), - ("Default", True), - ("DeepSeekV4", False), - ("Renormalize", False), - ], -) -def test_flashinfer_provider_turns_down_only_the_routing_it_cannot_fuse( - routing, fused_by_the_native_runner -): - """The routing half of the old ``_check_flashinfer_backend_support``. - - It has to name the algorithm rather than the kernel encoding: V4 shares - V3's encoding while sharing none of the fused routing this gate objects - to, and the provider serves it. - """ - verdict = check_flashinfer_provider( - FlashinferTrtllmGenNvfp4Impl, - MoEProblem(quant="NVFP4", dtype_act=torch.bfloat16, routing=routing), - _single_rank_deployment(_trtllm_gen_environment(flashinfer=True, use_flashinfer_flag=True)), - ) - - if fused_by_the_native_runner: - assert verdict is not None - assert verdict.reject_reason is MoERejectReason.ROUTING_UNSUPPORTED - else: - assert verdict is None - - -def test_pinned_flashinfer_leaf_is_turned_down_when_not_opted_in(): - """A pin fails hard, and says it was a policy call rather than a capability.""" - with override_moe_environment(_trtllm_gen_environment(flashinfer=True)): - report = resolve_moe_impl( - _trtllm_gen_model_config(), impl_id="flashinfer.trtllm_gen.fused_moe.nvfp4" - ) - assert report.winner is None - assert [rejection.reason for rejection in report.rejected] == [MoERejectReason.PATH_NOT_ENABLED] - - def test_flashinfer_leaves_report_dep_missing_without_the_wheel(): """On a host with no FlashInfer, every FlashInfer leaf declines for that reason.""" - deployment = _single_rank_deployment(_trtllm_gen_environment(use_flashinfer_flag=True)) + deployment = _single_rank_deployment(_trtllm_gen_environment()) for impl_id in _TRTLLM_GEN_IDS: impl = MOE_IMPL_REGISTRY.lookup(MoEImplId.parse(impl_id)) if impl.provider != "flashinfer": @@ -746,46 +677,11 @@ def test_flashinfer_leaves_report_dep_missing_without_the_wheel(): assert verdict.reject_reason is MoERejectReason.DEP_MISSING -@pytest.mark.parametrize( - "activation_constants", - [frozenset(), frozenset({"alpha"})], - ids=["unpadded_method", "padded_method"], -) -def test_flashinfer_nvfp4_declines_expert_bias_under_either_quant_method(activation_constants): - """Bias is a provider capability, so the shape gate cannot be its only home. - - Which quant method this leaf picks depends on the activation constants, - and only one of the two goes through the shape gate. A bias answered there - alone is answered for half the problems, and the other half reaches an - ``__init__`` assert after resolution already named a winner. - - Distinct from ``swiglu_gptoss_style``: that is the bias-plus-alpha-beta - package, and neither case here declares it. - """ - verdict = FlashinferTrtllmGenNvfp4Impl.can_implement( - MoEProblem( - quant="NVFP4", - dtype_act=torch.bfloat16, - bias=True, - activation_constants=activation_constants, - ), - _single_rank_deployment(_trtllm_gen_environment(flashinfer=True, use_flashinfer_flag=True)), - ) - - assert not verdict.eligible - assert verdict.reject_reason is MoERejectReason.ACTIVATION_UNSUPPORTED - - -# The leaves whose runner reads no expert bias. Three because the FlashInfer -# provider passes none for any quantized format, three because the fp8/fp4 -# activation cubins take none on either provider, and the unquantized one -# because its FlashInfer runner is the same. +# The leaves whose runner reads no expert bias: two because the fp8/fp4 +# activation cubins take none, and the unquantized one because its FlashInfer +# runner takes none either. _NO_EXPERT_BIAS_IDS = ( - "flashinfer.trtllm_gen.fused_moe.nvfp4", - "flashinfer.trtllm_gen.fused_moe.w4a16_mxfp4", - "flashinfer.trtllm_gen.fused_moe.w4a8_mxfp4_mxfp8", "trtllm.trtllm_gen.fused_moe.fp8_block_scales", - "flashinfer.trtllm_gen.fused_moe.fp8_block_scales", "trtllm.trtllm_gen.fused_moe.w4a8_nvfp4_fp8", "flashinfer.trtllm_gen.fused_moe.none", ) @@ -798,7 +694,7 @@ def test_trtllm_gen_leaves_without_a_bias_runner_decline_during_selection(impl_i ``capabilities.supports_expert_bias`` is one declaration shared by the whole family, and ``create_moe`` reads it against the class resolution already picked -- so a leaf answering only there is chosen and then - refuses, with no chance for a sibling that can serve to be tried instead. + refuses, with no chance for a backend that can serve to be tried instead. Distinct from ``swiglu_gptoss_style``, which is the bias-plus-alpha-beta package: none of these problems declares it, and a checkpoint can carry a @@ -812,7 +708,7 @@ def test_trtllm_gen_leaves_without_a_bias_runner_decline_during_selection(impl_i dtype_act=torch.bfloat16, bias=True, ), - _single_rank_deployment(_trtllm_gen_environment(flashinfer=True, use_flashinfer_flag=True)), + _single_rank_deployment(_trtllm_gen_environment(flashinfer=True)), ) assert not verdict.eligible @@ -861,7 +757,7 @@ def test_leaves_without_constant_support_decline_during_selection(impl, constant dtype_act=torch.bfloat16, activation_constants=frozenset({constant}), ), - _single_rank_deployment(_trtllm_gen_environment(flashinfer=True, use_flashinfer_flag=True)), + _single_rank_deployment(_trtllm_gen_environment(flashinfer=True)), ) assert not verdict.eligible @@ -885,22 +781,16 @@ def test_nvfp4_leaves_admit_the_calibration_aliases(alias): assert verdict.eligible, verdict.detail -def test_resolution_order_lists_flashinfer_before_its_native_sibling(): - """Filtering callers need resolution's order, not a lookup's. +def test_trtllm_gen_lookup_finds_the_one_leaf_per_format(): + """Each format has exactly one leaf, on whichever provider serves it. - ``IMPL_PRIORITY`` ranks each FlashInfer leaf ahead of its native sibling - and the opt-in flag is what makes the FlashInfer ones reject, so a caller - asking "would a run serve this?" has to walk both and fall through on a - rejection. Provider-exclusive formats yield the one leaf that exists. + The quantized formats land on the native provider and the unquantized one + on FlashInfer; a format no leaf publishes is an ordinary ``None``. """ - assert trtllm_gen_leaves_in_resolution_order(QuantAlgo.NVFP4) == ( - FlashinferTrtllmGenNvfp4Impl, - TrtllmTrtllmGenNvfp4Impl, - ) - assert trtllm_gen_leaves_in_resolution_order(None) == (FlashinferTrtllmGenBf16Impl,) - assert trtllm_gen_leaves_in_resolution_order(QuantAlgo.W4A8_NVFP4_FP8) == ( - TrtllmTrtllmGenW4a8Nvfp4Fp8Impl, - ) + assert find_trtllm_gen_leaf(QuantAlgo.NVFP4) is TrtllmTrtllmGenNvfp4Impl + assert find_trtllm_gen_leaf(QuantAlgo.W4A8_NVFP4_FP8) is TrtllmTrtllmGenW4a8Nvfp4Fp8Impl + assert find_trtllm_gen_leaf(None) is FlashinferTrtllmGenBf16Impl + assert find_trtllm_gen_leaf(QuantAlgo.W8A16) is None # The two native leaves whose cubin family has a fused SiTu FC1 epilogue: @@ -924,9 +814,8 @@ def test_trtllm_gen_situ_admitted_only_by_the_leaves_with_a_fused_cubin(impl_id) raises, which reads as a crash rather than as a format that was never eligible. - The FlashInfer leaves are turned down here even for the two capable - formats: SiTu is a native TRTLLM-Gen cubin and is absent from FlashInfer's - activation enum. + The FlashInfer bf16 leaf is turned down too: SiTu is a native TRTLLM-Gen + cubin and is absent from FlashInfer's activation enum. """ impl = MOE_IMPL_REGISTRY.lookup(MoEImplId.parse(impl_id)) quant = impl.descriptor.identity.quant @@ -936,7 +825,7 @@ def test_trtllm_gen_situ_admitted_only_by_the_leaves_with_a_fused_cubin(impl_id) dtype_act=torch.bfloat16, activation="SiTu", ), - _single_rank_deployment(_trtllm_gen_environment(flashinfer=True, use_flashinfer_flag=True)), + _single_rank_deployment(_trtllm_gen_environment(flashinfer=True)), ) if impl_id in _SITU_CAPABLE_IDS: assert verdict.eligible, verdict.detail @@ -2652,7 +2541,6 @@ def test_cutlass_communication_selection(monkeypatch, is_oss, method, force_sour "AllGatherReduceScatter", "NVLinkOneSided", "NVLinkTwoSided", - "NVLinkTwoSidedFlashinfer", "NcclEP", "DeepEP", "DeepEPLowLatency", @@ -2668,7 +2556,6 @@ def test_cutlass_communication_selection(monkeypatch, is_oss, method, force_sour num_slots=4, routing_method=SimpleNamespace(experts_per_token=1), expert_size_per_partition=2, - use_flashinfer=False, hidden_size=128, communication_method=method if force_source == "model" else None, )