[Worker] Load weights without the split limit under expandable segments (canonical twin of #980)

[voipmonitor]

Canonical twin of #980 (merged into the Karmic beta as 00a33e2314), on dev/karmic-kraken. The same change and release fragment; the cherry-pick conflicted only in test context (this branch lacks beta's startup-plan test), and the added and removed lines are identical.

Summary

Upstream vllm-project#41268 loads weights under max_split_size_mb:20 so that cached blocks do not fragment. With expandable_segments:True, which our presets set, the allocator releases free pages rather than whole cached blocks, and the split limit instead strands pages that persistent weights share with freed loading temporaries. empty_cache() cannot unmap them, and the memory profile then counts them as consumed: a smaller KV cache, or with a fixed KV size, less headroom.

_scoped_allocator_max_split now leaves the allocator settings alone when PYTORCH_CUDA_ALLOC_CONF enables expandable segments. The classic allocator keeps the limit.

Measurements

Instrumented gpu_worker.py (allocator stats after loading, before the profile run), rank 0:

model / config	reserved − allocated, with limit	without limit	effect
GLM-5.3-Flash glm53-tp2, NVFP4-CSF stage-read scales (default)	0.59 GiB	0.08 GiB	0.52 GiB/GPU more headroom (fixed 8 GiB KV)
GLM-5.3-Flash glm53-tp2, per-layer scale expansion	0.42 GiB	0.12 GiB	0.31 GiB/GPU
DeepSeek-V4.1-Flash TP4, b12x#479 inline scales before its own allocator fix	1.01 GiB	0.04 GiB	KV 11.96M → 12.86M tokens, load 143 → 120 s

GLM on 2× RTX PRO 6000 Max-Q (.4, GPUs 5-8, beta code b12x d4b73caf / vLLM ff292a9c63); DS4.1 on 4× RTX PRO 6000 (aiserver, GPUs 10-13). Both GLM servers started healthy without the limit (341 s), no OOM; model memory unchanged (83.76 / 83.21 GiB).

Tests

tests/v1/worker/test_gpu_worker.py in the beta image with the beta file overlaid: 79 passed; with this branch's gpu_worker.py and test file: the 4 new cases pass. The new test_weight_loading_split_limit_skips_expandable_segments (4 cases: expandable on, expandable with other options, empty, expandable off) fails on the two expandable cases without the change.

Notes

• Not a duplicate: no open PR in this fork touches _scoped_allocator_max_split; upstream still applies the limit unconditionally.
• No model output change: only the allocator configuration during weight loading differs.
• AI assistance (Claude Code) was used for the investigation, the change, the test and this description.

🤖 Generated with Claude Code

Summary by CodeRabbit

• Performance
  • Reduced stranded GPU memory after model loading when expandable memory segments are enabled: from 0.59 GiB to 0.08 GiB per GPU for GLM-5.3-Flash, and from 1.01 GiB to 0.04 GiB per GPU for DeepSeek-V4.1-Flash.
  • DeepSeek-V4.1-Flash can now use more memory for its KV cache, increasing capacity from 11.96M to 12.86M tokens. No user action is required.

[coderabbitai]

Navigate logical layers of code changes, visualize relationships, and explore their blast radius.

🧰 Additional context used

📚 Code guidelines (1)

AGENTS.md — auto-discovered

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

• Configuration used: Organization UI
• Review profile: CHILL
• Plan: Advanced
• Run ID: a66ec041-8b33-4f49-90d0-031590b507a9

📥 Commits

Reviewing files that changed from the base of the PR and between ab86b70 and 1ac34e0.

📒 Files selected for processing (3)

• .lil/changes/vllm-load-split-expandable-segments.json
• tests/v1/worker/test_gpu_worker.py
• vllm/v1/worker/gpu_worker.py

Included review availability: This review used your included allowance. Your plan provides up to 1 included review per hour; 0 remain after this review.

---

📝 Walkthrough

Walkthrough

The worker now skips setting a temporary allocator split limit when PYTORCH_CUDA_ALLOC_CONF enables expandable_segments. Tests cover enabled and disabled configurations. A release entry reports model loading memory measurements.

Changes

Allocator loading

Layer / File(s)	Summary
Conditional split limit and validation vllm/v1/worker/gpu_worker.py, tests/v1/worker/test_gpu_worker.py, .lil/changes/vllm-load-split-expandable-segments.json	The worker skips the temporary split limit when expandable segments are enabled and retains the existing behavior for other CUDA configurations. Tests cover both paths and restoration of allocator settings. The release entry reports memory and KV-cache measurements for GLM-5.3-Flash and DeepSeek-V4.1-Flash.

Priority: ➖ Normal

Estimated code review effort: 2 (Simple) | ~10 minutes

Change: Bug fix

Suggested reviewers: lukealonso

Merge Risk: ⚪ Minimal · up to 1ac34

The change preserves classic allocator behavior and avoids applying the temporary split limit to expandable segments. No concrete current-head merge risk remains.

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly describes the main change: skipping the allocator split limit during weight loading when expandable segments are enabled.
Docstring Coverage	✅ Passed	Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 3 functions across 2 files. (1 skipped: 1 …
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

✨ Finishing Touches

📝 Generate docstrings

• Commit to this branch
• Create a new PR

🧪 Generate unit tests (beta)

• Commit to this branch
• Create a new PR

• Autopilot · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[voipmonitor]

Closing: dev/karmic-kraken has its own implementation of loading without the split limit under expandable segments (gpu_worker, allocator settings read from torch). This twin is no longer needed.