diff --git a/.buildkite/test-merge.yml b/.buildkite/test-merge.yml index 318b872b5..8d47539d3 100644 --- a/.buildkite/test-merge.yml +++ b/.buildkite/test-merge.yml @@ -107,7 +107,7 @@ steps: depends_on: upload-merge-pipeline steps: - label: "Entrypoint Test with H100" - timeout_in_minutes: 35 + timeout_in_minutes: 60 commands: - pytest -s -v tests/entrypoints/ -m "advanced_model and cuda and H100" --run-level "advanced_model" agents: diff --git a/.buildkite/test-ready.yml b/.buildkite/test-ready.yml index 15188fde5..055b9b0e3 100644 --- a/.buildkite/test-ready.yml +++ b/.buildkite/test-ready.yml @@ -183,7 +183,7 @@ steps: depends_on: upload-ready-pipeline commands: - | - timeout 15m bash -c " + timeout 20m bash -c " pytest -sv tests/distributed/ -m 'core_model and cuda' --run-level 'core_model' " agents: diff --git a/benchmarks/diffusion/__init__.py b/benchmarks/diffusion/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/benchmarks/tts/model_configs.yaml b/benchmarks/tts/model_configs.yaml index 17254e37c..804e60be8 100644 --- a/benchmarks/tts/model_configs.yaml +++ b/benchmarks/tts/model_configs.yaml @@ -14,6 +14,7 @@ models: supported_tasks: [default_voice, voice_design] backend: openai-audio-speech endpoint: /v1/audio/speech + trust_remote_code: true task_extra_body: default_voice: voice: Vivian @@ -27,6 +28,7 @@ models: supported_tasks: [voice_clone] backend: openai-audio-speech endpoint: /v1/audio/speech + trust_remote_code: true task_extra_body: voice_clone: task_type: Base @@ -36,6 +38,7 @@ models: trust_remote_code: true backend: openai-audio-speech endpoint: /v1/audio/speech + trust_remote_code: true task_extra_body: voice_clone: {} default_voice: diff --git a/docker/Dockerfile.ci b/docker/Dockerfile.ci index a823b5bc5..31e2529b1 100644 --- a/docker/Dockerfile.ci +++ b/docker/Dockerfile.ci @@ -1,5 +1,8 @@ +# Target vllm version: v0.21.0 (releases/v0.21.0 branch). +# The pinned VLLM_PRECOMPILED_WHEEL_COMMIT and VLLM_BASE_TAG below must +# be updated on every rebase to match the new upstream vllm commit/tag. ARG VLLM_BASE_IMAGE=vllm/vllm-openai -ARG VLLM_BASE_TAG=v0.21.0 +ARG VLLM_BASE_TAG=v0.22.0 FROM ${VLLM_BASE_IMAGE}:${VLLM_BASE_TAG} ARG APP_DIR=/workspace/vllm-omni WORKDIR ${APP_DIR} diff --git a/docs/user_guide/quantization/gguf.md b/docs/user_guide/quantization/gguf.md index cd2e8b875..e70d2d26e 100644 --- a/docs/user_guide/quantization/gguf.md +++ b/docs/user_guide/quantization/gguf.md @@ -74,7 +74,7 @@ Online: vllm serve Qwen/Qwen-Image \ --omni \ --port 8000 \ - --quantization-config '{"method":"gguf","gguf_model":"QuantStack/Qwen-Image-GGUF/Qwen_Image-Q4_K_M.gguf"}' + --diffusion-quantization-config '{"method":"gguf","gguf_model":"QuantStack/Qwen-Image-GGUF/Qwen_Image-Q4_K_M.gguf"}' ``` ## Parameters diff --git a/recipes/README.md b/recipes/README.md index 20248b59c..ccf0efd53 100644 --- a/recipes/README.md +++ b/recipes/README.md @@ -47,6 +47,7 @@ recipes/ | [`Wan-AI/Wan2.2-S2V.md`](./Wan-AI/Wan2.2-S2V.md) | Speech-to-video serving (Wan2.2 14B) | 2x A100/H100 80GB | | [`zai-org/GLM-TTS.md`](./zai-org/GLM-TTS.md) | Online serving for Chinese/English zero-shot voice-cloned TTS | 1x A40 48GB | | [`GLM/GLM-Image.md`](./GLM/GLM-Image.md) | Online serving for image generation | 1x A800 80GB / 2x A800 80GB | +| Within a single recipe file, include different hardware support sections such as `GPU`, `ROCm`, and `NPU`, and add concrete tested configurations like diff --git a/requirements/common.txt b/requirements/common.txt index 84e0fa49c..e649450af 100644 --- a/requirements/common.txt +++ b/requirements/common.txt @@ -1,4 +1,10 @@ # Common dependencies for all platforms +# transformers: keep aligned with upstream vLLM constraint. +# vLLM pins ">= 4.56.0, != 5.0.*, != 5.1.*, != 5.2.*, != 5.3.*, != 5.4.*, != 5.5.0". +# 5.9.0 breaks models that define _keys_to_ignore_on_load_unexpected as a list +# (e.g. OpenMOSS-Team/MOSS-TTS-Nano) because it changed +# _adjust_missing_and_unexpected_keys from list+list to set|set. +transformers >= 4.56.0, != 5.0.*, != 5.1.*, != 5.2.*, != 5.3.*, != 5.4.*, != 5.5.0, < 5.9.0 av>=14.0.0 omegaconf>=2.3.0 diffusers==0.38.0 diff --git a/tests/dfx/perf/scripts/run_diffusion_benchmark.py b/tests/dfx/perf/scripts/run_diffusion_benchmark.py index 92ce2740d..62ce59606 100644 --- a/tests/dfx/perf/scripts/run_diffusion_benchmark.py +++ b/tests/dfx/perf/scripts/run_diffusion_benchmark.py @@ -223,8 +223,12 @@ def _get_open_port() -> int: return s.getsockname()[1] -def _wait_for_port(host: str, port: int, timeout: int = 1200) -> None: - """Block until the given host:port accepts connections or timeout expires.""" +def _wait_for_port(host: str, port: int, timeout: int = 1200, proc: subprocess.Popen | None = None) -> None: + """Block until the given host:port accepts connections or timeout expires. + + If *proc* is provided, also monitors the process; raises RuntimeError + immediately if the server process exits before the port becomes available. + """ start = time.time() while time.time() - start < timeout: try: @@ -234,6 +238,10 @@ def _wait_for_port(host: str, port: int, timeout: int = 1200) -> None: return except Exception: pass + if proc is not None: + ret = proc.poll() + if ret is not None: + raise RuntimeError(f"Server process exited with code {ret} before port {host}:{port} became ready") time.sleep(2) raise RuntimeError(f"Server did not start on {host}:{port} within {timeout}s") @@ -331,7 +339,7 @@ class DiffusionServer: env=env, cwd=str(Path(__file__).parent.parent.parent.parent), ) - _wait_for_port(self.host, self.port) + _wait_for_port(self.host, self.port, proc=self.proc) print(f"DiffusionServer ready on {self.host}:{self.port}") def __enter__(self): diff --git a/tests/dfx/perf/tests/test_qwen3_omni_vllm_text.json b/tests/dfx/perf/tests/test_qwen3_omni_vllm_text.json index 2f580aac7..76ac1d858 100644 --- a/tests/dfx/perf/tests/test_qwen3_omni_vllm_text.json +++ b/tests/dfx/perf/tests/test_qwen3_omni_vllm_text.json @@ -5,7 +5,8 @@ "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct", "use_omni": false, "extra_cli_args": [ - "--no-enable-prefix-caching" + "--no-enable-prefix-caching", + "--skip-mm-profiling" ] }, "benchmark_params": [ diff --git a/tests/diffusion/offloader/test_diffusion_cpu_offload.py b/tests/diffusion/offloader/test_diffusion_cpu_offload.py index 142e6b437..92e2abc7e 100644 --- a/tests/diffusion/offloader/test_diffusion_cpu_offload.py +++ b/tests/diffusion/offloader/test_diffusion_cpu_offload.py @@ -17,7 +17,7 @@ AUDIO_MODEL = { } IMAGE_VIDEO_MODELS = { - "riverclouds/qwen_image_random": {"cuda": 2500, "rocm": 2100}, + "riverclouds/qwen_image_random": {"cuda": 2200, "rocm": 2100}, } MODELS = {**AUDIO_MODEL, **IMAGE_VIDEO_MODELS} diff --git a/tests/diffusion/test_diffusion_model_runner.py b/tests/diffusion/test_diffusion_model_runner.py index 4205d8280..960dd072b 100644 --- a/tests/diffusion/test_diffusion_model_runner.py +++ b/tests/diffusion/test_diffusion_model_runner.py @@ -301,8 +301,4 @@ def test_vllm_set_forward_context_implementation(monkeypatch): "vllm_config", ), ), - ("ir_op_priority", None), - ("enable_torch_wrap", vllm_config.compilation_config.ir_enable_torch_wrap), - ("enable_torch_wrap_exit", vllm_config.compilation_config.ir_enable_torch_wrap), - ("ir_op_priority_exit", None), ], ERROR_MESSAGE diff --git a/tests/diffusion/test_stage_diffusion_proc.py b/tests/diffusion/test_stage_diffusion_proc.py index 0419bb813..dcefafb8d 100644 --- a/tests/diffusion/test_stage_diffusion_proc.py +++ b/tests/diffusion/test_stage_diffusion_proc.py @@ -85,7 +85,7 @@ class MockOmniRequestOutput: BASE_HEIGHT = 512 BASE_WIDTH = 512 BASE_INFER_STEPS = 10 -DELAY_BASE = 2 +DELAY_BASE = 0.01 class MockDiffusionEngine: @@ -93,7 +93,7 @@ class MockDiffusionEngine: def simulate_step_delay(height, width, num_inference_steps) -> float: return (height / BASE_HEIGHT) * (width / BASE_WIDTH) * (num_inference_steps / BASE_INFER_STEPS) - DELAY_BASE = 2 + DELAY_BASE = 0.01 delay_scale = simulate_step_delay( request.sampling_params.height, request.sampling_params.width, request.sampling_params.num_inference_steps ) @@ -139,7 +139,7 @@ async def test_proc_process_request_with_batching_async_output(): assert len(results) == len(test_requests) base_time = DELAY_BASE time_gap_std = DELAY_BASE * 2 * 2 * 1 # height/width/steps infer time scale - eps = 0.5 + eps = 0.1 for i, (res, elapsed_time) in enumerate(results): assert res.request_id == test_requests[i]["request_id"] assert isinstance(res, MockOmniRequestOutput) diff --git a/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py b/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py index 8bfe837e2..aa72133c9 100644 --- a/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py +++ b/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py @@ -703,6 +703,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=None, routed_experts_dict=None, ) @@ -789,6 +790,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=None, routed_experts_dict=None, ) diff --git a/tests/distributed/omni_coordinator/test_omni_coordinator.py b/tests/distributed/omni_coordinator/test_omni_coordinator.py index 81bd12454..c7b1b597e 100644 --- a/tests/distributed/omni_coordinator/test_omni_coordinator.py +++ b/tests/distributed/omni_coordinator/test_omni_coordinator.py @@ -68,13 +68,13 @@ def test_omni_coordinator_pub_coalescing_on_rapid_queue_updates(): sub_ctx = zmq.Context.instance() sub = sub_ctx.socket(zmq.SUB) - sub.connect(pub_addr) + sub.connect(coordinator.pub_zmq_addr) sub.setsockopt(zmq.SUBSCRIBE, b"") time.sleep(0.3) # PUB/SUB slow-joiner client = OmniCoordClientForStage( - router_addr, + coordinator.router_zmq_addr, "tcp://stage:coalesce", "tcp://stage:coalesce-out", 0, @@ -131,7 +131,7 @@ def test_omni_coordinator_registration_broadcast(): sub_ctx = zmq.Context.instance() sub = sub_ctx.socket(zmq.SUB) - sub.connect(pub_addr) + sub.connect(coordinator.pub_zmq_addr) sub.setsockopt(zmq.SUBSCRIBE, b"") # ZMQ PUB/SUB slow-joiner: allow SUB to connect before clients register. @@ -139,9 +139,9 @@ def test_omni_coordinator_registration_broadcast(): # Create 3 stage clients; each auto-registers on init. clients = [ - OmniCoordClientForStage(router_addr, "tcp://stage:10001", "tcp://stage:10001-out", 0), - OmniCoordClientForStage(router_addr, "tcp://stage:10002", "tcp://stage:10002-out", 0), - OmniCoordClientForStage(router_addr, "tcp://stage:10003", "tcp://stage:10003-out", 1), + OmniCoordClientForStage(coordinator.router_zmq_addr, "tcp://stage:10001", "tcp://stage:10001-out", 0), + OmniCoordClientForStage(coordinator.router_zmq_addr, "tcp://stage:10002", "tcp://stage:10002-out", 0), + OmniCoordClientForStage(coordinator.router_zmq_addr, "tcp://stage:10003", "tcp://stage:10003-out", 1), ] msg = _wait_for_replica_list(sub, expected_count=3) @@ -183,19 +183,19 @@ def test_omni_coordinator_heartbeat_timeout_handling(): sub_ctx = zmq.Context.instance() sub = sub_ctx.socket(zmq.SUB) - sub.connect(pub_addr) + sub.connect(coordinator.pub_zmq_addr) sub.setsockopt(zmq.SUBSCRIBE, b"") time.sleep(0.3) # A and B: real clients that send heartbeats every 5s. - client_a = OmniCoordClientForStage(router_addr, "tcp://stage:a", "tcp://stage:a-out", 0) - client_b = OmniCoordClientForStage(router_addr, "tcp://stage:b", "tcp://stage:b-out", 0) + client_a = OmniCoordClientForStage(coordinator.router_zmq_addr, "tcp://stage:a", "tcp://stage:a-out", 0) + client_b = OmniCoordClientForStage(coordinator.router_zmq_addr, "tcp://stage:b", "tcp://stage:b-out", 0) # C: raw DEALER that sends only registration, no heartbeat. dealer_ctx = zmq.Context.instance() dealer_c = dealer_ctx.socket(zmq.DEALER) - dealer_c.connect(router_addr) + dealer_c.connect(coordinator.router_zmq_addr) reg_event = { "input_addr": "tcp://stage:c", "output_addr": "tcp://stage:c-out", @@ -254,12 +254,12 @@ def test_omni_coordinator_replica_shutdown_handling(): sub_ctx = zmq.Context.instance() sub = sub_ctx.socket(zmq.SUB) - sub.connect(pub_addr) + sub.connect(coordinator.pub_zmq_addr) sub.setsockopt(zmq.SUBSCRIBE, b"") time.sleep(0.3) # PUB/SUB slow-joiner - client = OmniCoordClientForStage(router_addr, "tcp://stage:shutdown", "tcp://stage:shutdown-out", 0) + client = OmniCoordClientForStage(coordinator.router_zmq_addr, "tcp://stage:shutdown", "tcp://stage:shutdown-out", 0) msg = _wait_for_replica_list(sub, expected_count=1) assert msg is not None diff --git a/tests/e2e/online_serving/test_qwen_image_expansion.py b/tests/e2e/online_serving/test_qwen_image_expansion.py index 0830e4944..f1cf449f7 100644 --- a/tests/e2e/online_serving/test_qwen_image_expansion.py +++ b/tests/e2e/online_serving/test_qwen_image_expansion.py @@ -97,7 +97,7 @@ def _get_diffusion_feature_cases(model: str): "--vae-patch-parallel-size", "2", "--vae-use-tiling", - "--quantization-config", + "--diffusion-quantization-config", '{"method":"fp8"}', ], ), diff --git a/tests/entrypoints/test_stage_utils.py b/tests/entrypoints/test_stage_utils.py index 4bc26bff9..f31bf3f20 100644 --- a/tests/entrypoints/test_stage_utils.py +++ b/tests/entrypoints/test_stage_utils.py @@ -4,7 +4,7 @@ import sys import pytest from pytest_mock import MockerFixture -from vllm_omni.entrypoints.stage_utils import set_stage_devices +from vllm_omni.entrypoints.stage_utils import _map_device_list, set_stage_devices def _make_dummy_torch(call_log): @@ -130,3 +130,90 @@ def test_set_stage_devices_npu_platform(mocker: MockerFixture, monkeypatch: pyte set_stage_devices(stage_id=0, devices="0,1") assert os.environ["ASCEND_RT_VISIBLE_DEVICES"] == "4,5" + + +# ---- _map_device_list unit tests ---- + + +@pytest.mark.core_model +@pytest.mark.cpu +def test_map_device_list_idempotency(): + """Device IDs already in the visible set are returned as-is (idempotency).""" + result = _map_device_list(0, ["0", "1"], ["0", "1", "2", "3"]) + assert result == ["0", "1"] + + +@pytest.mark.core_model +@pytest.mark.cpu +def test_map_device_list_idempotency_single(): + """Single device ID in visible set passes through.""" + result = _map_device_list(1, ["1"], ["0", "1"]) + assert result == ["1"] + + +@pytest.mark.core_model +@pytest.mark.cpu +def test_map_device_list_index_mapping(): + """Device IDs < num_visible not in set are treated as indices.""" + result = _map_device_list(0, ["0", "1"], ["6", "7", "8", "9"]) + assert result == ["6", "7"] + + +@pytest.mark.core_model +@pytest.mark.cpu +def test_map_device_list_physical_fallback(): + """Device IDs >= num_visible are treated as physical IDs and passed through.""" + result = _map_device_list(1, ["2"], ["0", "1"]) + assert result == ["2"] + + +@pytest.mark.core_model +@pytest.mark.cpu +def test_map_device_list_physical_fallback_multi(): + """Multiple device IDs >= num_visible all pass through as physical IDs.""" + result = _map_device_list(1, ["2", "3"], ["0", "1"]) + assert result == ["2", "3"] + + +@pytest.mark.core_model +@pytest.mark.cpu +def test_map_device_list_physical_fallback_mixed(): + """When some devices are >= num_visible but some < num_visible, index mapping applies for the subset.""" + # "1" < num_visible(2) so index mapping applies: visible[1] = "5" + # "2" >= num_visible(2) so it's dropped via the partial mapping path + result = _map_device_list(1, ["1", "2"], ["0", "5"]) + assert result == ["5"] + + +@pytest.mark.core_model +@pytest.mark.cpu +def test_map_device_list_index_with_gaps(): + """Index-based mapping works with non-contiguous visible device lists.""" + result = _map_device_list(0, ["0", "1"], ["4", "5", "7"]) + assert result == ["4", "5"] + + +@pytest.mark.core_model +@pytest.mark.cpu +def test_map_device_list_index_mapping_no_idempotency(): + """Device IDs < num_visible not in visible set map via index even when no ID matches literally.""" + # "0","1" with visible ["2","3"]: idempotency check fails (none in set), + # but index mapping succeeds: visible[0]="2", visible[1]="3" + result = _map_device_list(0, ["0", "1"], ["2", "3"]) + assert result == ["2", "3"] + + +@pytest.mark.core_model +@pytest.mark.cpu +def test_map_device_list_partial_mapping(): + """When only a subset can map, returns the mapped subset (no error).""" + result = _map_device_list(0, ["0", "1", "2"], ["5", "6"]) + assert result == ["5", "6"] + + +@pytest.mark.core_model +@pytest.mark.cpu +def test_map_device_list_raises_on_non_numeric(): + """Non-numeric device IDs raise ValueError.""" + with pytest.raises(ValueError, match="must be non-negative integers"): + _map_device_list(0, ["a"], ["0", "1"]) diff --git a/tests/helpers/env.py b/tests/helpers/env.py index ef9f8730e..6dee5e610 100644 --- a/tests/helpers/env.py +++ b/tests/helpers/env.py @@ -160,9 +160,41 @@ def _print_gpu_processes() -> None: print("=" * 80) +def _cleanup_stale_device_locks() -> None: + """Remove stale device-initialization lock files whose recorded PID is dead. + + Lock files at ``/tmp/vllm_omni_device_*_init.lock`` may persist after a + crashed / killed test run and block subsequent orchestrator startups. + """ + import glob as _glob + + for lock_file in _glob.glob("/tmp/vllm_omni_device_*_init.lock"): + try: + with open(lock_file) as fh: + content = fh.read().strip() + if not content: + continue + pid = int(content) + except (OSError, ValueError): + continue + + try: + os.kill(pid, 0) + except ProcessLookupError: + print(f" Removing stale device lock {lock_file} (PID {pid} is dead)") + try: + os.unlink(lock_file) + except OSError: + pass + except PermissionError: + pass + + def run_pre_test_cleanup() -> None: print("Pre-test GPU status:") + _cleanup_stale_device_locks() + num_gpus = current_omni_platform.device_count() if num_gpus > 0: try: diff --git a/tests/helpers/mark.py b/tests/helpers/mark.py index 74dde5186..7d045e6df 100644 --- a/tests/helpers/mark.py +++ b/tests/helpers/mark.py @@ -21,7 +21,7 @@ def cuda_marks(*, res: str, num_cards: int): return marks test_distributed = pytest.mark.distributed_cuda(num_cards=num_cards) - test_skipif = pytest.mark.skipif_cuda( + test_skipif = pytest.mark.skipif( not current_platform.is_cuda() or (current_platform.device_count() < num_cards), reason=f"Need at least {num_cards} CUDA GPUs to run the test.", ) @@ -52,7 +52,7 @@ def xpu_marks(*, res: str, num_cards: int): return marks test_distributed = pytest.mark.distributed_xpu(num_cards=num_cards) - test_skipif = pytest.mark.skipif_xpu( + test_skipif = pytest.mark.skipif( not current_platform.is_xpu() or (current_platform.device_count() < num_cards), reason=f"Need at least {num_cards} XPUs to run the test.", ) diff --git a/tests/helpers/runtime.py b/tests/helpers/runtime.py index a3d337e1b..c0813537c 100644 --- a/tests/helpers/runtime.py +++ b/tests/helpers/runtime.py @@ -305,40 +305,132 @@ class OmniServer: time.sleep(2) raise RuntimeError(f"Server failed to start within {max_wait} seconds") + @staticmethod + def _reap_zombie(proc: "psutil.Process") -> bool: + """Reap a zombie child process via ``os.waitpid``. + + ``psutil.Process.wait()`` uses ``pidfd_open`` + ``poll()``, which + never fires for a zombie (the zombie is already dead and will not + change state). Since the test process is the parent, we can reap + the zombie directly with ``os.waitpid(pid, os.WNOHANG)`` and + retrieve its exit code. + + Returns True if the process was a zombie and was reaped. + """ + if proc.status() != psutil.STATUS_ZOMBIE: + return False + try: + os.waitpid(proc.pid, os.WNOHANG) + except ChildProcessError: + pass + return True + + @staticmethod + def _wait_or_reap(proc: "psutil.Process", timeout: int) -> None: + """Wait for *proc* to exit, handling zombie state transparently. + + When the process is already a zombie (e.g. orchestrator thread + hung during shutdown and the main process exited), ``pidfd_open`` + + ``poll()`` inside ``psutil`` will never see a state change. + Fall back to ``os.waitpid`` to reap the zombie. + """ + if OmniServer._reap_zombie(proc): + return + try: + proc.wait(timeout=timeout) + except psutil.TimeoutExpired: + OmniServer._reap_zombie(proc) + def _kill_process_tree(self, pid): + """Kill the process tree rooted at *pid*. + + Terminate the parent **first** so the OmniServer can gracefully shut + down its stage-engine children through the orchestrator. This avoids + the ``subprocess died unexpectedly`` ERROR that the APIServer monitor + thread logs when children are killed before the parent, which in turn + can cause CI watchdogs to false-trigger on the upstream ``Shutdown + initiated`` message. + + When the parent does not exit within the grace period (e.g. CPU- + offloaded workers stuck in CUDA D-state), the method falls back to + killing children first so the parent can be reaped cleanly. + """ try: parent = psutil.Process(pid) children = parent.children(recursive=True) - all_pids = [pid] + [child.pid for child in children] - - for child in children: - try: - child.terminate() - except psutil.NoSuchProcess: - pass - - _, still_alive = psutil.wait_procs(children, timeout=10) - - for child in still_alive: - try: - child.kill() - except psutil.NoSuchProcess: - pass + # 1. Terminate the parent first — let it run its graceful + # shutdown cascade (orchestrator → stage pools → engine cores). try: parent.terminate() - parent.wait(timeout=10) - except (psutil.NoSuchProcess, psutil.TimeoutExpired): + except psutil.NoSuchProcess: + pass + + # 2. Give the parent time to shut down its children cleanly. + parent_exited = False + try: + parent.wait(timeout=15) + parent_exited = True + except psutil.NoSuchProcess: + parent_exited = True + except psutil.TimeoutExpired: + parent_exited = OmniServer._reap_zombie(parent) + + if not parent_exited: + # Parent is stuck — children (e.g. CPU-offloaded CFG workers) + # are likely in uninterruptible sleep. Kill children first + # so the parent can be reaped without lingering as a zombie. + for child in children: + try: + child.kill() + except psutil.NoSuchProcess: + pass + psutil.wait_procs(children, timeout=5) try: parent.kill() except psutil.NoSuchProcess: pass + OmniServer._wait_or_reap(parent, timeout=5) + else: + # Parent exited cleanly — clean up any remaining children. + for child in children: + try: + if child.is_running(): + child.terminate() + except psutil.NoSuchProcess: + pass + gone, still_alive = psutil.wait_procs(children, timeout=10) + + for child in still_alive: + try: + child.kill() + except psutil.NoSuchProcess: + pass + + try: + if parent.is_running() and not OmniServer._reap_zombie(parent): + parent.kill() + parent.wait(timeout=10) + except psutil.NoSuchProcess: + pass + + # 3. Final sweep — ``kill -9`` anything that escaped. time.sleep(1) - alive_processes = [] - for check_pid in all_pids: - if psutil.pid_exists(check_pid): - alive_processes.append(check_pid) + alive_processes: list[int] = [] + for child in children: + try: + if child.is_running(): + alive_processes.append(child.pid) + except psutil.NoSuchProcess: + pass + # Only count the parent as alive if it is NOT a zombie + # (zombies are already dead — just waiting to be reaped). + try: + if parent.is_running() and parent.status() != psutil.STATUS_ZOMBIE: + alive_processes.append(parent.pid) + except psutil.NoSuchProcess: + pass if alive_processes: print(f"Warning: Processes still alive: {alive_processes}") diff --git a/tests/test_arg_utils.py b/tests/test_arg_utils.py index 8fc445a60..1f454783f 100644 --- a/tests/test_arg_utils.py +++ b/tests/test_arg_utils.py @@ -164,7 +164,8 @@ def test_split_mixed_real_world(): # orchestrator-only keys never reach engine assert "stage_init_timeout" not in engine assert "deploy_config" not in engine - assert "async_chunk" not in engine + # async_chunk is in SHARED_FIELDS so it flows to both + assert "async_chunk" in engine # ============================================================================ @@ -172,24 +173,47 @@ def test_split_mixed_real_world(): # ============================================================================ -def test_user_typed_unclassified_warns(caplog): +def test_user_typed_unclassified_warns(monkeypatch): """If the user types a flag we can't route, warn — don't silently drop.""" - raw = {"bogus_flag": "value", "max_num_seqs": 64} - with caplog.at_level(logging.WARNING, logger="vllm_omni.engine.arg_utils"): + import io + + buf = io.StringIO() + vllm_logger = logging.getLogger("vllm") + orig_handlers = list(vllm_logger.handlers) + vllm_logger.handlers.clear() + handler = logging.StreamHandler(buf) + handler.setLevel(logging.WARNING) + vllm_logger.addHandler(handler) + try: + raw = {"bogus_flag": "value", "max_num_seqs": 64} split_kwargs(raw, engine_cls=_FakeEngineArgs, user_typed={"bogus_flag"}) - assert any("bogus_flag" in rec.message for rec in caplog.records), ( - f"Expected warning mentioning 'bogus_flag', got: {[rec.message for rec in caplog.records]}" - ) + assert "bogus_flag" in buf.getvalue(), f"Expected 'bogus_flag' in log output, got: {buf.getvalue()!r}" + finally: + vllm_logger.handlers.clear() + for h in orig_handlers: + vllm_logger.addHandler(h) -def test_unclassified_without_user_typed_silent(caplog): +def test_unclassified_without_user_typed_silent(): """Without user_typed, unclassified keys drop silently (argparse defaults for server flags shouldn't spam logs on every launch).""" - raw = {"host": "0.0.0.0", "port": 8091, "max_num_seqs": 64} - with caplog.at_level(logging.WARNING, logger="vllm_omni.engine.arg_utils"): + import io + + vllm_logger = logging.getLogger("vllm") + orig_handlers = list(vllm_logger.handlers) + vllm_logger.handlers.clear() + buf = io.StringIO() + handler = logging.StreamHandler(buf) + handler.setLevel(logging.WARNING) + vllm_logger.addHandler(handler) + try: + raw = {"host": "0.0.0.0", "port": 8091, "max_num_seqs": 64} split_kwargs(raw, engine_cls=_FakeEngineArgs, user_typed=None) - # No warnings because we don't know these were user-typed. - assert not any("host" in rec.message or "port" in rec.message for rec in caplog.records) + assert "host" not in buf.getvalue() + finally: + vllm_logger.handlers.clear() + for h in orig_handlers: + vllm_logger.addHandler(h) # ============================================================================ @@ -290,7 +314,7 @@ def test_internal_blacklist_keys_derived_from_orchestrator(): # Spot-check expected entries assert "stage_init_timeout" in blacklist assert "deploy_config" in blacklist - assert "async_chunk" in blacklist + assert "async_chunk" not in blacklist # now in SHARED_FIELDS # Shared fields must NOT appear — they flow to both orchestrator and engine assert "model" not in blacklist assert "log_stats" not in blacklist @@ -339,18 +363,31 @@ def test_ambiguous_field_strict_raises(): split_kwargs({"deploy_config": "x"}, engine_cls=_AmbiguousEngine, strict=True) -def test_ambiguous_field_non_strict_routes_to_orchestrator(caplog): +def test_ambiguous_field_non_strict_routes_to_orchestrator(): """strict=False logs ERROR but routes the ambiguous field to orchestrator.""" @dataclass class _AmbiguousEngine: deploy_config: str | None = None - with caplog.at_level(logging.ERROR, logger="vllm_omni.engine.arg_utils"): + import io + + vllm_logger = logging.getLogger("vllm") + orig_handlers = list(vllm_logger.handlers) + vllm_logger.handlers.clear() + buf = io.StringIO() + handler = logging.StreamHandler(buf) + handler.setLevel(logging.ERROR) + vllm_logger.addHandler(handler) + try: orch, engine = split_kwargs({"deploy_config": "x"}, engine_cls=_AmbiguousEngine, strict=False) - assert orch.deploy_config == "x" - assert "deploy_config" not in engine - assert any("both OrchestratorArgs" in r.message for r in caplog.records) + assert orch.deploy_config == "x" + assert "deploy_config" not in engine + assert "both OrchestratorArgs" in buf.getvalue() + finally: + vllm_logger.handlers.clear() + for h in orig_handlers: + vllm_logger.addHandler(h) # Sentinel-default precedence invariants (#3035) @@ -426,24 +463,20 @@ def test_help_text_preserves_default_after_nullify(): _OMNIENGINEARGS_USER_INPUT_FIELDS = frozenset( { - "model_stage", "model_arch", "engine_output_type", "hf_config_name", "custom_process_next_stage_input_func", "subtalker_sampling_params", - "async_chunk", "omni_kv_config", "quantization_config", "worker_type", "task_type", "worker_cls", - "enable_sleep_mode", "omni_master_address", "omni_master_port", "stage_configs_path", "output_modalities", - "log_stats", "custom_pipeline_args", } ) diff --git a/tools/nightly/generate_nightly_perf_html.py b/tools/nightly/generate_nightly_perf_html.py index 73e245ab1..11e9b0961 100644 --- a/tools/nightly/generate_nightly_perf_html.py +++ b/tools/nightly/generate_nightly_perf_html.py @@ -250,7 +250,7 @@ def _collect_diffusion_records(input_dir: str) -> list[dict[str, Any]]: def _sort_omni_records(records: list[dict[str, Any]]) -> list[dict[str, Any]]: - by_date_desc = sorted(records, key=lambda r: (r.get("date") or ""), reverse=True) + by_date_desc = sorted(records, key=lambda r: r.get("date") or "", reverse=True) return sorted( by_date_desc, key=lambda r: ( @@ -264,8 +264,8 @@ def _sort_omni_records(records: list[dict[str, Any]]) -> list[dict[str, Any]]: def _sort_diffusion_records(records: list[dict[str, Any]]) -> list[dict[str, Any]]: - by_date_desc = sorted(records, key=lambda r: (r.get("date") or ""), reverse=True) - return sorted(by_date_desc, key=lambda r: (r.get("test_name") or "")) + by_date_desc = sorted(records, key=lambda r: r.get("date") or "", reverse=True) + return sorted(by_date_desc, key=lambda r: r.get("test_name") or "") def _ensure_parent_dir(path: str) -> None: diff --git a/vllm_omni/core/sched/omni_ar_scheduler.py b/vllm_omni/core/sched/omni_ar_scheduler.py index 5c0c489d1..1fac5518e 100644 --- a/vllm_omni/core/sched/omni_ar_scheduler.py +++ b/vllm_omni/core/sched/omni_ar_scheduler.py @@ -24,9 +24,11 @@ from vllm_omni.core.sched.omni_scheduling_coordinator import ( OmniSchedulingCoordinator, uses_full_payload_input_coordinator, ) +from vllm_omni.core.sched.utils import omni_routed_experts_for_request from vllm_omni.distributed.omni_connectors.transfer_adapter.chunk_transfer_adapter import ( OmniChunkTransferAdapter, ) +from vllm_omni.engine import OmniEngineCoreOutput from vllm_omni.engine.serialization import deserialize_additional_information from vllm_omni.outputs import OmniConnectorOutput @@ -321,7 +323,7 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): req = self.requests.get(req_id) if req is not None and not req.is_finished(): outputs[req.client_index].append( - EngineCoreOutput( + OmniEngineCoreOutput( request_id=req_id, new_token_ids=[], kv_transfer_params={"kv_ready": True}, @@ -417,11 +419,8 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): stopped = True if stopped: - 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] + if model_runner_output.routed_experts is not None: + routed_experts = omni_routed_experts_for_request(model_runner_output.routed_experts, request) # Capture finish_reason BEFORE _handle_stopped_request, which may # reset the status to WAITING for streaming requests that continue. @@ -452,7 +451,7 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): stopped_preempted_reqs.add(request) # Extract sample logprobs if needed. - if request.sampling_params is not None and request.sampling_params.logprobs is not None and logprobs: + if request.sampling_params is not None and request.sampling_params.num_logprobs is not None and logprobs: new_logprobs = logprobs.slice_request(req_index, len(new_token_ids)) if num_nans_in_logits is not None and req_id in num_nans_in_logits: @@ -463,7 +462,7 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): if new_token_ids or pooler_output is not None or kv_transfer_params or stopped: # Add EngineCoreOutput for this Request. outputs[request.client_index].append( - EngineCoreOutput( + OmniEngineCoreOutput( request_id=req_id, new_token_ids=new_token_ids, finish_reason=finish_reason, @@ -501,7 +500,7 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): self.finish_requests(failed_kv_load_req_ids, RequestStatus.FINISHED_ERROR) for request in requests: outputs[request.client_index].append( - EngineCoreOutput( + OmniEngineCoreOutput( request_id=request.request_id, new_token_ids=[], finish_reason=request.get_finished_reason(), @@ -633,6 +632,7 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): """ req_id = session.request_id self._new_prompt_len_snapshot[req_id] = len(update.prompt_token_ids) + session.discard_latest_async_tokens = False outstanding_async_tokens = getattr(session, "num_output_placeholders", 0) if outstanding_async_tokens > 0: # Async scheduling may already have sampled the previous diff --git a/vllm_omni/core/sched/omni_generation_scheduler.py b/vllm_omni/core/sched/omni_generation_scheduler.py index dae3be8ae..1c43b958a 100644 --- a/vllm_omni/core/sched/omni_generation_scheduler.py +++ b/vllm_omni/core/sched/omni_generation_scheduler.py @@ -29,9 +29,11 @@ from vllm_omni.core.sched.omni_scheduling_coordinator import ( uses_full_payload_input_coordinator, ) from vllm_omni.core.sched.output import OmniCachedRequestData, OmniNewRequestData +from vllm_omni.core.sched.utils import omni_routed_experts_for_request from vllm_omni.distributed.omni_connectors.transfer_adapter.chunk_transfer_adapter import ( OmniChunkTransferAdapter, ) +from vllm_omni.engine import OmniEngineCoreOutput from vllm_omni.outputs import OmniConnectorOutput, OmniModelRunnerOutput logger = init_logger(__name__) @@ -499,11 +501,8 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler): stopped = True if stopped: - 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] + if model_runner_output.routed_experts is not None: + routed_experts = omni_routed_experts_for_request(model_runner_output.routed_experts, request) finish_reason = request.get_finished_reason() finished = self._handle_stopped_request(request) if not finished: @@ -524,7 +523,7 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler): stopped_running_reqs.add(request) # Extract sample logprobs if needed. - if request.sampling_params is not None and request.sampling_params.logprobs is not None and logprobs: + if request.sampling_params is not None and request.sampling_params.num_logprobs is not None and logprobs: new_logprobs = logprobs.slice_request(req_index, len(new_token_ids)) if new_token_ids and self.structured_output_manager.should_advance(request): @@ -543,7 +542,7 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler): prompt_logprobs_tensors = prompt_logprobs_dict.get(req_id) if new_token_ids or pooler_output is not None or kv_transfer_params or stopped: outputs[request.client_index].append( - EngineCoreOutput( + OmniEngineCoreOutput( request_id=req_id, new_token_ids=new_token_ids, finish_reason=finish_reason, @@ -580,7 +579,7 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler): getattr(request, "external_req_id", None), ) outputs[request.client_index].append( - EngineCoreOutput( + OmniEngineCoreOutput( request_id=request.request_id, new_token_ids=[], finish_reason=finish_reason, @@ -607,7 +606,7 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler): self.finish_requests(failed_kv_load_req_ids, RequestStatus.FINISHED_ERROR) for request in requests: outputs[request.client_index].append( - EngineCoreOutput( + OmniEngineCoreOutput( request_id=request.request_id, new_token_ids=[], finish_reason=request.get_finished_reason(), diff --git a/vllm_omni/core/sched/utils.py b/vllm_omni/core/sched/utils.py new file mode 100644 index 000000000..7e0c16bd4 --- /dev/null +++ b/vllm_omni/core/sched/utils.py @@ -0,0 +1,23 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Shared utilities for omni schedulers.""" + +import numpy as np +from vllm.v1.outputs import RoutedExpertsLists + + +def omni_routed_experts_for_request(routed_experts: RoutedExpertsLists, request) -> np.ndarray | None: + """Extract per-request routed experts from RoutedExpertsLists using slot_mapping. + + Matches upstream RoutedExpertsManager.get() pattern — filters routing_data + rows whose slot_mapping entries belong to this request's block_table. + """ + if routed_experts is None: + return None + slots = getattr(request, "block_table", None) + if slots is None: + return None + slot_set = set(slots) + mask = np.isin(routed_experts.slot_mapping, list(slot_set)) + data = routed_experts.routing_data[mask] + return data if data.size > 0 else None diff --git a/vllm_omni/diffusion/diffusion_engine.py b/vllm_omni/diffusion/diffusion_engine.py index b4bfecb77..1805b5dc3 100644 --- a/vllm_omni/diffusion/diffusion_engine.py +++ b/vllm_omni/diffusion/diffusion_engine.py @@ -867,7 +867,7 @@ class DiffusionEngine: worker_thread.join(timeout=10) if worker_thread.is_alive(): logger.warning( - "Worker thread did not terminate within 10s; scheduler and executor shutdown will be skipped." + "Worker thread did not terminate within 10s; scheduler and executor shutdown will be deferred." ) return else: diff --git a/vllm_omni/diffusion/model_loader/hub_prefetch.py b/vllm_omni/diffusion/model_loader/hub_prefetch.py index c4efb0eff..d60e1734e 100644 --- a/vllm_omni/diffusion/model_loader/hub_prefetch.py +++ b/vllm_omni/diffusion/model_loader/hub_prefetch.py @@ -64,6 +64,7 @@ import contextlib import errno import logging import os +import time from collections.abc import Iterable, Iterator logger = logging.getLogger(__name__) @@ -112,9 +113,43 @@ def _safe_repo_filename(model: str) -> str: return model.replace("/", "__").replace(os.sep, "__") + ".lock" +def _dotfile_lock_acquire(lock_dir: str, model: str, timeout: float = 300.0, poll_interval: float = 0.5) -> bool: + """Acquire an exclusive lock via atomic directory creation. + + ``os.makedirs(path, exist_ok=False)`` is atomic on POSIX filesystems + (including Lustre and NFS), making it a reliable fallback when + ``fcntl.flock`` is unsupported (``ENOLCK`` on FSx/Lustre mounts + without the ``flock`` mount option). + + Returns ``True`` if the lock was acquired, ``False`` on timeout. + The caller must call ``os.rmdir(lock_path)`` to release. + """ + lock_path = os.path.join(lock_dir, _safe_repo_filename(model) + ".dir") + + deadline = time.monotonic() + timeout + while True: + try: + os.makedirs(lock_path, exist_ok=False) + logger.info("Acquired dotfile prefetch lock for %s at %s", model, lock_path) + return True + except FileExistsError: + if time.monotonic() >= deadline: + logger.warning( + "Timed out waiting for dotfile prefetch lock %s after %.0fs; proceeding unlocked", + lock_path, + timeout, + ) + return False + time.sleep(poll_interval) + + @contextlib.contextmanager def _repo_prefetch_lock(model: str) -> Iterator[None]: - """Hold an exclusive ``fcntl.flock`` keyed on the HF repo id. + """Hold an exclusive lock keyed on the HF repo id. + + Tries ``fcntl.flock`` first; falls back to an atomic directory-creation + dotfile lock when flock is unsupported (e.g. FSx for Lustre without the + ``flock`` mount option, which returns ``ENOLCK``). Why we still need a node-wide lock on top of ``snapshot_download``'s own per-blob ``.lock`` files: @@ -131,71 +166,93 @@ def _repo_prefetch_lock(model: str) -> Iterator[None]: * Multiple unrelated processes on the same node (concurrent CI jobs, multiple ``OmniServer`` instances, every ``DiffusionWorker`` subprocess from the multiproc executor) can enter this code path - in parallel. ``flock`` on a per-repo file makes the whole + in parallel. An exclusive lock makes the whole ``snapshot_download(...)`` call a true node-wide critical section, so the *first* entrant fully populates the cache and every subsequent ``from_pretrained`` sees a warm, complete tree. - Best-effort: if we cannot create / acquire the lock (no ``fcntl``, - read-only FS, NFS without flock support) we log and run the download - anyway. Worst case behaviour reverts to "snapshot_download + - transformers v5 race" which is exactly the pre-fix state. + Best-effort: if we cannot acquire any lock (no ``fcntl``, read-only + FS, dotfile lock timeout) we log and run the download anyway. Worst + case behaviour reverts to "snapshot_download + transformers v5 race" + which is exactly the pre-fix state. """ + lock_dir = None + lock_path = None + dotfile_held = None + fd = None + flock_held = False + + # --- fcntl.flock path --- try: import fcntl # type: ignore[import-not-found] except ImportError: # pragma: no cover - non-POSIX (Windows) - logger.debug("fcntl unavailable on this platform; running prefetch unlocked") - yield - return + fcntl = None - try: - lock_dir = _node_lock_dir() - except OSError as exc: - logger.warning("Could not allocate lock dir for prefetch of %s (%s); skipping flock", model, exc) - yield - return - - lock_path = os.path.join(lock_dir, _safe_repo_filename(model)) - - try: - fd = os.open(lock_path, os.O_RDWR | os.O_CREAT, 0o644) - except OSError as exc: - logger.warning("Could not open prefetch lock %s (%s); proceeding without flock", lock_path, exc) - yield - return - - locked = False - try: + if fcntl is not None: try: - fcntl.flock(fd, fcntl.LOCK_EX) - locked = True - logger.debug("Acquired prefetch flock for %s at %s", model, lock_path) + lock_dir = _node_lock_dir() except OSError as exc: - # FS without flock support (ENOLCK / EOPNOTSUPP / EINVAL). - # EACCES on some Lustre configs without ``flock`` mount opt. - if exc.errno in (errno.ENOLCK, errno.EOPNOTSUPP, errno.EACCES, errno.EINVAL): - logger.warning( - "fcntl.flock not supported on '%s' (%s: %s); running prefetch unlocked", - lock_path, - type(exc).__name__, - exc, - ) - else: - raise + logger.warning("Could not allocate lock dir for prefetch of %s (%s); skipping flock", model, exc) + fcntl = None # force dotfile fallback + + if fcntl is not None and lock_dir is not None: + lock_path = os.path.join(lock_dir, _safe_repo_filename(model)) + try: + fd = os.open(lock_path, os.O_RDWR | os.O_CREAT, 0o644) + except OSError as exc: + logger.warning("Could not open prefetch lock %s (%s); falling back to dotfile lock", lock_path, exc) + fd = None + + if fd is not None: + try: + fcntl.flock(fd, fcntl.LOCK_EX) + flock_held = True + logger.info("Acquired flock prefetch lock for %s at %s", model, lock_path) + except OSError as exc: + if exc.errno in (errno.ENOLCK, errno.EOPNOTSUPP, errno.EACCES, errno.EINVAL): + logger.warning( + "fcntl.flock not supported on '%s' (%s: %s); falling back to dotfile lock", + lock_path, + type(exc).__name__, + exc, + ) + os.close(fd) + fd = None + else: + raise + + # --- dotfile fallback --- + if not flock_held: + if lock_dir is None: + try: + lock_dir = _node_lock_dir() + except OSError as exc: + logger.warning("Could not allocate lock dir (%s); running prefetch unlocked", exc) + yield + return + + if _dotfile_lock_acquire(lock_dir, model): + dotfile_held = os.path.join(lock_dir, _safe_repo_filename(model) + ".dir") + + try: yield finally: - if locked: + if flock_held and fd is not None: with contextlib.suppress(OSError): fcntl.flock(fd, fcntl.LOCK_UN) - with contextlib.suppress(OSError): - os.close(fd) + if fd is not None: + with contextlib.suppress(OSError): + os.close(fd) + if dotfile_held is not None: + with contextlib.suppress(OSError): + os.rmdir(dotfile_held) def prefetch_subfolders( model: str, subfolders: Iterable[str], *, - local_files_only: bool = False, + local_files_only: bool | None = None, include_root_metadata: bool = True, ) -> None: """Materialise ``model``'s ``subfolders`` in the HF cache before loading. @@ -206,17 +263,21 @@ def prefetch_subfolders( subfolders: Iterable of subfolder names (e.g. ``["text_encoder", "vae"]``) whose contents need to be fully present before any worker calls ``from_pretrained(subfolder=...)``. - local_files_only: When True, skip the prefetch entirely. The caller - has explicitly promised the cache is already populated (as happens - for local model checkouts), so hitting the network would defeat - the intent and may fail in air-gapped environments. + local_files_only: When ``True``, skip the prefetch entirely. + When ``None`` (default), auto-detect: skip if *model* is a + local directory, run otherwise. include_root_metadata: When True, also pull ``*.json`` at the repo root so ``model_index.json`` / ``config.json`` resolution during ``from_pretrained`` also hits a warm cache. """ + if local_files_only is None: + local_files_only = os.path.isdir(model) + if local_files_only or not model or os.path.isdir(model): return + logger.info("Prefetching %s subfolders: %s", model, list(subfolders)) + try: from huggingface_hub import snapshot_download except ImportError: # pragma: no cover - huggingface_hub is a hard dep @@ -251,6 +312,7 @@ def prefetch_subfolders( repo_id=model, allow_patterns=allow_patterns, ) + logger.info("Prefetch complete for %s", model) except Exception as exc: # Best-effort: propagate only via logging. The subsequent # ``from_pretrained`` call will raise a clearer, call-site-specific @@ -312,3 +374,31 @@ def _looks_like_auth_error(exc: BaseException) -> bool: # only load-bearing signal. msg = str(exc).lower() return "401 client error" in msg or "403 client error" in msg or "gatedrepo" in msg + + +def retry_on_missing_shard(load_fn, *, max_retries: int = 3, base_delay: float = 5.0): + """Call *load_fn* with retry on the transformers v5 shard-resolution race. + + When the prefetch lock cannot be acquired (e.g. flock unsupported on + the filesystem and dotfile lock times out), ``from_pretrained`` may + still hit the ``cached_files`` race. This wrapper retries with + exponential backoff when the OSError message matches the specific + "does not appear to have a file named" pattern. + """ + for attempt in range(max_retries): + try: + return load_fn() + except OSError as exc: + if "does not appear to have a file" not in str(exc): + raise + if attempt == max_retries - 1: + raise + delay = base_delay * (attempt + 1) + logger.warning( + "from_pretrained failed with shard-resolution race (%s); retrying in %.1fs (attempt %d/%d)", + exc, + delay, + attempt + 1, + max_retries, + ) + time.sleep(delay) diff --git a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py index 6f7a6bb58..61e50b061 100644 --- a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py +++ b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py @@ -26,7 +26,7 @@ from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_qwenimage import DistributedAutoencoderKLQwenImage from vllm_omni.diffusion.distributed.utils import get_local_device from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader -from vllm_omni.diffusion.model_loader.hub_prefetch import prefetch_subfolders +from vllm_omni.diffusion.model_loader.hub_prefetch import prefetch_subfolders, retry_on_missing_shard from vllm_omni.diffusion.models.dmd2 import DMD2PipelineMixin from vllm_omni.diffusion.models.qwen_image.cfg_parallel import ( QwenImageCFGParallelMixin, @@ -275,21 +275,22 @@ class QwenImagePipeline(nn.Module, QwenImageCFGParallelMixin, DiffusionPipelineP self.device = get_local_device() model = od_config.model # Check if model is a local path - local_files_only = os.path.exists(model) + local_files_only = os.path.isdir(model) # See pipeline_qwen_image_edit_plus: guard against transformers v5 # multi-worker race on partial subfolder shard sets (Buildkite #1043). prefetch_subfolders( model, ["scheduler", "text_encoder", "vae", "tokenizer"], - local_files_only=local_files_only, ) self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained( model, subfolder="scheduler", local_files_only=local_files_only ) - self.text_encoder = Qwen2_5_VLForConditionalGeneration.from_pretrained( - model, subfolder="text_encoder", local_files_only=local_files_only + self.text_encoder = retry_on_missing_shard( + lambda: Qwen2_5_VLForConditionalGeneration.from_pretrained( + model, subfolder="text_encoder", local_files_only=local_files_only + ) ) # Qwen2.5-VL ships a vision tower that text-to-image does not use. # Drop it while the model is still on CPU, before moving to GPU, so diff --git a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit.py b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit.py index 136604bb4..c29fcb8ba 100644 --- a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit.py +++ b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit.py @@ -27,7 +27,7 @@ from vllm.model_executor.models.utils import AutoWeightsLoader from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig from vllm_omni.diffusion.distributed.utils import get_local_device from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader -from vllm_omni.diffusion.model_loader.hub_prefetch import prefetch_subfolders +from vllm_omni.diffusion.model_loader.hub_prefetch import prefetch_subfolders, retry_on_missing_shard from vllm_omni.diffusion.models.interface import SupportImageInput from vllm_omni.diffusion.models.qwen_image.cfg_parallel import ( QwenImageCFGParallelMixin, @@ -243,22 +243,23 @@ class QwenImageEditPipeline(nn.Module, SupportImageInput, QwenImageCFGParallelMi model = od_config.model # Check if model is a local path - local_files_only = os.path.exists(model) + local_files_only = os.path.isdir(model) # See pipeline_qwen_image_edit_plus: guard against transformers v5 # multi-worker race on partial subfolder shard sets (Buildkite #1043). prefetch_subfolders( model, ["scheduler", "text_encoder", "vae", "tokenizer", "processor"], - local_files_only=local_files_only, ) self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained( model, subfolder="scheduler", local_files_only=local_files_only ) - self.text_encoder = Qwen2_5_VLForConditionalGeneration.from_pretrained( - model, subfolder="text_encoder", local_files_only=local_files_only - ).to(self.device) + self.text_encoder = retry_on_missing_shard( + lambda: Qwen2_5_VLForConditionalGeneration.from_pretrained( + model, subfolder="text_encoder", local_files_only=local_files_only + ).to(self.device) + ) self.vae = AutoencoderKLQwenImage.from_pretrained(model, subfolder="vae", local_files_only=local_files_only).to( self.device @@ -266,8 +267,8 @@ class QwenImageEditPipeline(nn.Module, SupportImageInput, QwenImageCFGParallelMi transformer_kwargs = get_transformer_config_kwargs(od_config.tf_model_config, QwenImageTransformer2DModel) self.transformer = QwenImageTransformer2DModel(od_config=od_config, **transformer_kwargs) self.tokenizer = Qwen2Tokenizer.from_pretrained(model, subfolder="tokenizer", local_files_only=local_files_only) - self.processor = Qwen2VLProcessor.from_pretrained( - model, subfolder="processor", local_files_only=local_files_only + self.processor = retry_on_missing_shard( + lambda: Qwen2VLProcessor.from_pretrained(model, subfolder="processor", local_files_only=local_files_only) ) self.stage = None diff --git a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit_plus.py b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit_plus.py index 326f5d14f..ecdddc91f 100644 --- a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit_plus.py +++ b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit_plus.py @@ -25,7 +25,7 @@ from vllm.model_executor.models.utils import AutoWeightsLoader from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig from vllm_omni.diffusion.distributed.utils import get_local_device from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader -from vllm_omni.diffusion.model_loader.hub_prefetch import prefetch_subfolders +from vllm_omni.diffusion.model_loader.hub_prefetch import prefetch_subfolders, retry_on_missing_shard from vllm_omni.diffusion.model_metadata import QWEN_IMAGE_EDIT_PLUS_MAX_INPUT_IMAGES from vllm_omni.diffusion.models.interface import SupportImageInput from vllm_omni.diffusion.models.qwen_image.cfg_parallel import ( @@ -208,7 +208,7 @@ class QwenImageEditPlusPipeline( model = od_config.model # Check if model is a local path - local_files_only = os.path.exists(model) + local_files_only = os.path.isdir(model) # Defend against a transformers v5 multi-worker race where a peer # worker's partially-written shard set makes our from_pretrained @@ -219,15 +219,16 @@ class QwenImageEditPlusPipeline( prefetch_subfolders( model, ["scheduler", "text_encoder", "vae", "tokenizer", "processor"], - local_files_only=local_files_only, ) self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained( model, subfolder="scheduler", local_files_only=local_files_only ) - self.text_encoder = Qwen2_5_VLForConditionalGeneration.from_pretrained( - model, subfolder="text_encoder", local_files_only=local_files_only - ).to(self.device) + self.text_encoder = retry_on_missing_shard( + lambda: Qwen2_5_VLForConditionalGeneration.from_pretrained( + model, subfolder="text_encoder", local_files_only=local_files_only + ).to(self.device) + ) self.vae = AutoencoderKLQwenImage.from_pretrained(model, subfolder="vae", local_files_only=local_files_only).to( self.device @@ -236,8 +237,8 @@ class QwenImageEditPlusPipeline( transformer_kwargs = get_transformer_config_kwargs(od_config.tf_model_config, QwenImageTransformer2DModel) self.transformer = QwenImageTransformer2DModel(od_config=od_config, **transformer_kwargs) self.tokenizer = Qwen2Tokenizer.from_pretrained(model, subfolder="tokenizer", local_files_only=local_files_only) - self.processor = Qwen2VLProcessor.from_pretrained( - model, subfolder="processor", local_files_only=local_files_only + self.processor = retry_on_missing_shard( + lambda: Qwen2VLProcessor.from_pretrained(model, subfolder="processor", local_files_only=local_files_only) ) self.stage = None diff --git a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_layered.py b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_layered.py index dffb9e8f5..df44404e8 100644 --- a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_layered.py +++ b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_layered.py @@ -24,7 +24,7 @@ from vllm.model_executor.models.utils import AutoWeightsLoader from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig from vllm_omni.diffusion.distributed.utils import get_local_device from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader -from vllm_omni.diffusion.model_loader.hub_prefetch import prefetch_subfolders +from vllm_omni.diffusion.model_loader.hub_prefetch import prefetch_subfolders, retry_on_missing_shard from vllm_omni.diffusion.models.interface import SupportImageInput from vllm_omni.diffusion.models.qwen_image.autoencoder_kl_qwenimage import ( AutoencoderKLQwenImage, @@ -216,29 +216,30 @@ class QwenImageLayeredPipeline(nn.Module, SupportImageInput, QwenImageCFGParalle self.device = get_local_device() model = od_config.model # Check if model is a local path - local_files_only = os.path.exists(model) + local_files_only = os.path.isdir(model) # See pipeline_qwen_image_edit_plus: guard against transformers v5 # multi-worker race on partial subfolder shard sets (Buildkite #1043). prefetch_subfolders( model, ["scheduler", "text_encoder", "vae", "tokenizer", "processor"], - local_files_only=local_files_only, ) # modules keep same as transformers & diffusers self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained( model, subfolder="scheduler", local_files_only=local_files_only ) - self.text_encoder = Qwen2_5_VLForConditionalGeneration.from_pretrained( - model, subfolder="text_encoder", local_files_only=local_files_only - ).to(self.device) + self.text_encoder = retry_on_missing_shard( + lambda: Qwen2_5_VLForConditionalGeneration.from_pretrained( + model, subfolder="text_encoder", local_files_only=local_files_only + ).to(self.device) + ) self.vae = AutoencoderKLQwenImage.from_pretrained(model, subfolder="vae", local_files_only=local_files_only).to( self.device ) self.tokenizer = Qwen2Tokenizer.from_pretrained(model, subfolder="tokenizer", local_files_only=local_files_only) - self.processor = Qwen2VLProcessor.from_pretrained( - model, subfolder="processor", local_files_only=local_files_only + self.processor = retry_on_missing_shard( + lambda: Qwen2VLProcessor.from_pretrained(model, subfolder="processor", local_files_only=local_files_only) ) # modules re-implemented for vLLM-Omni diff --git a/vllm_omni/distributed/omni_coordinator/omni_coordinator.py b/vllm_omni/distributed/omni_coordinator/omni_coordinator.py index 8f75f841d..bbedcfeaa 100644 --- a/vllm_omni/distributed/omni_coordinator/omni_coordinator.py +++ b/vllm_omni/distributed/omni_coordinator/omni_coordinator.py @@ -43,17 +43,20 @@ class OmniCoordinator: heartbeat_timeout: Seconds before a replica is considered unhealthy if no heartbeat / update is received. """ - self._router_zmq_addr = router_zmq_addr - self._pub_zmq_addr = pub_zmq_addr self._heartbeat_timeout = heartbeat_timeout # Dedicated ZMQ context for this coordinator instance. self._ctx = zmq.Context() self._router = self._ctx.socket(zmq.ROUTER) - self._router.bind(self._router_zmq_addr) + self._router.bind(router_zmq_addr) + # Recover the actual bound address (port=0 → real port). + self.router_zmq_addr = self._router.getsockopt_string(zmq.LAST_ENDPOINT) + self._pub_zmq_addr = pub_zmq_addr # keep original for internal use self._pub = self._ctx.socket(zmq.PUB) - self._pub.bind(self._pub_zmq_addr) + self._pub.bind(pub_zmq_addr) + # Recover the actual bound address (port=0 → real port). + self.pub_zmq_addr = self._pub.getsockopt_string(zmq.LAST_ENDPOINT) self._replicas: dict[str, ReplicaInfo] = {} self._lock = threading.Lock() diff --git a/vllm_omni/engine/arg_utils.py b/vllm_omni/engine/arg_utils.py index c46222899..8d4388b5e 100644 --- a/vllm_omni/engine/arg_utils.py +++ b/vllm_omni/engine/arg_utils.py @@ -144,7 +144,8 @@ class OmniEngineArgs(EngineArgs): force_cutlass_fp8: bool | None = None worker_type: str | None = None task_type: str | None = None - worker_cls: str = None + worker_cls: str = None # type: ignore[assignment] # Upstream default is "auto"; omni resolves + # in __post_init__ based on worker_type (ar/generation), so None is safe here. enable_sleep_mode: bool = False omni: bool = False @@ -444,6 +445,48 @@ class OrchestratorArgs: # === Pre-built Objects === parallel_config: Any = None + # === Diffusion model config === + num_gpus: int | None = None + model_class_name: str | None = None + diffusion_load_format: str | None = None + diffusers_load_kwargs: str = "{}" + diffusers_call_kwargs: str = "{}" + ulysses_degree: int | None = None + ulysses_mode: str = "strict" + ring_degree: int | None = None + diffusion_quantization_config: str | None = None + use_hsdp: bool = False + hsdp_shard_size: int = -1 + hsdp_replicate_size: int = 1 + diffusion_attention_backend: str | None = None + diffusion_attention_config: str | None = None + cache_backend: str = "none" + cache_config: str | None = None + enable_cache_dit_summary: bool = False + step_execution: bool = False + vae_use_slicing: bool = False + vae_use_tiling: bool = False + enable_multithread_weight_load: bool = True + num_weight_load_threads: int = 4 + enable_cpu_offload: bool = False + enable_layerwise_offload: bool = False + boundary_ratio: float | None = None + flow_shift: float | None = None + diffusion_kv_cache_dtype: str | None = None + diffusion_kv_cache_skip_steps: str | None = None + diffusion_kv_cache_skip_layers: str | None = None + cfg_parallel_size: int = 1 + vae_patch_parallel_size: int = 1 + default_sampling_params: str | None = None + max_generated_image_size: int | None = None + tts_max_instructions_length: int | None = None + enable_diffusion_pipeline_profiler: bool = False + enable_ar_profiler: bool = False + auxiliary_text_encoder: str | None = None + log_file: str | None = None + replica_id: int | None = None + omni_replica_address: str | None = None + # === Multi-stage guards === # --tokenizer is captured by the orchestrator and forwarded to stages # only when the stage does not define tokenizer/tokenizer_subdir itself. @@ -459,6 +502,8 @@ SHARED_FIELDS: frozenset[str] = frozenset( "stage_id", # orch: route (headless); engine: identity "log_stats", # both want the flag "stage_configs_path", # orch: load legacy YAML; engine: may reference for validation + "async_chunk", # orch: read from CLI, redistribute; engine: per-stage flag + "tokenizer", # orch: detect model type; engine: tokenization } ) diff --git a/vllm_omni/engine/async_omni_engine.py b/vllm_omni/engine/async_omni_engine.py index f6eb77e30..939f9b224 100644 --- a/vllm_omni/engine/async_omni_engine.py +++ b/vllm_omni/engine/async_omni_engine.py @@ -2124,7 +2124,7 @@ class AsyncOmniEngine: kwargs.get("diffusion_attention_config"), attention_backend=kwargs.get("diffusion_attention_backend"), ) - quantization_config = kwargs.get("quantization_config") + quantization_config = kwargs.get("diffusion_quantization_config") if quantization_config is not None: if ( not hasattr(cfg.engine_args, "quantization_config") diff --git a/vllm_omni/engine/output_processor.py b/vllm_omni/engine/output_processor.py index 7f70037f8..5fb46ed0d 100644 --- a/vllm_omni/engine/output_processor.py +++ b/vllm_omni/engine/output_processor.py @@ -1,11 +1,7 @@ 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 @@ -172,7 +168,6 @@ class OmniRequestState(RequestState): finish_reason: FinishReason | None, stop_reason: int | str | None, kv_transfer_params: dict[str, Any] | None = None, - routed_experts: np.ndarray | None = None, ) -> OmniRequestOutput | PoolingRequestOutput | None: """Create a request output from generation results. @@ -199,7 +194,6 @@ class OmniRequestState(RequestState): finish_reason, stop_reason, kv_transfer_params, - routed_experts, ) finished = finish_reason is not None @@ -234,16 +228,7 @@ class OmniRequestState(RequestState): external_req_id = self.external_req_id - # 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) + output = self._new_completion_output(new_token_ids, finish_reason, stop_reason) if self.parent_req is None: outputs = [output] @@ -258,7 +243,6 @@ class OmniRequestState(RequestState): outputs, finished, kv_transfer_params, - prompt_routed_experts, ) def _new_completion_output( @@ -266,10 +250,11 @@ class OmniRequestState(RequestState): token_ids: list[int], finish_reason: FinishReason | None, stop_reason: int | str | None, - routed_experts: np.ndarray | None = None, ) -> Any: # Reuse base text/logprobs logic, then annotate with pooling_result. - base_output = super()._new_completion_output(token_ids, finish_reason, stop_reason, routed_experts) + # Note: upstream _new_completion_output no longer accepts routed_experts + # as a parameter; it reads from ``self.routed_experts_chunks`` internally. + base_output = super()._new_completion_output(token_ids, finish_reason, stop_reason) # Inter-stage processors need the full cumulative token sequence. # In DELTA mode, base_output.token_ids only has the latest step's diff --git a/vllm_omni/engine/stage_engine_core_proc.py b/vllm_omni/engine/stage_engine_core_proc.py index 0bb700369..32c92ff68 100644 --- a/vllm_omni/engine/stage_engine_core_proc.py +++ b/vllm_omni/engine/stage_engine_core_proc.py @@ -101,6 +101,11 @@ class StageEngineCoreProc(EngineCoreProc): set_death_signal(signal.SIGTERM) set_process_title(f"StageEngineCoreProc_{stage_label}_replica{omni_replica_id}_DP{dp_rank}") decorate_logs() + # Workaround for flashinfer/jit-cache version mismatch in CI. + # The parent process handles this gracefully via ring_globals.py, + # but the subprocess hits an unprotected import in TopKTopPSampler. + # Setting this env var allows the same graceful fallback to work. + os.environ.setdefault("FLASHINFER_DISABLE_VERSION_CHECK", "1") os.environ["VLLM_OMNI_REPLICA_ID"] = str(max(int(omni_replica_id), 0)) engine_core = StageEngineCoreProc( diff --git a/vllm_omni/engine/stage_init_utils.py b/vllm_omni/engine/stage_init_utils.py index 50e750c68..75024abdb 100644 --- a/vllm_omni/engine/stage_init_utils.py +++ b/vllm_omni/engine/stage_init_utils.py @@ -860,6 +860,45 @@ def build_stage0_input_processor(stage_vllm_config: Any) -> InputProcessor: return input_processor +def _cleanup_stale_lock_if_dead(lock_file: str) -> bool: + """If *lock_file* exists and its recorded PID is dead, unlink the file. + + Returns ``True`` if the stale lock was cleaned up (caller should retry), + ``False`` otherwise (lock holder appears alive, or file could not be read). + """ + try: + with open(lock_file) as fh: + content = fh.read().strip() + if not content: + return False + pid = int(content) + except (OSError, ValueError): + return False + + # Check whether the PID is still alive. + try: + os.kill(pid, 0) + except ProcessLookupError: + # PID does not exist — stale lock. + logger.info( + "Removing stale device lock %s (PID %s is dead)", + lock_file, + pid, + ) + try: + os.unlink(lock_file) + return True + except OSError: + logger.debug("Failed to unlink stale lock %s", lock_file) + return False + except PermissionError: + # PID exists but we cannot signal it (different user) — treat as alive. + return False + + # PID is alive — legitimate lock holder. + return False + + def acquire_device_locks( stage_id: int, engine_args_dict: dict[str, Any], @@ -938,6 +977,7 @@ def acquire_device_locks( for device_id in devices_to_lock: lock_file = f"/tmp/vllm_omni_device_{device_id}_init.lock" lock_acquired = False + already_cleaned_stale = False # only try stale cleanup once per device while not lock_acquired: try: @@ -952,6 +992,11 @@ def acquire_device_locks( logger.debug("Acquired exclusive lock for device %s", device_id) except BlockingIOError: os.close(lock_fd) + # Detect and clean stale locks from dead processes. + if not already_cleaned_stale: + already_cleaned_stale = True + if _cleanup_stale_lock_if_dead(lock_file): + continue # retry flock immediately if time.time() - wait_start > stage_init_timeout: logger.warning( "Timeout waiting for device %s initialization lock, proceeding anyway", diff --git a/vllm_omni/engine/stage_pool.py b/vllm_omni/engine/stage_pool.py index effe64705..8305c23d4 100644 --- a/vllm_omni/engine/stage_pool.py +++ b/vllm_omni/engine/stage_pool.py @@ -22,6 +22,7 @@ from vllm_omni.engine.stage_client import ( StagePoolDiffusionClient, StagePoolLLMClient, ) +from vllm_omni.inputs.data import OmniDiffusionSamplingParams from vllm_omni.metrics.stats import StageRequestStats as StageRequestMetrics from vllm_omni.metrics.stats import StageStats from vllm_omni.metrics.utils import count_tokens_from_outputs @@ -533,6 +534,10 @@ class StagePool: ) -> int: """Submit a stage-entry request into this pool.""" params = params_override if params_override is not None else req_state.sampling_params_list[self.stage_id] + # Convert plain vllm SamplingParams for single-stage diffusion models + # that receive sampling params from the user/caller directly. + if self.stage_type == "diffusion" and not isinstance(params, OmniDiffusionSamplingParams): + params = OmniDiffusionSamplingParams() submit_kwargs = dict(submit_kwargs or {}) if self.stage_type == "diffusion": replica_id = await self._pick_or_select( @@ -593,6 +598,8 @@ class StagePool: ) -> int: """Submit a streaming update to an already admitted request.""" params = req_state.sampling_params_list[self.stage_id] + if self.stage_type == "diffusion" and not isinstance(params, OmniDiffusionSamplingParams): + params = OmniDiffusionSamplingParams() replica_id = self.get_bound_replica_id(request_id) if replica_id is None or self.clients[replica_id] is None: replica_id = await self._pick_or_select(request_id) diff --git a/vllm_omni/entrypoints/cli/serve.py b/vllm_omni/entrypoints/cli/serve.py index a9f3210f0..4be487415 100644 --- a/vllm_omni/entrypoints/cli/serve.py +++ b/vllm_omni/entrypoints/cli/serve.py @@ -436,7 +436,7 @@ class OmniServeCommand(CLISubcommand): "Equivalent to setting DiffusionParallelConfig.ring_degree.", ) omni_config_group.add_argument( - "--quantization-config", + "--diffusion-quantization-config", type=json.loads, default=None, help=( diff --git a/vllm_omni/entrypoints/omni_base.py b/vllm_omni/entrypoints/omni_base.py index de4ef327c..9ddd32e9d 100644 --- a/vllm_omni/entrypoints/omni_base.py +++ b/vllm_omni/entrypoints/omni_base.py @@ -463,7 +463,7 @@ class OmniBase(PDDisaggregationMixin): peak_memory_mb=peak_memory_mb, ) - def shutdown(self) -> None: + def shutdown(self, timeout: float | None = None) -> None: logger.info("[%s] Shutting down", self.__class__.__name__) self._shutdown_base() diff --git a/vllm_omni/entrypoints/openai/api_server.py b/vllm_omni/entrypoints/openai/api_server.py index c13001063..4db31efbb 100644 --- a/vllm_omni/entrypoints/openai/api_server.py +++ b/vllm_omni/entrypoints/openai/api_server.py @@ -55,13 +55,8 @@ from vllm.entrypoints.openai.engine.protocol import ( from vllm.entrypoints.openai.models.protocol import BaseModelPath from vllm.entrypoints.openai.models.serving import OpenAIServingModels from vllm.entrypoints.openai.orca_metrics import metrics_header -from vllm.entrypoints.openai.realtime.serving import OpenAIServingRealtime from vllm.entrypoints.openai.responses.serving import OpenAIServingResponses from vllm.entrypoints.openai.server_utils import get_uvicorn_log_config -from vllm.entrypoints.openai.speech_to_text.serving import ( - OpenAIServingTranscription, - OpenAIServingTranslation, -) from vllm.entrypoints.openai.utils import validate_json_request from vllm.entrypoints.pooling.classify.serving import ServingClassification from vllm.entrypoints.pooling.embed.serving import ServingEmbedding as OpenAIServingEmbedding @@ -74,6 +69,13 @@ from vllm.entrypoints.serve.disagg.serving import ServingTokens from vllm.entrypoints.serve.instrumentator.basic import base from vllm.entrypoints.serve.render.serving import OpenAIServingRender from vllm.entrypoints.serve.tokenize.serving import OpenAIServingTokenization +from vllm.entrypoints.speech_to_text.realtime.serving import OpenAIServingRealtime +from vllm.entrypoints.speech_to_text.transcription.serving import ( + OpenAIServingTranscription, +) +from vllm.entrypoints.speech_to_text.translation.serving import ( + OpenAIServingTranslation, +) from vllm.entrypoints.utils import ( create_error_response, load_aware_call, @@ -355,7 +357,7 @@ async def omni_run_server(args, **uvicorn_kwargs) -> None: warnings_module.filterwarnings("ignore", message=".*PydanticSerializationUnexpectedValue.*", category=UserWarning) # Add process-specific prefix to stdout and stderr. - decorate_logs("APIServer") + decorate_logs("APIServer", skip_if_decorated=True) listen_address, sock = setup_openai_server(args) @@ -646,6 +648,19 @@ async def omni_init_app_state( state.vllm_config = vllm_config + # Propagate enable_in_reasoning to the API-server process. The engine core + # runs in a separate process, so the contextvar that backs + # `get_current_vllm_config_or_none()` is None on this stack. Tool parsers + # call `get_enable_structured_outputs_in_reasoning()` during request + # handling and need to see the real flag, otherwise they silently fall + # back to False and mismatch the engine-side bitmask gating. + if vllm_config is not None: + from vllm.tool_parsers.structural_tag_registry import ( + set_enable_structured_outputs_in_reasoning, + ) + + set_enable_structured_outputs_in_reasoning(vllm_config.structured_outputs_config.enable_in_reasoning) + # Get supported tasks supported_tasks: set[str] = {"generate"} if hasattr(engine_client, "get_supported_tasks"): @@ -1531,6 +1546,7 @@ async def generate_images(request: ImageGenerationRequest, raw_request: Request) prompt=request.prompt, extra_body=extra_body, request_id=f"img_gen-{random_uuid()}", + raw_request=raw_request, ) if isinstance(generation_result, ErrorResponse): return JSONResponse( @@ -1937,6 +1953,7 @@ async def edit_images( output_format=output_format, output_compression=output_compression, size=size_str, + raw_request=raw_request, ) if stream and not isinstance(generation_result, ErrorResponse): return StreamingResponse( diff --git a/vllm_omni/entrypoints/openai/realtime_connection.py b/vllm_omni/entrypoints/openai/realtime_connection.py index e04e8cdf9..1bc3df64e 100644 --- a/vllm_omni/entrypoints/openai/realtime_connection.py +++ b/vllm_omni/entrypoints/openai/realtime_connection.py @@ -9,8 +9,8 @@ from uuid import uuid4 import numpy as np from vllm.entrypoints.openai.engine.protocol import UsageInfo -from vllm.entrypoints.openai.realtime.connection import RealtimeConnection as VllmRealtimeConnection -from vllm.entrypoints.openai.realtime.protocol import TranscriptionDelta, TranscriptionDone +from vllm.entrypoints.speech_to_text.realtime.connection import RealtimeConnection as VllmRealtimeConnection +from vllm.entrypoints.speech_to_text.realtime.protocol import TranscriptionDelta, TranscriptionDone from vllm.logger import init_logger from vllm_omni.entrypoints.async_omni import AsyncOmni diff --git a/vllm_omni/entrypoints/openai/serving_chat.py b/vllm_omni/entrypoints/openai/serving_chat.py index f98732984..2abfefe82 100644 --- a/vllm_omni/entrypoints/openai/serving_chat.py +++ b/vllm_omni/entrypoints/openai/serving_chat.py @@ -2574,6 +2574,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): output_format: str = "png", output_compression: int = 100, size: str = "auto", + raw_request: Request | None = None, ) -> tuple[list[Image.Image], dict[str, Any], float, str | None] | ErrorResponse | AsyncIterator[str]: """Generate diffusion images and return raw images plus generation stats.""" if request_id is None: @@ -2635,6 +2636,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): output_format=output_format, output_compression=output_compression, size=size, + raw_request=raw_request, ) result = None @@ -2659,6 +2661,16 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): peak_memory_mb = result.peak_memory_mb cot_output = None + req_out = getattr(result, "request_output", None) + if req_out: + prompt_obj = getattr(req_out, "prompt", None) + if isinstance(prompt_obj, dict): + extra = prompt_obj.get("extra", {}) + if isinstance(extra, dict): + ar_text = extra.get("ar_generated_text") + if isinstance(ar_text, str) and ar_text.strip(): + cot_output = ar_text + req_out = getattr(result, "request_output", None) if req_out: prompt_obj = getattr(req_out, "prompt", None) @@ -2679,6 +2691,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): output_format: str, output_compression: int, size: str, + raw_request: Request | None = None, ) -> AsyncIterator[str]: """Yield image edit SSE chunks from multi-stage diffusion outputs.""" created = int(time.time()) @@ -2730,7 +2743,16 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): logger.error("EngineDeadError during streaming image edit: %s", exc) data = self.create_streaming_error_response(exc) yield f"data: {data}\n\n" - terminate_if_errored(engine=self.engine_client) + if raw_request is not None: + terminate_if_errored( + server=raw_request.app.state.server, + engine=self.engine_client, + ) + else: + logger.warning( + "[OmniOpenAIServingChat] Engine dead during streaming image edit, " + "but cannot terminate server (no raw_request context)." + ) except Exception as exc: logger.exception("Streaming image edit failed: %s", exc) chunk = ImageEditStreamError( diff --git a/vllm_omni/entrypoints/stage_utils.py b/vllm_omni/entrypoints/stage_utils.py index b86c59fcb..e7c91cf4d 100644 --- a/vllm_omni/entrypoints/stage_utils.py +++ b/vllm_omni/entrypoints/stage_utils.py @@ -114,6 +114,24 @@ def _map_device_list(stage_id: int, device_list: list[str], visible_device_list: return list(device_list) logical_ids = [int(device) for device in device_list] + + # If ALL logical device IDs are >= num_visible, none can be 0-based indices + # into the visible list — they must be physical device IDs. This happens + # during multi-replica stage initialization where split_devices_for_replicas + # assigns physical device IDs (e.g. "2") that are not part of the current + # process's CUDA_VISIBLE_DEVICES but are valid physical GPUs. Pass them + # through directly so the replica subprocess inherits the correct device. + if logical_ids and all(d >= num_visible for d in logical_ids): + logger.warning( + "Stage %s has device IDs %s, none of which are < the visible device count %d " + "(visible=%s). Treating them as physical device IDs and passing through.", + stage_id, + device_list, + num_visible, + visible_device_list, + ) + return list(device_list) + mapped_devices = [visible_device_list[idx] for idx in logical_ids if idx < num_visible] mapping_pairs = [ f"{logical_id}->{visible_device_list[logical_id]}" for logical_id in logical_ids if logical_id < num_visible diff --git a/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py b/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py index 2dc81174c..02326f456 100644 --- a/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py +++ b/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py @@ -2,7 +2,6 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import os from collections.abc import Iterable, Mapping, Sequence -from dataclasses import replace from functools import partial from threading import Lock @@ -365,6 +364,12 @@ class CosyVoice3Model( else: raise ValueError(f"Model stage not supported {self.model_stage}") + def get_language_model(self) -> "nn.Module": + """Return the language model for upstream MoE detection.""" + if hasattr(self.model, "get_language_model"): + return self.model.get_language_model() + return self.model + def _create_llm_vllm_config(self, parent_config: VllmConfig) -> VllmConfig: """Create VllmConfig for the inner Qwen2 LLM. @@ -542,17 +547,21 @@ class CosyVoice3Model( if self.model_stage != "cosyvoice3_talker": return None - sampler = getattr(self, "_talker_sampler", None) - if sampler is None: - sampler = Sampler() - self._talker_sampler = sampler - if not self._cosyvoice3_ras_enabled(sampling_metadata): + sampler = getattr(self, "_talker_sampler", None) + if sampler is None: + sampler = Sampler() + self._talker_sampler = sampler return sampler(logits=logits, sampling_metadata=sampling_metadata) logits = logits.to(torch.float32) - sampling_for_processors = replace(sampling_metadata, no_penalties=True) - logits = sampler.apply_logits_processors(logits, sampling_for_processors, predict_bonus_token=False) + # Apply logits processors directly — RAS handles its own repetition + # logic. We avoid instantiating Sampler() here because its import + # chain pulls in flashinfer / GPU deps that fail in CPU-only tests. + if sampling_metadata.allowed_token_ids_mask is not None: + logits.masked_fill_(sampling_metadata.allowed_token_ids_mask, float("-inf")) + for processor in sampling_metadata.logitsprocs.non_argmax_invariant: + logits = processor.apply(logits) sampling_cfg = dict(self.config.llm.get("sampling", {})) default_top_p = float(sampling_cfg.get("top_p", 0.8)) diff --git a/vllm_omni/model_executor/models/covo_audio/covo_audio.py b/vllm_omni/model_executor/models/covo_audio/covo_audio.py index 42f80bb74..efbe7b0d6 100644 --- a/vllm_omni/model_executor/models/covo_audio/covo_audio.py +++ b/vllm_omni/model_executor/models/covo_audio/covo_audio.py @@ -243,6 +243,12 @@ class CovoAudioForConditionalGeneration( else lambda *args, **kwargs: None ) + def get_language_model(self) -> "nn.Module": + """Return the language model for upstream MoE detection.""" + if hasattr(self.model, "get_language_model"): + return self.model.get_language_model() + return self.model + def fused_thinker_talker_preprocess( self, input_ids: torch.Tensor | None, diff --git a/vllm_omni/model_executor/models/ming_flash_omni/ming_flash_omni.py b/vllm_omni/model_executor/models/ming_flash_omni/ming_flash_omni.py index 4e3da735a..4da2455ae 100644 --- a/vllm_omni/model_executor/models/ming_flash_omni/ming_flash_omni.py +++ b/vllm_omni/model_executor/models/ming_flash_omni/ming_flash_omni.py @@ -143,6 +143,12 @@ class MingFlashOmniForConditionalGeneration( f"For image generation use stage_type: diffusion with model_arch: MingImagePipeline." ) + def get_language_model(self) -> "nn.Module": + """Return the language model for upstream MoE detection.""" + if hasattr(self.model, "get_language_model"): + return self.model.get_language_model() + return self.model + def forward( self, input_ids: torch.Tensor, 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 0ae1c63e3..81023c4a3 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 @@ -400,6 +400,13 @@ class MossTTSNanoForGeneration(nn.Module): runtime_additional_information: list[dict[str, Any]] | None = None, **kwargs: Any, ) -> OmniOutput: + if self._lm is None or self._audio_tokenizer is None: + # load_weights() loads the real LM from HuggingFace. When + # load_format=dummy the vllm DummyModelLoader skips calling + # model.load_weights(), so we do it lazily on first forward. + # Must happen before _make_dummy_hidden so self._device is set. + self.load_weights([]) + sr = getattr(self.config, "audio_tokenizer_sample_rate", 48000) sr_tensor = torch.tensor(sr, dtype=torch.int32) empty = torch.zeros((0,), dtype=torch.float32) @@ -418,9 +425,6 @@ class MossTTSNanoForGeneration(nn.Module): }, ) - if self._lm is None or self._audio_tokenizer is None: - raise RuntimeError("MOSS-TTS-Nano model not loaded. Was load_weights() called?") - outputs: list[torch.Tensor] = [] srs: list[torch.Tensor] = [] last_chunk_flags: list[bool] = [] diff --git a/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni.py b/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni.py index 712515f27..8d5341f3e 100644 --- a/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni.py +++ b/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni.py @@ -91,18 +91,17 @@ class Qwen2_5OmniForConditionalGeneration( elif self.model_stage == "talker": multimodal_config.skip_mm_profiling = True - # register the process function for the talker stage self.has_preprocess = True self.set_custom_preprocess(self.talker_preprocess) self.thinker = None # Initialize talker model wrapper (handles projection + LM) - self.talker = init_vllm_registered_model( - vllm_config=vllm_config, - prefix=maybe_prefix(prefix, "talker"), - hf_config=talker_config, - # Use registry architecture key - architectures=["Qwen2_5OmniTalkerModel"], - ) + with self._mark_language_model(vllm_config=vllm_config): + self.talker = init_vllm_registered_model( + vllm_config=vllm_config, + prefix=maybe_prefix(prefix, "talker"), + hf_config=talker_config, + architectures=["Qwen2_5OmniTalkerModel"], + ) self.model = self.talker self.token2wav = None # set suppress start id according to token2wav @@ -130,12 +129,13 @@ class Qwen2_5OmniForConditionalGeneration( self.token2wav_config = getattr(config, "token2wav_config", None) self.token2wav = None if self.token2wav_config is not None: - self.token2wav = init_vllm_registered_model( - vllm_config=vllm_config, - prefix=maybe_prefix(prefix, "token2wav"), - hf_config=self.token2wav_config, - architectures=["Qwen2_5OmniToken2WavModel"], - ) + with self._mark_language_model(vllm_config=vllm_config): + self.token2wav = init_vllm_registered_model( + vllm_config=vllm_config, + prefix=maybe_prefix(prefix, "token2wav"), + hf_config=self.token2wav_config, + architectures=["Qwen2_5OmniToken2WavModel"], + ) # voice resources (loaded on demand) self._token2wav_conds: dict[str, torch.Tensor] = {} self._token2wav_ref_mels: dict[str, torch.Tensor] = {} @@ -149,6 +149,18 @@ class Qwen2_5OmniForConditionalGeneration( (self.thinker.make_empty_intermediate_tensors) if self.model_stage == "thinker" else lambda: None ) + def get_language_model(self) -> nn.Module: + """Return the language model for MoE detection in upstream load_model. + + Upstream GPUModelRunner.load_model() calls get_language_model() to find + the MoE model inside VLM wrappers. For thinker (which IS MoE), delegate + to the child. For talker/code2wav (not MoE, not SupportsMultiModal), + return self.model directly to avoid NotImplementedError. + """ + if hasattr(self.model, "get_language_model"): + return self.model.get_language_model() + return self.model + # -------------------- Device utilities -------------------- @staticmethod def _module_device(module: nn.Module) -> torch.device: diff --git a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py index 6bd8bf305..3966f2140 100644 --- a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py +++ b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py @@ -278,6 +278,12 @@ class Qwen3OmniMoeForConditionalGeneration( return self.model.sampler return Sampler() + def get_language_model(self) -> torch.nn.Module: + """Delegate to the active stage's language model for upstream MoE resolution.""" + if hasattr(self.model, "get_language_model"): + return self.model.get_language_model() + return self.model + def embed_input_ids( self, input_ids: torch.Tensor, diff --git a/vllm_omni/model_executor/models/registry.py b/vllm_omni/model_executor/models/registry.py index ff61d6e6c..9ef3c853a 100644 --- a/vllm_omni/model_executor/models/registry.py +++ b/vllm_omni/model_executor/models/registry.py @@ -2,6 +2,7 @@ from vllm.model_executor.models.registry import ( _VLLM_MODELS, _LazyRegisteredModel, _ModelRegistry, + _resolve_module_name, ) _OMNI_MODELS = { @@ -261,7 +262,7 @@ OmniModelRegistry = _ModelRegistry( { **{ model_arch: _LazyRegisteredModel( - module_name=f"vllm.model_executor.models.{mod_relname}", + module_name=_resolve_module_name(mod_relname), class_name=cls_name, ) for model_arch, (mod_relname, cls_name) in _VLLM_MODELS.items() diff --git a/vllm_omni/model_executor/models/voxtral_tts/voxtral_tts.py b/vllm_omni/model_executor/models/voxtral_tts/voxtral_tts.py index 4d2421eea..eb8b5c14a 100644 --- a/vllm_omni/model_executor/models/voxtral_tts/voxtral_tts.py +++ b/vllm_omni/model_executor/models/voxtral_tts/voxtral_tts.py @@ -159,6 +159,13 @@ class VoxtralTTSForConditionalGeneration( raise ValueError("Invalid model stage") # -------------------- CUDA Graph for acoustic transformer -------------------- + + def get_language_model(self) -> "nn.Module": + """Return the language model for upstream MoE detection.""" + if hasattr(self.model, "get_language_model"): + return self.model.get_language_model() + return self.model + def _enable_acoustic_transformer_cudagraph(self): """Initialize and capture CUDA graphs for compute_mm_logits.""" if self.model_stage != "audio_generation" or not hasattr(self, "_vllm_config"): diff --git a/vllm_omni/patch.py b/vllm_omni/patch.py index b78323c21..af774fa21 100644 --- a/vllm_omni/patch.py +++ b/vllm_omni/patch.py @@ -128,3 +128,25 @@ for module_name, module in list(sys.modules.items()): module.StreamingUpdate = OmniStreamingUpdate if hasattr(module, "EngineCoreRequest") and module.EngineCoreRequest == _OriginalEngineCoreRequest: module.EngineCoreRequest = OmniEngineCoreRequest + + +# Patch: add qwen3_omni_moe to vllm's chat template fallback registry. +# Qwen/Qwen3-Omni-30B-A3B-Instruct stores its chat_template in a standalone +# chat_template.json (not in tokenizer_config.json). transformers < 5.9.0 +# does not load this file, so the tokenizer has no chat_template attribute. +# vllm's resolve_chat_template falls back to MODEL_TYPE_TO_CHAT_TEMPLATE +# which has "qwen" but not "qwen3_omni_moe". Register the same fallback. +def _patch_chat_template_registry(): + try: + from vllm.transformers_utils.chat_templates.registry import ( + _MODEL_TYPE_TO_CHAT_TEMPLATE_FALLBACK, + _get_qwen_chat_template_fallback, + ) + + if "qwen3_omni_moe" not in _MODEL_TYPE_TO_CHAT_TEMPLATE_FALLBACK: + _MODEL_TYPE_TO_CHAT_TEMPLATE_FALLBACK["qwen3_omni_moe"] = _get_qwen_chat_template_fallback + except ImportError: + pass + + +_patch_chat_template_registry() diff --git a/vllm_omni/platforms/cuda/platform.py b/vllm_omni/platforms/cuda/platform.py index 63002a98d..c96da69fc 100644 --- a/vllm_omni/platforms/cuda/platform.py +++ b/vllm_omni/platforms/cuda/platform.py @@ -229,4 +229,4 @@ class CudaOmniPlatform(OmniPlatform, CudaPlatformBase): if envs.VLLM_USE_OINK_OPS: rms_norm = ["oink"] + default - return IrOpPriorityConfig.with_default(default, rms_norm=rms_norm) + return IrOpPriorityConfig.with_default(default, rms_norm=rms_norm, fused_add_rms_norm=rms_norm) diff --git a/vllm_omni/platforms/interface.py b/vllm_omni/platforms/interface.py index 48c513da3..c1660ccb7 100644 --- a/vllm_omni/platforms/interface.py +++ b/vllm_omni/platforms/interface.py @@ -11,6 +11,7 @@ from vllm.config import CUDAGraphMode, VllmConfig from vllm.forward_context import BatchDescriptor from vllm.logger import init_logger from vllm.platforms import Platform +from vllm.platforms.interface import PlatformEnum logger = init_logger(__name__) @@ -239,6 +240,7 @@ class OmniPlatform(Platform): class UnspecifiedOmniPlatform(OmniPlatform): _omni_enum = OmniPlatformEnum.UNSPECIFIED + _enum = PlatformEnum.UNSPECIFIED device_type = "" @classmethod 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 3a8556273..a85682b8a 100644 --- a/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py +++ b/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py @@ -16,11 +16,6 @@ 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 ( - 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 ( @@ -257,7 +252,7 @@ class NPUARModelRunner(OmniNPUModelRunner): intermediate_tensors: IntermediateTensors | None = None, ) -> OmniModelRunnerOutput | IntermediateTensors | None: if self.vllm_config.model_config.enable_return_routed_experts: - capturer = get_global_experts_capturer() + capturer = self.routed_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: @@ -728,12 +723,7 @@ class NPUARModelRunner(OmniNPUModelRunner): 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, - ) + self._omni_routed_experts_d2h(scheduler_output) return None @@ -911,17 +901,11 @@ class NPUARModelRunner(OmniNPUModelRunner): if self.speculative_config is not None: self.finalize_kv_connector() - routed_experts_dict = None + routed_experts_lists = None if self.model_config.enable_return_routed_experts: - capturer = get_global_experts_capturer() + capturer = self.routed_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, - ) + routed_experts_lists = self._omni_extract_routed_experts(scheduler_output) # -------------------------------------- Omni-new ------------------------------------------------- engine_output_type, downstream_req_ids = self._resolve_pooler_payload_req_ids(req_ids_output_copy) @@ -1042,9 +1026,9 @@ 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 + model_runner_output.routed_experts = routed_experts_lists # -------------------------------------- Omni-new ------------------------------------------------- if self.ascend_config.profiling_chunk_config.enabled and hasattr(self, "_execution_start_time"): 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 f6af3de52..36125447d 100644 --- a/vllm_omni/platforms/npu/worker/npu_generation_model_runner.py +++ b/vllm_omni/platforms/npu/worker/npu_generation_model_runner.py @@ -15,11 +15,6 @@ 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 ( - 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 @@ -78,7 +73,7 @@ class NPUGenerationModelRunner(OmniNPUModelRunner): intermediate_tensors: IntermediateTensors | None = None, ) -> OmniModelRunnerOutput | IntermediateTensors | None: if self.vllm_config.model_config.enable_return_routed_experts: - capturer = get_global_experts_capturer() + capturer = self.routed_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: @@ -387,12 +382,7 @@ class NPUGenerationModelRunner(OmniNPUModelRunner): 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, - ) + self._omni_routed_experts_d2h(scheduler_output) return None @@ -479,17 +469,11 @@ 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 + routed_experts_lists = 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, - ) + routed_experts_lists = self._omni_extract_routed_experts(scheduler_output) output = OmniModelRunnerOutput( req_ids=req_ids_output_copy, req_id_to_index=req_id_to_index_output_copy, @@ -500,8 +484,8 @@ 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, ) + output.routed_experts = routed_experts_lists # -------------------------------------- Omni-new ------------------------------------------------- if self.speculative_config is not None: diff --git a/vllm_omni/platforms/rocm/platform.py b/vllm_omni/platforms/rocm/platform.py index c2e3fae67..1b50c328b 100644 --- a/vllm_omni/platforms/rocm/platform.py +++ b/vllm_omni/platforms/rocm/platform.py @@ -172,4 +172,4 @@ class RocmOmniPlatform(OmniPlatform, RocmPlatform): else: rms_norm = default - return IrOpPriorityConfig.with_default(default, rms_norm=rms_norm) + return IrOpPriorityConfig.with_default(default, rms_norm=rms_norm, fused_add_rms_norm=rms_norm) diff --git a/vllm_omni/quantization/factory.py b/vllm_omni/quantization/factory.py index 2a0326c2e..d9a7a31bf 100644 --- a/vllm_omni/quantization/factory.py +++ b/vllm_omni/quantization/factory.py @@ -10,19 +10,73 @@ different QuantizationConfig subclass in the OMNI context (e.g. GGUF). from __future__ import annotations import inspect +import sys from collections.abc import Callable, Mapping +from types import ModuleType from typing import Any from vllm.logger import init_logger -from vllm.model_executor.layers.quantization import ( + + +# --------------------------------------------------------------------------- +# Stub the ``humming`` package so that vLLM's lazy import inside +# ``get_quantization_config()`` (which unconditionally does +# ``from .humming import HummingConfig``) does not crash when the real +# ``humming`` wheel is not installed. Only populate the bare-minimum +# names that ``humming.py`` accesses at module level. +# --------------------------------------------------------------------------- +def _register_humming_stubs() -> None: + """Register stub ``humming`` sub-modules so that the optional + humming quantization backend can be imported without the real wheel.""" + if "humming" in sys.modules: + return # already present (real or stub) + + # --- sub-modules --- + submodules: dict[str, tuple[str, ...]] = { + "humming": (), + "humming.config": ("GemmType",), + "humming.dtypes": ("DataType",), + "humming.layer": ("HummingLayerMeta", "HummingMethod"), + "humming.schema": ( + "BaseInputSchema", + "BaseWeightSchema", + "HummingInputSchema", + "HummingWeightSchema", + ), + "humming.utils": (), + "humming.utils.weight": ("quantize_weight",), + } + + registry: dict[str, ModuleType] = {} + for name, attrs in submodules.items(): + mod = ModuleType(name) + for attr in attrs: + setattr(mod, attr, type(attr, (), {})) + registry[name] = mod + + # wire parent references + registry["humming"].config = registry["humming.config"] + registry["humming"].dtypes = registry["humming.dtypes"] + registry["humming"].layer = registry["humming.layer"] + registry["humming"].schema = registry["humming.schema"] + registry["humming"].utils = registry["humming.utils"] + registry["humming.utils"].weight = registry["humming.utils.weight"] + + for name, mod in registry.items(): + sys.modules[name] = mod + + +_register_humming_stubs() + +from vllm.model_executor.layers.quantization import ( # noqa: E402 QUANTIZATION_METHODS, get_quantization_config, ) -from vllm.model_executor.layers.quantization.base_config import ( +from vllm.model_executor.layers.quantization.base_config import ( # noqa: E402 QuantizationConfig, ) -from .component_config import ComponentQuantizationConfig +from .component_config import ComponentQuantizationConfig # noqa: E402 logger = init_logger(__name__) diff --git a/vllm_omni/worker/gpu_ar_model_runner.py b/vllm_omni/worker/gpu_ar_model_runner.py index 78e2acdb5..6f959cf46 100644 --- a/vllm_omni/worker/gpu_ar_model_runner.py +++ b/vllm_omni/worker/gpu_ar_model_runner.py @@ -19,11 +19,6 @@ 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 set_forward_context from vllm.logger import init_logger -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.v1.core.sched.output import GrammarOutput, SchedulerOutput from vllm.v1.outputs import AsyncModelRunnerOutput, make_empty_encoder_model_runner_output from vllm.v1.spec_decode.dflash import DFlashProposer @@ -363,6 +358,9 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin): if self.execute_model_state is not None: raise RuntimeError("State error: sample_tokens() must be called after execute_model() returns None.") + if self.routed_experts_initialized: + self.routed_experts_capturer.clear_buffer() + if not getattr(self, "_warmup_state_cleared", False): self._warmup_state_cleared = True if hasattr(self.model, "_clear_warmup_state"): @@ -414,7 +412,7 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin): ) if self.routed_experts_initialized: - capturer = get_global_experts_capturer() + capturer = self.routed_experts_capturer if capturer is not None and hasattr(capturer, "finalize_pending_copy"): capturer.finalize_pending_copy() @@ -679,7 +677,6 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin): if not get_pp_group().is_last_rank: # Return the intermediate tensors. assert isinstance(hidden_states, IntermediateTensors) - hidden_states.kv_connector_output = kv_connector_output self.kv_connector_output = kv_connector_output return hidden_states @@ -692,7 +689,7 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin): kv_connector_output, ) - sample_hidden_states = hidden_states[logits_indices] + sample_hidden_states = hidden_states[logits_indices.to(hidden_states.device)] # Try with sampling_metadata first; fall back to without for models that don't support it try: logits = self.model.compute_logits( @@ -704,7 +701,7 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin): # Rare case. assert not self.is_pooling_model - sample_hidden_states = hidden_states[logits_indices] + sample_hidden_states = hidden_states[logits_indices.to(hidden_states.device)] if not get_pp_group().is_last_rank: all_gather_tensors = { "residual": not is_residual_scattered_for_sp(self.vllm_config, num_tokens_padded) @@ -753,12 +750,7 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin): 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, - ) + self._omni_routed_experts_d2h(scheduler_output) return None @@ -830,18 +822,14 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin): if self.execute_model_state is None: kv_connector_output = self.kv_connector_output self.kv_connector_output = None - # Nothing to do (PP non-final rank case), output isn't used. - if not kv_connector_output: - return None # type: ignore[return-value] - + # receive sampled token ids from the last PP rank. + if self.use_async_scheduling and not get_pp_group().is_last_rank: + self._pp_receive_prev_sampled_token_ids_to_input_batch() # In case of PP with kv transfer, we need to pass through the # kv_connector_output - if kv_connector_output.is_empty(): - return self.attach_omni_connector_output(EMPTY_MODEL_RUNNER_OUTPUT) - - output = copy(EMPTY_MODEL_RUNNER_OUTPUT) - output.kv_connector_output = kv_connector_output - return self.attach_omni_connector_output(output) + return self.attach_omni_connector_output( + OmniModelRunnerOutput.with_kv_conn_output_only(kv_connector_output) + ) # Unpack ephemeral state. ( @@ -1108,15 +1096,9 @@ 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 + routed_experts_lists = 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, - ) + routed_experts_lists = self._omni_extract_routed_experts(scheduler_output) output = OmniModelRunnerOutput( req_ids=req_ids_output_copy, req_id_to_index=req_id_to_index_output_copy, @@ -1128,10 +1110,10 @@ 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() + output.routed_experts = routed_experts_lists if not self.use_async_scheduling: return output diff --git a/vllm_omni/worker/gpu_generation_model_runner.py b/vllm_omni/worker/gpu_generation_model_runner.py index 9f1060ed1..1dcdd31b7 100644 --- a/vllm_omni/worker/gpu_generation_model_runner.py +++ b/vllm_omni/worker/gpu_generation_model_runner.py @@ -8,7 +8,6 @@ from __future__ import annotations import gc import logging -from copy import copy from dataclasses import replace import numpy as np @@ -18,11 +17,6 @@ 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.forward_context import set_forward_context -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.utils.math_utils import cdiv from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput from vllm.v1.outputs import AsyncModelRunnerOutput, make_empty_encoder_model_runner_output @@ -112,6 +106,9 @@ class GPUGenerationModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin if self.execute_model_state is not None: raise RuntimeError("State error: sample_tokens() must be called after execute_model() returns None.") + if self.routed_experts_initialized: + self.routed_experts_capturer.clear_buffer() + if hasattr(self, "_omni_connector"): for request in getattr(scheduler_output, "pending_input_registrations", []): self.register_chunk_recv(request) @@ -123,7 +120,7 @@ class GPUGenerationModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin self.flush_full_payload_outputs(flush_ids) if self.routed_experts_initialized: - capturer = get_global_experts_capturer() + capturer = self.routed_experts_capturer if capturer is not None and hasattr(capturer, "finalize_pending_copy"): capturer.finalize_pending_copy() @@ -369,12 +366,7 @@ class GPUGenerationModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin 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, - ) + self._omni_routed_experts_d2h(scheduler_output) return None @@ -390,18 +382,14 @@ class GPUGenerationModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin self.kv_connector_output = None if self.execute_model_state is None: - # Nothing to do (PP non-final rank case), output isn't used. - if not kv_connector_output: - return None # type: ignore[return-value] - + # receive sampled token ids from the last PP rank. + if self.use_async_scheduling and not get_pp_group().is_last_rank: + self._pp_receive_prev_sampled_token_ids_to_input_batch() # In case of PP with kv transfer, we need to pass through the # kv_connector_output - if kv_connector_output.is_empty(): - return self.attach_omni_connector_output(EMPTY_MODEL_RUNNER_OUTPUT) - - output = copy(EMPTY_MODEL_RUNNER_OUTPUT) - output.kv_connector_output = kv_connector_output - return self.attach_omni_connector_output(output) + return self.attach_omni_connector_output( + OmniModelRunnerOutput.with_kv_conn_output_only(kv_connector_output) + ) # Unpack ephemeral state. ( @@ -462,15 +450,9 @@ 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 + routed_experts_lists = 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, - ) + routed_experts_lists = self._omni_extract_routed_experts(scheduler_output) if self._should_accumulate_full_payload_output(): for i, rid in enumerate(req_ids_output_copy): req_state = self.requests.get(rid) @@ -488,9 +470,9 @@ 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() + output.routed_experts = routed_experts_lists if not self.use_async_scheduling: return output @@ -544,9 +526,9 @@ class GPUGenerationModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin ) @torch.inference_mode() - def _dummy_sampler_run(self, hidden_states: torch.Tensor) -> None: + def _dummy_sampler_run(self, hidden_states: torch.Tensor) -> torch.Tensor: logger.warning("Dummy sampler run is not implemented for generation model") - return None + return torch.tensor([]) @torch.inference_mode() def _dummy_run( diff --git a/vllm_omni/worker/gpu_model_runner.py b/vllm_omni/worker/gpu_model_runner.py index 0663673d2..e74997618 100644 --- a/vllm_omni/worker/gpu_model_runner.py +++ b/vllm_omni/worker/gpu_model_runner.py @@ -36,6 +36,7 @@ from vllm_omni.platforms import current_omni_platform if TYPE_CHECKING: from vllm.v1.core.sched.output import SchedulerOutput + from vllm.v1.outputs import RoutedExpertsLists else: xgr = LazyLoader("xgr", globals(), "xgrammar") xgr_torch_compile = LazyLoader( @@ -57,6 +58,52 @@ class OmniGPUModelRunner(GPUModelRunner): # when we initialize the metadata builders if enabled self.omni_prefix_cache = None + def _omni_routed_experts_d2h(self, scheduler_output) -> None: + """Issue routed-experts D2H copy matching upstream GPUModelRunner pattern. + + Upstream does this inline in ``execute_model``: + buf = self.routed_experts_capturer.get_device_buffer() + total = scheduler_output.total_num_scheduled_tokens + self.routed_experts_cpu[:total].copy_(buf[:total], non_blocking=True) + self.routed_experts_slot_mapping_cpu[:total].copy_( + self.routed_experts_slot_mapping_device[:total], non_blocking=True) + """ + if not self.routed_experts_initialized: + return + buf = self.routed_experts_capturer.get_device_buffer() + total = scheduler_output.total_num_scheduled_tokens + self.routed_experts_cpu[:total].copy_(buf[:total], non_blocking=True) + if hasattr(self, "routed_experts_slot_mapping_device"): + self.routed_experts_slot_mapping_cpu[:total].copy_( + self.routed_experts_slot_mapping_device[:total], + non_blocking=True, + ) + + def _omni_extract_routed_experts(self, scheduler_output) -> "RoutedExpertsLists | None": + """Extract routed experts matching upstream GPUModelRunner pattern. + + Upstream (sync path, sample_tokens): + total = scheduler_output.total_num_scheduled_tokens + output.routed_experts = RoutedExpertsLists( + routing_data=self.routed_experts_cpu[:total].numpy(), + slot_mapping=self.routed_experts_slot_mapping_cpu[:total].numpy(), + ) + + Returns RoutedExpertsLists (batch-level, with slot_mapping) so that + downstream schedulers can use slot_mapping to map back to requests. + """ + from vllm.v1.outputs import RoutedExpertsLists + + if not self.routed_experts_initialized: + return None + total = scheduler_output.total_num_scheduled_tokens + if total <= 0: + return None + return RoutedExpertsLists( + routing_data=self.routed_experts_cpu[:total].numpy(), + slot_mapping=self.routed_experts_slot_mapping_cpu[:total].numpy(), + ) + def initialize_metadata_builders(self, kv_cache_config, kernel_block_sizes): """Override to fix scheduler_metadata buffer size for FA3 + CUDA graph. @@ -747,6 +794,15 @@ class OmniGPUModelRunner(GPUModelRunner): raise ValueError(f"Invalid hidden states type: {type(hidden_states)}") return text_hidden_states, multimodal_outputs + def _dummy_sampler_run(self, hidden_states: torch.Tensor) -> torch.Tensor: + # Models loaded with load_format=dummy (e.g. MossTTSNano) may + # produce CPU hidden_states while the upstream sampler warmup + # expects GPU tensors. Skip the warmup in that case — a + # meaningful warmup requires real weights anyway. + if hidden_states.device != self.device: + return torch.tensor([]) + return super()._dummy_sampler_run(hidden_states=hidden_states) + @torch.inference_mode() def _dummy_run( self, @@ -1100,7 +1156,7 @@ class OmniGPUModelRunner(GPUModelRunner): self.eplb_step(is_dummy=True, is_profile=is_profile) logit_indices = np.cumsum(num_scheduled_tokens) - 1 - logit_indices_device = torch.from_numpy(logit_indices).to(self.device, non_blocking=True) + logit_indices_device = torch.from_numpy(logit_indices).to(hidden_states.device, non_blocking=True) return hidden_states, hidden_states[logit_indices_device] # ------------------------------------------------------------------