Skip to content

Fix(spec decode): catch up trim bug#335

Merged
LorrinWWW merged 3 commits into
lightseekorg:mainfrom
rjzhb:fix/spec-decode-catch-up-trim-bug
Jun 2, 2026
Merged

Fix(spec decode): catch up trim bug#335
LorrinWWW merged 3 commits into
lightseekorg:mainfrom
rjzhb:fix/spec-decode-catch-up-trim-bug

Conversation

@rjzhb

@rjzhb rjzhb commented Jun 1, 2026

Copy link
Copy Markdown
Collaborator

Summary

fix #217 that broke spec-decode catch-up on Qwen NextN and DeepSeek V3 NextN:

  • catch-up seq_lens trim in drafter/eagle.py
  • Llama Eagle3 pre-attention q-slice (dispatch B) in llama_eagle3.py

The trim assumed q was sliced to 1 row before the kernel. Llama Eagle3 had the matching q-slice; Qwen NextN and DeepSeek V3 NextN didn't — trim alone left q_len > cache_seqlens, giving negative kernel positions, all-masked softmax, NaN, and accept_rate=0.

Solution is to add a temp flag: pre_attention_trim.

NOTICE: this flag would be removed in the future refactor PR once adding pre_attention_trim feature to both Qwen NextN and DeepSeek V3 NextN.

Test Plan

  • Qwen3.5-35B-A3B-NVFP4 + MTP local serve: no NaN in catch-up attention output, accept_rate > 0 restored.
  • Llama Eagle3 (MiniMax-M2.5-NVFP4 target + thoughtworks/MiniMax-M2.5-Eagle3 draft): serves without regression.

rjzhb added 2 commits June 1, 2026 21:00
… NextN NaN)

Signed-off-by: rjzhb <rjzhb222@163.com>
…ch-up-trim-bug

Signed-off-by: rjzhb <rjzhb222@163.com>
@rjzhb
rjzhb requested a review from a team as a code owner June 1, 2026 21:07

@chatgpt-codex-connector chatgpt-codex-connector Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: c87743f7ec

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

  • Open a pull request for review
  • Mark a draft as ready
  • Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

Comment thread python/tokenspeed/runtime/models/llama_eagle3.py Outdated

@chatgpt-codex-connector chatgpt-codex-connector Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 0c7d1ef3b5

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

  • Open a pull request for review
  • Mark a draft as ready
  • Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

Comment thread python/tokenspeed/runtime/execution/drafter/eagle.py Outdated
@rjzhb
rjzhb force-pushed the fix/spec-decode-catch-up-trim-bug branch 2 times, most recently from a0a7a59 to 20d1c61 Compare June 1, 2026 22:30
…ion_trim flag

Signed-off-by: rjzhb <rjzhb222@163.com>
@rjzhb
rjzhb force-pushed the fix/spec-decode-catch-up-trim-bug branch from 20d1c61 to edac557 Compare June 1, 2026 22:41

@LorrinWWW LorrinWWW left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

This is a correctness fix for spec. And we will have a follow-up refactoring to reduce / remove the extra if-else branching.

@LorrinWWW
LorrinWWW merged commit 15385b2 into lightseekorg:main Jun 2, 2026
77 of 92 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants