[None][feat] Add simple optimizations for MTP 2-model (#9176)

Signed-off-by: Mike Iovine <[email protected]>

Author: mikeiovine

tensorrt_llm/_torch/speculative/drafting_loops.py

@@ -59,6 +59,7 @@ def save_metadata_state(attn_metadata: AttentionMetadata,
             spec_metadata.eagle3_resource_manager.is_first_draft = True
 
 
+@torch.compile(options={'max-autotune': True})
 def prepare_for_generation(attn_metadata: AttentionMetadata,
                            spec_metadata: SpecMetadata,
                            position_ids: torch.Tensor) -> torch.Tensor:

tensorrt_llm/_torch/speculative/interface.py

@@ -67,10 +67,6 @@ def needs_kv_cache_rewind(self):
         ) or self.is_ngram()
 
     def support_overlap_scheduler(self):
-        # TODO: fix accuracy issue
-        if self.is_mtp_eagle():
-            return False
-
         return self.is_mtp_one_model() or self.is_eagle3_one_model(
         ) or self.has_draft_model()