Skip to content

refactor(spec-decode): simplify Llama Eagle3 attention path for #217 (1/3)#390

Merged
LorrinWWW merged 12 commits into
lightseekorg:mainfrom
rjzhb:refactor/llama-attention-hooks
Jun 11, 2026
Merged

refactor(spec-decode): simplify Llama Eagle3 attention path for #217 (1/3)#390
LorrinWWW merged 12 commits into
lightseekorg:mainfrom
rjzhb:refactor/llama-attention-hooks

Conversation

@rjzhb

@rjzhb rjzhb commented Jun 9, 2026

Copy link
Copy Markdown
Collaborator

Summary

First PR in a series refactoring the spec-decode attention path introduced in #217.

  • Lift fused KV pre-write into base llama.LlamaAttention._attn (= pure perf win for target Llama when backend supports it).
  • Eagle3 LlamaAttention collapses the 60-line monolith forward + 4 _maybe_* helpers + 80-line __init__ copy into a 2-method subclass: _attn overrides the draft-first-step dispatch B (correction + q-slice + DECODE / post-slice), inactive steps delegate to super()._attn.
  • Correction (= spec_num_tokens - accept_lengths trim of draft_seq_lens_buf) moves from the drafter and the standalone _maybe_trim_cache_seqlens helper into a single _apply_correction method, called atomically alongside the q-slice in the active branch. Drafter plumbs accept_lengths / draft_seq_lens_buf via ForwardContext.
  • Drops the temporary pre_attention_trim class flag + drafter-side isinstance gate now that the trim lives next to its only consumer.

ctx.draft_first_step_reduce stays for now (= Qwen NextN / DeepSeek V3 NextN / comm_manager still read it); follow-up PRs in the series migrate those families and then delete the field.

no behaviour change for target Llama or Eagle3 draft. Builds on #217.

@rjzhb rjzhb changed the title refactor(spec-decode): wrap Llama Eagle3 attention via base _attn hook (1/3) refactor(spec-decode): simplify Llama Eagle3 attention path (1/3) Jun 9, 2026
@rjzhb rjzhb changed the title refactor(spec-decode): simplify Llama Eagle3 attention path (1/3) refactor(spec-decode): simplify Llama Eagle3 attention path for #217 (1/3) Jun 9, 2026
@rjzhb
rjzhb force-pushed the refactor/llama-attention-hooks branch 2 times, most recently from a1f27ec to 04e406f Compare June 9, 2026 01:11
@rjzhb rjzhb changed the title refactor(spec-decode): simplify Llama Eagle3 attention path for #217 (1/3) [WIP] refactor(spec-decode): simplify Llama Eagle3 attention path for #217 (1/3) Jun 9, 2026
rjzhb added 4 commits June 9, 2026 04:20
Signed-off-by: rjzhb <rjzhb222@163.com>
Signed-off-by: rjzhb <rjzhb222@163.com>
@rjzhb
rjzhb force-pushed the refactor/llama-attention-hooks branch from f385dc5 to 90ec73a Compare June 9, 2026 04:20
Comment thread python/tokenspeed/runtime/execution/drafter/eagle.py Outdated
Comment thread python/tokenspeed/runtime/execution/drafter/eagle.py
@rjzhb
rjzhb changed the base branch from main to refactor/draft-trim June 9, 2026 19:49
@rjzhb
rjzhb changed the base branch from refactor/draft-trim to main June 9, 2026 20:19
@rjzhb
rjzhb force-pushed the refactor/llama-attention-hooks branch from 3e03507 to cc698bf Compare June 9, 2026 22:44
@rjzhb
rjzhb marked this pull request as ready for review June 10, 2026 02:06
@rjzhb
rjzhb requested a review from a team as a code owner June 10, 2026 02:06

@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: bad01b8ffe

ℹ️ 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.py Outdated
@rjzhb rjzhb changed the title [WIP] refactor(spec-decode): simplify Llama Eagle3 attention path for #217 (1/3) refactor(spec-decode): simplify Llama Eagle3 attention path for #217 (1/3) Jun 10, 2026

@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: e06b280564

ℹ️ 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

@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.

temporal if-condition will be removed when 3/3 refactoring PRs are all in.

@LorrinWWW
LorrinWWW merged commit a7d4137 into lightseekorg:main Jun 11, 2026
33 of 36 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