From 7119d5974798d157bc3a16cc9074e84c2b9c4c9a Mon Sep 17 00:00:00 2001 From: xieminghe1 <141820649+xieminghe1@users.noreply.github.com> Date: Mon, 30 Mar 2026 22:27:28 +0800 Subject: [PATCH] DeepSeek-R1-0528-w4a8: DeepEP Low Latency Dispatch Adopts FP8 Communication (#14162) Co-authored-by: undefined --- .../sglang/srt/layers/moe/cutlass_w4a8_moe.py | 22 ++++-- .../sglang/srt/layers/moe/ep_moe/kernels.py | 73 +++++++++++++++++++ python/sglang/srt/layers/moe/ep_moe/layer.py | 6 +- .../srt/layers/moe/token_dispatcher/deepep.py | 2 +- .../sglang/srt/layers/quantization/w4afp8.py | 3 +- 5 files changed, 94 insertions(+), 12 deletions(-) diff --git a/python/sglang/srt/layers/moe/cutlass_w4a8_moe.py b/python/sglang/srt/layers/moe/cutlass_w4a8_moe.py index 6ce5a3ddc..7e63d5c88 100644 --- a/python/sglang/srt/layers/moe/cutlass_w4a8_moe.py +++ b/python/sglang/srt/layers/moe/cutlass_w4a8_moe.py @@ -25,6 +25,7 @@ from sglang.srt.layers.moe.ep_moe.kernels import ( deepep_permute_triton_kernel, deepep_post_reorder_triton_kernel, deepep_run_moe_deep_preprocess, + fp8_per_token_to_per_tensor_quant_triton, post_reorder_for_cutlass_moe, pre_reorder_for_cutlass_moe, silu_and_mul_masked_post_per_tensor_quant_fwd, @@ -411,7 +412,8 @@ def cutlass_w4a8_moe_deepep_normal( def cutlass_w4a8_moe_deepep_ll( - a: torch.Tensor, + a_states: torch.Tensor, + a_scales: torch.Tensor, w1_q: torch.Tensor, w2_q: torch.Tensor, w1_scale: torch.Tensor, @@ -473,7 +475,7 @@ def cutlass_w4a8_moe_deepep_ll( """ assert w1_q.dtype == torch.int8 assert w2_q.dtype == torch.int8 - assert a.shape[2] // 2 == w1_q.shape[2], "Hidden size mismatch w1" + assert a_states.shape[2] // 2 == w1_q.shape[2], "Hidden size mismatch w1" assert w1_q.shape[2] * 2 == w2_q.shape[1], "Hidden size mismatch w2" assert w1_q.shape[0] == w2_q.shape[0], "Expert number mismatch" assert w1_q.shape[0] == w1_scale.shape[0], "w1 scales expert number mismatch" @@ -484,12 +486,12 @@ def cutlass_w4a8_moe_deepep_ll( assert a_strides2.shape[0] == w2_q.shape[0], "A Strides 2 expert number mismatch" assert b_strides2.shape[0] == w2_q.shape[0], "B Strides 2 expert number mismatch" num_experts = w1_q.size(0) - m = a.size(1) + m = a_states.size(1) k = w1_q.size(2) * 2 # w1_q is transposed and packed n = w2_q.size(2) * 2 # w2_q is transposed and packed topk = topk_ids_.size(1) - device = a.device + device = a_states.device problem_sizes1, problem_sizes2 = deepep_ll_get_cutlass_w4a8_moe_mm_data( masked_m, @@ -500,8 +502,14 @@ def cutlass_w4a8_moe_deepep_ll( k, ) - gateup_input = torch.empty(a.shape, dtype=torch.float8_e4m3fn, device=device) - per_tensor_quant_fp8(a, gateup_input, a1_scale.float(), True) + gateup_input = torch.empty(a_states.shape, dtype=torch.float8_e4m3fn, device=device) + fp8_per_token_to_per_tensor_quant_triton( + x=a_states, + x_scale=a_scales, + masked_m=masked_m, + output_scale=a1_scale, + output=gateup_input, + ) c1 = torch.empty((num_experts, m, n * 2), device=device, dtype=torch.bfloat16) c2 = torch.empty((num_experts, m, k), device=device, dtype=torch.bfloat16) @@ -522,7 +530,7 @@ def cutlass_w4a8_moe_deepep_ll( ) intermediate_q = torch.empty( - (num_experts, m, n), device=a.device, dtype=torch.float8_e4m3fn + (num_experts, m, n), device=a_states.device, dtype=torch.float8_e4m3fn ) silu_and_mul_masked_post_per_tensor_quant_fwd( c1, intermediate_q, masked_m, a2_scale diff --git a/python/sglang/srt/layers/moe/ep_moe/kernels.py b/python/sglang/srt/layers/moe/ep_moe/kernels.py index 044c590f2..4cd4b4f81 100644 --- a/python/sglang/srt/layers/moe/ep_moe/kernels.py +++ b/python/sglang/srt/layers/moe/ep_moe/kernels.py @@ -1381,3 +1381,76 @@ def silu_and_mul_masked_post_per_tensor_quant_fwd( NUM_STAGE=NUM_STAGES, ) return output + + +@triton.jit +def _fp8_per_token_quant_to_per_tensor_quant_kernel( + x_ptr, + x_scale_ptr, + x_scale_stride0, + x_scale_stride1, + x_scale_stride2, + masked_m_ptr, + output_scale_ptr, + output_ptr, + m, + k, + K_SCALE_BLOCK_SIZE: tl.constexpr, + K_BLOCK_SIZE: tl.constexpr, +): + pid_k, pid_m, pid_e = ( + tl.program_id(axis=0), + tl.program_id(axis=1), + tl.program_id(axis=2), + ) + pid_m_dim = tl.num_programs(1) + + token_id = pid_m + last_effective_id = tl.load(masked_m_ptr + pid_e) + + if token_id >= last_effective_id: + return + output_scale_val_inv = 1.0 / tl.load(output_scale_ptr).to(tl.float32) + k_offsets = pid_k * K_BLOCK_SIZE + tl.arange(0, K_BLOCK_SIZE) + scale_offsets = (k_offsets // K_SCALE_BLOCK_SIZE) * x_scale_stride2 + + x_ptrs = x_ptr + pid_e * m * k + k_offsets + output_ptrs = output_ptr + pid_e * m * k + k_offsets + x_scale_ptrs = x_scale_ptr + pid_e * x_scale_stride0 + scale_offsets + + for tok_idx in tl.range(token_id, last_effective_id, pid_m_dim): + hidden = tl.load(x_ptrs + tok_idx * k).to(tl.float32) + scale_fp32 = tl.load(x_scale_ptrs + tok_idx * x_scale_stride1).to(tl.float32) + hidden = hidden * scale_fp32 * output_scale_val_inv + tl.store(output_ptrs + tok_idx * k, hidden.to(output_ptr.dtype.element_ty)) + + +def fp8_per_token_to_per_tensor_quant_triton( + x: torch.Tensor, + x_scale: torch.Tensor, + masked_m: torch.Tensor, + output_scale: torch.Tensor, + output: torch.Tensor, +): + K_SCALE_BLOCK_SIZE = 128 + assert len(x.shape) == 3 and x.size(2) % K_SCALE_BLOCK_SIZE == 0 + assert x.is_contiguous() + assert x_scale.size(2) == x.size(2) // K_SCALE_BLOCK_SIZE + assert output_scale.numel() == 1 + + K_BLOCK_SIZE = 1024 + assert x.size(2) % K_BLOCK_SIZE == 0 + grid = (x.size(2) // K_BLOCK_SIZE, 32, x.size(0)) + _fp8_per_token_quant_to_per_tensor_quant_kernel[grid]( + x, + x_scale, + *x_scale.stride(), + masked_m, + output_scale, + output, + x.size(1), + x.size(2), + K_SCALE_BLOCK_SIZE=K_SCALE_BLOCK_SIZE, + K_BLOCK_SIZE=K_BLOCK_SIZE, + num_warps=8, + ) diff --git a/python/sglang/srt/layers/moe/ep_moe/layer.py b/python/sglang/srt/layers/moe/ep_moe/layer.py index ef2225ecc..10f825d4f 100644 --- a/python/sglang/srt/layers/moe/ep_moe/layer.py +++ b/python/sglang/srt/layers/moe/ep_moe/layer.py @@ -331,6 +331,9 @@ class DeepEPMoE(FusedMoE): ): assert self.moe_runner_config.activation == "silu" assert isinstance(self.quant_method, W4AFp8MoEMethod) + assert ( + envs.SGLANG_DEEPEP_BF16_DISPATCH.get() + ), "W4AFP8 does not support FP8 normal dispatch; please set SGLANG_DEEPEP_BF16_DISPATCH=1." return self.quant_method.apply_deepep_normal( layer=self, dispatch_output=dispatch_output, @@ -342,9 +345,6 @@ class DeepEPMoE(FusedMoE): ): assert self.moe_runner_config.activation == "silu" assert isinstance(self.quant_method, W4AFp8MoEMethod) - assert ( - envs.SGLANG_DEEPEP_BF16_DISPATCH.get() - ), "W4AFP8 does not support FP8 dispatch; please set SGLANG_DEEPEP_BF16_DISPATCH=1." return self.quant_method.apply_deepep_ll( layer=self, dispatch_output=dispatch_output, diff --git a/python/sglang/srt/layers/moe/token_dispatcher/deepep.py b/python/sglang/srt/layers/moe/token_dispatcher/deepep.py index 8539639d5..a4339808b 100644 --- a/python/sglang/srt/layers/moe/token_dispatcher/deepep.py +++ b/python/sglang/srt/layers/moe/token_dispatcher/deepep.py @@ -609,7 +609,7 @@ class _DeepEPDispatcherImplLowLatency(_DeepEPDispatcherImplBase): input_global_scale = self.quant_config.get("input_global_scale", None) if input_global_scale is not None: use_nvfp4 = True - elif not envs.SGLANG_DEEPEP_BF16_DISPATCH.get(): + else: use_fp8 = True buffer = self._get_buffer() diff --git a/python/sglang/srt/layers/quantization/w4afp8.py b/python/sglang/srt/layers/quantization/w4afp8.py index 83869289e..42b1dc24c 100644 --- a/python/sglang/srt/layers/quantization/w4afp8.py +++ b/python/sglang/srt/layers/quantization/w4afp8.py @@ -334,10 +334,11 @@ class W4AFp8MoEMethod(FusedMoEMethodBase): from sglang.srt.layers.moe.cutlass_w4a8_moe import cutlass_w4a8_moe_deepep_ll - hidden_states, _, topk_ids, _, masked_m, _ = dispatch_output + hidden_states, hidden_scales, topk_ids, _, masked_m, _ = dispatch_output output = cutlass_w4a8_moe_deepep_ll( hidden_states, + hidden_scales, layer.w13_weight, layer.w2_weight, layer.w13_weight_scale_inv,