Repository navigation
[Bugfix][SM70] Release unused AWQ/NVFP4 MoE conversion cache between layers - #509
Merged
yangzhuxinyzx merged 1 commit intoSep 6, 2026
Merged
Conversation
Bound load-time caching allocator growth in AWQ and ModelOpt NVFP4 without changing prepared weights or inference kernels. Add CPU lifecycle regression coverage for ordering, once-per-layer cleanup and failed conversion. Co-authored-by: OpenAI Codex <noreply@openai.com> Signed-off-by: Leonccaa <166551845+Leonccaa@users.noreply.github.com>
Leonccaa
marked this pull request as ready for review
September 5, 2026 15:55
This was referenced Sep 6, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Purpose
Reduce the loading-time GPU memory peak of the SM70 AWQ and ModelOpt NVFP4 MoE paths without changing prepared weights or inference kernels.
After each layer finishes conversion, workspace allocation, and source-parameter removal, call
torch.accelerator.empty_cache()once. This releases unused caching-allocator blocks between layers; live prepared weights and inference buffers remain allocated. The call is restricted to the two weight-loading methods, not the inference path.Controlled loading probes identified cross-layer retention of temporary conversion allocations as the dominant observed peak. Output-bank preallocation alone had a much smaller effect, so this change keeps the existing conversion algorithm and avoids a larger preallocation/workspace rewrite. It is not an FP8 conversion port or a change to NVFP4 quantization semantics.
Duplicate check: no open 1Cat PR was found for these two layer-boundary cleanup calls. Earlier model/FP8 integration PRs #345 and #212 are merged, and current main still lacks the calls. This is separate from active-grouped-decode dispatch changes. Original validation/review record: Leonccaa/1Cat-vLLM#9.
Test Plan
Tested head:
ebf621e9f607cce5960843ae9926ff2c741a4b00, against baseline7e3efe7443cfd25285c31137cc9d767b48e84751. At submission, upstream main is755baae1d075ee04fa9096b23fc0225b23589a86; its intervening changes are documentation-only, with no differences invllm/,csrc/, ortests/.CPU lifecycle regression commands (project virtual-environment path normalized):
The tests call the real loading methods with native converters and GPU workspaces stubbed. They check once-per-layer cleanup after source removal, preservation of prepared tensors/buffers, and no success cleanup after a failed conversion. Removing both cleanup calls makes the two success-path tests fail.
GPU acceptance: two independent eager starts per variant and one CUDA Graph start per variant, for both AWQ and NVFP4: 12 model startups / 6 A/B pairs. Record prepared weight/scale hashes for all 4 ranks × 48 MoE layers, ready memory, source/native provenance, and four frozen short prompts per run. Compare prompt/output token IDs and stop reasons with natural EOS (
ignore_eos=false, temperature 0, seed 0, max output 512); require nonempty output and finite selected-token logprobs.Frozen configuration:
expandable_segments:True, MRv2, PYNCCL, matching hybrid PLE offload in both arms.FULL_AND_PIECEWISEgraph mode, capture sizes 1 and 2. Comparisons are within the same format and mode, not across AWQ/NVFP4 or eager/graph.Test Result
All 12 starts and 48 natural-EOS requests pass. Every A/B pair has identical prepared-weight/scale hashes, prompt/generated token IDs and stop reasons, and exactly equal ready live allocations. All four eager repeatability groups pass. Both formats complete actual CUDA Graph capture and inference in both arms.
Every GPU has a lower startup peak in every pair. Worst-GPU reductions are 7.52–7.81 GiB. Ready live allocations are exactly equal before rounding: this reduces startup headroom requirements, not steady-state weight storage.
Four CPU regression tests pass, including a fresh rerun before upstream submission; all applicable pre-commit hooks passed on the unchanged commit. These local results are separate from upstream CI status.
Scope and limitations
The GPU matrix used the same frozen
_Cbinary in both arms. Its default AWQ indexed-prefill capability fallback remains active. NVFP4 used a matching QPN sidecar compiled from the unmodified current CUDA source through the existing registration hook, identically in both arms. Its batch-versus-serial bitwise and CUDA Graph replay checks pass for M=1/2/4/8/16. Other missing optional native capabilities retain their ordinary fallback. The sidecar and experiment harness are not product changes in this PR; this is not validation of every latest native kernel.For the two eager repeats, maximum-rank
Model loading tooktimes were AWQ baseline 121.46–121.93 s versus fix 117.25–118.29 s; NVFP4 baseline 150.50–176.49 s versus fix 172.51–178.50 s. NVFP4's earlier checkpoint-read phase itself varied from 108.55–131.73 s versus 130.28–140.29 s. These samples do not isolate cleanup overhead; no startup-speedup or broad latency guarantee is claimed.The short natural-EOS tests are a regression gate, not broad model-quality or throughput evaluation. MTP, larger batches, long contexts, and other models are outside this acceptance matrix. No model deployment or merge is included.
AI assistance was used for implementation, testing, and PR preparation. On 2026-09-05, Leon confirmed that he had reviewed the code and authorized marking this PR ready for review. The reported test runs were agent-executed; independent human test execution is not asserted.