Fix CI

Signed-off-by: Zheyu Fu <zheyuf@NVIDIA.com>

Author: zheyuf

tensorrt_llm/_torch/pyexecutor/model_engine.py

@@ -874,7 +874,7 @@ def _get_graphs_to_capture(
             return graphs
 
         # Case 3: Target model (two-model) or one-model without dynamic draft
-        draft_lengths = [self.max_total_draft_tokens]
+        draft_lengths = [self.max_draft_len]
         should_capture_no_spec = (
             self.max_total_draft_tokens > 0
             and not self.spec_config.spec_dec_mode.use_one_engine()

tests/unittest/_torch/speculative/test_mtp.py

@@ -312,6 +312,7 @@ def test_sample_and_accept_draft_tokens(self, test_case_name,
                                         max_total_draft_tokens=mtp_num_modules,
                                         mtp_num_modules=mtp_num_modules)
         spec_metadata.draft_tokens = draft_tokens
+        spec_metadata.runtime_draft_len = mtp_num_modules
 
         # mtp worker
         mtpworker = MTPWorker(spec_config)
@@ -901,6 +902,7 @@ def test_mtp_update_mtp_hidden_states(
 
         spec_metadata.mtp_hidden_states_manager.mtp_past_hidden_states_pool = mtp_hidden_states_tensor_pool
         spec_metadata.mtp_hidden_states_manager.mtp_past_tokens_pool = mtp_tokens_tensor_pool
+        spec_metadata.runtime_draft_len = num_nextn_predict_layers
         spec_metadata.prepare()
 
         mtpworker = MTPWorker(spec_config)
@@ -1397,6 +1399,7 @@ def test_prepare_drafter_inputs(
 
         spec_metadata.mtp_hidden_states_manager.mtp_past_hidden_states_pool = mtp_hidden_states_tensor_pool
         spec_metadata.mtp_hidden_states_manager.mtp_past_tokens_pool = mtp_tokens_tensor_pool
+        spec_metadata.runtime_draft_len = num_nextn_predict_layers
         spec_metadata.prepare()
 
         mtpworker = MTPWorker(spec_config)