diff --git a/python/sglang/srt/layers/moe/moe_runner/flashinfer_cutlass.py b/python/sglang/srt/layers/moe/moe_runner/flashinfer_cutlass.py index 950591ca1..1c4ad4404 100644 --- a/python/sglang/srt/layers/moe/moe_runner/flashinfer_cutlass.py +++ b/python/sglang/srt/layers/moe/moe_runner/flashinfer_cutlass.py @@ -1,9 +1,9 @@ """FlashInfer CUTLASS MoE fused funcs. This module owns the FlashInfer ``cutlass_fused_moe`` calls used by the -unquantized, ModelOpt FP8, ModelOpt NVFP4, and MXFP4 MoE paths. -Quantization methods prepare a small quant_info payload and route through -``MoeRunner``. +unquantized, ModelOpt FP8, ModelOpt NVFP4, and CUTLASS MXFP4 MoE paths, plus +the shared ``flashinfer_mxfp4`` dispatcher. Quantization methods prepare a +small quant_info payload and route through ``MoeRunner``. """ from __future__ import annotations @@ -312,14 +312,40 @@ def fused_experts_none_to_flashinfer_mxfp4( quant_info: MoeQuantInfo, runner_config: MoeRunnerConfig, ) -> StandardCombineInput: - """Run the FlashInfer CUTLASS MXFP4 fused experts.""" - from sglang.srt.layers.moe.token_dispatcher.standard import StandardCombineInput - from sglang.srt.layers.moe.topk import TopKOutputChecker + """Dispatch flashinfer_mxfp4 by quant-info type. - assert isinstance(quant_info, FlashInferCutlassMxfp4MoeQuantInfo), ( + Both mxfp4 paths register under this single ``("none", "flashinfer_mxfp4")`` + key but call different kernels. + """ + if isinstance(quant_info, FlashInferCutlassMxfp4MoeQuantInfo): + return _fused_experts_flashinfer_mxfp4_cutlass( + dispatch_output, quant_info, runner_config + ) + + # Keep one fused-op registration for the shared backend while loading the + # TRT-LLM implementation only when its quant-info type is dispatched. + from sglang.srt.layers.moe.moe_runner.flashinfer_trtllm import ( + FlashInferTrtllmGenMxfp4MoeQuantInfo, + _fused_experts_flashinfer_mxfp4_sm100_trtllm_gen, + ) + + if isinstance(quant_info, FlashInferTrtllmGenMxfp4MoeQuantInfo): + return _fused_experts_flashinfer_mxfp4_sm100_trtllm_gen( + dispatch_output, quant_info, runner_config + ) + raise TypeError( f"Unexpected quant_info type for flashinfer_mxfp4: {type(quant_info)}" ) + +def _fused_experts_flashinfer_mxfp4_cutlass( + dispatch_output: StandardDispatchOutput, + quant_info: FlashInferCutlassMxfp4MoeQuantInfo, + runner_config: MoeRunnerConfig, +) -> StandardCombineInput: + from sglang.srt.layers.moe.token_dispatcher.standard import StandardCombineInput + from sglang.srt.layers.moe.topk import TopKOutputChecker + flashinfer_cutlass_fused_moe, ActivationType = _flashinfer_cutlass_fused_moe() x = dispatch_output.hidden_states diff --git a/python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py b/python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py index caf6b6fe9..e1914d287 100644 --- a/python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py +++ b/python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py @@ -958,6 +958,242 @@ def fused_experts_none_to_flashinfer_trtllm_fp8( return StandardCombineInput(hidden_states=output) +@dataclass +class FlashInferTrtllmGenMxfp4MoeQuantInfo(MoeQuantInfo): + """Payload for the SM100 (Blackwell) trtllm-gen MXFP4 MoE path.""" + + # Packed MXFP4 weights: uint8, e2m1 x 2. + w13_weight: torch.Tensor + w2_weight: torch.Tensor + w13_weight_scale: torch.Tensor + w2_weight_scale: torch.Tensor + + # fp32 per expert. GPT-OSS sets these. + w13_weight_bias: torch.Tensor + w2_weight_bias: torch.Tensor + gemm1_alpha: torch.Tensor + gemm1_beta: torch.Tensor + gemm1_clamp_limit: torch.Tensor + + global_num_experts: int + local_expert_offset: int + local_num_experts: int + intermediate_size_per_partition: int + + hidden_size: int + flashinfer_mxfp4_moe_precision: str + routing_bias: Optional[torch.Tensor] = None + + +def _fused_experts_flashinfer_mxfp4_sm100_trtllm_gen( + dispatch_output: StandardDispatchOutput, + quant_info: FlashInferTrtllmGenMxfp4MoeQuantInfo, + runner_config: MoeRunnerConfig, +) -> StandardCombineInput: + """SM100 (Blackwell) trtllm-gen MXFP4 fused experts.""" + from sglang.srt.layers.moe.token_dispatcher.standard import StandardCombineInput + from sglang.srt.layers.moe.topk import TopKOutputChecker + + x = dispatch_output.hidden_states + topk_output = dispatch_output.topk_output + + origin_hidden_states_dim = x.shape[-1] + prepared_packed_topk = None + if quant_info.flashinfer_mxfp4_moe_precision == "bf16": + assert x.dtype == torch.bfloat16 + x_quant = x + x_scale = None + if quant_info.hidden_size != origin_hidden_states_dim: + x_quant = torch.nn.functional.pad( + x_quant, + (0, quant_info.hidden_size - origin_hidden_states_dim), + mode="constant", + value=0.0, + ) + elif quant_info.flashinfer_mxfp4_moe_precision == "default": + # Deferred import: mxfp4 dispatches into this module, and the helper + # stays there so its registered unit test keeps its import path. + from sglang.srt.layers.quantization.mxfp4 import ( + _prepare_flashinfer_mxfp8_activations, + ) + + x, prepared_packed_topk, x_quant, x_scale = ( + _prepare_flashinfer_mxfp8_activations(x, quant_info.hidden_size) + ) + else: + raise NotImplementedError( + f"Unsupported flashinfer_mxfp4_moe_precision: " + f"{quant_info.flashinfer_mxfp4_moe_precision}" + ) + + assert x_quant.shape[-1] == quant_info.hidden_size + is_standard = TopKOutputChecker.format_is_standard(topk_output) + assert is_standard or TopKOutputChecker.format_is_bypassed(topk_output), ( + f"unsupported topk format: {topk_output.format}" + ) + if is_standard: + assert runner_config.activation == "situ", ( + "standard topk output only wired for the situ path" + ) + top_k = topk_output.topk_ids.shape[1] + router_logits = None + else: + top_k = topk_output.topk_config.top_k + router_logits = topk_output.router_logits + + num_tokens = x_quant.shape[0] + from sglang.srt.layers import zero_copy_context + + symm_output = zero_copy_context.get_moe_output_spec( + torch.Size((num_tokens, origin_hidden_states_dim)), + torch.bfloat16, + x_quant.device, + ) + if symm_output is None: + with use_symmetric_memory( + get_tp_group(), disabled=not is_allocation_symmetric() + ): + symm_output = torch.empty( + num_tokens, + origin_hidden_states_dim, + dtype=torch.bfloat16, + device=x_quant.device, + ) + + if runner_config.activation == "situ": + from flashinfer import trtllm_fp4_block_scale_moe + from flashinfer.fused_moe import trtllm_fp4_block_scale_routed_moe + from flashinfer.tllm_enums import ActivationType, RoutingMethodType + + if is_standard: + if prepared_packed_topk is not None: + packed_topk = prepared_packed_topk + else: + packed_topk = PackTopkIds.execute( + topk_output.topk_ids, topk_output.topk_weights + ) + + defer_finalize = _deferred_finalize_enabled.get() + result = trtllm_fp4_block_scale_routed_moe( + topk_ids=packed_topk, + routing_bias=None, + hidden_states=x_quant, + hidden_states_scale=x_scale, + gemm1_weights=quant_info.w13_weight, + gemm1_weights_scale=quant_info.w13_weight_scale, + gemm1_bias=None, + gemm1_alpha=quant_info.gemm1_alpha, + gemm1_beta=quant_info.gemm1_clamp_limit, + gemm1_clamp_limit=None, + gemm2_weights=quant_info.w2_weight, + gemm2_weights_scale=quant_info.w2_weight_scale, + gemm2_bias=None, + output1_scale_scalar=None, + output1_scale_gate_scalar=None, + output2_scale_scalar=None, + num_experts=quant_info.global_num_experts, + top_k=packed_topk.shape[1], + n_group=None, + topk_group=None, + intermediate_size=quant_info.intermediate_size_per_partition, + local_expert_offset=quant_info.local_expert_offset, + local_num_experts=quant_info.local_num_experts, + routed_scaling_factor=None, + routing_method_type=RoutingMethodType.TopK.value, + activation_type=ActivationType.Situ.value, + tune_max_num_tokens=next_power_of_2(x_quant.shape[0]), + output=symm_output, + do_finalize=not defer_finalize, + enable_pdl=trtllm_moe_enable_pdl(x_quant.shape[0]), + ) + if defer_finalize: + gemm2_out, topk_weights, expanded_idx = result + result = FlashInferTrtllmDeferredFinalizeOutput( + gemm2_out=gemm2_out, + expert_weights=topk_weights, + expanded_idx_to_permuted_idx=expanded_idx, + top_k=packed_topk.shape[1], + ) + return StandardCombineInput(hidden_states=result) + # The finalized kernel writes to its explicit output argument. Do + # not propagate the FFI return tensor: some SiTU runner versions + # return a distinct wrapper/allocation even though symm_output + # contains the published result. Returning the destination makes + # the pointer contract explicit for K3's zero-copy latent buffer. + return StandardCombineInput(hidden_states=symm_output) + + trtllm_fp4_block_scale_moe( + routing_logits=router_logits.to(torch.bfloat16).contiguous(), + routing_bias=quant_info.routing_bias, + hidden_states=x_quant, + hidden_states_scale=x_scale, + gemm1_weights=quant_info.w13_weight, + gemm1_weights_scale=quant_info.w13_weight_scale, + gemm1_bias=None, + gemm1_alpha=quant_info.gemm1_alpha, + gemm1_beta=quant_info.gemm1_clamp_limit, + gemm1_clamp_limit=None, + gemm2_weights=quant_info.w2_weight, + gemm2_weights_scale=quant_info.w2_weight_scale, + gemm2_bias=None, + output1_scale_scalar=None, + output1_scale_gate_scalar=None, + output2_scale_scalar=None, + num_experts=quant_info.global_num_experts, + top_k=top_k, + n_group=topk_output.topk_config.num_expert_group, + topk_group=topk_output.topk_config.topk_group, + intermediate_size=quant_info.intermediate_size_per_partition, + routed_scaling_factor=( + topk_output.topk_config.routed_scaling_factor or 1.0 + ), + routing_method_type=RoutingMethodType.DeepSeekV3.value, + activation_type=ActivationType.Situ.value, + norm_topk_prob=topk_output.topk_config.renormalize, + local_expert_offset=quant_info.local_expert_offset, + local_num_experts=quant_info.local_num_experts, + tune_max_num_tokens=next_power_of_2(x_quant.shape[0]), + output=symm_output, + enable_pdl=trtllm_moe_enable_pdl(x_quant.shape[0]), + ) + return StandardCombineInput(hidden_states=symm_output) + + from flashinfer import trtllm_fp4_block_scale_moe + + trtllm_gen_output = trtllm_fp4_block_scale_moe( + router_logits.to(torch.bfloat16), + None, # routing_bias + x_quant, + x_scale, + quant_info.w13_weight, # uint8 (e2m1 x 2) + quant_info.w13_weight_scale, # uint8 (e4m3 x 2) + quant_info.w13_weight_bias, # fp32 per expert per channel + quant_info.gemm1_alpha, # fp32 per expert + quant_info.gemm1_beta, # fp32 per expert + quant_info.gemm1_clamp_limit, # fp32 per expert + quant_info.w2_weight, # uint8 (e2m1 x 2) + quant_info.w2_weight_scale, # ue8m0 + quant_info.w2_weight_bias, # fp32 per expert per channel + None, # output1_scale_scalar + None, # output1_scale_gate_scalar + None, # output2_scale_scalar + quant_info.global_num_experts, + top_k, + None, # n_group # TODO: support n_group + None, # topk_group # TODO: support topk_group + quant_info.intermediate_size_per_partition, # padded to multiple of 128 + quant_info.local_expert_offset, + quant_info.local_num_experts, + None, # routed_scaling_factor + 1, # routing_method_type, renormalize + True, # do finalize + tune_max_num_tokens=next_power_of_2(x_quant.shape[0]), + output=symm_output, + enable_pdl=trtllm_moe_enable_pdl(x_quant.shape[0]), + )[0] + return StandardCombineInput(hidden_states=trtllm_gen_output) + + @dataclass class FlashInferTrtllmFp4MoeQuantInfo(MoeQuantInfo): """Quantization payload consumed by FlashInfer TRT-LLM FP4 MoE kernels.""" diff --git a/python/sglang/srt/layers/quantization/mxfp4.py b/python/sglang/srt/layers/quantization/mxfp4.py index babe1de5d..b94894ea7 100644 --- a/python/sglang/srt/layers/quantization/mxfp4.py +++ b/python/sglang/srt/layers/quantization/mxfp4.py @@ -29,17 +29,11 @@ from torch.nn.parameter import Parameter # cutlass_fused_moe. Its C++ logger reads TLLM_LOG_LEVEL on first kernel launch; # setdefault preserves any explicit user override. os.environ.setdefault("TLLM_LOG_LEVEL", "INFO") -from sglang.srt.distributed import get_tp_group -from sglang.srt.distributed.device_communicators.pynccl_allocator import ( - use_symmetric_memory, -) from sglang.srt.environ import envs -from sglang.srt.layers import zero_copy_context from sglang.srt.layers.amx_utils import ( CPUQuantMethod, _amx_process_weight_after_loading, ) -from sglang.srt.layers.dp_attention import is_allocation_symmetric from sglang.srt.layers.moe import MoeRunner, MoeRunnerBackend, MoeRunnerConfig from sglang.srt.layers.moe.moe_runner.triton import TritonMoeQuantInfo from sglang.srt.layers.moe.utils import get_moe_a2a_backend, get_moe_runner_backend @@ -62,7 +56,6 @@ from sglang.srt.utils import ( is_hip, is_triton_kernels_available, is_xpu, - next_power_of_2, round_up, set_weight_attrs, use_intel_amx_backend, @@ -127,13 +120,10 @@ def _prepare_flashinfer_mxfp8_activations( if is_flashinfer_available(): from flashinfer import ( nvfp4_block_scale_interleave, - trtllm_fp4_block_scale_moe, ) - from flashinfer.fused_moe import ( - trtllm_fp4_block_scale_routed_moe, + from flashinfer.fused_moe.core import ( + get_w2_permute_indices_with_cache, ) - from flashinfer.fused_moe.core import get_w2_permute_indices_with_cache - from flashinfer.tllm_enums import ActivationType, RoutingMethodType # SM90 mixed-input helpers landed in FlashInfer #3084 (post-0.6.10). Older # versions don't ship them; gate at import so unrelated code paths still load. @@ -729,6 +719,7 @@ class Mxfp4MoEMethod(FusedMoEMethodBase): torch.tensor([_limit] * E, dtype=torch.float32).cuda(), requires_grad=False, ) + layer._situ_routing_bias_bf16 = None sf_block_size = 32 # mxfp4 block size assert ( @@ -1415,6 +1406,11 @@ class Mxfp4MoEMethod(FusedMoEMethodBase): ): self.moe_runner_config = moe_runner_config moe_runner_backend = get_moe_runner_backend() + if self.use_mega_moe and moe_runner_backend.is_flashinfer_mxfp4(): + # MegaMoE uses this method for weight preparation only and calls + # DeepGEMM directly instead of dispatching through a MoeRunner. + return + if moe_runner_backend.is_auto(): # Must match apply() priority: _use_aiter before use_triton_kernels. if _use_aiter and get_moe_a2a_backend().supports_aiter(): @@ -1442,16 +1438,14 @@ class Mxfp4MoEMethod(FusedMoEMethodBase): elif moe_runner_backend.is_flashinfer_mxfp4() and self._fi_kernel in ( "cutlass_sm90", "cutlass_sm120", + "trtllm_sm100", ): - # Register the fused func at runner construction so the FusedOpPool - # lookup at `MoeRunner.__init__` finds it. - import sglang.srt.layers.moe.moe_runner.flashinfer_cutlass # noqa: F401 - self.runner = MoeRunner(moe_runner_backend, moe_runner_config) else: - # Legacy bypass path (e.g. SM100 trtllm-gen under flashinfer_mxfp4) - # routes through `apply` without a MoeRunner. TODO(cwan): migrate. - pass + raise NotImplementedError( + f"Mxfp4MoEMethod has no MoeRunner for backend={moe_runner_backend} " + f"/ _fi_kernel={self._fi_kernel}." + ) def _apply_sm90_cutlass(self, layer, dispatch_output): """SM90 MXFP4 x BF16/FP8 MoE via FlashInfer's mixed-input kernels. @@ -1534,6 +1528,44 @@ class Mxfp4MoEMethod(FusedMoEMethodBase): ) return self.runner.run(dispatch_output, quant_info) + def _apply_sm100_trtllm_gen(self, layer, dispatch_output): + from sglang.srt.layers.moe.moe_runner.flashinfer_trtllm import ( + FlashInferTrtllmGenMxfp4MoeQuantInfo, + ) + from sglang.srt.layers.moe.topk import TopKOutputChecker + + routing_bias = layer._situ_routing_bias_bf16 + topk_output = dispatch_output.topk_output + if ( + self.moe_runner_config.activation == "situ" + and routing_bias is None + and TopKOutputChecker.format_is_bypassed(topk_output) + ): + correction_bias = topk_output.topk_config.correction_bias + if correction_bias is not None: + routing_bias = correction_bias.to(torch.bfloat16) + layer._situ_routing_bias_bf16 = routing_bias + + quant_info = FlashInferTrtllmGenMxfp4MoeQuantInfo( + w13_weight=layer.w13_weight, + w2_weight=layer.w2_weight, + w13_weight_scale=layer.w13_weight_scale, + w2_weight_scale=layer.w2_weight_scale, + w13_weight_bias=layer.w13_weight_bias, + w2_weight_bias=layer.w2_weight_bias, + gemm1_alpha=layer.gemm1_alpha, + gemm1_beta=layer.gemm1_beta, + gemm1_clamp_limit=layer.gemm1_clamp_limit, + global_num_experts=layer.num_experts, + local_expert_offset=layer.moe_ep_rank * layer.num_local_experts, + local_num_experts=layer.num_local_experts, + intermediate_size_per_partition=self.intermediate_size_per_partition, + hidden_size=self.hidden_size, + flashinfer_mxfp4_moe_precision=self.flashinfer_mxfp4_moe_precision, + routing_bias=routing_bias, + ) + return self.runner.run(dispatch_output, quant_info) + def apply( self, layer: torch.nn.Module, @@ -1649,237 +1681,7 @@ class Mxfp4MoEMethod(FusedMoEMethodBase): if self._fi_kernel == "cutlass_sm120": return self._apply_sm120_cutlass(layer, dispatch_output) if self.use_flashinfer: - from sglang.srt.layers.moe.moe_runner.flashinfer_trtllm import ( - trtllm_moe_enable_pdl, - ) - - # When bf16 mode is enabled, we don't need to quantize the input, - # TRT-LLM automatically handles quantization in the kernel implementation and pipelines it with GEMM operations, - # which can theoretically improve performance - origin_hidden_states_dim = x.shape[-1] - # Filled by the staged K3 route+pack+quant fusion below; the pack - # site further down falls back to PackTopkIds when it is None. - prepared_packed_topk = None - if self.flashinfer_mxfp4_moe_precision == "bf16": - assert x.dtype == torch.bfloat16 - x_quant = x - x_scale = None - - # May be fused later if this code branch is frequently needed - if self.hidden_size != origin_hidden_states_dim: - x_quant = torch.nn.functional.pad( - x_quant, - (0, self.hidden_size - origin_hidden_states_dim), - mode="constant", - value=0.0, - ) - elif self.flashinfer_mxfp4_moe_precision == "default": - x, prepared_packed_topk, x_quant, x_scale = ( - _prepare_flashinfer_mxfp8_activations(x, self.hidden_size) - ) - else: - raise NotImplementedError() - - assert x_quant.shape[-1] == self.hidden_size - is_standard = TopKOutputChecker.format_is_standard(topk_output) - # The situ path accepts precomputed (standard) routing; the - # public path below is logits-only. - assert is_standard or TopKOutputChecker.format_is_bypassed(topk_output), ( - f"unsupported topk format: {topk_output.format}" - ) - if is_standard: - assert self.moe_runner_config.activation == "situ", ( - "standard topk output only wired for the situ path" - ) - top_k = topk_output.topk_ids.shape[1] - router_logits = None - else: - top_k = topk_output.topk_config.top_k - router_logits = topk_output.router_logits - - num_tokens = x_quant.shape[0] - hidden_size = origin_hidden_states_dim - # The K3 fused-front path publishes its [latent | shared] buffer - # slice as the output destination (zero_copy_context); writing - # the finalize output there directly skips this allocation and - # the copy_ back in _forward_fused. The slice lives in the same - # symmetric buffer the caller all-reduces. - symm_output = zero_copy_context.get_moe_output_spec( - torch.Size((num_tokens, hidden_size)), - torch.bfloat16, - x_quant.device, - ) - if symm_output is None: - with use_symmetric_memory( - get_tp_group(), disabled=not is_allocation_symmetric() - ): - symm_output = torch.empty( - num_tokens, - hidden_size, - dtype=torch.bfloat16, - device=x_quant.device, - ) - - if self.moe_runner_config.activation == "situ": - # FlashInfer 0.6.17+ ships the SiTU TRT-LLM-gen kernels. - # Routing must be noaux_tc (sigmoid + correction bias, - # DeepSeekV3), not the renormalize-softmax default below. - # EP is cubin-internal: each rank computes its local expert slice - # [offset, +num_local) and the caller all-reduces. ep=1 -> TP path. - local_expert_offset = layer.moe_ep_rank * layer.num_local_experts - if TopKOutputChecker.format_is_standard(topk_output): - # Precomputed routing (radix router upstream): skip the - # in-op routing kernels entirely. At small T the in-op - # single-CTA routing costs ~22 us/layer vs ~6 us for the - # external radix router. - if prepared_packed_topk is not None: - packed_topk = prepared_packed_topk - else: - from sglang.kernels.ops.moe.pack_topk_ids import PackTopkIds - - packed_topk = PackTopkIds.execute( - topk_output.topk_ids, topk_output.topk_weights - ) - # Deferred finalize (K3 forward_deferred_finalize): return - # the finalize inputs instead of the finalized output. - from sglang.srt.layers.moe.moe_runner.flashinfer_trtllm import ( - _deferred_finalize_enabled, - ) - - defer_finalize = _deferred_finalize_enabled.get() - result = trtllm_fp4_block_scale_routed_moe( - topk_ids=packed_topk, - routing_bias=None, - hidden_states=x_quant, - hidden_states_scale=x_scale, - gemm1_weights=layer.w13_weight, - gemm1_weights_scale=layer.w13_weight_scale, - gemm1_bias=None, - gemm1_alpha=layer.gemm1_alpha, - # SiTU beta is the linear-half tanh clip; K3 stores it - # in gemm1_clamp_limit. - gemm1_beta=layer.gemm1_clamp_limit, - gemm1_clamp_limit=None, - gemm2_weights=layer.w2_weight, - gemm2_weights_scale=layer.w2_weight_scale, - gemm2_bias=None, - output1_scale_scalar=None, - output1_scale_gate_scalar=None, - output2_scale_scalar=None, - num_experts=layer.num_experts, - top_k=packed_topk.shape[1], - n_group=None, - topk_group=None, - intermediate_size=self.intermediate_size_per_partition, - local_expert_offset=local_expert_offset, - local_num_experts=layer.num_local_experts, - routed_scaling_factor=None, - routing_method_type=RoutingMethodType.TopK.value, - activation_type=ActivationType.Situ.value, - tune_max_num_tokens=next_power_of_2(x_quant.shape[0]), - output=symm_output, - do_finalize=not defer_finalize, - enable_pdl=trtllm_moe_enable_pdl(x_quant.shape[0]), - ) - if defer_finalize: - from sglang.srt.layers.moe.moe_runner.flashinfer_trtllm import ( - FlashInferTrtllmDeferredFinalizeOutput, - ) - - gemm2_out, topk_weights, expanded_idx = result - result = FlashInferTrtllmDeferredFinalizeOutput( - gemm2_out=gemm2_out, - expert_weights=topk_weights, - expanded_idx_to_permuted_idx=expanded_idx, - top_k=packed_topk.shape[1], - ) - return StandardCombineInput(hidden_states=result) - # The finalized kernel writes to its explicit output - # argument. Do not propagate the FFI return tensor: some - # SiTU runner versions return a distinct wrapper/allocation - # even though symm_output contains the published result. - # Returning the destination makes the pointer contract - # explicit for K3's zero-copy latent buffer. - return StandardCombineInput(hidden_states=symm_output) - - # Bypassed topk: route from logits inside the op. - correction_bias = topk_output.topk_config.correction_bias - bias_bf16 = getattr(layer, "_situ_routing_bias_bf16", None) - if bias_bf16 is None and correction_bias is not None: - bias_bf16 = correction_bias.to(torch.bfloat16) - layer._situ_routing_bias_bf16 = bias_bf16 - trtllm_fp4_block_scale_moe( - # router_logits is a row-strided slice of the K3 fused - # front GEMM output; the FFI reads it as dense. - routing_logits=router_logits.to(torch.bfloat16).contiguous(), - routing_bias=bias_bf16, - hidden_states=x_quant, - hidden_states_scale=x_scale, - gemm1_weights=layer.w13_weight, - gemm1_weights_scale=layer.w13_weight_scale, - gemm1_bias=None, - gemm1_alpha=layer.gemm1_alpha, - # SiTU beta is the linear-half tanh clip; K3 stores it in - # gemm1_clamp_limit (situ_linear_beta). - gemm1_beta=layer.gemm1_clamp_limit, - gemm1_clamp_limit=None, - gemm2_weights=layer.w2_weight, - gemm2_weights_scale=layer.w2_weight_scale, - gemm2_bias=None, - output1_scale_scalar=None, - output1_scale_gate_scalar=None, - output2_scale_scalar=None, - num_experts=layer.num_experts, - top_k=top_k, - n_group=topk_output.topk_config.num_expert_group, - topk_group=topk_output.topk_config.topk_group, - intermediate_size=self.intermediate_size_per_partition, - routed_scaling_factor=( - topk_output.topk_config.routed_scaling_factor or 1.0 - ), - routing_method_type=RoutingMethodType.DeepSeekV3.value, - activation_type=ActivationType.Situ.value, - norm_topk_prob=topk_output.topk_config.renormalize, - local_expert_offset=local_expert_offset, - local_num_experts=layer.num_local_experts, - tune_max_num_tokens=next_power_of_2(x_quant.shape[0]), - output=symm_output, - enable_pdl=trtllm_moe_enable_pdl(x_quant.shape[0]), - ) - return StandardCombineInput(hidden_states=symm_output) - - trtllm_gen_output = trtllm_fp4_block_scale_moe( - router_logits.to(torch.bfloat16), - None, # routing_bias - x_quant, - x_scale, - layer.w13_weight, # uint8 (e2m1 x 2) - layer.w13_weight_scale, # uint8 (e4m3 x 2) - layer.w13_weight_bias, # fp32 per expert per channel - layer.gemm1_alpha, # fp32 per expert - layer.gemm1_beta, # fp32 per expert - layer.gemm1_clamp_limit, # fp32 per expert - layer.w2_weight, # uint8 (e2m1 x 2) - layer.w2_weight_scale, # ue8m0 - layer.w2_weight_bias, # fp32 per expert per channel - None, # output1_scale_scalar - None, # output1_scale_gate_scalar - None, # output2_scale_scalar - layer.num_experts, - top_k, - None, # n_group # TODO: support n_group - None, # topk_group # TODO: support topk_group - self.intermediate_size_per_partition, # padded to multiple of 256 - layer.moe_ep_rank * layer.num_local_experts, # local_expert_offset - layer.num_local_experts, # local num experts - None, # routed_scaling_factor - 1, # routing_method_type, renormalize - True, # do finalize - tune_max_num_tokens=next_power_of_2(x_quant.shape[0]), - output=symm_output, - enable_pdl=trtllm_moe_enable_pdl(x_quant.shape[0]), - )[0] - return StandardCombineInput(hidden_states=trtllm_gen_output) + return self._apply_sm100_trtllm_gen(layer, dispatch_output) if _use_aiter: return self._apply_aiter(layer, dispatch_output) diff --git a/test/registered/unit/layers/quantization/test_mxfp4_situ_output.py b/test/registered/unit/layers/quantization/test_mxfp4_situ_output.py index f5ed4c6bb..88be7f014 100644 --- a/test/registered/unit/layers/quantization/test_mxfp4_situ_output.py +++ b/test/registered/unit/layers/quantization/test_mxfp4_situ_output.py @@ -1,5 +1,5 @@ import sys -from types import SimpleNamespace +from types import ModuleType, SimpleNamespace from unittest.mock import patch import pytest @@ -16,9 +16,19 @@ pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires def test_situ_routed_moe_returns_published_output_buffer(): - from sglang.srt.layers.moe import route_quant_handoff + # Import mxfp4 before flashinfer_trtllm to avoid the pre-existing + # compressed_tensors circular import. + # isort: off from sglang.srt.layers.quantization import mxfp4 as mxfp4_module - from sglang.srt.layers.quantization.mxfp4 import Mxfp4MoEMethod + from sglang.srt.layers.moe.moe_runner import ( + flashinfer_trtllm as flashinfer_trtllm_module, + ) + from sglang.srt.layers.moe.moe_runner.base import MoeRunnerConfig + from sglang.srt.layers.moe.moe_runner.flashinfer_trtllm import ( + FlashInferTrtllmGenMxfp4MoeQuantInfo, + _fused_experts_flashinfer_mxfp4_sm100_trtllm_gen, + ) + # isort: on tokens, hidden, top_k = 3, 128, 2 x = torch.randn(tokens, hidden, dtype=torch.bfloat16, device="cuda") @@ -38,27 +48,23 @@ def test_situ_routed_moe_returns_published_output_buffer(): topk_output=topk_output, ) - method = Mxfp4MoEMethod.__new__(Mxfp4MoEMethod) - method.use_deep_gemm = False - method.use_marlin = False - method.use_flashinfer = True - method._fi_kernel = None - method.flashinfer_mxfp4_moe_precision = "default" - method.hidden_size = hidden - method.intermediate_size_per_partition = 128 - method.moe_runner_config = SimpleNamespace(activation="situ") - dummy = torch.empty(1, dtype=torch.uint8, device="cuda") - layer = SimpleNamespace( - moe_ep_rank=0, - num_local_experts=1, - num_experts=1, + quant_info = FlashInferTrtllmGenMxfp4MoeQuantInfo( w13_weight=dummy, - w13_weight_scale=dummy, - gemm1_alpha=None, - gemm1_clamp_limit=None, w2_weight=dummy, + w13_weight_scale=dummy, w2_weight_scale=dummy, + w13_weight_bias=dummy, + w2_weight_bias=dummy, + gemm1_alpha=dummy, + gemm1_beta=dummy, + gemm1_clamp_limit=dummy, + global_num_experts=1, + local_expert_offset=0, + local_num_experts=1, + intermediate_size_per_partition=128, + hidden_size=hidden, + flashinfer_mxfp4_moe_precision="default", ) expected = ( torch.arange(tokens * hidden, dtype=torch.float32, device="cuda") @@ -74,33 +80,40 @@ def test_situ_routed_moe_returns_published_output_buffer(): returned_ptr = ffi_result.data_ptr() return ffi_result + flashinfer = ModuleType("flashinfer") + flashinfer.__path__ = [] + flashinfer.trtllm_fp4_block_scale_moe = None + fused_moe = ModuleType("flashinfer.fused_moe") + fused_moe.trtllm_fp4_block_scale_routed_moe = fake_routed_moe + tllm_enums = ModuleType("flashinfer.tllm_enums") + tllm_enums.RoutingMethodType = SimpleNamespace(TopK=SimpleNamespace(value=0)) + tllm_enums.ActivationType = SimpleNamespace(Situ=SimpleNamespace(value=0)) + latent = torch.empty_like(x) with ( patch.object( - route_quant_handoff, - "take", - return_value=(packed_topk, x_quant, x_scale), + mxfp4_module, + "_prepare_flashinfer_mxfp8_activations", + return_value=(x, packed_topk, x_quant, x_scale), ), - patch( - "sglang.srt.layers.quantization.mxfp4.trtllm_fp4_block_scale_routed_moe", - side_effect=fake_routed_moe, - create=True, + patch.dict( + "sys.modules", + { + "flashinfer": flashinfer, + "flashinfer.fused_moe": fused_moe, + "flashinfer.tllm_enums": tllm_enums, + }, ), patch.object( - mxfp4_module, - "RoutingMethodType", - SimpleNamespace(TopK=SimpleNamespace(value=0)), - create=True, - ), - patch.object( - mxfp4_module, - "ActivationType", - SimpleNamespace(Situ=SimpleNamespace(value=0)), - create=True, + flashinfer_trtllm_module, "trtllm_moe_enable_pdl", return_value=False ), zero_copy_context.set_moe_output(latent), ): - combine_input = method.apply(layer, dispatch_output) + combine_input = _fused_experts_flashinfer_mxfp4_sm100_trtllm_gen( + dispatch_output, + quant_info, + MoeRunnerConfig(activation="situ"), + ) assert returned_ptr is not None assert returned_ptr != latent.data_ptr() diff --git a/test/registered/unit/layers/quantization/test_mxfp4_sm100_trtllm_gen.py b/test/registered/unit/layers/quantization/test_mxfp4_sm100_trtllm_gen.py new file mode 100644 index 000000000..d6912c27b --- /dev/null +++ b/test/registered/unit/layers/quantization/test_mxfp4_sm100_trtllm_gen.py @@ -0,0 +1,335 @@ +"""Focused SM100 trtllm-gen MXFP4 MoE regression test. + +``Mxfp4MoEMethod.apply`` (SM100 branch, via the unified MoeRunner) must feed +``trtllm_fp4_block_scale_moe`` the same args a direct kernel call does, so the +two outputs stay bit-exact. + +Fixtures are raw checkpoint-order MXFP4, converted by the production +``process_weights_after_loading`` so the kernel sees the interleaved weights and +float8_e4m3fn block scales it gets in a real run. +""" + +from __future__ import annotations + +from contextlib import nullcontext + +import pytest +import torch +from flashinfer import trtllm_fp4_block_scale_moe + +from sglang.srt.utils import is_sm100_supported +from sglang.test.ci.ci_register import register_cuda_ci + +register_cuda_ci(est_time=120, stage="base-b", runner_config="4-gpu-b200") + +if not is_sm100_supported(): + pytest.skip( + reason="trtllm-gen MXFP4 requires SM100 (Blackwell).", + allow_module_level=True, + ) + +GROUP_SIZE = 32 # MXFP4 block size + + +class _MockLayer: + """Hand-built ``FusedMoE`` stand-in (avoids distributed init).""" + + +def _make_random_mxfp4(num_experts, hidden, inter, seed=0): + g = torch.Generator(device="cuda").manual_seed(seed) + w13 = torch.randint( + 0, + 256, + (num_experts, 2 * inter, hidden // 2), + dtype=torch.uint8, + device="cuda", + generator=g, + ) + w2 = torch.randint( + 0, + 256, + (num_experts, hidden, inter // 2), + dtype=torch.uint8, + device="cuda", + generator=g, + ) + # E8M0 scales centered around 127 (= 2^0); narrow band keeps dequant values + # sane so the SwiGLU clamp doesn't dominate. + w13_s = torch.randint( + 125, + 130, + (num_experts, 2 * inter, hidden // GROUP_SIZE), + dtype=torch.uint8, + device="cuda", + generator=g, + ) + w2_s = torch.randint( + 125, + 130, + (num_experts, hidden, inter // GROUP_SIZE), + dtype=torch.uint8, + device="cuda", + generator=g, + ) + w13_b = ( + torch.randn( + num_experts, 2 * inter, dtype=torch.float32, device="cuda", generator=g + ) + * 0.01 + ) + w2_b = ( + torch.randn( + num_experts, hidden, dtype=torch.float32, device="cuda", generator=g + ) + * 0.01 + ) + return w13, w2, w13_s, w2_s, w13_b, w2_b + + +def _build_mock_layer(num_experts, hidden, inter, fixtures): + """Raw checkpoint-order weights; ``process_weights_after_loading`` converts + them in place and seeds the SwiGLU scalars, so nothing is pre-applied here.""" + w13, w2, w13_s, w2_s, w13_b, w2_b = fixtures + layer = _MockLayer() + layer.w13_weight = torch.nn.Parameter(w13.clone(), requires_grad=False) + layer.w2_weight = torch.nn.Parameter(w2.clone(), requires_grad=False) + layer.w13_weight_scale = torch.nn.Parameter(w13_s.clone(), requires_grad=False) + layer.w2_weight_scale = torch.nn.Parameter(w2_s.clone(), requires_grad=False) + layer.w13_weight_bias = torch.nn.Parameter(w13_b.clone(), requires_grad=False) + layer.w2_weight_bias = torch.nn.Parameter(w2_b.clone(), requires_grad=False) + layer.num_experts = num_experts + layer.num_local_experts = num_experts # tests run with EP size = 1 + layer.moe_ep_rank = 0 + return layer + + +def _build_method(num_experts, hidden, inter, precision): + from sglang.srt.layers.quantization.mxfp4 import Mxfp4MoEMethod + + method = Mxfp4MoEMethod.__new__(Mxfp4MoEMethod) + method._fi_kernel = "trtllm_sm100" + method.use_flashinfer = True + method.use_marlin = False + method.use_deep_gemm = False + method.use_mega_moe = False + method.num_experts = num_experts + method.hidden_size = hidden + method.intermediate_size_per_partition = inter + method.flashinfer_mxfp4_moe_precision = precision + method.runner = _build_flashinfer_mxfp4_runner(num_experts, hidden, inter) + method.moe_runner_config = method.runner.config + return method + + +def _build_flashinfer_mxfp4_runner(num_experts, hidden, inter): + # Bypass create_moe_runner (needs a live server arg context); the fused func + # only reads dispatch_output / quant_info, so a minimal config suffices. + import sglang.srt.layers.moe.moe_runner.flashinfer_cutlass # noqa: F401 + from sglang.srt.layers.moe.moe_runner.base import MoeRunnerConfig + from sglang.srt.layers.moe.moe_runner.runner import MoeRunner + from sglang.srt.layers.moe.utils import MoeRunnerBackend + + cfg = MoeRunnerConfig( + num_experts=num_experts, + num_local_experts=num_experts, + hidden_size=hidden, + intermediate_size_per_partition=inter, + top_k=None, + activation="silu", + is_gated=True, + ) + return MoeRunner(MoeRunnerBackend.FLASHINFER_MXFP4, cfg) + + +@pytest.mark.parametrize("use_mega_moe", [False, True]) +def test_create_moe_runner_handles_flashinfer_for_megamoe(monkeypatch, use_mega_moe): + import sglang.srt.layers.quantization.mxfp4 as mxfp4_mod + from sglang.srt.layers.moe.moe_runner.base import MoeRunnerConfig + from sglang.srt.layers.moe.utils import MoeRunnerBackend + + runner = object() + + def build_runner(backend, config): + assert not use_mega_moe + assert backend == MoeRunnerBackend.FLASHINFER_MXFP4 + assert config is runner_config + return runner + + monkeypatch.setattr( + mxfp4_mod, + "get_moe_runner_backend", + lambda: MoeRunnerBackend.FLASHINFER_MXFP4, + ) + monkeypatch.setattr(mxfp4_mod, "MoeRunner", build_runner) + + method = mxfp4_mod.Mxfp4MoEMethod.__new__(mxfp4_mod.Mxfp4MoEMethod) + method._fi_kernel = "trtllm_sm100" + method.use_mega_moe = use_mega_moe + runner_config = MoeRunnerConfig() + + method.create_moe_runner(object(), runner_config) + + assert method.moe_runner_config is runner_config + if use_mega_moe: + # FusedMoEMethodBase declares ``runner: MoeRunner | None = None``, so the + # early return leaves the class default rather than no attribute at all. + assert method.runner is None + else: + assert method.runner is runner + + +class _MockDispatchOutput: + # SM100 keeps BYPASSED topk (kernel routes from router_logits), so the + # dispatch output must carry a real BypassedTopKOutput. + def __init__(self, hidden_states, router_logits, top_k): + from sglang.srt.layers.moe.topk import BypassedTopKOutput, TopKConfig + + self.hidden_states = hidden_states + self.topk_output = BypassedTopKOutput( + hidden_states=hidden_states, + router_logits=router_logits, + topk_config=TopKConfig(top_k=top_k, renormalize=True), + ) + + +def _quant_input(x, precision, hidden_size): + # Mirror the SM100 helper's input-quant branch so the reference feeds the + # kernel the same x_quant / x_scale the SGLang path does. + origin = x.shape[-1] + if precision == "bf16": + x_quant = x + x_scale = None + if hidden_size != origin: + x_quant = torch.nn.functional.pad( + x_quant, (0, hidden_size - origin), mode="constant", value=0.0 + ) + elif precision == "default": + if x.shape[-1] == hidden_size: + if x.dim() > 2: + x = x.view(-1, x.shape[-1]) + from sglang.kernels.ops.quantization.per_token_group_quant import ( + per_token_group_quant, + ) + + x_quant, x_scale = per_token_group_quant(x, group_size=32, scale_ue8m0=True) + x_scale = x_scale.view(torch.float8_e4m3fn) + else: + from sglang.srt.layers.quantization.fp8_utils import ( + flashinfer_mxfp8_quantize, + ) + + x_quant, x_scale = flashinfer_mxfp8_quantize( + x, False, alignment=hidden_size + ) + x_scale = x_scale.view(torch.float8_e4m3fn).reshape(*x.shape[:-1], -1) + else: + raise AssertionError(precision) + return x_quant, x_scale + + +def _ref_trtllm(x, layer, method, precision, top_k, router_logits): + # Direct kernel call mirroring the SM100 helper's arg list. + from sglang.srt.layers.moe.moe_runner.flashinfer_trtllm import ( + trtllm_moe_enable_pdl, + ) + from sglang.srt.utils.common import next_power_of_2 + + x_quant, x_scale = _quant_input(x, precision, method.hidden_size) + # zeros, not empty: the output is compared bit-exact, so any row the kernel + # leaves unwritten must not carry allocator garbage. + out = torch.zeros( + x_quant.shape[0], x.shape[-1], dtype=torch.bfloat16, device=x_quant.device + ) + return trtllm_fp4_block_scale_moe( + router_logits.to(torch.bfloat16), + None, + x_quant, + x_scale, + layer.w13_weight, + layer.w13_weight_scale, + layer.w13_weight_bias, + layer.gemm1_alpha, + layer.gemm1_beta, + layer.gemm1_clamp_limit, + layer.w2_weight, + layer.w2_weight_scale, + layer.w2_weight_bias, + None, + None, + None, + layer.num_experts, + top_k, + None, + None, + method.intermediate_size_per_partition, + layer.moe_ep_rank * layer.num_local_experts, + layer.num_local_experts, + None, + 1, + True, + tune_max_num_tokens=next_power_of_2(x_quant.shape[0]), + output=out, + enable_pdl=trtllm_moe_enable_pdl(x_quant.shape[0]), + )[0] + + +@pytest.mark.parametrize("precision", ["default", "bf16"]) +@pytest.mark.parametrize( + "tokens,num_experts,hidden,inter,top_k", + [ + (4, 4, 256, 256, 2), + (16, 8, 512, 512, 2), + (32, 8, 1024, 1024, 4), + ], +) +def test_apply_trtllm_gen_matches_flashinfer_direct( + tokens, num_experts, hidden, inter, top_k, precision, monkeypatch +): + """``Mxfp4MoEMethod.apply`` (SM100 branch) must produce the same output as a + direct ``trtllm_fp4_block_scale_moe`` call fed the same inputs. + + Turns red if apply mis-wires a ``FlashInferTrtllmGenMxfp4MoeQuantInfo`` field + into the kernel (e.g. swapping ``local_expert_offset`` / ``local_num_experts`` + or dropping the bf16-vs-default input-quant branch).""" + method = _build_method(num_experts, hidden, inter, precision) + + import sglang.srt.layers.moe.moe_runner.flashinfer_trtllm as fi_trtllm_mod + + # Bypass symmetric-memory / TP-group in the fused-func module, where the + # kernel call now lives. + monkeypatch.setattr( + fi_trtllm_mod, "use_symmetric_memory", lambda *a, **kw: nullcontext() + ) + monkeypatch.setattr(fi_trtllm_mod, "is_allocation_symmetric", lambda: False) + monkeypatch.setattr(fi_trtllm_mod, "get_tp_group", lambda: None) + + fixtures = _make_random_mxfp4(num_experts, hidden, inter) + x = torch.randn(tokens, hidden, dtype=torch.bfloat16, device="cuda") * 0.1 + g = torch.Generator(device="cuda").manual_seed(1234) + router_logits = torch.randn( + tokens, num_experts, dtype=torch.float32, device="cuda", generator=g + ) + layer = _build_mock_layer(num_experts, hidden, inter, fixtures) + layer.moe_runner_config = method.moe_runner_config + + # Convert via the production path so the fixtures can't drift from it. + method.process_weights_after_loading(layer) + + # ---- FlashInfer-direct reference ---- + out_ref = _ref_trtllm(x, layer, method, precision, top_k, router_logits) + + # ---- SGLang path (same x + router_logits) ---- + out_sglang = method.apply( + layer, _MockDispatchOutput(x.clone(), router_logits, top_k) + ).hidden_states + + assert torch.equal(out_sglang, out_ref), ( + f"SGLang vs FlashInfer-direct mismatch (precision={precision}); " + f"max abs diff = {(out_sglang.float() - out_ref.float()).abs().max().item():.4g}" + ) + + +if __name__ == "__main__": + import sys + + sys.exit(pytest.main([__file__, "-v"])) diff --git a/test/registered/unit/layers/quantization/test_mxfp4_sm90_cutlass.py b/test/registered/unit/layers/quantization/test_mxfp4_sm90_cutlass.py index f07503795..e6cb81f2e 100644 --- a/test/registered/unit/layers/quantization/test_mxfp4_sm90_cutlass.py +++ b/test/registered/unit/layers/quantization/test_mxfp4_sm90_cutlass.py @@ -345,15 +345,9 @@ def test_apply_sm90_cutlass_matches_flashinfer_direct( here we just verify that ``apply`` calls the kernel with the right arguments (incl. input padding + output trim).""" import sglang.srt.layers.moe.moe_runner.flashinfer_cutlass as fi_cutlass_mod - import sglang.srt.layers.quantization.mxfp4 as mxfp4_mod - # Bypass symmetric-memory / TP-group in both the legacy quant_method and - # the new fused-func module (where the kernel call now lives). - monkeypatch.setattr( - mxfp4_mod, "use_symmetric_memory", lambda *a, **kw: nullcontext() - ) - monkeypatch.setattr(mxfp4_mod, "is_allocation_symmetric", lambda: False) - monkeypatch.setattr(mxfp4_mod, "get_tp_group", lambda: None) + # Bypass symmetric-memory / TP-group in the fused-func module, which is where + # the kernel call lives. monkeypatch.setattr( fi_cutlass_mod, "use_symmetric_memory", lambda *a, **kw: nullcontext() )