Skip to content

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

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

Merged
LorrinWWW merged 3 commits into from
Jun 2, 2026
Merged

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
a461586
fix(spec-decode): drop catch-up seq_lens trim and Llama q-slice (Qwen…
rjzhb Jun 1, 2026
c87743f
Merge remote-tracking branch 'upstream/main' into fix/spec-decode-cat…
rjzhb Jun 1, 2026
edac557
fix(spec-decode): gate catch-up seq_lens trim on per-model pre_attent…
rjzhb Jun 1, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
14 changes: 11 additions & 3 deletions python/tokenspeed/runtime/execution/drafter/eagle.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -36,6 +36,7 @@
CaptureHiddenMode,
ForwardMode,
)
from tokenspeed.runtime.models.base import BaseCausalLM
from tokenspeed.runtime.multimodal.inputs import maybe_substitute_mm_pad
from tokenspeed.runtime.utils import get_colorful_logger
from tokenspeed.runtime.utils.nvtx import nvtx_range
Expand Down Expand Up @@ -226,9 +227,16 @@ def _run_first_step(
)
input_ids = maybe_substitute_mm_pad(input_ids, self.mm_pad_substitute_id)
draft_first_step_reduce = forward_mode.is_decode()
if draft_first_step_reduce and self.attn_backend.support_kv_cache_prewrite:
# Trim seq_lens by rejected-draft count so the sliced decode
# query does not attend to dead positions.

# TODO: remove the isinstance/flag gate together with pre_attention_trim
# once Qwen NextN and DeepSeek V3 NextN also pre-slice q.
draft_model = self.draft_model_runner.model
if (
draft_first_step_reduce
and self.attn_backend.support_kv_cache_prewrite
and isinstance(draft_model, BaseCausalLM)
and draft_model.pre_attention_trim
):
correction = (self.spec_num_tokens - draft_input.accept_lengths).to(
self.draft_seq_lens_buf.dtype
)
Expand Down
4 changes: 4 additions & 0 deletions python/tokenspeed/runtime/models/base/causal_lm.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -44,6 +44,10 @@ class BaseCausalLM(nn.Module):

model_cls: type[BaseTransformerModel]

# TODO: temporary; remove in the follow-up refactoring that extends
# pre-attn q-slice to Qwen NextN and DeepSeek V3 NextN.
pre_attention_trim: bool = False

def __init__(
self,
config: PretrainedConfig,
Expand Down
5 changes: 5 additions & 0 deletions python/tokenspeed/runtime/models/llama_eagle3.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -583,6 +583,11 @@ class LlamaForCausalLMEagle3(BaseCausalLM):

model_cls = Eagle3LlamaModel

# Eagle3 catch-up pre-slices q to active row before attn; trim must pair.
# TODO: remove together with the base flag once Qwen NextN / DeepSeek V3
# NextN also pre-slice.
pre_attention_trim: bool = True

def __init__(
self,
config: LlamaConfig,
Expand Down
Loading