Merge remote-tracking branch 'origin/main' into dev/vllm-align

Made-with: Cursor

# Conflicts:
#	vllm_omni/entrypoints/openai/serving_chat.py
#	vllm_omni/worker/gpu_model_runner.py

Signed-off-by: tzhouam <tzhouam@connect.ust.hk>

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.

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:

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))
 

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

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

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

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
 
 

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

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"

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

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()

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 = (

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

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
 

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

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:

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

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",

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)
 

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

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",

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)

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(

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:

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

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"

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)
 

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]

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])

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] = []

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

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

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|>"

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
 

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:

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:

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)