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.

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

[Core] Refactor Worker and ModelRunner to consolidate control plane communication #5408

Merged
merged 64 commits into from
Jun 26, 2024
Merged
Show file tree
Hide file tree
Changes from 56 commits
Commits
Show all changes
64 commits
Select commit Hold shift + click to select a range
725b0b2
tmp
stephanie-wang Jun 11, 2024
b74eb10
fix
stephanie-wang Jun 11, 2024
38b0ddf
ray and mp backends work
stephanie-wang Jun 11, 2024
0d11e92
embedding model runner works
stephanie-wang Jun 11, 2024
2cdc218
GPU executor works
stephanie-wang Jun 11, 2024
c728512
remove comment
stephanie-wang Jun 11, 2024
2bf752b
use the right ModelInput class
stephanie-wang Jun 11, 2024
f35a23f
CPU worker
stephanie-wang Jun 11, 2024
11133fe
remove commented
stephanie-wang Jun 11, 2024
174bdb1
lint
stephanie-wang Jun 11, 2024
c0e98ca
Worker.execute_model vs execute_model_local
stephanie-wang Jun 11, 2024
dccec95
lint
stephanie-wang Jun 11, 2024
dad94ba
neuron model runner
stephanie-wang Jun 11, 2024
fca606e
disallow distributed comms
stephanie-wang Jun 11, 2024
6ed3c2a
disable communication
stephanie-wang Jun 12, 2024
1803e33
Update worker.py
stephanie-wang Jun 12, 2024
dde799e
fix tests
stephanie-wang Jun 12, 2024
0398631
update
stephanie-wang Jun 12, 2024
5c41cc6
Merge branch 'control-refactor-2' of github.com:stephanie-wang/vllm i…
stephanie-wang Jun 12, 2024
72f0383
Merge remote-tracking branch 'upstream/main' into control-refactor-2
stephanie-wang Jun 12, 2024
eef6623
merge
stephanie-wang Jun 12, 2024
3004ceb
update
Jun 13, 2024
8d852e9
Merge remote-tracking branch 'upstream/main' into control-refactor-2
stephanie-wang Jun 13, 2024
9380ed8
fix
stephanie-wang Jun 13, 2024
3c4de6d
fix
stephanie-wang Jun 13, 2024
5053f30
fix
stephanie-wang Jun 13, 2024
db38556
x
stephanie-wang Jun 14, 2024
456185d
rm
stephanie-wang Jun 14, 2024
e860652
lint
stephanie-wang Jun 14, 2024
3d4f242
add missing
stephanie-wang Jun 14, 2024
11304cb
revert
stephanie-wang Jun 14, 2024
99f532e
refactor
stephanie-wang Jun 15, 2024
797a7cf
doc
stephanie-wang Jun 15, 2024
6ad2513
revert spec decode and doc
stephanie-wang Jun 15, 2024
97ec303
Merge remote-tracking branch 'upstream/main' into control-refactor-2
stephanie-wang Jun 15, 2024
e10bace
typing
stephanie-wang Jun 15, 2024
ce087ae
fix
stephanie-wang Jun 18, 2024
f851b00
Merge remote-tracking branch 'upstream/main' into control-refactor-2
stephanie-wang Jun 18, 2024
0e2acc4
XPU worker and rename
stephanie-wang Jun 18, 2024
d318ec8
lint
stephanie-wang Jun 18, 2024
b48f783
lint
stephanie-wang Jun 18, 2024
c93afc1
Merge remote-tracking branch 'upstream/main' into control-refactor-2
stephanie-wang Jun 18, 2024
30ac400
fix
stephanie-wang Jun 18, 2024
01688d5
x
stephanie-wang Jun 18, 2024
7dbb646
fix
stephanie-wang Jun 18, 2024
d2e4c41
fix
stephanie-wang Jun 19, 2024
3e46253
lint
stephanie-wang Jun 19, 2024
0a2890a
Merge remote-tracking branch 'upstream/main' into control-refactor-2
stephanie-wang Jun 21, 2024
36dfce1
merge
stephanie-wang Jun 21, 2024
ea5412e
Merge remote-tracking branch 'upstream/main' into control-refactor-2
stephanie-wang Jun 21, 2024
dc2f103
x
stephanie-wang Jun 21, 2024
fbf074d
x
stephanie-wang Jun 22, 2024
660a8d5
rename ModelInput -> ModelInputBase, override as_broadcastable_tensor…
stephanie-wang Jun 23, 2024
8cca634
fixes
stephanie-wang Jun 23, 2024
0a25c19
rename
stephanie-wang Jun 23, 2024
0b26877
fix
stephanie-wang Jun 24, 2024
e7052d5
do not filter Nones
stephanie-wang Jun 24, 2024
df5551f
dupe
stephanie-wang Jun 24, 2024
6745b3b
update
stephanie-wang Jun 25, 2024
ebae970
lint
stephanie-wang Jun 25, 2024
5763621
revert
stephanie-wang Jun 25, 2024
46d5b18
Merge branch 'main' into control-refactor-2
stephanie-wang Jun 25, 2024
d16d5fe
rm
stephanie-wang Jun 25, 2024
f6c6234
fix
stephanie-wang Jun 25, 2024
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
150 changes: 150 additions & 0 deletions tests/worker/test_model_input.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,150 @@
import dataclasses
from typing import List, Tuple, Type

import torch

from vllm.attention import AttentionMetadata
from vllm.attention.backends.abstract import AttentionBackend
from vllm.model_executor import SamplingMetadata
from vllm.model_executor.pooling_metadata import PoolingMetadata
from vllm.worker.embedding_model_runner import (
ModelInputForGPUWithPoolingMetadata)
from vllm.worker.model_runner import ModelInputForGPUWithSamplingMetadata


class MockAttentionBackend(AttentionBackend):

@staticmethod
def get_name() -> str:
raise NotImplementedError

@staticmethod
def get_impl_cls():
raise NotImplementedError

@staticmethod
def get_metadata_cls() -> Type["AttentionMetadata"]:
return AttentionMetadata

@staticmethod
def get_kv_cache_shape(
num_blocks: int,
block_size: int,
num_kv_heads: int,
head_size: int,
) -> Tuple[int, ...]:
raise NotImplementedError

@staticmethod
def swap_blocks(
src_kv_cache: torch.Tensor,
dst_kv_cache: torch.Tensor,
src_to_dst: torch.Tensor,
) -> None:
pass

@staticmethod
def copy_blocks(
kv_caches: List[torch.Tensor],
src_to_dists: torch.Tensor,
) -> None:
pass


def test_model_runner_input():
sampling_metadata = SamplingMetadata(
["seq_group"],
"selected_token_indices",
"categorized_sample_indices",
"num_prompts",
)
attn_metadata = AttentionMetadata(
num_prefills=1,
num_prefill_tokens=2,
num_decode_tokens=3,
slot_mapping=torch.zeros(1),
)
model_input = ModelInputForGPUWithSamplingMetadata.new(
input_tokens=torch.ones(10),
input_positions=torch.ones(10),
sampling_metadata=sampling_metadata,
attn_metadata=attn_metadata)

assert isinstance(model_input, ModelInputForGPUWithSamplingMetadata)

# Test round trip serialization.
tensor_dict = model_input.as_broadcastable_tensor_dict()
attn_backend = MockAttentionBackend()
received_model_input = ModelInputForGPUWithSamplingMetadata.new(
attn_backend=attn_backend, **tensor_dict)
# Check that received copy has correct values.
assert isinstance(received_model_input,
ModelInputForGPUWithSamplingMetadata)
assert received_model_input.input_tokens is not None
assert (
received_model_input.input_tokens == model_input.input_tokens).all()
assert received_model_input.input_positions is not None
assert (received_model_input.input_positions == model_input.input_positions
).all()
assert received_model_input.multi_modal_kwargs is None
assert (received_model_input.multi_modal_kwargs ==
model_input.multi_modal_kwargs)
assert received_model_input.lora_requests is None
assert received_model_input.lora_requests == model_input.lora_requests
assert received_model_input.lora_mapping is None
assert received_model_input.lora_mapping == model_input.lora_mapping
for field in dataclasses.fields(AttentionMetadata):
assert getattr(received_model_input.attn_metadata, field.name,
None) == getattr(attn_metadata, field.name, None)
# For sampling metadata, only selected_token_indices is copied.
assert (received_model_input.sampling_metadata.selected_token_indices ==
sampling_metadata.selected_token_indices)
assert received_model_input.sampling_metadata.seq_groups is None


def test_embedding_model_runner_input():
pooling_metadata = PoolingMetadata(
seq_groups=[[0]],
seq_data={},
prompt_lens=[1],
)
attn_metadata = AttentionMetadata(
num_prefills=1,
num_prefill_tokens=2,
num_decode_tokens=3,
slot_mapping=torch.zeros(1),
)
model_input = ModelInputForGPUWithPoolingMetadata.new(
input_tokens=torch.ones(10),
input_positions=torch.ones(10),
pooling_metadata=pooling_metadata,
attn_metadata=attn_metadata)

assert isinstance(model_input, ModelInputForGPUWithPoolingMetadata)

# Test round trip serialization.
tensor_dict = model_input.as_broadcastable_tensor_dict()
attn_backend = MockAttentionBackend()
received_model_input = ModelInputForGPUWithPoolingMetadata.new(
attn_backend=attn_backend, **tensor_dict)
# Check that received copy has correct values.
assert isinstance(received_model_input,
ModelInputForGPUWithPoolingMetadata)
assert received_model_input.input_tokens is not None
assert (
received_model_input.input_tokens == model_input.input_tokens).all()
assert received_model_input.input_positions is not None
assert (received_model_input.input_positions == model_input.input_positions
).all()
assert received_model_input.multi_modal_kwargs is None
assert (received_model_input.multi_modal_kwargs ==
model_input.multi_modal_kwargs)
assert received_model_input.lora_requests is None
assert received_model_input.lora_requests == model_input.lora_requests
assert received_model_input.lora_mapping is None
assert received_model_input.lora_mapping == model_input.lora_mapping
for field in dataclasses.fields(AttentionMetadata):
assert getattr(received_model_input.attn_metadata, field.name,
None) == getattr(attn_metadata, field.name, None)
# Pooling metadata is not broadcast.
assert received_model_input.pooling_metadata is None
57 changes: 30 additions & 27 deletions tests/worker/test_model_runner.py
Original file line number Diff line number Diff line change
Expand Up @@ -61,12 +61,13 @@ def test_prepare_prompt(batch_size):
expected_selected_token_indices.append(selected_token_start_idx +
seq_len - 1)
selected_token_start_idx += seq_len
model_input = model_runner._prepare_model_input(seq_group_metadata_list)
model_input = model_runner._prepare_model_input_tensors(
seq_group_metadata_list)
input_tokens = model_input.input_tokens
input_positions = model_input.input_positions
attn_metadata = model_input.attn_metadata
return_seq_lens = model_input.seq_lens
slot_mapping = model_input.slot_mapping
slot_mapping = attn_metadata.slot_mapping
assert return_seq_lens == seq_lens
assert len(slot_mapping) == len(input_tokens)

Expand Down Expand Up @@ -174,10 +175,11 @@ def test_prepare_decode_cuda_graph(batch_size):
assert seq_group_metadata.token_chunk_size == 1
seq_group_metadata_list.append(seq_group_metadata)

model_input = model_runner._prepare_model_input(seq_group_metadata_list)
model_input = model_runner._prepare_model_input_tensors(
seq_group_metadata_list)
input_tokens, input_positions, attn_metadata, slot_mapping = (
model_input.input_tokens, model_input.input_positions,
model_input.attn_metadata, model_input.slot_mapping)
model_input.attn_metadata, model_input.attn_metadata.slot_mapping)
assert len(slot_mapping) == len(input_tokens)

expected_bs = _get_graph_batch_size(len(seq_group_metadata_list))
Expand Down Expand Up @@ -259,32 +261,29 @@ def test_empty_seq_group():
enforce_eager=False,
)
seq_group_metadata_list: List[SequenceGroupMetadata] = []
model_input = model_runner._prepare_model_input(seq_group_metadata_list)
input_tokens, input_positions, attn_metadata, slot_mapping = (
model_input = model_runner._prepare_model_input_tensors(
seq_group_metadata_list)
input_tokens, input_positions, attn_metadata = (
model_input.input_tokens,
model_input.input_positions,
model_input.attn_metadata,
model_input.slot_mapping,
)
assert len(input_tokens) == 0
assert len(input_positions) == 0
assert input_tokens is None
assert input_positions is None
assert attn_metadata is None
assert len(slot_mapping) == 0

model_input = model_runner._prepare_model_input(seq_group_metadata_list)
(input_tokens, input_positions, attn_metadata, slot_mapping,
return_seq_lens) = (
model_input.input_tokens,
model_input.input_positions,
model_input.attn_metadata,
model_input.slot_mapping,
model_input.seq_lens,
)
assert len(input_tokens) == 0
assert len(input_positions) == 0

model_input = model_runner._prepare_model_input_tensors(
seq_group_metadata_list)
(input_tokens, input_positions, attn_metadata, return_seq_lens) = (
model_input.input_tokens,
model_input.input_positions,
model_input.attn_metadata,
model_input.seq_lens,
)
assert input_tokens is None
assert input_positions is None
assert attn_metadata is None
assert len(slot_mapping) == 0
assert len(return_seq_lens) == 0
assert return_seq_lens is None


@pytest.fixture
Expand Down Expand Up @@ -353,8 +352,12 @@ def test_hybrid_batches(batch_size, enforce_eager, distributed_init):
seq_group_metadata_list.append(seq_group_metadata)
decode_metadata_list.append(seq_group_metadata)

(input_tokens, input_positions, attn_metadata, _, _, _,
_) = model_runner.prepare_input_tensors(seq_group_metadata_list)
model_input = model_runner.prepare_model_input(seq_group_metadata_list)
(input_tokens, input_positions, attn_metadata) = (
model_input.input_tokens,
model_input.input_positions,
model_input.attn_metadata,
)

prefill_meta_actual = attn_metadata.prefill_metadata
decode_meta_actual = attn_metadata.decode_metadata
Expand All @@ -367,7 +370,7 @@ def test_hybrid_batches(batch_size, enforce_eager, distributed_init):

# Verify attn metadata is consistent. We don't need to test individual
# values here because they are tested above.
attn_metadata = model_runner._prepare_model_input(
attn_metadata = model_runner._prepare_model_input_tensors(
seq_group_metadata_list).attn_metadata

for attr_expected, attr_actual in zip(vars(attn_metadata.prefill_metadata),
Expand Down
6 changes: 5 additions & 1 deletion vllm/attention/backends/abstract.py
Original file line number Diff line number Diff line change
Expand Up @@ -21,9 +21,13 @@ def get_impl_cls() -> Type["AttentionImpl"]:

@staticmethod
@abstractmethod
def make_metadata(*args, **kwargs) -> "AttentionMetadata":
def get_metadata_cls() -> Type["AttentionMetadata"]:
raise NotImplementedError

@classmethod
def make_metadata(cls, *args, **kwargs) -> "AttentionMetadata":
return cls.get_metadata_cls()(*args, **kwargs)

@staticmethod
@abstractmethod
def get_kv_cache_shape(
Expand Down
4 changes: 2 additions & 2 deletions vllm/attention/backends/blocksparse_attn.py
Original file line number Diff line number Diff line change
Expand Up @@ -90,8 +90,8 @@ def get_impl_cls() -> Type["BlocksparseFlashAttentionImpl"]:
return BlocksparseFlashAttentionImpl

@staticmethod
def make_metadata(*args, **kwargs) -> "BlocksparseFlashAttentionMetadata":
return BlocksparseFlashAttentionMetadata(*args, **kwargs)
def get_metadata_cls() -> Type["AttentionMetadata"]:
return BlocksparseFlashAttentionMetadata

@staticmethod
def get_kv_cache_shape(
Expand Down
4 changes: 2 additions & 2 deletions vllm/attention/backends/flash_attn.py
Original file line number Diff line number Diff line change
Expand Up @@ -25,8 +25,8 @@ def get_impl_cls() -> Type["FlashAttentionImpl"]:
return FlashAttentionImpl

@staticmethod
def make_metadata(*args, **kwargs) -> "FlashAttentionMetadata":
return FlashAttentionMetadata(*args, **kwargs)
def get_metadata_cls() -> Type["AttentionMetadata"]:
return FlashAttentionMetadata

@staticmethod
def get_kv_cache_shape(
Expand Down
4 changes: 2 additions & 2 deletions vllm/attention/backends/flashinfer.py
Original file line number Diff line number Diff line change
Expand Up @@ -22,8 +22,8 @@ def get_impl_cls() -> Type["FlashInferImpl"]:
return FlashInferImpl

@staticmethod
def make_metadata(*args, **kwargs) -> "FlashInferMetadata":
return FlashInferMetadata(*args, **kwargs)
def get_metadata_cls() -> Type["AttentionMetadata"]:
return FlashInferMetadata

@staticmethod
def get_kv_cache_shape(
Expand Down
4 changes: 2 additions & 2 deletions vllm/attention/backends/rocm_flash_attn.py
Original file line number Diff line number Diff line change
Expand Up @@ -25,8 +25,8 @@ def get_impl_cls() -> Type["ROCmFlashAttentionImpl"]:
return ROCmFlashAttentionImpl

@staticmethod
def make_metadata(*args, **kwargs) -> "ROCmFlashAttentionMetadata":
return ROCmFlashAttentionMetadata(*args, **kwargs)
def get_metadata_cls() -> Type["AttentionMetadata"]:
return ROCmFlashAttentionMetadata

@staticmethod
def get_kv_cache_shape(
Expand Down
4 changes: 2 additions & 2 deletions vllm/attention/backends/torch_sdpa.py
Original file line number Diff line number Diff line change
Expand Up @@ -31,8 +31,8 @@ def get_impl_cls() -> Type["TorchSDPABackendImpl"]:
return TorchSDPABackendImpl

@staticmethod
def make_metadata(*args, **kwargs) -> "TorchSDPAMetadata":
return TorchSDPAMetadata(*args, **kwargs)
def get_metadata_cls() -> Type["AttentionMetadata"]:
return TorchSDPAMetadata

@staticmethod
def get_kv_cache_shape(
Expand Down
4 changes: 2 additions & 2 deletions vllm/attention/backends/xformers.py
Original file line number Diff line number Diff line change
Expand Up @@ -28,8 +28,8 @@ def get_impl_cls() -> Type["XFormersImpl"]:
return XFormersImpl

@staticmethod
def make_metadata(*args, **kwargs) -> "XFormersMetadata":
return XFormersMetadata(*args, **kwargs)
def get_metadata_cls() -> Type["AttentionMetadata"]:
return XFormersMetadata

@staticmethod
def get_kv_cache_shape(
Expand Down
16 changes: 8 additions & 8 deletions vllm/executor/distributed_gpu_executor.py
Original file line number Diff line number Diff line change
Expand Up @@ -64,8 +64,8 @@ def initialize_cache(self, num_gpu_blocks: int,
num_cpu_blocks=num_cpu_blocks)

def execute_model(
self,
execute_model_req: ExecuteModelRequest) -> List[SamplerOutput]:
self, execute_model_req: ExecuteModelRequest
) -> Optional[List[SamplerOutput]]:
if self.parallel_worker_tasks is None:
self.parallel_worker_tasks = self._run_workers(
"start_worker_execution_loop",
Expand All @@ -79,7 +79,7 @@ def stop_remote_worker_execution_loop(self) -> None:
if self.parallel_worker_tasks is None:
return

self._driver_execute_model()
self._driver_execute_model(execute_model_req=None)
parallel_worker_tasks = self.parallel_worker_tasks
self.parallel_worker_tasks = None
# Ensure that workers exit model loop cleanly
Expand Down Expand Up @@ -123,13 +123,13 @@ def save_sharded_state(

@abstractmethod
def _driver_execute_model(
self,
execute_model_req: Optional[ExecuteModelRequest] = None
) -> List[SamplerOutput]:
self, execute_model_req: Optional[ExecuteModelRequest]
) -> Optional[List[SamplerOutput]]:
"""Run execute_model in the driver worker.

Passing None will cause the driver to stop the model execution
loop running in each of the remote workers.
Passing None will cause the driver to stop the model execution loop
running in each of the remote workers. In this case, this method
returns None. Otherwise, this method returns the model output.
"""
raise NotImplementedError

Expand Down
4 changes: 2 additions & 2 deletions vllm/executor/executor_base.py
Original file line number Diff line number Diff line change
Expand Up @@ -69,8 +69,8 @@ def initialize_cache(self, num_gpu_blocks: int,

@abstractmethod
def execute_model(
self,
execute_model_req: ExecuteModelRequest) -> List[SamplerOutput]:
self, execute_model_req: ExecuteModelRequest
) -> Optional[List[SamplerOutput]]:
"""Executes at least one model step on the given sequences."""
raise NotImplementedError

Expand Down
2 changes: 1 addition & 1 deletion vllm/executor/gpu_executor.py
Original file line number Diff line number Diff line change
Expand Up @@ -87,7 +87,7 @@ def initialize_cache(self, num_gpu_blocks: int, num_cpu_blocks) -> None:

def execute_model(
self, execute_model_req: ExecuteModelRequest
) -> List[Union[SamplerOutput, PoolerOutput]]:
) -> Optional[List[Union[SamplerOutput, PoolerOutput]]]:
output = self.driver_worker.execute_model(execute_model_req)
return output

Expand Down
Loading
Loading