diff --git a/python/sglang/srt/model_loader/weight_utils.py b/python/sglang/srt/model_loader/weight_utils.py index 6bc867c35..2afd8a8c2 100644 --- a/python/sglang/srt/model_loader/weight_utils.py +++ b/python/sglang/srt/model_loader/weight_utils.py @@ -1180,6 +1180,43 @@ def _drop_file_cache_after_load(path: str) -> None: os.close(fd) +# [k3] Skip routed-expert tensors that this EP rank does not own BEFORE +# reading them (mmap -> untouched bytes are never read). Enabled by +# SGLANG_EXPERT_LOAD_FILTER= (e.g. 896 for Kimi-K3). +_B2_EXPERT_RE = None +_B2_EXPERT_RANGE = None + + +def _b2_skip_tensor(name: str) -> bool: + global _B2_EXPERT_RE, _B2_EXPERT_RANGE + n_exp = int(os.environ.get("SGLANG_EXPERT_LOAD_FILTER", "0") or 0) + if n_exp <= 0: + return False + if _B2_EXPERT_RANGE is None: + import re as _re + + from sglang.srt.distributed import ( + get_moe_expert_parallel_rank, + get_moe_expert_parallel_world_size, + ) + + _B2_EXPERT_RE = _re.compile(r"(?