File size: 3,088 Bytes
78a49cb | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 | 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=<num_routed_experts> (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"(?<!shared_)experts\.(\d+)\.")
+ ep = get_moe_expert_parallel_world_size()
+ r = get_moe_expert_parallel_rank()
+ per = n_exp // ep
+ _B2_EXPERT_RANGE = (r * per, (r + 1) * per)
+ logger.info(
+ "[k3] expert load filter: ep_rank %d/%d keeps experts [%d, %d)",
+ r, ep, *_B2_EXPERT_RANGE,
+ )
+ m = _B2_EXPERT_RE.search(name)
+ if m is None:
+ return False
+ e = int(m.group(1))
+ lo, hi = _B2_EXPERT_RANGE
+ return not (lo <= e < hi)
+
+
def safetensors_weights_iterator(
hf_weights_files: List[str],
disable_mmap: bool = False,
@@ -1212,6 +1249,8 @@ def safetensors_weights_iterator(
else:
with safetensors.safe_open(st_file, framework="pt", device="cpu") as f:
for name in f.keys():
+ if _b2_skip_tensor(name):
+ continue
yield name, f.get_tensor(name)
if drop_cache_after_load:
_drop_file_cache_after_load(st_file)
@@ -1294,7 +1333,9 @@ def multi_thread_safetensors_weights_iterator(
result = safetensors.torch.load(f.read())
else:
with safetensors.safe_open(st_file, framework="pt", device="cpu") as f:
- result = {k: f.get_tensor(k) for k in f.keys()}
+ result = {
+ k: f.get_tensor(k) for k in f.keys() if not _b2_skip_tensor(k)
+ }
return st_file, result
@@ -1349,7 +1390,9 @@ def buffered_multi_thread_safetensors_weights_iterator(
result = safetensors.torch.load(f.read())
else:
with safetensors.safe_open(st_file, framework="pt", device="cpu") as f:
- result = {k: f.get_tensor(k) for k in f.keys()}
+ result = {
+ k: f.get_tensor(k) for k in f.keys() if not _b2_skip_tensor(k)
+ }
return result
# Sliding window: max_workers loading + 1 prefetched.
|