Skip to content

Conversation

@chenaoxuan
Copy link

@chenaoxuan chenaoxuan commented Oct 27, 2025

What this PR does / why we need it?

MagicMTP Key Project: Introduce the 'block verify' algorithm into the draft token rejection sampling logic, taking into account the correlation between generated draft tokens to improve draft token acceptance rate and increase vllm_ascend runtime throughput.

@github-actions
Copy link

👋 Hi! Thank you for contributing to the vLLM Ascend project. The following points will speed up your PR merge:‌‌

  • A PR should do only one thing, smaller PRs enable faster reviews.
  • Every PR should include unit tests and end-to-end tests ‌to ensure it works and is not broken by other future PRs.
  • Write the commit message by fulfilling the PR description to help reviewer and future developers understand.

If CI fails, you can run linting and testing checks locally according Contributing and Testing.

Copy link
Contributor

@gemini-code-assist gemini-code-assist bot left a comment

Choose a reason for hiding this comment

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

Code Review

This pull request introduces a 'block verify' algorithm for rejection sampling to improve performance. The changes are mainly in rejection_random_sample_pytorch.

My review identified two critical issues in the new implementation that could lead to runtime errors:

  1. A potential UnboundLocalError due to an uninitialized variable last_accepted_token_pos.
  2. A potential ZeroDivisionError when calculating pi as draft_prob can be zero.

I have provided code suggestions to fix both issues. Please review them carefully.

Comment on lines +430 to +453
pi = min(pi * target_prob / draft_prob, 1.0)

output_token_ids[req_idx, pos] = token_id
if draft_prob > 0 and pi >= uniform_prob:
last_accepted_token_pos = pos
rejected = False
else:
rejected = True
Copy link
Contributor

Choose a reason for hiding this comment

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

critical

There is a potential ZeroDivisionError at line 430. draft_prob can be 0, but it is used as a divisor before being checked. The division should be guarded by a check for draft_prob > 0.

            if draft_prob > 0:
                pi = min(pi * target_prob / draft_prob, 1.0)
            else:
                pi = 0.0

            if pi >= uniform_prob:
                last_accepted_token_pos = pos
                rejected = False
            else:
                rejected = True

@chenaoxuan chenaoxuan changed the title [MTP++] New Feature. [MagicMTP] New Feature. Oct 27, 2025
@wangxiyuan wangxiyuan added the ready read for review label Oct 28, 2025
@chenaoxuan
Copy link
Author

chenaoxuan commented Oct 29, 2025

MTP>=3时强制开启MagicMTP优化逻辑:块校验(Block Verify)。
吞吐提升10%,保持测试精度不变。
绕过greedy校验逻辑(当前版本greedy校验导致性能劣化)。

@wangxiyuan wangxiyuan added the ready-for-test start test by label for PR label Oct 29, 2025
Signed-off-by: Aoxuan Chen <[email protected]>
Signed-off-by: chenaoxuan <[email protected]>
Signed-off-by: Aoxuan Chen <[email protected]>
Signed-off-by: chenaoxuan <[email protected]>
Signed-off-by: Aoxuan Chen <[email protected]>
Signed-off-by: chenaoxuan <[email protected]>
Signed-off-by: Aoxuan Chen <[email protected]>
Signed-off-by: chenaoxuan <[email protected]>
Signed-off-by: Aoxuan Chen <[email protected]>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

ready read for review ready-for-test start test by label for PR

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants