Skip to content

[Bugfix] Fix torch.compile error for DP + MoE on CPU Backend#31650

Merged
bigPYJ1151 merged 3 commits intovllm-project:mainfrom
kzwrime:fix-cpu-dp-moe
Jan 6, 2026
Merged

[Bugfix] Fix torch.compile error for DP + MoE on CPU Backend#31650
bigPYJ1151 merged 3 commits intovllm-project:mainfrom
kzwrime:fix-cpu-dp-moe

Conversation

@kzwrime
Copy link
Copy Markdown
Contributor

@kzwrime kzwrime commented Jan 3, 2026

Purpose

Fix #31648 (torch.compile fails for MoE models on CPU backend with -dp 2).

Fix the error by reordering has_flashinfer_trtllm_fused_moe() in post_quant_allgather.

Test Plan

vllm serve "Qwen/Qwen1.5-MoE-A2.7B-Chat" -dp 2

Test Result

The service can start successfully.


Essential Elements of an Effective PR Description Checklist
  • The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
  • The test plan, such as providing test command.
  • The test results, such as pasting the results comparison before and after, or e2e results
  • (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.
  • (Optional) Release notes update. If your change is user facing, please update the release notes draft in the Google Doc.

Copy link
Copy Markdown
Contributor

@gemini-code-assist gemini-code-assist bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Code Review

This pull request effectively resolves a torch.compile error that occurred with MoE models on CPU backends using data parallelism. The core of the fix is in vllm/model_executor/layers/fused_moe/layer.py, where the conditions for the post_quant_allgather flag are reordered. By moving the has_flashinfer_trtllm_fused_moe() check to the end of the expression, the change leverages short-circuiting to prevent torch.compile from tracing this non-traceable function on CPU platforms. This is a clean and correct solution. The accompanying changes in base_device_communicator.py and xpu_communicator.py are necessary API updates to plumb the extra_tensors parameter, which is part of the feature enabled by post_quant_allgather. These changes are implemented safely and do not pose a risk to other backends. The pull request is well-executed and improves the framework's robustness.

Signed-off-by: kunzh <zhikun.wu@outlook.com>
@bigPYJ1151 bigPYJ1151 enabled auto-merge (squash) January 6, 2026 08:36
@github-actions github-actions bot added the ready ONLY add when PR is ready to merge/full CI is needed label Jan 6, 2026
@bigPYJ1151 bigPYJ1151 merged commit bf0f3a4 into vllm-project:main Jan 6, 2026
51 checks passed
LucasWilkinson pushed a commit to neuralmagic/vllm that referenced this pull request Jan 6, 2026
yugong333 pushed a commit to yugong333/vllm that referenced this pull request Jan 9, 2026
akh64bit pushed a commit to akh64bit/vllm that referenced this pull request Jan 16, 2026
dsuhinin pushed a commit to dsuhinin/vllm that referenced this pull request Jan 21, 2026
…oject#31650)

Signed-off-by: kunzh <zhikun.wu@outlook.com>
Signed-off-by: dsuhinin <suhinin.dmitriy@gmail.com>
ItzDEXX pushed a commit to ItzDEXX/vllm that referenced this pull request Feb 19, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

ready ONLY add when PR is ready to merge/full CI is needed

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[Bug][CPU Backend]: torch.compile fails for MoE models on CPU backend with -dp 2

2 participants