[https://nvbugs/5590408][fix] Fallback to greedy sampling in two-model overlap scheduler (#9321)

Signed-off-by: ziyixiong-nv <[email protected]>

Author: ziyixiong-nv

tensorrt_llm/_torch/pyexecutor/py_executor_creator.py

@@ -414,6 +414,11 @@ def drafting_loop_wrapper(model):
         model_engine_max_seq_len += get_num_extra_kv_tokens(spec_config)
         model_engine_max_seq_len += spec_config.max_total_draft_tokens
 
+    if has_draft_model_engine and not llm_args.disable_overlap_scheduler:
+        logger.warning(
+            "Overlap scheduler is enabled for two-model speculative decoding. Rejection sampling will fallback to greedy sampling."
+        )
+
     max_seq_len = model_engine_max_seq_len
     max_num_tokens = model_engine.max_num_tokens
     sparse_attention_config = model_engine.sparse_attention_config

tensorrt_llm/_torch/speculative/model_drafter.py

@@ -625,7 +625,10 @@ def process_static_draft_outputs(self, outputs: dict[str, torch.Tensor]
                 target_model_req.py_draft_tokens.append(
                     draft_tokens_host[token_idx][req_idx])
                 py_draft_logits.append(draft_logits[token_idx][req_idx])
-            target_model_req.py_draft_logits = torch.stack(py_draft_logits)
+
+            # The overlap scheduler doesn't support rejection sampling yet, so we don't update the py_draft_logits to get it fallback to greedy sampling.
+            if self.disable_overlap_scheduler:
+                target_model_req.py_draft_logits = torch.stack(py_draft_logits)
 
     def process_dynamic_draft_outputs(
             self,