Skip to content

Commit 4141608

Browse files
authored
[Hardware][intel GPU] add async output process for xpu (#8897)
1 parent dfe43a2 commit 4141608

File tree

2 files changed

+8
-4
lines changed

2 files changed

+8
-4
lines changed

vllm/config.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -361,9 +361,9 @@ def verify_async_output_proc(self, parallel_config, speculative_config,
361361

362362
# Reminder: Please update docs/source/serving/compatibility_matrix.rst
363363
# If the feature combo become valid
364-
if device_config.device_type not in ("cuda", "tpu"):
364+
if device_config.device_type not in ("cuda", "tpu", "xpu"):
365365
logger.warning(
366-
"Async output processing is only supported for CUDA or TPU. "
366+
"Async output processing is only supported for CUDA, TPU, XPU. "
367367
"Disabling it for other platforms.")
368368
self.use_async_output_proc = False
369369
return

vllm/worker/xpu_model_runner.py

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -2,8 +2,8 @@
22
import time
33
import weakref
44
from dataclasses import dataclass
5-
from typing import (TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Type,
6-
TypeVar)
5+
from typing import (TYPE_CHECKING, Any, Callable, Dict, List, Optional, Tuple,
6+
Type, TypeVar)
77

88
import torch
99
import torch.nn as nn
@@ -57,6 +57,7 @@ class ModelInputForXPU(ModelRunnerInputBase):
5757
virtual_engine: Optional[int] = None
5858
seq_lens: Optional[List[int]] = None
5959
query_lens: Optional[List[int]] = None
60+
async_callback: Optional[Callable] = None
6061

6162
def as_broadcastable_tensor_dict(self) -> Dict[str, Any]:
6263
tensor_dict = {
@@ -582,6 +583,9 @@ def execute_model(
582583
if not self.is_driver_worker:
583584
return []
584585

586+
if model_input.async_callback is not None:
587+
model_input.async_callback()
588+
585589
# Sample the next token.
586590
output: SamplerOutput = self.model.sample(
587591
logits=logits,

0 commit comments

Comments
 (0)