diff --git a/python/sglang/kernels/ops/moe/ep_moe_kernels.py b/python/sglang/kernels/ops/moe/ep_moe_kernels.py index 93246893e..0aca8a654 100644 --- a/python/sglang/kernels/ops/moe/ep_moe_kernels.py +++ b/python/sglang/kernels/ops/moe/ep_moe_kernels.py @@ -1,3 +1,4 @@ +import os import logging from functools import lru_cache from typing import Optional, Tuple @@ -188,7 +189,53 @@ def deepep_run_moe_deep_preprocess(topk_ids: torch.Tensor, num_experts: int): return reorder_topk_ids, src2dst, seg_indptr -def cutlass_w4_run_moe_ep_preproess(topk_ids: torch.Tensor): +# [k3] torch.sort(stable) of the m*topk (2400 at c=50/gamma 2) expert ids +# runs as ONE-CTA radixSortKVInPlace: ~20 us per MoE layer (1.8 ms per spec step, +# profile/sp3). The only consumer is src2dst = inverse permutation of the stable +# sort, which is offset[id] + rank-of-i-among-equal-ids (stable => rank in +# original order). One program per expert value computes exactly that with a +# count + cumsum: bit-identical src2dst, ~3 us. K3_W4_SRC2DST=sort restores stock. +@triton.jit +def _k3_count_src2dst_kernel(ids_ptr, src2dst_ptr, n, BLOCK: tl.constexpr): + v = tl.program_id(0) + off = tl.zeros((), dtype=tl.int32) + for start in range(0, n, BLOCK): + idx = start + tl.arange(0, BLOCK) + x = tl.load(ids_ptr + idx, mask=idx < n, other=2147483647) + off += tl.sum((x < v).to(tl.int32), axis=0) + for start in range(0, n, BLOCK): + idx = start + tl.arange(0, BLOCK) + m = idx < n + x = tl.load(ids_ptr + idx, mask=m, other=-2147483647) + sel = (x == v) & m + c = tl.cumsum(sel.to(tl.int32), axis=0) + tl.store(src2dst_ptr + idx, off + c - 1, mask=sel) + off += tl.sum(sel.to(tl.int32), axis=0) + + +_K3_SRC2DST_LOGGED = [] + + +def _k3_count_src2dst(topk_ids: torch.Tensor, num_values: int) -> torch.Tensor: + ids = topk_ids.view(-1) + if not _K3_SRC2DST_LOGGED: + _K3_SRC2DST_LOGGED.append(1) + logging.getLogger(__name__).info( + "[k3] cutlass_w4a8_moe src2dst via count kernel (n=%d, values=%d)", ids.numel(), num_values) + n = ids.numel() + src2dst = torch.empty(n, device=ids.device, dtype=torch.int32) + BLOCK = min(4096, max(128, triton.next_power_of_2(n))) + _k3_count_src2dst_kernel[(num_values,)](ids, src2dst, n, BLOCK=BLOCK, num_warps=4) + return src2dst + + +def cutlass_w4_run_moe_ep_preproess(topk_ids: torch.Tensor, num_values: int = 0): + # num_values: ids are in [0, num_values) (cutlass_w4a8_moe maps -1 -> num_local_experts, + # so num_values = num_local_experts + 1); 0 = unknown -> stock sort path. + # decode/verify sizes only (<= 16384 pairs; prefill chunks keep the multi-CTA sort) + if (num_values > 0 and topk_ids.numel() <= 16384 and topk_ids.dtype == torch.int32 + and os.environ.get("K3_W4_SRC2DST", "count") == "count"): + return _k3_count_src2dst(topk_ids, num_values) _, reorder_ids = torch.sort(topk_ids.view(-1), stable=True) BLOCK_SIZE = 512 diff --git a/python/sglang/srt/layers/moe/cutlass_w4a8_moe.py b/python/sglang/srt/layers/moe/cutlass_w4a8_moe.py index 0efbd7878..aebe3c6da 100644 --- a/python/sglang/srt/layers/moe/cutlass_w4a8_moe.py +++ b/python/sglang/srt/layers/moe/cutlass_w4a8_moe.py @@ -131,8 +131,11 @@ def cutlass_w4a8_moe( if get_parallel().moe_ep_size > 1: topk_ids = torch.where(topk_ids == -1, num_local_experts, topk_ids) + # [k3] ids are in [0, num_local_experts] here (non-local -> num_local_experts + # under EP); pass the value range so the count-based src2dst can replace the sort. src2dst = cutlass_w4_run_moe_ep_preproess( topk_ids, + num_local_experts + 1, ) gateup_input = torch.empty(