[Reland] DeepSeek-R1-0528-w4a8: DeepEP Low Latency Dispatch Adopts FP8 Communication (#22316)

Co-authored-by: undefined <zhouchen.arrebol@jd.com>
Co-authored-by: xq25478 <xq25478@qq.com>

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

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,
+    )

python/sglang/srt/layers/moe/ep_moe/layer.py

@@ -342,9 +342,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,

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()

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,