Skip to content

[BugFix] fix possible shape mismatch when using ROPE - #4655

Merged
wtomin merged 2 commits into
vllm-project:mainfrom
Semmer2:FixROPEIssue
Jul 7, 2026
Merged

[BugFix] fix possible shape mismatch when using ROPE#4655
wtomin merged 2 commits into
vllm-project:mainfrom
Semmer2:FixROPEIssue

Conversation

@Semmer2

@Semmer2 Semmer2 commented Jun 23, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor

Fix possible shape mismatch when using RotaryEmbedding.forward_native. Now code use the first sin and cos for ROPE and broadcast automatically, which also aligned with other forward func, like forward_cuda, forward_hip and forward_npu::apply_rotary_emb_mindiesd

PLEASE FILL IN THE PR DESCRIPTION HERE.

Purpose

Test Plan

enable step_execution and set max_num_seqs = 4, then launch omni with HunyuanImage3.0.
curl with 2 request at same time to trigger step-wise batching.
vLLM Version:

Name: vllm
Version: 0.23.0+empty

vLLM-Omni Commit:
c6ee3f7 + this PR

Test Result

the results with and without this PR generated identical output:
image

BEFORE SUBMITTING: read CONTRIBUTING.md and run the precheck-pr skill with the code agent for a self-check against project conventions.
(anything written below this line will be removed by GitHub Actions)

@chatgpt-codex-connector

chatgpt-codex-connector Bot commented Jun 23, 2026

Copy link
Copy Markdown

Codex usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits.
Credits must be used to enable repository wide code reviews.

hsliuustc0106
hsliuustc0106 reviewed Jun 23, 2026
View reviewed changes

@hsliuustc0106 hsliuustc0106 left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

Looks reasonable for review.

@wtomin

wtomin commented Jun 25, 2026

Copy link
Copy Markdown
Collaborator

LGTM. Could you provide some regression test to prevent the same problem in the future forward_xxx functions?

And, please resolve the DCO errors. @Semmer2

@Semmer2
Semmer2 requested a review from yenuo26 as a code owner June 27, 2026 03:33
@Semmer2

Semmer2 commented Jun 27, 2026

Copy link
Copy Markdown
Contributor Author

@hsliuustc0106 @wtomin Thank you for review. Test results updated, test cases added.

@hsliuustc0106

hsliuustc0106 commented Jun 29, 2026

Copy link
Copy Markdown
Collaborator

is there any other case? 4/5?

@hsliuustc0106 hsliuustc0106 added the diffusion codes related to diffusion models label Jun 29, 2026
@Semmer2

Semmer2 commented Jun 30, 2026

Copy link
Copy Markdown
Contributor Author

is there any other case? 4/5?

You mean test and compare the results with bs=4 and 5?
As I checked the images got same md5 as well.
79C5B0CB-E8A3-4D11-E17D-3D9C5D61867D
1F2C3F00-96C1-4BF0-8F58-603F93E0F2E0

@wtomin wtomin added the ready label to trigger buildkite CI label Jul 5, 2026
@wtomin

wtomin commented Jul 5, 2026

Copy link
Copy Markdown
Collaborator

@Semmer2 Please fix the following error in the CI:

[2026-07-05T14:07:55Z] =========================== short test summary info ============================
--
  | [2026-07-05T14:07:55Z] FAILED tests/diffusion/layers/test_rope_broadcast.py::TestRotaryEmbeddingNativeShapeRegression::test_wan_forward_native_unchanged - RuntimeError: The size of tensor a (8) must match the size of tensor b (16) at non-singleton dimension 2

@Semmer2
Semmer2 requested a review from xuechendi as a code owner July 6, 2026 07:17
@Semmer2

Semmer2 commented Jul 6, 2026

Copy link
Copy Markdown
Contributor Author

@Semmer2 Please fix the following error in the CI:

[2026-07-05T14:07:55Z] =========================== short test summary info ============================
--
  | [2026-07-05T14:07:55Z] FAILED tests/diffusion/layers/test_rope_broadcast.py::TestRotaryEmbeddingNativeShapeRegression::test_wan_forward_native_unchanged - RuntimeError: The size of tensor a (8) must match the size of tensor b (16) at non-singleton dimension 2

Sure, fixed and tested again on the latest main code.

Semmer2 added 2 commits July 6, 2026 20:51
@Semmer2
[BugFix] fix possible shape mismatch when using ROPE
f60ff47 
Fix possible shape mismatch when using RotaryEmbedding.forward_native.
Now code use the first sin and cos for ROPE and broadcast automatically,
which also aligned with other forward func, like forward_cuda, forward_hip
and forward_npu::apply_rotary_emb_mindiesd

Signed-off-by: Semmer <semmer@live.cn>
@Semmer2
[Test] Add test case for RotaryEmbedding::forward_native
bf79584 
Signed-off-by: Semmer <semmer@live.cn>
@Semmer2

Semmer2 commented Jul 6, 2026

Copy link
Copy Markdown
Contributor Author

@Semmer2 Please fix the following error in the CI:

[2026-07-05T14:07:55Z] =========================== short test summary info ============================
--
  | [2026-07-05T14:07:55Z] FAILED tests/diffusion/layers/test_rope_broadcast.py::TestRotaryEmbeddingNativeShapeRegression::test_wan_forward_native_unchanged - RuntimeError: The size of tensor a (8) must match the size of tensor b (16) at non-singleton dimension 2

Sure, fixed and tested again on the latest main code.

@wtomin Hi could you help to check the npu-ci failed case? I think it has nothing to do with this PR. Thank you.

@wtomin wtomin removed the ready label to trigger buildkite CI label Jul 7, 2026
@wtomin

wtomin commented Jul 7, 2026

Copy link
Copy Markdown
Collaborator

@wtomin Hi could you help to check the npu-ci failed case? I think it has nothing to do with this PR. Thank you.
I agree. The npu-ci failed at docker image launching. I think it's not relevant to this PR.

@wtomin wtomin added the ready label to trigger buildkite CI label Jul 7, 2026
@wtomin
wtomin merged commit 815abce into vllm-project:main Jul 7, 2026
8 of 9 checks passed
ZhengWG pushed a commit to ZhengWG/vllm-omni that referenced this pull request Jul 13, 2026
@Semmer2 @ZhengWG
[BugFix] fix possible shape mismatch when using ROPE (vllm-project#4655)
5ebca24 
Signed-off-by: Semmer <semmer@live.cn>
Signed-off-by: ZhengWG <zwg0606@gmail.com>
wkutak pushed a commit to wkutak/vllm-omni-cosmos-genai-nim that referenced this pull request Jul 13, 2026
@Semmer2 @wkutak
[BugFix] fix possible shape mismatch when using ROPE (vllm-project#4655)
8acb667 
Signed-off-by: Semmer <semmer@live.cn>
Signed-off-by: Wojciech Kutak <wkutak@nvidia.com>
khairulkabir1661 pushed a commit to khairulkabir1661/vllm-omni that referenced this pull request Sep 25, 2026
@Semmer2
[BugFix] fix possible shape mismatch when using ROPE (vllm-project#4655)
e02f7af 
Signed-off-by: Semmer <semmer@live.cn>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@hsliuustc0106 hsliuustc0106 hsliuustc0106 left review comments
@wtomin wtomin wtomin approved these changes
@Isotr0py Isotr0py Awaiting requested review from Isotr0py Isotr0py is a code owner
@princepride princepride Awaiting requested review from princepride princepride is a code owner
@SamitHuang SamitHuang Awaiting requested review from SamitHuang SamitHuang is a code owner
@ZJY0516 ZJY0516 Awaiting requested review from ZJY0516 ZJY0516 is a code owner
@RuixiangMa RuixiangMa Awaiting requested review from RuixiangMa RuixiangMa is a code owner
@david6666666 david6666666 Awaiting requested review from david6666666 david6666666 is a code owner
@yenuo26 yenuo26 Awaiting requested review from yenuo26 yenuo26 is a code owner
@xuechendi xuechendi Awaiting requested review from xuechendi xuechendi is a code owner

Assignees

No one assigned

Labels

diffusion codes related to diffusion models ready label to trigger buildkite CI

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

3 participants

@Semmer2 @wtomin @hsliuustc0106