Skip to content

[Bugfix][Quantization] Reduce NVFP4 MoE load memory under sleep mode#48701

Draft
hubunt wants to merge 1 commit into
vllm-project:mainfrom
hubunt:fix/nvfp4-sleep-mode-load-oom
Draft

[Bugfix][Quantization] Reduce NVFP4 MoE load memory under sleep mode#48701
hubunt wants to merge 1 commit into
vllm-project:mainfrom
hubunt:fix/nvfp4-sleep-mode-load-oom

Conversation

@hubunt

@hubunt hubunt commented Jul 15, 2026

Copy link
Copy Markdown

Purpose

Fixes #48680.

NVFP4 ModelOpt MoE conversion retained four lists of per-expert transformed tensors and then created four additional full-size outputs with torch.stack. Under --enable-sleep-mode, those overlapping allocations increase pressure on the CuMemAllocator private pool during weight loading.

This change preallocates the final W13/W2 weight and scale tensors, writes permuted weights directly with torch.index_select(..., out=...), and reuses one scale scratch buffer per scale type. It removes the per-expert result lists and final stack copies without changing the output layout.

This is not a duplicate of #45589: that change addressed the FlashInfer TRTLLM BF16 MoE conversion path, while this PR addresses the separate ModelOpt NVFP4 conversion path. It is intentionally a path-specific fix and does not change CuMemAllocator globally.

AI assistance (OpenAI Codex) was used; the submitter reviewed the resulting changes.

Test Plan

Test Result

  • .venv/bin/python -m pytest tests/kernels/moe/test_flashinfer_fp4_moe_utils.py -q: 1 passed.
  • Ruff check and format checks passed.
  • Python compilation and git diff --check passed.
  • Applicable pre-commit hooks passed.
  • Manual old/new algorithm comparison produced byte-for-byte identical outputs.
  • A synthetic Nemotron-like conversion shape reduced measured transient peak allocation from 0.996 GiB to 0.793 GiB (20.4%).

End-to-end SM120 model validation is requested from the issue reporter before this Draft is marked ready.

Co-authored-by: OpenAI Codex <noreply@openai.com>
Signed-off-by: hubunt <150658615+hubunt@users.noreply.github.com>
@github-actions

Copy link
Copy Markdown

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

PRs do not trigger a full CI run by default. Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging.

To run CI, PR reviewers can either: Add ready label to the PR or enable auto-merge.

If you have any questions, please reach out to us on Slack at https://slack.vllm.ai.

Agent Guidelines

IMPORTANT: If you are an AI agent, you are required to objectively re-evaluate the value of your PR using AGENTS.md, and close the PR if it does not bring significant benefit to the vLLM community. Failure to do so may result in an immediate ban.

🚀

@mergify mergify Bot added nvidia bug Something isn't working labels Jul 15, 2026
@hubunt

hubunt commented Jul 15, 2026

Copy link
Copy Markdown
Author

@JasonVranek, could you please validate this Draft on the SM120 setup from #48680? In particular, please check that the NVFP4 model loads with --enable-sleep-mode using the original reproduction configuration, and that a basic inference request returns correct output. The branch is hubunt:fix/nvfp4-sleep-mode-load-oom. Once that validation is available, this PR can be marked ready and a reviewer can trigger the full CI run.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

bug Something isn't working nvidia

Projects

Status: No status

1 participant