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

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

Author: mikeiovine

tensorrt_llm/_torch/speculative/interface.py

@@ -5,7 +5,6 @@
 
 import torch
 
-from ..._utils import get_sm_version
 from ..attention_backend.trtllm import AttentionBackend, TrtllmAttention
 from ..pyexecutor.resource_manager import BaseResourceManager
 
@@ -117,17 +116,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,
@@ -137,9 +130,10 @@ def attention_need_spec_dec_mode(
         If true, the attention backend kernel needs to run in spec-dec mode (multi-token query mode).
         """
         is_trtllm_attention = issubclass(attention_backend, TrtllmAttention)
-        return self.is_eagle3_one_model() or (
+
+        return self.is_eagle3_one_model() or not is_draft_model or (
             self.is_eagle3() and spec_resource_manager.is_first_draft
-            and is_trtllm_attention and use_chain_drafter and is_draft_model)
+            and is_trtllm_attention)
 
     @staticmethod
     def from_string(name: Optional[str]) -> "SpeculativeDecodingMode":

tests/unittest/_torch/speculative/test_eagle3.py

@@ -136,7 +136,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)
@@ -182,7 +182,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))]