sglang 0.5.0rc0__py3-none-any.whl → 0.5.0rc2__py3-none-any.whl

sglang/srt/layers/dp_attention.py

@@ -4,7 +4,7 @@ import functools
 import logging
 from contextlib import contextmanager
 from enum import IntEnum, auto
-from typing import TYPE_CHECKING, List, Tuple
+from typing import TYPE_CHECKING, List, Optional, Tuple
 
 import torch
 import triton
@@ -18,21 +18,26 @@ from sglang.srt.distributed import (
     tensor_model_parallel_all_reduce,
 )
 
+if TYPE_CHECKING:
+    from sglang.srt.configs.model_config import ModelConfig
+    from sglang.srt.server_args import ServerArgs
+
 logger = logging.getLogger(__name__)
 
 if TYPE_CHECKING:
     from sglang.srt.model_executor.forward_batch_info import ForwardBatch
 
-_ATTN_TP_GROUP = None
-_ATTN_TP_RANK = None
-_ATTN_TP_SIZE = None
-_ATTN_DP_RANK = None
-_ATTN_DP_SIZE = None
-_LOCAL_ATTN_DP_SIZE = None
-_LOCAL_ATTN_DP_RANK = None
+_ATTN_TP_GROUP: Optional[GroupCoordinator] = None
+_ATTN_TP_RANK: Optional[int] = None
+_ATTN_TP_SIZE: Optional[int] = None
+_ATTN_DP_RANK: Optional[int] = None
+_ATTN_DP_SIZE: Optional[int] = None
+_LOCAL_ATTN_DP_SIZE: Optional[int] = None
+_LOCAL_ATTN_DP_RANK: Optional[int] = None
+_ENABLE_DP_ATTENTION_FLAG: bool = False
 
 
-class DPPaddingMode(IntEnum):
+class DpPaddingMode(IntEnum):
 
     # Padding tokens to max length and then gather tokens using `all_gather_into_tensor`
     MAX_LEN = auto()
@@ -40,13 +45,13 @@ class DPPaddingMode(IntEnum):
     SUM_LEN = auto()
 
     def is_max_len(self):
-        return self == DPPaddingMode.MAX_LEN
+        return self == DpPaddingMode.MAX_LEN
 
     def is_sum_len(self):
-        return self == DPPaddingMode.SUM_LEN
+        return self == DpPaddingMode.SUM_LEN
 
     @classmethod
-    def get_dp_padding_mode(cls, global_num_tokens: List[int]) -> DPPaddingMode:
+    def get_dp_padding_mode(cls, global_num_tokens: List[int]) -> DpPaddingMode:
         # we choose the mode that minimizes the communication cost
         max_len = max(global_num_tokens)
         sum_len = sum(global_num_tokens)
@@ -56,10 +61,76 @@ class DPPaddingMode(IntEnum):
             return cls.SUM_LEN
 
     @classmethod
-    def get_default_mode_in_cuda_graph(cls) -> DPPaddingMode:
+    def get_default_mode_in_cuda_graph(cls) -> DpPaddingMode:
         return cls.MAX_LEN
 
 
+class _DpGatheredBufferWrapper:
+
+    _hidden_size: int
+    _dtype: torch.dtype
+    _device: torch.device
+    _global_dp_buffer_len: int
+    _local_dp_buffer_len: int
+
+    @classmethod
+    def set_metadata(cls, hidden_size: int, dtype: torch.dtype, device: torch.device):
+        cls._hidden_size = hidden_size
+        cls._dtype = dtype
+        cls._device = device
+
+    @classmethod
+    def set_dp_buffer_len(cls, global_dp_buffer_len: int, local_dp_buffer_len: int):
+        cls._global_dp_buffer_len = global_dp_buffer_len
+        cls._local_dp_buffer_len = local_dp_buffer_len
+
+    @classmethod
+    def get_global_dp_buffer(cls) -> torch.Tensor:
+        return torch.empty(
+            (cls._global_dp_buffer_len, cls._hidden_size),
+            dtype=cls._dtype,
+            device=cls._device,
+        )
+
+    @classmethod
+    def get_local_dp_buffer(cls) -> torch.Tensor:
+        return torch.empty(
+            (cls._local_dp_buffer_len, cls._hidden_size),
+            dtype=cls._dtype,
+            device=cls._device,
+        )
+
+    @classmethod
+    def get_global_dp_buffer_len(cls) -> int:
+        return cls._global_dp_buffer_len
+
+    @classmethod
+    def get_local_dp_buffer_len(cls) -> int:
+        return cls._local_dp_buffer_len
+
+
+def set_dp_buffer_len(global_dp_buffer_len: int, local_dp_buffer_len: int):
+    _DpGatheredBufferWrapper.set_dp_buffer_len(
+        global_dp_buffer_len, local_dp_buffer_len
+    )
+
+
+def get_global_dp_buffer() -> torch.Tensor:
+    return _DpGatheredBufferWrapper.get_global_dp_buffer()
+
+
+def get_local_dp_buffer() -> torch.Tensor:
+    return _DpGatheredBufferWrapper.get_local_dp_buffer()
+
+
+def get_global_dp_buffer_len() -> int:
+    return _DpGatheredBufferWrapper.get_global_dp_buffer_len()
+
+
+def get_local_dp_buffer_len() -> int:
+    return _DpGatheredBufferWrapper.get_local_dp_buffer_len()
+
+
 def compute_dp_attention_world_info(enable_dp_attention, tp_rank, tp_size, dp_size):
     if not enable_dp_attention:
         return tp_rank, tp_size, 0
@@ -89,18 +160,24 @@ def compute_dp_attention_local_info(
 
 
 def initialize_dp_attention(
-    enable_dp_attention: bool,
-    tp_rank: int,
-    tp_size: int,
-    dp_size: int,
-    moe_dense_tp_size: int,
-    pp_size: int,
+    server_args: ServerArgs,
+    model_config: ModelConfig,
 ):
     global _ATTN_TP_GROUP, _ATTN_TP_RANK, _ATTN_TP_SIZE, _ATTN_DP_RANK, _ATTN_DP_SIZE
-    global _LOCAL_ATTN_DP_SIZE, _LOCAL_ATTN_DP_RANK
+    global _LOCAL_ATTN_DP_SIZE, _LOCAL_ATTN_DP_RANK, _ENABLE_DP_ATTENTION_FLAG
 
     from sglang.srt.layers.sampler import SYNC_TOKEN_IDS_ACROSS_TP
 
+    enable_dp_attention = server_args.enable_dp_attention
+    tp_size = server_args.tp_size
+    dp_size = server_args.dp_size
+    moe_dense_tp_size = server_args.moe_dense_tp_size
+    pp_size = server_args.pp_size
+
+    tp_rank = get_tensor_model_parallel_rank()
+
+    _ENABLE_DP_ATTENTION_FLAG = enable_dp_attention
+
     _ATTN_TP_RANK, _ATTN_TP_SIZE, _ATTN_DP_RANK = compute_dp_attention_world_info(
         enable_dp_attention, tp_rank, tp_size, dp_size
     )
@@ -135,38 +212,48 @@ def initialize_dp_attention(
         group_name="attention_tp",
     )
 
+    _DpGatheredBufferWrapper.set_metadata(
+        hidden_size=model_config.hidden_size,
+        dtype=model_config.dtype,
+        device=torch.device("cuda"),
+    )
 
-def get_attention_tp_group():
+
+def is_dp_attention_enabled() -> bool:
+    return _ENABLE_DP_ATTENTION_FLAG
+
+
+def get_attention_tp_group() -> GroupCoordinator:
     assert _ATTN_TP_GROUP is not None, "dp attention not initialized!"
     return _ATTN_TP_GROUP
 
 
-def get_attention_tp_rank():
+def get_attention_tp_rank() -> int:
     assert _ATTN_TP_RANK is not None, "dp attention not initialized!"
     return _ATTN_TP_RANK
 
 
-def get_attention_tp_size():
+def get_attention_tp_size() -> int:
     assert _ATTN_TP_SIZE is not None, "dp attention not initialized!"
     return _ATTN_TP_SIZE
 
 
-def get_attention_dp_rank():
+def get_attention_dp_rank() -> int:
     assert _ATTN_DP_RANK is not None, "dp attention not initialized!"
     return _ATTN_DP_RANK
 
 
-def get_attention_dp_size():
+def get_attention_dp_size() -> int:
     assert _ATTN_DP_SIZE is not None, "dp attention not initialized!"
     return _ATTN_DP_SIZE
 
 
-def get_local_attention_dp_rank():
+def get_local_attention_dp_rank() -> int:
     assert _LOCAL_ATTN_DP_RANK is not None, "dp attention not initialized!"
     return _LOCAL_ATTN_DP_RANK
 
 
-def get_local_attention_dp_size():
+def get_local_attention_dp_size() -> int:
     assert _LOCAL_ATTN_DP_SIZE is not None, "dp attention not initialized!"
     return _LOCAL_ATTN_DP_SIZE
 
@@ -292,6 +379,10 @@ def _dp_gather_via_all_gather(
     forward_batch: ForwardBatch,
     is_partial: bool,
 ):
+    if get_attention_tp_size() == 1:
+        get_tp_group().all_gather_into_tensor(global_tokens, local_tokens)
+        return
+
     if not is_partial:
         if get_attention_tp_rank() != 0:
             local_tokens.fill_(0)

sglang/srt/layers/flashinfer_comm_fusion.py

@@ -1,5 +1,5 @@
 import logging
-from typing import Tuple
+from typing import Optional, Tuple
 
 import torch
 import torch.distributed as dist
@@ -92,7 +92,7 @@ _workspace_manager = FlashInferWorkspaceManager()
 
 
 def ensure_workspace_initialized(
-    max_token_num: int = 128, hidden_dim: int = 4096, use_fp32_lamport: bool = False
+    max_token_num: int = 2048, hidden_dim: int = 4096, use_fp32_lamport: bool = False
 ):
     """Ensure workspace is initialized"""
     if not is_flashinfer_available() or _flashinfer_comm is None:
@@ -124,8 +124,8 @@ def flashinfer_allreduce_residual_rmsnorm(
     residual: torch.Tensor,
     weight: torch.Tensor,
     eps: float = 1e-6,
-    max_token_num: int = 128,
-    use_oneshot: bool = True,
+    max_token_num: int = 2048,
+    use_oneshot: Optional[bool] = None,
     trigger_completion_at_end: bool = False,
     fp32_acc: bool = False,
 ) -> Tuple[torch.Tensor, torch.Tensor]:

sglang/srt/layers/linear.py

@@ -1294,6 +1294,7 @@ class RowParallelLinear(LinearBase):
         with use_symmetric_memory(parallel_state.get_tp_group()) as sm:
             output_parallel = self.quant_method.apply(self, input_parallel, bias=bias_)
             sm.tag(output_parallel)
+
         if self.reduce_results and self.tp_size > 1 and not skip_all_reduce:
             output = tensor_model_parallel_all_reduce(output_parallel)
         else:

sglang/srt/layers/logits_processor.py

@@ -27,7 +27,7 @@ from sglang.srt.distributed import (
     tensor_model_parallel_all_gather,
 )
 from sglang.srt.layers.dp_attention import (
-    DPPaddingMode,
+    DpPaddingMode,
     attn_tp_all_gather,
     attn_tp_all_gather_into_tensor,
     dp_gather_replicate,
@@ -35,7 +35,9 @@ from sglang.srt.layers.dp_attention import (
     get_attention_dp_rank,
     get_attention_dp_size,
     get_attention_tp_size,
+    get_global_dp_buffer,
     get_local_attention_dp_size,
+    set_dp_buffer_len,
 )
 from sglang.srt.layers.vocab_parallel_embedding import VocabParallelEmbedding
 from sglang.srt.managers.schedule_batch import global_server_args_dict
@@ -108,14 +110,12 @@ class LogitsMetadata:
     # The start position of local hidden states.
     dp_local_start_pos: Optional[torch.Tensor] = None
     dp_local_num_tokens: Optional[torch.Tensor] = None
-    gathered_buffer: Optional[torch.Tensor] = None
-    # Buffer to gather logits from all ranks.
-    forward_batch_gathered_buffer: Optional[torch.Tensor] = None
+    global_dp_buffer_len: Optional[int] = None
     # Number of tokens to sample per DP rank
     global_num_tokens_for_logprob_cpu: Optional[torch.Tensor] = None
     global_num_tokens_for_logprob_gpu: Optional[torch.Tensor] = None
     # The gather mode for DP attention
-    dp_padding_mode: Optional[DPPaddingMode] = None
+    dp_padding_mode: Optional[DpPaddingMode] = None
     # for padding
     padded_static_len: int = -1
 
@@ -164,11 +164,10 @@ class LogitsMetadata:
             global_num_tokens_gpu=forward_batch.global_num_tokens_gpu,
             dp_local_start_pos=forward_batch.dp_local_start_pos,
             dp_local_num_tokens=forward_batch.dp_local_num_tokens,
-            gathered_buffer=forward_batch.gathered_buffer,
-            forward_batch_gathered_buffer=forward_batch.gathered_buffer,
+            global_dp_buffer_len=forward_batch.global_dp_buffer_len,
             global_num_tokens_for_logprob_cpu=forward_batch.global_num_tokens_for_logprob_cpu,
             global_num_tokens_for_logprob_gpu=forward_batch.global_num_tokens_for_logprob_gpu,
-            dp_padding_mode=DPPaddingMode.SUM_LEN,
+            dp_padding_mode=DpPaddingMode.SUM_LEN,
         )
 
     def compute_dp_attention_metadata(self):
@@ -188,16 +187,11 @@ class LogitsMetadata:
 
         if self.global_num_tokens_for_logprob_cpu is not None:
             # create a smaller buffer to reduce peak memory usage
-            self.gathered_buffer = torch.empty(
-                (
-                    sum(self.global_num_tokens_for_logprob_cpu),
-                    self.gathered_buffer.shape[1],
-                ),
-                dtype=self.gathered_buffer.dtype,
-                device=self.gathered_buffer.device,
-            )
+            self.global_dp_buffer_len = sum(self.global_num_tokens_for_logprob_cpu)
         else:
-            self.gathered_buffer = torch.empty_like(self.gathered_buffer)
+            self.global_dp_buffer_len = self.global_dp_buffer_len
+
+        set_dp_buffer_len(self.global_dp_buffer_len, self.dp_local_num_tokens)
 
 
 class LogitsProcessor(nn.Module):
@@ -443,7 +437,7 @@ class LogitsProcessor(nn.Module):
         if self.do_tensor_parallel_all_gather_dp_attn:
             logits_metadata.compute_dp_attention_metadata()
             hidden_states, local_hidden_states = (
-                logits_metadata.gathered_buffer,
+                get_global_dp_buffer(),
                 hidden_states,
             )
             dp_gather_replicate(hidden_states, local_hidden_states, logits_metadata)

sglang/srt/layers/moe/cutlass_moe.py

@@ -9,7 +9,7 @@ from typing import Any, Callable, Dict, List, Optional, Tuple
 import torch
 
 from sglang.srt.layers.moe.cutlass_moe_params import CutlassMoEParams
-from sglang.srt.layers.utils import is_sm100_supported
+from sglang.srt.layers.utils import is_sm90_supported, is_sm100_supported
 from sglang.srt.utils import is_cuda
 
 _is_cuda = is_cuda()
@@ -124,6 +124,7 @@ def cutlass_fused_experts_fp8(
 
     if is_cuda:
         from sglang.srt.layers.quantization.fp8_kernel import (
+            per_group_transpose,
             per_token_group_quant_fp8_hopper_moe_mn_major,
             sglang_per_token_group_quant_fp8,
         )
@@ -152,15 +153,12 @@ def cutlass_fused_experts_fp8(
         k,
     )
 
-    if is_sm100_supported():
-        a_q, a1_scale = sglang_per_token_group_quant_fp8(a, 128)
-        rep_a_q = shuffle_rows(a_q, a_map, (m * topk, k))
-        rep_a1_scales = shuffle_rows(a1_scale, a_map, (m * topk, int(k / 128)))
-    else:
-        rep_a = shuffle_rows(a, a_map, (m * topk, k))
-        rep_a_q, rep_a1_scales = per_token_group_quant_fp8_hopper_moe_mn_major(
-            rep_a, expert_offsets, problem_sizes1, 128
-        )
+    a_q, a1_scale = sglang_per_token_group_quant_fp8(a, 128)
+    rep_a_q = shuffle_rows(a_q, a_map, (m * topk, k))
+    rep_a1_scales = shuffle_rows(a1_scale, a_map, (m * topk, int(k / 128)))
+
+    if not is_sm100_supported():
+        rep_a1_scales = per_group_transpose(rep_a1_scales, expert_offsets)
         w1_scale = w1_scale.contiguous()
 
     c1 = torch.empty((m * topk, n * 2), device=device, dtype=out_dtype)
@@ -193,12 +191,9 @@ def cutlass_fused_experts_fp8(
     intermediate = torch.empty((m * topk, n), device=device, dtype=out_dtype)
     silu_and_mul(c1, intermediate)
 
-    if is_sm100_supported():
-        intemediate_q, a2_scale = sglang_per_token_group_quant_fp8(intermediate, 128)
-    else:
-        intemediate_q, a2_scale = per_token_group_quant_fp8_hopper_moe_mn_major(
-            intermediate, expert_offsets, problem_sizes2, 128
-        )
+    intemediate_q, a2_scale = sglang_per_token_group_quant_fp8(intermediate, 128)
+    if not is_sm100_supported():
+        a2_scale = per_group_transpose(a2_scale, expert_offsets)
         w2_scale = w2_scale.contiguous()
 
     fp8_blockwise_scaled_grouped_mm(

sglang/srt/layers/moe/cutlass_w4a8_moe.py

@@ -11,7 +11,7 @@ from sgl_kernel import (
 )
 
 from sglang.srt.layers.moe.ep_moe.kernels import (
-    post_reorder_triton_kernel,
+    post_reorder_triton_kernel_for_cutlass_moe,
     pre_reorder_triton_kernel_for_cutlass_moe,
     run_cutlass_moe_ep_preproess,
 )
@@ -199,14 +199,13 @@ def cutlass_w4a8_moe(
     )
 
     output = torch.empty_like(a)
-    post_reorder_triton_kernel[(m,)](
+    post_reorder_triton_kernel_for_cutlass_moe[(m,)](
         c2,
         output,
         src2dst,
-        topk_ids_,
+        local_topk_ids,
         topk_weights,
-        start_expert_id,
-        end_expert_id,
+        num_experts,
         topk,
         k,
         0,

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

@@ -581,6 +581,49 @@ def post_reorder_triton_kernel(
             )
 
 
+@triton.jit
+def post_reorder_triton_kernel_for_cutlass_moe(
+    down_output_ptr,
+    output_ptr,
+    src2dst_ptr,
+    topk_ids_ptr,
+    topk_weights_ptr,
+    num_experts,
+    topk,
+    hidden_size,
+    dst_start,
+    BLOCK_SIZE: tl.constexpr,
+):
+    InDtype = down_output_ptr.dtype.element_ty
+
+    src_idx_int32 = tl.program_id(0)
+    src_idx = src_idx_int32.to(tl.int64)
+    src2dst_ptr = src2dst_ptr + src_idx * topk
+    topk_ids_ptr = topk_ids_ptr + src_idx * topk
+    topk_weights_ptr = topk_weights_ptr + src_idx * topk
+
+    store_ptr = output_ptr + src_idx * hidden_size
+
+    vec = tl.arange(0, BLOCK_SIZE)
+
+    for start_offset in tl.range(0, hidden_size, BLOCK_SIZE):
+        offset = start_offset + vec
+        mask = offset < hidden_size
+
+        sum_vec = tl.zeros([BLOCK_SIZE], dtype=InDtype)
+        for idx in range(topk):
+            expert_id = tl.load(topk_ids_ptr + idx)
+            if expert_id != num_experts:
+                dst_idx_int32 = tl.load(src2dst_ptr + idx)
+                dst_idx = dst_idx_int32.to(tl.int64)
+                dst_idx = dst_idx - dst_start
+                weigh_scale = tl.load(topk_weights_ptr + idx).to(InDtype)
+                load_ptr = down_output_ptr + dst_idx * hidden_size
+                in_data = tl.load(load_ptr + offset, mask=mask)
+                sum_vec += in_data * weigh_scale
+        tl.store(store_ptr + offset, sum_vec, mask=mask)
+
+
 @triton.jit
 def compute_m_range(
     pid,

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

@@ -34,6 +34,7 @@ from sglang.srt.utils import ceil_div, dispose_tensor, get_bool_env_var, is_hip,
 
 if TYPE_CHECKING:
     from sglang.srt.layers.moe.token_dispatcher import (
+        AscendDeepEPLLOutput,
         DeepEPLLOutput,
         DeepEPNormalOutput,
         DispatchOutput,
@@ -387,7 +388,8 @@ class DeepEPMoE(EPMoE):
             return_recv_hook=True,
         )
 
-        if self.deepep_mode.enable_low_latency():
+        if self.deepep_mode.enable_low_latency() and not _is_npu:
+            # NPU supports low_latency deepep without deepgemm
             assert (
                 deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM
             ), f"DeepEP {self.deepep_mode} mode requires deep_gemm"
@@ -404,7 +406,7 @@ class DeepEPMoE(EPMoE):
             )
             # the last one is invalid rank_id
             self.expert_mask[:-1] = 1
-        else:
+        elif not _is_npu:
             self.w13_weight_fp8 = (
                 self.w13_weight,
                 (
@@ -459,6 +461,8 @@ class DeepEPMoE(EPMoE):
         if _use_aiter:
             # in forward_aiter, we skip token permutation and unpermutation, which have been fused inside aiter kernel
             return self.forward_aiter(dispatch_output)
+        if _is_npu:
+            return self.forward_npu(dispatch_output)
         if dispatch_output.format.is_deepep_normal():
             assert deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM and self.use_fp8_w8a8
             return self.forward_deepgemm_contiguous(dispatch_output)
@@ -723,6 +727,60 @@ class DeepEPMoE(EPMoE):
 
         return down_output
 
+    def forward_npu(
+        self,
+        dispatch_output: DeepEPLLOutput,
+    ):
+        if TYPE_CHECKING:
+            assert isinstance(dispatch_output, AscendDeepEPLLOutput)
+        hidden_states, topk_idx, topk_weights, _, seg_indptr, _ = dispatch_output
+        assert self.quant_method is not None
+        assert self.activation == "silu"
+
+        # NOTE: Ascend's Dispatch & Combine does not support FP16
+        output_dtype = torch.bfloat16
+
+        pertoken_scale = hidden_states[1]
+        hidden_states = hidden_states[0]
+
+        group_list_type = 1
+        seg_indptr = seg_indptr.to(torch.int64)
+
+        import torch_npu
+
+        # gmm1: gate_up_proj
+        hidden_states = torch_npu.npu_grouped_matmul(
+            x=[hidden_states],
+            weight=[self.w13_weight],
+            scale=[self.w13_weight_scale.to(output_dtype)],
+            per_token_scale=[pertoken_scale],
+            split_item=2,
+            group_list_type=group_list_type,
+            group_type=0,
+            group_list=seg_indptr,
+            output_dtype=output_dtype,
+        )[0]
+
+        # act_fn: swiglu
+        hidden_states = torch_npu.npu_swiglu(hidden_states)
+
+        hidden_states, swiglu_out_scale = torch_npu.npu_dynamic_quant(hidden_states)
+
+        # gmm2: down_proj
+        hidden_states = torch_npu.npu_grouped_matmul(
+            x=[hidden_states],
+            weight=[self.w2_weight],
+            scale=[self.w2_weight_scale.to(output_dtype)],
+            per_token_scale=[swiglu_out_scale],
+            split_item=2,
+            group_list_type=group_list_type,
+            group_type=0,
+            group_list=seg_indptr,
+            output_dtype=output_dtype,
+        )[0]
+
+        return hidden_states
+
 
 def get_moe_impl_class():
     if global_server_args_dict["moe_a2a_backend"].is_deepep():