[None][feat] Make 2-model spec dec use the 1-model kernels (Hopper)

Signed-off-by: Mike Iovine <6158008+mikeiovine@users.noreply.github.com>

Author: mikeiovine

tensorrt_llm/_torch/speculative/interface.py

@@ -136,17 +136,11 @@ def extend_ctx(self, attention_backend: Type[AttentionBackend]):
             # 1-model has separate logic for handling draft tokens
             return False
 
-        if issubclass(attention_backend,
-                      TrtllmAttention) and self.is_mtp_eagle():
-            # TRTLLM MLA does not work with the chunked context mode.
-            return False
-
-        return not issubclass(attention_backend,
-                              TrtllmAttention) or get_sm_version() != 100
+        return not issubclass(attention_backend, TrtllmAttention)
 
     def attention_need_spec_dec_mode(
         self,
-        spec_resource_manager: BaseResourceManager,
+        spec_resource_manager: Optional[BaseResourceManager],
         is_draft_model: bool,
         attention_backend: Type[AttentionBackend],
         use_chain_drafter: bool,  # CDL
@@ -164,9 +158,9 @@ def attention_need_spec_dec_mode(
         is_trtllm_attention = issubclass(attention_backend, TrtllmAttention)
         # Case 1: one model
         use_case_1 = self.is_eagle3_one_model()
-        # Case 2: eagle3 two model + draft model + CDL + is_first_draft + TRTLLM attention
+        # Case 2: eagle3 two model + is_first_draft + TRTLLM attention
         use_case_2 = self.is_eagle3(
-        ) and spec_resource_manager.is_first_draft and use_chain_drafter and is_draft_model and is_trtllm_attention
+        ) and spec_resource_manager.is_first_draft and is_trtllm_attention
         # Case 3: eagle3 two model + tree decoding + draft model + CDL + TRTLLM attention
         use_case_3 = self.is_eagle3(
         ) and is_spec_dec_tree and is_draft_model and use_chain_drafter and is_trtllm_attention

tests/unittest/_torch/speculative/test_eagle3.py

@@ -206,7 +206,7 @@ def test_llama_eagle3(use_cuda_graph: bool, attn_backend: str,
             num_tokens = len(new_tokens)
 
         accept_rate = num_accepted / num_drafted
-        assert accept_rate > 0.15
+        assert accept_rate > 0.10
 
     # Output tests
     sampling_params = SamplingParams(max_tokens=10, temperature=0)
@@ -252,7 +252,7 @@ def test_llama_eagle3_long_prompt(use_cuda_graph):
                    speculative_config=spec_config,
                    max_batch_size=1,
                    cuda_graph_config=cuda_graph_config,
-                   disable_overlap_scheduler=False)
+                   disable_overlap_scheduler=True)
 
     prompt = [", ".join(str(i) for i in range(1000))]