sglang 0.4.4__py3-none-any.whl → 0.4.4.post2__py3-none-any.whl

sglang/srt/layers/moe/router.py

@@ -0,0 +1,342 @@
+from typing import Tuple
+
+import torch
+import triton
+import triton.language as tl
+
+from sglang.srt.layers.moe.topk import fused_topk
+
+
+@triton.jit
+def fused_moe_router_kernel(
+    input_ptr,  # input (bs, hidden_dim)
+    moe_router_weight_ptr,  # input (num_experts, hidden_dim)
+    topk_weights_ptr,  # output (bs, topk)
+    topk_ids_ptr,  # output (bs, topk)
+    num_experts: tl.constexpr,
+    topk: tl.constexpr,
+    moe_softcapping: tl.constexpr,
+    moe_renormalize: tl.constexpr,  # not supported
+    hidden_dim: tl.constexpr,
+    BLOCK_SIZE: tl.constexpr,
+):
+    pid = tl.program_id(axis=0)
+
+    offsets = tl.arange(0, BLOCK_SIZE)
+    mask = offsets < hidden_dim
+
+    # moe_router_weight is k major
+    expert_offsets = tl.arange(0, num_experts)[:, None]
+    router_mask = mask[None, :]
+    w_router = tl.load(
+        moe_router_weight_ptr + expert_offsets * hidden_dim + offsets[None, :],
+        mask=router_mask,
+        other=0.0,
+    )
+
+    x = tl.load(input_ptr + pid * hidden_dim + offsets, mask=mask, other=0.0)
+
+    # todo: tl.dot?
+    logits = tl.sum((w_router.to(tl.float32) * x[None, :].to(tl.float32)), axis=-1)
+
+    # logit softcap
+    logits_scaled = logits / moe_softcapping
+    exped = tl.exp(2 * logits_scaled)
+    top = exped - 1
+    bottom = exped + 1
+    logits_softcapped = top / bottom * moe_softcapping
+
+    # topk
+    # assert 1 <= topk <= num_experts
+
+    # 5.38 us
+
+    top1 = tl.argmax(logits_softcapped, axis=0)
+    tl.store(topk_ids_ptr + pid * topk + 0, top1)  # 5.63 us
+
+    top1_v = tl.max(logits_softcapped, axis=0)
+    invsumexp = 1.0 / tl.sum(tl.exp(logits_softcapped - top1_v), axis=0)
+
+    tl.store(
+        topk_weights_ptr + pid * topk + 0,
+        invsumexp,
+    )  # 5.73 us
+
+    if topk >= 2:
+        top2 = tl.argmax(
+            tl.where(
+                tl.arange(0, num_experts) != top1, logits_softcapped, float("-inf")
+            ),
+            axis=0,
+        )
+        tl.store(topk_ids_ptr + pid * topk + 1, top2)
+        top2_v = tl.sum(logits_softcapped * (tl.arange(0, num_experts) == top2), axis=0)
+        tl.store(
+            topk_weights_ptr + pid * topk + 1,
+            tl.exp(top2_v - top1_v) * invsumexp,
+        )  # 5.95us
+
+    # probably slow
+    if topk > 2:
+        topk_mask = tl.full(logits_softcapped.shape, 1.0, dtype=logits_softcapped.dtype)
+        topk_mask = tl.where(
+            tl.arange(0, num_experts) != top1, topk_mask, float("-inf")
+        )
+        topk_mask = tl.where(
+            tl.arange(0, num_experts) != top2, topk_mask, float("-inf")
+        )
+        for i in range(2, topk):
+            topi = tl.argmax(logits_softcapped + topk_mask, axis=0)
+            topk_mask = tl.where(
+                tl.arange(0, num_experts) != topi, topk_mask, float("-inf")
+            )
+            tl.store(topk_ids_ptr + pid * topk + i, topi)
+            topi_v = tl.sum(
+                logits_softcapped * (tl.arange(0, num_experts) == topi), axis=0
+            )
+            tl.store(
+                topk_weights_ptr + pid * topk + i,
+                tl.exp(topi_v - top1_v) * invsumexp,
+            )
+    # assert not moe_renormalize, "moe weight renormalization not implemented"
+
+
+def fused_moe_router_impl(
+    x: torch.Tensor,
+    router_weight: torch.Tensor,
+    topk: int,
+    moe_softcapping: float,
+):
+    assert len(x.shape) == 2 and x.shape[1] == router_weight.shape[1]
+    bs, hidden_dim = x.shape
+    num_experts = router_weight.shape[0]
+
+    # router_logits = torch.empty((bs, num_experts), dtype=torch.float32, device=x.device)
+    topk_weights = torch.empty((bs, topk), dtype=torch.float32, device=x.device)
+    topk_ids = torch.empty((bs, topk), dtype=torch.int32, device=x.device)
+
+    grid = lambda meta: (bs,)
+    config = {
+        "BLOCK_SIZE": triton.next_power_of_2(hidden_dim),
+        "num_warps": max(
+            min(triton.next_power_of_2(triton.cdiv(hidden_dim, 256)), 32), 4
+        ),
+    }
+
+    fused_moe_router_kernel[grid](
+        x,
+        router_weight,
+        topk_weights,
+        topk_ids,
+        num_experts=num_experts,
+        topk=topk,
+        moe_softcapping=moe_softcapping,
+        moe_renormalize=False,
+        hidden_dim=hidden_dim,
+        **config,
+    )
+
+    return topk_weights, topk_ids
+
+
+@triton.jit
+def fused_moe_router_large_bs_kernel(
+    a_ptr,  # input (bs, hidden_dim)
+    b_ptr,  # input (num_experts, hidden_dim)
+    topk_weights_ptr,  # output (bs, topk)
+    topk_ids_ptr,  # output (bs, topk)
+    bs,
+    num_experts: tl.constexpr,
+    topk: tl.constexpr,  # only support topk == 1
+    moe_softcapping: tl.constexpr,
+    moe_renormalize: tl.constexpr,  # not supported
+    K: tl.constexpr,
+    BLOCK_SIZE_M: tl.constexpr,
+    BLOCK_SIZE_N: tl.constexpr,
+    BLOCK_SIZE_K: tl.constexpr,
+    stride_am: tl.constexpr,
+    stride_bn: tl.constexpr,
+):
+
+    # 1. get block id
+    pid = tl.program_id(axis=0)
+
+    # 2. create pointers for the first block of A and B
+    # 2.1. setup a_ptrs with offsets in m and k
+    offs_m = pid * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)[:, None]
+    bs_mask = offs_m < bs
+    offs_k = tl.arange(0, BLOCK_SIZE_K)[None, :]
+    a_ptrs = a_ptr + (offs_m * stride_am + offs_k)
+
+    # 2.2. setup b_ptrs with offsets in k and n.
+    #      Note: b matrix is k-major.
+    offs_k = tl.arange(0, BLOCK_SIZE_K)[None, :]
+    offs_n = tl.arange(0, BLOCK_SIZE_N)[:, None]
+    expert_mask = offs_n < num_experts
+    b_ptrs = b_ptr + (offs_n * stride_bn + offs_k)
+
+    # 3. Create an accumulator of float32 of size [BLOCK_SIZE_M, BLOCK_SIZE_N]
+    #    3.1. iterate in K dimension
+    #    3.2. transpose tile B
+    acc = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
+    for k in range(0, K // BLOCK_SIZE_K):  # hidden_dim % BLOCK_SIZE_K == 0
+        a = tl.load(
+            a_ptrs,
+            mask=bs_mask,
+            other=0.0,
+        ).to(tl.float32)
+        b = tl.load(b_ptrs, mask=expert_mask, other=0.0).to(tl.float32).T
+        acc += tl.dot(a, b)
+
+        # Advance the ptrs to the next K block.
+        a_ptrs += BLOCK_SIZE_K
+        b_ptrs += BLOCK_SIZE_K
+
+    # 4. logit softcap
+    logits_scaled = acc / moe_softcapping
+    exped = tl.exp(2 * logits_scaled)
+    logits_softcapped = (exped - 1) / (exped + 1) * moe_softcapping
+
+    # 5. top1
+    cond = tl.arange(0, BLOCK_SIZE_N)[None, :] < num_experts
+    top1 = tl.argmax(tl.where(cond, logits_softcapped, float("-inf")), axis=1)
+    top1_v = tl.max(
+        tl.where(cond, logits_softcapped, float("-inf")), axis=1, keep_dims=True
+    )
+    invsumexp = 1.0 / tl.sum(
+        tl.where(cond, tl.exp(logits_softcapped - top1_v), 0.0), axis=1
+    )
+
+    # 6. store to output
+    offs_topk = pid * topk * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
+    topk_mask = offs_topk < bs
+    tl.store(topk_ids_ptr + offs_topk, top1, mask=topk_mask)
+    tl.store(
+        topk_weights_ptr + offs_topk,
+        invsumexp,
+        mask=topk_mask,
+    )
+
+
+def fused_moe_router_large_bs_impl(
+    x: torch.Tensor,
+    router_weight: torch.Tensor,
+    topk: int,
+    moe_softcapping: float,
+    BLOCK_SIZE_M: int,
+    BLOCK_SIZE_N: int,
+    BLOCK_SIZE_K: int,
+):
+    assert len(x.shape) == 2 and x.shape[1] == router_weight.shape[1]
+    bs, hidden_dim = x.shape
+    num_experts = router_weight.shape[0]
+
+    assert num_experts <= BLOCK_SIZE_N
+    assert hidden_dim % BLOCK_SIZE_K == 0
+    assert topk == 1
+
+    topk_weights = torch.empty((bs, topk), dtype=torch.float32, device=x.device)
+    topk_ids = torch.empty((bs, topk), dtype=torch.int32, device=x.device)
+
+    grid = (triton.cdiv(bs, BLOCK_SIZE_M) * triton.cdiv(num_experts, BLOCK_SIZE_N),)
+
+    fused_moe_router_large_bs_kernel[grid](
+        a_ptr=x,
+        b_ptr=router_weight,
+        topk_weights_ptr=topk_weights,
+        topk_ids_ptr=topk_ids,
+        bs=bs,
+        num_experts=num_experts,
+        topk=topk,
+        moe_softcapping=moe_softcapping,
+        moe_renormalize=False,
+        K=hidden_dim,
+        BLOCK_SIZE_M=BLOCK_SIZE_M,
+        BLOCK_SIZE_N=BLOCK_SIZE_N,
+        BLOCK_SIZE_K=BLOCK_SIZE_K,
+        stride_am=hidden_dim,
+        stride_bn=hidden_dim,
+    )
+
+    return topk_weights, topk_ids
+
+
+def fused_moe_router_shim(
+    moe_softcapping,
+    hidden_states,
+    gating_output,
+    topk,
+    renormalize,
+):
+    assert not renormalize
+    assert (
+        len(hidden_states.shape) == 2
+        and hidden_states.shape[1] == gating_output.shape[1]
+    )
+    bs, hidden_dim = hidden_states.shape
+    num_experts = gating_output.shape[0]
+    BLOCK_SIZE_M = 32
+    BLOCK_SIZE_N = 16
+    BLOCK_SIZE_K = 256
+    if (
+        bs >= 512
+        and topk == 1
+        and num_experts <= BLOCK_SIZE_N
+        and hidden_dim % BLOCK_SIZE_K == 0
+    ):
+        return fused_moe_router_large_bs_impl(
+            x=hidden_states,
+            router_weight=gating_output,
+            topk=topk,
+            moe_softcapping=moe_softcapping,
+            BLOCK_SIZE_M=BLOCK_SIZE_M,
+            BLOCK_SIZE_N=BLOCK_SIZE_N,
+            BLOCK_SIZE_K=BLOCK_SIZE_K,
+        )
+    else:
+        return fused_moe_router_impl(
+            x=hidden_states,
+            router_weight=gating_output,
+            topk=topk,
+            moe_softcapping=moe_softcapping,
+        )
+
+
+class FusedMoeRouter:
+    def __init__(self, router_linear, topk, moe_softcapping) -> None:
+        self.router_linear = router_linear
+        self.topk = topk
+        self.moe_softcapping = moe_softcapping
+
+    def __call__(self, *args, **kwargs):
+        return self.forward(*args, **kwargs)
+
+    def forward(
+        self, x: torch.Tensor, residual: torch.Tensor
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        if x.is_cuda:
+            return self.forward_cuda(x, residual)
+        else:
+            return self.forward_vllm(x, residual)
+
+    def forward_cuda(
+        self, x: torch.Tensor, autotune=False
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        return fused_moe_router_shim(
+            moe_softcapping=self.moe_softcapping,
+            hidden_states=x,
+            gating_output=self.router_linear.weight,
+            topk=self.topk,
+            renormalize=False,
+        )
+
+    def forward_vllm(
+        self,
+        x: torch.Tensor,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        # g, _ = self.router_linear.forward(x)
+        g = x.float() @ self.router_linear.weight.T.float()
+
+        g = torch.tanh(g.float() / self.moe_softcapping) * self.moe_softcapping
+
+        return fused_topk(x, g, self.topk, False)

sglang/srt/layers/moe/topk.py

@@ -17,7 +17,14 @@ from typing import Callable, Optional
 import torch
 import torch.nn.functional as F
 
-from sglang.srt.utils import get_compiler_backend
+from sglang.srt.utils import get_compiler_backend, is_cuda, is_hip
+
+_is_cuda = is_cuda()
+_is_hip = is_hip()
+
+from sglang.srt.managers.expert_distribution import ExpertDistributionRecorder
+
+expert_distribution_recorder = ExpertDistributionRecorder()
 
 
 def fused_topk_native(
@@ -47,7 +54,10 @@ def fused_topk(
     topk: int,
     renormalize: bool,
 ):
-    from vllm import _custom_ops as ops
+    if _is_cuda or _is_hip:
+        from sgl_kernel import topk_softmax
+    else:
+        from vllm import _custom_ops as vllm_ops
 
     assert hidden_states.shape[0] == gating_output.shape[0], "Number of tokens mismatch"
 
@@ -61,12 +71,20 @@ def fused_topk(
         M, topk, dtype=torch.int32, device=hidden_states.device
     )
 
-    ops.topk_softmax(
-        topk_weights,
-        topk_ids,
-        token_expert_indicies,
-        gating_output.float(),
-    )
+    if _is_cuda or _is_hip:
+        topk_softmax(
+            topk_weights,
+            topk_ids,
+            token_expert_indicies,
+            gating_output.float(),
+        )
+    else:
+        vllm_ops.topk_softmax(
+            topk_weights,
+            topk_ids,
+            token_expert_indicies,
+            gating_output.float(),
+        )
     del token_expert_indicies
 
     if renormalize:
@@ -75,6 +93,7 @@ def fused_topk(
     return topk_weights, topk_ids
 
 
+# This is used by the Deepseek V2/V3/R1 series models
 @torch.compile(dynamic=True, backend=get_compiler_backend())
 def grouped_topk(
     hidden_states: torch.Tensor,
@@ -83,17 +102,10 @@ def grouped_topk(
     renormalize: bool,
     num_expert_group: int = 0,
     topk_group: int = 0,
-    scoring_func: str = "softmax",
 ):
     assert hidden_states.shape[0] == gating_output.shape[0], "Number of tokens mismatch"
 
-    if scoring_func == "softmax":
-        scores = torch.softmax(gating_output, dim=-1)
-    elif scoring_func == "sigmoid":
-        scores = gating_output.sigmoid()
-    else:
-        raise ValueError(f"Scoring function '{scoring_func}' is not supported.")
-
+    scores = torch.softmax(gating_output, dim=-1)
     num_token = scores.shape[0]
     group_scores = (
         scores.view(num_token, num_expert_group, -1).max(dim=-1).values
@@ -117,7 +129,6 @@ def grouped_topk(
     return topk_weights.to(torch.float32), topk_ids.to(torch.int32)
 
 
-# DeepSeek V2/V3/R1 uses biased_grouped_top
 @torch.compile(dynamic=True, backend=get_compiler_backend())
 def biased_grouped_topk(
     hidden_states: torch.Tensor,
@@ -172,7 +183,7 @@ def select_experts(
     correction_bias: Optional[torch.Tensor] = None,
     torch_native: bool = False,
 ):
-    # DeepSeek V2/V3/R1 uses biased_grouped_top
+    # DeekSeekv2 uses grouped_top_k
     if use_grouped_topk:
         assert topk_group is not None
         assert num_expert_group is not None
@@ -217,4 +228,6 @@ def select_experts(
             renormalize=renormalize,
         )
 
+    expert_distribution_recorder.record_new_token(topk_ids)
+
     return topk_weights, topk_ids

sglang/srt/layers/parameter.py

@@ -105,6 +105,7 @@ class _ColumnvLLMParameter(BasevLLMParameter):
 
         shard_offset = kwargs.get("shard_offset")
         shard_size = kwargs.get("shard_size")
+        tp_rank = kwargs.get("tp_rank")
         use_presharded_weights = kwargs.get("use_presharded_weights")
         if (
             isinstance(self, (PackedColumnParameter, PackedvLLMParameter))
@@ -116,7 +117,6 @@ class _ColumnvLLMParameter(BasevLLMParameter):
 
         param_data = self.data
 
-        tp_rank = get_tensor_model_parallel_rank()
         param_data = param_data.narrow(self.output_dim, shard_offset, shard_size)
         if not use_presharded_weights:
             loaded_weight = loaded_weight.narrow(