62 lines
2.1 KiB
Python
62 lines
2.1 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
|
|
from types import SimpleNamespace
|
|
|
|
import pytest
|
|
from vllm.compilation.cuda_graph import CUDAGraphMode
|
|
|
|
from vllm_omni.worker_v2.omni_model_runner import OmniGPUModelRunner
|
|
|
|
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
|
|
|
|
|
@pytest.mark.parametrize("dp_size", [1, 2])
|
|
def test_dispatch_eager_and_manager_paths(monkeypatch, dp_size):
|
|
runner = object.__new__(OmniGPUModelRunner)
|
|
runner.cudagraph_manager = None
|
|
runner.parallel_config = SimpleNamespace(data_parallel_size=dp_size, data_parallel_rank=1)
|
|
runner.dp_size, runner.dp_rank = dp_size, 1
|
|
|
|
import vllm.v1.worker.gpu.dp_utils as dp_utils
|
|
|
|
expected_tokens = object()
|
|
sync_calls = []
|
|
|
|
def sync_padding(*args, **kwargs):
|
|
sync_calls.append((args, kwargs))
|
|
return args[1], expected_tokens
|
|
|
|
monkeypatch.setattr(
|
|
dp_utils,
|
|
"sync_cudagraph_and_dp_padding",
|
|
sync_padding,
|
|
)
|
|
|
|
batch_desc, num_tokens_across_dp = runner._dispatch_batch_descriptor(
|
|
num_reqs=1, num_toks=8, uniform_tok_count=8, num_active_loras=0, use_eager=True, max_query_len=8
|
|
)
|
|
assert batch_desc.cg_mode == CUDAGraphMode.NONE
|
|
if dp_size == 1:
|
|
assert num_tokens_across_dp is None
|
|
assert sync_calls == []
|
|
else:
|
|
assert num_tokens_across_dp is expected_tokens
|
|
|
|
# Non-eager dispatch goes through the cudagraph manager.
|
|
expected = SimpleNamespace(cg_mode=CUDAGraphMode.PIECEWISE, num_tokens=8, num_reqs=1)
|
|
dispatch_calls = []
|
|
|
|
def dispatch(*args, **kwargs):
|
|
dispatch_calls.append((args, kwargs))
|
|
return expected
|
|
|
|
runner.cudagraph_manager = SimpleNamespace(dispatch=dispatch)
|
|
runner.parallel_config.data_parallel_size = 1
|
|
runner.dp_size = 1
|
|
batch_desc, _ = runner._dispatch_batch_descriptor(
|
|
num_reqs=1, num_toks=8, uniform_tok_count=8, num_active_loras=0, use_eager=False, max_query_len=8
|
|
)
|
|
assert batch_desc is expected
|
|
assert dispatch_calls == [((1, 8, 8), {"num_active_loras": 0, "max_query_len": 8})]
|