[Hrithvik-Alex]

Motivation

With --load-format instanttensor (#40453), GPU memory after weight loading is higher than with the default loader, and the KV cache shrinks by the difference. On GLM-5.3 NVFP4, TP4 on B200, "Load weight end" mem usage was 108.89 GB against 106.88 GB for the default loader. That cost 812k KV tokens against 836k.

There are two causes:

1. A lazily created NCCL communicator. InstantTensor all-gathers over get_parallel().world_group.device_group. SGLang never otherwise runs a collective on that group (its world-group uses are CPU-group object collectives and barriers), so the loader's first all-gather creates its NCCL communicator. That communicator's buffers then stay resident for the life of the process: 0.77 GB/GPU on H100 TP4 without NVLS, and more with NVLS.
2. Allocator fragmentation from the clones. copy=True clones every tensor, unsharded, through the shared caching allocator. Long-lived allocations made during and after loading (e.g. post_load_weights' w_kc/w_vc) are carved out of the freed blocks of the largest clone (the 1.8 GB embedding for GLM-5 dimensions). empty_cache() then can't return that segment: 1.74 GB/GPU in a synthetic run.

Modifications

instanttensor_weights_iterator:

• When loading over the world group (the default load group), it creates a temporary NCCL group spanning the world, passes it to InstantTensor, and destroys it when iteration ends. Other load groups, such as a PP stage's TP group, keep using their own device group, since new_group must be called by every rank.
• It pulls each tensor from InstantTensor inside torch.cuda.use_mem_pool with a private MemPool, so only the clones use it, and the pool's segments are freed once the consumer drops the last tensor. It falls back to the shared allocator when use_mem_pool is missing or CUDA is unavailable.

test_instanttensor_loader.py: the existing world-group test now expects the temporary group as InstantTensor's process_group, with new_group / destroy_process_group mocked. That's the only test change; no new tests.

Accuracy / memory results

GLM-5.3 NVFP4 TP4 + MTP, B200	default loader	instanttensor before	instanttensor after
target mem usage after load	106.88 GB	108.89 GB	106.76 GB
KV tokens	836,352	812,544	835,648
target load time	91.7 s	28.8 s	31.3 s

The fix costs about 0.5 s to destroy the communicator, plus a new communicator (~0.5 s) for a second load in the same process, such as the MTP draft. The pool switch costs ~3 µs per tensor.

H100 TP4, synthetic DeepSeek-V3-architecture checkpoint with GLM-5 dimensions:

	mem usage	KV tokens
default loader	9.459 GB	6.114M
instanttensor before	10.549 GB	5.982M
instanttensor after	9.434 GB	6.117M

Two consecutive loads over the default (world) load group, the target + draft pattern: tensors are byte-identical to safetensors, and GPU memory returns to baseline after each load (+2.28 GB before the fix).

End-to-end accuracy with this change (GLM-5.3 NVFP4 TP4 + MTP on B200, --load-format instanttensor): GSM8K (300, temp 0) 99.0% / 98.7% over two runs; IFEval (n=150) 0.88, against a reference of 0.85–0.87 from production serving with the default loader.

Checklist

• test/registered/unit/model_loader/test_instanttensor_loader.py: 9 passed (CPU, as registered) against current main (d67d1ae).
• 2-GPU check on current main, default load path: 0 mismatches against safetensors, temporary group destroyed, memory back to baseline.
• ruff check passes on the touched files.

---

CI States

Latest PR Test (Base): ⏳ Run #42156153742
Latest PR Test (Extra): ❌ Run #42156153624
Latest PR Test (AMD ROCm 10): ⏳ Run #42156153669

Tip

Base CI only covers the core cases, and its workload may not exercise your change. Use /rerun-test to run the tests your change touches, or copy this prompt for an agent:

Find the registered tests this PR affects and run them with /rerun-test, following https://github.com/sgl-project/sglang/blob/main/.agents/skills/rerun-test/SKILL.md

[b8zhong]

Delete all the AI comments that aren't necessary

[Jiminator]

/tag-and-rerun-ci