DeepSeek-R1-0528-w4a8: DeepEP Low Latency Dispatch Adopts FP8 Communication (#14162)

Co-authored-by: undefined <zhouchen.arrebol@jd.com>
2026-03-30 22:27:28 +08:00
parent 673ffb3116
commit 7119d59747
5 changed files with 94 additions and 12 deletions
@@ -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
@@ -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,
+    )
@@ -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,
@@ -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()
@@ -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,