Skip to content

fix: fix accuracy and illegal memory access issues when using mtp + attention dp #4379

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom
Merged
merged 11 commits into from
Jun 1, 2025
Merged

fix: fix accuracy and illegal memory access issues when using mtp + attention dp #4379

merged 11 commits into from
Jun 1, 2025

Conversation

lfr-0531
Copy link
Collaborator

@lfr-0531 lfr-0531 commented May 16, 2025 •
edited
Loading

Description

Fixed two issues:

  • MTP + attention dp illegal memory access
  • MTP + attention dp + overlap scheduler accuracy drop

The illegal memory access is caused by the token_num in add_dummy_requests. With token_num=1, the prompt_len will be zero. But our attention kernel cannot handle such a special case, so we got illegal memory access. The fix in this PR should be a workaround to make Attention DP+MTP work.

The accuracy drop is because we didn't correctly prepare the inputs for the dummy requests when enabling overlap scheduler and speculative decoding.

Other changes:

  • Change the default py_draft_tokens to an empty list.

NOTE: there is still an illegal memory access issue when running DeepSeek-V3-Lite MTP with tp + autotuner + cuda graph + overlap scheduler. We track this issue in nvbug 5304040.

Test Coverage

The waived tests are added back.

@lfr-0531
Copy link
Collaborator Author

/bot run

@tensorrt-cicd
Copy link
Collaborator

PR_Github #5431 [ run ] triggered by Bot

zhhuang-nv
zhhuang-nv reviewed May 16, 2025
View reviewed changes
@lfr-0531
Copy link
Collaborator Author

/bot kill

@tensorrt-cicd
Copy link
Collaborator

PR_Github #5441 [ kill ] triggered by Bot

@tensorrt-cicd
Copy link
Collaborator

PR_Github #5431 [ run ] completed with state ABORTED

@tensorrt-cicd
Copy link
Collaborator

PR_Github #5441 [ kill ] completed with state SUCCESS
Successfully killed previous jobs for commit f10d532

@lfr-0531
Copy link
Collaborator Author

/bot run

@tensorrt-cicd
Copy link
Collaborator

PR_Github #5442 [ run ] triggered by Bot

@tensorrt-cicd
Copy link
Collaborator

PR_Github #5442 [ run ] completed with state SUCCESS
/LLM/main/L0_MergeRequest_PR pipeline #3972 completed with status: 'FAILURE'

@lfr-0531 lfr-0531 force-pushed the user/fanrongl/fix_input_ids_in_dummy_req branch from 7072ea3 to 7437b2b Compare May 22, 2025 05:38
@lfr-0531 lfr-0531 requested a review from a team as a code owner May 22, 2025 05:38
@lfr-0531 lfr-0531 requested review from mikeiovine and lucaslie May 22, 2025 05:38
@lfr-0531 lfr-0531 mentioned this pull request May 23, 2025
Closed
@lfr-0531 lfr-0531 force-pushed the user/fanrongl/fix_input_ids_in_dummy_req branch from 0a2ffda to 3ee3361 Compare May 23, 2025 10:27
@lfr-0531
Copy link
Collaborator Author

/bot run --disable-fail-fast

@tensorrt-cicd
Copy link
Collaborator

PR_Github #6297 [ run ] triggered by Bot

@lfr-0531 lfr-0531 changed the title fix: fix input ids when adding dummy requests fix: fix accuracy and illegal memory access issues when using mtp + attention dp May 23, 2025
@lfr-0531 lfr-0531 requested a review from PerkzZheng May 23, 2025 12:39
@tensorrt-cicd
Copy link
Collaborator

PR_Github #6297 [ run ] completed with state SUCCESS
/LLM/main/L0_MergeRequest_PR pipeline #4602 completed with status: 'FAILURE'

@PerkzZheng
Copy link
Collaborator

@lfr-0531 have all the illegal memory access issues been resolved or the one you mentioned still exists ? thanks.

@lfr-0531
Copy link
Collaborator Author

@lfr-0531 have all the illegal memory access issues been resolved or the one you mentioned still exists ? thanks.

No, we haven't fixed it. But looks like we didn't cover this special case in our test list before. I mean the setting mtp_nextn=2-attention_dp=False-cuda_graph=True-overlap_scheduler=True. I want to merge this PR first so that we can have other waived tests back. And we can track this special case in a new nvbug and fix it later.

@PerkzZheng, what do you think?

@lfr-0531
Copy link
Collaborator Author

/bot run

@PerkzZheng
Copy link
Collaborator

@lfr-0531 have all the illegal memory access issues been resolved or the one you mentioned still exists ? thanks.

No, we haven't fixed it. But looks like we didn't cover this special case in our test list before. I mean the setting mtp_nextn=2-attention_dp=False-cuda_graph=True-overlap_scheduler=True. I want to merge this PR first so that we can have other waived tests back. And we can track this special case in a new nvbug and fix it later.

@PerkzZheng, what do you think?

that makes sense to me.

@tensorrt-cicd
Copy link
Collaborator

PR_Github #6382 [ run ] triggered by Bot

@tensorrt-cicd
Copy link
Collaborator

PR_Github #6382 [ run ] completed with state SUCCESS
/LLM/main/L0_MergeRequest_PR pipeline #4664 completed with status: 'FAILURE'

lfr-0531 added 11 commits June 1, 2025 01:16
@lfr-0531
fix input ids when adding dummy requests.
e799438 
Signed-off-by: Fanrong Li <[email protected]>
@lfr-0531
fix token_nums.
06c2d03 
Signed-off-by: Fanrong Li <[email protected]>
@lfr-0531
fix the dummy requests + overlap scheduler + spec decoding issue.
f7f5c37 
Signed-off-by: Fanrong Li <[email protected]>
@lfr-0531
fix prepare inputs for the dummy requests in dp+mtp.
bdd28b5 
Signed-off-by: Fanrong Li <[email protected]>
@lfr-0531
fix prompt_len.
2d4ace6 
Signed-off-by: Fanrong Li <[email protected]>
@lfr-0531
add extra tokens.
49cd272 
Signed-off-by: Fanrong Li <[email protected]>
@lfr-0531
add waived tests back.
a7ef090 
Signed-off-by: Fanrong Li <[email protected]>
@lfr-0531
update the py_draft_tokens default to an empty list and add some expl…
cc11e44 
…anatory comments.

Signed-off-by: Fanrong Li <[email protected]>
@lfr-0531
fix py_draft_tokens in dis-agg.
7b974aa 
Signed-off-by: Fanrong Li <[email protected]>
@lfr-0531
waive test_fp8_block_scales tests due to unstable errors.
2511aec 
Signed-off-by: Fanrong Li <[email protected]>
@lfr-0531
waive test_bfloat16 tests due to unstable errors in post-merge.
bead11c 
Signed-off-by: Fanrong Li <[email protected]>
@tensorrt-cicd
Copy link
Collaborator

PR_Github #7154 [ run ] completed with state SUCCESS
/LLM/main/L0_MergeRequest_PR pipeline #5170 completed with status: 'SUCCESS'

@lfr-0531 lfr-0531 force-pushed the user/fanrongl/fix_input_ids_in_dummy_req branch from 7175803 to bead11c Compare June 1, 2025 08:24
@lfr-0531
Copy link
Collaborator Author

lfr-0531 commented Jun 1, 2025

/bot run

@lfr-0531
Copy link
Collaborator Author

lfr-0531 commented Jun 1, 2025

Got another conflict. Rerun the pipeline.

@tensorrt-cicd
Copy link
Collaborator

PR_Github #7160 [ run ] triggered by Bot

@tensorrt-cicd
Copy link
Collaborator

PR_Github #7160 [ run ] completed with state SUCCESS
/LLM/main/L0_MergeRequest_PR pipeline #5175 completed with status: 'FAILURE'

@lfr-0531
Copy link
Collaborator Author

lfr-0531 commented Jun 1, 2025

/bot run

@tensorrt-cicd
Copy link
Collaborator

PR_Github #7170 [ run ] triggered by Bot

@tensorrt-cicd
Copy link
Collaborator

PR_Github #7170 [ run ] completed with state SUCCESS
/LLM/main/L0_MergeRequest_PR pipeline #5185 completed with status: 'SUCCESS'

@lfr-0531 lfr-0531 merged commit 7d356ef into NVIDIA:main Jun 1, 2025
3 checks passed
lfr-0531 added a commit to lfr-0531/TensorRT-LLM that referenced this pull request Jun 2, 2025
@lfr-0531
fix: fix accuracy and illegal memory access issues when using mtp + a…
e0dffa3 
…ttention dp (NVIDIA#4379)

Signed-off-by: Fanrong Li <[email protected]>
@lfr-0531 lfr-0531 mentioned this pull request Jun 2, 2025
Merged
@mikeiovine mikeiovine mentioned this pull request Jun 2, 2025
Merged
lfr-0531 added a commit that referenced this pull request Jun 3, 2025
@lfr-0531
Cherry-pick #4379 (#4833)
6e46e13 
Signed-off-by: Fanrong Li <[email protected]>
darraghdog pushed a commit to darraghdog/TensorRT-LLM that referenced this pull request Jun 3, 2025
@lfr-0531 @darraghdog
fix: fix accuracy and illegal memory access issues when using mtp + a…
72ad3e8 
…ttention dp (NVIDIA#4379)

Signed-off-by: Fanrong Li <[email protected]>
Signed-off-by: darraghdog <[email protected]>
@lfr-0531 lfr-0531 mentioned this pull request Jun 20, 2025
Merged
k-l-lambda pushed a commit to k-l-lambda/TensorRT-LLM that referenced this pull request Jun 23, 2025
@lfr-0531 @k-l-lambda
Cherry-pick NVIDIA#4379 (NVIDIA#4833)
9935911 
Signed-off-by: Fanrong Li <[email protected]>
@lfr-0531 lfr-0531 deleted the user/fanrongl/fix_input_ids_in_dummy_req branch June 27, 2025 12:43
@symphonylyh symphonylyh mentioned this pull request Aug 8, 2025
Closed
4 tasks
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@liji-nv liji-nv liji-nv left review comments

@syuoni syuoni syuoni approved these changes

@QiJune QiJune QiJune approved these changes

@schetlur-nv schetlur-nv Awaiting requested review from schetlur-nv

@mikeiovine mikeiovine Awaiting requested review from mikeiovine mikeiovine was automatically assigned from NVIDIA/trt-llm-torch-devs

@PerkzZheng PerkzZheng Awaiting requested review from PerkzZheng

+1 more reviewer

@zhhuang-nv zhhuang-nv zhhuang-nv approved these changes

Reviewers whose approvals may not affect merge requirements
Assignees
No one assigned
Labels
None yet
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
7 participants
@lfr-0531 @tensorrt-cicd @PerkzZheng @QiJune @syuoni @liji-nv @zhhuang-nv