feat: efficient eagle3 with cross attn and flex attn

Signed-off-by: h-guo18 <[email protected]>

Author: h-guo18

examples/speculative_decoding/main.py

@@ -203,6 +203,8 @@ def train():
                     "draft_vocab_size": custom_config["draft_vocab_size"]
                     if eagle_args.eagle_config and "draft_vocab_size" in custom_config
                     else model.config.vocab_size,
+                    # pass in the seq length for flex attention mask compilation
+                    "training_seq_len": training_args.training_seq_len,
                 }
             )