diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/fused_marlin_moe.py b/python/sglang/srt/layers/moe/fused_moe_triton/fused_marlin_moe.py index 7f76b2616..2d6ced31d 100644 --- a/python/sglang/srt/layers/moe/fused_moe_triton/fused_marlin_moe.py +++ b/python/sglang/srt/layers/moe/fused_moe_triton/fused_marlin_moe.py @@ -1,3 +1,4 @@ +import os from typing import Optional import torch @@ -10,6 +11,7 @@ from sglang.srt.utils import is_cuda from sglang.srt.utils.custom_op import register_custom_op _is_cuda = is_cuda() +_K3_MARLIN_LOGGED = set() if _is_cuda: from sgl_kernel import moe_sum_reduce @@ -237,9 +239,27 @@ def fused_marlin_moe( # M block size selection logic # TODO: tune this further for specific models + # [k3] Under EP the topk ids are global-routed (most are -1 = + # non-local), so M * topk / E_local overestimates the rows per local expert + # by ep_size (TP32/EP32, c=50, gamma 2: 150*16/28 = 86 -> block_m 64, the + # kernel then runs 64-row MMA tiles for ~2.7 real rows -> compute-bound). + # Estimate with the global expert count instead (150*16/896 = 2.7 -> 8). + # SGLANG_MARLIN_EP_BLOCK=0 restores the stock choice; =8/16/32/48/64 forces. + _tok_per_expert = M * topk / E + _ep_block = os.environ.get("SGLANG_MARLIN_EP_BLOCK", "1") + if expert_map is not None and global_num_experts > E and _ep_block == "1": + _tok_per_expert = M * topk / global_num_experts for block_size_m in [8, 16, 32, 48, 64]: - if M * topk / E / block_size_m < 0.9: + if _tok_per_expert / block_size_m < 0.9: break + if _ep_block in ("8", "16", "32", "48", "64"): + block_size_m = int(_ep_block) + if expert_map is not None and M > 1 and (M, block_size_m) not in _K3_MARLIN_LOGGED and len(_K3_MARLIN_LOGGED) < 8: + _K3_MARLIN_LOGGED.add((M, block_size_m)) + import logging as _lg + _lg.getLogger(__name__).info( + "[k3] marlin EP block_m=%d (M=%d topk=%d E_local=%d E_global=%d mode=%s)", + block_size_m, M, topk, E, global_num_experts, _ep_block) if global_num_experts == -1: global_num_experts = E