update execute_model

Author: yuhao-zh

src/parallax/vllm/model_runner.py

@@ -150,6 +150,34 @@ def __init__(
             f"ParallaxVLLMModelRunner initialized: layers [{start_layer}, {end_layer}), "
             f"is_first={self.is_first_peer}, is_last={self.is_last_peer}"
         )
+    
+    def execute_model(
+        self,
+        scheduler_output: Any,
+        intermediate_tensors: Optional[Any] = None,
+    ) -> Any:
+        """
+        Execute the model with proper intermediate tensors handling for pipeline parallelism.
+        
+        For vLLM v1, intermediate_tensors must be set as an instance variable
+        AND passed as a parameter for proper pipeline parallelism support.
+        """
+        # Set intermediate_tensors as instance variable if provided
+        # This is needed for sync_and_slice_intermediate_tensors to work
+        if intermediate_tensors is not None:
+            self.intermediate_tensors = intermediate_tensors
+            logger.debug(
+                f"Set intermediate_tensors on model_runner for PP "
+                f"(is_first_peer={self.is_first_peer}, "
+                f"hidden_states shape={intermediate_tensors.tensors.get('hidden_states').shape if hasattr(intermediate_tensors, 'tensors') else 'N/A'})"
+            )
+        
+        # Call parent execute_model with intermediate_tensors parameter
+        # Both the instance variable and parameter are needed for vLLM v1
+        return super().execute_model(
+            scheduler_output=scheduler_output,
+            intermediate_tensors=intermediate_tensors,
+        )
 
     def _create_kv_cache_config(self, kv_cache_memory_fraction: float = None) -> KVCacheConfig:
         logger.debug("Generating KV cache configuration from model...")