tbo debug fix error

vllm/forward_context.py

@@ -145,7 +145,8 @@ def set_forward_context(attn_metadata: Any,
         virtual_engine=virtual_engine,
         attn_metadata=attn_metadata,
         dp_metadata=dp_metadata)
-
+    if envs.VLLM_ENABLE_TBO:
+        set_tbo_forward_context(_forward_context)
     try:
         yield
     finally:

vllm/two_batch_overlap/two_batch_overlap.py

@@ -270,6 +270,7 @@ def tbo_model_executable(
                                     vllm_config, virtual_engine):
             hidden_or_intermediate_states = model_executable(
                 input_ids=model_input.input_tokens,
+                inputs_embeds=model_input.inputs_embeds,
                 positions=model_input.input_positions,
                 intermediate_tensors=intermediate_tensors,
                 **MultiModalKwargs.as_kwargs(multi_modal_kwargs,