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