Add max_draft_tokens more tokens to kv cache seq_len for memory

estimation

Signed-off-by: Hui Gao <[email protected]>

Author: HuiGao-NV

tensorrt_llm/_torch/pyexecutor/_util.py

@@ -153,10 +153,12 @@ def _get_token_num_for_estimation(self) -> int:
         num_cache_blocks = 0
         num_extra_tokens_per_seq = 1  # account for generated tokens
         pytorch_backend_config = executor_config.pytorch_backend_config
+        spec_cfg = executor_config.speculative_config
         if not pytorch_backend_config.disable_overlap_scheduler:
             num_extra_tokens_per_seq = num_extra_tokens_per_seq + 1
+            if spec_cfg is not None:
+                num_extra_tokens_per_seq += spec_cfg.max_draft_tokens
 
-        spec_cfg = executor_config.speculative_config
         if spec_cfg is not None:
             num_extra_tokens_per_seq += spec_cfg.max_draft_tokens
             num_extra_tokens_per_seq += spec_cfg.num_extra_kv_tokens