162 lines
6.7 KiB
Python
162 lines
6.7 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
"""CPU contract coverage for NPU routing capture and the native data plane."""
|
|
|
|
from types import SimpleNamespace
|
|
|
|
import numpy as np
|
|
import pytest
|
|
import torch
|
|
from vllm.config import AuxOutputConfig
|
|
from vllm.distributed.aux_output_connector import worker
|
|
from vllm.distributed.aux_output_connector.connector import (
|
|
AuxOutputConnectorMetadata,
|
|
AuxOutputSchedulerConnector,
|
|
PackedBlockHashes,
|
|
)
|
|
from vllm.model_executor.layers.fused_moe import routed_experts_capturer
|
|
from vllm.v1.kv_cache_interface import KVCacheConfig
|
|
|
|
from vllm_omni.platforms.npu.worker.aux_output import NPUAuxOutputMixin
|
|
|
|
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
|
|
|
|
|
class _Router:
|
|
def _select_experts(self, logical_ids):
|
|
return torch.ones_like(logical_ids, dtype=torch.float32), logical_ids
|
|
|
|
|
|
class _AscendMoE(torch.nn.Module):
|
|
__module__ = "vllm_ascend.ops.fused_moe"
|
|
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.layer_id = 0
|
|
self.router = _Router()
|
|
|
|
|
|
class _Runner(NPUAuxOutputMixin):
|
|
def _sync_device(self):
|
|
self.synchronized += 1
|
|
|
|
|
|
@pytest.fixture
|
|
def runner(monkeypatch):
|
|
# Exercise the real capturer, block store, worker and scheduler connectors
|
|
# with CPU tensors; importing vllm-ascend or owning NPU hardware is unnecessary.
|
|
monkeypatch.setattr(routed_experts_capturer, "current_platform", SimpleNamespace(device_type="cpu"))
|
|
monkeypatch.setattr(routed_experts_capturer, "get_forward_context", lambda: SimpleNamespace(dp_metadata=None))
|
|
monkeypatch.setattr(worker, "get_tp_group", lambda: SimpleNamespace(is_first_rank=True))
|
|
runner = _Runner()
|
|
runner.model = torch.nn.Sequential(_AscendMoE())
|
|
runner.vllm_config = SimpleNamespace(
|
|
aux_output_config=AuxOutputConfig(enable_return_routed_experts=True),
|
|
model_config=SimpleNamespace(
|
|
get_total_num_hidden_layers=lambda: 1,
|
|
get_num_experts=lambda: 16,
|
|
get_num_experts_per_tok=lambda: 2,
|
|
hf_text_config=SimpleNamespace(model_type="test_moe"),
|
|
),
|
|
scheduler_config=SimpleNamespace(max_num_batched_tokens=16, max_num_seqs=4),
|
|
parallel_config=SimpleNamespace(data_parallel_rank=0, tensor_parallel_size=1, decode_context_parallel_size=1),
|
|
cache_config=SimpleNamespace(block_size=4, prefix_match_unit=None),
|
|
max_concurrent_batches=2,
|
|
)
|
|
runner.input_batch = SimpleNamespace(req_ids=["r"], num_computed_tokens_cpu=np.array([0]))
|
|
runner.query_start_loc = SimpleNamespace(np=np.array([0, 4]))
|
|
runner.synchronized = 0
|
|
runner._init_omni_aux_output(KVCacheConfig(8, [], []))
|
|
try:
|
|
yield runner
|
|
finally:
|
|
runner._close_omni_aux_output()
|
|
|
|
|
|
def _step(runner, *, emit_start=0, hashes=None, finished=()):
|
|
scheduler_output = SimpleNamespace(
|
|
aux_output_connector_metadata=AuxOutputConnectorMetadata(
|
|
generation=0,
|
|
requests={} if finished else {"r": emit_start},
|
|
block_hashes=hashes or {},
|
|
finished_requests=finished,
|
|
),
|
|
scheduled_spec_decode_tokens={},
|
|
)
|
|
runner._begin_omni_aux_output_step(scheduler_output)
|
|
return scheduler_output
|
|
|
|
|
|
def _capture(runner, ids):
|
|
logical_ids = torch.tensor(ids, dtype=torch.int32)
|
|
weights, returned = runner.model[0].router._select_experts(logical_ids)
|
|
assert torch.equal(returned, logical_ids)
|
|
assert weights.shape == logical_ids.shape
|
|
return logical_ids.numpy()[:, None, :].astype(np.uint8)
|
|
|
|
|
|
def test_routing_reaches_scheduler_and_owns_legacy_batch_snapshot(runner):
|
|
step = _step(runner)
|
|
expected = _capture(runner, [[0, 1], [1, 2], [2, 3], [3, 4]])
|
|
pending = runner._prepare_omni_aux_output()
|
|
# Simulate bookkeeping and the next batch reusing the runner's buffers.
|
|
runner.input_batch.num_computed_tokens_cpu[:] = 99
|
|
runner.query_start_loc.np[:] = 99
|
|
_capture(runner, [[7, 8]] * 4)
|
|
output = runner._finish_omni_aux_output(pending, step, torch.tensor([[5]]))
|
|
request = SimpleNamespace(request_id="r", num_tokens=5)
|
|
rows = AuxOutputSchedulerConnector().take_output(request, output)
|
|
np.testing.assert_array_equal(rows, expected)
|
|
assert output["r"].token_start == 0
|
|
assert runner.synchronized == 1
|
|
|
|
|
|
def test_speculative_rejection_drops_unaccepted_routing_suffix(runner):
|
|
step = _step(runner)
|
|
step.scheduled_spec_decode_tokens = {"r": [7, 8, 9]}
|
|
expected = _capture(runner, [[0, 1], [1, 2], [2, 3], [3, 4]])
|
|
pending = runner._prepare_omni_aux_output()
|
|
output = runner._finish_omni_aux_output(pending, step, torch.tensor([[5, 6, -1, -1]]))
|
|
np.testing.assert_array_equal(output["r"].rows, expected[:2])
|
|
|
|
|
|
def test_partial_prefill_accumulates_rows_and_no_forward_step_tears_down(runner):
|
|
step = _step(runner)
|
|
prefix = _capture(runner, [[0, 1], [1, 2], [2, 3], [3, 4]])
|
|
pending = runner._prepare_omni_aux_output()
|
|
assert runner._finish_omni_aux_output(pending, step, torch.tensor([[5]]), [0]) == {}
|
|
|
|
# Both completed four-token blocks need their scheduler-provided hashes.
|
|
step = _step(runner, hashes={"r": PackedBlockHashes(b"h0h1", 2)})
|
|
runner.input_batch.num_computed_tokens_cpu[:] = 4
|
|
suffix = _capture(runner, [[4, 5], [5, 6], [6, 7], [7, 8]])
|
|
pending = runner._prepare_omni_aux_output()
|
|
output = runner._finish_omni_aux_output(pending, step, torch.tensor([[5]]))
|
|
np.testing.assert_array_equal(output["r"].rows, np.concatenate((prefix, suffix)))
|
|
_step(runner, finished=("r",))
|
|
assert not runner.aux_output_connector._requests
|
|
store = runner.aux_output_connector._store
|
|
runner._close_omni_aux_output()
|
|
assert not store._thread.is_alive()
|
|
assert runner.model[0].capture_fn is None
|
|
runner._close_omni_aux_output() # shutdown is idempotent
|
|
|
|
|
|
def test_reinitialization_does_not_wrap_router_twice(runner):
|
|
selection = runner.model[0].router._select_experts
|
|
runner._init_omni_aux_output(KVCacheConfig(8, [], []))
|
|
assert runner.model[0].router._select_experts is selection
|
|
step = _step(runner)
|
|
expected = _capture(runner, [[0, 1]] * 4)
|
|
output = runner._finish_omni_aux_output(runner._prepare_omni_aux_output(), step, torch.tensor([[5]]))
|
|
np.testing.assert_array_equal(output["r"].rows, expected)
|
|
|
|
|
|
def test_disabled_capture_preserves_normal_runner_path(runner):
|
|
runner._close_omni_aux_output()
|
|
runner.vllm_config.aux_output_config.enable_return_routed_experts = False
|
|
runner._init_omni_aux_output(KVCacheConfig(8, [], []))
|
|
runner._begin_omni_aux_output_step(SimpleNamespace())
|
|
assert runner._prepare_omni_aux_output() is None
|
|
assert runner._finish_omni_aux_output(None, SimpleNamespace()) is None
|