Add test cases for should_use_spec_decode.

Signed-off-by: Zheyu Fu <[email protected]>

Author: zheyuf

tensorrt_llm/_torch/pyexecutor/py_executor.py

@@ -923,7 +923,7 @@ def _prepare_and_schedule_batch(self):
             self.use_spec_decode = self.drafter.should_use_spec_decode(
                 self.active_requests, self.max_batch_size,
                 self.model_engine.max_num_tokens,
-                self.model_engine.max_draft_len)
+                self.model_engine.spec_config.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

@@ -35,10 +35,13 @@ def should_use_spec_decode(self, requests: List[LlmRequest],
         assumes that speculation is always on if max_concurrency
         is not specified by the user's spec config.
         """
+
+        # Inputs validated upstream: max_batch_size>0, max_num_tokens>0, max_draft_len>0
+
         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
+        token_cap = max_num_tokens // (1 + max_draft_len)
+        num_effective_requests = min(len(requests), max_batch_size, token_cap)
+
+        return 0 < num_effective_requests <= self.max_concurrency

tests/unittest/_torch/speculative/test_dynamic_spec_decode.py

@@ -88,5 +88,62 @@ def mock_should_use_spec_decode(self, requests, max_batch_size,
         assert text_spec == text_ref
 
 
+def test_should_use_spec_decode():
+    from tensorrt_llm._torch.speculative.drafter import Drafter
+
+    class _DummyDrafter(Drafter):
+
+        def prepare_draft_tokens(self,
+                                 scheduled_requests,
+                                 resource_manager=None) -> None:
+            return
+
+    drafter = _DummyDrafter(max_concurrency=6)
+
+    # Compare min(len(requests), max_batch_size, token_cap) with max_concurrency
+
+    # Small active_requests ON case: num_effective_requests = min(5, 8, very_large) = 5 <= 6 → True
+    active_requests = [object()] * 5
+    assert drafter.should_use_spec_decode(active_requests,
+                                          max_batch_size=8,
+                                          max_num_tokens=4096 * 8,
+                                          max_draft_len=4) is True
+
+    # Small batch size ON case: num_effective_requests = min(12, 5, very_large) = 5 <= 6 → True
+    active_requests = [object()] * 12
+    assert drafter.should_use_spec_decode(active_requests,
+                                          max_batch_size=5,
+                                          max_num_tokens=4096 * 8,
+                                          max_draft_len=4) is True
+
+    # Small token budget ON case: token_cap = 28 // (1+4) = 5 → min(8, 12, 5) = 5 <= 6 → True
+    active_requests = [object()] * 12
+    assert drafter.should_use_spec_decode(active_requests,
+                                          max_batch_size=8,
+                                          max_num_tokens=28,
+                                          max_draft_len=4) is True
+
+    # Generic OFF case: num_effective_requests = min(12, 8, very_large) = 8 > 6 → False
+    active_requests = [object()] * 12
+    assert drafter.should_use_spec_decode(active_requests,
+                                          max_batch_size=8,
+                                          max_num_tokens=4096 * 8,
+                                          max_draft_len=4) is False
+
+    # Edge case - None active requests OFF case: num_effective_requests = min(0, 8, very_large) = 0 <= 6 → False
+    active_requests = []
+    assert drafter.should_use_spec_decode(active_requests,
+                                          max_batch_size=8,
+                                          max_num_tokens=4096 * 8,
+                                          max_draft_len=4) is False
+
+    # Edge case - Token cap equals 0 OFF case: token_cap = 4 // (1+4) = 0 → min(12, 8, 0) = 0 <= 6 → False
+    active_requests = [object()] * 12
+    assert drafter.should_use_spec_decode(active_requests,
+                                          max_batch_size=8,
+                                          max_num_tokens=4,
+                                          max_draft_len=4) is False
+
+
 if __name__ == "__main__":
     unittest.main()