Support DeepSeek-R1 w4a8 low latency deepep

python/sglang/srt/layers/moe/cutlass_w4a8_moe.py

@@ -11,6 +11,7 @@
 )
 
 from sglang.srt.layers.moe.ep_moe.kernels import (
+    deepep_ll_get_cutlass_w4a8_moe_mm_data,
     post_reorder_triton_kernel,
     pre_reorder_triton_kernel_for_cutlass_moe,
     run_cutlass_moe_ep_preproess,
@@ -43,6 +44,7 @@ def cutlass_w4a8_moe(
     a1_scale: Optional[torch.Tensor] = None,
     a2_scale: Optional[torch.Tensor] = None,
     apply_router_weight_on_input: bool = False,
+    ep_mode: str = "ep",
 ) -> torch.Tensor:
     """
     This function computes a w4a8-quantized Mixture of Experts (MoE) layer
@@ -83,10 +85,14 @@ def cutlass_w4a8_moe(
     Returns:
     - torch.Tensor: The fp8 output tensor after applying the MoE layer.
     """
-    assert topk_weights.shape == topk_ids_.shape, "topk shape mismatch"
+    assert (
+        topk_weights.shape == topk_ids_.shape if topk_weights is not None else True
+    ), "topk shape mismatch"
     assert w1_q.dtype == torch.int8
     assert w2_q.dtype == torch.int8
-    assert a.shape[1] // 2 == w1_q.shape[2], "Hidden size mismatch w1"
+    assert (
+        a.shape[1] // 2 == w1_q.shape[2] if ep_mode != "deepep_ll" else True
+    ), "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"
@@ -108,52 +114,79 @@ def cutlass_w4a8_moe(
     m = a.size(0)
     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)
+    topk = topk_ids_.size(1) if ep_mode == "ep" else 1
 
     if apply_router_weight_on_input:
         assert topk == 1, "apply_router_weight_on_input is only implemented for topk=1"
 
     device = a.device
 
-    _, src2dst, _ = run_cutlass_moe_ep_preproess(
-        local_topk_ids,
-        num_experts,
-    )
+    if ep_mode == "ep":
+        _, src2dst, _ = run_cutlass_moe_ep_preproess(
+            local_topk_ids,
+            num_experts,
+        )
 
-    gateup_input = torch.empty(
-        (m * topk, k),
-        device=device,
-        dtype=torch.float8_e4m3fn,
-    )
+        gateup_input = torch.empty(
+            (m * topk, k),
+            device=device,
+            dtype=torch.float8_e4m3fn,
+        )
 
-    pre_reorder_triton_kernel_for_cutlass_moe[(m,)](
-        a,
-        gateup_input,
-        src2dst,
-        local_topk_ids,
-        a1_scale,
-        total_num_experts,
-        topk,
-        k,
-        BLOCK_SIZE=512,
-    )
+        pre_reorder_triton_kernel_for_cutlass_moe[(m,)](
+            a,
+            gateup_input,
+            src2dst,
+            local_topk_ids,
+            a1_scale,
+            total_num_experts,
+            topk,
+            k,
+            BLOCK_SIZE=512,
+        )
+    elif ep_mode == "deepep_ll":
+        num_tokens = a.size(1)
+
+    else:
+        raise ValueError(f"Invalid ep_mode: {ep_mode}")
 
     # NOTE: a_map and c_map are not used in the get_cutlass_w4a8_moe_mm_data kernel,
     # they are kept to allow for a quick switch of the permutation logic
     # from the current triton kernel implementation to the cutlass-based one if needed.
-    a_map = torch.empty((local_topk_ids.numel()), dtype=torch.int32, device=device)
-    c_map = torch.empty((local_topk_ids.numel()), dtype=torch.int32, device=device)
-    get_cutlass_w4a8_moe_mm_data(
-        local_topk_ids,
-        expert_offsets,
-        problem_sizes1,
-        problem_sizes2,
-        a_map,
-        c_map,
-        num_experts,
-        n,
-        k,
-    )
+    if ep_mode == "deepep_ll":
+        gateup_input_origin, expert_offsets, problem_sizes1, problem_sizes2 = (
+            deepep_ll_get_cutlass_w4a8_moe_mm_data(
+                a,
+                local_topk_ids,
+                expert_offsets,
+                problem_sizes1,
+                problem_sizes2,
+                num_experts,
+                n,
+                k,
+            )
+        )
+        gateup_input = torch.empty(
+            gateup_input_origin.shape, dtype=torch.float8_e4m3fn, device=device
+        )
+        sgl_per_tensor_quant_fp8(
+            gateup_input_origin, gateup_input, a1_scale.float(), True
+        )
+
+    else:
+        a_map = torch.empty((local_topk_ids.numel()), dtype=torch.int32, device=device)
+        c_map = torch.empty((local_topk_ids.numel()), dtype=torch.int32, device=device)
+        get_cutlass_w4a8_moe_mm_data(
+            local_topk_ids,
+            expert_offsets,
+            problem_sizes1,
+            problem_sizes2,
+            a_map,
+            c_map,
+            num_experts,
+            n,
+            k,
+        )
 
     c1 = torch.empty((m * topk, n * 2), device=device, dtype=torch.half)
     c2 = torch.zeros((m * topk, k), device=device, dtype=torch.half)
@@ -197,19 +230,33 @@ def cutlass_w4a8_moe(
         128,
         topk,
     )
-
-    output = torch.empty_like(a)
-    post_reorder_triton_kernel[(m,)](
-        c2,
-        output,
-        src2dst,
-        topk_ids_,
-        topk_weights,
-        start_expert_id,
-        end_expert_id,
-        topk,
-        k,
-        0,
-        BLOCK_SIZE=512,
-    )
+    if ep_mode == "ep":
+        output = torch.empty_like(a)
+        post_reorder_triton_kernel[(m,)](
+            c2,
+            output,
+            src2dst,
+            local_topk_ids,
+            topk_weights,
+            start_expert_id,
+            end_expert_id,
+            topk,
+            k,
+            0,
+            BLOCK_SIZE=512,
+        )
+    elif ep_mode == "deepep_ll":
+        output = torch.zeros(
+            (len(local_topk_ids), num_tokens, k), device=device, dtype=c2.dtype
+        )
+        non_zero_indices = torch.nonzero(local_topk_ids, as_tuple=True)[0]
+        c2_index = 0
+        for expert_idx in non_zero_indices:
+            num_non_zero_rows = local_topk_ids[expert_idx].item()
+            output[expert_idx, :num_non_zero_rows] = c2[
+                c2_index : c2_index + num_non_zero_rows
+            ]
+            c2_index += num_non_zero_rows
+    else:
+        output = c2
     return output

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

@@ -1319,3 +1319,75 @@ def moe_ep_deepgemm_preprocess(
         gateup_input,
         gateup_input_scale,
     )
+
+
+@triton.jit
+def compute_problem_sizes_w4a8_kernel(
+    masked_m_ptr,
+    problem_sizes1_ptr,
+    problem_sizes2_ptr,
+    n,
+    k,
+    num_experts,
+    BLOCK_SIZE: tl.constexpr,
+):
+    pid = tl.program_id(axis=0) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+    mask = pid < num_experts
+    final_occurrences = tl.load(masked_m_ptr + pid, mask=mask, other=0)
+    tl.store(problem_sizes1_ptr + pid * 3, 2 * n)
+    tl.store(problem_sizes1_ptr + pid * 3 + 1, final_occurrences)
+    tl.store(problem_sizes1_ptr + pid * 3 + 2, k)
+    tl.store(problem_sizes2_ptr + pid * 3, k)
+    tl.store(problem_sizes2_ptr + pid * 3 + 1, final_occurrences)
+    tl.store(problem_sizes2_ptr + pid * 3 + 2, n)
+
+
+def compute_problem_sizes_w4a8(
+    masked_m, problem_sizes1, problem_sizes2, n, k, num_experts
+):
+    BLOCK_SIZE = 256
+    grid = lambda meta: (triton.cdiv(num_experts, meta["BLOCK_SIZE"]),)
+    compute_problem_sizes_w4a8_kernel[grid](
+        masked_m,
+        problem_sizes1,
+        problem_sizes2,
+        n,
+        k,
+        num_experts,
+        BLOCK_SIZE=BLOCK_SIZE,
+    )
+    return problem_sizes1, problem_sizes2
+
+
+def deepep_ll_get_cutlass_w4a8_moe_mm_data(
+    hidden_states,
+    masked_m,
+    expert_offsets,
+    problem_sizes1,
+    problem_sizes2,
+    num_experts,
+    n,
+    k,
+):
+    problem_sizes1, problem_sizes2 = compute_problem_sizes_w4a8(
+        masked_m, problem_sizes1, problem_sizes2, n, k, num_experts
+    )
+    masked_m_with_zero = torch.cat(
+        [torch.tensor([0], device=masked_m.device, dtype=masked_m.dtype), masked_m],
+        dim=0,
+    )
+    expert_offsets = torch.cumsum(masked_m_with_zero, dim=0)
+    expert_indices = torch.nonzero(masked_m, as_tuple=True)[0]
+    hidden_states_real = hidden_states[expert_indices]
+    hidden_states_real_reshaped = hidden_states_real.view(
+        -1, hidden_states_real.size(-1)
+    )
+    non_zero_rows_mask = (hidden_states_real_reshaped != 0).any(dim=-1)
+    hidden_states_real_reshaped = hidden_states_real_reshaped[non_zero_rows_mask]
+    logger.info(f"masked_m {masked_m}")
+    return (
+        hidden_states_real_reshaped,
+        expert_offsets.to(torch.int32),
+        problem_sizes1.to(torch.int32),
+        problem_sizes2.to(torch.int32),
+    )

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

@@ -47,13 +47,15 @@
     get_bool_env_var,
     is_hip,
     is_npu,
+    set_weight_attrs,
 )
 
 _is_hip = is_hip()
 _is_npu = is_npu()
 _is_fp8_fnuz = is_fp8_fnuz()
 _use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip
 
+
 if not (_is_npu or _is_hip):
     from sgl_kernel import silu_and_mul
 
@@ -954,10 +956,52 @@ def forward(
             else:
                 return self.forward_normal(hidden_states, reorder_topk_ids, seg_indptr)
         elif resolved_deepep_mode == DeepEPMode.low_latency:
-            return self.forward_deepgemm_masked(hidden_states, masked_m, expected_m)
+            if self.use_w4afp8:
+                return self.forward_cutlass_w4a8_masked(
+                    hidden_states, masked_m, ep_mode="deepep_ll"
+                )
+            else:
+                return self.forward_deepgemm_masked(hidden_states, masked_m, expected_m)
         else:
             raise ValueError(f"Invalid deepep_mode: {self.deepep_mode}")
 
+    def forward_cutlass_w4a8_masked(
+        self, hidden_states: torch.Tensor, masked_m: torch.Tensor, ep_mode: str
+    ):
+
+        total_m = torch.sum(masked_m)
+        if total_m > 0:
+            output = cutlass_w4a8_moe(
+                self.start_expert_id,
+                self.end_expert_id,
+                self.num_experts,
+                hidden_states,
+                self.w13_weight,
+                self.w2_weight,
+                self.w13_weight_scale_inv,
+                self.w2_weight_scale_inv,
+                None,
+                None,
+                masked_m,
+                self.quant_method.a_strides1,
+                self.quant_method.b_strides1,
+                self.quant_method.c_strides1,
+                self.quant_method.a_strides2,
+                self.quant_method.b_strides2,
+                self.quant_method.c_strides2,
+                self.quant_method.s_strides13,
+                self.quant_method.s_strides2,
+                self.quant_method.expert_offsets,
+                self.quant_method.problem_sizes1,
+                self.quant_method.problem_sizes2,
+                self.w13_input_scale,
+                self.w2_input_scale,
+                ep_mode=ep_mode,
+            )
+            return output.to(torch.bfloat16)
+        else:
+            return hidden_states.to(torch.bfloat16)
+
     def forward_normal(
         self,
         hidden_states: torch.Tensor,