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..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`. +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. @@ -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/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 659a8d8b3881..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 @@ -25,7 +25,8 @@ 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 .cutlass.lora import moe_lora_active from .impl_contract import ( MoEDeployment, MoEEligibility, @@ -50,7 +51,10 @@ # 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 +# 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" @@ -154,8 +158,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 +174,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 +193,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 +242,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 +273,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 +344,58 @@ 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 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 -- with + 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. + """ + 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: + # 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; 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( + 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 " + "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( + 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 +411,16 @@ 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." ) + # 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; " + 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 +460,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 +506,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 @@ -472,8 +561,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 @@ -493,4 +584,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..a5406cc14416 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/impl_contract.py +++ b/tensorrt_llm/_torch/moe/fused_moe/impl_contract.py @@ -79,6 +79,13 @@ 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. ``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. @@ -327,6 +334,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..33b9aaf14488 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,10 @@ def _forward_multiple_chunks( row_offset += chunk_size_list[idx_chunk] if chunked_used[idx_chunk]: - outputs_list.append(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: @@ -1142,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 @@ -1188,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/tensorrt_llm/_torch/moe/fused_moe/quantization.py b/tensorrt_llm/_torch/moe/fused_moe/quantization.py index 764d45775cfb..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() @@ -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 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. 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..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 @@ -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 @@ -40,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 @@ -70,6 +75,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 +86,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 +170,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 +194,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 +403,246 @@ 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"] + +_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.""" + + 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 layer index, 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.layer_idx = _LAYER_IDX + 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=_moe_lora_params(_LAYER_IDX)), 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( + "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( + "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, 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 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=lora_params)) + 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..5f45ffeca896 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, @@ -2673,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, @@ -2689,6 +2692,7 @@ def _fused_comm_scheduler( "reduce_results", "accepts_global_routed_output", "mapping", + "backend", ] ) moe.parallel_size = tp_size @@ -2696,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) @@ -2796,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", @@ -5635,6 +5663,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:")