Skip to content

[TRTLLM-17328][chore] Remove quantized FlashInfer TRTLLM-Gen MoE leaves and FlashInfer two-sided alltoall #21013

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
rosong11 wants to merge 2 commits into
base: main
Choose a base branch
from
+192 −1,496
Open

[TRTLLM-17328][chore] Remove quantized FlashInfer TRTLLM-Gen MoE leaves and FlashInfer two-sided alltoall #21013

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
fffe861
[TRTLLM-17328][chore] Remove quantized FlashInfer TRTLLM-Gen MoE leav…
rosong11 Oct 26, 2026
81ee87a
[TRTLLM-17328][chore] Drop the dead hidden-size gate and the test-onl…
rosong11 Oct 29, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion .claude/skills/trtllm-moe-develop/references/moe-canonical-code-examples.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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_<quant>.py` / `flashinfer_<quant>.py` leaves.
`trtllm_<quant>.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.<kernel>.<quant>`.
Expand Down
54 changes: 22 additions & 32 deletions tensorrt_llm/_torch/moe/fused_moe/MOE_DEVELOPER_GUIDE.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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/`)

Expand Down Expand Up @@ -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_<quant>.py` | one registered native leaf each: `nvfp4`, `fp8_block_scales`, `w4a16_mxfp4`, `w4a8_mxfp4_mxfp8`, `w4a8_nvfp4_fp8`, `w4a8_mxfp4_fp8` |
| `flashinfer_<quant>.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
Expand Down Expand Up @@ -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
Expand Down
56 changes: 16 additions & 40 deletions tensorrt_llm/_torch/moe/fused_moe/communication/communication_factory.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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(
Expand Down Expand Up @@ -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(

@sunnyqgg sunnyqgg Oct 30, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Required (R1). Worth a line in the description: besides lifting the top_k % 8 restriction, moving the BF16 leaf to native NVLinkTwoSided means it now honours use_low_precision_moe_combine (an NVFP4 combine of the BF16 output, as CUTLASS BF16 layers already get on this strategy), where NVLinkTwoSidedFlashinfer asserted it off. So TRTLLM_FORCE_COMM_METHOD=NVLINK_TWO_SIDED plus the knob used to fail at startup for the BF16 leaf and now runs, and auto-selection no longer falls through to the next strategy. In practice this reaches e.g. the unquantized MTP layer of a DeepSeek NVFP4 checkpoint with moe_backend: TRTLLM. Whether the knob should touch unquantized layers at all (it is documented as 'only for NVFP4 quantization') is a cross-backend question, maybe for the per-layer quant config in #20768, not something this PR has to settle.

@rosong11 rosong11 Oct 30, 2026

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Good catch, BF16 TRTLLM over NVLINK_TWO_SIDED (4 GPUs, test_configurable_moe_multi_gpu, TRTLLM_TEST_MOE_CI=0): I ran them locally with that skip removed and only quant=None cases in the process (the cascade needs a block-scale case first). Native NVLinkTwoSided: 40/40 passed (DEP/TEP/DTP/TTP). Base with NVLinkTwoSidedFlashinfer: 16/16 passed (only the top_k=8 shapes were eligible there).

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:
Expand Down Expand Up @@ -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:
"""
Expand Down Expand Up @@ -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(
Expand Down
Loading
Loading