Skip to content

[TRTLLM-12229][fix] Fix MTP by scratch reuse rewind - #14403

Merged
jiaganc merged 5 commits into
NVIDIA:feat/deepseek_v4from
jiaganc:dsv4-scratch-mtp-fix
May 22, 2026
Merged

[TRTLLM-12229][fix] Fix MTP by scratch reuse rewind#14403
jiaganc merged 5 commits into
NVIDIA:feat/deepseek_v4from
jiaganc:dsv4-scratch-mtp-fix

Conversation

@jiaganc

@jiaganc jiaganc commented May 21, 2026 •
edited
Loading

Copy link
Copy Markdown
Collaborator

@coderabbitai summary

Description

This PR merges jiaganc#1 to feature branch. Thanks @lowsfer!

This PR fixes DeepSeek V4 SWA scratch reuse rewind handling. It preserves the scratch allocation rewind state and accounts for extra KV tokens when rewinding scratch allocations, which is needed for MTP paths using SWA scratch reuse.

Scratch reuse is enabled by default for DeepSeek V4. Set TRTLLM_DSV4_ENABLE_SWA_SCRATCH_REUSE=0 only when scratch reuse needs to be disabled explicitly.

Test Coverage

  • Clean build on b200-02 with CUDA_ARCHITECTURES=100-real
  • Installed TensorRT-LLM in container trtllm-dsv4-scratch-mtp-fix-jenkins-jiaganc
  • Ran pytest tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_cache_manager.py on b200-02 GPU 4: 47 passed
  • Ran dep_perf_flash.sh on GPUs 4,5,6,7 with max_num_tokens=8192, MTP1, SWA1: completed successfully
  • Ran dep_perf_flash.sh on GPUs 4,5,6,7 with max_num_tokens=8192, MTP3, SWA1: completed successfully
  • Ran gsm8k.sh on GPUs 4,5,6,7: average accuracy 96.59

PR Checklist

Please review the following before submitting your PR:

  • PR description clearly explains what and why. If using CodeRabbit summary, please make sure it makes sense.

  • PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.

  • Test cases are provided for new code paths (see test instructions)

  • Any new dependencies have been scanned for license and vulnerabilities

  • CODEOWNERS updated if ownership changes

  • Documentation updated as needed

  • Update tava architecture diagram if there is a significant design change in PR.

  • The reviewers assigned automatically/manually are appropriate for the PR.

  • Please check this after reviewing the above items as appropriate for this PR.

GitHub Bot Help

To see a list of available CI bot commands, please comment /bot help.

@jiaganc
jiaganc requested review from a team as code owners May 21, 2026 09:56
@jiaganc
jiaganc requested review from yuxianq and removed request for a team May 21, 2026 09:56
@jiaganc
jiaganc requested review from lfr-0531 and lowsfer May 21, 2026 10:05
@jiaganc jiaganc changed the title [None][fix] support SWA scratch reuse rewind [TRTLLM-12229][fix] support SWA scratch reuse rewind May 21, 2026
@jiaganc jiaganc changed the title [TRTLLM-12229][fix] support SWA scratch reuse rewind [TRTLLM-12229][fix] Fix MTP by scratch reuse rewind May 21, 2026
@jiaganc

jiaganc commented May 21, 2026

Copy link
Copy Markdown
Collaborator Author

/bot run

@tensorrt-cicd

tensorrt-cicd commented May 21, 2026

Copy link
Copy Markdown
Collaborator

PR_Github #49680 [ run ] triggered by Bot. Commit: 19bf5ba Link to invocation

@tensorrt-cicd

tensorrt-cicd commented May 21, 2026

Copy link
Copy Markdown
Collaborator

PR_Github #49680 [ run ] completed with state SUCCESS. Commit: 19bf5ba
/LLM/main/L0_MergeRequest_PR pipeline #39287 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

  • Please check the failed tests and fix your PR
  • If you cannot view the failures, ask the CI triggerer to share details
  • Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

@jiaganc

jiaganc commented May 22, 2026

Copy link
Copy Markdown
Collaborator Author

/bot run

lowsfer and others added 4 commits May 21, 2026 19:59
@lowsfer @jiaganc
[None][fix] support SWA scratch reuse rewind
c7816b4 
Signed-off-by: Yao Yao <lowsfer@users.noreply.github.com>
@jiaganc
[None][fix] use extra KV tokens for scratch rewind
d8f5b9b 
Signed-off-by: Jiagan Cheng <jiaganc@nvidia.com>
@jiaganc
[None][fix] enable DSV4 scratch reuse by default
eb67fa4 
Signed-off-by: Jiagan Cheng <jiaganc@nvidia.com>
@jiaganc
test: update KV cache scratch resize coverage
d05ac24 
Signed-off-by: Jiagan Cheng <jiaganc@nvidia.com>
@jiaganc jiaganc mentioned this pull request May 22, 2026
Closed
@jiaganc
jiaganc force-pushed the dsv4-scratch-mtp-fix branch from f91f116 to d05ac24 Compare May 22, 2026 03:14
@jiaganc

jiaganc commented May 22, 2026

Copy link
Copy Markdown
Collaborator Author

/bot run

@tensorrt-cicd

tensorrt-cicd commented May 22, 2026

Copy link
Copy Markdown
Collaborator

PR_Github #49839 [ run ] triggered by Bot. Commit: d05ac24 Link to invocation

@tensorrt-cicd

tensorrt-cicd commented May 22, 2026

Copy link
Copy Markdown
Collaborator

PR_Github #49839 [ run ] completed with state SUCCESS. Commit: d05ac24
/LLM/main/L0_MergeRequest_PR pipeline #39423 completed with status: 'SUCCESS'

CI Report

Link to invocation

@jiaganc
style: format resource manager import
b0209a4 
Signed-off-by: Jiagan Cheng <jiaganc@nvidia.com>
@jiaganc
jiaganc merged commit 7a4b113 into NVIDIA:feat/deepseek_v4 May 22, 2026
5 checks passed
@jiaganc
jiaganc deleted the dsv4-scratch-mtp-fix branch May 22, 2026 06:57
@lancelly lancelly mentioned this pull request May 25, 2026
Merged
lfr-0531 pushed a commit to lfr-0531/TensorRT-LLM that referenced this pull request Jun 10, 2026
@jiaganc @lowsfer @lfr-0531
[TRTLLM-12229][fix] Fix MTP by scratch reuse rewind (NVIDIA#14403)
6772cd9 
Signed-off-by: Yao Yao <lowsfer@users.noreply.github.com>
Signed-off-by: Jiagan Cheng <jiaganc@nvidia.com>
Co-authored-by: Yao Yao <lowsfer@users.noreply.github.com>
lfr-0531 pushed a commit to lfr-0531/TensorRT-LLM that referenced this pull request Jun 11, 2026
@jiaganc @lowsfer @lfr-0531
[TRTLLM-12229][fix] Fix MTP by scratch reuse rewind (NVIDIA#14403)
771383b 
Signed-off-by: Yao Yao <lowsfer@users.noreply.github.com>
Signed-off-by: Jiagan Cheng <jiaganc@nvidia.com>
Co-authored-by: Yao Yao <lowsfer@users.noreply.github.com>
lfr-0531 pushed a commit to lfr-0531/TensorRT-LLM that referenced this pull request Jun 12, 2026
@jiaganc @lowsfer @lfr-0531
[TRTLLM-12229][fix] Fix MTP by scratch reuse rewind (NVIDIA#14403)
effe6fb 
Signed-off-by: Yao Yao <lowsfer@users.noreply.github.com>
Signed-off-by: Jiagan Cheng <jiaganc@nvidia.com>
Co-authored-by: Yao Yao <lowsfer@users.noreply.github.com>
jiaganc added a commit to jiaganc/TensorRT-LLM that referenced this pull request Jun 29, 2026
@jiaganc
[TRTLLM-12229][fix] Fix MTP by scratch reuse rewind (NVIDIA#14403)
ee8c0e5 
Source-Commit: 7a4b113
Signed-off-by: Jiagan Cheng <jiaganc@nvidia.com>
jiaganc added a commit to jiaganc/TensorRT-LLM that referenced this pull request Jun 29, 2026
@jiaganc
[TRTLLM-12229][fix] Fix MTP by scratch reuse rewind (NVIDIA#14403)
814d272 
Source-Commit: 7a4b113
Signed-off-by: Jiagan Cheng <jiaganc@nvidia.com>
jiaganc added a commit to jiaganc/TensorRT-LLM that referenced this pull request Jul 1, 2026
@jiaganc
[TRTLLM-12229][fix] Fix MTP by scratch reuse rewind (NVIDIA#14403)
7a3cfa5 
Source-Commit: 7a4b113
Signed-off-by: Jiagan Cheng <jiaganc@nvidia.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@lfr-0531 lfr-0531 lfr-0531 approved these changes
@yuxianq yuxianq Awaiting requested review from yuxianq yuxianq is a code owner automatically assigned from NVIDIA/trt-llm-torch-attention-devs
@lowsfer lowsfer Awaiting requested review from lowsfer

Assignees

@jiaganc jiaganc

Labels

deepseek-v4

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

4 participants

@jiaganc @tensorrt-cicd @lfr-0531 @lowsfer