Conversation
Summary of ChangesHello, I'm Gemini Code Assist1! 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! This pull request optimizes the Expert Parallelism (EP) communication path by adopting a consistent padded layout for gather and reduce operations. By keeping tensors in a padded state (ep_size, max_local_size) throughout the process, the implementation removes redundant host-side operations like slicing and concatenation. While this approach may introduce minor redundant computations on faster ranks, it significantly improves overall performance by eliminating overhead in the hot path of MoE layers. Highlights
New Features🧠 You can now enable Memory (public preview) to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console. Using Gemini Code AssistThe 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 by creating a comment using either
Customization To customize the 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 Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counterproductive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for GitHub and other Google products, sign up here. Footnotes
|
|
This pull request has conflicts, please resolve those before we can evaluate the pull request. |
|
👋 Hi! Thank you for contributing to the vLLM Ascend project. The following points will speed up your PR merge:
If CI fails, you can run linting and testing checks locally according Contributing and Testing. Tip 💡 Consider Linking a Related Issue or RFCYour PR title contains the [BugFix] tag, indicating a bug fix or new feature. Linking a related issue or RFC in the PR description is strongly encouraged — it gives reviewers helpful context and speeds up the review. You can use any of these keywords:
🙏 Thanks for helping us keep the project well-organized! |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
Suggested PR Title:
[Ops][Feature] Optimize EP communication by keeping padded layout in MoESuggested PR Summary:
### What this PR does / why we need it?
This PR optimizes the Expert Parallel (EP) communication path in MoE by keeping the tensors in a padded layout throughout the MoE execution (v2 padded-layout semantics). Instead of unpadding after `all_gather` and repadding before `reduce_scatter`, the implementation now keeps the padded layout. This completely eliminates the host overhead associated with per-shard Python slicing and concatenation loops, which is particularly beneficial for deep networks (e.g., 61 layers in prefill hot paths).
Additionally, two critical issues were identified during the review:
1. Potential runtime errors or silent data corruption in `all_gather` if the padded tensor is non-contiguous. It is recommended to call `.contiguous()` on the padded tensor.
2. Potential `RuntimeError` when calling `.view()` on a non-contiguous tensor `x` in `_maybe_pad_and_reduce_impl`. It is recommended to use `.reshape(...).contiguous()` before passing the tensor to `reduce_scatter`.
### Does this PR introduce _any_ user-facing change?
No. This is an internal performance optimization for EP communication.
### How was this patch tested?
The changes were verified by updating the existing unit tests in `tests/ut/ops/test_register_custom_ops.py` to match the new padded-layout shapes and assertions.| # gather 后再按各 rank 真实的 local_sizes 截回。 | ||
| # all_gather 要求各 rank 输入等长: pad 到 max_local_size。 | ||
| # 等分片时 _pad_to_ep_local_size 零开销原样返回。 | ||
| x = _pad_to_ep_local_size(x, max_local_size) |
There was a problem hiding this comment.
If x is already of size max_local_size, _pad_to_ep_local_size returns x directly. If x is non-contiguous (e.g., from sequence parallel slicing), passing it directly to all_gather can cause silent data corruption or runtime errors. Ensuring contiguity here is highly recommended.
| x = _pad_to_ep_local_size(x, max_local_size) | |
| x = _pad_to_ep_local_size(x, max_local_size).contiguous() |
| reduced = ep_group.reduce_scatter(padded_x.view(-1, *x.shape[1:]), 0) | ||
| # 与 v2 gather 配对: x 行数 = len(local_sizes) * max(local_sizes), | ||
| # 每 rank 的真实 token 位于其分片头部, pad 尾部会被 slice 丢弃。 | ||
| reduced = ep_group.reduce_scatter(x.view(-1, *x.shape[1:]), 0) |
There was a problem hiding this comment.
Since x is the output of MoE layers, it may not be contiguous. Calling .view() on a non-contiguous tensor will raise a RuntimeError. Furthermore, collective communication operations like reduce_scatter require contiguous input tensors to avoid silent data corruption or runtime crashes. Use .reshape() and ensure contiguity by calling .contiguous() before passing to reduce_scatter.
| reduced = ep_group.reduce_scatter(x.view(-1, *x.shape[1:]), 0) | |
| reduced = ep_group.reduce_scatter(x.reshape(-1, *x.shape[1:]).contiguous(), 0) |
…tom ops The EP/SP dispatch path unpads the all_gather output with a per-shard python slicing loop + torch.cat, and the finalize path re-pads it with new_zeros + a copy loop before reduce_scatter. On the MoE prefill hot path (3 gather calls per MoE layer) this host-side packing/unpacking costs ~10% extra host time per forward and congests the collective queue (measured on DeepSeek-V3.1 W4A8 TP8/SP: -4.1% overall TPS, TTFT +6.8%). Restore the rc1 sp_by_pass semantics: once the input is padded to max_local_size for the all_gather, keep the padded layout all the way to the paired reduce_scatter, which then runs zero-alloc / zero-copy and only slices the pad tail off its own shard. Extra pad rows on fast ranks are wall-clock free (they wait for the slow rank at the collective anyway). Notes: - the gather fake shape now matches len(local_sizes) * max(local_sizes) - draft models and PCP-expanded layouts keep the same code path - the DP (non-SP) unpad/repack paths are untouched Signed-off-by: Xuyzhen <958522639@qq.com>
833c691 to
269a7d3
Compare
1 participant
…(rc1 sp_by_pass restore)
What this PR does / why we need it?
Restores the rc1
sp_by_passsemantics for the EP/SP MoE dispatch/finalize custom ops, fixing a prefill-side performance regression.Problem. In
_maybe_all_gather_and_maybe_unpad_impl, after padding each rank's input tomax_local_sizeand all-gathering, the SP path unpacks the gathered tensor with a per-shard python slicing loop +torch.cat. In_maybe_pad_and_reduce_impl, the finalize path re-packs the tensor withnew_zeros+ a copy loop beforereduce_scatter. This pack/unpack pair runs on the MoE hot path (3 gather calls per MoE layer, ~60 layers per prefill) and is pure host-side overhead. Measured on DeepSeek V3.1-terminated W4A8, TP8+SP, 16K/1K mixed workload (Atlas A2, ~100s profile window):Fix. Once the input is padded to
max_local_sizefor the all_gather, keep the padded layout all the way to the pairedreduce_scatter:reduce_scatter(x.view(-1, ...))directly (zero-alloc, zero-copy), then slice the pad tail off this rank's shardlen(local_sizes) * max(local_sizes)The extra pad rows computed by fast ranks are wall-clock free (they already wait for the slow rank at the collective), and the numerical output is unchanged: pad rows are dropped before the result leaves the reduce op.
Consumer audit (why the padded layout is safe downstream):
all_gather_input_ids()gathersinput_idsthrough the same op, so its rows stay row-aligned with the padded hidden states; hash routing is row-independent and masks pad rows (input_ids == -1)_EXTRA_CTX.num_tokens(attention-side metadata), not the prepare-side token count, so it is unaffectedDoes this PR introduce any user-facing change?
No.
How was this patch tested?
tests/ut/ops/test_register_custom_ops.pyupdated to the padded-layout semantics, including the draft-model EP-layout regression test and the PCP (DP x PCP x TP) token-order round-trip testRebase note
Originally targeted v0.27.1 and rebased onto current main to resolve conflicts. Main-side changes audited during the port:
[BugFix][Core] Fix MTP MoE finalize layout with upstream SP #16081 removed the draft-VL TP allreduce bypass: not reintroduced here; its draft-model regression UT is kept and updated to the padded input
[BugFix] Fix SP/PCP token layout in MoE AllGather #16477 added the PCP expansion in
_get_ep_local_sizes: kept unchanged; the round-trip UT expectation is updated to the padded layoutvLLM main: vllm-project/vllm@ced6857