[Deepseek] Fix OOM during DeepSeek R1 startup

[MatthewBonanni]

Purpose

Starting up DeepSeek R1 DP8/EP on 8xH200 currently OOMs at the default --gpu-memory-utilization (0.9). This PR prevents an unnecessary 4 GiB allocation during the post-CG-capture dummy run, allowing it to start up without having to reduce --gpu-memory-utilization. It still OOMs when prompted, though, so while this improves the situation and reflects the original intent of the code, it doesn't solve the problem.

Test Plan

vllm serve deepseek-ai/DeepSeek-R1 -dp 8 --enable-expert-parallel

Test Result

main: OOM during server startup

PR branch: no longer OOMs during startup

---

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.

[gemini-code-assist]

Code Review

This pull request addresses an out-of-memory (OOM) error during the startup of DeepSeek R1 models. The fix is achieved by making a large, unnecessary memory allocation conditional. This allocation, intended for worst-case memory profiling or CUDA graph capture, is now skipped during other dummy runs, such as the warmup phase before graph capture.

The changes are implemented by introducing an is_memory_profile flag in the ForwardContext and using it, along with the cudagraph_runtime_mode, to control the allocation in vllm/v1/attention/backends/mla/common.py.

My review of the changes indicates that the logic is sound and correctly targets the source of the OOM issue. The modifications are clean, well-contained, and effectively resolve the problem without introducing any apparent side effects. The code quality is good, and I have no further suggestions for improvement.

[MatthewBonanni]

cc @LucasWilkinson

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

[zhuohan123]

Should we avoid adding too many things to the forward_context? It is becoming increasingly complicated and I am increasingly worried about this class getting more and more bloated. cc @WoosukKwon @youkaichao

[zhuohan123]

Also seems like this is probably no longer needed with model runner v2?