From 8d5b347edd9b192c3a339d32a5af666af4400fe0 Mon Sep 17 00:00:00 2001 From: YAMY <74099316+YAMY1234@users.noreply.github.com> Date: Thu, 14 May 2026 19:49:14 -0700 Subject: [PATCH] Support Qwen3.5 NVFP4 MTP DeepEP (#24906) --- .../layers/attention/linear/gdn_backend.py | 8 ++- .../linear/kernels/gdn_flashinfer.py | 1 + python/sglang/srt/layers/moe/ep_moe/layer.py | 49 ++++++++++++++++++- .../srt/layers/moe/token_dispatcher/deepep.py | 5 +- 4 files changed, 58 insertions(+), 5 deletions(-) diff --git a/python/sglang/srt/layers/attention/linear/gdn_backend.py b/python/sglang/srt/layers/attention/linear/gdn_backend.py index 1f463430e..0e5453ebd 100644 --- a/python/sglang/srt/layers/attention/linear/gdn_backend.py +++ b/python/sglang/srt/layers/attention/linear/gdn_backend.py @@ -105,8 +105,12 @@ class GDNKernelDispatcher: else: raise ValueError(f"Unsupported GDN prefill backend: {prefill_backend}") - # Verify kernel: use FlashInfer if either decode or prefill selected it - if decode_backend.is_flashinfer() or prefill_backend.is_flashinfer(): + # Verify kernel: use FlashInfer only when the selected FlashInfer kernel + # supports MTP verify. On SM100+ FlashInfer GDN decode is supported, but + # its MTP verify path is not, so keep Triton as the verify fallback. + if ( + decode_backend.is_flashinfer() or prefill_backend.is_flashinfer() + ) and flashinfer_kernel.supports_target_verify: self.verify_kernel = flashinfer_kernel else: self.verify_kernel = triton_kernel diff --git a/python/sglang/srt/layers/attention/linear/kernels/gdn_flashinfer.py b/python/sglang/srt/layers/attention/linear/kernels/gdn_flashinfer.py index f0bf4a04a..a2324747f 100644 --- a/python/sglang/srt/layers/attention/linear/kernels/gdn_flashinfer.py +++ b/python/sglang/srt/layers/attention/linear/kernels/gdn_flashinfer.py @@ -98,6 +98,7 @@ class FlashInferGDNKernel(LinearAttnKernelBase): sm_major = torch.cuda.get_device_capability()[0] self.use_state_pool = sm_major != 9 + self.supports_target_verify = sm_major == 9 if sm_major == 9: if self._prefill_fn is None: diff --git a/python/sglang/srt/layers/moe/ep_moe/layer.py b/python/sglang/srt/layers/moe/ep_moe/layer.py index 15f5d0847..19ebc34e4 100644 --- a/python/sglang/srt/layers/moe/ep_moe/layer.py +++ b/python/sglang/srt/layers/moe/ep_moe/layer.py @@ -4,6 +4,7 @@ import logging from typing import TYPE_CHECKING, Any, Dict, Optional, Union import torch +import torch.nn.functional as F from sglang.srt.compilation.piecewise_context_manager import is_in_piecewise_cuda_graph from sglang.srt.environ import envs @@ -137,17 +138,23 @@ class DeepEPMoE(FusedMoE): self.deepep_mode = get_deepep_mode() + if quant_config is None and hasattr(self.dispatcher, "set_quant_config"): + self.dispatcher.set_quant_config({"bf16_dispatch": True}) + if ( self.deepep_mode.enable_low_latency() and not _is_npu and not _is_hip and not ( get_moe_runner_backend().is_flashinfer_cutedsl() + and self.quant_config is not None and self.quant_config.get_name() == "modelopt_fp4" ) + and quant_config is not None ): # AMD HIP, NPU supports low_latency deepep without deepgemm # NV FP4 quantization with flashinfer_cutedsl also supports low_latency deepep without deepgemm + # Unquantized draft MoE uses BF16 DeepEP dispatch and a local fallback. assert ( deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM ), f"DeepEP {self.deepep_mode} mode requires deep_gemm" @@ -237,13 +244,20 @@ class DeepEPMoE(FusedMoE): assert DispatchOutputChecker.format_is_deepep(dispatch_output) output = self.forward_npu(dispatch_output) elif DispatchOutputChecker.format_is_deepep_normal(dispatch_output): - if self.use_w4afp8: + if self.quant_config is None: + raise NotImplementedError( + "Unquantized DeepEP MoE currently supports low_latency mode only" + ) + elif self.use_w4afp8: output = self.forward_cutlass_w4afp8(dispatch_output) else: assert False, "forward_deepgemm_contiguous is deprecated" elif DispatchOutputChecker.format_is_deepep_ll(dispatch_output): - if ( + if self.quant_config is None: + output = self.forward_unquantized_deepep_ll(dispatch_output) + elif ( get_moe_runner_backend().is_flashinfer_cutedsl() + and self.quant_config is not None and self.quant_config.get_name() == "modelopt_fp4" ): output = self.forward_flashinfer_cutedsl(dispatch_output) @@ -314,6 +328,37 @@ class DeepEPMoE(FusedMoE): expert_mask=self.expert_mask, ) + def forward_unquantized_deepep_ll( + self, + dispatch_output: DeepEPLLDispatchOutput, + ): + hidden_states, hidden_states_scale, _, _, masked_m, _ = dispatch_output + assert hidden_states_scale is None + assert self.moe_runner_config.activation == "silu" + assert self.moe_runner_config.is_gated + assert hidden_states.dim() == 3 + + num_experts, max_tokens, _ = hidden_states.shape + token_offsets = torch.arange(max_tokens, device=hidden_states.device) + valid_mask = ( + token_offsets.unsqueeze(0) < masked_m[:num_experts].unsqueeze(1) + ).unsqueeze(-1) + hidden_states = hidden_states.masked_fill(~valid_mask, 0) + + gate_up = torch.bmm(hidden_states, self.w13_weight.transpose(1, 2)) + w13_bias = getattr(self, "w13_weight_bias", None) + if w13_bias is not None: + gate_up = gate_up + w13_bias.unsqueeze(1) + + gate, up = gate_up.chunk(2, dim=-1) + hidden_states = F.silu(gate) * up + + output = torch.bmm(hidden_states, self.w2_weight.transpose(1, 2)) + w2_bias = getattr(self, "w2_weight_bias", None) + if w2_bias is not None: + output = output + w2_bias.unsqueeze(1) + return output.masked_fill(~valid_mask, 0) + def forward_flashinfer_cutedsl( self, dispatch_output: DeepEPLLDispatchOutput, diff --git a/python/sglang/srt/layers/moe/token_dispatcher/deepep.py b/python/sglang/srt/layers/moe/token_dispatcher/deepep.py index a6d0d754c..25206b2bb 100644 --- a/python/sglang/srt/layers/moe/token_dispatcher/deepep.py +++ b/python/sglang/srt/layers/moe/token_dispatcher/deepep.py @@ -625,16 +625,19 @@ class _DeepEPDispatcherImplLowLatency(_DeepEPDispatcherImplBase): ): use_nvfp4 = use_fp8 = False input_global_scale = self.quant_config.get("input_global_scale", None) + bf16_dispatch = self.quant_config.get("bf16_dispatch", False) if input_global_scale is not None: use_nvfp4 = True else: backend = get_moe_runner_backend() # BF16 dispatch is needed when: + # - quant_config requests BF16 dispatch explicitly # - flashinfer_cutedsl: kernel quantizes to NVFP4 internally # - NPU with SGLANG_DEEPEP_BF16_DISPATCH: INT8 input + BF16 weight GMM not supported # - deep_gemm with SGLANG_DEEPEP_BF16_DISPATCH: user requests BF16 dispatch need_bf16_dispatch = ( - backend.is_flashinfer_cutedsl() + bf16_dispatch + or backend.is_flashinfer_cutedsl() or (_is_npu and envs.SGLANG_DEEPEP_BF16_DISPATCH.get()) or (backend.is_deep_gemm() and envs.SGLANG_DEEPEP_BF16_DISPATCH.get()) )