[GH-ISSUE #2230] [Performance] NativeMoEWrapper rebuilds and retains a full checkpoint index per layer #1310

New issue

Open

opened 2026-10-29 14:51:06 +00:00 by bsakaguchi · 0 comments

bsakaguchi commented

2026-10-29 14:51:06 +00:00

Owner

Originally created by @ZenAlexa on GitHub (Oct 7, 2026).
Original GitHub issue: https://github.com/kvcache-ai/ktransformers/issues/2230

Observed behavior

At a5d7ad90479c9e8bdee491c16bfd563fa9d70262, NativeMoEWrapper.load_weights() rebuilds the complete safetensors key-to-file index after each layer. _release_loader() closes mmap handles and clears the class singleton, while each completed wrapper keeps its self.loader and its complete tensor_file_map.

The SGLang consumer constructs all wrappers before loading their weights. The wrappers initially share one loader; the per-layer recreation then leaves a separate full index on each completed layer. For L loaded layers and K checkpoint keys, both index construction and retained key entries grow as O(LK).

Component reproduction

A checkpoint fixture with 16 real safetensors shards, 8,192 keys and 64 layers exercises the current SafeTensorLoader and BF16SafeTensorLoader, with the wrapper's loader assignment/release sequence replayed around each layer load:

Loaded layers	`safe_open` calls	Enumerated keys	Retained loader indexes
1	16	8,192	1
8	128	65,536	8
32	512	262,144	32
64	1,024	524,288	64

The fixture totals 773,384 bytes. All 768 requested tensors match their expected values, and every file-handle map is empty after layer release. The retained key maps remain live until the simulated model holders are deleted. Environment: Python 3.12.14, PyTorch 2.8.0, safetensors 0.8.0, CPU. This measures Python loader behavior; native MoE packing and full-model startup remain to be measured.

Proposed boundary

Keep the index shared by the wrappers that still need to load, reopen only the shards needed by the current layer, and preserve per-layer handle release after native consumers finish. Drop each completed wrapper's index reference so the final layer releases the metadata too. Loader reuse must retain checkpoint/method identity, failure recovery and same-instance reload behavior.

This preserves the mmap-release intent of #1997. It leaves the RAWINT4 tensor-owner fix in #2225 and the byte-range prefetch/overlap work described in #2227 in their respective contributions. Header parsing on shard reopen remains part of the measured cost; the proposed reduction concerns full-directory scans, full key-map construction and retained metadata.

Originally created by @ZenAlexa on GitHub (Oct 7, 2026). Original GitHub issue: https://github.com/kvcache-ai/ktransformers/issues/2230 ### Observed behavior At `a5d7ad90479c9e8bdee491c16bfd563fa9d70262`, `NativeMoEWrapper.load_weights()` rebuilds the complete safetensors key-to-file index after each layer. `_release_loader()` closes mmap handles and clears the class singleton, while each completed wrapper keeps its `self.loader` and its complete `tensor_file_map`. The SGLang consumer constructs all wrappers before loading their weights. The wrappers initially share one loader; the per-layer recreation then leaves a separate full index on each completed layer. For `L` loaded layers and `K` checkpoint keys, both index construction and retained key entries grow as `O(LK)`. ### Component reproduction A checkpoint fixture with 16 real safetensors shards, 8,192 keys and 64 layers exercises the current `SafeTensorLoader` and `BF16SafeTensorLoader`, with the wrapper's loader assignment/release sequence replayed around each layer load: | Loaded layers | `safe_open` calls | Enumerated keys | Retained loader indexes | | --- | ---: | ---: | ---: | | 1 | 16 | 8,192 | 1 | | 8 | 128 | 65,536 | 8 | | 32 | 512 | 262,144 | 32 | | 64 | 1,024 | 524,288 | 64 | The fixture totals 773,384 bytes. All 768 requested tensors match their expected values, and every file-handle map is empty after layer release. The retained key maps remain live until the simulated model holders are deleted. Environment: Python 3.12.14, PyTorch 2.8.0, safetensors 0.8.0, CPU. This measures Python loader behavior; native MoE packing and full-model startup remain to be measured. ### Proposed boundary Keep the index shared by the wrappers that still need to load, reopen only the shards needed by the current layer, and preserve per-layer handle release after native consumers finish. Drop each completed wrapper's index reference so the final layer releases the metadata too. Loader reuse must retain checkpoint/method identity, failure recovery and same-instance reload behavior. This preserves the mmap-release intent of #1997. It leaves the RAWINT4 tensor-owner fix in #2225 and the byte-range prefetch/overlap work described in #2227 in their respective contributions. Header parsing on shard reopen remains part of the measured cost; the proposed reduction concerns full-directory scans, full key-map construction and retained metadata.