diff --git a/.buildkite/test-merge.yml b/.buildkite/test-merge.yml index 130e9e59c..9e5027272 100644 --- a/.buildkite/test-merge.yml +++ b/.buildkite/test-merge.yml @@ -1,5 +1,6 @@ env: VLLM_WORKER_MULTIPROC_METHOD: spawn + VLLM_USE_FLASHINFER_MOE_FP16: 0 HF_HUB_DOWNLOAD_TIMEOUT: 300 HF_HUB_ETAG_TIMEOUT: 60 diff --git a/.buildkite/test-nightly.yml b/.buildkite/test-nightly.yml index 9a3a0c4b5..779e09efc 100644 --- a/.buildkite/test-nightly.yml +++ b/.buildkite/test-nightly.yml @@ -1,5 +1,6 @@ env: VLLM_WORKER_MULTIPROC_METHOD: spawn + VLLM_USE_FLASHINFER_MOE_FP16: 0 HF_HUB_DOWNLOAD_TIMEOUT: 300 HF_HUB_ETAG_TIMEOUT: 60 diff --git a/.buildkite/test-ready.yml b/.buildkite/test-ready.yml index a9aea9f62..15188fde5 100644 --- a/.buildkite/test-ready.yml +++ b/.buildkite/test-ready.yml @@ -1,5 +1,6 @@ env: VLLM_WORKER_MULTIPROC_METHOD: spawn + VLLM_USE_FLASHINFER_MOE_FP16: 0 HF_HUB_DOWNLOAD_TIMEOUT: 300 HF_HUB_ETAG_TIMEOUT: 60 diff --git a/.buildkite/test-weekly.yml b/.buildkite/test-weekly.yml index dba648996..00c9b25f2 100644 --- a/.buildkite/test-weekly.yml +++ b/.buildkite/test-weekly.yml @@ -1,5 +1,6 @@ env: VLLM_WORKER_MULTIPROC_METHOD: spawn + VLLM_USE_FLASHINFER_MOE_FP16: 0 HF_HUB_DOWNLOAD_TIMEOUT: 300 HF_HUB_ETAG_TIMEOUT: 60 diff --git a/docker/Dockerfile.ci b/docker/Dockerfile.ci index a263c12e2..a823b5bc5 100644 --- a/docker/Dockerfile.ci +++ b/docker/Dockerfile.ci @@ -1,5 +1,5 @@ ARG VLLM_BASE_IMAGE=vllm/vllm-openai -ARG VLLM_BASE_TAG=v0.20.0 +ARG VLLM_BASE_TAG=v0.21.0 FROM ${VLLM_BASE_IMAGE}:${VLLM_BASE_TAG} ARG APP_DIR=/workspace/vllm-omni WORKDIR ${APP_DIR} diff --git a/docs/getting_started/installation/gpu.md b/docs/getting_started/installation/gpu.md index 8e53eab5c..e78f3d797 100644 --- a/docs/getting_started/installation/gpu.md +++ b/docs/getting_started/installation/gpu.md @@ -34,7 +34,7 @@ vLLM-Omni is a Python library that supports the following GPU variants. The libr ### Pre-built wheels -Note: Pre-built wheels are currently available for vLLM-Omni 0.11.0rc1, 0.12.0rc1, 0.14.0rc1, 0.14.0, 0.16.0, 0.18.0, and 0.20.0. If you need a newer unreleased revision, please [build from source](https://docs.vllm.ai/projects/vllm-omni/en/latest/getting_started/installation/gpu/#build-wheel-from-source). +Note: Pre-built wheels are currently available for vLLM-Omni 0.11.0rc1, 0.12.0rc1, 0.14.0rc1, 0.14.0, 0.16.0, 0.18.0, 0.20.0, and 0.21.0. If you need a newer unreleased revision, please [build from source](https://docs.vllm.ai/projects/vllm-omni/en/latest/getting_started/installation/gpu/#build-wheel-from-source). === "NVIDIA CUDA" diff --git a/docs/getting_started/installation/gpu/cuda.inc.md b/docs/getting_started/installation/gpu/cuda.inc.md index 840cb2b4a..69fc05351 100644 --- a/docs/getting_started/installation/gpu/cuda.inc.md +++ b/docs/getting_started/installation/gpu/cuda.inc.md @@ -20,7 +20,7 @@ Therefore, it is recommended to install vLLM and vLLM-Omni with a **fresh new** vLLM-Omni is built based on vLLM. Please install it with command below. ```bash -uv pip install vllm==0.20.0 --torch-backend=auto +uv pip install vllm==0.21.0 --torch-backend=auto ``` #### Installation of vLLM-Omni @@ -39,13 +39,13 @@ uv pip install 'vllm-omni[demo]' # --8<-- [start:build-wheel-from-source] #### Installation of vLLM -If you do not need to modify source code of vLLM, you can directly install the stable 0.20.0 release version of the library +If you do not need to modify source code of vLLM, you can directly install the stable 0.21.0 release version of the library ```bash -uv pip install vllm==0.20.0 --torch-backend=auto +uv pip install vllm==0.21.0 --torch-backend=auto ``` -The 0.20.0 release of vLLM ships CUDA 13.0-compatible binaries by default. If you need a different CUDA variant or want to reuse an existing PyTorch installation, build vLLM from source instead. +The 0.21.0 release of vLLM ships CUDA 13.0-compatible binaries by default. If you need a different CUDA variant or want to reuse an existing PyTorch installation, build vLLM from source instead. #### Installation of vLLM-Omni Since vllm-omni is rapidly evolving, it's recommended to install it from source @@ -66,12 +66,12 @@ If you want to check, modify or debug with source code of vLLM, install the libr ```bash git clone https://github.com/vllm-project/vllm.git cd vllm -git checkout v0.20.0 +git checkout v0.21.0 ``` Set up environment variables to get pre-built wheels. If there are internet problems, just download the whl file manually. And set `VLLM_PRECOMPILED_WHEEL_LOCATION` as your local absolute path of whl file. ```bash -#For CUDA 13.0 (the default for v0.20.0; the wheel filename has no `+cu130` suffix) -export VLLM_PRECOMPILED_WHEEL_LOCATION=https://github.com/vllm-project/vllm/releases/download/v0.20.0/vllm-0.20.0-cp38-abi3-manylinux_2_35_x86_64.whl +#For CUDA 13.0 (the default for v0.21.0; the wheel filename has no `+cu130` suffix) +export VLLM_PRECOMPILED_WHEEL_LOCATION=https://github.com/vllm-project/vllm/releases/download/v0.21.0/vllm-0.21.0-cp38-abi3-manylinux_2_35_x86_64.whl ``` Install vllm with command below (If you have no existing PyTorch). ```bash @@ -102,7 +102,7 @@ docker run --runtime nvidia --gpus 2 \ --env "HF_TOKEN=$HF_TOKEN" \ -p 8091:8091 \ --ipc=host \ - vllm/vllm-omni:v0.20.0 \ + vllm/vllm-omni:v0.21.0 \ --model Qwen/Qwen3-Omni-30B-A3B-Instruct --port 8091 ``` diff --git a/docs/getting_started/installation/gpu/rocm.inc.md b/docs/getting_started/installation/gpu/rocm.inc.md index 608c259af..4d0fcfdf3 100644 --- a/docs/getting_started/installation/gpu/rocm.inc.md +++ b/docs/getting_started/installation/gpu/rocm.inc.md @@ -13,7 +13,7 @@ vLLM-Omni current recommends the steps in under setup through Docker Images. vLLM-Omni is built based on vLLM. Please install it with command below. ```bash -uv pip install vllm==0.20.0+rocm721 --extra-index-url https://wheels.vllm.ai/rocm/0.20.0/rocm721 +uv pip install vllm==0.21.0+rocm721 --extra-index-url https://wheels.vllm.ai/rocm/0.21.0/rocm721 ``` #### Installation of vLLM-Omni @@ -34,13 +34,13 @@ uv pip install onnxruntime-rocm # --8<-- [start:build-wheel-from-source] #### Installation of vLLM -If you do not need to modify source code of vLLM, you can directly install the stable 0.20.0 release version of the library +If you do not need to modify source code of vLLM, you can directly install the stable 0.21.0 release version of the library ```bash -uv pip install vllm==0.20.0+rocm721 --extra-index-url https://wheels.vllm.ai/rocm/0.20.0/rocm721 +uv pip install vllm==0.21.0+rocm721 --extra-index-url https://wheels.vllm.ai/rocm/0.21.0/rocm721 ``` -The pre-built 0.20.0 vLLM wheel targets ROCm 7.2.1. If you need a different ROCm stack or want to reuse an existing PyTorch installation, build vLLM from source instead. +The pre-built 0.21.0 vLLM wheel targets ROCm 7.2.1. If you need a different ROCm stack or want to reuse an existing PyTorch installation, build vLLM from source instead. #### Installation of vLLM-Omni Since vllm-omni is rapidly evolving, it's recommended to install it from source @@ -58,7 +58,7 @@ If you want to check, modify or debug with source code of vLLM, install the libr ```bash git clone https://github.com/vllm-project/vllm.git cd vllm -git checkout v0.20.0 +git checkout v0.21.0 python3 -m pip install -r requirements/rocm.txt python3 setup.py develop ``` @@ -130,7 +130,7 @@ docker run --rm \ -v ~/.cache/huggingface:/root/.cache/huggingface \ --env "HF_TOKEN=$HF_TOKEN" \ -p 8091:8091 \ - vllm/vllm-omni-rocm:v0.20.0 \ + vllm/vllm-omni-rocm:v0.21.0 \ --model Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 ``` @@ -149,7 +149,7 @@ docker run --rm -it \ -v ~/.cache/huggingface:/root/.cache/huggingface \ --env "HF_TOKEN=$HF_TOKEN" \ --entrypoint bash \ - vllm/vllm-omni-rocm:v0.20.0 + vllm/vllm-omni-rocm:v0.21.0 ``` # --8<-- [end:pre-built-images] diff --git a/docs/getting_started/quickstart.md b/docs/getting_started/quickstart.md index dfebe8b61..10665f2c2 100644 --- a/docs/getting_started/quickstart.md +++ b/docs/getting_started/quickstart.md @@ -19,10 +19,10 @@ uv venv --python 3.12 --seed source .venv/bin/activate # On CUDA -uv pip install vllm==0.20.0 --torch-backend=auto +uv pip install vllm==0.21.0 --torch-backend=auto # On ROCm -uv pip install vllm==0.20.0+rocm721 --extra-index-url https://wheels.vllm.ai/rocm/0.20.0/rocm721 +uv pip install vllm==0.21.0+rocm721 --extra-index-url https://wheels.vllm.ai/rocm/0.21.0/rocm721 git clone https://github.com/vllm-project/vllm-omni.git cd vllm-omni @@ -35,7 +35,7 @@ For additional installation methods — please see the [installation guide](inst !!! note It is important to install the same major & minor version of vLLM and vLLM Omni, otherwise things may not work as expected. If the versions are misaligned, you will see a warning when you import vLLM Omni. - If you are seeing strange behavior with the `vllm` command not handling the `--omni` flag correctly, you most likely have a version mismatch with vLLM < `0.20.0` and vLLM Omni `0.20.0`, as vLLM Omni no longer hijacks the vLLM entrypoint. Updating vLLM should resolve this issue. + If you are seeing strange behavior with the `vllm` command not handling the `--omni` flag correctly, you most likely have a version mismatch with vLLM < `0.21.0` and vLLM Omni `0.21.0`, as vLLM Omni no longer hijacks the vLLM entrypoint. Updating vLLM should resolve this issue. ## Offline Inference diff --git a/tests/diffusion/attention/test_attention_config.py b/tests/diffusion/attention/test_attention_config.py index 35be5e6f0..5dea43b91 100644 --- a/tests/diffusion/attention/test_attention_config.py +++ b/tests/diffusion/attention/test_attention_config.py @@ -394,6 +394,9 @@ class TestAttentionInitUsesCurrentDiffusionConfig: per_role={"cross": AttentionSpec(backend="TORCH_SDPA", extra={"block_size": 128})}, ), parallel_config=SimpleNamespace(ring_degree=2), + diffusion_kv_cache_dtype=None, + diffusion_kv_cache_skip_step_indices=None, + diffusion_kv_cache_skip_layer_indices=None, ) with set_current_diffusion_config(od_config): diff --git a/tests/diffusion/offloader/test_diffusion_cpu_offload.py b/tests/diffusion/offloader/test_diffusion_cpu_offload.py index e7a8a44cb..142e6b437 100644 --- a/tests/diffusion/offloader/test_diffusion_cpu_offload.py +++ b/tests/diffusion/offloader/test_diffusion_cpu_offload.py @@ -5,6 +5,7 @@ import pytest import torch from vllm.distributed.parallel_state import cleanup_dist_env_and_memory +from tests.helpers import skip_if_gated_repo_inaccessible from tests.helpers.env import DeviceMemoryMonitor from tests.helpers.mark import hardware_test from tests.helpers.runtime import OmniRunner @@ -21,6 +22,11 @@ IMAGE_VIDEO_MODELS = { MODELS = {**AUDIO_MODEL, **IMAGE_VIDEO_MODELS} +_GATED_MODELS = {"stabilityai/stable-audio-open-1.0"} + +# Aliased for backward compatibility (imported by test_diffusion_layerwise_offload.py). +_skip_if_gated_repo_inaccessible = skip_if_gated_repo_inaccessible + AUDIO_MODEL_PARAMS = { "runner_params": {}, "sampler_params": {}, @@ -94,10 +100,23 @@ def check_audio_determinism(audio1, audio2, atol=1e-2): @hardware_test(res={"cuda": "L4", "rocm": "MI325"}) @pytest.mark.parametrize("model_name", list(MODELS.keys())) def test_cpu_offload_diffusion_model(model_name: str): + if model_name in _GATED_MODELS: + _skip_if_gated_repo_inaccessible(model_name) try: offload_peak_memory, output_offload = inference(model_name, offload=True) cleanup_dist_env_and_memory() no_offload_peak_memory, output_no_offload = inference(model_name, offload=False) + except ValueError as exc: + # omni_snapshot_download wraps GatedRepoError in a ValueError. + # If the pre-flight guard above did not catch it (e.g. partial + # HF_TOKEN where config.json is accessible but weight shards are + # blocked), skip instead of failing. + if "Access to model" in str(exc) and "is restricted" in str(exc): + pytest.skip( + f"Skipping: gated HF repo {model_name!r} inaccessible " + f"({exc}). See docs/contributing/ci/hf_credentials.md." + ) + pytest.fail(f"Inference failed: {exc}") except Exception: pytest.fail("Inference failed") print(f"Offload peak memory: {offload_peak_memory} MB") diff --git a/tests/diffusion/offloader/test_diffusion_layerwise_offload.py b/tests/diffusion/offloader/test_diffusion_layerwise_offload.py index a66c32563..f02da6fab 100644 --- a/tests/diffusion/offloader/test_diffusion_layerwise_offload.py +++ b/tests/diffusion/offloader/test_diffusion_layerwise_offload.py @@ -3,6 +3,10 @@ import pytest import torch from vllm.distributed.parallel_state import cleanup_dist_env_and_memory +from tests.diffusion.offloader.test_diffusion_cpu_offload import ( + _GATED_MODELS, + _skip_if_gated_repo_inaccessible, +) from tests.helpers.env import DeviceMemoryMonitor from tests.helpers.runtime import OmniRunner from vllm_omni.inputs.data import OmniDiffusionSamplingParams @@ -96,6 +100,8 @@ def test_layerwise_offload_diffusion_model(model_name: str): offloader keeps only a single transformer block on GPU at a time, with prefetching for compute-memory overlap. """ + if model_name in _GATED_MODELS: + _skip_if_gated_repo_inaccessible(model_name) try: # Run without layerwise offloading (baseline) no_offload_peak_memory, output_no_offload = run_inference(model_name, layerwise_offload=False) @@ -104,6 +110,17 @@ def test_layerwise_offload_diffusion_model(model_name: str): # Run with layerwise offloading (1 layer on device) layerwise_offload_peak_memory, output_offload = run_inference(model_name, layerwise_offload=True) cleanup_dist_env_and_memory() + except ValueError as exc: + # omni_snapshot_download wraps GatedRepoError in a ValueError. + # If the pre-flight guard above did not catch it (e.g. partial + # HF_TOKEN where config.json is accessible but weight shards are + # blocked), skip instead of failing. + if "Access to model" in str(exc) and "is restricted" in str(exc): + pytest.skip( + f"Skipping: gated HF repo {model_name!r} inaccessible " + f"({exc}). See docs/contributing/ci/hf_credentials.md." + ) + pytest.fail(f"Inference failed: {exc}") except Exception: pytest.fail("Inference failed") diff --git a/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py b/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py index 87430fd7a..58e885d46 100644 --- a/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py +++ b/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py @@ -540,6 +540,7 @@ def test_generation_scheduler_calls_cleanup_on_finished(monkeypatch, mocker: Moc kv_connector_output=None, cudagraph_stats=None, req_id_to_index={"req-s1": 0}, + routed_experts_dict=None, ) OmniGenerationScheduler.update_from_output(scheduler, scheduler_output, model_runner_output) @@ -625,6 +626,7 @@ def test_ar_scheduler_defers_cleanup_and_queues_save_on_finished(mocker: MockerF cudagraph_stats=None, req_id_to_index={"req-ar": 0}, kv_extracted_req_ids=None, + routed_experts_dict=None, ) OmniARScheduler.update_from_output(scheduler, scheduler_output, model_runner_output) diff --git a/tests/distributed/omni_connectors/test_mooncake_transfer_engine_rdma.py b/tests/distributed/omni_connectors/test_mooncake_transfer_engine_rdma.py index e5781f719..60653b3cd 100755 --- a/tests/distributed/omni_connectors/test_mooncake_transfer_engine_rdma.py +++ b/tests/distributed/omni_connectors/test_mooncake_transfer_engine_rdma.py @@ -25,7 +25,11 @@ from vllm_omni.distributed.omni_connectors.connectors.mooncake_transfer_engine_c ) # All tests in this file require Mooncake TransferEngine and an RDMA environment. -pytestmark = [pytest.mark.parallel, pytest.mark.core_model] +pytestmark = [ + pytest.mark.parallel, + pytest.mark.core_model, + pytest.mark.skipif(TransferEngine is None, reason="Mooncake TransferEngine not installed"), +] # --------------------------------------------------------------------------- # Shared helpers diff --git a/tests/e2e/offline_inference/test_moss_tts_nano_expansion.py b/tests/e2e/offline_inference/test_moss_tts_nano_expansion.py index 6b9304f48..7592fdc40 100644 --- a/tests/e2e/offline_inference/test_moss_tts_nano_expansion.py +++ b/tests/e2e/offline_inference/test_moss_tts_nano_expansion.py @@ -106,18 +106,21 @@ def _build_request( def _collect_audio(omni: Omni, request: dict) -> tuple[torch.Tensor, int]: """Run a single request and return (waveform, sample_rate).""" - # Omni.generate returns list[OmniRequestOutput] by default (py_generator=False). - for omni_out in omni.generate(request, DEFAULT_SAMPLING): - mm = omni_out.multimodal_output - assert mm is not None, "Expected multimodal_output to be non-None" - audio = mm.get("audio") - sr = mm.get("sr") - assert audio is not None, "Expected 'audio' key in multimodal_output" - assert isinstance(audio, torch.Tensor), f"audio should be Tensor, got {type(audio)}" - return audio.cpu(), int(sr.item()) if sr is not None else SAMPLE_RATE + for stage_outputs in omni.generate(request, DEFAULT_SAMPLING): + req_output = stage_outputs.request_output + if req_output is not None: + mm = req_output.outputs[0].multimodal_output + assert mm is not None, "Expected multimodal_output to be non-None" + audio = mm.get("audio") + sr = mm.get("sr") + assert audio is not None, "Expected 'audio' key in multimodal_output" + assert isinstance(audio, torch.Tensor), f"audio should be Tensor, got {type(audio)}" + return audio.cpu(), int(sr.item()) if sr is not None else SAMPLE_RATE raise AssertionError("No stage outputs received") +@pytest.mark.advanced_model +@pytest.mark.omni @hardware_test(res={"cuda": "L4"}) def test_moss_tts_nano_english(omni_runner: OmniRunner, ref_audio_path) -> None: """English TTS produces non-empty 48 kHz stereo audio.""" @@ -129,6 +132,8 @@ def test_moss_tts_nano_english(omni_runner: OmniRunner, ref_audio_path) -> None: assert not torch.all(audio == 0), "Audio should not be all-zeros (silence)" +@pytest.mark.advanced_model +@pytest.mark.omni @hardware_test(res={"cuda": "L4"}) def test_moss_tts_nano_chinese(omni_runner: OmniRunner, ref_audio_path) -> None: """Chinese TTS produces non-empty audio.""" @@ -140,6 +145,8 @@ def test_moss_tts_nano_chinese(omni_runner: OmniRunner, ref_audio_path) -> None: assert not torch.all(audio == 0) +@pytest.mark.advanced_model +@pytest.mark.omni @hardware_test(res={"cuda": "L4"}) def test_moss_tts_nano_deterministic(omni_runner: OmniRunner, ref_audio_path) -> None: """Same seed produces identical waveforms.""" @@ -151,6 +158,8 @@ def test_moss_tts_nano_deterministic(omni_runner: OmniRunner, ref_audio_path) -> assert torch.allclose(audio1, audio2, atol=1e-4), "Waveforms should match with same seed" +@pytest.mark.advanced_model +@pytest.mark.omni @hardware_test(res={"cuda": "L4"}) def test_moss_tts_nano_batch(omni_runner: OmniRunner, ref_audio_path) -> None: """Batch of two requests returns audio for each.""" @@ -159,11 +168,13 @@ def test_moss_tts_nano_batch(omni_runner: OmniRunner, ref_audio_path) -> None: _build_request("Second request.", ref_audio_path), ] results = [] - # Single-stage pipeline: one SamplingParams object is broadcast to all prompts. - for omni_out in omni_runner.omni.generate(requests, DEFAULT_SAMPLING): - mm = omni_out.multimodal_output - assert mm is not None - results.append(mm["audio"].cpu()) + # Single-stage model (num_stages=1): one sampling param for all requests. + for stage_outputs in omni_runner.omni.generate(requests, [DEFAULT_SAMPLING]): + req_output = stage_outputs.request_output + if req_output is not None: + mm = req_output.outputs[0].multimodal_output + assert mm is not None + results.append(mm["audio"].cpu()) assert len(results) == 2, f"Expected 2 outputs, got {len(results)}" for i, audio in enumerate(results): diff --git a/tests/e2e/offline_inference/test_stable_audio_expansion.py b/tests/e2e/offline_inference/test_stable_audio_expansion.py index 24ea78721..96365c7d3 100644 --- a/tests/e2e/offline_inference/test_stable_audio_expansion.py +++ b/tests/e2e/offline_inference/test_stable_audio_expansion.py @@ -15,6 +15,7 @@ import numpy as np import pytest import torch +from tests.helpers import skip_if_gated_repo_inaccessible from tests.helpers.assertions import assert_audio_valid from tests.helpers.mark import hardware_test from vllm_omni import Omni @@ -24,6 +25,8 @@ from vllm_omni.platforms import current_omni_platform pytestmark = [pytest.mark.full_model, pytest.mark.diffusion] +_MODEL_REPO = "stabilityai/stable-audio-open-1.0" + _SAMPLE_RATE = 44100 _CLIP_DURATION_S = 2.0 @@ -77,6 +80,7 @@ def test_stable_audio_quantization_and_teacache() -> None: CI should provide ``HF_TOKEN`` if the checkpoint is gated. """ + skip_if_gated_repo_inaccessible(_MODEL_REPO) # ``model_class_name`` must be passed explicitly: the default-stage-cfg # factory in ``async_omni_engine.py`` reads it out of ``kwargs`` when # deciding ``final_output_type`` (#2077), and at construction time the diff --git a/tests/engine/test_orchestrator.py b/tests/engine/test_orchestrator.py index f0d45a1e6..c84256f04 100644 --- a/tests/engine/test_orchestrator.py +++ b/tests/engine/test_orchestrator.py @@ -79,7 +79,7 @@ class FakeStageClient: self._diffusion_outputs = queue.Queue() # Orchestrator-facing interface. - async def add_request_async(self, *args, **_kwargs) -> None: + async def add_request_async(self, *args, **kwargs) -> None: self.add_request_calls.append(args) async def get_output_async(self): @@ -448,7 +448,7 @@ def test_stage_engine_core_client_shutdown_cleans_children_if_base_shutdown_fail assert recursive return [fake_child] - def fail_base_shutdown(self): + def fail_base_shutdown(self, **kwargs): raise RuntimeError("base shutdown failed") monkeypatch.setattr(psutil, "Process", FakePsutilProcess) @@ -480,7 +480,7 @@ def test_stage_engine_core_client_shutdown_kills_stubborn_children(monkeypatch): monkeypatch.setattr(psutil, "Process", FakePsutilProcess) monkeypatch.setattr(psutil, "wait_procs", lambda procs, timeout: ([], list(procs))) - monkeypatch.setattr(AsyncMPClient, "shutdown", lambda self: None) + monkeypatch.setattr(AsyncMPClient, "shutdown", lambda self, **kwargs: None) client = object.__new__(StageEngineCoreClient) client._proc = fake_proc @@ -994,7 +994,7 @@ async def test_handle_streaming_update_passes_prompt_text_to_stage_pool() -> Non @pytest.mark.asyncio async def test_stage_pool_submit_initial_rolls_back_output_processor_when_client_submit_fails() -> None: class FailingStageClient(FakeStageClient): - async def add_request_async(self, *args, **_kwargs) -> None: + async def add_request_async(self, *args, **kwargs) -> None: raise RuntimeError("submit failed") class TrackingOutputProcessor(FakeOutputProcessor): diff --git a/tests/entrypoints/openai_api/test_serving_audio_generate.py b/tests/entrypoints/openai_api/test_serving_audio_generate.py index fcf22bcef..ffb23bf4f 100644 --- a/tests/entrypoints/openai_api/test_serving_audio_generate.py +++ b/tests/entrypoints/openai_api/test_serving_audio_generate.py @@ -389,7 +389,7 @@ class TestErrorHandling: async def empty_gen(*args, **kwargs): return - yield # noqa: unreachable – makes this an async generator + yield # unreachable – makes this an async generator engine.generate = MagicMock(side_effect=empty_gen) server = _make_server(engine_client=engine) @@ -448,7 +448,7 @@ class TestErrorHandling: async def gen_value_error(*args, **kwargs): raise ValueError("bad value") - yield # noqa: unreachable + yield # unreachable engine.generate = MagicMock(side_effect=gen_value_error) server = _make_server(engine_client=engine) @@ -463,7 +463,7 @@ class TestErrorHandling: async def gen_runtime_error(*args, **kwargs): raise RuntimeError("something went wrong") - yield # noqa: unreachable + yield # unreachable engine.generate = MagicMock(side_effect=gen_runtime_error) server = _make_server(engine_client=engine) diff --git a/tests/entrypoints/openai_api/test_serving_chat_speaker.py b/tests/entrypoints/openai_api/test_serving_chat_speaker.py index 97c05e45b..7d6119b6f 100644 --- a/tests/entrypoints/openai_api/test_serving_chat_speaker.py +++ b/tests/entrypoints/openai_api/test_serving_chat_speaker.py @@ -22,6 +22,7 @@ def serving_chat(): instance = object.__new__(OmniOpenAIServingChat) instance._supported_speakers = None + instance.has_kv_connector = False return instance diff --git a/tests/helpers/__init__.py b/tests/helpers/__init__.py index a3348b07f..04d51fb90 100644 --- a/tests/helpers/__init__.py +++ b/tests/helpers/__init__.py @@ -6,3 +6,32 @@ refactor only inside the old monolithic ``conftest``; a greedy ``__init__`` changes import order and can affect in-process Omni (``OmniRunner`` / offline e2e) vs subprocess-based ``OmniServer`` tests. """ + +import pytest + + +def skip_if_gated_repo_inaccessible(repo_id: str) -> None: + """Skip the test if a gated HuggingFace repo is not accessible. + + Tries to download the model's config.json via ``hf_hub_download``, + which performs an actual file-access check (unlike ``HfApi().model_info()`` + that only checks metadata). If the token has metadata access but not + file-download access, ``hf_hub_download`` will raise ``GatedRepoError`` + and we skip cleanly. + """ + try: + from huggingface_hub import hf_hub_download + from huggingface_hub.errors import GatedRepoError, RepositoryNotFoundError + except Exception: + return + try: + hf_hub_download(repo_id=repo_id, filename="config.json") + except GatedRepoError as exc: + pytest.skip( + f"Skipping: gated HF repo {repo_id!r} inaccessible to the current " + f"HF_TOKEN ({exc}). See docs/contributing/ci/hf_credentials.md." + ) + except RepositoryNotFoundError as exc: + pytest.skip(f"Skipping: HF repo {repo_id!r} not found ({exc}).") + except Exception: + return diff --git a/tests/model_executor/models/qwen3_tts/test_code_predictor_dtype.py b/tests/model_executor/models/qwen3_tts/test_code_predictor_dtype.py index 3c76f4210..12de38e23 100644 --- a/tests/model_executor/models/qwen3_tts/test_code_predictor_dtype.py +++ b/tests/model_executor/models/qwen3_tts/test_code_predictor_dtype.py @@ -48,6 +48,7 @@ def _build_mock_modules(mocker: MockerFixture) -> dict[str, object]: """Build the dict of modules to inject into sys.modules.""" platforms_mock = mocker.MagicMock() platforms_mock.current_omni_platform.supports_torch_inductor.return_value = False + platforms_mock.current_omni_platform.is_npu.return_value = False logger_mock = mocker.MagicMock() logger_mock.init_logger = lambda name: mocker.MagicMock() @@ -211,6 +212,10 @@ class TestCodePredictorDtypeAlignment: predictor._model_dtype = torch.float16 predictor._compiled_model_fwd = predictor.model.forward + # Ensure NPU path is not taken on non-NPU hardware + common_mod = sys.modules["vllm_omni.model_executor.models.common.qwen3_code_predictor"] + mocker.patch.object(common_mod.current_omni_platform, "is_npu", return_value=False) + # _warmup_buckets should fix the dtype mismatch predictor._warmup_buckets() @@ -230,6 +235,9 @@ class TestCodePredictorDtypeAlignment: predictor = predictor.to(torch.float16) assert predictor._model_dtype is None + # Ensure NPU path is not taken on non-NPU hardware + common_mod = sys.modules["vllm_omni.model_executor.models.common.qwen3_code_predictor"] + mocker.patch.object(common_mod.current_omni_platform, "is_npu", return_value=False) predictor._setup_compile() assert predictor._model_dtype == torch.float16 diff --git a/tests/worker/test_gpu_generation_model_runner.py b/tests/worker/test_gpu_generation_model_runner.py index c8637c8bd..15880d7a9 100644 --- a/tests/worker/test_gpu_generation_model_runner.py +++ b/tests/worker/test_gpu_generation_model_runner.py @@ -35,6 +35,7 @@ def _make_runner(multimodal_outputs): runner.device = torch.device("cpu") runner.supports_mm_inputs = False runner.speculative_config = None + runner.routed_experts_initialized = False return runner diff --git a/vllm_omni/core/sched/omni_ar_scheduler.py b/vllm_omni/core/sched/omni_ar_scheduler.py index 0a7fd431c..61ee24547 100644 --- a/vllm_omni/core/sched/omni_ar_scheduler.py +++ b/vllm_omni/core/sched/omni_ar_scheduler.py @@ -203,6 +203,14 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): return False def schedule(self) -> SchedulerOutput: # type: ignore[override] + # Remove FINISHED_ABORTED requests before the upstream scheduler sees + # them. Upstream vllm raises RuntimeError on this status; omni allows + # async abort (e.g. client disconnect during TTS streaming) to leave + # requests in the waiting/running queues temporarily. + for queue in (self.waiting, self.running): + for req in list(queue): + if getattr(req, "status", None) == RequestStatus.FINISHED_ABORTED: + queue.remove(req) connector_output = self._latest_omni_connector_output self._latest_omni_connector_output = None if self.input_coordinator: @@ -250,6 +258,7 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): lora_request=nr.lora_request, # Enrich with omni payloads from the live request object prompt_embeds=(getattr(request, "prompt_embeds", None) if request else None), + prompt_is_token_ids=nr.prompt_is_token_ids, additional_information=(getattr(request, "additional_information", None) if request else None), ) new_list.append(omni_nr) @@ -419,7 +428,11 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): stopped = True if stopped: - routed_experts = self._get_routed_experts(request) + if ( + model_runner_output.routed_experts_dict is not None + and req_id in model_runner_output.routed_experts_dict + ): + routed_experts = model_runner_output.routed_experts_dict[req_id] # Capture finish_reason BEFORE _handle_stopped_request, which may # reset the status to WAITING for streaming requests that continue. diff --git a/vllm_omni/core/sched/omni_generation_scheduler.py b/vllm_omni/core/sched/omni_generation_scheduler.py index 5fc20d2da..e290d3ef6 100644 --- a/vllm_omni/core/sched/omni_generation_scheduler.py +++ b/vllm_omni/core/sched/omni_generation_scheduler.py @@ -337,6 +337,7 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler): lora_request=nr.lora_request, # Enrich with omni payloads from the live request object prompt_embeds=(getattr(request, "prompt_embeds", None) if request else None), + prompt_is_token_ids=nr.prompt_is_token_ids, additional_information=(getattr(request, "additional_information", None) if request else None), ) new_list.append(omni_nr) @@ -520,7 +521,11 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler): stopped = True if stopped: - routed_experts = self._get_routed_experts(request) + if ( + model_runner_output.routed_experts_dict is not None + and req_id in model_runner_output.routed_experts_dict + ): + routed_experts = model_runner_output.routed_experts_dict[req_id] finish_reason = request.get_finished_reason() finished = self._handle_stopped_request(request) if finished: diff --git a/vllm_omni/core/sched/output.py b/vllm_omni/core/sched/output.py index c016cbd81..800881d9f 100644 --- a/vllm_omni/core/sched/output.py +++ b/vllm_omni/core/sched/output.py @@ -54,6 +54,7 @@ class OmniNewRequestData(NewRequestData): num_computed_tokens=request.num_computed_tokens, lora_request=request.lora_request, prompt_embeds=getattr(request, "prompt_embeds", None), + prompt_is_token_ids=getattr(request, "prompt_is_token_ids", None), prefill_token_ids=prefill_token_ids, additional_information=getattr(request, "additional_information", None), ) diff --git a/vllm_omni/deploy/mimo_audio.yaml b/vllm_omni/deploy/mimo_audio.yaml index a92e905f7..1e4fd7505 100644 --- a/vllm_omni/deploy/mimo_audio.yaml +++ b/vllm_omni/deploy/mimo_audio.yaml @@ -23,7 +23,7 @@ connectors: stages: - stage_id: 0 max_num_seqs: 1 - gpu_memory_utilization: 0.3 + gpu_memory_utilization: 0.5 enforce_eager: true trust_remote_code: true enable_prefix_caching: false @@ -42,7 +42,7 @@ stages: - stage_id: 1 max_num_seqs: 1 - gpu_memory_utilization: 0.2 + gpu_memory_utilization: 0.3 enforce_eager: true trust_remote_code: true enable_prefix_caching: false diff --git a/vllm_omni/diffusion/stage_diffusion_client.py b/vllm_omni/diffusion/stage_diffusion_client.py index 11809f6ef..476b44b5b 100644 --- a/vllm_omni/diffusion/stage_diffusion_client.py +++ b/vllm_omni/diffusion/stage_diffusion_client.py @@ -391,6 +391,8 @@ class StageDiffusionClient(StageClientBase): return self._output_queue.get_nowait() except asyncio.QueueEmpty: if self._engine_dead: + if self._shutting_down: + return None raise EngineDeadError() if not self._shutting_down and self._owns_process and self._proc is not None and not self._proc.is_alive(): self._engine_dead = True diff --git a/vllm_omni/diffusion/worker/diffusion_worker.py b/vllm_omni/diffusion/worker/diffusion_worker.py index 3909fd2d4..aa0c99cec 100644 --- a/vllm_omni/diffusion/worker/diffusion_worker.py +++ b/vllm_omni/diffusion/worker/diffusion_worker.py @@ -97,7 +97,7 @@ def _make_diffusion_vllm_model_config(od_config: OmniDiffusionConfig) -> _Diffus @contextmanager def _force_cutlass_fp8_linear_kernel(quant_config: object | None) -> Iterator[None]: - from vllm.model_executor.layers.quantization import modelopt as vllm_modelopt + import vllm.model_executor.layers.quantization.modelopt as vllm_modelopt linear_method_cls = getattr(quant_config, "LinearMethodCls", None) if linear_method_cls in { diff --git a/vllm_omni/engine/__init__.py b/vllm_omni/engine/__init__.py index 55d119ad6..aa08fcde6 100644 --- a/vllm_omni/engine/__init__.py +++ b/vllm_omni/engine/__init__.py @@ -94,6 +94,7 @@ class OmniEngineCoreRequest(EngineCoreRequest): return cls( request_id=request.request_id, prompt_token_ids=request.prompt_token_ids, + prompt_is_token_ids=request.prompt_is_token_ids, mm_features=request.mm_features, sampling_params=request.sampling_params, pooling_params=request.pooling_params, @@ -109,6 +110,8 @@ class OmniEngineCoreRequest(EngineCoreRequest): resumable=request.resumable, external_req_id=request.external_req_id, reasoning_ended=request.reasoning_ended, + reasoning_parser_kwargs=request.reasoning_parser_kwargs, + abort_immediately=request.abort_immediately, additional_information=additional_information, ) diff --git a/vllm_omni/engine/orchestrator.py b/vllm_omni/engine/orchestrator.py index 4bc032ae2..ea615d579 100644 --- a/vllm_omni/engine/orchestrator.py +++ b/vllm_omni/engine/orchestrator.py @@ -32,7 +32,6 @@ from vllm_omni.engine.messages import ( EngineQueueMessage, ErrorMessage, OutputMessage, - ShutdownRequestMessage, StageMetricsMessage, StageSubmissionMessage, ) @@ -128,8 +127,8 @@ class Orchestrator: def __init__( self, request_async_queue: janus.AsyncQueue[EngineQueueMessage], - output_async_queue: janus.AsyncQueue[EngineQueueMessage], - rpc_async_queue: janus.AsyncQueue[EngineQueueMessage], + output_async_queue: janus.AsyncQueue[dict[str, Any]], + rpc_async_queue: janus.AsyncQueue[dict[str, Any]], stage_pools: list[StagePool], *, async_chunk: bool = False, @@ -174,6 +173,13 @@ class Orchestrator: await asyncio.gather(request_task, output_task) except asyncio.CancelledError: raise + except EngineDeadError as e: + # EngineDeadError from _orchestration_loop means the diffusion + # engine died. All pending requests were already notified and + # _shutdown_event was already set by the loop's handler. + # During teardown this is expected; the finally block handles + # proper cleanup. Do not re-raise. + logger.info("[Orchestrator] Engine dead during shutdown: %s", e) except Exception: logger.exception("[Orchestrator] Fatal error in orchestrator tasks") raise @@ -208,27 +214,32 @@ class Orchestrator: """Read messages from the main thread via request_async_queue.""" while True: msg = await self.request_async_queue.get() + msg_type = msg.type - if isinstance(msg, StageSubmissionMessage): - if msg.type == "add_request": - await self._handle_add_request(msg) - elif msg.type == "streaming_update": - await self._handle_streaming_update(msg) - else: - logger.warning("[Orchestrator] Unknown stage submission type: %s", msg.type) - elif isinstance(msg, AddCompanionRequestMessage): + if msg_type == "add_request": + await self._handle_add_request(msg) + elif msg_type == "streaming_update": + await self._handle_streaming_update(msg) + elif msg_type == "add_companion_request": await self._handle_add_companion(msg) - elif isinstance(msg, AbortRequestMessage): + elif msg_type == "abort": await self._handle_abort(msg) - elif isinstance(msg, CollectiveRPCRequestMessage): + elif msg_type == "collective_rpc": await self._handle_collective_rpc(msg) - elif isinstance(msg, ShutdownRequestMessage): + elif msg_type == "shutdown": logger.info("[Orchestrator] Received shutdown signal") self._shutdown_event.set() + # Pre-mark stage clients as shutting down to prevent + # proc_monitor daemon threads from flagging normal + # process exit as EngineDeadError during teardown. + for pool in self.stage_pools: + for client in pool.clients: + if hasattr(client, "_shutting_down"): + client._shutting_down = True self._shutdown_stages() break else: - logger.warning("[Orchestrator] Unknown message type: %s", type(msg).__name__) + logger.warning("[Orchestrator] Unknown message type: %s", msg_type) async def _handle_add_request(self, msg: StageSubmissionMessage) -> None: """Handle an add_request message from the main thread.""" @@ -387,8 +398,8 @@ class Orchestrator: rpc_id = msg.rpc_id method = msg.method timeout = msg.timeout - args = msg.args - kwargs = dict(msg.kwargs) + args = tuple(msg.args) + kwargs = dict(msg.kwargs or {}) requested_stage_ids = msg.stage_ids target_pools: list[StagePool] = [] @@ -793,120 +804,78 @@ class Orchestrator: """Forward output from the current logical stage to the next one.""" next_logical = src_stage_id + 1 next_pool = self.stage_pools[next_logical] - - if next_pool.stage_type == "diffusion": - await self._forward_to_diffusion_stage( - req_id, - src_stage_id, - output, - req_state, - ) - else: - await self._forward_to_llm_stage( - req_id, - src_stage_id, - output, - req_state, - is_streaming_session=is_streaming_session, - is_final_update=is_final_update, - ) - - async def _forward_to_diffusion_stage( - self, - req_id: str, - src_stage_id: int, - output: Any, - req_state: OrchestratorRequestState, - ) -> None: - """Forward output from an LLM stage to the next diffusion stage.""" - next_logical = src_stage_id + 1 - next_pool = self.stage_pools[next_logical] - params = req_state.sampling_params_list[next_logical] - source_outputs = [output] - already_submitted = self._next_stage_already_submitted(src_stage_id, req_state) - - diffusion_client = next_pool.stage_client - if diffusion_client.custom_process_input_func is not None: - _t_ar2d = _time.perf_counter() - diffusion_prompt = diffusion_client.custom_process_input_func( - source_outputs, - req_state.prompt, - getattr(diffusion_client, "requires_multimodal_data", False), - ) - _dt_ar2d = (_time.perf_counter() - _t_ar2d) * 1000 - req_state.pipeline_timings["ar2diffusion_ms"] = _dt_ar2d - logger.info( - "[Orchestrator] ar2diffusion req=%s wall_time=%.3fms stage=%d->%d", - req_id, - _dt_ar2d, - src_stage_id, - next_logical, - ) - if isinstance(diffusion_prompt, list): - if not diffusion_prompt: - error_output = OmniRequestOutput.from_error( - req_id, - f"Stage-{src_stage_id} produced no valid inputs for diffusion stage-{next_logical}", - ) - logger.warning( - "[Orchestrator] req=%s stage=%d produced empty diffusion inputs for stage=%d; " - "routing terminal error output", - req_id, - src_stage_id, - next_logical, - ) - await self.output_async_queue.put( - OutputMessage( - request_id=req_id, - stage_id=next_logical, - engine_outputs=error_output, - metrics=None, - finished=True, - ) - ) - await self._cleanup_request_ids( - [req_id, *self._cfg_tracker.cleanup_parent(req_id)], - ) - return - if already_submitted and len(diffusion_prompt) == 1: - diffusion_prompt = diffusion_prompt[0] - else: - diffusion_prompt = req_state.prompt - - if already_submitted: - await next_pool.submit_update(req_id, req_state, diffusion_prompt) - else: - await next_pool.submit_initial( - req_id, - req_state, - diffusion_prompt, - submit_kwargs={ - "kv_sender_info": self._build_kv_sender_info( - list(getattr(diffusion_client, "engine_input_source", None) or [src_stage_id]), - request_id=req_id, - ) - }, - params_override=self._maybe_clone_diffusion_params_for_cfg(req_id, params), - ) - req_state.stage_submit_ts[next_logical] = _time.time() - - async def _forward_to_llm_stage( - self, - req_id: str, - src_stage_id: int, - output: Any, - req_state: OrchestratorRequestState, - *, - is_streaming_session: bool = False, - is_final_update: bool = False, - ) -> None: - """Forward output from the current stage to the next LLM stage.""" - next_logical = src_stage_id + 1 - next_pool = self.stage_pools[next_logical] + next_client = next_pool.stage_client params = req_state.sampling_params_list[next_logical] source_outputs = [output] next_stage_resumable = is_streaming_session and not is_final_update already_submitted = self._next_stage_already_submitted(src_stage_id, req_state) + requires_multimodal_data = getattr(next_client, "requires_multimodal_data", False) + + if next_pool.stage_type == "diffusion": + if next_client.custom_process_input_func is not None: + _t_ar2d = _time.perf_counter() + diffusion_prompt = next_client.custom_process_input_func( + source_outputs, + req_state.prompt, + requires_multimodal_data, + ) + _dt_ar2d = (_time.perf_counter() - _t_ar2d) * 1000 + req_state.pipeline_timings["ar2diffusion_ms"] = _dt_ar2d + logger.info( + "[Orchestrator] ar2diffusion req=%s wall_time=%.3fms stage=%d->%d", + req_id, + _dt_ar2d, + src_stage_id, + next_logical, + ) + if isinstance(diffusion_prompt, list): + if not diffusion_prompt: + error_output = OmniRequestOutput.from_error( + req_id, + f"Stage-{src_stage_id} produced no valid inputs for diffusion stage-{next_logical}", + ) + logger.warning( + "[Orchestrator] req=%s stage=%d produced empty diffusion inputs for stage=%d; " + "routing terminal error output", + req_id, + src_stage_id, + next_logical, + ) + await self.output_async_queue.put( + OutputMessage( + request_id=req_id, + stage_id=next_logical, + engine_outputs=error_output, + metrics=None, + finished=True, + ) + ) + await self._cleanup_request_ids( + [req_id, *self._cfg_tracker.cleanup_parent(req_id)], + ) + return + if already_submitted and len(diffusion_prompt) == 1: + diffusion_prompt = diffusion_prompt[0] + else: + diffusion_prompt = req_state.prompt + + if already_submitted: + await next_pool.submit_update(req_id, req_state, diffusion_prompt) + else: + await next_pool.submit_initial( + req_id, + req_state, + diffusion_prompt, + submit_kwargs={ + "kv_sender_info": self._build_kv_sender_info( + list(getattr(next_client, "engine_input_source", None) or [src_stage_id]), + request_id=req_id, + ) + }, + params_override=self._maybe_clone_diffusion_params_for_cfg(req_id, params), + ) + req_state.stage_submit_ts[next_logical] = _time.time() + return # PD disaggregation: prefill → decode routing uses original prompt + KV transfer params if self._pd_pair is not None and (src_stage_id, next_logical) == self._pd_pair: @@ -953,7 +922,6 @@ class Orchestrator: {}, )[req_id] = req_state.pd_prefill_multimodal_output - next_client = next_pool.llm_stage_client try: next_inputs = next_client.process_engine_inputs( source_outputs, @@ -972,7 +940,7 @@ class Orchestrator: for next_input in next_inputs: # Only AR thinker stages consume encoder mm_features; downstream # (talker/code2wav/…) must not see them (avoids encoder-cache misses). - model_stage = next_client.model_stage + model_stage = getattr(next_client, "model_stage", None) mm_features = req_state.mm_features if model_stage == "thinker" else None request = build_engine_core_request_from_tokens( request_id=req_id, @@ -1022,7 +990,7 @@ class Orchestrator: req_state.prompt, submit_kwargs={ "kv_sender_info": self._build_kv_sender_info( - list(next_pool.stage_client.engine_input_source or [next_stage_id - 1]), + list(getattr(next_pool.stage_client, "engine_input_source", None) or [next_stage_id - 1]), request_id=request_id, ) }, @@ -1069,14 +1037,14 @@ class Orchestrator: continue sender_pool = self.stage_pools[sender_stage_id] - if sender_pool.stage_type == "diffusion": + sender_stage = sender_pool.get_bound_client(request_id) if request_id is not None else None + if sender_stage is None: + sender_stage = sender_pool.stage_client + get_sender_info = getattr(sender_stage, "get_kv_sender_info", None) + if not callable(get_sender_info): continue - sender_stage = sender_pool.get_bound_llm_client(request_id) if request_id is not None else None - if sender_stage is None: - sender_stage = sender_pool.llm_stage_client - - sender_info = sender_stage.get_kv_sender_info() + sender_info = get_sender_info() if not sender_info: logger.warning( "[Orchestrator] Stage-%s has no KV sender info available", @@ -1110,7 +1078,7 @@ class Orchestrator: msg = self.request_async_queue.get_nowait() except Exception: break - if isinstance(msg, StageSubmissionMessage) and msg.type == "add_request": + if msg.type == "add_request": req_id = msg.request_id await self.output_async_queue.put( ErrorMessage( diff --git a/vllm_omni/engine/output_processor.py b/vllm_omni/engine/output_processor.py index 6689ac05b..7f70037f8 100644 --- a/vllm_omni/engine/output_processor.py +++ b/vllm_omni/engine/output_processor.py @@ -3,6 +3,9 @@ from typing import Any import numpy as np import torch from vllm.logger import init_logger +from vllm.model_executor.layers.fused_moe.routed_experts_capturer import ( + split_routed_experts, +) from vllm.outputs import PoolingRequestOutput from vllm.sampling_params import RequestOutputKind from vllm.tokenizers import TokenizerLike @@ -230,7 +233,17 @@ class OmniRequestState(RequestState): self.sent_tokens_offset = self.detokenizer.num_output_tokens() external_req_id = self.external_req_id - output = self._new_completion_output(new_token_ids, finish_reason, stop_reason, routed_experts) + + # Split routing data into prompt and generation portions, matching + # upstream make_request_output behaviour. + prompt_routed_experts = None + gen_routed_experts = None + if routed_experts is not None: + prompt_len = len(self.prompt_token_ids) if self.prompt_token_ids else 0 + num_gen = self.detokenizer.num_output_tokens() if self.detokenizer is not None else None + prompt_routed_experts, gen_routed_experts = split_routed_experts(routed_experts, prompt_len, num_gen) + + output = self._new_completion_output(new_token_ids, finish_reason, stop_reason, gen_routed_experts) if self.parent_req is None: outputs = [output] @@ -240,7 +253,13 @@ class OmniRequestState(RequestState): return None external_req_id = self.parent_req.external_req_id - return self._new_request_output(external_req_id, outputs, finished, kv_transfer_params) + return self._new_request_output( + external_req_id, + outputs, + finished, + kv_transfer_params, + prompt_routed_experts, + ) def _new_completion_output( self, diff --git a/vllm_omni/engine/stage_engine_core_client.py b/vllm_omni/engine/stage_engine_core_client.py index f2d4779a9..25e4294e5 100644 --- a/vllm_omni/engine/stage_engine_core_client.py +++ b/vllm_omni/engine/stage_engine_core_client.py @@ -80,7 +80,8 @@ class StageEngineCoreClientBase(StageClientBase): def make_async_mp_client( vllm_config: Any, executor_class: type, - metadata: StageMetadata, + log_stats: bool = False, + metadata: StageMetadata | None = None, client_addresses: dict[str, str] | None = None, proc: Any = None, engine_manager: Any = None, @@ -93,6 +94,7 @@ class StageEngineCoreClientBase(StageClientBase): client_args = dict( vllm_config=vllm_config, executor_class=executor_class, + log_stats=log_stats, metadata=metadata, client_addresses=client_addresses, proc=proc, @@ -441,7 +443,7 @@ class StageEngineCoreClientBase(StageClientBase): kwargs=kwargs, ) - def shutdown(self) -> None: + def shutdown(self, timeout: float | None = None) -> None: """Shutdown managed resources and any externally spawned subprocess.""" child_procs: list[psutil.Process] = [] if self._proc is not None and self._proc.pid is not None: @@ -451,7 +453,7 @@ class StageEngineCoreClientBase(StageClientBase): child_procs = [] try: - super().shutdown() + super().shutdown(timeout=timeout) finally: if self._proc is not None and self._proc.is_alive(): self._proc.terminate() diff --git a/vllm_omni/entrypoints/async_omni.py b/vllm_omni/entrypoints/async_omni.py index c22e78096..3485f98f7 100644 --- a/vllm_omni/entrypoints/async_omni.py +++ b/vllm_omni/entrypoints/async_omni.py @@ -232,6 +232,7 @@ class AsyncOmni(EngineClient, OmniBase): priority: int = 0, data_parallel_rank: int | None = None, reasoning_ended: bool | None = None, + reasoning_parser_kwargs: dict[str, Any] | None = None, ) -> AsyncGenerator[OmniRequestOutput, None]: """Generate outputs for the given prompt(s) asynchronously. @@ -938,6 +939,37 @@ class AsyncOmni(EngineClient, OmniBase): """Check if tracing is enabled.""" return False + async def notify_kv_transfer_request_rejected( + self, + request_id: str, + kv_transfer_params: dict[str, Any], + *, + data_parallel_rank: int | None = None, + ) -> None: + """Notify engine that a KV-transfer request was rejected before admission. + + Omni does not currently use KV-transfer pre-admission resources, + so this is a no-op. + """ + logger.debug( + "KV-transfer request rejected (no-op in omni): request_id=%s", + request_id, + ) + + async def start_weight_update(self, is_checkpoint_format: bool = True) -> None: + """Start a new weight update. + + Omni does not currently support weight transfer, so this is a no-op. + """ + logger.debug("Weight update start requested (no-op in omni)") + + async def finish_weight_update(self) -> None: + """Finish the current weight update. + + Omni does not currently support weight transfer, so this is a no-op. + """ + logger.debug("Weight update finish requested (no-op in omni)") + async def do_log_stats(self) -> None: """Log statistics. diff --git a/vllm_omni/entrypoints/openai/serving_chat.py b/vllm_omni/entrypoints/openai/serving_chat.py index 2c375fa29..3296c0b4d 100644 --- a/vllm_omni/entrypoints/openai/serving_chat.py +++ b/vllm_omni/entrypoints/openai/serving_chat.py @@ -148,6 +148,8 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): instance._diffusion_mode = True instance._diffusion_engine = diffusion_engine instance._diffusion_model_name = model_name + instance.engine_client = None + instance.has_kv_connector = False return instance def _get_diffusion_extra_body_params(self) -> frozenset[str]: @@ -216,6 +218,15 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): For diffusion models, this generates images and returns them in a chat completion response format. """ + return await self._with_kv_transfer_rejection_cleanup( + self._create_chat_completion(request, raw_request), request, raw_request + ) + + async def _create_chat_completion( + self, + request: ChatCompletionRequest, + raw_request: Request | None = None, + ) -> AsyncGenerator[str, None] | ChatCompletionResponse | ErrorResponse: # Handle diffusion mode if self._diffusion_mode: return await self._create_diffusion_chat_completion(request, raw_request) @@ -996,6 +1007,9 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): # when handling audio/image responses role = self.get_chat_request_role(request) + # Compute prompt_text once at first iteration (upstream #42052) + prompt_text = getattr(res, "prompt", None) if getattr(request, "return_prompt_text", None) else None + # We need to do it here, because if there are exceptions in # the result_generator, it needs to be sent as the FIRST # response (by the try...catch). @@ -1022,6 +1036,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): choices=[choice_data], model=model_name, prompt_token_ids=(res.prompt_token_ids if request.return_token_ids else None), + prompt_text=prompt_text, modality=final_output_type, ) @@ -1740,6 +1755,13 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): response_metrics.update(extra) choices.extend(choices_data) + # Compute prompt_text for non-streaming response (upstream #42052) + prompt_text = ( + getattr(final_res, "prompt", None) + if final_res is not None and getattr(request, "return_prompt_text", None) + else None + ) + response = OmniChatCompletionResponse( id=request_id, created=created_time, @@ -1748,6 +1770,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): usage=usage, prompt_logprobs=prompt_logprobs, prompt_token_ids=prompt_token_ids, + prompt_text=prompt_text, kv_transfer_params=kv_transfer_params, metrics=response_metrics, ) @@ -2034,13 +2057,17 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): mm_output = final_res.outputs[0].multimodal_output audio_data = mm_output.get("audio") if isinstance(audio_data, list): - if stream: + if not audio_data: + audio_tensor = None + elif stream: audio_tensor = audio_data[-1] else: audio_tensor = torch.cat(audio_data, dim=-1) else: audio_tensor = audio_data - audio_tensor = audio_tensor.float().detach().cpu().numpy() + if audio_tensor is None: + return self._create_error_response("Audio generation completed but no audio was produced.") + audio_tensor = audio_tensor.detach().cpu().float().numpy() # Ensure audio is 1D (flatten if needed) if audio_tensor.ndim > 1: diff --git a/vllm_omni/entrypoints/openai/serving_speech.py b/vllm_omni/entrypoints/openai/serving_speech.py index 256b9a56b..e5bfb1329 100644 --- a/vllm_omni/entrypoints/openai/serving_speech.py +++ b/vllm_omni/entrypoints/openai/serving_speech.py @@ -2103,6 +2103,13 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): voice_lower = request.voice.lower() tts_params["voice_name"] = [voice_lower] tts_params["voice_created_at"] = [self._voice_created_at(voice_lower)] + # Propagate seed from sampling params for deterministic generation. + # MOSS-TTS-Nano uses its own internal sampling inside + # inference_stream(), which reads seed from additional_information + # (not SamplingParams). Without this the model uses ambient RNG + # state and produces non-deterministic output. + if sampling_params_list and getattr(sampling_params_list[0], "seed", None) is not None: + tts_params["seed"] = [sampling_params_list[0].seed] prompt = tokens_input(prompt_token_ids=[1]) prompt["additional_information"] = tts_params else: diff --git a/vllm_omni/model_executor/models/common/qwen3_code_predictor.py b/vllm_omni/model_executor/models/common/qwen3_code_predictor.py index d36068d46..b3074ed79 100644 --- a/vllm_omni/model_executor/models/common/qwen3_code_predictor.py +++ b/vllm_omni/model_executor/models/common/qwen3_code_predictor.py @@ -351,8 +351,10 @@ class CodePredictorBaseModel(nn.Module): # native bf16 support (Turing, Volta). The RMSNorm and RoPE # layers already upcast internally; this extends the same # treatment to attention and MLP. + # autocast to float32 is unsupported on CPU; skip fp32 upcast there + # (CPU uses full-precision intermediates internally). input_dtype = inputs_embeds.dtype - use_fp32 = input_dtype == torch.float16 + use_fp32 = input_dtype == torch.float16 and inputs_embeds.device.type != "cpu" if use_fp32: inputs_embeds = inputs_embeds.float() hidden_states = inputs_embeds diff --git a/vllm_omni/model_executor/models/common/snake_activation.py b/vllm_omni/model_executor/models/common/snake_activation.py index 0be93d74a..35bb3384a 100644 --- a/vllm_omni/model_executor/models/common/snake_activation.py +++ b/vllm_omni/model_executor/models/common/snake_activation.py @@ -61,8 +61,9 @@ class SnakeBeta(nn.Module): x = tl.load(x_ptr + bid * stride_b + cid * stride_c + t_off, mask=mask, other=0.0) ea = tl.load(exp_alpha_ptr + cid) ib = tl.load(inv_beta_ptr + cid) - sin_val = tl.sin(x * ea) - result = x + ib * sin_val * sin_val + x_float = x.to(tl.float32) + sin_val = tl.sin(x_float * ea) + result = x + (ib * sin_val * sin_val).to(x.dtype) tl.store(out_ptr + bid * stride_b + cid * stride_c + t_off, result, mask=mask) diff --git a/vllm_omni/model_executor/models/hunyuan_image3/hunyuan_image3.py b/vllm_omni/model_executor/models/hunyuan_image3/hunyuan_image3.py index a59b30e44..4596d1145 100644 --- a/vllm_omni/model_executor/models/hunyuan_image3/hunyuan_image3.py +++ b/vllm_omni/model_executor/models/hunyuan_image3/hunyuan_image3.py @@ -27,17 +27,8 @@ from vllm.distributed import ( ) from vllm.inputs import MultiModalDataDict from vllm.logger import init_logger +from vllm.model_executor.layers.fused_moe import FusedMoE as SharedFusedMoE from vllm.model_executor.layers.fused_moe import fused_moe_make_expert_params_mapping - -try: - from vllm.model_executor.layers.fused_moe.shared_fused_moe import SharedFusedMoE -except ImportError: - # PyPI vllm 0.20.x neither exports `SharedFusedMoE` from the package top-level - # nor ships a `shared_fused_moe.py` submodule. The functionality lives on - # `FusedMoE` directly (which gained a `shared_experts` parameter), so alias - # the symbol — call sites only use the classmethod `make_expert_params_mapping` - # and `__init__(shared_experts=..., ...)` which are present on `FusedMoE`. - from vllm.model_executor.layers.fused_moe import FusedMoE as SharedFusedMoE from vllm.model_executor.layers.linear import ( ColumnParallelLinear, ReplicatedLinear, diff --git a/vllm_omni/model_executor/models/moss_tts_nano/modeling_moss_tts_nano.py b/vllm_omni/model_executor/models/moss_tts_nano/modeling_moss_tts_nano.py index 01d2cf139..0ae1c63e3 100644 --- a/vllm_omni/model_executor/models/moss_tts_nano/modeling_moss_tts_nano.py +++ b/vllm_omni/model_executor/models/moss_tts_nano/modeling_moss_tts_nano.py @@ -160,7 +160,7 @@ class MossTTSNanoForGeneration(nn.Module): def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: with self._lock: if self._lm is not None: - return set() + return None _patch_torchaudio_load() try: @@ -233,7 +233,7 @@ class MossTTSNanoForGeneration(nn.Module): # iterator so the underlying stream is closed cleanly. for _ in weights: pass - return set() + return None # ------------------------------------------------------------------ # Dummy run support diff --git a/vllm_omni/model_executor/models/qwen3_tts/tokenizer_12hz/modeling_qwen3_tts_tokenizer_v2.py b/vllm_omni/model_executor/models/qwen3_tts/tokenizer_12hz/modeling_qwen3_tts_tokenizer_v2.py index e71dbc091..3cd8b75b3 100644 --- a/vllm_omni/model_executor/models/qwen3_tts/tokenizer_12hz/modeling_qwen3_tts_tokenizer_v2.py +++ b/vllm_omni/model_executor/models/qwen3_tts/tokenizer_12hz/modeling_qwen3_tts_tokenizer_v2.py @@ -534,6 +534,10 @@ class Qwen3TTSTokenizerV2DecoderTransformerModel(Qwen3TTSTokenizerV2DecoderPreTr cache_position=None, **kwargs, ) -> BaseModelOutputWithPast: + r""" + cache_position (`torch.LongTensor` of shape `(sequence_length)`, *optional*): + Indices depicting the position of the input sequence tokens in the sequence. Used to update the cache. + """ if input_ids is not None: raise ValueError("input_ids is not expected") if (input_ids is None) ^ (inputs_embeds is not None): diff --git a/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py b/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py index 84ac562cc..82d628ae9 100644 --- a/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py +++ b/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py @@ -16,7 +16,11 @@ from vllm.distributed.kv_transfer import get_kv_transfer_group, has_kv_transfer_ from vllm.distributed.parallel_state import get_pp_group, get_tp_group from vllm.forward_context import BatchDescriptor from vllm.logger import logger -from vllm.model_executor.layers.fused_moe.routed_experts_capturer import RoutedExpertsCapturer +from vllm.model_executor.layers.fused_moe.routed_experts_capturer import ( + extract_routed_experts_for_current_batch, + get_global_experts_capturer, + issue_routing_d2h_copy, +) from vllm.sequence import IntermediateTensors from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput from vllm.v1.outputs import ( @@ -28,8 +32,8 @@ from vllm.v1.outputs import ( from vllm.v1.spec_decode.metadata import SpecDecodeMetadata from vllm.v1.structured_output.utils import apply_grammar_bitmask from vllm.v1.utils import record_function_or_nullcontext -from vllm.v1.worker import mamba_utils from vllm.v1.worker.gpu_model_runner import AsyncGPUModelRunnerOutput, PerLayerAttnMetadata +from vllm.v1.worker.mamba_utils import preprocess_mamba from vllm.v1.worker.ubatch_utils import maybe_create_ubatch_slices from vllm_ascend.ascend_forward_context import set_ascend_forward_context from vllm_ascend.attention.utils import AscendCommonAttentionMetadata @@ -246,11 +250,9 @@ class NPUARModelRunner(OmniNPUModelRunner): intermediate_tensors: IntermediateTensors | None = None, ) -> OmniModelRunnerOutput | IntermediateTensors | None: if self.vllm_config.model_config.enable_return_routed_experts: - capturer = RoutedExpertsCapturer.get_instance() - if capturer is not None: - capturer.clear_buffer() - else: - logger.warning("RoutedExpertsCapturer is not initialized.") + capturer = get_global_experts_capturer() + if capturer is not None and hasattr(capturer, "finalize_pending_copy"): + capturer.finalize_pending_copy() if self.ascend_config.profiling_chunk_config.enabled: self._sync_device() self._execution_start_time = time.perf_counter() @@ -299,7 +301,9 @@ class NPUARModelRunner(OmniNPUModelRunner): # TODO(Ronald1995): deepcopy is expensive when there is a large # number of requests, optimize it later. if (( - self.use_async_scheduling and self.num_spec_tokens and self._draft_token_ids is None # type: ignore[has-type] + self.use_async_scheduling + and self.num_spec_tokens + and self._draft_token_ids is None # type: ignore[has-type] ) or ( # NOTE: This branch specifically triggers a deepcopy during the prefill phase # only for PCP (Parallel Context Processing) + Multi-Modal (MM) scenarios. @@ -455,7 +459,7 @@ class NPUARModelRunner(OmniNPUModelRunner): if deferred_state_corrections_fn: deferred_state_corrections_fn() deferred_state_corrections_fn = None - mamba_utils.preprocess_mamba( + preprocess_mamba( scheduler_output, self.kv_cache_config, self.cache_config, @@ -715,6 +719,15 @@ class NPUARModelRunner(OmniNPUModelRunner): # previous model forward without breaking async scheduling. if deferred_state_corrections_fn: deferred_state_corrections_fn() + + if self.vllm_config.model_config.enable_return_routed_experts and hasattr(self, "_positions_cpu"): + issue_routing_d2h_copy( + input_batch_req_ids=self.input_batch.req_ids, + num_scheduled_tokens=scheduler_output.num_scheduled_tokens, + positions=self.positions, + positions_cpu=self._positions_cpu, + ) + return None @torch.inference_mode() @@ -858,12 +871,17 @@ class NPUARModelRunner(OmniNPUModelRunner): if self.speculative_config is not None: self.finalize_kv_connector() + routed_experts_dict = None if self.model_config.enable_return_routed_experts: - capturer = RoutedExpertsCapturer.get_instance() - if capturer is not None: - capturer.save_captured_experts(indices=self.cpu_slot_mapping) - else: - logger.warning("RoutedExpertsCapturer is not initialized.") + capturer = get_global_experts_capturer() + if capturer is not None and hasattr(self.input_batch, "num_tokens_no_spec"): + routed_experts_dict = extract_routed_experts_for_current_batch( + req_ids=req_ids_output_copy, + requests=self.requests, + req_id_to_index=self.input_batch.req_id_to_index, + num_tokens_no_spec=self.input_batch.num_tokens_no_spec, + max_model_len=self.max_model_len, + ) # -------------------------------------- Omni-new ------------------------------------------------- engine_output_type, downstream_req_ids = self._resolve_pooler_payload_req_ids(req_ids_output_copy) @@ -984,6 +1002,7 @@ class NPUARModelRunner(OmniNPUModelRunner): kv_connector_output=kv_connector_output, ec_connector_output=ec_connector_output if self.supports_mm_inputs else None, cudagraph_stats=cudagraph_stats, + routed_experts_dict=routed_experts_dict, ) model_runner_output.kv_extracted_req_ids = kv_extracted_req_ids # -------------------------------------- Omni-new ------------------------------------------------- diff --git a/vllm_omni/platforms/npu/worker/npu_generation_model_runner.py b/vllm_omni/platforms/npu/worker/npu_generation_model_runner.py index 0b9b6ecac..f6af3de52 100644 --- a/vllm_omni/platforms/npu/worker/npu_generation_model_runner.py +++ b/vllm_omni/platforms/npu/worker/npu_generation_model_runner.py @@ -15,14 +15,18 @@ from vllm.distributed.ec_transfer import get_ec_transfer, has_ec_transfer from vllm.distributed.kv_transfer import get_kv_transfer_group, has_kv_transfer_group from vllm.distributed.parallel_state import get_pp_group from vllm.logger import logger -from vllm.model_executor.layers.fused_moe.routed_experts_capturer import RoutedExpertsCapturer +from vllm.model_executor.layers.fused_moe.routed_experts_capturer import ( + extract_routed_experts_for_current_batch, + get_global_experts_capturer, + issue_routing_d2h_copy, +) from vllm.sequence import IntermediateTensors from vllm.utils.math_utils import cdiv from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput from vllm.v1.outputs import EMPTY_MODEL_RUNNER_OUTPUT, AsyncModelRunnerOutput, make_empty_encoder_model_runner_output from vllm.v1.utils import record_function_or_nullcontext -from vllm.v1.worker import mamba_utils from vllm.v1.worker.gpu_model_runner import AsyncGPUModelRunnerOutput, PerLayerAttnMetadata +from vllm.v1.worker.mamba_utils import preprocess_mamba from vllm.v1.worker.ubatch_utils import maybe_create_ubatch_slices from vllm.v1.worker.utils import sanity_check_mm_encoder_outputs from vllm_ascend.ascend_forward_context import set_ascend_forward_context @@ -74,11 +78,9 @@ class NPUGenerationModelRunner(OmniNPUModelRunner): intermediate_tensors: IntermediateTensors | None = None, ) -> OmniModelRunnerOutput | IntermediateTensors | None: if self.vllm_config.model_config.enable_return_routed_experts: - capturer = RoutedExpertsCapturer.get_instance() - if capturer is not None: - capturer.clear_buffer() - else: - logger.warning("RoutedExpertsCapturer is not initialized.") + capturer = get_global_experts_capturer() + if capturer is not None and hasattr(capturer, "finalize_pending_copy"): + capturer.finalize_pending_copy() if self.ascend_config.profiling_chunk_config.enabled: self._sync_device() self._execution_start_time = time.perf_counter() @@ -227,7 +229,7 @@ class NPUGenerationModelRunner(OmniNPUModelRunner): if deferred_state_corrections_fn: deferred_state_corrections_fn() deferred_state_corrections_fn = None - mamba_utils.preprocess_mamba( + preprocess_mamba( scheduler_output, self.kv_cache_config, self.cache_config, @@ -383,6 +385,15 @@ class NPUGenerationModelRunner(OmniNPUModelRunner): # previous model forward without breaking async scheduling. if deferred_state_corrections_fn: deferred_state_corrections_fn() + + if self.vllm_config.model_config.enable_return_routed_experts and hasattr(self, "_positions_cpu"): + issue_routing_d2h_copy( + input_batch_req_ids=self.input_batch.req_ids, + num_scheduled_tokens=scheduler_output.num_scheduled_tokens, + positions=self.positions, + positions_cpu=self._positions_cpu, + ) + return None @torch.inference_mode() @@ -468,6 +479,17 @@ class NPUGenerationModelRunner(OmniNPUModelRunner): # [Omni] Copy req_id mappings to avoid async scheduling mutation. req_ids_output_copy = self.input_batch.req_ids.copy() req_id_to_index_output_copy = self.input_batch.req_id_to_index.copy() + routed_experts_dict = None + if self.vllm_config.model_config.enable_return_routed_experts and hasattr( + self.input_batch, "num_tokens_no_spec" + ): + routed_experts_dict = extract_routed_experts_for_current_batch( + req_ids=req_ids_output_copy, + requests=self.requests, + req_id_to_index=self.input_batch.req_id_to_index, + num_tokens_no_spec=self.input_batch.num_tokens_no_spec, + max_model_len=self.max_model_len, + ) output = OmniModelRunnerOutput( req_ids=req_ids_output_copy, req_id_to_index=req_id_to_index_output_copy, @@ -478,6 +500,7 @@ class NPUGenerationModelRunner(OmniNPUModelRunner): kv_connector_output=kv_connector_output, ec_connector_output=ec_connector_output if self.supports_mm_inputs else None, cudagraph_stats=cudagraph_stats, + routed_experts_dict=routed_experts_dict, ) # -------------------------------------- Omni-new ------------------------------------------------- diff --git a/vllm_omni/request.py b/vllm_omni/request.py index 48cbf9b31..f936069be 100644 --- a/vllm_omni/request.py +++ b/vllm_omni/request.py @@ -82,6 +82,7 @@ class OmniRequest(Request): client_index=request.client_index, prompt_token_ids=request.prompt_token_ids, prompt_embeds=request.prompt_embeds, + prompt_is_token_ids=request.prompt_is_token_ids, mm_features=request.mm_features, sampling_params=request.sampling_params, pooling_params=request.pooling_params, @@ -94,6 +95,8 @@ class OmniRequest(Request): additional_information=request.additional_information, resumable=request.resumable, reasoning_ended=request.reasoning_ended, + reasoning_parser_kwargs=request.reasoning_parser_kwargs, + abort_immediately=request.abort_immediately, ) diff --git a/vllm_omni/worker/gpu_ar_model_runner.py b/vllm_omni/worker/gpu_ar_model_runner.py index b36072bec..c4ba29a35 100644 --- a/vllm_omni/worker/gpu_ar_model_runner.py +++ b/vllm_omni/worker/gpu_ar_model_runner.py @@ -20,7 +20,9 @@ from vllm.distributed.parallel_state import get_pp_group, get_tp_group from vllm.forward_context import set_forward_context from vllm.logger import init_logger from vllm.model_executor.layers.fused_moe.routed_experts_capturer import ( - RoutedExpertsCapturer, + extract_routed_experts_for_current_batch, + get_global_experts_capturer, + issue_routing_d2h_copy, ) from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput from vllm.v1.outputs import AsyncModelRunnerOutput, make_empty_encoder_model_runner_output @@ -28,6 +30,7 @@ from vllm.v1.spec_decode.dflash import DFlashProposer from vllm.v1.spec_decode.draft_model import DraftModelProposer from vllm.v1.spec_decode.eagle import EagleProposer from vllm.v1.spec_decode.extract_hidden_states import ExtractHiddenStatesProposer +from vllm.v1.spec_decode.gemma4 import Gemma4Proposer from vllm.v1.structured_output.utils import apply_grammar_bitmask from vllm.v1.utils import record_function_or_nullcontext from vllm.v1.worker.gpu_model_runner import ( @@ -358,11 +361,21 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin): self.flush_full_payload_outputs(flush_ids) if self.routed_experts_initialized: - capturer = RoutedExpertsCapturer.get_instance() - if capturer is not None: - capturer.clear_buffer() # noqa - else: - logger.error("RoutedExpertsCapturer not initialized.") + capturer = get_global_experts_capturer() + if capturer is not None and hasattr(capturer, "finalize_pending_copy"): + capturer.finalize_pending_copy() + + # If ngram_gpu is used, we need to copy the scheduler_output to avoid + # the modification has influence on the scheduler_output in engine core process. + # The replace is much faster than deepcopy. + if self.speculative_config is not None and self.speculative_config.use_ngram_gpu(): + num_scheduled_tokens_copy = scheduler_output.num_scheduled_tokens.copy() + spec_decode_tokens_copy = scheduler_output.scheduled_spec_decode_tokens.copy() + scheduler_output = replace( + scheduler_output, + num_scheduled_tokens=num_scheduled_tokens_copy, + scheduled_spec_decode_tokens=spec_decode_tokens_copy, + ) if has_kv_transfer_group(): kv_connector_metadata = scheduler_output.kv_connector_metadata @@ -678,6 +691,14 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin): if deferred_state_corrections_fn: deferred_state_corrections_fn() + if self.routed_experts_initialized and hasattr(self, "_positions_cpu"): + issue_routing_d2h_copy( + input_batch_req_ids=self.input_batch.req_ids, + num_scheduled_tokens=scheduler_output.num_scheduled_tokens, + positions=self.positions, + positions_cpu=self._positions_cpu, + ) + return None def _sample( @@ -829,7 +850,7 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin): if use_gpu_toks: assert isinstance( self.drafter, - EagleProposer | DFlashProposer | DraftModelProposer | ExtractHiddenStatesProposer, + EagleProposer | DFlashProposer | DraftModelProposer | ExtractHiddenStatesProposer | Gemma4Proposer, ) sampled_token_ids = sampler_output.sampled_token_ids if input_fits_in_drafter: @@ -1013,12 +1034,15 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin): self.accumulate_full_payload_output(rid, pooler_output[i], req_state) with record_function_or_nullcontext("gpu_model_runner: ModelRunnerOutput"): + routed_experts_dict = None if self.routed_experts_initialized: - capturer = RoutedExpertsCapturer.get_instance() - if capturer is not None: - capturer.save_captured_experts(indices=self.slot_mapping) # noqa - else: - logger.error("RoutedExpertsCapturer not initialized.") + routed_experts_dict = extract_routed_experts_for_current_batch( + req_ids=req_ids_output_copy, + requests=self.requests, + req_id_to_index=self.input_batch.req_id_to_index, + num_tokens_no_spec=self.input_batch.num_tokens_no_spec, + max_model_len=self.max_model_len, + ) output = OmniModelRunnerOutput( req_ids=req_ids_output_copy, req_id_to_index=req_id_to_index_output_copy, @@ -1030,6 +1054,7 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin): ec_connector_output=ec_connector_output if self.supports_mm_inputs else None, num_nans_in_logits=num_nans_in_logits, cudagraph_stats=cudagraph_stats, + routed_experts_dict=routed_experts_dict, ) output.kv_extracted_req_ids = kv_extracted_req_ids output.omni_connector_output = self.get_omni_connector_output() diff --git a/vllm_omni/worker/gpu_generation_model_runner.py b/vllm_omni/worker/gpu_generation_model_runner.py index 3a7411459..5a5175dc0 100644 --- a/vllm_omni/worker/gpu_generation_model_runner.py +++ b/vllm_omni/worker/gpu_generation_model_runner.py @@ -9,6 +9,7 @@ from __future__ import annotations import gc import logging from copy import copy +from dataclasses import replace import numpy as np import torch @@ -18,7 +19,9 @@ from vllm.distributed.kv_transfer import get_kv_transfer_group, has_kv_transfer_ from vllm.distributed.parallel_state import get_pp_group from vllm.forward_context import set_forward_context from vllm.model_executor.layers.fused_moe.routed_experts_capturer import ( - RoutedExpertsCapturer, + extract_routed_experts_for_current_batch, + get_global_experts_capturer, + issue_routing_d2h_copy, ) from vllm.utils.math_utils import cdiv from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput @@ -27,6 +30,7 @@ from vllm.v1.spec_decode.dflash import DFlashProposer from vllm.v1.spec_decode.draft_model import DraftModelProposer from vllm.v1.spec_decode.eagle import EagleProposer from vllm.v1.spec_decode.extract_hidden_states import ExtractHiddenStatesProposer +from vllm.v1.spec_decode.gemma4 import Gemma4Proposer from vllm.v1.utils import record_function_or_nullcontext from vllm.v1.worker.gpu_model_runner import ( EMPTY_MODEL_RUNNER_OUTPUT, @@ -111,11 +115,21 @@ class GPUGenerationModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin self.flush_full_payload_outputs(flush_ids) if self.routed_experts_initialized: - capturer = RoutedExpertsCapturer.get_instance() - if capturer is not None: - capturer.clear_buffer() # noqa - else: - logger.error("RoutedExpertsCapturer not initialized.") + capturer = get_global_experts_capturer() + if capturer is not None and hasattr(capturer, "finalize_pending_copy"): + capturer.finalize_pending_copy() + + # If ngram_gpu is used, we need to copy the scheduler_output to avoid + # the modification has influence on the scheduler_output in engine core process. + # The replace is much faster than deepcopy. + if self.speculative_config is not None and self.speculative_config.use_ngram_gpu(): + num_scheduled_tokens_copy = scheduler_output.num_scheduled_tokens.copy() + spec_decode_tokens_copy = scheduler_output.scheduled_spec_decode_tokens.copy() + scheduler_output = replace( + scheduler_output, + num_scheduled_tokens=num_scheduled_tokens_copy, + scheduled_spec_decode_tokens=spec_decode_tokens_copy, + ) if has_kv_transfer_group(): kv_connector_metadata = scheduler_output.kv_connector_metadata @@ -339,6 +353,14 @@ class GPUGenerationModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin if deferred_state_corrections_fn: deferred_state_corrections_fn() + if self.routed_experts_initialized and hasattr(self, "_positions_cpu"): + issue_routing_d2h_copy( + input_batch_req_ids=self.input_batch.req_ids, + num_scheduled_tokens=scheduler_output.num_scheduled_tokens, + positions=self.positions, + positions_cpu=self._positions_cpu, + ) + return None @torch.inference_mode() @@ -425,6 +447,15 @@ class GPUGenerationModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin # [Omni] Copy req_id mappings to avoid async scheduling mutation. req_ids_output_copy = self.input_batch.req_ids.copy() req_id_to_index_output_copy = self.input_batch.req_id_to_index.copy() + routed_experts_dict = None + if self.routed_experts_initialized: + routed_experts_dict = extract_routed_experts_for_current_batch( + req_ids=req_ids_output_copy, + requests=self.requests, + req_id_to_index=self.input_batch.req_id_to_index, + num_tokens_no_spec=self.input_batch.num_tokens_no_spec, + max_model_len=self.max_model_len, + ) if self._should_accumulate_full_payload_output(): for i, rid in enumerate(req_ids_output_copy): req_state = self.requests.get(rid) @@ -442,6 +473,7 @@ class GPUGenerationModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin num_nans_in_logits={}, cudagraph_stats=cudagraph_stats, ec_connector_output=ec_connector_output if self.supports_mm_inputs else None, + routed_experts_dict=routed_experts_dict, ) output.omni_connector_output = self.get_omni_connector_output() @@ -751,7 +783,7 @@ class GPUGenerationModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin device=self.device, ) - intermediate_tensors = self.sync_and_slice_intermediate_tensors(num_tokens_padded, None, False) + intermediate_tensors = self.sync_and_gather_intermediate_tensors(num_tokens_padded, None, False) if ubatch_slices_padded is not None: # Adjust values to reflect a single ubatch. @@ -794,7 +826,7 @@ class GPUGenerationModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin ): assert isinstance( self.drafter, - EagleProposer | DFlashProposer | DraftModelProposer | ExtractHiddenStatesProposer, + EagleProposer | DFlashProposer | DraftModelProposer | ExtractHiddenStatesProposer | Gemma4Proposer, ) assert self.speculative_config is not None # Eagle currently only supports PIECEWISE cudagraphs. diff --git a/vllm_omni/worker/gpu_model_runner.py b/vllm_omni/worker/gpu_model_runner.py index b3f746eb4..af2a6cac4 100644 --- a/vllm_omni/worker/gpu_model_runner.py +++ b/vllm_omni/worker/gpu_model_runner.py @@ -1,4 +1,5 @@ import contextlib +from collections.abc import Callable from typing import TYPE_CHECKING, Any, cast import numpy as np @@ -18,6 +19,11 @@ from vllm.v1.spec_decode.dflash import DFlashProposer from vllm.v1.spec_decode.draft_model import DraftModelProposer from vllm.v1.spec_decode.eagle import EagleProposer from vllm.v1.spec_decode.extract_hidden_states import ExtractHiddenStatesProposer +from vllm.v1.spec_decode.gemma4 import Gemma4Proposer +from vllm.v1.spec_decode.ngram_proposer_gpu import ( + update_ngram_gpu_tensors_incremental, + update_scheduler_for_invalid_drafts, +) from vllm.v1.worker.gpu_input_batch import CachedRequestState from vllm.v1.worker.gpu_model_runner import GPUModelRunner, IntermediateTensors, PerLayerAttnMetadata from vllm.v1.worker.ubatch_utils import maybe_create_ubatch_slices @@ -250,7 +256,7 @@ class OmniGPUModelRunner(GPUModelRunner): mrope_pos_ptr += completion_part_len - def _update_states(self, scheduler_output: "SchedulerOutput"): + def _update_states(self, scheduler_output: "SchedulerOutput") -> Callable | None: """Update the cached states and the persistent batch with the scheduler output. @@ -284,8 +290,7 @@ class OmniGPUModelRunner(GPUModelRunner): if cleanup_finished_request is not None: cleanup_finished_request(req_id) - if hasattr(self, "late_interaction_runner"): - self.late_interaction_runner.on_requests_finished(scheduler_output.finished_req_ids) + self.late_interaction_runner.on_requests_finished(scheduler_output.finished_req_ids) # Remove the finished requests from the persistent batch. # NOTE(woosuk): There could be an edge case where finished_req_ids and # scheduled_req_ids overlap. This happens when a request is aborted and @@ -326,8 +331,9 @@ class OmniGPUModelRunner(GPUModelRunner): for req_id in unscheduled_req_ids: self.input_batch.remove_request(req_id) - if self.use_async_spec_decode: - self.prev_num_draft_tokens.np.fill(0) + is_ngram_gpu = self.speculative_config is not None and self.speculative_config.use_ngram_gpu() + if is_ngram_gpu: + ngram_gpu_new_reqs: list[CachedRequestState] = [] reqs_to_add: list[CachedRequestState] = [] deferred_spec_decode_corrections = [] @@ -379,8 +385,7 @@ class OmniGPUModelRunner(GPUModelRunner): lora_request=new_req_data.lora_request, ) self.requests[req_id] = req_state - if hasattr(self, "late_interaction_runner"): - self.late_interaction_runner.register_request(req_id, pooling_params) + self.late_interaction_runner.register_request(req_id, pooling_params) # If prompt embeddings are provided, decode and attach to inter_data try: @@ -429,12 +434,30 @@ class OmniGPUModelRunner(GPUModelRunner): self._init_xdrope_positions(req_state) reqs_to_add.append(self.requests[req_id]) + # Track new requests for ngram_gpu full tensor copy + if is_ngram_gpu: + ngram_gpu_new_reqs.append(req_state) # Update the states of the running/resumed requests. is_last_rank = get_pp_group().is_last_rank req_data = scheduler_output.scheduled_cached_reqs scheduled_spec_tokens = scheduler_output.scheduled_spec_decode_tokens + # Save scheduler-allocated spec lengths before trimming so + # prev_num_draft_len keeps the optimistic count for rejection correction. + original_num_spec_per_req: dict[str, int] = {} + if self.speculative_config is not None and self.speculative_config.use_ngram_gpu(): + for req_id, toks in scheduled_spec_tokens.items(): + original_num_spec_per_req[req_id] = len(toks) + update_scheduler_for_invalid_drafts( + self._num_valid_draft_tokens_event, + self._num_valid_draft_tokens_cpu, + scheduler_output, + self.input_batch.req_id_to_index, + ) + if self.use_async_spec_decode: + self.prev_num_draft_tokens.np.fill(0) + for i, req_id in enumerate(req_data.req_ids): req_state = self.requests[req_id] num_computed_tokens = req_data.num_computed_tokens[i] @@ -473,6 +496,9 @@ class OmniGPUModelRunner(GPUModelRunner): if prev_req_index is not None: self.prev_num_draft_tokens.np[prev_req_index] = optimistic_num_accepted + if is_ngram_gpu and optimistic_num_accepted > 0: + self.input_batch.num_tokens_no_spec[req_index] += optimistic_num_accepted + # Update the cached states. req_state.num_computed_tokens = num_computed_tokens @@ -521,6 +547,9 @@ class OmniGPUModelRunner(GPUModelRunner): req_state.output_token_ids = resumed_token_ids[-num_output_tokens:] reqs_to_add.append(req_state) + # Track resumed requests for ngram_gpu full tensor copy + if is_ngram_gpu: + ngram_gpu_new_reqs.append(req_state) continue # Update the persistent batch. @@ -539,6 +568,11 @@ class OmniGPUModelRunner(GPUModelRunner): # Add spec_token_ids to token_ids_cpu. self.input_batch.update_req_spec_token_ids(req_state, scheduled_spec_tokens) + # Restore scheduler-side draft count after ngram trimming. + if original_num_spec_per_req: + orig = original_num_spec_per_req.get(req_id, 0) + if orig != req_state.prev_num_draft_len: + req_state.prev_num_draft_len = orig # Add the new or resumed requests to the persistent batch. # The smaller empty indices are filled first. @@ -553,6 +587,18 @@ class OmniGPUModelRunner(GPUModelRunner): # Refresh batch metadata with any pending updates. self.input_batch.refresh_metadata() + # Incrementally update ngram_gpu tensors after batch is stable + if is_ngram_gpu: + update_ngram_gpu_tensors_incremental( + self.input_batch, + self.token_ids_gpu_tensor, + self.num_tokens_no_spec_gpu, + ngram_gpu_new_reqs, + self.device, + _pinned_idx_buf=self._ngram_pinned_idx_buf, + _pinned_val_buf=self._ngram_pinned_val_buf, + ) + if deferred_spec_decode_corrections: def correct_spec_decode_token_counts(): @@ -577,6 +623,9 @@ class OmniGPUModelRunner(GPUModelRunner): if cur_req_index is None: continue self.input_batch.num_computed_tokens_cpu[cur_req_index] -= correction + if is_ngram_gpu and correction > 0: + self.input_batch.num_tokens_no_spec[cur_req_index] -= correction + self.num_tokens_no_spec_gpu[cur_req_index] -= correction return correct_spec_decode_token_counts else: @@ -846,7 +895,7 @@ class OmniGPUModelRunner(GPUModelRunner): device=self.device, ) - intermediate_tensors = self.sync_and_slice_intermediate_tensors(num_tokens_padded, None, False) + intermediate_tensors = self.sync_and_gather_intermediate_tensors(num_tokens_padded, None, False) if ubatch_slices_padded is not None: # Adjust values to reflect a single ubatch. @@ -900,7 +949,7 @@ class OmniGPUModelRunner(GPUModelRunner): ): assert isinstance( self.drafter, - EagleProposer | DFlashProposer | DraftModelProposer | ExtractHiddenStatesProposer, + EagleProposer | DFlashProposer | DraftModelProposer | ExtractHiddenStatesProposer | Gemma4Proposer, ) assert self.speculative_config is not None # Eagle currently only supports PIECEWISE cudagraphs. @@ -1300,7 +1349,7 @@ class OmniGPUModelRunner(GPUModelRunner): intermediate_tensors = None else: assert intermediate_tensors is not None - intermediate_tensors = self.sync_and_slice_intermediate_tensors( + intermediate_tensors = self.sync_and_gather_intermediate_tensors( num_input_tokens, intermediate_tensors, True )