A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 15, 22:15 UTC
[misc] Clean up ray compiled graph type hints (#13731)
Rui Qiao committed
Feb 25, 2025 at 00:37 UTC
3173c3b34ef8245164c63be47a1207c4db1369c3
2 files changed
+17
-6
vllm/executor/ray_distributed_executor.py
+12
-4
@@ -528,10 +528,18 @@ class RayDistributedExecutor(DistributedExecutorBase):
528
envs.VLLM_USE_RAY_COMPILED_DAG_OVERLAP_COMM)
529
with InputNode() as input_data:
530
# Example DAG: PP=2, TP=4
531
- # (ExecuteModelReq, None) -> 0 -> (ExecuteModelReq, IntermediateOutput) -> 4 -> SamplerOutput # noqa: E501
532
- # -> 1 -> (ExecuteModelReq, IntermediateOutput) -> 5 -> SamplerOutput # noqa: E501
533
- # -> 2 -> (ExecuteModelReq, IntermediateOutput) -> 6 -> SamplerOutput # noqa: E501
534
- # -> 3 -> (ExecuteModelReq, IntermediateOutput) -> 7 -> SamplerOutput # noqa: E501
531
+ #
532
+ # For V0:
533
+ # ExecuteModelRequest -> 0 -> (ExecuteModelReq, IntermediateTensors) -> 4 -> SamplerOutput # noqa: E501
534
+ # ExecuteModelRequest -> 1 -> (ExecuteModelReq, IntermediateTensors) -> 5 -> SamplerOutput # noqa: E501
535
+ # ExecuteModelRequest -> 2 -> (ExecuteModelReq, IntermediateTensors) -> 6 -> SamplerOutput # noqa: E501
536
+ # ExecuteModelRequest -> 3 -> (ExecuteModelReq, IntermediateTensors) -> 7 -> SamplerOutput # noqa: E501
537
+ #
538
+ # For V1:
539
+ # SchedulerOutput -> 0 -> (SchedulerOutput, IntermediateTensors) -> 4 -> ModelRunnerOutput # noqa: E501
540
+ # SchedulerOutput -> 1 -> (SchedulerOutput, IntermediateTensors) -> 5 -> ModelRunnerOutput # noqa: E501
541
+ # SchedulerOutput -> 2 -> (SchedulerOutput, IntermediateTensors) -> 6 -> ModelRunnerOutput # noqa: E501
542
+ # SchedulerOutput -> 3 -> (SchedulerOutput, IntermediateTensors) -> 7 -> ModelRunnerOutput # noqa: E501
543
544
# All workers in the first TP group will take in the
545
# ExecuteModelRequest as input.
vllm/executor/ray_utils.py
+5
-2
@@ -114,8 +114,11 @@ try:
114
115
def execute_model_ray(
116
self,
117
- scheduler_output: "SchedulerOutput",
118
- ) -> "ModelRunnerOutput":
117
+ scheduler_output: Union["SchedulerOutput",
118
+ Tuple["SchedulerOutput",
119
+ "IntermediateTensors"]],
120
+ ) -> Union["ModelRunnerOutput", Tuple["SchedulerOutput",
121
+ "IntermediateTensors"]]:
122
# this method is used to compile ray CG,
123
# and it needs a special logic of self.setup_device_if_necessary()
124
self.setup_device_if_necessary()