Taichang ZhouandClaude Opus 4.8
9678ae58cf
Squashed migration rebased onto current main. Collapses the 46-commit dev/migrate-MR-v2 history (10 merge commits + superseded pre-rebase work) into a single commit, so the diff reflects only the genuine MR-V2 feature (worker_v2/ runners + model_states, omni schedulers, stage engine pipeline, pooling_output serialization bridge). Content already merged to main separately (step_audio2, tts_adapters, npu/_310p, etc.) drops out. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
96 lines
3.1 KiB
Python
96 lines
3.1 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
from types import SimpleNamespace
|
|
|
|
from vllm.compilation.cuda_graph import CUDAGraphMode
|
|
|
|
from vllm_omni.worker_v2.omni_model_runner import OmniGPUModelRunner
|
|
|
|
|
|
def test_eager_dispatch_uses_eager_without_cudagraph_manager():
|
|
runner = object.__new__(OmniGPUModelRunner)
|
|
runner.cudagraph_manager = None
|
|
runner.dp_size = 1
|
|
runner.dp_rank = 0
|
|
|
|
batch_desc, num_tokens_across_dp = runner._dispatch_batch_descriptor(
|
|
num_reqs=1,
|
|
num_toks=8,
|
|
uniform_tok_count=8,
|
|
use_eager=True,
|
|
)
|
|
|
|
assert batch_desc.cg_mode == CUDAGraphMode.NONE
|
|
assert batch_desc.num_tokens == 8
|
|
assert batch_desc.num_reqs == 1
|
|
assert num_tokens_across_dp is None
|
|
|
|
|
|
def test_eager_dispatch_syncs_dp_padding(monkeypatch):
|
|
runner = object.__new__(OmniGPUModelRunner)
|
|
runner.cudagraph_manager = None
|
|
runner.dp_size = 2
|
|
runner.dp_rank = 1
|
|
expected_tokens_across_dp = object()
|
|
calls = []
|
|
|
|
def sync_dp(cudagraph_manager, batch_desc, num_tokens, num_reqs, uniform_token_count, dp_size, dp_rank):
|
|
calls.append((cudagraph_manager, batch_desc, num_tokens, num_reqs, uniform_token_count, dp_size, dp_rank))
|
|
return batch_desc, expected_tokens_across_dp
|
|
|
|
import vllm.v1.worker.gpu.dp_utils as dp_utils
|
|
|
|
monkeypatch.setattr(dp_utils, "sync_cudagraph_and_dp_padding", sync_dp)
|
|
|
|
batch_desc, num_tokens_across_dp = runner._dispatch_batch_descriptor(
|
|
num_reqs=1,
|
|
num_toks=8,
|
|
uniform_tok_count=8,
|
|
use_eager=True,
|
|
)
|
|
|
|
assert batch_desc.cg_mode == CUDAGraphMode.NONE
|
|
assert batch_desc.num_tokens == 8
|
|
assert batch_desc.num_reqs == 1
|
|
assert num_tokens_across_dp is expected_tokens_across_dp
|
|
assert calls == [(None, batch_desc, 8, 1, 8, 2, 1)]
|
|
|
|
|
|
def test_non_eager_dispatch_uses_cudagraph_manager():
|
|
runner = object.__new__(OmniGPUModelRunner)
|
|
expected = SimpleNamespace(cg_mode=CUDAGraphMode.PIECEWISE, num_tokens=8, num_reqs=1)
|
|
calls = []
|
|
runner.cudagraph_manager = SimpleNamespace(dispatch=lambda *args: calls.append(args) or expected)
|
|
runner.dp_size = 1
|
|
runner.dp_rank = 0
|
|
|
|
batch_desc, num_tokens_across_dp = runner._dispatch_batch_descriptor(
|
|
num_reqs=1,
|
|
num_toks=8,
|
|
uniform_tok_count=8,
|
|
use_eager=False,
|
|
)
|
|
|
|
assert batch_desc is expected
|
|
assert calls == [(1, 8, 8)]
|
|
assert num_tokens_across_dp is None
|
|
|
|
|
|
if __name__ == "__main__":
|
|
test_eager_dispatch_uses_eager_without_cudagraph_manager()
|
|
test_non_eager_dispatch_uses_cudagraph_manager()
|
|
|
|
|
|
def test_skip_compiled_uses_eager_dispatch_without_overwriting_batch_descriptor():
|
|
source = __import__("inspect").getsource(OmniGPUModelRunner.execute_model)
|
|
|
|
assert "use_eager=is_profile or skip_compiled" in source
|
|
assert "if self.is_encoder_decoder and scheduler_output.scheduled_encoder_inputs" not in source
|
|
|
|
|
|
def test_fullgraph_requires_cudagraph_manager_contract():
|
|
source = __import__("inspect").getsource(OmniGPUModelRunner.execute_model)
|
|
|
|
assert "assert self.cudagraph_manager is not None" in source
|