[Perf] Reuse MoE workspace for DeepGEMM warmup by LucasWilkinson · Pull Request #57268 · vllm-project/vllm · GitHub
Skip to content

[Perf] Reuse MoE workspace for DeepGEMM warmup - #57268

Merged
LucasWilkinson merged 1 commit into
vllm-project:mainfrom
LucasWilkinson:fix/deepgemm-moe-warmup-workspace
Sep 17, 2026
Merged

LucasWilkinson merged 1 commit into
vllm-project:mainfrom
LucasWilkinson:fix/deepgemm-moe-warmup-workspace

Conversation

@LucasWilkinson

@LucasWilkinson LucasWilkinson commented Sep 17, 2026 •

Copy link
Copy Markdown
Contributor

Reuse the existing MoE workspace for DeepGEMM grouped-GEMM warmup activations, scales, and output, avoiding a second large scratch allocation after profiling. Ordinary linear warmup keeps temporary buffers so dense models do not acquire a persistent workspace.

B300 validation with two-layer dummy block-FP8 Qwen3 and Qwen3-MoE models (hidden size 1024, intermediate size 3072, 1024 batched tokens, eager execution):

  • Both completed startup and generated the same four tokens as the baseline.
  • Final workspace stayed at 0 MiB (dense) and 54 MiB (MoE).
  • Additional MoE warmup peak fell from 39.25 MiB to 7.04 MiB.
PYTEST_DISABLE_PLUGIN_AUTOLOAD=1 .venv/bin/python -m pytest \
  tests/model_executor/test_deep_gemm_warmup.py -q
# 6 passed

.venv/bin/pre-commit run --files \
  vllm/model_executor/warmup/deep_gemm_warmup.py
# Passed

Alternative to #57258, which moves warmup before profiling. Related #49469 removes DP scaling; this change preserves warmup shapes and ordering. Duplicate-work searches found no existing PR implementing this workspace reuse.

Codex assisted with implementation and validation. Reviewed by submitter.

Assisted-by: OpenAI Codex
Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>
@robertgshaw2-redhat robertgshaw2-redhat added the ready ONLY add when PR is ready to merge/full CI is needed label Sep 17, 2026
@robertgshaw2-redhat

Copy link
Copy Markdown
Collaborator

@LucasWilkinson
LucasWilkinson marked this pull request as ready for review September 17, 2026 01:41

@claude claude Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

@github-actions

Copy link
Copy Markdown

❌ This PR is 5 commits behind upstream main. Your branch must contain every commit currently on upstream main. No new CI build was started. Merge or rebase onto the latest main, then rerun /ci run. To test this branch at your own risk, use /ci run --allow-stale.

@LucasWilkinson

Copy link
Copy Markdown
Contributor Author

/ci run --allow-stale

@github-actions

Copy link
Copy Markdown

@LucasWilkinson
LucasWilkinson merged commit bd2d7ae into vllm-project:main Sep 17, 2026
123 checks passed
michael-gpu pushed a commit to practical-inference/vllm-ampere-extended that referenced this pull request Oct 20, 2026
Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>
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.

2 participants