Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[Neuron] Enable async output processing for neuron #9553

Draft
wants to merge 1 commit into
base: main
Choose a base branch
from
Draft

[Neuron] Enable async output processing for neuron #9553

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 3 additions & 3 deletions vllm/config.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -416,10 +416,10 @@ def verify_async_output_proc(self, parallel_config, speculative_config,

# Reminder: Please update docs/source/serving/compatibility_matrix.rst
# If the feature combo become valid
if device_config.device_type not in ("cuda", "tpu", "xpu"):
if device_config.device_type not in ("cuda", "tpu", "xpu", "neuron"):
logger.warning(
"Async output processing is only supported for CUDA, TPU, XPU. "
"Disabling it for other platforms.")
"Async output processing is only supported for CUDA, TPU, XPU,"
" Neuron. Disabling it for other platforms.")
self.use_async_output_proc = False
return

Expand Down
7 changes: 6 additions & 1 deletion vllm/worker/neuron_model_runner.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,7 +1,8 @@
import os
from dataclasses import dataclass
from importlib.util import find_spec
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union
from typing import (TYPE_CHECKING, Any, Callable, Dict, List, Optional, Tuple,
Union)

import torch
from torch import nn
Expand Down Expand Up @@ -35,6 +36,7 @@ class ModelInputForNeuron(ModelRunnerInputBase):
input_block_ids: Optional[torch.Tensor] = None
sampling_metadata: Optional["SamplingMetadata"] = None
multi_modal_kwargs: Optional[BatchedTensorInputs] = None
async_callback: Optional[Callable] = None

def as_broadcastable_tensor_dict(
self) -> Dict[str, Union[int, torch.Tensor]]:
Expand Down Expand Up @@ -334,6 +336,9 @@ def execute_model(
else:
logits = hidden_states

if model_input.async_callback is not None:
model_input.async_callback()

# Sample the next token.
output = self.model.sample(
logits=logits,
Expand Down