[Misc] Fixes and Optimizations for DeepEP + DeepGEMM combination.

[varun-sundar-rabindranath]

Purpose

Issue1:

The engine fails to initialize with DeepseekR1 + data-parallel-size 32 + expert-parallel + VLLM_ALL2ALL_BACKEND="deepep_high_throughput" + VLLM_USE_DEEP_GEMM=1

Cause:

During profile runs we execute model forward pass on all GPUs with maximum batch size. We do this with dummy input_ids all set to zeros. This has an effect of routing all tokens from all GPUs to a specific set of experts / GPU Ranks in the All2All dispatch call. The GPU receiving all the tokens has an enormous batch-size to process. The GPUs OOM as a result.

Issue2:

The engine fails to initialize with DeepseekR1 + data-parallel-size 128 + expert-parallel + VLLM_ALL2ALL_BACKEND="deepep_low_latency" + VLLM_USE_DEEP_GEMM=1

Cause:

RuntimeError: Failed: Assertion error /mnt/data/home/smo/vllm/tools/ep_kernels/ep_kernels_workspace/DeepEP/csrc/deep_ep.cpp:1040 'layout.total_bytes <= num_rdma_bytes'

Issue 3:

CUDA illegal memory access in FP8 block-quant triton kernel

Cause:

Integer overflow issues in the block-quant triton kernel.

Changes / Fixes:

• [issue1 - fix] Randomize _dummy_run input_ids()
  • Randomize dummy run input IDs. This is so all experts of the model and consequently the GPU ranks, receive a balanced number of tokens.
• [issue1 - fix] Reduce memory usage
  • DeepGemm MOE: Reuse preallocated workspaces for Quantization and for "inverse permutation" outputs.
  • DeepEP High Throughput Prepare/Finalize : Remove expensive torch ops in finalize function
• [issue2 - fix] Ask DeepEP for num_rdma_bytes hint
  • DeepEP should give us the right requirements.
• [issue3 - fix] Use int64 in offset calculations (stride multiplication) in the fp8 block-quant kernels.

Test Plan

Tested locally with :

VLLM_ALL2ALL_BACKEND="deepep_high_throughput" VLLM_USE_DEEP_GEMM=1  vllm serve Qwen/Qwen3-30B-A3B-FP8  --trust-remote-code  --data-parallel-size 2 --enable-expert-parallel --port 9010  --no-enable-prefix-caching

and

lm_eval --model local-completions --tasks gsm8k --model_args model=Qwen/Qwen3-30B-A3B-FP8,base_url=http://127.0.0.1:9010/v1/completions,num_concurrent=30,max_retries=3,tokenized_requests=False --limit 100

Test Result

|Tasks|Version|     Filter     |n-shot|  Metric   |   |Value|   |Stderr|
|-----|------:|----------------|-----:|-----------|---|----:|---|-----:|
|gsm8k|      3|flexible-extract|     5|exact_match|↑  | 0.81|±  |0.0394|
|     |       |strict-match    |     5|exact_match|↑  | 0.90|±  |0.0302|

[github-actions]

👋 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.

Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run fastcheck CI which starts running only a small and essential subset of CI tests to quickly catch errors. You can run other CI tests on top of those by going to your fastcheck build on Buildkite UI (linked in the PR checks section) and unblock them. If you do not have permission to unblock, ping simon-mo or khluu to add you in our Buildkite org.

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.

🚀

[gemini-code-assist]

Hello @varun-sundar-rabindranath, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

Summary of Changes

Hello! Gemini here, providing a summary of this pull request to help everyone quickly understand the changes and context.

This PR addresses two specific issues encountered when using the DeepEP and DeepGEMM backends, particularly with high data parallelism and the DeepseekR1 model. The first issue involved Out-of-Memory (OOM) errors during profile runs with dummy zero inputs, which led to unbalanced token distribution across experts and GPUs. The second issue was an assertion failure related to RDMA buffer size when using the low-latency DeepEP backend with high data parallelism. To fix these, the PR introduces randomization for dummy run input IDs to balance expert load, optimizes memory usage in DeepGEMM by reusing preallocated workspaces, and replaces expensive PyTorch operations with potentially more efficient custom ops in the DeepEP high-throughput finalize function.

Highlights

• DeepGEMM Memory Optimization: Modifies the DeepGEMM MOE layer to reuse preallocated workspaces for intermediate results, specifically for the activation output, the second quantization step's output, and the final output after inverse permutation. This reduces overall memory footprint.
• DeepEP High Throughput Optimization: Replaces the torch.sum(dim=1) operation in the DeepEP high-throughput finalize function with a custom ops.moe_sum kernel. This aims to improve performance and potentially reduce memory overhead during the final combination step.
• New Environment Variable: Adds a new environment variable VLLM_RANDOMIZE_DP_DUMMY_INPUTS to control whether dummy inputs should be randomized during data parallel dummy runs.

Changelog

Click here to see the changelog

• vllm/envs.py
  • Added VLLM_RANDOMIZE_DP_DUMMY_INPUTS boolean environment variable (default: False) at line 113.
  • Added logic to parse the VLLM_RANDOMIZE_DP_DUMMY_INPUTS environment variable (checking for '1') at line 765.
• vllm/model_executor/layers/fused_moe/deep_gemm_moe.py
  • Adjusted the calculation for workspace2 size in workspace_shapes to M_sum * max(N, K) from M_sum * N at line 87.
  • Modified the apply method to reuse workspace13 for mm1_out and quant_out (viewed as float8) and workspace2 for act_out and mm2_out at lines 140-145.
  • Replaced the tensor indexing workspace3[inv_perm, ...] with torch.index_select(mm2_out, 0, inv_perm, out=out) to perform the inverse permutation directly into the preallocated out tensor at line 161.
• vllm/model_executor/layers/fused_moe/deepep_ht_prepare_finalize.py
  • Imported vllm._custom_ops as ops at line 8.
  • Replaced fused_expert_output.sum(dim=1).to(output_dtype) with creating an output tensor out and calling ops.moe_sum(fused_expert_output, out) at lines 208-211.
• vllm/model_executor/layers/quantization/utils/fp8_utils.py
  • Added an optional out_q: Optional[torch.Tensor] = None parameter to the per_token_group_quant_fp8 function signature at line 314.
  • Updated the docstring to describe the new out_q parameter at line 326.
  • Modified the function to use the provided out_q tensor if available, otherwise create a new one, at lines 341-344.
• vllm/v1/worker/gpu_model_runner.py
  • Imported contextmanager from contextlib at line 8.
  • Imported vllm.envs at line 16.
  • Added a new context manager maybe_randomize_inputs at lines 1726-1753, which randomizes input_ids if VLLM_RANDOMIZE_DP_DUMMY_INPUTS is true and DP size > 1, and resets them to zeros upon exiting the context.
  • Applied the maybe_randomize_inputs context manager around the _dummy_run call at line 1835.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in issue comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist is currently in preview and may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments to provide feedback.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request introduces crucial fixes and optimizations for DeepEP and DeepGEMM, addressing OOM errors and enhancing memory efficiency. The randomization of dummy inputs is a key fix for expert load balancing, and the memory optimizations in both DeepGEMM and DeepEP modules are well-implemented. The code is clear and the changes are well-justified by the PR description.

Summary of Findings

• Correctness of Workspace Sizing: In deep_gemm_moe.py, the workspace2 sizing was updated. While the new sizing M_sum * max(N, K) correctly fixes issues when K > N, a slightly tighter bound could be M_sum * max(N // 2, K) based on its direct usages for act_out and mm2_out. The current approach is safe, however.
• Memory Optimizations: Significant memory optimizations were made by reusing tensors (e.g., workspace13, workspace2 in DeepGEMM, out_q in quantization) and using in-place operations or custom ops (e.g., mul_ and ops.moe_sum in DeepEP HT finalize). These are excellent for performance and reducing OOM risks.
• Dummy Run Input Randomization: The new VLLM_RANDOMIZE_DP_DUMMY_INPUTS flag and the maybe_randomize_inputs context manager effectively address the expert load imbalance OOM during dummy/profile runs. The implementation is clear and robust.

Merge Readiness

The pull request introduces important fixes and optimizations. The changes are well-implemented and address the described issues effectively. After considering the feedback, I believe this PR is in good shape for merging. I am unable to approve the pull request myself, so please ensure it is reviewed and approved by other maintainers before merging.

[gemini-code-assist]

The in-place multiplication fused_expert_output.mul_(topk_weights.view(m_x_topk, -1, 1)) is a good optimization, replacing the previous type casting and out-of-place multiplication. This assumes that the dtypes of fused_expert_output and topk_weights are compatible for in-place multiplication, or that PyTorch's type promotion handles it correctly without unexpected side effects. Given this is for performance, this is likely fine, but it's good to be mindful of the dtypes involved. Could you confirm if fused_expert_output's dtype (coming from DeepEP kernels) is consistently compatible with topk_weights (usually float/bfloat16) for this in-place operation?

[varun-sundar-rabindranath]

topk_weights is float32 - and fused_expert_output is bfloat16 - the multiplication relies on type promotion/demotion.

[gemini-code-assist]

The addition of the out_q parameter to allow for in-place quantization is a good memory optimization. The assertion assert out_q is None or out_q.shape == x.shape and the conditional assignment x_q = out_q if out_q is not None else torch.empty_like(...) are correctly implemented. This enhances flexibility and efficiency.

[gemini-code-assist]

The maybe_randomize_inputs context manager correctly randomizes input_ids for dummy runs when the VLLM_RANDOMIZE_DP_DUMMY_INPUTS flag is set. The use of torch.randint_like and slicing ensures appropriate random token IDs are generated and copied.

Restoring input_ids to zeros with input_ids.fill_(0) after the yield is crucial for ensuring that subsequent operations or dummy runs (if randomization is off or for other purposes) start with a clean, expected state. This is well handled.

[varun-sundar-rabindranath]

It is incorrect to do this when we are doing profile runs -- during profile runs, we do want the system to be stress tested (i.e. all tokens reaching the same set of GPU ranks).
However, at the moment, without this we are not able to run large scale DP -- things OOM during profile_run.

#19168 should fix the OOM - then we can remove this logic for the profile run case.

[varun-sundar-rabindranath]

This could be optimized. We dont have to fill the input_ids and then set it to zeros every time. For eager-mode runs (i.e. batch size > 512) - we could just use the rand tensor in the place of input-ids .. I plan to do this in a follow up PR.

[varun-sundar-rabindranath]

memory optimization to prevent inv_perm from making a brand-new tensor.

[mgoin]

Is DP important here? I think you would want this for any EP case, so maybe just VLLM_RANDOMIZE_DUMMY_INPUTS

[varun-sundar-rabindranath]

Investigated it for a bit and I think it is better to call out DP in the name. It is only in the context of DP that some DP Ranks execute dummy-runs so we can synchronize with the DP Ranks that run the model with actual tokens.

— The other way we could do expert parallel is with DP=1 and TP > 1 - with this, all the ranks run with actual data (the input data is replicated across all ranks)

also, I have this statement in code,
randomize_inputs = envs.VLLM_RANDOMIZE_DP_DUMMY_INPUTS and dp_size > 1

But I see what you are saying, we could do,
VLLM_RANDOMIZE_DUMMY_INPUTS -> VLLM_RANDOMIZE_DUMMY_INPUTS and randomize_inputs = envs.VLLM_RANDOMIZE_DP_DUMMY_INPUTS and randomize if the env var is just set.
let's do it when more use cases for randomizing dummy runs come up ? What do you think ?

[varun-sundar-rabindranath]

^ rebase on to main