Skip to content

[Bugfix][SM70] Release unused AWQ/NVFP4 MoE conversion cache between layers - #509

Merged
yangzhuxinyzx merged 1 commit into
1CatAI:mainfrom
Leonccaa:fix/sm70-moe-load-cache-release
Sep 6, 2026
Merged

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

Conversation

@Leonccaa

@Leonccaa Leonccaa commented Sep 5, 2026 •

Copy link
Copy Markdown
Contributor

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 baseline 7e3efe7443cfd25285c31137cc9d767b48e84751. At submission, upstream main is 755baae1d075ee04fa9096b23fc0225b23589a86; its intervening changes are documentation-only, with no differences in vllm/, csrc/, or tests/.

CPU lifecycle regression commands (project virtual-environment path normalized):

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

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:

  • Qwen3.8 Flash Next, one fixed checkpoint per format across A/B; 4×V100, TP4, MTP off, FP16 activations and KV.
  • Context and batched-token limit 4096, max sequences 1, fixed KV allocation 512 MiB per GPU; prefix caching and async scheduling disabled.
  • PyTorch 2.10.0+cu128 / CUDA 12.8, expandable_segments:True, MRv2, PYNCCL, matching hybrid PLE offload in both arms.
  • Eager versus FULL_AND_PIECEWISE graph mode, capture sizes 1 and 2. Comparisons are within the same format and mode, not across AWQ/NVFP4 or eager/graph.
  • NVML sampling every 50 ms measures startup through ready. This is a sampled startup peak, not an instantaneous converter-only peak. Ready allocated values measure live PyTorch allocations, not cached reservations.

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.

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

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 _C binary 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 took times 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.

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.

2 participants