debug and fix error

vllm/two_batch_overlap/v1/model_input_split_v1.py

@@ -306,12 +306,16 @@ def tbo_split_and_execute_model(
     intermediate_tensors: Optional[IntermediateTensors] = None,
 ) -> Union[ModelRunnerOutput, IntermediateTensors]:
     use_tbo = False
-    if len(scheduler_output.num_scheduled_tokens) > 1:
-        split_scheduler_output(runner, scheduler_output)
-        if input_split.scheduler_output_left.total_num_scheduled_tokens >= envs.VLLM_TBO_MIN_TOKENS and \
-            input_split.scheduler_output_right.total_num_scheduled_tokens >= envs.VLLM_TBO_MIN_TOKENS:
-            use_tbo = True
 
+    if isinstance(runner.attn_metadata_builders[0], MLACommonMetadataBuilder) and \
+        runner.attn_metadata_builders[0]._num_decodes > 0: #is mla decode
+        use_tbo = False
+    else:
+        if len(scheduler_output.num_scheduled_tokens) > 1:
+            split_scheduler_output(runner, scheduler_output)
+            if input_split.scheduler_output_left.total_num_scheduled_tokens >= envs.VLLM_TBO_MIN_TOKENS and \
+                input_split.scheduler_output_right.total_num_scheduled_tokens >= envs.VLLM_TBO_MIN_TOKENS:
+                use_tbo = True
     if use_tbo:
         num_input_tokens_left = input_split.scheduler_output_left.total_num_scheduled_tokens
         num_input_tokens_right = num_input_tokens - num_input_tokens_left