Skip to content

[Bugfix][SM70] Release unused MoE conversion cache between layers - #9

Merged
Leonccaa merged 1 commit into
mainfrom
fix/sm70-moe-load-cache-release
Sep 9, 2026
Merged

Leonccaa merged 1 commit into
mainfrom
fix/sm70-moe-load-cache-release

Conversation

@Leonccaa

@Leonccaa Leonccaa commented Sep 5, 2026 •

Copy link
Copy Markdown
Owner

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.

Format Mode / repeat Worst-GPU startup peak, baseline → fix (MiB) Ready live allocation per GPU, unchanged (MiB)
AWQ eager 1 32,494 → 24,724 22,111.96
AWQ eager 2 32,488 → 24,724 22,111.96
AWQ graph 32,492 → 24,658 22,196.23
NVFP4 eager 1 32,394 → 24,626 22,204.28
NVFP4 eager 2 32,328 → 24,626 22,204.28
NVFP4 graph 32,494 → 24,494 22,220.55

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 one FULL_AND_PIECEWISE graph 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 _C binary 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 took times 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):

VLLM_TARGET_DEVICE=cpu .venv/bin/python -m pytest tests/quantization/test_sm70_moe_load_cache_release.py -q
pre-commit run --files vllm/model_executor/layers/quantization/awq_sm70_moe.py vllm/model_executor/layers/quantization/nvfp4_sm70_moe.py tests/quantization/test_sm70_moe_load_cache_release.py

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.

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>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant