[ZenAlexa]

What does this PR do?

Reuse the checkpoint tensor index across pending NativeMoEWrapper layers and release each layer's shard handles after extracting its tensors. This removes repeated whole-checkpoint indexing, the completed wrappers' retained copies of that index, and a full-heap collection after every layer.

Fixes #2230.

Ownership and implementation

The SGLang load path constructs the layer wrappers before loading their weights. Each pending wrapper now owns a reference to the same loader for its (method, absolute requested path). A single weak cache connects construction of those wrappers without extending the loader's lifetime. Completed wrappers drop their loader reference; the last pending owner releases the index. Reloading a completed wrapper discovers the checkpoint again.

SafeTensorLoader retains its key-to-full-shard-path index when closing handles and reopens only the shards requested by subsequent tensor reads. FP8 and BF16 reuse this common implementation. Returned tensors retain their mapped storage through native loading and CPUInfer.sync(). Copying backends keep their post-sync weight cleanup. AVX2 RAWINT4 retains its normalized weight tensors for the native MoE lifetime, preserving the borrowed mmap owners introduced by #2225; its copied scales are released after sync. Clearing shard handles now releases their references directly; process-wide cyclic garbage collection stays with Python and the consuming runtime.

This preserves per-layer shard-handle release, including failure paths, and separates the metadata lifetime from the tensor storage lifetime. It changes no checkpoint format, serving API, prefetch policy, or native kernel. Full-path shard identity also supports nested shards with the same basename.

Measured metadata-path impact

The following paired measurements use candidate cc1d127d; the loader and wrapper files are byte-identical at 89dd0944. The current integration adds the upstream AVX2 RAWINT4 owner retention described above.

CPU component measurement on macOS arm64, Python 3.12.14, Torch 2.8.0, and safetensors 0.8.0: real safetensors files and Torch CPU tensors, 64 wrappers, 65,536 keys, 8 experts per layer, 16 shards, and a 6.18 MB checkpoint. Complete loader/wrapper source was exercised with substitutes for the base wrapper, CPUInfer, and native backend. Timings include construction, all layer loads, and destruction; the metadata-byte census is outside the timed region. Each arm ran in seven fresh processes with warm filesystem cache, in alternating order. Both arms include one explicit collection at final destruction. The value oracle uses one fixed ctypes array type throughout the run.

Metric	Baseline	Patched
Directory walks	64	1
Enumerated keys	4,194,304	65,536
Peak simultaneously live indexes	64	1
Peak Python index containers and strings	410.65 MiB	6.42 MiB
Retained index bytes after all loads	410.65 MiB	0
Layer-local shard layout: median elapsed	3.5619 s	0.1712 s
Layer-local shard layout: elapsed range	3.3987–3.6623 s	0.1640–0.1880 s
All-shard layer layout: median elapsed	3.5388 s	1.6374 s
All-shard layer layout: elapsed range	3.4249–3.7155 s	1.5836–5.3579 s

Shard opens fall from 1,024 to 79 for the layer-local layout. For the all-shard layout they remain at 1,024, and the median safe_open time is 1.4915 s for baseline and 1.4731 s for patched. Reopening a shard parses its header. The optimization removes repeated index enumeration and retained Python dictionaries; shard-header reads remain in the measured call boundary.

The 5.3579 s patched all-shard run remains in the seven-run sample; it includes 4.1320 s in safe_open and 0.9413 s in final scene collection. The source of that variation is unassigned.

For the construct-all consumer, key enumeration and peak index storage fall from O(LK) to O(K) for L layers and K checkpoint keys. Reopening work depends on the headers of the shards each layer actually touches. Interleaved construct/load callers retain correct behavior and may receive less sharing.

Handle release and collection

With automatic GC disabled, real checkpoint mappings fell from 16 to 1 when shard handles were cleared. The 12 returned tensors retained the current layer's storage and expected values. Releasing the final tensor and loader references removed the last mapping, with zero collections throughout. The safe_open objects were outside cyclic-GC tracking.

A separate three-process-per-arm comparison held the shared-index implementation fixed and changed only explicit collection. For 64 layers, 8,192 keys and 16 shards, the median lifecycle fell from 1.0142 s to 0.0279 s. The 64 explicit collections took 0.9699 s and reclaimed zero objects. Both arms performed 79 opens and 768 checked tensor-pointer reads, with zero live tensors, handles or indexes at completion. This component experiment isolates collection cost; the native experiment below covers numerical behavior and loader ownership.

Native BF16 validation

The native results below were measured at cc1d127d; the loader and wrapper files were verified byte-identical at 89dd0944. These measurements describe that BF16 revision pair. Current integration c712b6d3 also retains the upstream AVX2 RAWINT4 owners from #2225.

Built the real AVX2 extension on Ubuntu 24.04 x86_64, Intel Xeon Platinum 8468, Python 3.12, and PyTorch 2.9.1+cu128. The experiment used two CPUInfer threads on NUMA node 0; an H100 supplied the CUDA-backed pinned allocator. MoE computation ran on the CPU. Baseline and patched Python code ran in separate processes against the same native extension.

The real NativeMoEWrapper → CPUInfer.submit/sync → AVX2BF16_MOE.forward_task path passed 42 forward cases per arm. The fixture used four layers, eight experts, hidden size 256, intermediate size 512, top-k 2, and query lengths 1 and 16. Every result met the existing FP32 MoE reference, rounded to BF16, under the normalized mean absolute error criterion mean(abs(output - reference)) / (mean(abs(reference)) + 1e-8) < 0.02; the maximum observed value was 0.0038790. All 42 patched outputs were also bitwise equal to their baseline counterparts.

The scenario covered full-cohort loading, atomic checkpoint replacement and same-wrapper reload, forward after explicit loader release and collection, abandoned pending wrappers, and missing-layer failure followed by checkpoint repair. Candidate recovery succeeded directly; the baseline needed its explicit loader-release operation to discard the stale index. The four-layer cohort reduced index builds from four to one and retained indexes after loading from four to zero. All wrapper/index state was released at final collection. No wrapper, loader, native backend, allocator, or safetensors function was substituted in this run. Profiling was enabled for ownership counts; native elapsed times are not used as performance claims.

Current-main integration and validation

Integrated main 5d048e51 through a normal merge at c712b6d3 on October 10. This includes #2225's AVX2 RAWINT4 owner fix. SGLang constructs the layer cohort before loading. The contribution remains confined to checkpoint index and shard-handle ownership in amx.py and loader.py.

Current native validation

Rebuilt the AVX2 extension from current kt-kernel source at 5d048e51, then ran main and c712b6d3 in separate processes against that same extension. The October 10 run used Ubuntu 24.04 x86_64, Intel Xeon Platinum 8468, Python 3.12, PyTorch 2.9.1+cu128, two CPUInfer threads, and NUMA node 0. An H100 supplied the CUDA-backed pinned allocator; MoE execution ran on the CPU. The full wrapper constructor, safetensors loader, native backend, and allocator ran without substitution.

Native format	Forward cases per arm	Maximum normalized mean absolute error	Existing reference limit	Main / candidate outputs
RAWINT4, uint8 and int32 storage	36	0.0038353	< 0.12	All 36 pairs bitwise equal
BF16	42	0.0038790	< 0.02	All 42 pairs bitwise equal

The RAWINT4 fixture used eight experts, hidden size 256, intermediate size 512, top-k 2, group size 128, and query lengths 1 and 16. Signed int4 values and exact BF16 scales defined the independent FP32 reference before packing or loading. Both storage formats completed native forward after explicit loader release and garbage collection, atomic checkpoint replacement and same-wrapper reload, and missing-layer failure followed by repair. Each loaded wrapper retained its 24 normalized weight tensors through forward; the copied source scales had been released. Reload released the superseded owners, and wrapper destruction released all remaining owners.

Re-running the BF16 lifecycle above on the current revision pair reproduced the four-layer cohort's reduction from four index builds to one and from four retained indexes after loading to zero. Candidate failure recovery succeeded directly; current main needed explicit index release. Ownership profiling was active in both native scenarios; their elapsed times are excluded from performance claims.

Component and existing regression coverage

• Re-ran the 20 real-file lifecycle/format scenarios at c712b6d3. Compared scaling at 1, 8, 32, and 64 layers with current main, and compared values for nine loader/format configurations before and after candidate shard reopening. This run used macOS arm64, Python 3.12.14, Torch 2.14.1, and safetensors 0.8.0. These scenarios exercise complete loader/wrapper source with substitutes for BaseMoEWrapper, CPUInfer, and the native backend.
• Ran a real-safetensors RAWINT4 ownership scene for both uint8 and int32 storage. After handle release and collection, current main and the integration retain all six normalized weight tensors and preserve 192 bytes of borrowed data. Atomic checkpoint replacement and same-wrapper reload preserve the new values and release the superseded owners; destroying the wrapper releases the remaining owners. Pre-merge 89dd0944 releases all six owners after load. The scene leaves its dangling addresses unread and executes no native forward.
• The 20-scenario coverage includes nine loader/format configurations, BF16 per-expert and packed layouts, FP8 scale views, RAWINT4 byte conversion, path/method isolation, interleaved checkpoints, explicit handle release, abandoned construction, failure/retry, same-wrapper reload, nested same-name shards, and a changed working directory.
• The historical 28 large-scale component runs passed 43,008 tensor-pointer reads in total, 1,536 per run. Those timing and size results remain bound to cc1d127d / 89dd0944.
• Re-ran the four existing FP8 transport checks, three expert-count checks, and nine existing auto-release checks at c712b6d3; these retain their existing source-inspection or mocked scope. Python 3.11 grammar and git diff --check also passed.

AMX, other native quantized formats, nonzero NUMA placement, the complete SGLang serving process, and multi-TB checkpoint startup remain outside the current measurements.

Before submitting

• Contributor guideline applied.
• Focused behavior and existing regression checks completed at the scope described above.
• Native RAWINT4 and BF16 loading and forward validation completed at the current revision pair.

Current mainline validation

Merged main a694e718 into efa4384a on 2026-10-19, retaining the new SYCL XMX selector and the complete loader lifecycle change. With the project-pinned PyTorch 2.9.1, 54 existing checks passed and two Linux ELF checks skipped. The existing mocked release checks retain their mock scope.

Twenty current-source scenarios using real safetensors and CPU tensors passed, covering pending-layer sharing, format/path isolation, handle release/reopen, repeated shard names, borrowed RAWINT4 storage, failed loading, and retry. The native execution boundary was substituted. The 1/8/32/64-layer comparison and changed-file formatting checks passed. Native ISA, SYCL execution, and the full serving process were not rerun; earlier hardware results retain their recorded revisions.

Project CI: Please run the current-head native kernel CI and checkpoint-loading integration on the project runners, including shard-handle release and forward execution after loader release.

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

[yyj6666667]

Reviewed the loader ownership and shard reopen path. A small real safetensors fixture passed for duplicate shard basenames, handle release, and tensor lifetime.

[ZenAlexa]

@yyj6666667 could you approve the workflow run and add run-ci so the kernel checks can run?

Retracted at the requester’s direction; this review was submitted before their approval.

[ZenAlexa]

@yyj6666667 could you rerun the run triggered by run-ci? the older run cancelled it, then failed because it used the labels saved from before run-ci was added.