[PR #103] [perf] KT: the shard loaders do not read the CPU experts' checkpoint tensors #101

New issue

Open

opened 2026-10-29 21:07:14 +00:00 by bsakaguchi · 0 comments

bsakaguchi commented

2026-10-29 21:07:14 +00:00

Owner

📋 Pull Request Information

Original PR: https://github.com/kvcache-ai/sglang/pull/103
Author: @05yuki
Created: 10/13/2026
Status: 🔄 Open

Base: main ← Head: launch-skip-cpu-experts

📝 Commits (1)

503c270 [perf] KT: the shard loaders do not read the CPU experts' checkpoint tensors

📊 Changes

2 files changed (+44 additions, -1 deletions)

View changed files

📝 python/sglang/srt/layers/moe/kt_ep_wrapper.py (+16 -0)
📝 python/sglang/srt/model_loader/weight_utils.py (+28 -1)

📄 Description

With a KT EP wrapper the GPU MoE weight_loader drops a routed expert's tensors unread when the GPU-experts masks place that expert on the CPU (fused_moe_triton/layer.py). The shard loaders had already materialised them: multi_thread_safetensors_weights_iterator and the buffered one build {k: f.get_tensor(k) for k in f.keys()} per shard, and the single-thread iterator get_tensors every key, so the whole expert set goes through RAM once per rank for nothing. Qwen3.8-Flash-Next NVFP4 (512 experts a layer, 32 on the GPU): 158 GB read per TP rank, 107 s of a 293 s weight load. DeepSeek V4.1: 60 GB, the shard phase 102 → 18 s.

The loaders now ask before get_tensor. kt_ep_wrapper.checkpoint_expert_is_cpu(layer_idx, expert_id) answers from the same masks the weight_loader uses; it is False whenever the wrapper is disabled (a draft model's load) or no masks exist, so nothing changes for those. The key pattern (KT_LOAD_SKIP_CPU_EXPERTS_RE, default [model.]layers.<L>.<...>.experts.<E>.) covers mlp / ffn / block_sparse_moe containers with or without a model. prefix; shared experts and packed expert tensors do not match it. KT_LOAD_SKIP_CPU_EXPERTS=0 restores the full read.

Prompt NLL on three fixed texts (teacher-forced input logprobs) is identical to four decimals with and without the skip on Flash-Next NVFP4, DeepSeek V4.1, V4-Flash-Vision and MiMo-V2.6-Flash.

🤖 Generated with Claude Code

https://claude.ai/code/session_01C2UBZdnrDZSBh8yQz8rc9g

_{🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.}

## 📋 Pull Request Information **Original PR:** https://github.com/kvcache-ai/sglang/pull/103 **Author:** [@05yuki](https://github.com/05yuki) **Created:** 10/13/2026 **Status:** 🔄 Open **Base:** `main` ← **Head:** `launch-skip-cpu-experts` --- ### 📝 Commits (1) - [`503c270`](https://github.com/kvcache-ai/sglang/commit/503c270e4c2986963f73828fef2a32b40dc78eea) [perf] KT: the shard loaders do not read the CPU experts' checkpoint tensors ### 📊 Changes **2 files changed** (+44 additions, -1 deletions) <details> <summary>View changed files</summary> 📝 `python/sglang/srt/layers/moe/kt_ep_wrapper.py` (+16 -0) 📝 `python/sglang/srt/model_loader/weight_utils.py` (+28 -1) </details> ### 📄 Description With a KT EP wrapper the GPU MoE `weight_loader` drops a routed expert's tensors unread when the GPU-experts masks place that expert on the CPU (`fused_moe_triton/layer.py`). The shard loaders had already materialised them: `multi_thread_safetensors_weights_iterator` and the buffered one build `{k: f.get_tensor(k) for k in f.keys()}` per shard, and the single-thread iterator `get_tensor`s every key, so the whole expert set goes through RAM once per rank for nothing. Qwen3.8-Flash-Next NVFP4 (512 experts a layer, 32 on the GPU): 158 GB read per TP rank, 107 s of a 293 s weight load. DeepSeek V4.1: 60 GB, the shard phase 102 → 18 s. The loaders now ask before `get_tensor`. `kt_ep_wrapper.checkpoint_expert_is_cpu(layer_idx, expert_id)` answers from the same masks the `weight_loader` uses; it is `False` whenever the wrapper is disabled (a draft model's load) or no masks exist, so nothing changes for those. The key pattern (`KT_LOAD_SKIP_CPU_EXPERTS_RE`, default `[model.]layers.<L>.<...>.experts.<E>.`) covers `mlp` / `ffn` / `block_sparse_moe` containers with or without a `model.` prefix; shared experts and packed expert tensors do not match it. `KT_LOAD_SKIP_CPU_EXPERTS=0` restores the full read. Prompt NLL on three fixed texts (teacher-forced input logprobs) is identical to four decimals with and without the skip on Flash-Next NVFP4, DeepSeek V4.1, V4-Flash-Vision and MiMo-V2.6-Flash. 🤖 Generated with [Claude Code](https://claude.com/claude-code) https://claude.ai/code/session_01C2UBZdnrDZSBh8yQz8rc9g --- <sub>🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.</sub>

bsakaguchi added the

pull-request

label

2026-10-29 21:07:14 +00:00

bsakaguchi referenced this issue

2026-10-29 21:07:14 +00:00

[PR #101] [CLOSED] [feat](kt): load STOCK ModelOpt MIXED_PRECISION NVFP4 checkpoints #102