refactor(spec-decode): simplify Llama Eagle3 attention path for #217 (1/3)#390
Conversation
a1f27ec to
04e406f
Compare
Signed-off-by: rjzhb <rjzhb222@163.com>
Signed-off-by: rjzhb <rjzhb222@163.com>
f385dc5 to
90ec73a
Compare
Signed-off-by: rjzhb <rjzhb222@163.com>
Signed-off-by: rjzhb <rjzhb222@163.com>
3e03507 to
cc698bf
Compare
There was a problem hiding this comment.
💡 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".
Signed-off-by: rjzhb <rjzhb222@163.com>
There was a problem hiding this comment.
💡 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".
Signed-off-by: rjzhb <rjzhb222@163.com>
LorrinWWW
left a comment
There was a problem hiding this comment.
temporal if-condition will be removed when 3/3 refactoring PRs are all in.
Summary
First PR in a series refactoring the spec-decode attention path introduced in #217.
llama.LlamaAttention._attn(= pure perf win for target Llama when backend supports it).LlamaAttentioncollapses the 60-line monolith forward + 4_maybe_*helpers + 80-line__init__copy into a 2-method subclass:_attnoverrides the draft-first-step dispatch B (correction + q-slice + DECODE / post-slice), inactive steps delegate tosuper()._attn.spec_num_tokens - accept_lengthstrim ofdraft_seq_lens_buf) moves from the drafter and the standalone_maybe_trim_cache_seqlenshelper into a single_apply_correctionmethod, called atomically alongside the q-slice in the active branch. Drafter plumbsaccept_lengths/draft_seq_lens_bufviaForwardContext.pre_attention_trimclass flag + drafter-side isinstance gate now that the trim lives next to its only consumer.ctx.draft_first_step_reducestays for now (= Qwen NextN / DeepSeek V3 NextN /comm_managerstill 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.