Skip to content

Commit 8938a09

Browse files
committed
feat: efficient eagle3 with cross attn and flex attn
Signed-off-by: h-guo18 <[email protected]>
1 parent be95a10 commit 8938a09

File tree

2 files changed

+88
-307
lines changed

2 files changed

+88
-307
lines changed

examples/speculative_decoding/main.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -182,6 +182,8 @@ def train():
182182
"draft_vocab_size": custom_config["draft_vocab_size"]
183183
if eagle_args.eagle_config and "draft_vocab_size" in custom_config
184184
else model.config.vocab_size,
185+
# pass in the seq length for flex attention mask compilation
186+
"traing_seq_len": training_args.training_seq_len,
185187
}
186188
)
187189

0 commit comments

Comments
 (0)