[None][feat] Make the should_use_spec_decode logic a bit smarter

tensorrt_llm/_torch/pyexecutor/py_executor.py

@@ -921,7 +921,9 @@ def _prepare_and_schedule_batch(self):
 
         if self.drafter is not None:
             self.use_spec_decode = self.drafter.should_use_spec_decode(
-                self.active_requests)
+                self.active_requests, self.max_batch_size,
+                self.model_engine.max_num_tokens,
+                self.model_engine.max_draft_len)
             self.model_engine.enable_spec_decode = self.use_spec_decode
             # If speculation is off, this function sets py_draft_tokens to None
             # for all active requests. If it's on, we initialize py_draft_tokens

tensorrt_llm/_torch/speculative/drafter.py

@@ -27,12 +27,18 @@ def prepare_draft_tokens(
         raise NotImplementedError
 
     @final
-    def should_use_spec_decode(self, requests: List[LlmRequest]) -> bool:
+    def should_use_spec_decode(self, requests: List[LlmRequest],
+                               max_batch_size: int, max_num_tokens: int,
+                               max_draft_len: int) -> bool:
         """
         You probably don't want to override this. ModelEngine
         assumes that speculation is always on if max_concurrency
         is not specified by the user's spec config.
         """
-        if self.max_concurrency is not None:
-            return len(requests) <= self.max_concurrency
-        return True
+        if self.max_concurrency is None:
+            return True
+
+        tokens_per_request = 1 + max_draft_len
+        token_cap = max_num_tokens // tokens_per_request
+        num_effective_requests = min(max_batch_size, len(requests), token_cap)
+        return num_effective_requests <= self.max_concurrency

tests/unittest/_torch/speculative/test_dynamic_spec_decode.py

@@ -52,7 +52,8 @@ def test_dynamic_spec_decode():
     )
 
     # Mock should_use_spec_decode to return True for first two calls, then False
-    def mock_should_use_spec_decode(self, requests):
+    def mock_should_use_spec_decode(self, requests, max_batch_size,
+                                    max_num_tokens, draft_len):
         if not hasattr(mock_should_use_spec_decode, 'call_count'):
             mock_should_use_spec_decode.call_count = 0
         mock_should_use_spec_decode.call_count += 1