diff --git a/docs/configuration/stage_configs.md b/docs/configuration/stage_configs.md index 0e671bebf..412f3d524 100644 --- a/docs/configuration/stage_configs.md +++ b/docs/configuration/stage_configs.md @@ -135,6 +135,14 @@ Each stage in the `stage_args` list contains the following configuration options A unique identifier for each stage in the multi-stage pipeline. Stages are numbered sequentially starting from 0, and this ID is used to reference stages in inter-stage dependencies (e.g., `engine_input_source`). +### `prompt_expand_func` (Optional) + +A custom Python function hook for the LLM stage (Stage 0) that expands a single incoming prompt object into multiple prompts. This is primarily used for multi-modal Classifier-Free Guidance (CFG), where it generates the necessary companion requests (like a negative text prompt) and tags them with internal roles (e.g., `cfg_text`). This ensures the upstream LLM generates the needed contextual hidden states for both the conditional and unconditional generations simultaneously. + +### `cfg_kv_collect_func` (Optional) + +A custom Python function hook for downstream diffusion stages (Stage 1+) to collect, map, and process the KV caches transferred from the companion requests fired by `prompt_expand_func`. It aggregates the hidden condition states cleanly (e.g., binding them as `cfg_text_past_key_values` and `cfg_text_kv_metadata`), allowing the diffusion runtime to perform CFG smoothly without redundantly evaluating text paths on the DiT workers. + ### `runtime` Configuration for disaggregated execution of the stage, controlling how the stage is deployed and executed. diff --git a/docs/design/architecture_overview.md b/docs/design/architecture_overview.md index 6793895cd..2e58320ad 100644 --- a/docs/design/architecture_overview.md +++ b/docs/design/architecture_overview.md @@ -92,6 +92,13 @@ The framework achieves high performance through several optimization techniques: * **Quantization:** Supports various quantization implementations including FP8 and AWQ. * **FusedOps:** Allows for custom and third-party integration. +### Classifier-Free Guidance (CFG) Companion Flow + +vLLM-Omni natively models Classifier-Free Guidance (CFG) across disaggregated multi-stage setups via a "companion request" paradigm, eliminating redundant textual/multimodal context computation boundaries: +1. **Prompt Expansion:** In the initial autoregressive (AR) stage, a customized `prompt_expand_func` hook intercepts incoming generation prompts and pairs them directly with negative companion prompts (e.g., a default negative prompt) on the fly, tagging the secondary prompt with a specific internal role (`cfg_text`). +2. **Synchronized KV Cache Transfer:** The AR stage evaluates both the primary and companion sequence batches concurrently. The `OmniConnector` captures these specific structural dependencies and reliably passes the positive and negative outcome KV caches seamlessly across stage boundaries via shared memory or network protocols. +3. **KV Cache Collection & Injection:** Upon reaching the downstream Diffusion (DiT) Engine, an assigned `cfg_kv_collect_func` automatically intercepts the mapped companion caches (`cfg_text_past_key_values`). These auxiliary dependencies are natively gathered and seamlessly bound to the primary generation sequence variables, enabling the DiT Engine to cleanly implement cross-attention CFG guidance over accurate conditioning and unconditioning structures in parallel. + ### Flexibility and Usability vLLM-Omni is designed to be flexible and straightforward for users: diff --git a/examples/offline_inference/bagel/end2end.py b/examples/offline_inference/bagel/end2end.py index bf9968ffa..04ae4c15d 100644 --- a/examples/offline_inference/bagel/end2end.py +++ b/examples/offline_inference/bagel/end2end.py @@ -49,7 +49,7 @@ def parse_args(): parser.add_argument("--cfg-text-scale", type=float, default=4.0, help="Text CFG scale (default: 4.0)") parser.add_argument("--cfg-img-scale", type=float, default=1.5, help="Image CFG scale (default: 1.5)") parser.add_argument( - "--negative-prompt", type=str, default=None, help="Negative prompt (not yet supported, reserved for future)" + "--negative-prompt", type=str, default=None, help="Negative prompt for CFG (default: empty prompt)" ) args = parser.parse_args() @@ -162,6 +162,8 @@ def main(): # text2img final_prompt_text = f"<|im_start|>{p}<|im_end|>" prompt_dict = {"prompt": final_prompt_text, "modalities": ["image"]} + if args.negative_prompt is not None: + prompt_dict["negative_prompt"] = args.negative_prompt formatted_prompts.append(prompt_dict) params_list = omni.default_sampling_params_list @@ -170,10 +172,13 @@ def main(): if len(params_list) > 1: diffusion_params = params_list[1] diffusion_params.num_inference_steps = args.steps # type: ignore - diffusion_params.extra_args = { # type: ignore + extra = { "cfg_text_scale": args.cfg_text_scale, "cfg_img_scale": args.cfg_img_scale, } + if args.negative_prompt is not None: + extra["negative_prompt"] = args.negative_prompt + diffusion_params.extra_args = extra # type: ignore omni_outputs = list(omni.generate(prompts=formatted_prompts, sampling_params_list=params_list)) diff --git a/tests/diffusion/test_diffusion_model_runner.py b/tests/diffusion/test_diffusion_model_runner.py index 6ff74154f..2f9aadfd8 100644 --- a/tests/diffusion/test_diffusion_model_runner.py +++ b/tests/diffusion/test_diffusion_model_runner.py @@ -56,7 +56,10 @@ def _make_runner(cache_backend, cache_backend_name: str, enable_cache_dit_summar enable_cache_dit_summary=enable_cache_dit_summary, parallel_config=SimpleNamespace(use_hsdp=False), ) - runner.kv_transfer_manager = SimpleNamespace(receive_kv_cache=lambda req, target_device: None) + runner.kv_transfer_manager = SimpleNamespace( + receive_kv_cache=lambda req, target_device=None: None, + receive_multi_kv_cache=lambda req, cfg_kv_collect_func=None, target_device=None: None, + ) return runner @@ -105,8 +108,8 @@ def test_execute_model_emits_cache_summary_with_active_cache_dit_backend(monkeyp def test_load_model_clears_cache_backend_for_unsupported_pipeline(monkeypatch): class _DummyLoader: - def __init__(self, load_config): - del load_config + def __init__(self, load_config, od_config=None): + del load_config, od_config def load_model(self, **kwargs): del kwargs diff --git a/tests/e2e/offline_inference/stage_configs/bagel_mooncake_ci.yaml b/tests/e2e/offline_inference/stage_configs/bagel_mooncake_ci.yaml index 8c170dffd..79e825d5d 100644 --- a/tests/e2e/offline_inference/stage_configs/bagel_mooncake_ci.yaml +++ b/tests/e2e/offline_inference/stage_configs/bagel_mooncake_ci.yaml @@ -4,6 +4,7 @@ stage_args: - stage_id: 0 stage_type: llm + prompt_expand_func: vllm_omni.model_executor.stage_input_processors.bagel.expand_cfg_prompts runtime: devices: "0" max_batch_size: 1 @@ -39,6 +40,7 @@ stage_args: to_stage_1: mooncake_connector - stage_id: 1 stage_type: diffusion + cfg_kv_collect_func: vllm_omni.model_executor.stage_input_processors.bagel.collect_cfg_kv_caches runtime: devices: "0" max_batch_size: 1 diff --git a/tests/e2e/offline_inference/stage_configs/bagel_sharedmemory_ci.yaml b/tests/e2e/offline_inference/stage_configs/bagel_sharedmemory_ci.yaml index ef605a3a3..eeefab236 100644 --- a/tests/e2e/offline_inference/stage_configs/bagel_sharedmemory_ci.yaml +++ b/tests/e2e/offline_inference/stage_configs/bagel_sharedmemory_ci.yaml @@ -4,6 +4,7 @@ stage_args: - stage_id: 0 stage_type: llm + prompt_expand_func: vllm_omni.model_executor.stage_input_processors.bagel.expand_cfg_prompts runtime: devices: "0" max_batch_size: 1 @@ -38,6 +39,7 @@ stage_args: - stage_id: 1 stage_type: diffusion + cfg_kv_collect_func: vllm_omni.model_executor.stage_input_processors.bagel.collect_cfg_kv_caches runtime: devices: "0" max_batch_size: 1 diff --git a/tests/e2e/offline_inference/test_bagel_text2img.py b/tests/e2e/offline_inference/test_bagel_text2img.py index 9aef47a74..20f788895 100644 --- a/tests/e2e/offline_inference/test_bagel_text2img.py +++ b/tests/e2e/offline_inference/test_bagel_text2img.py @@ -37,16 +37,16 @@ from vllm_omni.entrypoints.omni import Omni # "Generated with seed=52, num_inference_steps=15, # prompt='A futuristic city skyline at twilight, cyberpunk style'" REFERENCE_PIXELS = [ - {"position": (100, 100), "rgb": (68, 107, 134)}, - {"position": (400, 50), "rgb": (95, 139, 166)}, - {"position": (700, 100), "rgb": (99, 122, 151)}, - {"position": (150, 400), "rgb": (111, 125, 153)}, - {"position": (512, 512), "rgb": (97, 107, 131)}, - {"position": (700, 400), "rgb": (48, 64, 98)}, - {"position": (100, 700), "rgb": (79, 63, 84)}, - {"position": (400, 700), "rgb": (40, 58, 79)}, - {"position": (700, 700), "rgb": (60, 75, 103)}, - {"position": (256, 256), "rgb": (97, 128, 156)}, + {"position": (100, 100), "rgb": (49, 96, 134)}, + {"position": (400, 50), "rgb": (63, 127, 167)}, + {"position": (700, 100), "rgb": (70, 101, 141)}, + {"position": (150, 400), "rgb": (115, 90, 150)}, + {"position": (512, 512), "rgb": (98, 86, 119)}, + {"position": (700, 400), "rgb": (29, 42, 91)}, + {"position": (100, 700), "rgb": (47, 50, 88)}, + {"position": (400, 700), "rgb": (36, 52, 91)}, + {"position": (700, 700), "rgb": (45, 58, 99)}, + {"position": (256, 256), "rgb": (62, 94, 135)}, ] # Maximum allowed difference per color channel @@ -80,6 +80,10 @@ def _configure_sampling_params(omni: Omni, max_tokens: int = 1, num_inference_st params_list[0].max_tokens = max_tokens # type: ignore if len(params_list) > 1: params_list[1].num_inference_steps = num_inference_steps # type: ignore + params_list[1].extra_args = { # type: ignore + "cfg_text_scale": 4.0, + "cfg_img_scale": 1.5, + } return params_list diff --git a/tests/entrypoints/test_cfg_companion_tracker.py b/tests/entrypoints/test_cfg_companion_tracker.py new file mode 100644 index 000000000..941ead41f --- /dev/null +++ b/tests/entrypoints/test_cfg_companion_tracker.py @@ -0,0 +1,114 @@ +import time +from types import SimpleNamespace + +import pytest + +from vllm_omni.entrypoints.cfg_companion_tracker import CfgCompanionTracker + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def dummy_expand_func(prompt, sp0): + if prompt == "expand_me": + return [SimpleNamespace(prompt={"prompt": "neg"}, role="cfg_text", request_id_suffix="__cfg_text")] + return [] + + +@pytest.fixture +def tracker(): + sp0 = SimpleNamespace() + return CfgCompanionTracker(prompt_expand_func=dummy_expand_func, stage0_sampling_params=sp0, timeout_s=0.1) + + +def test_companion_tracker_initialization(tracker): + assert not tracker.is_active + assert tracker.num_companions == 0 + + +def test_expand_prompts_registers_companions(tracker): + request_id_to_prompt = {"req1": "expand_me", "req2": "do_not_expand"} + + pairs = tracker.expand_prompts(request_id_to_prompt) + + assert len(pairs) == 1 + companion_id, prompt = pairs[0] + assert companion_id == "req1__cfg_text" + assert prompt == {"prompt": "neg"} + + assert tracker.is_active + assert tracker.num_companions == 1 + assert tracker.is_companion("req1__cfg_text") + assert not tracker.is_companion("req2__cfg_text") + assert tracker.has_companions("req1") + assert not tracker.has_companions("req2") + + comp_map = tracker.get_companion_request_ids("req1") + assert comp_map == {"cfg_text": "req1__cfg_text"} + + +def test_companion_lifecycle_success(tracker): + request_id_to_prompt = {"req1": "expand_me"} + tracker.expand_prompts(request_id_to_prompt) + + # Defer parent + engine_outputs = {"out": 123} + tracker.defer_parent("req1", engine_outputs, stage_id=0) + + # Initially not done + assert not tracker.all_companions_done("req1") + + # Companion completes + parent_id = tracker.on_companion_completed("req1__cfg_text") + + # Parent should be returned since all companions are done and it is pending + assert parent_id == "req1" + assert tracker.all_companions_done("req1") + + # Pop pending parent + popped = tracker.pop_pending_parent("req1") + assert popped is not None + assert popped["engine_outputs"] == engine_outputs + assert popped["stage_id"] == 0 + + +def test_companion_lifecycle_failure(tracker): + request_id_to_prompt = {"req1": "expand_me"} + tracker.expand_prompts(request_id_to_prompt) + + tracker.defer_parent("req1", {"out": 123}, stage_id=0) + + # Companion fails + parent_id, aborted = tracker.on_companion_error("req1__cfg_text") + + assert parent_id == "req1" + assert aborted is True + assert tracker.is_parent_failed("req1") + + # Parent should be removed from pending list + assert tracker.pop_pending_parent("req1") is None + + # Consume failure + tracker.consume_parent_failure("req1") + assert not tracker.is_parent_failed("req1") + + +def test_companion_lifecycle_timeout(tracker): + request_id_to_prompt = {"req1": "expand_me"} + tracker.expand_prompts(request_id_to_prompt) + + tracker.defer_parent("req1", {"out": 123}, stage_id=0) + + # Initially no timeouts + timeouts = tracker.check_timeouts() + assert len(timeouts) == 0 + + # Wait for timeout + time.sleep(0.15) + + # Check timeouts again + timeouts = tracker.check_timeouts() + assert len(timeouts) == 1 + assert timeouts[0] == "req1" + + # Should be removed from pending + assert tracker.pop_pending_parent("req1") is None diff --git a/tests/worker/test_omni_gpu_model_runner.py b/tests/worker/test_omni_gpu_model_runner.py index aab2781f0..b2d619315 100644 --- a/tests/worker/test_omni_gpu_model_runner.py +++ b/tests/worker/test_omni_gpu_model_runner.py @@ -64,6 +64,7 @@ def _make_runner(req_ids=("r1", "r2"), hidden_size=4): # Minimal attributes used by OmniGPUModelRunner._talker_mtp_forward runner.input_batch = DummyInputBatch(list(req_ids)) runner.requests = {rid: DummyReqState() for rid in req_ids} + runner.model_intermediate_buffer = {} # query_start_loc.cpu[req_index] is used to locate the token position # in the flattened `inputs_embeds`. @@ -167,6 +168,61 @@ def test_talker_mtp_forward_cpu_empty_batch_noop(monkeypatch): assert torch.allclose(inputs_embeds, before) +def test_update_intermediate_buffer_writes_to_buffer_and_setattr(monkeypatch): + """Validate that _update_intermediate_buffer writes to model_intermediate_buffer + (forward path) and mirrors to additional_information_cpu setattr (backward compat).""" + import vllm_omni.worker.gpu_model_runner as mod + + monkeypatch.setattr(mod, "set_forward_context", _noop_forward_context) + + runner = _make_runner(req_ids=("r1",), hidden_size=4) + + update = {"my_tensor": torch.tensor([1.0, 2.0]), "my_list": [3, 4]} + OmniGPUModelRunner._update_intermediate_buffer(runner, "r1", update) + + # Forward: buffer is populated + assert "r1" in runner.model_intermediate_buffer + buf = runner.model_intermediate_buffer["r1"] + assert torch.allclose(buf["my_tensor"], torch.tensor([1.0, 2.0])) + assert buf["my_list"] == [3, 4] + + # Backward compat: setattr is also populated + info_cpu = runner.requests["r1"].additional_information_cpu + assert torch.allclose(info_cpu["my_tensor"], torch.tensor([1.0, 2.0])) + assert info_cpu["my_list"] == [3, 4] + + +def test_update_intermediate_buffer_accumulates(): + """Validate that successive merges accumulate keys in the buffer.""" + runner = _make_runner(req_ids=("r1",), hidden_size=4) + + OmniGPUModelRunner._update_intermediate_buffer(runner, "r1", {"a": torch.tensor([1.0])}) + OmniGPUModelRunner._update_intermediate_buffer(runner, "r1", {"b": torch.tensor([2.0])}) + + buf = runner.model_intermediate_buffer["r1"] + assert "a" in buf and "b" in buf + assert torch.allclose(buf["a"], torch.tensor([1.0])) + assert torch.allclose(buf["b"], torch.tensor([2.0])) + + +def test_update_intermediate_buffer_skips_empty_update(): + """Validate that an empty update dict is a no-op.""" + runner = _make_runner(req_ids=("r1",), hidden_size=4) + + OmniGPUModelRunner._update_intermediate_buffer(runner, "r1", {}) + + assert "r1" not in runner.model_intermediate_buffer + + +def test_update_intermediate_buffer_skips_unknown_req_id(): + """Validate that merge is a no-op when req_id is not in self.requests.""" + runner = _make_runner(req_ids=("r1",), hidden_size=4) + + OmniGPUModelRunner._update_intermediate_buffer(runner, "unknown_req", {"key": torch.tensor([1.0])}) + + assert "unknown_req" not in runner.model_intermediate_buffer + + def test_maybe_attach_mimo_audio_req_infos_enriches_dict(): runner = _make_runner_for_mimo() req_id = "r_mimo" diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index bdf877a6e..bd5af2263 100644 --- a/vllm_omni/diffusion/data.py +++ b/vllm_omni/diffusion/data.py @@ -379,6 +379,10 @@ class OmniDiffusionConfig: # Compilation enforce_eager: bool = False + # Parallel weight loading (for faster diffusion model startup) + enable_multithread_weight_load: bool = True + num_weight_load_threads: int = 4 + # Enable sleep mode enable_sleep_mode: bool = False @@ -442,6 +446,9 @@ class OmniDiffusionConfig: # Omni configuration (injected from stage config) omni_kv_config: dict[str, Any] = field(default_factory=dict) + # Model-specific function for collecting CFG KV caches (set at runtime) + cfg_kv_collect_func: Any | None = None + # Quantization settings # Supported methods: "fp8" (FP8 W8A8 on Ada/Hopper, weight-only on older GPUs) quantization: str | None = None diff --git a/vllm_omni/diffusion/model_loader/diffusers_loader.py b/vllm_omni/diffusion/model_loader/diffusers_loader.py index c4271c338..78d9ffd03 100644 --- a/vllm_omni/diffusion/model_loader/diffusers_loader.py +++ b/vllm_omni/diffusion/model_loader/diffusers_loader.py @@ -3,6 +3,7 @@ import dataclasses import glob import os +import re import time from collections.abc import Generator, Iterable from pathlib import Path @@ -22,6 +23,7 @@ from vllm.model_executor.model_loader.weight_utils import ( filter_duplicate_safetensors_files, filter_files_not_needed_for_inference, maybe_download_from_modelscope, + multi_thread_safetensors_weights_iterator, safetensors_weights_iterator, ) from vllm.utils.import_utils import resolve_obj_by_qualname @@ -35,6 +37,12 @@ from vllm_omni.diffusion.registry import initialize_model logger = init_logger(__name__) +def _natural_sort_key(filepath: str) -> list: + """Natural sort key for filenames with numeric components, e.g. + model-00001-of-00005.safetensors -> ['model-', 1, '-of-', 5, '.safetensors'].""" + return [int(s) if s.isdigit() else s for s in re.split(r"(\d+)", os.path.basename(filepath))] + + MODEL_INDEX = "model_index.json" DIFFUSION_MODEL_WEIGHTS_INDEX = "diffusion_pytorch_model.safetensors.index.json" @@ -42,9 +50,6 @@ DIFFUSION_MODEL_WEIGHTS_INDEX = "diffusion_pytorch_model.safetensors.index.json" class DiffusersPipelineLoader: """Model loader that can load diffusers pipeline components from disk.""" - # default number of thread when enable multithread weight loading - DEFAULT_NUM_THREADS = 8 - @dataclasses.dataclass class ComponentSource: """A source for weights.""" @@ -70,18 +75,9 @@ class DiffusersPipelineLoader: counter_before_loading_weights: float = 0.0 counter_after_loading_weights: float = 0.0 - def __init__(self, load_config: LoadConfig): + def __init__(self, load_config: LoadConfig, od_config: OmniDiffusionConfig | None = None): self.load_config = load_config - - # TODO(Isotr0py): Enable multithreaded weight loading - # extra_config = load_config.model_loader_extra_config - # allowed_keys = {"enable_multithread_load", "num_threads"} - # unexpected_keys = set(extra_config.keys()) - allowed_keys - - # if unexpected_keys: - # raise ValueError( - # f"Unexpected extra config keys for load format {load_config.load_format}: {unexpected_keys}" - # ) + self.od_config = od_config def _prepare_weights( self, @@ -175,18 +171,36 @@ class DiffusersPipelineLoader: def _get_weights_iterator(self, source: "ComponentSource") -> Generator[tuple[str, torch.Tensor], None, None]: """Get an iterator for the model weights based on the load format.""" - hf_folder, hf_weights_files, use_safetensors = self._prepare_weights( + _, hf_weights_files, use_safetensors = self._prepare_weights( source.model_or_path, source.subfolder, source.revision, source.fall_back_to_pt, source.allow_patterns_overrides, ) - weights_iterator = safetensors_weights_iterator( - hf_weights_files, - self.load_config.use_tqdm_on_load, - self.load_config.safetensors_load_strategy, + + od_config = self.od_config + use_multithread = ( + use_safetensors + and od_config is not None + and getattr(od_config, "enable_multithread_weight_load", False) + and self.load_config.safetensors_load_strategy != "torchao" ) + if use_multithread: + num_threads = getattr(od_config, "num_weight_load_threads", 4) + # Keep deterministic shard order before passing to vLLM helper. + sorted_hf_weights_files = sorted(hf_weights_files, key=_natural_sort_key) + weights_iterator = multi_thread_safetensors_weights_iterator( + sorted_hf_weights_files, + self.load_config.use_tqdm_on_load, + max_workers=num_threads, + ) + else: + weights_iterator = safetensors_weights_iterator( + hf_weights_files, + self.load_config.use_tqdm_on_load, + self.load_config.safetensors_load_strategy, + ) if self.counter_before_loading_weights == 0.0: self.counter_before_loading_weights = time.perf_counter() diff --git a/vllm_omni/diffusion/models/bagel/pipeline_bagel.py b/vllm_omni/diffusion/models/bagel/pipeline_bagel.py index cf8ade329..46ce4b6bd 100644 --- a/vllm_omni/diffusion/models/bagel/pipeline_bagel.py +++ b/vllm_omni/diffusion/models/bagel/pipeline_bagel.py @@ -327,16 +327,40 @@ class BagelPipeline(nn.Module): gen_context["past_key_values"] = injected_kv seq_len = injected_kv.key_cache[0].shape[0] gen_context["kv_lens"] = [seq_len] - gen_context["ropes"] = [seq_len] + if req.sampling_params.kv_metadata and "ropes" in req.sampling_params.kv_metadata: + gen_context["ropes"] = req.sampling_params.kv_metadata["ropes"] + else: + gen_context["ropes"] = [seq_len] - # Disable CFG: single KV cache cannot support 3-branch CFG - logger.warning("CFG is disabled when using injected KV Cache") - gen_params = BagelGenParams( - num_timesteps=gen_params.num_timesteps, - timestep_shift=gen_params.timestep_shift, - cfg_text_scale=1.0, - cfg_img_scale=1.0, - ) + cfg_text_kv = getattr(req.sampling_params, "cfg_text_past_key_values", None) + if cfg_text_kv is not None: + logger.info("CFG enabled with multi-KV: using injected cfg_text KV Cache") + cfg_text_seq_len = cfg_text_kv.key_cache[0].shape[0] + cfg_text_context["past_key_values"] = cfg_text_kv + cfg_text_context["kv_lens"] = [cfg_text_seq_len] + cfg_text_metadata = getattr(req.sampling_params, "cfg_text_kv_metadata", None) + if cfg_text_metadata and "ropes" in cfg_text_metadata: + cfg_text_context["ropes"] = cfg_text_metadata["ropes"] + else: + cfg_text_context["ropes"] = [cfg_text_seq_len] + + cfg_img_kv = getattr(req.sampling_params, "cfg_img_past_key_values", None) or injected_kv + cfg_img_seq_len = cfg_img_kv.key_cache[0].shape[0] + cfg_img_context["past_key_values"] = cfg_img_kv + cfg_img_context["kv_lens"] = [cfg_img_seq_len] + cfg_img_metadata = getattr(req.sampling_params, "cfg_img_kv_metadata", None) + if cfg_img_metadata and "ropes" in cfg_img_metadata: + cfg_img_context["ropes"] = cfg_img_metadata["ropes"] + else: + cfg_img_context["ropes"] = [cfg_img_seq_len] + else: + logger.warning("CFG is disabled: only single KV cache available") + gen_params = BagelGenParams( + num_timesteps=gen_params.num_timesteps, + timestep_shift=gen_params.timestep_shift, + cfg_text_scale=1.0, + cfg_img_scale=1.0, + ) else: image_input = ( diff --git a/vllm_omni/diffusion/models/hunyuan_image_3/hunyuan_image_3_transformer.py b/vllm_omni/diffusion/models/hunyuan_image_3/hunyuan_image_3_transformer.py index 04b961c9a..0c2f9e290 100644 --- a/vllm_omni/diffusion/models/hunyuan_image_3/hunyuan_image_3_transformer.py +++ b/vllm_omni/diffusion/models/hunyuan_image_3/hunyuan_image_3_transformer.py @@ -1,11 +1,9 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -import glob import inspect import logging import math -import os from collections.abc import Callable, Iterable from dataclasses import dataclass from typing import Any, cast @@ -22,7 +20,6 @@ from diffusers.utils.outputs import BaseOutput from diffusers.utils.torch_utils import randn_tensor from einops import rearrange from PIL import Image -from safetensors.torch import load_file from torch import nn from torchvision import transforms from transformers import PretrainedConfig, Siglip2ImageProcessorFast @@ -357,14 +354,6 @@ def build_batch_2d_rope( return stacked_cos, stacked_sin -def get_full_state_dict(model_path): - files = glob.glob(os.path.join(model_path, "*.safetensors")) - full_sd = {} - for f in files: - full_sd.update(load_file(f)) - return full_sd - - def rotate_half(x): """Rotates half the hidden dims of the input.""" x1 = x[..., : x.shape[-1] // 2] @@ -1744,7 +1733,6 @@ class HunyuanImage3Model(nn.Module): lora_vocab = (lora_config.lora_extra_vocab_size * (lora_config.max_loras or 1)) if lora_config else 0 self.vocab_size = config.vocab_size + lora_vocab self.org_vocab_size = config.vocab_size - self.wte = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx) if get_pp_group().is_first_rank or (config.tie_word_embeddings and get_pp_group().is_last_rank): self.embed_tokens = VocabParallelEmbedding( self.vocab_size, @@ -2047,7 +2035,7 @@ class HunyuanImage3Model(nn.Module): return_dict = return_dict if return_dict is not None else self.config.use_return_dict if inputs_embeds is None: - inputs_embeds = self.wte(input_ids) + inputs_embeds = self.embed_tokens(input_ids) # embed positions hidden_states = inputs_embeds diff --git a/vllm_omni/diffusion/models/hunyuan_image_3/pipeline_hunyuan_image_3.py b/vllm_omni/diffusion/models/hunyuan_image_3/pipeline_hunyuan_image_3.py index 9721f2edb..e4b717a69 100644 --- a/vllm_omni/diffusion/models/hunyuan_image_3/pipeline_hunyuan_image_3.py +++ b/vllm_omni/diffusion/models/hunyuan_image_3/pipeline_hunyuan_image_3.py @@ -37,7 +37,6 @@ from .hunyuan_image_3_transformer import ( UNetDown, UNetUp, build_batch_2d_rope, - get_full_state_dict, real_batched_index_select, ) @@ -113,10 +112,10 @@ class HunyuanImage3Pipeline(HunyuanImage3PreTrainedModel, GenerationMixin): self.vllm_config = get_current_vllm_config() self.post_init() - def pre_load(self): - tp_rank = get_tensor_model_parallel_rank() - state_dict = get_full_state_dict(self.od_config.model) - non_layer_prefixes = [ + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + skip_prefixes = ["lm_head."] if self.hf_config.tie_word_embeddings else [] + # List of unexpected keywords in weight names + non_model_layer_prefixes = [ "vae", "vision_model", "vision_aligner", @@ -129,32 +128,15 @@ class HunyuanImage3Pipeline(HunyuanImage3PreTrainedModel, GenerationMixin): "time_embed_2", "final_layer.model", ] - filtered_sd = {} - for k, v in state_dict.items(): - if any(k.startswith(prefix) for prefix in non_layer_prefixes): - filtered_sd[k] = v + tp_rank = get_tensor_model_parallel_rank() + device_str = f"{self.model.device.type}:{tp_rank}" + named_modules = dict(self.named_modules()) + for prefix in non_model_layer_prefixes: + mod = named_modules.get(prefix) + if mod: + mod.to(device_str) - missing, unexpected = self.load_state_dict(filtered_sd, strict=False) - - for prefix in non_layer_prefixes: - if hasattr(self, prefix.split(".")[0]): - module = dict(self.named_modules()).get(prefix) - if module: - module.to(f"{self.model.device.type}:{tp_rank}") - - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - self.pre_load() - skip_prefixes = ["lm_head."] if self.hf_config.tie_word_embeddings else [] - # List of unexpected keywords in weight names unexpected_keywords = [ - "vae", - "vision_aligner", - "vision_model", - "final_layer", - "patch_embed", - "timestep_emb", - "time_embed", - "time_embed_2", "guidance_emb", "timestep_r_emb", ] @@ -892,7 +874,7 @@ class HunyuanImage3Pipeline(HunyuanImage3PreTrainedModel, GenerationMixin): custom_pos_emb = self.get_pos_emb(custom_pos_emb, position_ids) - inputs_embeds = self.model.wte(input_ids) + inputs_embeds = self.model.embed_tokens(input_ids) bsz, seq_len, n_embd = inputs_embeds.shape diff --git a/vllm_omni/diffusion/models/longcat_image/longcat_image_transformer.py b/vllm_omni/diffusion/models/longcat_image/longcat_image_transformer.py index 2f6d54aef..329fc648b 100644 --- a/vllm_omni/diffusion/models/longcat_image/longcat_image_transformer.py +++ b/vllm_omni/diffusion/models/longcat_image/longcat_image_transformer.py @@ -517,17 +517,18 @@ class LongCatImageTransformer2DModel(nn.Module): def __init__( self, od_config: OmniDiffusionConfig, - patch_size: int = 1, - in_channels: int = 64, - num_layers: int = 19, - num_single_layers: int = 38, - attention_head_dim: int = 128, - num_attention_heads: int = 24, - joint_attention_dim: int = 3584, - pooled_projection_dim: int = 3584, - axes_dims_rope: list[int] = [16, 56, 56], ): super().__init__() + model_config = od_config.tf_model_config + patch_size = model_config.patch_size + in_channels = model_config.in_channels + num_layers = model_config.num_layers + num_single_layers = model_config.num_single_layers + attention_head_dim = model_config.attention_head_dim + num_attention_heads = model_config.num_attention_heads + joint_attention_dim = model_config.joint_attention_dim + pooled_projection_dim = model_config.pooled_projection_dim + axes_dims_rope = getattr(model_config, "axes_dims_rope", [16, 56, 56]) self.out_channels = in_channels self.inner_dim = num_attention_heads * attention_head_dim self.pooled_projection_dim = pooled_projection_dim diff --git a/vllm_omni/diffusion/worker/diffusion_model_runner.py b/vllm_omni/diffusion/worker/diffusion_model_runner.py index a9fd48261..263ce3126 100644 --- a/vllm_omni/diffusion/worker/diffusion_model_runner.py +++ b/vllm_omni/diffusion/worker/diffusion_model_runner.py @@ -102,7 +102,7 @@ class DiffusionModelRunner: # Load model within forward context load_config = LoadConfig() - model_loader = DiffusersPipelineLoader(load_config) + model_loader = DiffusersPipelineLoader(load_config, od_config=self.od_config) time_before_load = time.perf_counter() with get_memory_context(): @@ -191,7 +191,11 @@ class DiffusionModelRunner: grad_context = torch.no_grad() if use_hsdp else torch.inference_mode() with grad_context: # The manager handles the check for need_recv_cache internally - self.kv_transfer_manager.receive_kv_cache(req, target_device=getattr(self.pipeline, "device", None)) + self.kv_transfer_manager.receive_multi_kv_cache( + req, + cfg_kv_collect_func=getattr(self.od_config, "cfg_kv_collect_func", None), + target_device=getattr(self.pipeline, "device", None), + ) if req.sampling_params.generator is None and req.sampling_params.seed is not None: if req.sampling_params.generator_device is not None: diff --git a/vllm_omni/distributed/omni_connectors/kv_transfer_manager.py b/vllm_omni/distributed/omni_connectors/kv_transfer_manager.py index 7c7d90ca1..21f09c951 100644 --- a/vllm_omni/distributed/omni_connectors/kv_transfer_manager.py +++ b/vllm_omni/distributed/omni_connectors/kv_transfer_manager.py @@ -522,3 +522,50 @@ class OmniKVTransferManager: self.apply_kv_cache_to_request(req, data) return True return False + + def receive_multi_kv_cache( + self, + req: Any, + cfg_kv_collect_func: Callable | None = None, + target_device: torch.device | None = None, + ) -> bool: + """Receive primary KV cache and optional CFG companion KV caches. + + First receives the primary KV cache (existing logic). Then, if the + request carries cfg_kv_request_ids and a model-specific + cfg_kv_collect_func is provided, calls it to fetch and attach the + companion KV caches to sampling_params. + + Args: + req: Request object with request_id and sampling_params. + cfg_kv_collect_func: Model-specific function for collecting + CFG KV caches. Signature: + (request_id, cfg_request_ids, kv_transfer_manager, target_device) + -> dict[str, Any] + target_device: Device to move tensors to. + + Returns: + True if primary KV cache was received successfully. + """ + primary_ok = self.receive_kv_cache(req, target_device) + + cfg_ids = getattr(getattr(req, "sampling_params", None), "cfg_kv_request_ids", None) + if cfg_ids and cfg_kv_collect_func: + request_id = getattr(req, "request_id", None) or ( + req.request_ids[0] if hasattr(req, "request_ids") and req.request_ids else None + ) + try: + cfg_kvs = cfg_kv_collect_func( + request_id, + cfg_ids, + self, + target_device, + ) + if cfg_kvs and hasattr(req, "sampling_params") and req.sampling_params is not None: + for key, value in cfg_kvs.items(): + setattr(req.sampling_params, key, value) + logger.info("Applied CFG KV caches: %s", list(cfg_kvs.keys())) + except Exception: + logger.exception("Failed to collect CFG KV caches for %s", request_id) + + return primary_ok diff --git a/vllm_omni/entrypoints/async_omni.py b/vllm_omni/entrypoints/async_omni.py index 19c8b07ac..4af60b1a2 100644 --- a/vllm_omni/entrypoints/async_omni.py +++ b/vllm_omni/entrypoints/async_omni.py @@ -190,6 +190,8 @@ class AsyncOmni(OmniBase): "enforce_eager": kwargs.get("enforce_eager", False), "diffusion_load_format": kwargs.get("diffusion_load_format", "default"), "custom_pipeline_args": kwargs.get("custom_pipeline_args", None), + "enable_multithread_weight_load": kwargs.get("enable_multithread_weight_load", True), + "num_weight_load_threads": kwargs.get("num_weight_load_threads", 4), }, "final_output": True, "final_output_type": "image", diff --git a/vllm_omni/entrypoints/async_omni_diffusion.py b/vllm_omni/entrypoints/async_omni_diffusion.py index 34187ec14..b3c02892c 100644 --- a/vllm_omni/entrypoints/async_omni_diffusion.py +++ b/vllm_omni/entrypoints/async_omni_diffusion.py @@ -60,6 +60,7 @@ class AsyncOmniDiffusion: # Capture stage info from kwargs before they might be filtered out stage_id = kwargs.get("stage_id") engine_input_source = kwargs.get("engine_input_source") + cfg_kv_collect_func = kwargs.pop("cfg_kv_collect_func", None) # Build config if od_config is None: @@ -115,6 +116,9 @@ class AsyncOmniDiffusion: else: raise + if cfg_kv_collect_func is not None: + od_config.cfg_kv_collect_func = cfg_kv_collect_func + # Initialize engine self.engine: DiffusionEngine = DiffusionEngine.make_engine(od_config) diff --git a/vllm_omni/entrypoints/cfg_companion_tracker.py b/vllm_omni/entrypoints/cfg_companion_tracker.py new file mode 100644 index 000000000..9c2e835f0 --- /dev/null +++ b/vllm_omni/entrypoints/cfg_companion_tracker.py @@ -0,0 +1,233 @@ +"""CFG companion request tracker for the Omni orchestrator. + +Encapsulates all bookkeeping for Classifier-Free Guidance companion +requests (prompt expansion, parent/companion ID mapping, completion +tracking, deferred forwarding, failure propagation, and timeouts) +so that ``Omni._run_generation`` stays clean. +""" + +from __future__ import annotations + +import copy +import logging +import os +import time +from collections.abc import Callable, Sequence +from typing import Any + +from vllm_omni.distributed.omni_connectors.adapter import try_send_via_connector +from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniSamplingParams + +logger = logging.getLogger(__name__) + + +class CfgCompanionTracker: + """Manages CFG companion request lifecycle in the orchestrator scheduling loop.""" + + def __init__( + self, + prompt_expand_func: Callable[..., Any] | None, + stage0_sampling_params: Any, + timeout_s: float | None = None, + ) -> None: + self._expand_func = prompt_expand_func + self._sp0 = stage0_sampling_params + self._timeout_s = ( + timeout_s if timeout_s is not None else float(os.environ.get("VLLM_CFG_PENDING_TIMEOUT_S", "120")) + ) + + self._companion_map: dict[str, dict[str, str]] = {} # parent -> {role: companion_id} + self._companion_ids: set[str] = set() + self._companion_to_parent: dict[str, str] = {} # companion -> parent + self._done: dict[str, set[str]] = {} # parent -> completed companion ids + self._pending_parents: dict[str, dict[str, Any]] = {} # parent -> deferred result + self._failed_parents: set[str] = set() + + @property + def is_active(self) -> bool: + return bool(self._companion_ids) + + @property + def num_companions(self) -> int: + return len(self._companion_ids) + + @property + def stage0_sampling_params(self) -> Any: + return self._sp0 + + def expand_prompts( + self, + request_id_to_prompt: dict[str, Any], + ) -> list[tuple[str, Any]]: + """Expand user prompts into ``(companion_id, prompt)`` pairs via model-specific func.""" + if not self._expand_func: + return [] + + pairs: list[tuple[str, Any]] = [] + for rid, prompt in request_id_to_prompt.items(): + expanded = self._expand_func(prompt, self._sp0) + if not expanded: + continue + role_map: dict[str, str] = {} + for ep in expanded: + cid = f"{rid}{ep.request_id_suffix}" + role_map[ep.role] = cid + self._companion_ids.add(cid) + self._companion_to_parent[cid] = rid + pairs.append((cid, ep.prompt)) + self._companion_map[rid] = role_map + self._done[rid] = set() + + logger.debug( + "CFG expansion: %d parent(s) -> %d companion(s)", + len(self._companion_map), + len(self._companion_ids), + ) + return pairs + + def is_companion(self, req_id: str) -> bool: + return req_id in self._companion_ids + + def has_companions(self, parent_id: str) -> bool: + return parent_id in self._companion_map + + def all_companions_done(self, parent_id: str) -> bool: + role_map = self._companion_map.get(parent_id, {}) + done_set = self._done.get(parent_id, set()) + return all(cid in done_set for cid in role_map.values()) + + def get_companion_request_ids(self, parent_id: str) -> dict[str, str]: + """Return ``{role: companion_request_id}`` for a parent.""" + return self._companion_map.get(parent_id, {}) + + def is_parent_failed(self, parent_id: str) -> bool: + return parent_id in self._failed_parents + + # -- Lifecycle events -- + + def on_companion_error(self, companion_id: str) -> tuple[str | None, bool]: + """Record failure. Returns ``(parent_id, parent_was_aborted)``.""" + parent_id = self._companion_to_parent.get(companion_id) + if parent_id is None: + return None, False + self._failed_parents.add(parent_id) + logger.error("CFG companion %s failed; marking parent %s as failed", companion_id, parent_id) + aborted = parent_id in self._pending_parents + if aborted: + self._pending_parents.pop(parent_id, None) + return parent_id, aborted + + def on_companion_completed(self, companion_id: str) -> str | None: + """Mark done. Returns parent_id only if parent is pending and all companions finished.""" + parent_id = self._companion_to_parent.get(companion_id) + if parent_id is None: + return None + self._done[parent_id].add(companion_id) + logger.debug("CFG companion %s completed (parent=%s)", companion_id, parent_id) + if parent_id in self._pending_parents and self.all_companions_done(parent_id): + return parent_id + return None + + def consume_parent_failure(self, parent_id: str) -> None: + self._failed_parents.discard(parent_id) + + # -- Deferred parent management -- + + def defer_parent(self, parent_id: str, engine_outputs: Any, stage_id: int) -> None: + """Hold parent result while waiting for companions to finish.""" + self._pending_parents[parent_id] = { + "engine_outputs": engine_outputs, + "stage_id": stage_id, + "pending_since": time.time(), + } + logger.debug("Parent %s deferred, waiting for CFG companions", parent_id) + + def pop_pending_parent(self, parent_id: str) -> dict[str, Any] | None: + return self._pending_parents.pop(parent_id, None) + + def check_timeouts(self) -> list[str]: + """Return and remove parent IDs that exceeded the pending timeout.""" + if not self._pending_parents: + return [] + now = time.time() + timed_out: list[str] = [] + for pid in list(self._pending_parents): + pending_since = self._pending_parents[pid].get("pending_since", now) + if now - pending_since > self._timeout_s: + self._pending_parents.pop(pid) + self._failed_parents.discard(pid) + timed_out.append(pid) + logger.error("Parent %s timed out waiting for CFG companions (>%.0fs)", pid, self._timeout_s) + return timed_out + + # -- Forward parent with CFG KV -- + + def forward_parent_with_cfg( + self, + req_id: str, + parent_result: dict[str, Any], + stage_list: Sequence[Any], + connectors: dict[tuple[str, str], Any], + sampling_params_list: Sequence[OmniSamplingParams], + request_id_to_prompt: dict[str, Any], + final_stage_id_to_prompt: dict[str, int], + metrics: Any, + remaining_by_stage: list[int], + ) -> bool: + """Forward a parent request to the next stage with CFG KV request IDs attached.""" + stage_id = parent_result["stage_id"] + next_stage_id = stage_id + 1 + if next_stage_id > final_stage_id_to_prompt.get(req_id, 0): + return True + + next_stage = stage_list[next_stage_id] + try: + with metrics.stage_postprocess_timer(stage_id, req_id): + next_inputs = next_stage.process_engine_inputs( + stage_list, + [request_id_to_prompt[req_id]], + source_outputs_override=parent_result["engine_outputs"], + ) + except Exception as e: + logger.exception( + "Process engine inputs error for req %s at stage %d: %s", + req_id, + next_stage_id, + e, + ) + return False + + sp_next = copy.deepcopy(sampling_params_list[next_stage_id]) + if isinstance(sp_next, OmniDiffusionSamplingParams): + sp_next.cfg_kv_request_ids = self.get_companion_request_ids(req_id) + logger.info( + "Attaching cfg_kv_request_ids=%s to request %s", + sp_next.cfg_kv_request_ids, + req_id, + ) + + connector_key = (str(stage_id), str(next_stage_id)) + connector = connectors.get(connector_key) + sent_via_connector = False + if connector: + sent_via_connector = try_send_via_connector( + connector=connector, + stage_id=stage_id, + next_stage_id=next_stage_id, + req_id=req_id, + next_inputs=next_inputs, + sampling_params=sp_next, + original_prompt=request_id_to_prompt[req_id], + next_stage_queue_submit_fn=stage_list[next_stage_id].submit, + metrics=metrics, + ) + + if not sent_via_connector: + raise RuntimeError( + f"Failed to send CFG request {req_id} to stage-{next_stage_id} via connector. " + "Configure a connector for this edge or inspect connector logs for details." + ) + + logger.debug("Forwarded CFG-enabled request %s to stage-%d", req_id, next_stage_id) + remaining_by_stage[next_stage_id] += 1 + return True diff --git a/vllm_omni/entrypoints/cli/serve.py b/vllm_omni/entrypoints/cli/serve.py index c9ed7150f..f99e6c77a 100644 --- a/vllm_omni/entrypoints/cli/serve.py +++ b/vllm_omni/entrypoints/cli/serve.py @@ -272,6 +272,21 @@ class OmniServeCommand(CLISubcommand): help="Enable VAE tiling for memory optimization (useful for mitigating OOM issues).", ) + # Parallel weight loading (faster diffusion startup) + omni_config_group.add_argument( + "--disable-multithread-weight-load", + action="store_false", + dest="enable_multithread_weight_load", + default=True, + help="Disable multi-threaded safetensors loading (default: enabled with 4 threads).", + ) + omni_config_group.add_argument( + "--num-weight-load-threads", + type=int, + default=4, + help="Number of threads for parallel weight loading (default: 4).", + ) + # diffusion model offload parameters omni_config_group.add_argument( "--enable-cpu-offload", diff --git a/vllm_omni/entrypoints/omni.py b/vllm_omni/entrypoints/omni.py index 5880ec594..58dd88eb0 100644 --- a/vllm_omni/entrypoints/omni.py +++ b/vllm_omni/entrypoints/omni.py @@ -34,6 +34,7 @@ from vllm_omni.distributed.ray_utils.utils import ( get_ray_queue_class, try_close_ray, ) +from vllm_omni.entrypoints.cfg_companion_tracker import CfgCompanionTracker from vllm_omni.entrypoints.omni_stage import OmniStage from vllm_omni.entrypoints.stage_utils import SHUTDOWN_TASK, OmniStageTaskType from vllm_omni.entrypoints.stage_utils import maybe_load_from_ipc as _load @@ -935,6 +936,13 @@ class Omni(OmniBase): _req_start_ts: dict[str, float] = {} _wall_start_ts: float = time.time() + # CFG companion tracking (prompt expansion + lifecycle management) + cfg = CfgCompanionTracker( + prompt_expand_func=getattr(self.stage_list[0], "prompt_expand_func", None), + stage0_sampling_params=sampling_params_list[0], + ) + expanded_companions = cfg.expand_prompts(request_id_to_prompt) + # Determine the final stage for E2E stats (highest stage_id with final_output=True; fallback to last stage) final_stage_id_to_prompt: dict[str, int] = {} for rid, prompt in request_id_to_prompt.items(): @@ -976,6 +984,18 @@ class Omni(OmniBase): _req_start_ts[req_id] = time.time() logger.debug(f"[{self._name}] Enqueued request {req_id} to stage-0") + # Submit CFG companion requests to stage-0 + if cfg.is_active: + for companion_id, companion_prompt in expanded_companions: + task = { + "request_id": companion_id, + "engine_inputs": companion_prompt, + "sampling_params": cfg.stage0_sampling_params, + } + self.stage_list[0].submit(task) + _req_start_ts[companion_id] = time.time() + logger.debug(f"[{self._name}] Enqueued CFG companion {companion_id} to stage-0") + pbar = None if use_tqdm: tqdm_func = use_tqdm if callable(use_tqdm) else tqdm @@ -987,7 +1007,7 @@ class Omni(OmniBase): ) # For each stage, forward results to next stage; collect finals at the end # We pipeline by continually polling output queues in stage order - remaining_by_stage: list[int] = [len(request_prompts)] + [0] * (num_stages - 1) + remaining_by_stage: list[int] = [len(request_prompts) + cfg.num_companions] + [0] * (num_stages - 1) completed_requests = 0 total_requests = len(request_prompts) @@ -1007,6 +1027,14 @@ class Omni(OmniBase): logger.error( f"[{self._name}] Stage {stage_id} error on request {req_id}: {result['error']}", ) + if cfg.is_companion(req_id) and stage_id == 0: + parent_id, parent_aborted = cfg.on_companion_error(req_id) + if parent_aborted: + completed_requests += 1 + logger.error( + f"[{self._name}] Parent {parent_id} aborted due to " + f"companion failure ({completed_requests}/{total_requests})", + ) continue if result.get("type") == "stage_ready": @@ -1015,6 +1043,30 @@ class Omni(OmniBase): time.sleep(0.05) continue + # CFG: companion requests only run through Stage-0 + if cfg.is_companion(req_id) and stage_id == 0: + ready_parent = cfg.on_companion_completed(req_id) + if ready_parent is not None: + success = cfg.forward_parent_with_cfg( + ready_parent, + cfg.pop_pending_parent(ready_parent), + self.stage_list, + self.connectors, + sampling_params_list, + request_id_to_prompt, + final_stage_id_to_prompt, + metrics, + remaining_by_stage, + ) + if not success: + cfg.consume_parent_failure(ready_parent) + completed_requests += 1 + logger.error( + f"[{self._name}] Parent {ready_parent} dropped due to CFG forwarding failure " + f"({completed_requests}/{total_requests})", + ) + continue + engine_outputs = _load(result, obj_key="engine_outputs", shm_key="engine_outputs_shm") # Mark last output time for this stage whenever we receive outputs metrics.stage_last_ts[stage_id] = max(metrics.stage_last_ts[stage_id] or 0.0, time.time()) @@ -1104,6 +1156,40 @@ class Omni(OmniBase): next_stage_id = stage_id + 1 if next_stage_id <= final_stage_id_to_prompt[req_id]: + # CFG: if this parent has companions, defer forwarding + if cfg.has_companions(req_id) and stage_id == 0: + if cfg.is_parent_failed(req_id): + cfg.consume_parent_failure(req_id) + completed_requests += 1 + logger.error( + f"[{self._name}] Parent {req_id} skipped CFG forwarding due to " + f"companion failure ({completed_requests}/{total_requests})", + ) + continue + + if cfg.all_companions_done(req_id): + success = cfg.forward_parent_with_cfg( + req_id, + {"engine_outputs": engine_outputs, "stage_id": stage_id}, + self.stage_list, + self.connectors, + sampling_params_list, + request_id_to_prompt, + final_stage_id_to_prompt, + metrics, + remaining_by_stage, + ) + if not success: + cfg.consume_parent_failure(req_id) + completed_requests += 1 + logger.error( + f"[{self._name}] Parent {req_id} dropped due to CFG forwarding failure " + f"({completed_requests}/{total_requests})", + ) + else: + cfg.defer_parent(req_id, engine_outputs, stage_id) + continue + next_stage: OmniStage = self.stage_list[next_stage_id] try: # Derive inputs for the next stage, record preprocess time @@ -1112,9 +1198,10 @@ class Omni(OmniBase): self.stage_list, [request_id_to_prompt[req_id]] ) except Exception as e: + completed_requests += 1 logger.exception( f"[{self._name}] Process engine inputs error for req {req_id}" - f" at stage {next_stage_id}: {e}", + f" at stage {next_stage_id}: {e} ({completed_requests}/{total_requests})", ) continue sp_next = sampling_params_list[next_stage_id] # type: ignore[index] @@ -1141,6 +1228,7 @@ class Omni(OmniBase): f"[{self._name}] Failed to send request {req_id} to stage-{next_stage_id} via connector. " "Configure a connector for this edge or inspect connector logs for details." ) + logger.debug( f"[{self._name}] Forwarded request {req_id} to stage-{next_stage_id}", ) @@ -1154,6 +1242,12 @@ class Omni(OmniBase): logger.debug( f"[{self._name}] Request {req_id} fully completed ({completed_requests}/{total_requests})", ) + for timed_out_id in cfg.check_timeouts(): + completed_requests += 1 + logger.error( + f"[{self._name}] Parent {timed_out_id} timed out; counting as failed " + f"({completed_requests}/{total_requests})", + ) if not made_progress: time.sleep(0.005) diff --git a/vllm_omni/entrypoints/omni_diffusion.py b/vllm_omni/entrypoints/omni_diffusion.py index 0163f8988..12db495c5 100644 --- a/vllm_omni/entrypoints/omni_diffusion.py +++ b/vllm_omni/entrypoints/omni_diffusion.py @@ -28,6 +28,7 @@ class OmniDiffusion: # Capture stage info from kwargs before they might be filtered out stage_id = kwargs.get("stage_id") engine_input_source = kwargs.get("engine_input_source") + cfg_kv_collect_func = kwargs.pop("cfg_kv_collect_func", None) if od_config is None: od_config = OmniDiffusionConfig.from_kwargs(**kwargs) @@ -94,6 +95,9 @@ class OmniDiffusion: od_config.tf_model_config = TransformerConfig() od_config.update_multimodal_support() + if cfg_kv_collect_func is not None: + od_config.cfg_kv_collect_func = cfg_kv_collect_func + self.engine: DiffusionEngine = DiffusionEngine.make_engine(od_config) def generate( diff --git a/vllm_omni/entrypoints/omni_stage.py b/vllm_omni/entrypoints/omni_stage.py index 0d2d9121e..4114ef3fb 100644 --- a/vllm_omni/entrypoints/omni_stage.py +++ b/vllm_omni/entrypoints/omni_stage.py @@ -47,6 +47,7 @@ from vllm_omni.entrypoints.stage_utils import ( _resolve_model_tokenizer_paths, _to_dict, is_profiler_task, + load_func_from_config, maybe_dump_to_shm, set_stage_devices, ) @@ -280,6 +281,7 @@ class OmniStage: else: self.custom_process_input_func = None + self.prompt_expand_func = load_func_from_config(getattr(stage_config, "prompt_expand_func", None)) self.final_output = getattr(stage_config, "final_output", False) self.final_output_type = getattr(stage_config, "final_output_type", None) self.tts_args = _to_dict(getattr(stage_config, "tts_args", {})) @@ -460,6 +462,7 @@ class OmniStage: "connectors_config": connectors_config or {}, "stage_type": self.stage_type, "engine_input_source": self.engine_input_source, + "cfg_kv_collect_func": getattr(self.stage_config, "cfg_kv_collect_func", None), "final_output": self.final_output, "final_output_type": self.final_output_type, } @@ -612,7 +615,11 @@ class OmniStage: return None def process_engine_inputs( - self, stage_list: list[Any], prompt: OmniTokensPrompt | TextPrompt = None + self, + stage_list: list[Any], + prompt: OmniTokensPrompt | TextPrompt = None, + *, + source_outputs_override: Any = None, ) -> list[OmniTokensPrompt | TextPrompt]: """Process engine inputs for this stage from upstream stage outputs. @@ -623,6 +630,8 @@ class OmniStage: Args: stage_list: List of all stages in the pipeline prompt: Optional original prompt (for multimodal data preservation) + source_outputs_override: Use these outputs instead of reading from + the source stage's ``engine_outputs`` (for deferred CFG requests). Returns: List of processed engine inputs ready for this stage @@ -635,7 +644,11 @@ class OmniStage: if len(self.engine_input_source) == 0: raise ValueError("engine_input_source is empty") source_stage_id = self.engine_input_source[0] - source_outputs = stage_list[source_stage_id].engine_outputs + source_outputs = ( + source_outputs_override + if source_outputs_override is not None + else stage_list[source_stage_id].engine_outputs + ) if not isinstance(prompt, list): prompt = [prompt] multi_modal_data = { @@ -657,6 +670,18 @@ class OmniStage: else: engine_input_source = self.engine_input_source + if source_outputs_override is not None and engine_input_source: + # Temporarily swap engine_outputs so custom_process_input_func + # (which reads stage_list directly) sees the correct data. + _source_id = engine_input_source[0] + _orig_outputs = stage_list[_source_id].engine_outputs + stage_list[_source_id].engine_outputs = source_outputs_override + try: + return self.custom_process_input_func( + stage_list, engine_input_source, prompt, self.requires_multimodal_data + ) + finally: + stage_list[_source_id].engine_outputs = _orig_outputs return self.custom_process_input_func( stage_list, engine_input_source, prompt, self.requires_multimodal_data ) @@ -702,6 +727,8 @@ def _stage_worker( connectors_config = stage_payload.get("connectors_config", {}) stage_type: Literal["llm", "diffusion"] = stage_payload.get("stage_type", "llm") + cfg_kv_collect_func = load_func_from_config(stage_payload.get("cfg_kv_collect_func")) + if stage_type != "diffusion": _resolve_worker_cls(engine_args) @@ -761,6 +788,7 @@ def _stage_worker( model=model, stage_id=stage_id, engine_input_source=stage_payload.get("engine_input_source", []), + cfg_kv_collect_func=cfg_kv_collect_func, **engine_args, ) else: @@ -1109,6 +1137,7 @@ async def _stage_worker_async( final_output = stage_payload.get("final_output", False) final_output_type = stage_payload.get("final_output_type", None) + cfg_kv_collect_func = load_func_from_config(stage_payload.get("cfg_kv_collect_func")) # Handle non-standard model directory structures (e.g., tokenizer in root, model in subdir) model = _resolve_model_tokenizer_paths(model, engine_args) @@ -1187,10 +1216,13 @@ async def _stage_worker_async( od_config["omni_kv_config"]["engine_input_source"] = stage_payload.get("engine_input_source", []) logger.debug(f"[Stage-%s] Initializing diffusion engine with config: {od_config}", stage_id) + _diffusion_kwargs = {k: v for k, v in engine_args.items() if k not in {"od_config", "model"}} + if cfg_kv_collect_func is not None: + _diffusion_kwargs["cfg_kv_collect_func"] = cfg_kv_collect_func stage_engine = AsyncOmniDiffusion( model=model, od_config=od_config, - **{k: v for k, v in engine_args.items() if k not in {"od_config", "model"}}, + **_diffusion_kwargs, ) vllm_config = None # Diffusion doesn't use vllm_config else: diff --git a/vllm_omni/entrypoints/openai/serving_chat.py b/vllm_omni/entrypoints/openai/serving_chat.py index b9622124d..71bb618ca 100644 --- a/vllm_omni/entrypoints/openai/serving_chat.py +++ b/vllm_omni/entrypoints/openai/serving_chat.py @@ -13,7 +13,6 @@ import torch from fastapi import Request from PIL import Image from pydantic import TypeAdapter -from vllm.renderers import BaseRenderer from vllm_omni.entrypoints.async_omni import AsyncOmni from vllm_omni.entrypoints.openai.protocol.chat_completion import OmniChatCompletionResponse @@ -67,7 +66,7 @@ from vllm.inputs.data import PromptType from vllm.logger import init_logger from vllm.outputs import RequestOutput from vllm.reasoning import ReasoningParser -from vllm.renderers import merge_kwargs +from vllm.renderers import BaseRenderer, merge_kwargs from vllm.renderers.inputs import TokPrompt from vllm.sampling_params import SamplingParams from vllm.tokenizers import TokenizerLike diff --git a/vllm_omni/entrypoints/stage_utils.py b/vllm_omni/entrypoints/stage_utils.py index 5eeb1ab0f..fade4c844 100644 --- a/vllm_omni/entrypoints/stage_utils.py +++ b/vllm_omni/entrypoints/stage_utils.py @@ -1,9 +1,11 @@ from __future__ import annotations import enum +import importlib import json import logging import os +from collections.abc import Callable from multiprocessing import shared_memory as _shm from typing import Any @@ -12,6 +14,22 @@ from omegaconf import OmegaConf logger = logging.getLogger(__name__) +def load_func_from_config(func_path: str | None) -> Callable[..., Any] | None: + """Dynamically import a callable from a fully-qualified dotted path. + + Args: + func_path: Dotted path such as ``"pkg.module.func_name"``, or *None*. + + Returns: + The imported callable, or *None* when *func_path* is falsy. + """ + if not func_path: + return None + module_path, func_name = func_path.rsplit(".", 1) + module = importlib.import_module(module_path) + return getattr(module, func_name) + + class OmniStageTaskType(enum.Enum): GENERATE = "generate" ABORT = "abort" diff --git a/vllm_omni/inputs/data.py b/vllm_omni/inputs/data.py index bf8563bd2..438abb971 100644 --- a/vllm_omni/inputs/data.py +++ b/vllm_omni/inputs/data.py @@ -229,6 +229,13 @@ class OmniDiffusionSamplingParams: kv_metadata: dict[str, Any] | None = None # Metadata for KV Cache (e.g., kv_lens, ropes) need_kv_receive: bool = True # Flag to indicate if this request expects KV transfer + # [Omni] Multi-KV for CFG: populated by model-specific cfg_kv_collect_func + cfg_text_past_key_values: Any | None = None + cfg_img_past_key_values: Any | None = None + cfg_text_kv_metadata: dict[str, Any] | None = None + cfg_img_kv_metadata: dict[str, Any] | None = None + cfg_kv_request_ids: dict[str, str] | None = None + # Component modules modules: dict[str, Any] = field(default_factory=dict) 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 46b07f9de..f798a21d2 100644 --- a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py +++ b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py @@ -410,11 +410,16 @@ class Qwen3OmniMoeForConditionalGeneration( talker_hidden = model_outputs # merge the code_predictor_codes from the info_dict list into a single tensor multimodal_outputs: dict = None - # Here is the only place to use runtime_additional_information. After MTP in the + # Here is the only place to use model_intermediate_buffer. After MTP in the # preprocess function, the code_predictor_codes are stored in the info_dict list. # We need to merge the tensors from different requests into a single tensor. # In the future, we may allow user to custom an aggregated function. - info_dicts = kwargs.get("runtime_additional_information") + info_dicts = kwargs.get("model_intermediate_buffer") + if info_dicts is None: + info_dicts = kwargs.get("runtime_additional_information") + + if "runtime_additional_information" in kwargs and "model_intermediate_buffer" not in kwargs: + logger.warning_once("runtime_additional_information is deprecated, use model_intermediate_buffer") code_predictor_codes = [info.get("code_predictor_codes") for info in info_dicts] multimodal_outputs = {"code_predictor_codes": torch.cat(code_predictor_codes, dim=0)} span_len = multimodal_outputs["code_predictor_codes"].shape[0] diff --git a/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts.py b/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts.py index 8c17fa50b..028d4a53d 100644 --- a/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts.py +++ b/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts.py @@ -13,6 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. import base64 +import copy import io import urllib.request from collections.abc import Iterable @@ -151,9 +152,14 @@ class Qwen3TTSModelForGeneration(nn.Module): # Extract additional parameters from kwargs that the generation methods expect - runtime_additional_information = kwargs.get("runtime_additional_information", [{}]) + runtime_additional_information = kwargs.get("model_intermediate_buffer") + if runtime_additional_information is None: + runtime_additional_information = kwargs.get("runtime_additional_information", [{}]) + if "runtime_additional_information" in kwargs and "model_intermediate_buffer" not in kwargs: + logger.warning_once("runtime_additional_information is deprecated, use model_intermediate_buffer") if isinstance(runtime_additional_information, list) and len(runtime_additional_information) > 0: runtime_additional_information = runtime_additional_information[0] + runtime_additional_information = copy.deepcopy(runtime_additional_information) text = runtime_additional_information.pop("text", [""])[0] # Extract task_type from kwargs, default to self.task_type task_type = _normalize_task_type(runtime_additional_information.pop("task_type", [self.task_type])[0]) diff --git a/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_talker.py b/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_talker.py index a39eded3a..1bdcd566e 100644 --- a/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_talker.py +++ b/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_talker.py @@ -332,7 +332,6 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module): self.have_multimodal_outputs = True self.has_preprocess = True self.has_postprocess = True - # Used by OmniGPUModelRunner for the GPU-side MTP fast-path. self.mtp_hidden_size = int(self.talker_config.hidden_size) # OmniGPUModelRunner will store talker_mtp output under this key in @@ -441,7 +440,11 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module): return model_outputs hidden = model_outputs - info_dicts = kwargs.get("runtime_additional_information") or [] + info_dicts = kwargs.get("model_intermediate_buffer") + if info_dicts is None: + info_dicts = kwargs.get("runtime_additional_information") or [] + if "runtime_additional_information" in kwargs and "model_intermediate_buffer" not in kwargs: + logger.warning_once("runtime_additional_information is deprecated, use model_intermediate_buffer") audio_codes_list: list[torch.Tensor] = [] ref_code_len_list: list[torch.Tensor] = [] codec_streaming_list: list[torch.Tensor] = [] diff --git a/vllm_omni/model_executor/stage_configs/bagel.yaml b/vllm_omni/model_executor/stage_configs/bagel.yaml index c60f3b700..23bbb441f 100644 --- a/vllm_omni/model_executor/stage_configs/bagel.yaml +++ b/vllm_omni/model_executor/stage_configs/bagel.yaml @@ -3,9 +3,12 @@ stage_args: - stage_id: 0 stage_type: llm + prompt_expand_func: vllm_omni.model_executor.stage_input_processors.bagel.expand_cfg_prompts runtime: devices: "0" - max_batch_size: 1 + # 2 = 1 user prompt + 1 CFG companion (text-unconditional). + # For multi-prompt batches this should scale as batch_size × 2. + max_batch_size: 2 engine_args: model_stage: thinker model_arch: BagelForConditionalGeneration @@ -37,6 +40,7 @@ stage_args: - stage_id: 1 stage_type: diffusion + cfg_kv_collect_func: vllm_omni.model_executor.stage_input_processors.bagel.collect_cfg_kv_caches runtime: devices: "0" max_batch_size: 1 diff --git a/vllm_omni/model_executor/stage_configs/qwen3_omni_moe.yaml b/vllm_omni/model_executor/stage_configs/qwen3_omni_moe.yaml index 802f8dd24..eb5bf7101 100644 --- a/vllm_omni/model_executor/stage_configs/qwen3_omni_moe.yaml +++ b/vllm_omni/model_executor/stage_configs/qwen3_omni_moe.yaml @@ -72,7 +72,7 @@ stage_args: stage_type: llm # Use llm stage type to launch OmniLLM runtime: devices: "1" - max_batch_size: 64 + max_batch_size: 32 engine_args: model_stage: code2wav model_arch: Qwen3OmniMoeForConditionalGeneration diff --git a/vllm_omni/model_executor/stage_input_processors/bagel.py b/vllm_omni/model_executor/stage_input_processors/bagel.py new file mode 100644 index 000000000..80b28b230 --- /dev/null +++ b/vllm_omni/model_executor/stage_input_processors/bagel.py @@ -0,0 +1,156 @@ +"""Stage input processor for Bagel: CFG prompt expansion and KV cache collection. + +Bagel's 3-branch CFG requires multiple prompts through the AR stage: + - gen (conditional): user prompt + - cfg_text (text unconditional): negative/empty prompt + - cfg_img (image unconditional): user prompt without image (same as gen for text2img) + +This module provides model-specific functions referenced by bagel.yaml: + - prompt_expand_func -> expand_cfg_prompts + - cfg_kv_collect_func -> collect_cfg_kv_caches +""" + +from __future__ import annotations + +import logging +from dataclasses import dataclass +from types import SimpleNamespace +from typing import Any + +logger = logging.getLogger(__name__) + +CFG_TEXT_SUFFIX = "__cfg_text" + + +@dataclass +class ExpandedPrompt: + """A single expanded prompt produced by the prompt expansion function.""" + + prompt: dict[str, Any] | str + role: str + request_id_suffix: str + + +def expand_cfg_prompts( + prompt: dict[str, Any] | str, + sampling_params: Any, +) -> list[ExpandedPrompt]: + """Expand a user prompt into additional prompts needed for Bagel CFG. + + For text2img (modalities contains "image", no multi_modal_data with image): + - One extra prompt for cfg_text (negative/empty prompt) + - cfg_img reuses gen KV for text2img, so no extra prompt needed + + For text2text or img2text: returns empty list (no expansion needed). + + Args: + prompt: The original user prompt (dict or string). + sampling_params: The Stage-0 sampling params (may carry extra_args + from the diffusion sampling params for negative_prompt). + + Returns: + List of ExpandedPrompt. Empty if no expansion needed. + """ + if not isinstance(prompt, dict): + return [] + + modalities = prompt.get("modalities", []) + if "image" not in modalities and "img2img" not in modalities: + return [] + + neg_prompt = _get_negative_prompt(prompt, sampling_params) + + if "image" in modalities: + neg_prompt_dict = { + "prompt": neg_prompt, + "modalities": prompt.get("modalities", []), + } + return [ + ExpandedPrompt( + prompt=neg_prompt_dict, + role="cfg_text", + request_id_suffix=CFG_TEXT_SUFFIX, + ), + ] + + # img2img: more complex (3 distinct caches). Reserve for future. + return [] + + +def collect_cfg_kv_caches( + request_id: str, + cfg_request_ids: dict[str, str], + kv_transfer_manager: Any, + target_device: Any | None = None, +) -> dict[str, Any]: + """Collect KV caches for all CFG roles from the KV transfer manager. + + Called by the diffusion model runner after receiving the primary KV cache. + Uses the kv_transfer_manager to fetch companion KV caches by their + request IDs. + + Args: + request_id: The original (parent) request ID. + cfg_request_ids: Mapping of role -> companion request ID, + e.g. {"cfg_text": "req_0__cfg_text"}. + kv_transfer_manager: The OmniKVTransferManager instance. + target_device: Device to move tensors to. + + Returns: + Dict with keys like "cfg_text_past_key_values", + "cfg_text_kv_metadata", etc. + """ + result: dict[str, Any] = {} + + for role, companion_rid in cfg_request_ids.items(): + try: + data, size = kv_transfer_manager.receive_kv_cache_for_request(companion_rid, target_device) + if data and "layer_blocks" in data: + layer_blocks = data["layer_blocks"] + kv_obj = SimpleNamespace(**layer_blocks) + result[f"{role}_past_key_values"] = kv_obj + if "metadata" in data: + result[f"{role}_kv_metadata"] = data["metadata"] + logger.info( + "Collected CFG KV cache for role=%s, rid=%s, size=%d bytes", + role, + companion_rid, + size, + ) + else: + logger.warning( + "Failed to collect CFG KV cache for role=%s, rid=%s", + role, + companion_rid, + ) + except Exception as e: + logger.exception( + "Error collecting CFG KV cache for role=%s, rid=%s: %s", + role, + companion_rid, + e, + ) + + return result + + +def _get_negative_prompt( + prompt: dict[str, Any], + sampling_params: Any, +) -> str: + """Resolve the negative prompt for CFG from prompt or sampling params. + + An empty string is treated the same as absent (falls through to + the Bagel default token pair), because an empty negative prompt is + not meaningful for CFG guidance. + """ + neg = prompt.get("negative_prompt") + if neg: + return neg + + if hasattr(sampling_params, "extra_args") and sampling_params.extra_args: + neg = sampling_params.extra_args.get("negative_prompt") + if neg: + return neg + + return "<|im_start|><|im_end|>" diff --git a/vllm_omni/model_executor/stage_input_processors/mimo_audio.py b/vllm_omni/model_executor/stage_input_processors/mimo_audio.py index cb0657675..2b347ab2a 100644 --- a/vllm_omni/model_executor/stage_input_processors/mimo_audio.py +++ b/vllm_omni/model_executor/stage_input_processors/mimo_audio.py @@ -51,7 +51,10 @@ def _make_finished_sentinel() -> dict[str, Any]: def llm2code2wav_async_chunk( - transfer_manager: Any, pooling_output: dict[str, Any], request: Any + transfer_manager: Any, + pooling_output: dict[str, Any], + request: Any, + is_finished: bool = False, ) -> dict[str, Any] | None: """ Async chunk version: convert stage-0 pooling_output to code2wav payload (pooling / connector accumulation). @@ -60,19 +63,19 @@ def llm2code2wav_async_chunk( returns payload only when chunk_size is full or request is finished; returns None when waiting. """ if "code_predictor_codes" not in pooling_output: - if request.is_finished(): + if is_finished: return _make_finished_sentinel() return None code_predictor_codes = pooling_output["code_predictor_codes"] if code_predictor_codes is None: - if request.is_finished(): + if is_finished: return _make_finished_sentinel() return None if isinstance(code_predictor_codes, torch.Tensor): if code_predictor_codes.numel() == 0: - if request.is_finished(): + if is_finished: return _make_finished_sentinel() return None elif hasattr(code_predictor_codes, "__len__"): @@ -81,14 +84,14 @@ def llm2code2wav_async_chunk( if isinstance(code_predictor_codes, torch.Tensor): if not code_predictor_codes.any(): - if request.is_finished(): + if is_finished: return _make_finished_sentinel() return None code_tensor = code_predictor_codes.to(torch.long) else: code_tensor = torch.tensor(code_predictor_codes, dtype=torch.long) if not code_tensor.any(): - if request.is_finished(): + if is_finished: return _make_finished_sentinel() return None @@ -101,7 +104,7 @@ def llm2code2wav_async_chunk( code_final = prepend_and_flatten_colmajor(code_tensor, pad_vec) code_list = code_final.tolist() if sum(code_list) == 0: - if request.is_finished(): + if is_finished: return _make_finished_sentinel() return None @@ -113,7 +116,7 @@ def llm2code2wav_async_chunk( transfer_manager.code_prompt_token_ids[request_id].append(code_list) length = len(transfer_manager.code_prompt_token_ids[request_id]) chunk_length = length % chunk_size - if chunk_length != 0 and not request.is_finished(): + if chunk_length != 0 and not is_finished: return None context_length = chunk_length if chunk_length != 0 else chunk_size @@ -123,7 +126,7 @@ def llm2code2wav_async_chunk( "code_predictor_codes": ( torch.tensor(transfer_manager.code_prompt_token_ids[request_id][-end_index:]).reshape(-1).tolist() ), - "finished": torch.tensor(request.is_finished(), dtype=torch.bool), + "finished": torch.tensor(is_finished, dtype=torch.bool), } return info 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 b0b4ea09d..f1c18e543 100644 --- a/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py +++ b/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py @@ -586,7 +586,7 @@ class NPUARModelRunner(OmniNPUModelRunner): if not req_state: return req_id - add_info = getattr(req_state, "additional_information_cpu", {}) or {} + add_info = self.model_intermediate_buffer.get(req_id, {}) global_id = add_info.get("global_request_id") if global_id: if isinstance(global_id, list) and global_id: diff --git a/vllm_omni/worker/gpu_ar_model_runner.py b/vllm_omni/worker/gpu_ar_model_runner.py index 912c8703c..40a21a59c 100644 --- a/vllm_omni/worker/gpu_ar_model_runner.py +++ b/vllm_omni/worker/gpu_ar_model_runner.py @@ -421,6 +421,14 @@ class GPUARModelRunner(OmniGPUModelRunner): if grammar_output is not None: apply_grammar_bitmask(scheduler_output, grammar_output, self.input_batch, logits) + # Correct padding values of prompt_token_ids to match the logits vocabulary size + if logits is not None and not self.input_batch.sampling_metadata.no_penalties: + smd = self.input_batch.sampling_metadata + if smd.prompt_token_ids is not None: + logits_vocab = logits.shape[-1] + if self.input_batch.vocab_size > logits_vocab: + smd.prompt_token_ids = smd.prompt_token_ids.clamp(max=logits_vocab) + with record_function_or_nullcontext("gpu_model_runner: sample"): sampler_output = self._sample(logits, spec_decode_metadata) @@ -601,7 +609,7 @@ class GPUARModelRunner(OmniGPUModelRunner): if not req_state: return req_id - add_info = getattr(req_state, "additional_information_cpu", {}) or {} + add_info = self.model_intermediate_buffer.get(req_id, {}) global_id = add_info.get("global_request_id") if global_id: if isinstance(global_id, list) and global_id: diff --git a/vllm_omni/worker/gpu_model_runner.py b/vllm_omni/worker/gpu_model_runner.py index b16e21404..79181a0ed 100644 --- a/vllm_omni/worker/gpu_model_runner.py +++ b/vllm_omni/worker/gpu_model_runner.py @@ -37,7 +37,7 @@ logger = init_logger(__name__) class OmniGPUModelRunner(GPUModelRunner): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) - self._omni_per_req_additional_information: dict[str, dict] | None = None + self.model_intermediate_buffer: dict[str, dict[str, Any]] = {} self._omni_num_scheduled_tokens_np: np.ndarray | None = None self._omni_last_model_output: object | None = None @@ -234,6 +234,7 @@ class OmniGPUModelRunner(GPUModelRunner): # Remove finished requests from the cached states. for req_id in scheduler_output.finished_req_ids: self.requests.pop(req_id, None) + self.model_intermediate_buffer.pop(req_id, None) self.num_prompt_logprobs.pop(req_id, None) # Remove the finished requests from the persistent batch. # NOTE(woosuk): There could be an edge case where finished_req_ids and @@ -312,14 +313,52 @@ class OmniGPUModelRunner(GPUModelRunner): ) self.requests[req_id] = req_state - # Decode and store omni-specific payloads on CPU. - pe_cpu = self._resolve_prompt_embeds_cpu(new_req_data.prompt_embeds) - if pe_cpu is not None: - setattr(self.requests[req_id], "prompt_embeds_cpu", pe_cpu) + # If prompt embeddings are provided, decode and attach to inter_data + try: + if getattr(new_req_data, "prompt_embeds", None) is not None: + payload = new_req_data.prompt_embeds + dtype = getattr(np, payload.dtype) + arr = np.frombuffer(payload.data, dtype=dtype) + arr = arr.reshape(payload.shape) + pe_cpu = torch.from_numpy(arr) + setattr(self.requests[req_id], "prompt_embeds_cpu", pe_cpu) + try: + new_req_data.prompt_embeds = pe_cpu # type: ignore[assignment] + except Exception: + pass + except Exception as e: + logger.error(f"Error decoding prompt embeds: {e}") + # Decode additional_information payloads (dictionary) + try: + if getattr(new_req_data, "additional_information", None) is not None: + logger.warning_once( + "additional_information on request data is deprecated, use model_intermediate_buffer" + ) + payload_info = new_req_data.additional_information + info_dict = {} + if isinstance(payload_info, dict): + info_dict = payload_info + else: + from vllm_omni.engine import AdditionalInformationPayload - info_dict = self._resolve_additional_information(getattr(new_req_data, "additional_information", None)) - if info_dict: - setattr(self.requests[req_id], "additional_information_cpu", info_dict) + if isinstance(payload_info, AdditionalInformationPayload): + for k, entry in payload_info.entries.items(): + if entry.tensor_data is not None: + dt = np.dtype(getattr(entry, "tensor_dtype", "float32")) + arr = np.frombuffer(entry.tensor_data, dtype=dt) + arr = arr.reshape(entry.tensor_shape) + info_dict[k] = torch.from_numpy(arr.copy()) + else: + info_dict[k] = entry.list_data + if info_dict: + self.model_intermediate_buffer[req_id] = info_dict + setattr( + self.requests[req_id], + "additional_information_cpu", + info_dict, + ) + except Exception as e: + logger.error(f"Error decoding additional information: {e}") if sampling_params and sampling_params.prompt_logprobs is not None: self.num_prompt_logprobs[req_id] = ( @@ -884,6 +923,8 @@ class OmniGPUModelRunner(GPUModelRunner): newly scheduled requests and store them on CPU in the request state. """ new_reqs = getattr(scheduler_output, "scheduled_new_reqs", []) + if not new_reqs: + return for nr in new_reqs: req_id = getattr(nr, "req_id", None) or getattr(nr, "request_id", None) if req_id is None or req_id not in self.requests: @@ -891,17 +932,22 @@ class OmniGPUModelRunner(GPUModelRunner): pe_cpu = self._resolve_prompt_embeds_cpu(getattr(nr, "prompt_embeds", None)) if pe_cpu is not None: setattr(self.requests[req_id], "prompt_embeds_cpu", pe_cpu) - info_dict = self._resolve_additional_information(getattr(nr, "additional_information", None)) + info_payload = getattr(nr, "additional_information", None) + if info_payload is not None: + logger.warning_once( + "additional_information on request data is deprecated, use model_intermediate_buffer" + ) + info_dict = self._resolve_additional_information(info_payload) if info_dict: + self.model_intermediate_buffer[req_id] = info_dict setattr(self.requests[req_id], "additional_information_cpu", info_dict) def _gather_runtime_additional_information(self) -> list[dict]: - """Gather per-request additional_information stored in request state in batch order.""" + """Gather per-request model_intermediate_buffer in batch order.""" per_req_runtime_info = [] for req_id in self.input_batch.req_ids: - req_state = self.requests.get(req_id) - info = getattr(req_state, "additional_information_cpu", None) if req_state is not None else None - if info and isinstance(info, dict): + info = self.model_intermediate_buffer.get(req_id, {}) + if info: per_req_runtime_info.append(info) if "thinker_reply_part_per_request" in info: q = info["thinker_reply_part_per_request"] @@ -921,12 +967,13 @@ class OmniGPUModelRunner(GPUModelRunner): return req_token_spans def _build_model_kwargs_extra(self) -> dict: - """Build extra keyword arguments passed to the model for this step, including: - - runtime_additional_information: per-request additional information stored in request state - """ + """Build extra keyword arguments passed to the model for this step.""" model_kwargs_extra: dict[str, object] = {} try: - model_kwargs_extra["runtime_additional_information"] = self._gather_runtime_additional_information() + buffer_map = self._gather_runtime_additional_information() + model_kwargs_extra["model_intermediate_buffer"] = buffer_map + # Backward compatible: also emit old name + model_kwargs_extra["runtime_additional_information"] = buffer_map except Exception as e: logger.error(f"[OMNI DEBUG] Error building model_kwargs_extra: {e}") import traceback @@ -941,23 +988,20 @@ class OmniGPUModelRunner(GPUModelRunner): num_scheduled_tokens_np: np.ndarray, scheduler_output: "SchedulerOutput", ) -> None: - """Process model-provided per-request additional_information updates and merge into request state.""" + """Process model-provided per-request updates and merge into model_intermediate_buffer.""" try: # execute the custom postprocess function # TODO(Peiqi): do we have a more elegant way to do this? if hasattr(self.model, "has_postprocess") and self.model.has_postprocess: for req_index, req_id in enumerate(self.input_batch.req_ids): - req_state = self.requests.get(req_id) - req_infos = ( - getattr(req_state, "additional_information_cpu", None) if req_state is not None else None - ) + req_infos = self.model_intermediate_buffer.get(req_id, {}) start_offset = int(self.query_start_loc.cpu[req_index]) sched_tokens = int(num_scheduled_tokens_np[req_index]) s, e = start_offset, start_offset + sched_tokens # only consider to store data into update dict. hidden_states_slice = hidden_states[s:e] update_dict = self.model.postprocess(hidden_states_slice, **req_infos) - self._merge_additional_information_update(req_id, update_dict) + self._update_intermediate_buffer(req_id, update_dict) except Exception as e: logger.error( f"Error merging for requests:{self.input_batch.req_ids} " @@ -991,27 +1035,23 @@ class OmniGPUModelRunner(GPUModelRunner): start_offset = int(self.query_start_loc.cpu[req_index]) self.inputs_embeds[start_offset : start_offset + overlay_len].copy_(src) - def _update_request_information(self, request_id: str, payload_info: dict) -> None: - """Update per-request additional_information stored in request state.""" - req_state = self.requests.get(request_id) - if req_state is None: - return - - info_dict = getattr(req_state, "additional_information_cpu", None) - if isinstance(payload_info, dict) and info_dict is not None: - info_dict.update(payload_info) - def _update_additional_information(self, scheduler_output: "SchedulerOutput") -> None: for new_req in scheduler_output.scheduled_new_reqs: payload_info = getattr(new_req, "additional_information", None) if isinstance(payload_info, dict): - self._update_request_information(new_req.req_id, payload_info) + logger.warning_once( + "additional_information on request data is deprecated, use model_intermediate_buffer" + ) + self._update_intermediate_buffer(new_req.req_id, payload_info) if hasattr(scheduler_output.scheduled_cached_reqs, "additional_information"): + logger.warning_once( + "additional_information on scheduled_cached_reqs is deprecated, use model_intermediate_buffer" + ) cached_infos = getattr(scheduler_output.scheduled_cached_reqs, "additional_information", {}) if isinstance(cached_infos, dict): for req_id, req_infos in cached_infos.items(): - self._update_request_information(req_id, req_infos) + self._update_intermediate_buffer(req_id, req_infos) def _maybe_attach_mimo_audio_req_infos( self, @@ -1158,10 +1198,10 @@ class OmniGPUModelRunner(GPUModelRunner): if self.vllm_config.model_config.async_chunk: self._update_additional_information(scheduler_output) for req_index, req_id in enumerate(self.input_batch.req_ids): - req_state = self.requests.get(req_id) - req_infos = getattr(req_state, "additional_information_cpu", None) if req_state is not None else None + req_infos = self.model_intermediate_buffer.get(req_id, {}) # mimo-audio check + req_state = self.requests.get(req_id) req_infos = self._maybe_attach_mimo_audio_req_infos(req_state, req_infos, req_id) start_offset = int(self.query_start_loc.cpu[req_index]) @@ -1270,23 +1310,25 @@ class OmniGPUModelRunner(GPUModelRunner): self._omni_last_model_output = model_output return model_output - def _merge_additional_information_update(self, req_id: str, upd: dict | None) -> None: - if not isinstance(upd, dict): + def _update_intermediate_buffer(self, req_id: str, upd: dict) -> None: + if not isinstance(upd, dict) or not upd: return req_state = self.requests.get(req_id) if req_state is None: return - existing = getattr(req_state, "additional_information_cpu", {}) - if not isinstance(existing, dict): - existing = {} - merged = dict(existing) + existing = self.model_intermediate_buffer.setdefault(req_id, {}) for k, v in upd.items(): if isinstance(v, torch.Tensor): - merged[k] = v.detach().to("cpu").contiguous() + existing[k] = v.detach().to("cpu").contiguous() elif isinstance(v, list): - merged[k] = [ + existing[k] = [ (item.detach().to("cpu").contiguous() if isinstance(item, torch.Tensor) else item) for item in v ] else: - merged[k] = v - setattr(req_state, "additional_information_cpu", merged) + existing[k] = v + # Backward compatible: mirror to old setattr location + setattr(req_state, "additional_information_cpu", existing) + + def _merge_additional_information_update(self, req_id, upd): + logger.warning_once("_merge_additional_information_update is deprecated, use _update_intermediate_buffer") + return self._update_intermediate_buffer(req_id, upd)