| |
| |
| |
| |
| @@ -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 |
|
|
