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>

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