Repository navigation
[Bugfix][SM70] Release unused MoE conversion cache between layers - #9
Merged
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>
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.
Summary
Upstream submission: 1CatAI/1Cat-vLLM#509, at the same tested commit
ebf621e9f607cce5960843ae9926ff2c741a4b00. This personal-fork draft retains the original validation record; the upstream PR contains the submission-oriented description and review discussion.Release unused caching-allocator blocks once per MoE layer, after source parameters have been removed, in the SM70 AWQ and ModelOpt NVFP4 loaders. Prepared weights and inference buffers remain live. No inference hook, converter algorithm, quantization semantics, or kernel is changed.
The loading peak is not the model's steady-state requirement: temporary conversion allocations can leave cached segments retained across layers. Earlier controlled probes found output-bank preallocation alone had only a small effect; the layer-lifecycle cache release addresses the dominant observed peak without a larger converter rewrite.
Validation
All 12 model startups passed (6 matched A/B pairs). Across each pair, all prepared weight/scale hashes for 4 ranks × 48 MoE layers and all prompt/generated token IDs plus stop reasons match exactly. All 48 short requests finish naturally with finite selected-token logprobs. Both eager repeats match within each format/variant. CUDA Graph capture and replay complete in both arms of both formats.
Startup peak decreases on every GPU in every pair; worst-GPU reductions are 7.52–7.81 GiB. Ready live allocations are exactly equal before rounding. This reduces loading headroom requirements, not steady-state weight storage. Equality is evaluated within each format and execution mode, not AWQ against NVFP4 or eager against graph.
Frozen contract: four V100 GPUs, TP4, MTP0, FP16 activations and KV, max context/batched tokens 4096, max sequences 1, fixed 512 MiB KV per GPU,
expandable_segments:True, MRv2 and PYNCCL. AWQ and NVFP4 each use one frozen checkpoint across their A/B arms. Two independent eager repeats per arm, plus oneFULL_AND_PIECEWISEgraph run per arm (capture sizes 1 and 2). Peaks are NVML 50 ms sampled startup maxima through ready, not instantaneous converter-only peaks. Ready allocated memory is PyTorch live allocation, not NVML reserved memory.Native scope: the same frozen
_Cbinary is used in all arms. Its default AWQ indexed-prefill capability fallback remains active. NVFP4 uses an independently compiled, unmodified current-source QPN operator through the existing sidecar hook, identically in both arms; batch-versus-serial and graph replay checks pass for M=1/2/4/8/16. Other unavailable optional native paths retain their ordinary fallback. This is not validation of every latest native kernel or of MTP, larger batches, long contexts, or production throughput.Loading-time samples are retained separately from memory results. NVFP4 checkpoint-read times varied substantially before the cleanup call is reached; no startup speedup or broad latency guarantee is claimed. Short natural-EOS equality checks are a regression gate, not a broad model-quality evaluation.
For the two eager repeats, the maximum-rank
Model loading tooktimes were AWQ baseline 121.46–121.93 s versus fix 117.25–118.29 s, and NVFP4 baseline 150.50–176.49 s versus fix 172.51–178.50 s. NVFP4's earlier checkpoint-read phase itself ranged 108.55–131.73 s versus 130.28–140.29 s. These are observed samples, not an isolated measurement of cleanup overhead or a statistically established latency regression.CPU checks (from the feature worktree; project virtual-environment path normalized below):
Four CPU tests pass; all applicable pre-commit hooks pass. Removing both cleanup calls makes the two success-path tests fail as expected. CPU tests exercise the real loading methods with native conversion/workspaces stubbed; GPU evidence supplies the numerical and memory gate.
Scope and review
Duplicate-work checks found no existing personal-fork PR for these two layer-boundary cleanup calls. The AWQ active-grouped-decode PR changes inference dispatch instead. Related vLLM FP8/FlashInfer preallocation and upload scratch work targets different loaders/backends; it is not a duplicate of this fork-specific SM70 fix.
AI assistance was used for implementation and validation. The unchanged patch has now been submitted as upstream draft 1CatAI#509 at Leon's request. Human review and independent human test execution are not claimed complete. No merge or production deployment is included.