From db154374ed2e7ffb11bd2dc67e20249de9937012 Mon Sep 17 00:00:00 2001 From: xxi Date: Tue, 27 Oct 2026 00:49:54 +0000 Subject: [PATCH 1/7] [TRTLLM-14973][fix] Refuse or honour every input on the B12x MoE path The B12x MoE leaves (SM120/121, NVFP4 and W4A16_NVFP4) answered several requests they cannot serve with a different result instead of an error, and returned wrong output for a forward larger than moe_config.max_num_tokens. run_moe inputs: the b12x route handed the FlashInfer kernel x and the routing pair and read nothing else of MoERunContext. The shared mixin now admits each run_moe call before either route runs: it refuses do_finalize=False, lora_params, logits-only routing and all-to-all, none of which the leaf can honour on its b12x or its CUTLASS NVFP4 prefill route, and joins x_ready_event before a kernel reads x. The b12x route also refuses an output_dtype other than bfloat16. The served configuration (bfloat16, routing applied by the scheduler, no communication strategy, do_finalize=True) is unchanged. float16: B12xMoEWrapper reads and writes bfloat16 only, but the gate admitted float16. The mixin now declares supported_dtypes = BF16_ONLY, so a float16 layer is declined with DTYPE_UNSUPPORTED and falls back to the CUTLASS leaf. force_dynamic_quantization: the NVFP4 leaf honoured it only on its CUTLASS prefill route; decode chunks used the scales transform_weights snapshots at load time, so one layer ran two quantization regimes. FlashInfer 0.6.18 quantizes the FC2 input inside the fused kernel against a fixed scale, so the leaf declines the setting in can_implement with a new reason, DYNAMIC_QUANT_UNSUPPORTED, and the layer falls back to the CUTLASS NVFP4 leaf. Multi-chunk forwards: a B12xMoEWrapper built with use_cuda_graph=True writes each call into one preallocated output buffer (shared across layers by _SHARED_MOE_OUTPUT_BUF) and returns a view of its leading rows, and it keeps one workspace across calls. ExternalCommMoEScheduler kept each chunk's run_moe result and concatenated them after the last chunk, so the layer returned the last chunk's output in place of every earlier one; it also ran even chunks on the aux stream, unordered against the odd chunks, so two chunks of one layer could be in flight on the same buffers. On an RTX PRO 6000 every two-chunk W4A16_NVFP4 and NVFP4 b12x-route forward returned chunk 0 equal to chunk 1, and holding chunk 0 back by about 1 ms corrupted 10/10 W4A16 forwards even with the leaf output cloned. MoEInputRequirement gains run_moe_reuses_buffers, both B12x leaves set it, and for such a backend _forward_multiple_chunks runs every chunk on one stream and clones each chunk's output before the next chunk runs. Other backends and the single-chunk path are unchanged. Signed-off-by: xxi --- .../moe/fused_moe/MOE_DEVELOPER_GUIDE.md | 21 +- .../moe/fused_moe/fused_moe_cute_dsl_b12x.py | 87 ++++++- .../_torch/moe/fused_moe/impl_contract.py | 11 + .../_torch/moe/fused_moe/moe_scheduler.py | 9 +- .../_torch/moe/fused_moe/quantization.py | 6 +- .../moe/test_cute_dsl_b12x_moe_backend.py | 220 +++++++++++++++++- tests/unittest/_torch/moe/test_moe_backend.py | 161 +++++++++++++ tests/unittest/_torch/moe/test_moe_impl.py | 88 ++++++- 8 files changed, 569 insertions(+), 34 deletions(-) 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 57063ebb6368..1aa06610ba18 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/MOE_DEVELOPER_GUIDE.md +++ b/tensorrt_llm/_torch/moe/fused_moe/MOE_DEVELOPER_GUIDE.md @@ -126,7 +126,7 @@ Adaptive quantize/dispatch order (gated by comm.supports_post_quant_dispatch()): Pre-quant flow: comm.dispatch() → quantize_input() (send raw, quantize locally) ``` -EPLB hooks fire only at the first/last chunk of the first/last `repeat_idx`. Multi-stream chunk overlap is enabled when `not enable_alltoall and aux_stream is not None`. +EPLB hooks fire only at the first/last chunk of the first/last `repeat_idx`. Multi-stream chunk overlap is enabled when `not enable_alltoall and aux_stream is not None`, unless the backend declares `input_requirement.run_moe_reuses_buffers`: its `run_moe` returns a view of a buffer the next call overwrites and keeps scratch across calls, so the scheduler runs its chunks on one stream and copies each chunk's output before the next chunk runs (the B12x leaves, whose `B12xMoEWrapper` output buffer is also shared across layers). `unfinalized_combine_fn` is a constructor argument (`create_moe(..., unfinalized_combine_fn=...)`) for a model that needs a per-`(token, expert)` transform applied to each expert output *before* the top-k sum. Once registered it runs on **every** `do_finalize=False` forward, with or without a comm strategy, so the forward's return type (dense tensor vs. unfinalized triple) depends only on static config. Under a comm strategy the scheduler runs it on the dispatched rows and the combine reduces already-finalized partials — equal to finalizing after the combine, because each term depends on one row alone and non-local slots contribute exact zeros. That reordering is only sound for a combine that is an unweighted row sum over a dense per-token tensor; `DeepEPLowLatency.combine` is not one (it consumes expert-major rows and applies the routing weights itself), which is what `Communication.supports_finalize_before_combine()` declares. `ConfigurableMoE.validate_backend` rejects at construction: a backend without `supports_unfinalized_output()`, a comm without `supports_finalize_before_combine()`, and `apply_router_weight_on_input`. With no hook registered, `do_finalize=False` plus a comm strategy stays a named runtime refusal instead of a crash inside the comm layer. @@ -227,7 +227,7 @@ Still on old path (standalone, with embedded communication): | `fused_moe_deepgemm.py` | `DeepgemmCudaFp8BlockScalesImpl` (aliased as `DeepGemmFusedMoE`) | SM100/SM103/SM107 | FP8 Block Scales on Blackwell/Rubin | `EXTERNAL_COMM` | | `cutedsl/dense_gemm.py` (+ the `fused_moe_densegemm.py` alias module) | `TrtllmCutedslDenseGemmNvfp4Impl` (aliased as `DenseGEMMFusedMoE`) | SM100/SM103 | NVFP4 min-latency; CuTe DSL dense GEMM packs all experts into one matrix (vs Cutlass per-expert scatter), efficient for small token counts | `EXTERNAL_COMM` | | `cutedsl/grouped_gemm/` (+ the `fused_moe_cute_dsl.py` alias module) | the abstract `CuteDslFusedMoEBase` (aliased as `CuteDslFusedMoE`) and its two leaves, `TrtllmCutedslGroupedGemmNvfp4Impl` (`trtllm.cutedsl.grouped_gemm.nvfp4`) and `TrtllmCutedslGroupedGemmBf16Impl` (`trtllm.cutedsl.grouped_gemm.none`) | SM100/SM103/SM107; BF16 SM107 only | High throughput NVFP4, generally faster than Cutlass; unquantized BF16 SwiGLU on Rubin | `EXTERNAL_COMM` | -| `fused_moe_cute_dsl_b12x.py` | `CuteDslB12xNvfp4FusedMoE` and `CuteDslB12xW4a16Nvfp4FusedMoE` (shared mixin `CuteDslB12xFusedMoEBase`, aliased as `CuteDslB12xFusedMoE`) | SM120/SM121 | NVFP4 (CUTLASS prefill, FlashInfer b12x decode) and W4A16 NVFP4 (b12x for every call), best on RTX PRO 6000 and DGX Spark. They head the `CUTEDSL` family, so they win on SM120/121 when flashinfer is present and yield to the grouped-GEMM leaves otherwise. Single-GPU-shaped topology only: with no dispatch/combine kernel they reject `ep_size > 1` and attention-DP | `EXTERNAL_COMM` | +| `fused_moe_cute_dsl_b12x.py` | `CuteDslB12xNvfp4FusedMoE` and `CuteDslB12xW4a16Nvfp4FusedMoE` (shared mixin `CuteDslB12xFusedMoEBase`, aliased as `CuteDslB12xFusedMoE`) | SM120/SM121 | NVFP4 (CUTLASS prefill, FlashInfer b12x decode) and W4A16 NVFP4 (b12x for every call), best on RTX PRO 6000 and DGX Spark. They head the `CUTEDSL` family, so they win on SM120/121 when flashinfer is present and yield to the grouped-GEMM leaves otherwise. BF16 activations only (`B12xMoEWrapper` reads and writes bfloat16), so a float16 layer falls back to Cutlass. The NVFP4 leaf also declines `force_dynamic_quantization` (`DYNAMIC_QUANT_UNSUPPORTED`): b12x quantizes decode activations against load-time scales, so such a layer falls back to Cutlass too. Single-GPU-shaped topology only: with no dispatch/combine kernel they reject `ep_size > 1` and attention-DP | `EXTERNAL_COMM` | | `cutedsl/fused_fc12.py` (+ the `fused_moe_cute_dsl_fc12.py` alias module) | `TrtllmCutedslFusedFc12Nvfp4Impl` (`trtllm.cutedsl.fused_fc12.nvfp4`, aliased as `CuteDslFc12FusedMoE`) | SM107 (Rubin) | NVFP4 fused FC1+FC2 in one persistent kernel (`cute_dsl_nvfp4_fc12_fused_rubin`) with a trace-time SwiGLU or SiTU epilogue, keeping the FC1→FC2 intermediate on-chip; shares `TrtllmCutedslGroupedGemmNvfp4Impl`'s NVFP4 weight layout and autotune runner; reachable through the `CUTEDSL_FC12` family or a pinned `impl_id`; needs `MoEDep.CUTEDSL_RUBIN`; non-uGPU | `EXTERNAL_COMM` | | `mega_moe/mega_moe_deepgemm.py` | `DeepgemmCudaW4a8Mxfp4Mxfp8Impl` (aliased as `MegaMoEDeepGemm`) | SM100/SM103/SM107 | W4A8_MXFP4_MXFP8 via DeepGEMM `fp8_fp4_mega_moe` fused dispatch+GEMM+act+GEMM+combine kernel; requires `hidden_size % 512 == 0` | `FUSED_COMM` | | `mega_moe/mega_moe_cute_dsl.py` | `TrtllmCutedslMegaMoeNvfp4Impl` (aliased as `MegaMoECuteDsl`) | SM100/SM103/SM107 | NVFP4 fused dispatch+FC1+act+FC2+combine; internally selects one of its per-architecture kernel variants. Requires CUDA 13 Cutlass DSL and a symmetric-memory provider. Uses uniform activation constants and applies routing weights before FC2 quantization for DeepSeek-V4, after FC2 for other models | `FUSED_COMM` | @@ -561,7 +561,10 @@ implementations' rejections only (the fallback candidates are asked after them): the highest-priority one that declined for a reason other than the layer's format or the SM. On SM100 the SM120-only b12x leaves lead the `CUTEDSL` family, so without the SM exclusion their `SM_UNSUPPORTED` would stand in for -the CuteDSL grouped-GEMM leaf's rejection. +the CuteDSL grouped-GEMM leaf's rejection. The same exclusion means a leaf that +serves the format but not a deployment setting must not answer +`QUANT_UNSUPPORTED`, or the report blames another candidate: the B12x NVFP4 +leaf declines `force_dynamic_quantization` as `DYNAMIC_QUANT_UNSUPPORTED`. Two things fail hard, and they fail in different places. An unknown or retired backend literal raises `ValueError` before any candidate is considered, so there @@ -639,7 +642,7 @@ Each backend's `can_implement(p, d)` classmethod declares what it supports. Sour | FP8 QDQ | Y (SM89+) | N | N | N | N | N | N | Y (SM90) | N | Y | | FP8 Block Scales | Y (SM90, SM120) | Y (SM100/103/107) | Y (SM100/103/107) | N | N‡ | N | N | N | N | Y | | NVFP4 | Y (SM100/103/107/120/121) | Y (SM100/103/107) | N | Y (SM100/103) | Y (SM100/103/107/120/121)¶ | N | Y (SM100/103/107, cu13 cutlass-dsl + symmetric-memory provider; per-expert alpha/norm_const + SwiGLU clamp) | N | Y (SM89-SM99) | Y | -| W4A16 NVFP4 | Y (SM80+, dequant-on-the-fly) | N | N | N | Y (SM120/121 via `CuteDslB12xW4a16Nvfp4FusedMoE`, needs flashinfer) | N | N | N | Y (SM89-SM99, BF16) | Y | +| W4A16 NVFP4 | Y (SM80+, dequant-on-the-fly) | N | N | N | Y (SM120/121, BF16, via `CuteDslB12xW4a16Nvfp4FusedMoE`, needs flashinfer) | N | N | N | Y (SM89-SM99, BF16) | Y | | W4A8 NVFP4 FP8 | N | Y (SM100/103) | N | N | N | N | N | N | N | N | | W4A16 MXFP4 | Y (SM90) | Y (SM100/103/107) | N | N | N | N | N | Y (SM90) | N | N | | W4A8 MXFP4 FP8 | Y (SM100/103/107) | Y (SM100/103/107) | N | N | N | N | N | Y (SM90) | N | N | @@ -689,8 +692,9 @@ two run different kernels: `W4A16NVFP4CutlassFusedMoEMethod` dequantizes the FP4 weights into the activation dtype each forward and then calls the unquantized kernel, so it inherits that path's `SM80+` floor instead of needing NVFP4 tensor cores. This is what makes Cutlass the landing spot when a W4A16 NVFP4 -layer finds no specialized backend — the B12x W4A16 leaf without flashinfer, -or any SM outside Marlin's 89-99 and B12x's 120/121. +layer finds no specialized backend — the B12x W4A16 leaf without flashinfer +or with float16 activations, or any SM outside Marlin's 89-99 and B12x's +120/121. ‡ No CuteDSL leaf publishes `fp8_block_scales`; the DeepGEMM and TRTLLM-Gen leaves own the format on SM100/103/107 and Cutlass on SM90/SM120. Because a @@ -780,8 +784,9 @@ the formats whose leaf sets `requires_oss_cutlass_build` (`MXFP8`, `W4A8_MXFP4_MXFP8`, `W4A16_MXFP4`, `W4A8_AWQ`), dynamic FC2 scaling on the leaf that sets `uses_dynamic_fc2_scale` (`NVFP4`), post-SiLU clamping, and routed-expert LoRA. The B12x NVFP4 leaf runs the same gate for the prefill -chunks it routes through CUTLASS. W4A8 AWQ uses per-expert activation-scale -buffers and is unavailable in the legacy build. `CutlassFusedMoEBase` exposes +chunks it routes through CUTLASS; it declines dynamic quantization before +that gate on any build, because its decode route cannot honour it. W4A8 AWQ +uses per-expert activation-scale buffers and is unavailable in the legacy build. `CutlassFusedMoEBase` exposes `supports_alltoall` before communication selection; it is False only for a leaf that issues the op on a legacy build. Unsupported auto-selection falls back to AllGather/ReduceScatter, and diff --git a/tensorrt_llm/_torch/moe/fused_moe/fused_moe_cute_dsl_b12x.py b/tensorrt_llm/_torch/moe/fused_moe/fused_moe_cute_dsl_b12x.py index 659a8d8b3881..cd96c2c80c47 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/fused_moe_cute_dsl_b12x.py +++ b/tensorrt_llm/_torch/moe/fused_moe/fused_moe_cute_dsl_b12x.py @@ -25,7 +25,7 @@ from ...utils import ActivationType, Fp4QuantizedTensor from .activation import MoEActivationSupport from .cutlass import CutlassFusedMoEBase, TrtllmCutlassNvfp4Impl -from .cutlass.eligibility import check_oss_cutlass_build +from .cutlass.eligibility import BF16_ONLY, check_dtype, check_oss_cutlass_build from .impl_contract import ( MoEDeployment, MoEEligibility, @@ -50,7 +50,9 @@ # bytes of GPU memory holding identical-shape buffers that are written # sequentially. We fold them into a single shared buffer because MoE layers # run sequentially on the same CUDA stream, and the wrapper consumes its -# previous output before the next layer is dispatched. +# previous output before the next layer is dispatched. Within one layer, the +# scheduler runs the chunks of a multi-chunk forward on one stream and copies +# each chunk out, because the leaves declare ``run_moe_reuses_buffers``. _SHARED_MOE_OUTPUT_BUF: dict = {} # ActivationType -> b12x activation string. b12x currently exposes "relu2" @@ -154,8 +156,8 @@ def candidate_tile_fits( class CuteDslB12xFusedMoEBase: """What the two B12x SM120/121 leaves share: the gates, the format check, - the quant method and the decode call, keyed by each leaf's - ``b12x_quant_algo``. + the quant method, the ``run_moe`` admission and the decode call, keyed by + each leaf's ``b12x_quant_algo``. A mixin, because the NVFP4 leaf inherits the CUTLASS NVFP4 leaf (for prefill) while the W4A16 leaf inherits only ``CutlassFusedMoEBase``. @@ -170,6 +172,10 @@ class CuteDslB12xFusedMoEBase: #: layer runs on. Compared as ``MoEProblem.identity_quant`` spells it, so the #: NVFP4 leaf also admits the NVFP4_AWQ / NVFP4_ARC calibration aliases. b12x_quant_algo: QuantAlgo | None = None + #: Whether the b12x kernel quantizes activations. It does so against the + #: global scales ``NVFP4CuteDslB12xFusedMoEMethod.transform_weights`` fixes + #: at load time, so the gate declines ``force_dynamic_quantization``. + b12x_quantizes_activations: bool = False # The decode path reads no expert bias, fuses no LoRA and has no EPLB # support; the gates still decline EPLB explicitly because the inherited @@ -185,8 +191,17 @@ class CuteDslB12xFusedMoEBase: ) # Pinned rather than inherited, so a change in the CUTLASS parent cannot - # retarget these backends. - input_requirement = MoEInputRequirement(routing_scales_dtype=torch.float32) + # retarget these backends. ``B12xMoEWrapper`` keeps one workspace across + # calls (FlashInfer caches it by shape even without CUDA graphs) and, with + # CUDA graphs, returns a view of its output buffer, which + # ``_SHARED_MOE_OUTPUT_BUF`` also shares across layers. + input_requirement = MoEInputRequirement( + routing_scales_dtype=torch.float32, run_moe_reuses_buffers=True + ) + + # ``B12xMoEWrapper`` reads and writes bfloat16 only. Pinned, so the NVFP4 + # leaf does not answer with its CUTLASS parent's set. + supported_dtypes = BF16_ONLY # The kinds ``_ACTIVATION_MAP`` gates on; the decode kernel takes no # activation constants. @@ -225,11 +240,9 @@ def can_implement(cls, p: MoEProblem, d: MoEDeployment) -> MoEEligibility: MoERejectReason.SM_UNSUPPORTED, f"{cls.__name__} requires {sm_list}, got SM{sm_version}", ) - if p.dtype_act not in {torch.float16, torch.bfloat16}: - return _reject( - MoERejectReason.DTYPE_UNSUPPORTED, - f"{cls.__name__} requires float16 or bfloat16 activation dtype (got {p.dtype_act})", - ) + rejection = check_dtype(cls, p) + if rejection is not None: + return rejection if p.swiglu_gptoss_style: return _reject( MoERejectReason.ACTIVATION_UNSUPPORTED, @@ -258,6 +271,16 @@ def can_implement(cls, p: MoEProblem, d: MoEDeployment) -> MoEEligibility: MoERejectReason.EPLB_UNSUPPORTED, f"{cls.__name__} does not support the MoE load balancer", ) + # b12x quantizes the FC2 input inside the fused kernel against a scale + # passed in, with no per-call amax to derive it from. Per-call scales + # would therefore reach only the CUTLASS prefill chunks, and decode + # would keep the load-time ones. + if cls.b12x_quantizes_activations and d.force_dynamic_quantization: + return _reject( + MoERejectReason.DYNAMIC_QUANT_UNSUPPORTED, + f"{cls.__name__} quantizes decode activations with load-time " + "scales and cannot honour force_dynamic_quantization", + ) # No expert-parallel dispatch/combine kernel: EP must stay at 1. if d.ep_size != 1: return _reject( @@ -319,6 +342,38 @@ def _reject_fp4_input(self, x: torch.Tensor | Fp4QuantizedTensor) -> None: "on the b12x path; b12x performs its own input quantization." ) + def _admit_run_context(self, ctx: MoERunContext) -> None: + """Refuse what neither route of the leaf honours, then join the + input-ready event before any kernel reads ``ctx.x``. + + On both routes the leaf finalizes, fuses no LoRA, routes nothing itself + and holds every expert (its gate rejects ``ep_size > 1``), so a context + asking otherwise fails here instead of having the request dropped. + ``all_rank_num_tokens`` has no reader on either route: ``run_moe`` uses + it only for the all-to-all tuner shapes, and all-to-all is refused. + """ + name = type(self).__name__ + if not ctx.do_finalize: + raise NotImplementedError( + f"{name} always finalizes, so it cannot return the unfinalized " + "(gemm2_output, expert_weights, expanded_idx_to_permuted_idx) " + "triple that do_finalize=False asks for." + ) + if ctx.lora_params is not None: + raise NotImplementedError(f"{name} does not fuse routed-expert LoRA; got lora_params.") + if ctx.token_selected_experts is None or ctx.token_final_scales is None: + raise ValueError( + f"{name} needs precomputed token_selected_experts and " + "token_final_scales; it cannot route from router_logits." + ) + if require_comm_plan(self, ctx).enable_alltoall: + raise NotImplementedError( + f"{name} holds every expert (ep_size == 1) and cannot take " + "tokens dispatched by all-to-all." + ) + if ctx.x_ready_event is not None: + ctx.x_ready_event.wait() + @nvtx_range("[b12x] decode") def _run_b12x_decode(self, ctx: MoERunContext) -> torch.Tensor: """Run the FlashInfer b12x kernel. The path both leaves reach.""" @@ -334,6 +389,13 @@ def _run_b12x_decode(self, ctx: MoERunContext) -> torch.Tensor: f"{type(self).__name__} expects unquantized input (x_sf=None) " "on the b12x path; got a precomputed scale factor." ) + # ``B12xMoEWrapper`` writes bfloat16 and refuses any other + # ``output_dtype`` at construction. + if ctx.output_dtype not in (None, torch.bfloat16): + raise ValueError( + f"{type(self).__name__} writes bfloat16 on the b12x path; " + f"cannot honour output_dtype={ctx.output_dtype}." + ) # Annotate the kwargs spread + wrapper entry separately so we can # attribute the per-layer Python dispatch cost vs. the kernel cost. @@ -373,6 +435,7 @@ class CuteDslB12xNvfp4FusedMoE(CuteDslB12xFusedMoEBase, TrtllmCutlassNvfp4Impl): """ b12x_quant_algo = QuantAlgo.NVFP4 + b12x_quantizes_activations = True # Prefill chunks (``x.shape[0] >= threshold``) route via CUTLASS NVFP4 # GroupGEMM; decode (``x.shape[0] < threshold``) uses b12x. 64 cleanly @@ -418,6 +481,7 @@ def run_moe( *, workspace: dict | None = None, ) -> torch.Tensor: + self._admit_run_context(ctx) if self._route_to_cutlass(ctx.x): # The inherited ``run_moe`` forwards ``output_dtype`` straight into # the C++ ``trtllm::fused_moe`` op, which requires a concrete @@ -493,4 +557,5 @@ def run_moe( workspace: dict | None = None, ) -> torch.Tensor: del workspace # The b12x wrapper allocates its own intermediates. + self._admit_run_context(ctx) return self._run_b12x_decode(ctx) diff --git a/tensorrt_llm/_torch/moe/fused_moe/impl_contract.py b/tensorrt_llm/_torch/moe/fused_moe/impl_contract.py index 4f8dd3ac5d9e..32bf12026319 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/impl_contract.py +++ b/tensorrt_llm/_torch/moe/fused_moe/impl_contract.py @@ -79,6 +79,12 @@ class MoEInputRequirement: # Overrides the NVLink one-sided combine payload dtype; None means the # buffer follows the model output dtype. onesided_workspace_dtype: Optional[torch.dtype] = None + # ``run_moe`` works in buffers the backend keeps across calls: the tensor it + # returns may be a view of one, overwritten by the next call, and two calls + # in flight on different streams would share the same scratch. A multi-chunk + # forward therefore runs every chunk on one stream and copies each chunk's + # output before the next chunk runs. + run_moe_reuses_buffers: bool = False # No sentinel for invalid_token_expert_id: every Communication uses -1. @@ -327,6 +333,11 @@ class MoERejectReason(str, Enum): # EPLB is registered for this layer and the impl cannot lay out slots for # it. Distinct from TOPOLOGY_UNSUPPORTED: the parallel sizes are fine. EPLB_UNSUPPORTED = "eplb_unsupported" + # The deployment forces per-call activation scales + # (``force_dynamic_quantization``) and the impl quantizes against scales + # fixed at load time. Distinct from QUANT_UNSUPPORTED: the impl serves the + # format, so the report must name this rejection as the degradation cause. + DYNAMIC_QUANT_UNSUPPORTED = "dynamic_quant_unsupported" # The impl only has a fused-finalize FC2 epilogue, and the caller disabled # finalize fusion (explicitly, or implicitly by configuring LoRA). FINALIZE_FUSION_REQUIRED = "finalize_fusion_required" diff --git a/tensorrt_llm/_torch/moe/fused_moe/moe_scheduler.py b/tensorrt_llm/_torch/moe/fused_moe/moe_scheduler.py index 468c9d3d4f4a..a73a16cf2dfc 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/moe_scheduler.py +++ b/tensorrt_llm/_torch/moe/fused_moe/moe_scheduler.py @@ -852,7 +852,12 @@ def _forward_multiple_chunks( else [None] * num_chunks ) - use_multi_stream = not moe.enable_alltoall and moe.aux_stream is not None + # A backend that reuses its buffers across calls can have only one chunk + # in flight, and each chunk's output is gone once the next one runs. + reuses_buffers = moe.backend.input_requirement.run_moe_reuses_buffers + use_multi_stream = ( + not moe.enable_alltoall and moe.aux_stream is not None and not reuses_buffers + ) # ========== Setup auxiliary stream ========== if use_multi_stream: @@ -975,7 +980,7 @@ def _forward_multiple_chunks( row_offset += chunk_size_list[idx_chunk] if chunked_used[idx_chunk]: - outputs_list.append(outputs) + outputs_list.append(outputs.clone() if reuses_buffers else outputs) # ========== Wait for auxiliary stream to complete ========== if use_multi_stream: diff --git a/tensorrt_llm/_torch/moe/fused_moe/quantization.py b/tensorrt_llm/_torch/moe/fused_moe/quantization.py index 764d45775cfb..0172e374af45 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/quantization.py +++ b/tensorrt_llm/_torch/moe/fused_moe/quantization.py @@ -4068,8 +4068,10 @@ def transform_weights(self, module: torch.nn.Module) -> None: ) # Replace the wrapper's per-instance output buffer with a shared one. - # Layers run sequentially on a single stream, so a single buffer of the - # right shape is correct and saves + # Layers run sequentially on a single stream, and the scheduler copies + # each chunk of a multi-chunk forward out of it (the backend declares + # ``run_moe_reuses_buffers``), so a single buffer of the right shape is + # correct and saves # ``(num_moe_layers - 1) * max_num_tokens * hidden_size * 2`` bytes — # ~2.5 GB on Nemotron-Super-120B with ``max_num_tokens=2048``, # ``hidden=8192``, bf16, 80 MoE layers. diff --git a/tests/unittest/_torch/moe/test_cute_dsl_b12x_moe_backend.py b/tests/unittest/_torch/moe/test_cute_dsl_b12x_moe_backend.py index 855b509a57ab..af51feff8946 100644 --- a/tests/unittest/_torch/moe/test_cute_dsl_b12x_moe_backend.py +++ b/tests/unittest/_torch/moe/test_cute_dsl_b12x_moe_backend.py @@ -14,6 +14,7 @@ # limitations under the License. """B12x MoE backend gating and dispatch tests.""" +import dataclasses import sys import types from typing import Optional @@ -22,6 +23,7 @@ import pytest import torch +from tensorrt_llm._torch.moe.fused_moe.cutlass import TrtllmCutlassNvfp4Impl from tensorrt_llm._torch.moe.fused_moe.fused_moe_cute_dsl_b12x import ( CuteDslB12xFusedMoE, CuteDslB12xNvfp4FusedMoE, @@ -29,10 +31,12 @@ _patch_flashinfer_w4a16_ultra_wide_fc2_tile_validation, ) from tensorrt_llm._torch.moe.fused_moe.impl_contract import ( + MoECommPlan, MoEDeployment, MoEEnvironment, MoEProblem, MoERejectReason, + MoERunContext, canonical_quant, ) from tensorrt_llm._torch.moe.fused_moe.impl_environment import MoEDep @@ -70,6 +74,8 @@ def _deployment( use_dp: bool = False, parallel_size: Optional[int] = None, eplb: bool = False, + force_dynamic_quantization: bool = False, + oss_cutlass_moe: bool = True, ) -> MoEDeployment: """Declare the machine rather than patching the probes that read it.""" return MoEDeployment( @@ -79,8 +85,10 @@ def _deployment( use_dp=use_dp, num_slots=8, eplb_enabled=eplb, + force_dynamic_quantization=force_dynamic_quantization, env=MoEEnvironment( sm=sm, + oss_cutlass_moe=oss_cutlass_moe, available_deps=(MoEDep.FLASHINFER.value,) if flashinfer else (), ), ) @@ -161,9 +169,13 @@ def test_can_implement_rejects_swiglu_gptoss_style(): assert verdict.reject_reason is MoERejectReason.ACTIVATION_UNSUPPORTED -@pytest.mark.parametrize("dtype", [torch.float32, torch.float8_e4m3fn]) -def test_can_implement_rejects_unsupported_activation_dtype(dtype): - verdict = NVFP4_LEAF.can_implement(_problem(dtype=dtype), _deployment(120)) +@pytest.mark.parametrize("dtype", [torch.float16, torch.float32, torch.float8_e4m3fn]) +@pytest.mark.parametrize("leaf,quant_algo", B12X_LEAVES, ids=LEAF_IDS) +def test_can_implement_rejects_unsupported_activation_dtype(leaf, quant_algo, dtype): + """``B12xMoEWrapper`` reads and writes bfloat16 only, so float16 is turned + down too, on both leaves: the W4A16 kernel has a float16 path of its own, + but the wrapper hands it a bfloat16 output buffer.""" + verdict = leaf.can_implement(_problem(quant_algo, dtype=dtype), _deployment(120)) assert not verdict.eligible assert verdict.reject_reason is MoERejectReason.DTYPE_UNSUPPORTED @@ -181,6 +193,29 @@ def test_can_implement_rejects_eplb(): assert verdict.reject_reason is MoERejectReason.EPLB_UNSUPPORTED +@pytest.mark.parametrize("oss_cutlass_moe", [True, False], ids=["oss", "internal"]) +@pytest.mark.parametrize( + "leaf,quant_algo,eligible", + [ + pytest.param(CuteDslB12xNvfp4FusedMoE, QuantAlgo.NVFP4, False, id="nvfp4"), + pytest.param(CuteDslB12xW4a16Nvfp4FusedMoE, QuantAlgo.W4A16_NVFP4, True, id="w4a16"), + ], +) +def test_can_implement_declines_forced_dynamic_quantization( + leaf, quant_algo, eligible, oss_cutlass_moe +): + """The NVFP4 decode kernel quantizes against load-time scales, so per-call + scales would reach only the CUTLASS prefill chunks. Declined on every build, + for that reason rather than the prefill build gate's. The W4A16 kernel + quantizes no activation, so the setting asks nothing of it.""" + deployment = _deployment(120, force_dynamic_quantization=True, oss_cutlass_moe=oss_cutlass_moe) + verdict = leaf.can_implement(_problem(quant_algo), deployment) + assert verdict.eligible is eligible + if not eligible: + assert verdict.reject_reason is MoERejectReason.DYNAMIC_QUANT_UNSUPPORTED + assert "force_dynamic_quantization" in verdict.detail + + def test_can_implement_rejects_expert_parallelism(): verdict = NVFP4_LEAF.can_implement(_problem(), _deployment(120, ep_size=2)) assert not verdict.eligible @@ -367,6 +402,185 @@ def test_dispatch_rejects_non_tensor(): assert moe._route_to_cutlass(object()) is False +# --------------------------------------------------------------------------- +# run_moe call boundary: each MoERunContext field is honoured or refused +# --------------------------------------------------------------------------- + +_DECODE_TOKENS = CuteDslB12xNvfp4FusedMoE._PREFILL_VIA_CUTLASS_THRESHOLD - 1 +_PREFILL_TOKENS = CuteDslB12xNvfp4FusedMoE._PREFILL_VIA_CUTLASS_THRESHOLD + +# (leaf, chunk size, the kernel that chunk reaches): both routes of the NVFP4 +# leaf, and the W4A16 leaf, which has only the b12x one. +B12X_ROUTES = [ + pytest.param(CuteDslB12xNvfp4FusedMoE, _DECODE_TOKENS, "b12x", id="nvfp4-decode"), + pytest.param(CuteDslB12xNvfp4FusedMoE, _PREFILL_TOKENS, "cutlass", id="nvfp4-prefill"), + pytest.param(CuteDslB12xW4a16Nvfp4FusedMoE, _PREFILL_TOKENS, "b12x", id="w4a16"), +] +B12X_KERNEL_ROUTES = [route for route in B12X_ROUTES if route.values[2] == "b12x"] + + +class _B12xWrapperRecorder: + """Stands in for ``B12xMoEWrapper``, writing bfloat16 as its kernels do.""" + + def __init__(self, log: list) -> None: + self.log = log + self.output: Optional[torch.Tensor] = None + + def run(self, *, x, token_selected_experts, token_final_scales, **weights) -> torch.Tensor: + self.log.append("b12x") + self.output = torch.zeros(x.shape, dtype=torch.bfloat16) + return self.output + + +class _RecordingEvent: + def __init__(self, log: list) -> None: + self.log = log + + def wait(self, stream=None) -> None: + self.log.append("wait") + + +def _leaf_with_recorded_kernels(leaf_cls: type, log: list, monkeypatch) -> CuteDslB12xFusedMoE: + """A leaf whose two kernels -- b12x and the inherited CUTLASS NVFP4 + ``run_moe`` -- only record that they ran. + + ``__new__`` without ``__init__``: the real constructor wants a GPU, while + ``run_moe`` reads only the wrapper and its weights. + """ + + def cutlass_run_moe(self, ctx: MoERunContext, *, workspace=None) -> torch.Tensor: + log.append("cutlass") + return torch.zeros(ctx.x.shape, dtype=ctx.output_dtype) + + monkeypatch.setattr(TrtllmCutlassNvfp4Impl, "run_moe", cutlass_run_moe) + leaf = leaf_cls.__new__(leaf_cls) + leaf.b12x_wrapper = _B12xWrapperRecorder(log) + leaf._b12x_weights = {} + return leaf + + +def _run_context(num_tokens: int, **overrides) -> MoERunContext: + """What the scheduler hands a served B12x layer: routing already applied, + no communication strategy, bfloat16 activations.""" + ctx = MoERunContext( + token_selected_experts=torch.zeros((num_tokens, 2), dtype=torch.int32), + token_final_scales=torch.ones((num_tokens, 2), dtype=torch.float32), + x=torch.zeros((num_tokens, 8), dtype=torch.bfloat16), + x_sf=None, + output_dtype=torch.bfloat16, + comm_plan=MoECommPlan( + input_sf_swizzled=True, + enable_alltoall=False, + moe_output=None, + payload_in_workspace=False, + ), + ) + return dataclasses.replace(ctx, **overrides) + + +@pytest.mark.parametrize("leaf_cls,num_tokens,kernel", B12X_ROUTES) +def test_run_moe_hands_back_the_kernel_output_untouched(leaf_cls, num_tokens, kernel, monkeypatch): + """The served configuration reaches its kernel, and the b12x output comes + back as the very tensor the wrapper wrote.""" + log = [] + leaf = _leaf_with_recorded_kernels(leaf_cls, log, monkeypatch) + + out = leaf.run_moe(_run_context(num_tokens, all_rank_num_tokens=[num_tokens])) + + assert log == [kernel] + if kernel == "b12x": + assert out is leaf.b12x_wrapper.output + + +@pytest.mark.parametrize( + "overrides,error,match", + [ + pytest.param( + dict(do_finalize=False), NotImplementedError, "do_finalize=False", id="unfinalized" + ), + pytest.param(dict(lora_params={"num_seqs": 1}), NotImplementedError, "LoRA", id="lora"), + pytest.param( + dict( + token_selected_experts=None, + token_final_scales=None, + router_logits=torch.zeros(1, 8), + ), + ValueError, + "router_logits", + id="router-logits-only", + ), + pytest.param( + dict( + comm_plan=MoECommPlan( + input_sf_swizzled=True, + enable_alltoall=True, + moe_output=None, + payload_in_workspace=False, + ) + ), + NotImplementedError, + "all-to-all", + id="alltoall", + ), + ], +) +@pytest.mark.parametrize("leaf_cls,num_tokens,kernel", B12X_ROUTES) +def test_run_moe_refuses_what_neither_route_honours( + leaf_cls, num_tokens, kernel, overrides, error, match, monkeypatch +): + """Neither kernel the leaf runs returns the unfinalized triple, fuses LoRA, + routes from logits or holds only a slice of the experts, so asking for any + of these fails before a kernel launches instead of being ignored.""" + log = [] + leaf = _leaf_with_recorded_kernels(leaf_cls, log, monkeypatch) + + with pytest.raises(error, match=match): + leaf.run_moe(_run_context(num_tokens, **overrides)) + assert log == [] + + +@pytest.mark.parametrize("output_dtype", [torch.float16, torch.float32]) +@pytest.mark.parametrize("leaf_cls,num_tokens,kernel", B12X_KERNEL_ROUTES) +def test_b12x_route_refuses_an_output_dtype_it_cannot_write( + leaf_cls, num_tokens, kernel, output_dtype, monkeypatch +): + """``B12xMoEWrapper`` writes bfloat16 and refuses any other + ``output_dtype`` at construction, so a chunk asking for another one fails + rather than receiving bfloat16.""" + log = [] + leaf = _leaf_with_recorded_kernels(leaf_cls, log, monkeypatch) + + with pytest.raises(ValueError, match="bfloat16"): + leaf.run_moe(_run_context(num_tokens, output_dtype=output_dtype)) + assert log == [] + + +def test_cutlass_route_writes_the_requested_output_dtype(monkeypatch): + """The NVFP4 prefill route is the CUTLASS op, which writes any + high-precision dtype; only the b12x route is bound to bfloat16.""" + log = [] + leaf = _leaf_with_recorded_kernels(CuteDslB12xNvfp4FusedMoE, log, monkeypatch) + + out = leaf.run_moe(_run_context(_PREFILL_TOKENS, output_dtype=torch.float16)) + + assert log == ["cutlass"] + assert out.dtype == torch.float16 + + +@pytest.mark.parametrize("leaf_cls,num_tokens,kernel", B12X_ROUTES) +def test_run_moe_joins_the_input_ready_event_before_its_kernel( + leaf_cls, num_tokens, kernel, monkeypatch +): + """``x`` may still be written on another stream; no kernel reads it before + the event the scheduler handed over.""" + log = [] + leaf = _leaf_with_recorded_kernels(leaf_cls, log, monkeypatch) + + leaf.run_moe(_run_context(num_tokens, x_ready_event=_RecordingEvent(log))) + + assert log == ["wait", kernel] + + # --------------------------------------------------------------------------- # flashinfer W4A16 ultra-wide FC2 tile re-validation shim (nvbug 6721561) # --------------------------------------------------------------------------- diff --git a/tests/unittest/_torch/moe/test_moe_backend.py b/tests/unittest/_torch/moe/test_moe_backend.py index 03366c8fe7cb..4be2fc113b4c 100644 --- a/tests/unittest/_torch/moe/test_moe_backend.py +++ b/tests/unittest/_torch/moe/test_moe_backend.py @@ -14,6 +14,7 @@ # limitations under the License. """MoE backend unit tests.""" +import collections import dataclasses import importlib import itertools @@ -102,6 +103,7 @@ MoECommPlan, MoEDeployment, MoEEnvironment, + MoEInputRequirement, MoEProblem, MoERejection, MoERejectReason, @@ -5635,6 +5637,165 @@ def test_b12x_reaches_the_cutlass_op_only_for_an_nvfp4_prefill_chunk( assert out.shape == (num_tokens, 8) +# ===================================================================== +# Multi-chunk forwards: a backend whose run_moe reuses its buffers +# across calls. +# ===================================================================== + +_MAIN_STREAM = "main" +_AUX_STREAM = "aux" + + +def _record_streams(monkeypatch) -> list: + """Replace ``torch.cuda.stream`` with a CPU stand-in; the returned stack's + top is the stream work is being issued on.""" + stack = [_MAIN_STREAM] + + @contextmanager + def stream(s): + stack.append(s) + try: + yield + finally: + stack.pop() + + monkeypatch.setattr(torch.cuda, "stream", stream) + return stack + + +class _OneBufferB12xWrapper: + """``B12xMoEWrapper`` as built with ``use_cuda_graph=True``: every call + writes into one preallocated buffer and returns its leading rows, so + consecutive calls hand back the same storage.""" + + def __init__(self, max_num_tokens: int, hidden_size: int, stream_stack: list) -> None: + self._moe_output = torch.empty((max_num_tokens, hidden_size), dtype=torch.bfloat16) + self._stream_stack = stream_stack + self.streams = [] + + def run(self, *, x, token_selected_experts, token_final_scales, **weights) -> torch.Tensor: + self.streams.append(self._stream_stack[-1]) + out = self._moe_output[: x.shape[0]] + out.copy_(x * 2) # any per-token function of the input + return out + + +class _FreshOutputBackend: + """A backend whose ``run_moe`` allocates its output, as the CUTLASS op does.""" + + capabilities = MoEStaticCapability() + input_requirement = MoEInputRequirement() + + def __init__(self, stream_stack: list) -> None: + self._stream_stack = stream_stack + self.streams = [] + + def _supports_load_balancer(self) -> bool: + return True + + def try_fused_route_quant(self, x, router_logits): + return None + + def quantize_input(self, x, post_quant_comm=True): + return x, None + + def run_moe(self, ctx: MoERunContext, *, workspace=None) -> torch.Tensor: + self.streams.append(self._stream_stack[-1]) + return ctx.x * 2 + + +def _chunking_moe(backend, num_chunks: int) -> SimpleNamespace: + """What ``ExternalCommMoEScheduler`` reads off a single-rank layer with no + communication strategy and no EPLB, whose ``moe_max_num_tokens`` is below + the batch so it has an aux stream for chunk overlap.""" + + def no_op(*args, **kwargs): + return None + + def route(router_logits, input_ids=None): + num_tokens = router_logits.shape[0] + return ( + torch.zeros((num_tokens, 1), dtype=torch.int32), + torch.ones((num_tokens, 1), dtype=torch.float32), + ) + + return SimpleNamespace( + backend=backend, + comm=None, + use_dp=False, + enable_alltoall=False, + parallel_size=1, + reduce_results=False, + apply_router_weight_on_input=False, + layer_idx=0, + num_slots=4, + repeat_idx=0, + repeat_count=1, + layer_load_balancer=None, + aux_stream=_AUX_STREAM, + event_dict=collections.defaultdict(Mock), + routing_method=SimpleNamespace( + experts_per_token=1, requires_separated_routing=True, apply=route + ), + calculate_num_chunks=lambda all_rank_num_tokens: num_chunks, + split_chunk=lambda n, k: [n // k + (i < n % k) for i in range(k)], + determine_communication_method=no_op, + _using_load_balancer=lambda: False, + _load_balancer_start_wait_gpu_stage=no_op, + _load_balancer_start_set_cpu_stage=no_op, + _load_balancer_done_set_cpu_stage=no_op, + ) + + +def _forward_in_two_chunks(backend) -> Tuple[torch.Tensor, torch.Tensor]: + """Eight tokens through the external-comm scheduler, split into two chunks + of four. Integer-valued, so the expected ``x * 2`` is exact in bfloat16.""" + x = torch.arange(8 * 8, dtype=torch.bfloat16).reshape(8, 8) + out = ExternalCommMoEScheduler(_chunking_moe(backend, num_chunks=2)).forward( + x, + torch.zeros((8, 4)), + do_finalize=True, + output_dtype=torch.bfloat16, + all_rank_num_tokens=None, + use_dp_padding=None, + ) + return x, out + + +@pytest.mark.cpu_only +@pytest.mark.parametrize( + "leaf_cls", [CuteDslB12xNvfp4FusedMoE, CuteDslB12xW4a16Nvfp4FusedMoE], ids=["nvfp4", "w4a16"] +) +def test_b12x_multi_chunk_forward_keeps_every_chunk(leaf_cls, monkeypatch): + """Every b12x call returns the same storage, so a two-chunk forward has to + copy the first chunk out before the second overwrites it. The two calls + also share the wrapper's workspace, so they must not be in flight on two + streams at once. Chunks of four tokens stay below the NVFP4 leaf's + CUTLASS threshold, so both leaves run the b12x kernel.""" + stream_stack = _record_streams(monkeypatch) + leaf = leaf_cls.__new__(leaf_cls) + leaf.fc31_act_scale = None + leaf.b12x_wrapper = _OneBufferB12xWrapper(4, 8, stream_stack) + leaf._b12x_weights = {} + + x, out = _forward_in_two_chunks(leaf) + + torch.testing.assert_close(out, x * 2, rtol=0, atol=0) + assert leaf.b12x_wrapper.streams == [_MAIN_STREAM, _MAIN_STREAM] + + +@pytest.mark.cpu_only +def test_multi_chunk_forward_overlaps_chunks_of_a_backend_that_allocates_its_output(monkeypatch): + """The serialization is the declaration's doing: a backend that does not + declare ``run_moe_reuses_buffers`` keeps the aux-stream overlap.""" + backend = _FreshOutputBackend(_record_streams(monkeypatch)) + + x, out = _forward_in_two_chunks(backend) + + torch.testing.assert_close(out, x * 2, rtol=0, atol=0) + assert backend.streams == [_AUX_STREAM, _MAIN_STREAM] + + @pytest.mark.cpu_only @pytest.mark.parametrize("enable_alltoall", [False, True], ids=["allgather", "alltoall"]) def test_triton_fp8_block_scales_leaf_localizes_global_slot_ids(enable_alltoall, monkeypatch): diff --git a/tests/unittest/_torch/moe/test_moe_impl.py b/tests/unittest/_torch/moe/test_moe_impl.py index b89f8476b958..7699ff96f370 100644 --- a/tests/unittest/_torch/moe/test_moe_impl.py +++ b/tests/unittest/_torch/moe/test_moe_impl.py @@ -1501,6 +1501,74 @@ def test_cutedsl_degradation_is_not_blamed_on_the_sm120_leaf(): ) +def test_cutedsl_float16_nvfp4_on_sm120_falls_back_to_cutlass(): + """The b12x kernel reads and writes bfloat16 only, so on its own SM a + float16 NVFP4 layer goes to the CUTLASS leaf, which takes float16, and the + b12x leaf's dtype is named as the reason.""" + cfg = ModelConfig() + cfg.moe_backend = MoeBackendType.CUTEDSL.value + problem = MoEProblem( + quant=canonical_quant(QuantAlgo.NVFP4), + dtype_act=torch.float16, + hidden_size=1024, + intermediate_size=1024, + num_experts=8, + top_k=2, + swiglu_gptoss_style=False, + ) + deployment = MoEDeployment( + ep_size=1, + tp_size=1, + parallel_size=1, + use_dp=False, + num_slots=8, + env=MoEEnvironment(sm=120, available_deps=(MoEDep.FLASHINFER.value,)), + ) + + report = resolve_moe_impl(cfg, problem=problem, deployment=deployment) + + assert impl_class_for(report) is TrtllmCutlassNvfp4Impl + assert (report.degraded_from.legacy_backend, report.degraded_from.reason) == ( + CuteDslB12xNvfp4FusedMoE.__name__, + MoERejectReason.DTYPE_UNSUPPORTED, + ) + + +def test_cutedsl_dynamic_quant_nvfp4_on_sm120_falls_back_to_cutlass(): + """The b12x decode kernel quantizes against load-time scales, so a layer + that forces dynamic quantization goes to the CUTLASS leaf, which honours it + on every chunk. The b12x leaf is named as the reason: its rejection must not + read as a format mismatch, which ``degraded_from`` passes over.""" + cfg = ModelConfig() + cfg.moe_backend = MoeBackendType.CUTEDSL.value + problem = MoEProblem( + quant=canonical_quant(QuantAlgo.NVFP4), + dtype_act=torch.bfloat16, + hidden_size=1024, + intermediate_size=1024, + num_experts=8, + top_k=2, + swiglu_gptoss_style=False, + ) + deployment = MoEDeployment( + ep_size=1, + tp_size=1, + parallel_size=1, + use_dp=False, + num_slots=8, + force_dynamic_quantization=True, + env=MoEEnvironment(sm=120, available_deps=(MoEDep.FLASHINFER.value,)), + ) + + report = resolve_moe_impl(cfg, problem=problem, deployment=deployment) + + assert impl_class_for(report) is TrtllmCutlassNvfp4Impl + assert (report.degraded_from.legacy_backend, report.degraded_from.reason) == ( + CuteDslB12xNvfp4FusedMoE.__name__, + MoERejectReason.DYNAMIC_QUANT_UNSUPPORTED, + ) + + # ===================================================================== # A layer's format is fixed when it is built # ===================================================================== @@ -2656,16 +2724,20 @@ def test_b12x_cutlass_build_mode_resolution(is_oss, sm, quant_algo, activation, env=MoEEnvironment(sm=sm, oss_cutlass_moe=is_oss, available_deps=(MoEDep.FLASHINFER,)), ), ) - rejection = not is_oss and quant_algo == QuantAlgo.NVFP4 and (activation == "Relu2" or dynamic) - if rejection: + nvfp4 = quant_algo == QuantAlgo.NVFP4 + if nvfp4 and dynamic: + # Declined ahead of the build gate, since decode cannot honour it on + # any build; the CUTLASS fallback needs the open-source kernels for it. + b12x_rejection = _rejection_of(report, leaf) + assert b12x_rejection.reason is MoERejectReason.DYNAMIC_QUANT_UNSUPPORTED + if is_oss: + assert impl_class_for(report) is TrtllmCutlassNvfp4Impl + else: + assert report.winner is None + elif nvfp4 and not is_oss and activation == "Relu2": assert report.winner is None b12x_rejection = _rejection_of(report, leaf) - expected = ( - MoERejectReason.ACTIVATION_UNSUPPORTED - if activation == "Relu2" - else MoERejectReason.QUANT_UNSUPPORTED - ) - assert b12x_rejection.reason is expected + assert b12x_rejection.reason is MoERejectReason.ACTIVATION_UNSUPPORTED assert "USING_OSS_CUTLASS_MOE_GEMM=ON" in b12x_rejection.detail assert leaf.__name__ in b12x_rejection.detail assert b12x_rejection.detail.startswith("NVFP4 prefill routes through CUTLASS:") From ccbc23c0ee28d5412cb310c68b37958a4dab30bd Mon Sep 17 00:00:00 2001 From: xxi Date: Wed, 28 Oct 2026 11:10:39 +0000 Subject: [PATCH 2/7] [TRTLLM-14973][fix] Refuse a buffer-reusing backend on the fused-comm scheduler run_moe_reuses_buffers is declared on the shared MoEInputRequirement, but only ExternalCommMoEScheduler honours it. FusedCommMoEScheduler._run_chunks keeps the tensor each chunk returned and never copies it, so a FUSED_COMM backend declaring the field would hand torch.cat the same storage for every chunk. Reject that combination in forward() instead of leaving a silent wrong-value path, and record the restriction on the field and in the guide. While there, skip the clone of the last chunk on the external-comm path: no further run_moe runs before the torch.cat that follows, and the cat copies that chunk anyway. Signed-off-by: xxi --- .../moe/fused_moe/MOE_DEVELOPER_GUIDE.md | 2 +- .../_torch/moe/fused_moe/impl_contract.py | 3 ++- .../_torch/moe/fused_moe/moe_scheduler.py | 20 +++++++++++++- tests/unittest/_torch/moe/test_moe_backend.py | 26 +++++++++++++++++++ 4 files changed, 48 insertions(+), 3 deletions(-) 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 1aa06610ba18..6302ce8b6913 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/MOE_DEVELOPER_GUIDE.md +++ b/tensorrt_llm/_torch/moe/fused_moe/MOE_DEVELOPER_GUIDE.md @@ -126,7 +126,7 @@ Adaptive quantize/dispatch order (gated by comm.supports_post_quant_dispatch()): Pre-quant flow: comm.dispatch() → quantize_input() (send raw, quantize locally) ``` -EPLB hooks fire only at the first/last chunk of the first/last `repeat_idx`. Multi-stream chunk overlap is enabled when `not enable_alltoall and aux_stream is not None`, unless the backend declares `input_requirement.run_moe_reuses_buffers`: its `run_moe` returns a view of a buffer the next call overwrites and keeps scratch across calls, so the scheduler runs its chunks on one stream and copies each chunk's output before the next chunk runs (the B12x leaves, whose `B12xMoEWrapper` output buffer is also shared across layers). +EPLB hooks fire only at the first/last chunk of the first/last `repeat_idx`. Multi-stream chunk overlap is enabled when `not enable_alltoall and aux_stream is not None`, unless the backend declares `input_requirement.run_moe_reuses_buffers`: its `run_moe` returns a view of a buffer the next call overwrites and keeps scratch across calls, so the scheduler runs its chunks on one stream and copies each chunk's output before the next chunk runs (the B12x leaves, whose `B12xMoEWrapper` output buffer is also shared across layers). The last chunk is not copied: the `torch.cat` that follows it copies it anyway. Only `ExternalCommMoEScheduler` honours the declaration; `FusedCommMoEScheduler` keeps every chunk output uncopied and refuses a backend that declares it. `unfinalized_combine_fn` is a constructor argument (`create_moe(..., unfinalized_combine_fn=...)`) for a model that needs a per-`(token, expert)` transform applied to each expert output *before* the top-k sum. Once registered it runs on **every** `do_finalize=False` forward, with or without a comm strategy, so the forward's return type (dense tensor vs. unfinalized triple) depends only on static config. Under a comm strategy the scheduler runs it on the dispatched rows and the combine reduces already-finalized partials — equal to finalizing after the combine, because each term depends on one row alone and non-local slots contribute exact zeros. That reordering is only sound for a combine that is an unweighted row sum over a dense per-token tensor; `DeepEPLowLatency.combine` is not one (it consumes expert-major rows and applies the routing weights itself), which is what `Communication.supports_finalize_before_combine()` declares. `ConfigurableMoE.validate_backend` rejects at construction: a backend without `supports_unfinalized_output()`, a comm without `supports_finalize_before_combine()`, and `apply_router_weight_on_input`. With no hook registered, `do_finalize=False` plus a comm strategy stays a named runtime refusal instead of a crash inside the comm layer. diff --git a/tensorrt_llm/_torch/moe/fused_moe/impl_contract.py b/tensorrt_llm/_torch/moe/fused_moe/impl_contract.py index 32bf12026319..a5406cc14416 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/impl_contract.py +++ b/tensorrt_llm/_torch/moe/fused_moe/impl_contract.py @@ -83,7 +83,8 @@ class MoEInputRequirement: # returns may be a view of one, overwritten by the next call, and two calls # in flight on different streams would share the same scratch. A multi-chunk # forward therefore runs every chunk on one stream and copies each chunk's - # output before the next chunk runs. + # output before the next chunk runs. ``FusedCommMoEScheduler`` copies no + # chunk output, so it refuses a backend that declares this. run_moe_reuses_buffers: bool = False # No sentinel for invalid_token_expert_id: every Communication uses -1. diff --git a/tensorrt_llm/_torch/moe/fused_moe/moe_scheduler.py b/tensorrt_llm/_torch/moe/fused_moe/moe_scheduler.py index a73a16cf2dfc..33b9aaf14488 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/moe_scheduler.py +++ b/tensorrt_llm/_torch/moe/fused_moe/moe_scheduler.py @@ -980,7 +980,10 @@ def _forward_multiple_chunks( row_offset += chunk_size_list[idx_chunk] if chunked_used[idx_chunk]: - outputs_list.append(outputs.clone() if reuses_buffers else outputs) + # The last chunk needs no copy: no further ``run_moe`` runs + # before the ``torch.cat`` below, which copies it anyway. + needs_copy = reuses_buffers and idx_chunk != num_chunks - 1 + outputs_list.append(outputs.clone() if needs_copy else outputs) # ========== Wait for auxiliary stream to complete ========== if use_multi_stream: @@ -1147,6 +1150,8 @@ class FusedCommMoEScheduler(MoEScheduler): so peers can cross the in-kernel NVLink barrier. 8. No external Communication.dispatch / Communication.combine. 9. No multi-stream chunk overlap. + 10. Reject a backend declaring ``input_requirement.run_moe_reuses_buffers``: + chunk outputs are kept, never copied, so such a backend would alias. ``repeat_idx`` advancement is done by ``ConfigurableMoE.forward_impl`` after this scheduler returns. The scheduler must not rotate @@ -1193,6 +1198,19 @@ def forward( "quantization happens in backend.quantize_input." ) + # ``run_moe_reuses_buffers`` is an external-comm promise today: this + # scheduler already runs its chunks on one stream, but ``_run_chunks`` + # keeps the tensor each chunk returned, so a backend whose next call + # overwrites it would hand ``torch.cat`` the same storage num_chunks + # times -- the last chunk's rows, silently, for every chunk. Refuse the + # combination rather than honour half of it. + if self.moe.backend.input_requirement.run_moe_reuses_buffers: + raise NotImplementedError( + f"{type(self.moe.backend).__name__} declares " + "run_moe_reuses_buffers; the fused-comm (MegaMoE) scheduler " + "does not copy chunk outputs and cannot drive it." + ) + x_real, rl_real, input_ids_real, real_all_rank_num_tokens, ep_rank, had_meta = ( self._strip_adp_padding(x, router_logits, input_ids, all_rank_num_tokens) ) diff --git a/tests/unittest/_torch/moe/test_moe_backend.py b/tests/unittest/_torch/moe/test_moe_backend.py index 4be2fc113b4c..5f45ffeca896 100644 --- a/tests/unittest/_torch/moe/test_moe_backend.py +++ b/tests/unittest/_torch/moe/test_moe_backend.py @@ -2675,6 +2675,7 @@ def _fused_comm_scheduler( enable_attention_dp: bool = False, reduce_results: bool = False, accepts_global_routed_output: bool = False, + run_moe_reuses_buffers: bool = False, ) -> FusedCommMoEScheduler: mapping = Mapping( world_size=tp_size, @@ -2691,6 +2692,7 @@ def _fused_comm_scheduler( "reduce_results", "accepts_global_routed_output", "mapping", + "backend", ] ) moe.parallel_size = tp_size @@ -2698,6 +2700,10 @@ def _fused_comm_scheduler( moe.reduce_results = reduce_results moe.accepts_global_routed_output = accepts_global_routed_output moe.mapping = mapping + # ``forward`` reads the declaration off the backend before it chunks. + moe.backend = SimpleNamespace( + input_requirement=MoEInputRequirement(run_moe_reuses_buffers=run_moe_reuses_buffers) + ) return FusedCommMoEScheduler(moe) @@ -2798,6 +2804,26 @@ def test_fused_comm_forward_returns_what_the_caller_reduces( assert not output.any() +@pytest.mark.cpu_only +def test_fused_comm_forward_refuses_a_backend_that_reuses_its_buffers() -> None: + """``_run_chunks`` keeps the tensor each chunk returned, so a backend whose + next ``run_moe`` overwrites it would collect the same storage for every + chunk. Only ``ExternalCommMoEScheduler`` copies chunks out; here the + combination is a named refusal rather than the last chunk's rows repeated.""" + scheduler = _fused_comm_scheduler(rank=0, run_moe_reuses_buffers=True) + + with pytest.raises(NotImplementedError, match="run_moe_reuses_buffers"): + scheduler.forward( + torch.randn(5, 16, dtype=torch.bfloat16), + torch.randn(5, 8), + do_finalize=True, + output_dtype=torch.bfloat16, + all_rank_num_tokens=None, + use_dp_padding=None, + input_ids=None, + ) + + @pytest.mark.cpu_only @pytest.mark.parametrize( "scheduler_kind, accepts_global_routed_output, expected", From a7a17e4815502a5a3594f5d75745b59dbd25e027 Mon Sep 17 00:00:00 2001 From: xxi Date: Wed, 28 Oct 2026 11:11:12 +0000 Subject: [PATCH 3/7] [TRTLLM-14973][fix] Say why each B12x refusal fires Comment and message wording only, no behaviour change: - The token_final_scales refusal also fires on the None that apply_router_weight_on_input leaves behind, not only on a caller that wants the leaf to route; name both causes. - The output_dtype refusal cited a B12xMoEWrapper constructor argument this repository never passes. State the kernel fact and the leaf's own gate instead, and that this check is the backstop for a direct run_moe caller. - The W4A16 quantize_input docstring described the kernel's bf16 / fp16 range as if it were the leaf's; keep the kernel fact and add that the leaf admits bf16 only. Signed-off-by: xxi --- .../moe/fused_moe/fused_moe_cute_dsl_b12x.py | 17 ++++++++++++----- 1 file changed, 12 insertions(+), 5 deletions(-) diff --git a/tensorrt_llm/_torch/moe/fused_moe/fused_moe_cute_dsl_b12x.py b/tensorrt_llm/_torch/moe/fused_moe/fused_moe_cute_dsl_b12x.py index cd96c2c80c47..d4981b18bac2 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/fused_moe_cute_dsl_b12x.py +++ b/tensorrt_llm/_torch/moe/fused_moe/fused_moe_cute_dsl_b12x.py @@ -364,7 +364,9 @@ def _admit_run_context(self, ctx: MoERunContext) -> None: if ctx.token_selected_experts is None or ctx.token_final_scales is None: raise ValueError( f"{name} needs precomputed token_selected_experts and " - "token_final_scales; it cannot route from router_logits." + "token_final_scales; it cannot route from router_logits, and it " + "does not accept the None scales apply_router_weight_on_input " + "leaves behind once the scheduler has folded them into x." ) if require_comm_plan(self, ctx).enable_alltoall: raise NotImplementedError( @@ -389,8 +391,11 @@ def _run_b12x_decode(self, ctx: MoERunContext) -> torch.Tensor: f"{type(self).__name__} expects unquantized input (x_sf=None) " "on the b12x path; got a precomputed scale factor." ) - # ``B12xMoEWrapper`` writes bfloat16 and refuses any other - # ``output_dtype`` at construction. + # The b12x kernel writes bfloat16, which is why ``supported_dtypes`` + # admits bf16 activations alone: ``ConfigurableMoE`` derives + # ``output_dtype`` from the activation, so a layer built through the + # factory cannot get here. This is the backstop for a caller driving + # ``run_moe`` directly. if ctx.output_dtype not in (None, torch.bfloat16): raise ValueError( f"{type(self).__name__} writes bfloat16 on the b12x path; " @@ -536,8 +541,10 @@ def quantize_input( The wrapper runs in ``quant_mode="w4a16"``: it dequantizes the NVFP4 weights and leaves the bf16 / fp16 activation alone, so there is no - activation quantization here at all. The pre-quant scale is a separate - concern. An AWQ-style checkpoint stores ``W / s``, and + activation quantization here at all. That is the kernel's own range; + the leaf admits bf16 activations only (``supported_dtypes``), because + the wrapper hands the kernel a bfloat16 output buffer. The pre-quant + scale is a separate concern. An AWQ-style checkpoint stores ``W / s``, and ``load_quant_scales`` fills ``fc31_act_scale`` from its ``*.pre_quant_scale`` keys regardless of ``quant_algo``, so the weights the kernel dequantizes expect the scaled activation ``x * s``. Same as From 4ef7594c67e5223639eb8dce5c7f8fa85ee9a8f1 Mon Sep 17 00:00:00 2001 From: xxi Date: Wed, 28 Oct 2026 11:28:41 +0000 Subject: [PATCH 4/7] [TRTLLM-14973][fix] Keep the shared-buffer comments exact about the last chunk Self-review follow-up. Three places still said the scheduler copies EVERY chunk of a multi-chunk forward out of the reused buffer, which stopped being true once the last chunk's clone was dropped: the trailing torch.cat copies that one. Two of them are the correctness argument for folding every B12x layer onto one shared output buffer, so they have to be precise. Also state the rule once in the guide instead of appending a retraction to it. Signed-off-by: xxi --- tensorrt_llm/_torch/moe/fused_moe/MOE_DEVELOPER_GUIDE.md | 2 +- .../_torch/moe/fused_moe/fused_moe_cute_dsl_b12x.py | 3 ++- tensorrt_llm/_torch/moe/fused_moe/quantization.py | 8 ++++---- 3 files changed, 7 insertions(+), 6 deletions(-) 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 6302ce8b6913..302f1e1e62ec 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/MOE_DEVELOPER_GUIDE.md +++ b/tensorrt_llm/_torch/moe/fused_moe/MOE_DEVELOPER_GUIDE.md @@ -126,7 +126,7 @@ Adaptive quantize/dispatch order (gated by comm.supports_post_quant_dispatch()): Pre-quant flow: comm.dispatch() → quantize_input() (send raw, quantize locally) ``` -EPLB hooks fire only at the first/last chunk of the first/last `repeat_idx`. Multi-stream chunk overlap is enabled when `not enable_alltoall and aux_stream is not None`, unless the backend declares `input_requirement.run_moe_reuses_buffers`: its `run_moe` returns a view of a buffer the next call overwrites and keeps scratch across calls, so the scheduler runs its chunks on one stream and copies each chunk's output before the next chunk runs (the B12x leaves, whose `B12xMoEWrapper` output buffer is also shared across layers). The last chunk is not copied: the `torch.cat` that follows it copies it anyway. Only `ExternalCommMoEScheduler` honours the declaration; `FusedCommMoEScheduler` keeps every chunk output uncopied and refuses a backend that declares it. +EPLB hooks fire only at the first/last chunk of the first/last `repeat_idx`. Multi-stream chunk overlap is enabled when `not enable_alltoall and aux_stream is not None`, unless the backend declares `input_requirement.run_moe_reuses_buffers`: its `run_moe` returns a view of a buffer the next call overwrites and keeps scratch across calls (the B12x leaves, whose `B12xMoEWrapper` output buffer is also shared across layers). `ExternalCommMoEScheduler` then runs its chunks on one stream and copies every chunk but the last out before the next chunk runs — the trailing `torch.cat` copies the last one. `FusedCommMoEScheduler` copies no chunk output at all, so it refuses a backend that declares the field rather than aliasing its chunks. `unfinalized_combine_fn` is a constructor argument (`create_moe(..., unfinalized_combine_fn=...)`) for a model that needs a per-`(token, expert)` transform applied to each expert output *before* the top-k sum. Once registered it runs on **every** `do_finalize=False` forward, with or without a comm strategy, so the forward's return type (dense tensor vs. unfinalized triple) depends only on static config. Under a comm strategy the scheduler runs it on the dispatched rows and the combine reduces already-finalized partials — equal to finalizing after the combine, because each term depends on one row alone and non-local slots contribute exact zeros. That reordering is only sound for a combine that is an unweighted row sum over a dense per-token tensor; `DeepEPLowLatency.combine` is not one (it consumes expert-major rows and applies the routing weights itself), which is what `Communication.supports_finalize_before_combine()` declares. `ConfigurableMoE.validate_backend` rejects at construction: a backend without `supports_unfinalized_output()`, a comm without `supports_finalize_before_combine()`, and `apply_router_weight_on_input`. With no hook registered, `do_finalize=False` plus a comm strategy stays a named runtime refusal instead of a crash inside the comm layer. diff --git a/tensorrt_llm/_torch/moe/fused_moe/fused_moe_cute_dsl_b12x.py b/tensorrt_llm/_torch/moe/fused_moe/fused_moe_cute_dsl_b12x.py index d4981b18bac2..c1519514ebdf 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/fused_moe_cute_dsl_b12x.py +++ b/tensorrt_llm/_torch/moe/fused_moe/fused_moe_cute_dsl_b12x.py @@ -52,7 +52,8 @@ # run sequentially on the same CUDA stream, and the wrapper consumes its # previous output before the next layer is dispatched. Within one layer, the # scheduler runs the chunks of a multi-chunk forward on one stream and copies -# each chunk out, because the leaves declare ``run_moe_reuses_buffers``. +# every chunk but the last out of it -- the trailing ``torch.cat`` copies that +# one -- because the leaves declare ``run_moe_reuses_buffers``. _SHARED_MOE_OUTPUT_BUF: dict = {} # ActivationType -> b12x activation string. b12x currently exposes "relu2" diff --git a/tensorrt_llm/_torch/moe/fused_moe/quantization.py b/tensorrt_llm/_torch/moe/fused_moe/quantization.py index 0172e374af45..55b17d1883a1 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/quantization.py +++ b/tensorrt_llm/_torch/moe/fused_moe/quantization.py @@ -4068,10 +4068,10 @@ def transform_weights(self, module: torch.nn.Module) -> None: ) # Replace the wrapper's per-instance output buffer with a shared one. - # Layers run sequentially on a single stream, and the scheduler copies - # each chunk of a multi-chunk forward out of it (the backend declares - # ``run_moe_reuses_buffers``), so a single buffer of the right shape is - # correct and saves + # Layers run sequentially on a single stream, and a multi-chunk forward + # copies every chunk but the last out of it, the trailing ``torch.cat`` + # the last (the backend declares ``run_moe_reuses_buffers``), so a + # single buffer of the right shape is correct and saves # ``(num_moe_layers - 1) * max_num_tokens * hidden_size * 2`` bytes — # ~2.5 GB on Nemotron-Super-120B with ``max_num_tokens=2048``, # ``hidden=8192``, bf16, 80 MoE layers. From 7ce0f16840f1cdd3d2337497ffcce6e4e04c43ca Mon Sep 17 00:00:00 2001 From: xxi Date: Wed, 28 Oct 2026 06:07:21 +0000 Subject: [PATCH 5/7] [TRTLLM-14973][fix] Refuse only MoE LoRA on the layer in B12x run_moe lora_params describes every adapter of the forward, so refusing any non-None lora_params would also refuse a forward whose adapters target attention alone, which the B12x kernels leave untouched. Move the predicate behind CutlassMoELoraMixin._moe_lora_active into the module function moe_lora_active(lora_params, layer_idx), shared by the mixin and the B12x leaves, and refuse only when this layer carries routed-expert MoE LoRA tensors. Also drop the float16 wording from the W4A16 comment in the B12x weight method; the B12x leaves admit bfloat16 activations only. Signed-off-by: xxi --- .../_torch/moe/fused_moe/cutlass/lora.py | 52 +++++++++++-------- .../moe/fused_moe/fused_moe_cute_dsl_b12x.py | 16 ++++-- .../_torch/moe/fused_moe/quantization.py | 2 +- .../moe/test_cute_dsl_b12x_moe_backend.py | 42 ++++++++++++++- 4 files changed, 83 insertions(+), 29 deletions(-) diff --git a/tensorrt_llm/_torch/moe/fused_moe/cutlass/lora.py b/tensorrt_llm/_torch/moe/fused_moe/cutlass/lora.py index 67baf18c1ad8..25904b452231 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/cutlass/lora.py +++ b/tensorrt_llm/_torch/moe/fused_moe/cutlass/lora.py @@ -54,6 +54,36 @@ def raise_moe_lora_multichunk_unsupported(num_chunks: int) -> None: ) +def moe_lora_active(lora_params: dict | None, layer_idx: int | None) -> bool: + """Return True when ``lora_params`` carries routed-expert MoE LoRA tensors + for layer ``layer_idx``. + + ``lora_params`` describes every adapter of the forward, so a backend that + cannot fuse MoE LoRA asks this rather than refusing any ``lora_params``. + """ + if not lora_params or layer_idx is None: + return False + # CUDA-graph slot-indexed mode carries MoE LoRA in cuda_graph_params + # rather than a per-layer eager dict (mirrors _extract_moe_lora_tensors), + # so consult the graph layer map to keep the multi-chunk guard effective + # during capture/replay. + if lora_params.get("use_cuda_graph_mode", False): + cuda_graph_params = lora_params.get("cuda_graph_params") + if cuda_graph_params is None: + return False + layer_module2key = getattr(cuda_graph_params, "layer_module2key", {}) + return any( + (layer_idx, int(LoraModuleType.from_string(name))) in layer_module2key + for name in MOE_LORA_MODULE_NAMES + ) + layer_params = lora_params.get(layer_idx, {}) + if not layer_params: + return False + return any( + int(LoraModuleType.from_string(name)) in layer_params for name in MOE_LORA_MODULE_NAMES + ) + + class CutlassMoELoraMixin: """The LoRA half of a CUTLASS leaf that supports routed-expert LoRA. @@ -200,27 +230,7 @@ def _moe_lora_active(self, lora_params: dict | None) -> bool: """Return True when lora_params carries routed-expert MoE LoRA tensors for this layer, meaning run_moe would fuse a LoRA delta. """ - if not lora_params or self.layer_idx is None: - return False - # CUDA-graph slot-indexed mode carries MoE LoRA in cuda_graph_params - # rather than a per-layer eager dict (mirrors _extract_moe_lora_tensors), - # so consult the graph layer map to keep the multi-chunk guard effective - # during capture/replay. - if lora_params.get("use_cuda_graph_mode", False): - cuda_graph_params = lora_params.get("cuda_graph_params") - if cuda_graph_params is None: - return False - layer_module2key = getattr(cuda_graph_params, "layer_module2key", {}) - return any( - (self.layer_idx, int(LoraModuleType.from_string(name))) in layer_module2key - for name in MOE_LORA_MODULE_NAMES - ) - layer_params = lora_params.get(self.layer_idx, {}) - if not layer_params: - return False - return any( - int(LoraModuleType.from_string(name)) in layer_params for name in MOE_LORA_MODULE_NAMES - ) + return moe_lora_active(lora_params, self.layer_idx) @staticmethod def _empty_kernel_slot_dict() -> dict[str, torch.Tensor | None]: diff --git a/tensorrt_llm/_torch/moe/fused_moe/fused_moe_cute_dsl_b12x.py b/tensorrt_llm/_torch/moe/fused_moe/fused_moe_cute_dsl_b12x.py index c1519514ebdf..acfc12022944 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/fused_moe_cute_dsl_b12x.py +++ b/tensorrt_llm/_torch/moe/fused_moe/fused_moe_cute_dsl_b12x.py @@ -26,6 +26,7 @@ from .activation import MoEActivationSupport from .cutlass import CutlassFusedMoEBase, TrtllmCutlassNvfp4Impl from .cutlass.eligibility import BF16_ONLY, check_dtype, check_oss_cutlass_build +from .cutlass.lora import moe_lora_active from .impl_contract import ( MoEDeployment, MoEEligibility, @@ -347,9 +348,11 @@ def _admit_run_context(self, ctx: MoERunContext) -> None: """Refuse what neither route of the leaf honours, then join the input-ready event before any kernel reads ``ctx.x``. - On both routes the leaf finalizes, fuses no LoRA, routes nothing itself - and holds every expert (its gate rejects ``ep_size > 1``), so a context - asking otherwise fails here instead of having the request dropped. + On both routes the leaf finalizes, fuses no routed-expert LoRA, routes + nothing itself and holds every expert (its gate rejects + ``ep_size > 1``), so a context asking otherwise fails here instead of + having the request dropped. ``lora_params`` describes every adapter of + the forward, so only MoE LoRA on this layer is such a request. ``all_rank_num_tokens`` has no reader on either route: ``run_moe`` uses it only for the all-to-all tuner shapes, and all-to-all is refused. """ @@ -360,8 +363,11 @@ def _admit_run_context(self, ctx: MoERunContext) -> None: "(gemm2_output, expert_weights, expanded_idx_to_permuted_idx) " "triple that do_finalize=False asks for." ) - if ctx.lora_params is not None: - raise NotImplementedError(f"{name} does not fuse routed-expert LoRA; got lora_params.") + if ctx.lora_params is not None and moe_lora_active(ctx.lora_params, self.layer_idx): + raise NotImplementedError( + f"{name} does not fuse routed-expert LoRA; got MoE LoRA tensors " + f"for layer {self.layer_idx}." + ) if ctx.token_selected_experts is None or ctx.token_final_scales is None: raise ValueError( f"{name} needs precomputed token_selected_experts and " diff --git a/tensorrt_llm/_torch/moe/fused_moe/quantization.py b/tensorrt_llm/_torch/moe/fused_moe/quantization.py index 55b17d1883a1..b2caeebf34ad 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/quantization.py +++ b/tensorrt_llm/_torch/moe/fused_moe/quantization.py @@ -4016,7 +4016,7 @@ def transform_weights(self, module: torch.nn.Module) -> None: num_groups=num_local_experts) if is_w4a16_nvfp4: - # W4A16 path: BF16/FP16 activations multiplied by FP4 weights. + # W4A16 path: BF16 activations multiplied by FP4 weights. # FlashInfer's W4A16 packer expects the ModelOpt scale contract: # normalized FP8 block scales plus per-expert ``weight_global_scale``. w1_alpha_b12x = w1_w_scale_2.to(torch.float32).contiguous() diff --git a/tests/unittest/_torch/moe/test_cute_dsl_b12x_moe_backend.py b/tests/unittest/_torch/moe/test_cute_dsl_b12x_moe_backend.py index af51feff8946..074eec2cfc2c 100644 --- a/tests/unittest/_torch/moe/test_cute_dsl_b12x_moe_backend.py +++ b/tests/unittest/_torch/moe/test_cute_dsl_b12x_moe_backend.py @@ -44,6 +44,7 @@ NVFP4CuteDslB12xFusedMoEMethod, NVFP4CutlassFusedMoEMethod, ) +from tensorrt_llm._torch.peft.lora.layer import LoraModuleType from tensorrt_llm._torch.utils import ActivationType from tensorrt_llm.models.modeling_utils import QuantAlgo, QuantConfig @@ -418,6 +419,18 @@ def test_dispatch_rejects_non_tensor(): ] B12X_KERNEL_ROUTES = [route for route in B12X_ROUTES if route.values[2] == "b12x"] +_LAYER_IDX = 3 + + +def _moe_lora_params(layer_idx: int) -> dict: + """``lora_params`` carrying a routed-expert MoE LoRA module for + ``layer_idx``.""" + module_id = int(LoraModuleType.from_string("moe_h_to_4h")) + return { + "num_seqs": 1, + layer_idx: {module_id: {"adapter_size": None, "weight_pointers": None}}, + } + class _B12xWrapperRecorder: """Stands in for ``B12xMoEWrapper``, writing bfloat16 as its kernels do.""" @@ -445,7 +458,7 @@ def _leaf_with_recorded_kernels(leaf_cls: type, log: list, monkeypatch) -> CuteD ``run_moe`` -- only record that they ran. ``__new__`` without ``__init__``: the real constructor wants a GPU, while - ``run_moe`` reads only the wrapper and its weights. + ``run_moe`` reads only the layer index, the wrapper and its weights. """ def cutlass_run_moe(self, ctx: MoERunContext, *, workspace=None) -> torch.Tensor: @@ -454,6 +467,7 @@ def cutlass_run_moe(self, ctx: MoERunContext, *, workspace=None) -> torch.Tensor monkeypatch.setattr(TrtllmCutlassNvfp4Impl, "run_moe", cutlass_run_moe) leaf = leaf_cls.__new__(leaf_cls) + leaf.layer_idx = _LAYER_IDX leaf.b12x_wrapper = _B12xWrapperRecorder(log) leaf._b12x_weights = {} return leaf @@ -498,7 +512,9 @@ def test_run_moe_hands_back_the_kernel_output_untouched(leaf_cls, num_tokens, ke pytest.param( dict(do_finalize=False), NotImplementedError, "do_finalize=False", id="unfinalized" ), - pytest.param(dict(lora_params={"num_seqs": 1}), NotImplementedError, "LoRA", id="lora"), + pytest.param( + dict(lora_params=_moe_lora_params(_LAYER_IDX)), NotImplementedError, "LoRA", id="lora" + ), pytest.param( dict( token_selected_experts=None, @@ -539,6 +555,28 @@ def test_run_moe_refuses_what_neither_route_honours( assert log == [] +@pytest.mark.parametrize( + "lora_params", + [ + pytest.param({"num_seqs": 1}, id="no-layer-adapters"), + pytest.param(_moe_lora_params(_LAYER_IDX + 1), id="moe-lora-on-another-layer"), + ], +) +@pytest.mark.parametrize("leaf_cls,num_tokens,kernel", B12X_ROUTES) +def test_run_moe_admits_lora_params_without_moe_lora_on_its_layer( + leaf_cls, num_tokens, kernel, lora_params, monkeypatch +): + """``lora_params`` describes every adapter of the forward. Only + routed-expert LoRA on this layer asks the leaf for something it cannot + fuse; adapters elsewhere leave its output as it is.""" + log = [] + leaf = _leaf_with_recorded_kernels(leaf_cls, log, monkeypatch) + + leaf.run_moe(_run_context(num_tokens, lora_params=lora_params)) + + assert log == [kernel] + + @pytest.mark.parametrize("output_dtype", [torch.float16, torch.float32]) @pytest.mark.parametrize("leaf_cls,num_tokens,kernel", B12X_KERNEL_ROUTES) def test_b12x_route_refuses_an_output_dtype_it_cannot_write( From 4fac4c8f48b8df53be7c8e642768937af25dcf77 Mon Sep 17 00:00:00 2001 From: xxi Date: Wed, 28 Oct 2026 14:17:58 +0000 Subject: [PATCH 6/7] [TRTLLM-14973][fix] Refuse lora_params when the B12x leaf has no layer index ``moe_lora_active`` answers False for every ``lora_params`` when ``layer_idx`` is None (cutlass/lora.py:64-65), so narrowing the refusal to it turned the unknown-layer case into silent admission: a context that does carry MoE LoRA reached the kernel and came back without the delta, where the previous ``ctx.lora_params is not None`` test raised. ``layer_idx`` is None by default (create_moe.py:57, impl_base.py:112) and this refusal exists for the caller that drives ``run_moe`` directly, which is the caller most likely not to pass one. Keep the narrowing for a known layer and refuse any adapter when there is no index to read it by. An empty ``lora_params`` dict now passes instead of being refused: it asks for no fusion, which is what a known layer with no MoE adapter already does. Signed-off-by: xxi --- .../moe/fused_moe/fused_moe_cute_dsl_b12x.py | 24 ++++++++++++++----- .../moe/test_cute_dsl_b12x_moe_backend.py | 16 +++++++++++++ 2 files changed, 34 insertions(+), 6 deletions(-) diff --git a/tensorrt_llm/_torch/moe/fused_moe/fused_moe_cute_dsl_b12x.py b/tensorrt_llm/_torch/moe/fused_moe/fused_moe_cute_dsl_b12x.py index acfc12022944..3998e86770a1 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/fused_moe_cute_dsl_b12x.py +++ b/tensorrt_llm/_torch/moe/fused_moe/fused_moe_cute_dsl_b12x.py @@ -352,7 +352,8 @@ def _admit_run_context(self, ctx: MoERunContext) -> None: nothing itself and holds every expert (its gate rejects ``ep_size > 1``), so a context asking otherwise fails here instead of having the request dropped. ``lora_params`` describes every adapter of - the forward, so only MoE LoRA on this layer is such a request. + the forward, so only MoE LoRA on this layer is such a request -- with + no layer index to read it by, any adapter is. ``all_rank_num_tokens`` has no reader on either route: ``run_moe`` uses it only for the all-to-all tuner shapes, and all-to-all is refused. """ @@ -363,11 +364,22 @@ def _admit_run_context(self, ctx: MoERunContext) -> None: "(gemm2_output, expert_weights, expanded_idx_to_permuted_idx) " "triple that do_finalize=False asks for." ) - if ctx.lora_params is not None and moe_lora_active(ctx.lora_params, self.layer_idx): - raise NotImplementedError( - f"{name} does not fuse routed-expert LoRA; got MoE LoRA tensors " - f"for layer {self.layer_idx}." - ) + if ctx.lora_params: + # Reading the layer's adapters out of ``lora_params`` needs a layer + # index. With none the leaf cannot tell MoE LoRA from an adapter it + # may ignore, so it refuses rather than return output a caller + # reads as carrying a LoRA delta. + if self.layer_idx is None: + raise NotImplementedError( + f"{name} does not fuse routed-expert LoRA and has no " + "layer_idx to tell whether lora_params carries any; build " + "the leaf with layer_idx to narrow this to one layer." + ) + if moe_lora_active(ctx.lora_params, self.layer_idx): + raise NotImplementedError( + f"{name} does not fuse routed-expert LoRA; got MoE LoRA " + f"tensors for layer {self.layer_idx}." + ) if ctx.token_selected_experts is None or ctx.token_final_scales is None: raise ValueError( f"{name} needs precomputed token_selected_experts and " diff --git a/tests/unittest/_torch/moe/test_cute_dsl_b12x_moe_backend.py b/tests/unittest/_torch/moe/test_cute_dsl_b12x_moe_backend.py index 074eec2cfc2c..830769b23bd8 100644 --- a/tests/unittest/_torch/moe/test_cute_dsl_b12x_moe_backend.py +++ b/tests/unittest/_torch/moe/test_cute_dsl_b12x_moe_backend.py @@ -577,6 +577,22 @@ def test_run_moe_admits_lora_params_without_moe_lora_on_its_layer( assert log == [kernel] +@pytest.mark.parametrize("leaf_cls,num_tokens,kernel", B12X_ROUTES) +def test_run_moe_refuses_lora_params_when_it_has_no_layer_index( + leaf_cls, num_tokens, kernel, monkeypatch +): + """``layer_idx`` is what picks this layer's adapters out of ``lora_params``, + and it defaults to None. A leaf without one cannot tell MoE LoRA from an + adapter it may ignore, so it refuses instead of dropping the delta.""" + log = [] + leaf = _leaf_with_recorded_kernels(leaf_cls, log, monkeypatch) + leaf.layer_idx = None + + with pytest.raises(NotImplementedError, match="layer_idx"): + leaf.run_moe(_run_context(num_tokens, lora_params=_moe_lora_params(_LAYER_IDX))) + assert log == [] + + @pytest.mark.parametrize("output_dtype", [torch.float16, torch.float32]) @pytest.mark.parametrize("leaf_cls,num_tokens,kernel", B12X_KERNEL_ROUTES) def test_b12x_route_refuses_an_output_dtype_it_cannot_write( From b4416ef7d584773d58bcbbac89dcc7a9866160b4 Mon Sep 17 00:00:00 2001 From: xxi Date: Wed, 28 Oct 2026 14:30:46 +0000 Subject: [PATCH 7/7] [TRTLLM-14973][fix] Say what the no-layer-index refusal covers and where to fix it Three exactness fixes to the refusal added in the previous commit. The docstring and the comment said "any adapter", but the condition is the truthiness of ``lora_params``, which holds for a dict carrying only the per-forward ``num_seqs`` bookkeeping key and no adapter entry at all. Say "every ``lora_params``" so the prose and the condition agree. The message said "build the leaf with layer_idx". A layer built through ``ConfigurableMoE`` is handed ``layer_idx=None`` on purpose and gets the real index mirrored on afterwards (``_BACKEND_SYNC_ATTRS``), so that sentence points a reader at the one place it must not change. Name ``create_moe_backend``, the entry point that takes the index, and the other way out. The test pinned only the MoE-LoRA input. Parametrize it with the adapter-free dict too: that is the input the same leaf admits once it knows its layer, so it is where the two states differ. Signed-off-by: xxi --- .../moe/fused_moe/fused_moe_cute_dsl_b12x.py | 15 ++++++++------- .../_torch/moe/test_cute_dsl_b12x_moe_backend.py | 16 ++++++++++++---- 2 files changed, 20 insertions(+), 11 deletions(-) diff --git a/tensorrt_llm/_torch/moe/fused_moe/fused_moe_cute_dsl_b12x.py b/tensorrt_llm/_torch/moe/fused_moe/fused_moe_cute_dsl_b12x.py index 3998e86770a1..ecd71b4a1170 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/fused_moe_cute_dsl_b12x.py +++ b/tensorrt_llm/_torch/moe/fused_moe/fused_moe_cute_dsl_b12x.py @@ -353,7 +353,7 @@ def _admit_run_context(self, ctx: MoERunContext) -> None: ``ep_size > 1``), so a context asking otherwise fails here instead of having the request dropped. ``lora_params`` describes every adapter of the forward, so only MoE LoRA on this layer is such a request -- with - no layer index to read it by, any adapter is. + no layer index to read it by, any ``lora_params`` is. ``all_rank_num_tokens`` has no reader on either route: ``run_moe`` uses it only for the all-to-all tuner shapes, and all-to-all is refused. """ @@ -365,15 +365,16 @@ def _admit_run_context(self, ctx: MoERunContext) -> None: "triple that do_finalize=False asks for." ) if ctx.lora_params: - # Reading the layer's adapters out of ``lora_params`` needs a layer - # index. With none the leaf cannot tell MoE LoRA from an adapter it - # may ignore, so it refuses rather than return output a caller - # reads as carrying a LoRA delta. + # Reading the layer's adapters out of ``lora_params`` needs a + # layer index. With none the leaf cannot tell MoE LoRA from the + # adapters it may ignore, so it refuses every ``lora_params`` + # rather than return output a caller reads as adapted. if self.layer_idx is None: raise NotImplementedError( f"{name} does not fuse routed-expert LoRA and has no " - "layer_idx to tell whether lora_params carries any; build " - "the leaf with layer_idx to narrow this to one layer." + "layer_idx to tell whether lora_params carries any; pass " + "layer_idx to create_moe_backend, or hand run_moe no " + "lora_params." ) if moe_lora_active(ctx.lora_params, self.layer_idx): raise NotImplementedError( diff --git a/tests/unittest/_torch/moe/test_cute_dsl_b12x_moe_backend.py b/tests/unittest/_torch/moe/test_cute_dsl_b12x_moe_backend.py index 830769b23bd8..b5d9fecea6c2 100644 --- a/tests/unittest/_torch/moe/test_cute_dsl_b12x_moe_backend.py +++ b/tests/unittest/_torch/moe/test_cute_dsl_b12x_moe_backend.py @@ -577,19 +577,27 @@ def test_run_moe_admits_lora_params_without_moe_lora_on_its_layer( assert log == [kernel] +@pytest.mark.parametrize( + "lora_params", + [ + pytest.param(_moe_lora_params(_LAYER_IDX), id="moe-lora"), + pytest.param({"num_seqs": 1}, id="no-layer-adapters"), + ], +) @pytest.mark.parametrize("leaf_cls,num_tokens,kernel", B12X_ROUTES) def test_run_moe_refuses_lora_params_when_it_has_no_layer_index( - leaf_cls, num_tokens, kernel, monkeypatch + leaf_cls, num_tokens, kernel, lora_params, monkeypatch ): """``layer_idx`` is what picks this layer's adapters out of ``lora_params``, - and it defaults to None. A leaf without one cannot tell MoE LoRA from an - adapter it may ignore, so it refuses instead of dropping the delta.""" + and it defaults to None. A leaf without one cannot tell MoE LoRA from the + adapters it may ignore, so it refuses every ``lora_params`` -- including the + one the same leaf admits once it knows which layer it is.""" log = [] leaf = _leaf_with_recorded_kernels(leaf_cls, log, monkeypatch) leaf.layer_idx = None with pytest.raises(NotImplementedError, match="layer_idx"): - leaf.run_moe(_run_context(num_tokens, lora_params=_moe_lora_params(_LAYER_IDX))) + leaf.run_moe(_run_context(num_tokens, lora_params=lora_params)) assert log == []