[Dong1017]

Purpose

On Ascend, normal H3 component offload can release allocator cache even when the existing retention bounds allow it. In the tested runtime, torch.accelerator.memory_allocated() raises “Allocator for npu is not a DeviceAllocator”; the existing telemetry-failure path then calls empty_cache().

For NPU devices, query allocated and reserved bytes through the registered device module returned by torch.get_device_module(self.device). Keep the generic torch.accelerator path for other devices. The existing 25% cached-memory and 5% free-memory thresholds, force-release behavior, and telemetry-failure cleanup remain unchanged. The patch changes only vllm_omni/diffusion/offloader/module_residency.py.

Related to #8228; independent of #8577, #8579, #8129 and #8229.

Test Plan

vLLM Version: 0.29.0 (paired service startup log).

vLLM-Omni Commit: 0e6791e plus the module_residency.py change in this PR.

• Run existing module-residency, H3 lifecycle, and VAE-release CPU tests.
• On 8× Ascend 910B1, issue three consecutive two-step Ref2VA requests and observe normal retention and forced release on each rank.
• Run an isolated baseline/candidate 50-step pair with the same input: 15-second reference video, seed 2101, flow shifts 12/3, short edge 768, 16:9, and 24 fps. Use text-encoder TP8, DiT USP8/ring1 and DLO all-gather, resident VAE tile/patch8, and 1344×768, 362-frame output with audio. Warm all eight ranks before the measured request.

Test Result

• Nineteen existing module-residency, H3 lifecycle, and VAE-release CPU tests passed.
• Three two-step requests completed on all eight ranks. Across those requests, 24 normal cache checks retained memory when allowed and 48 forced calls released cache; no allocator-telemetry error or OOM occurred.
• In the isolated 50-step pair, all eight baseline normal calls released cache and all eight candidate calls retained it. Normal-call duration changed from 216.822–251.161 ms to 1.334–1.579 ms per rank. The 16 forced releases per arm still released cache.
• The isolated pair's normal-release markers intersected with shared all-rank no-device-task coverage for 221.296680 ms in baseline and 3.164550 ms in candidate, a reduction of 218.132130 ms. This is the direct eight-rank idle evidence for the normal cache decision.
• The isolated client E2E pair measured 2248.795837 to 2244.482967 seconds (-4.312870 seconds). This is one source-isolated observation.

Allocator API reproduction

In the tested torch-npu 2.10.0.post4 runtime, the generic query raises the reported DeviceAllocator error; the registered NPU query returns the allocation count:

python -c 'import torch, torch_npu; print(torch.accelerator.memory_allocated(torch.device("npu:0")))'
python -c 'import torch, torch_npu; d = torch.device("npu:0"); print(torch.get_device_module(d).memory_allocated(d))'

The API reproduction identifies the root cause. The consecutive requests and isolated paired profile validate the cache decision in the actual model path. The isolated cache-decision timing is the direct benefit of this one-file fix.

Full-stack context

The complete Ref2VA combination (this change, #8129, #8229, #8577, #8579, and resident VAE configuration) was measured on 8×Ascend 910B1 64 GB. The fixed workload requests 50 inference steps using a 15-second reference video with audio, seed 2101, short edge 768, aspect ratio 16:9, fps 24 and flow shifts 12/3; output is 1344×768, 362 frames with audio. Each service completed a same-input 50-step warmup on all eight ranks before measurement.

Three alternating unprofiled pairs give client E2E medians 2253.664743 → 2239.304066 s, saving 14.360677 s (0.637%). A separate profiled pair covers all eight ranks: time with no device task on any rank falls 6.631556 → 3.039535 s (−54.17%); the preparation window falls 3.670694 → 0.941927 s (−74.34%). These are results for the complete combination; individual attribution uses the targeted measurements above.

Measured baseline: 0e6791e9c4925c3539352cff2c2738bd48a1a65a; vLLM 0.29.0, PyTorch 2.10.0+cpu, torch-npu 2.10.0.post4. The publication branch is rebased onto main d5a3380103df2fc827f095c60dbfb6e1c5655fd3. Device results refer to the recorded experiment sources; source application, lint, types and CPU regression gates are checked on the publication branch.

Publication-baseline CPU validation (main d5a3380103 plus the three scoped candidates, vLLM 0.30.0 / PyTorch 2.13.0+cu130, accelerator devices hidden): 160 passed, 1 deselected, 17 warnings in 129.78 s, normal process exit. This covers the 23 preparation tests, 18 allocator/lifecycle/revert tests, and 119 KV cases. Tests run in isolation from the full repository plugin fixture, which requires vLLM 0.31. The one deselected existing example-CLI test imports that 0.31 engine interface and fails in the installed 0.30 runtime; its error is retained separately. Required cache-dit and pytest-mock packages are installed in task-local directories; the shared environment is unchanged.

[Dong1017]

Author self-review completed using the repository's precheck-pr, full mode, against main d5a3380103.

• Scope: one production file. NPU allocation telemetry uses PyTorch's registered device module; other devices keep their existing generic accelerator path.
• The existing 25% cached-memory and 5% physically-free bounds, forced release and telemetry-failure cleanup remain unchanged. No cache-policy change, retry mechanism, extra synchronization or dependency is introduced.
• The real torch-npu API reproduction identifies the generic DeviceAllocator failure. Eight-rank requests verify 24 normal retention decisions and 48 forced releases across three consecutive requests; the isolated 50-step pair verifies the same behavior and a 218.132 ms reduction in common device idle at the normal decision.
• Diff-scoped code-quality, examples-policy and simplification checks found no blocking issue. The broad exception handler is pre-existing and its forced-cleanup behavior remains intact.
• Local pre-commit gates and CPU regression results are recorded in the PR. No hooks were bypassed.
• The full-stack profile is reported separately: normal cleanup improves but force cleanup increases, so the roughly unchanged combined cleanup wall time is not presented as a second independent gain.

Ready for maintainer review.

Publication-baseline CPU validation (main d5a3380103 plus the three scoped candidates, vLLM 0.30.0 / PyTorch 2.13.0+cu130, accelerator devices hidden): 160 passed, 1 deselected, 17 warnings in 129.78 s, normal process exit. This covers the 23 preparation tests, 18 allocator/lifecycle/revert tests, and 119 KV cases. Tests run in isolation from the full repository plugin fixture, which requires vLLM 0.31. The one deselected existing example-CLI test imports that 0.31 engine interface and fails in the installed 0.30 runtime; its error is retained separately. Required cache-dit and pytest-mock packages are installed in task-local directories; the shared environment is unchanged.

@vllm-omni-review-bot Please review this PR.

[Dong1017]

@vllm-omni-review-bot Please review this PR.

@david6666666 Could you review this change as the MiniMax-H3 model owner, especially its model semantics, correctness evidence and performance attribution?

Scope: Registered NPU allocator telemetry, preserving existing cache thresholds and forced-release behavior.

The five changes have separate scopes: #8129, #8229, #8577, #8578 and #8579. The shared eight-rank Ref2VA evidence and per-change attribution are summarized in the Issue #8228 update. The existing Draft/Ready status is retained.

[david6666666]

Could we add a CPU regression test for the NPU branch? The current allocator tests use a CPU device and only mock torch.accelerator, so they still pass if this fix is reverted. A fake registered device module could cover retention, threshold release and forced cleanup.

[Dong1017]

Added in ac461aea1: seven CPU cases using the repository's mocker fixture and the existing core_model / cpu marks.

They cover retention, the exact 25% cached-memory / 5% free-memory boundaries, release above/below those thresholds, unavailable registered-device telemetry, and forced cleanup. The tests explicitly require torch.get_device_module(device) and the registered module's memory APIs, and assert that the generic torch.accelerator queries are not called. Forced cleanup must bypass all telemetry.

Validation:

• Full allocator/staging test file: 15 passed.
• Reverting only the NPU production branch, with the same tests: 6 failed, 9 passed. All six NPU telemetry cases detect the regression; forced cleanup remains valid.
• All applicable pre-commit checks, including mypy, CI marks, import restrictions and sign-off, passed.

This follow-up changes only tests/diffusion/offloader/test_module_residency.py; production code is unchanged.

Reproduce in the repository test environment:

python -m pytest -q tests/diffusion/offloader/test_module_residency.py -m "core_model and cpu" --run-level core_model

[chatgpt-codex-connector]

Codex usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits.
Credits must be used to enable repository wide code reviews.

[vllm-omni-review-bot]

This PR appears to belong to: docs/design/module/host_weight_runtime.md, docs/design/module/diffusion/index.md, docs/design/module/diffusion/offloader.md.

Module owners: @david6666666 @lishunyang12 @yuanheng-zhao

Routing: @david6666666 via module of the changed files, module named in the PR description, semantic router, CODEOWNERS; @lishunyang12 via module of the changed files, module named in the PR description, semantic router, CODEOWNERS; @yuanheng-zhao via module of the changed files, module named in the PR description, semantic router, CODEOWNERS

@Dong1017, please review your own changes and leave a short self-review comment describing what you checked. PRs without author self-review may not be assigned a reviewer.

Please take a look when you have a chance. If you would like an automated review, mention @vllm-omni-review-bot in a comment.

[vllm-omni-review-bot]

Omni ReviewBot triage note

Automated triage of commit 9feeaffb7def produced:

• Priority: high. Prompt maintainer attention is suggested.

These are automated triage suggestions only — the final decision belongs to the maintainers.

[Dong1017]

@vllm-omni-review-bot plz review this pr

[vllm-omni-review-bot]

Omni ReviewBot routing record

Assigned Strict on cursor (cursor-grok-4.6-high) under experiment fleet-strict-cursor-grok46-zcode-glm53flash-5050-c5-z10-20261002.

[Dong1017]

@david6666666 can you add the ready label to trigger the CI? thx a lot

[Dong1017]

CI check for head 9feeaffb7defbccdbb45ae716e54a8cccec28805:

NVIDIA 16996, AMD 13582 and NPU 8722 passed, as did pre-commit, both wheel builds, DCO and docs.

The remaining Intel 9861 failure is:
tests/diffusion/lora/test_images_generations_lora.py::test_images_generations_per_request_lora_switching[zimage_turbo_lora] (fixture setup).

The server's startup warmup failed in Transformers' Qwen3 text encoder, at rotate_half → torch.cat((-x2, x1), dim=-1), with level_zero backend failed with error: 40 (UR_RESULT_ERROR_OUT_OF_RESOURCES). The service then never became healthy and the fixture reported Server failed to start within 1200 seconds. The test summary reports 2 passed, 6 skipped, 1 setup error; the overall job exits 137. Exit 137 alone is not proof of an OOM kill.

The production change in this PR is gated on device.type == "npu". The XPU branch still uses the original torch.accelerator.memory_reserved/memory_allocated queries and the same cache policy. The new NPU branch cannot execute in this XPU test, and there are no Z-Image, Transformers, XPU or warmup changes in this PR. This is an XPU runtime/startup failure rather than a regression in the NPU allocator-telemetry change.

A clean-worker rerun of the failing Z-Image LoRA test is the appropriate follow-up. Scoped issue searches did not find a matching tracked report for this exact Level Zero error and test. No production or test changes were made for this failure.

[david6666666]

LGTM now