[Bagel]: Support multistage img2img (#1669)

Signed-off-by: princepride <wangzhipeng628@gmail.com>
Signed-off-by: 汪志鹏 <wangzhipeng628@gmail.com>

examples/offline_inference/bagel/end2end.py

@@ -1,8 +1,7 @@
 import argparse
 import os
-from typing import cast
 
-from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniPromptType
+from vllm_omni.inputs.data import OmniPromptType
 
 
 def parse_args():
@@ -58,6 +57,7 @@ def parse_args():
         choices=[1, 2, 3],
         help="CFG parallel size: 1=batched (single GPU), 2=parallel with 2 branches (text CFG only), 3=parallel (3 GPUs).",
     )
+    parser.add_argument("--seed", type=int, default=None, help="Random seed for generation.")
 
     args = parser.parse_args()
     return args
@@ -88,108 +88,80 @@ def main():
 
     from PIL import Image
 
-    if args.modality == "img2img":
-        from vllm_omni.entrypoints.omni_diffusion import OmniDiffusion
+    from vllm_omni.entrypoints.omni import Omni
 
-        print(f"[Info] Running in {args.modality} mode (Stage 1 only, cfg_parallel_size={args.cfg_parallel_size})")
+    omni_kwargs = {}
+    if args.stage_configs_path:
+        omni_kwargs["stage_configs_path"] = args.stage_configs_path
 
-        client = OmniDiffusion(
-            model=model_name,
-            parallel_config={"cfg_parallel_size": args.cfg_parallel_size},
-        )
+    omni_kwargs.update(
+        {
+            "log_stats": args.log_stats,
+            "init_sleep_seconds": args.init_sleep_seconds,
+            "batch_timeout": args.batch_timeout,
+            "init_timeout": args.init_timeout,
+            "shm_threshold_bytes": args.shm_threshold_bytes,
+            "worker_backend": args.worker_backend,
+            "ray_address": args.ray_address,
+        }
+    )
 
-        if args.image_path:
-            if os.path.exists(args.image_path):
-                loaded_image = Image.open(args.image_path).convert("RGB")
-                prompts = [
-                    {
-                        "prompt": cast(str, p),
-                        "multi_modal_data": {"image": loaded_image},
-                    }
-                    for p in prompts
-                ]
-            else:
-                print(f"[Warning] Image path {args.image_path} does not exist.")
+    omni = Omni(model=model_name, **omni_kwargs)
 
-        result = client.generate(
-            prompts,
-            OmniDiffusionSamplingParams(
-                seed=52,
-                need_kv_receive=False,
-                num_inference_steps=args.steps,
-                extra_args={
-                    "cfg_text_scale": args.cfg_text_scale,
-                    "cfg_img_scale": args.cfg_img_scale,
-                },
-            ),
-        )
-
-        # Ensure result is a list for iteration
-        if not isinstance(result, list):
-            omni_outputs = [result]
-        else:
-            omni_outputs = result
-
-    else:
-        from vllm_omni.entrypoints.omni import Omni
-
-        omni_kwargs = {}
-        if args.stage_configs_path:
-            omni_kwargs["stage_configs_path"] = args.stage_configs_path
-
-        omni_kwargs.update(
-            {
-                "log_stats": args.log_stats,
-                "init_sleep_seconds": args.init_sleep_seconds,
-                "batch_timeout": args.batch_timeout,
-                "init_timeout": args.init_timeout,
-                "shm_threshold_bytes": args.shm_threshold_bytes,
-                "worker_backend": args.worker_backend,
-                "ray_address": args.ray_address,
+    formatted_prompts = []
+    for p in prompts:
+        if args.modality == "img2img":
+            if not args.image_path or not os.path.exists(args.image_path):
+                raise ValueError(f"img2img requires --image-path pointing to an existing file, got: {args.image_path}")
+            loaded_image = Image.open(args.image_path).convert("RGB")
+            final_prompt_text = f"<|fim_middle|><|im_start|>{p}<|im_end|>"
+            prompt_dict = {
+                "prompt": final_prompt_text,
+                "multi_modal_data": {"img2img": loaded_image},
+                "modalities": ["img2img"],
             }
-        )
-
-        omni = Omni(model=model_name, **omni_kwargs)
-
-        formatted_prompts = []
-        for p in args.prompts:
-            if args.modality == "img2text":
-                if args.image_path:
-                    loaded_image = Image.open(args.image_path).convert("RGB")
-                    final_prompt_text = f"<|im_start|>user\n<|image_pad|>\n{p}<|im_end|>\n<|im_start|>assistant\n"
-                    prompt_dict = {
-                        "prompt": final_prompt_text,
-                        "multi_modal_data": {"image": loaded_image},
-                        "modalities": ["text"],
-                    }
-                    formatted_prompts.append(prompt_dict)
-            elif args.modality == "text2text":
-                final_prompt_text = f"<|im_start|>user\n{p}<|im_end|>\n<|im_start|>assistant\n"
-                prompt_dict = {"prompt": final_prompt_text, "modalities": ["text"]}
-                formatted_prompts.append(prompt_dict)
-            else:
-                # 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
-        if args.modality == "text2img":
-            params_list[0].max_tokens = 1  # type: ignore # The first stage is a SamplingParam (vllm)
-            if len(params_list) > 1:
-                diffusion_params = params_list[1]
-                diffusion_params.num_inference_steps = args.steps  # type: ignore
-                extra = {
-                    "cfg_text_scale": args.cfg_text_scale,
-                    "cfg_img_scale": args.cfg_img_scale,
+            if args.negative_prompt is not None:
+                prompt_dict["negative_prompt"] = args.negative_prompt
+            formatted_prompts.append(prompt_dict)
+        elif args.modality == "img2text":
+            if args.image_path:
+                loaded_image = Image.open(args.image_path).convert("RGB")
+                final_prompt_text = f"<|im_start|>user\n<|image_pad|>\n{p}<|im_end|>\n<|im_start|>assistant\n"
+                prompt_dict = {
+                    "prompt": final_prompt_text,
+                    "multi_modal_data": {"image": loaded_image},
+                    "modalities": ["text"],
                 }
-                if args.negative_prompt is not None:
-                    extra["negative_prompt"] = args.negative_prompt
-                diffusion_params.extra_args = extra  # type: ignore
+                formatted_prompts.append(prompt_dict)
+        elif args.modality == "text2text":
+            final_prompt_text = f"<|im_start|>user\n{p}<|im_end|>\n<|im_start|>assistant\n"
+            prompt_dict = {"prompt": final_prompt_text, "modalities": ["text"]}
+            formatted_prompts.append(prompt_dict)
+        else:
+            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)
 
-        omni_outputs = list(omni.generate(prompts=formatted_prompts, sampling_params_list=params_list))
+    params_list = omni.default_sampling_params_list
+    if args.modality in ("text2img", "img2img"):
+        params_list[0].max_tokens = 1  # type: ignore
+        if len(params_list) > 1:
+            diffusion_params = params_list[1]
+            diffusion_params.num_inference_steps = args.steps  # type: ignore
+            diffusion_params.cfg_parallel_size = args.cfg_parallel_size  # type: ignore
+            if args.seed is not None:
+                diffusion_params.seed = args.seed  # 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))
 
     for i, req_output in enumerate(omni_outputs):
         images = getattr(req_output, "images", None)

tests/distributed/omni_connectors/test_kv_flow.py

@@ -6,6 +6,7 @@ from vllm_omni.distributed.omni_connectors.kv_transfer_manager import (
     OmniKVCacheConfig,
     OmniKVTransferManager,
 )
+from vllm_omni.distributed.omni_connectors.utils.kv_utils import normalize_layer_kv
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.cache]
@@ -192,8 +193,7 @@ def test_normalize_layer_kv_rejects_invalid_inputs(kv_config, common_constants,
     else:
         layer_kv = (torch.randn(2, block_size, num_heads, head_dim), "not-a-tensor")
 
-    manager = OmniKVTransferManager(kv_config)
-    normalized = manager._normalize_layer_kv(layer_kv, req_id=req_id, layer_idx=0)
+    normalized = normalize_layer_kv(layer_kv, req_id=req_id, layer_idx=0)
     assert normalized is None
 
 

tests/e2e/offline_inference/stage_configs/bagel_mooncake_ci.yaml

@@ -13,7 +13,7 @@ stage_args:
       model_arch: BagelForConditionalGeneration
       worker_type: ar
       scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
-      gpu_memory_utilization: 0.35
+      gpu_memory_utilization: 0.45
       enforce_eager: true
       trust_remote_code: true
       engine_output_type: text
@@ -46,7 +46,7 @@ stage_args:
       max_batch_size: 1
     engine_args:
       model_stage: dit
-      gpu_memory_utilization: 0.55
+      gpu_memory_utilization: 0.45
       enforce_eager: true
       trust_remote_code: true
       engine_output_type: image

tests/e2e/offline_inference/stage_configs/bagel_sharedmemory_ci.yaml

@@ -13,7 +13,7 @@ stage_args:
       model_arch: BagelForConditionalGeneration
       worker_type: ar
       scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
-      gpu_memory_utilization: 0.35
+      gpu_memory_utilization: 0.45
       enforce_eager: true
       trust_remote_code: true
       engine_output_type: text
@@ -45,7 +45,7 @@ stage_args:
       max_batch_size: 1
     engine_args:
       model_stage: dit
-      gpu_memory_utilization: 0.55
+      gpu_memory_utilization: 0.45
       enforce_eager: true
       trust_remote_code: true
       engine_output_type: image

vllm_omni/diffusion/models/bagel/pipeline_bagel.py

@@ -332,6 +332,9 @@ class BagelPipeline(nn.Module):
             else:
                 gen_context["ropes"] = [seq_len]
 
+            if req.sampling_params.kv_metadata and "image_shape" in req.sampling_params.kv_metadata:
+                image_shape = tuple(req.sampling_params.kv_metadata["image_shape"])
+
             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")

vllm_omni/distributed/omni_connectors/kv_transfer_manager.py

@@ -12,6 +12,7 @@ from vllm.logger import init_logger
 
 from .factory import OmniConnectorFactory
 from .utils.config import ConnectorSpec
+from .utils.kv_utils import normalize_layer_kv
 
 logger = init_logger(__name__)
 
@@ -200,8 +201,12 @@ class OmniKVTransferManager:
                     logger.warning(f"Request {req_id} has no block IDs, skipping")
                     continue
 
+                custom_metadata = data.get("custom_metadata")
+
                 # Extract KV cache from GPU blocks -> CPU tensors
-                kv_data = self._extract_kv_cache(req_id, block_ids, seq_len, kv_caches, block_size, cache_dtype)
+                kv_data = self._extract_kv_cache(
+                    req_id, block_ids, seq_len, kv_caches, block_size, cache_dtype, custom_metadata
+                )
                 if kv_data:
                     # Resolve global request ID if available
                     transfer_req_id = request_id_resolver(req_id) if request_id_resolver else req_id
@@ -224,6 +229,7 @@ class OmniKVTransferManager:
         kv_caches: list[LayerKV],
         block_size: int,
         cache_dtype: str,
+        custom_metadata: dict[str, Any] | None = None,
     ) -> KVCacheTransferData | None:
         """Extract KV cache from GPU blocks for a single request.
 
@@ -234,6 +240,7 @@ class OmniKVTransferManager:
             kv_caches: List of KV cache (tensor or tuple) per layer
             block_size: Size of each cache block
             cache_dtype: Data type of the cache
+            custom_metadata: Optional custom metadata to include
 
         Note: If key/value block counts differ, extraction uses only the overlapping
         block range. Extra key/value blocks are ignored, so returned KV may be partial.
@@ -246,7 +253,7 @@ class OmniKVTransferManager:
         value_cache: list[torch.Tensor | None] = [None] * num_layers
 
         for layer_idx, layer_kv in enumerate(kv_caches):
-            kv_pair = self._normalize_layer_kv(layer_kv, req_id=req_id, layer_idx=layer_idx)
+            kv_pair = normalize_layer_kv(layer_kv, req_id=req_id, layer_idx=layer_idx)
             if kv_pair is None:
                 continue
             key_blocks, value_blocks = kv_pair
@@ -289,57 +296,10 @@ class OmniKVTransferManager:
                 "num_layers": num_layers,
                 "dtype": str(cache_dtype),
                 "seq_len": seq_len,
+                **(custom_metadata or {}),
             },
         )
 
-    def _normalize_layer_kv(
-        self,
-        layer_kv: LayerKV,
-        req_id: str,
-        layer_idx: int,
-    ) -> tuple[torch.Tensor, torch.Tensor] | None:
-        """Normalize one layer KV cache to a `(key_blocks, value_blocks)` tuple.
-
-        Args:
-            layer_kv: The raw KV cache (tensor or tuple) for the layer
-            req_id: Request ID for logging
-            layer_idx: Layer index for logging
-
-        Returns:
-            Tuple of (key_blocks, value_blocks) if valid, None otherwise
-        """
-        if isinstance(layer_kv, torch.Tensor):
-            if layer_kv.ndim < 3 or layer_kv.shape[0] != 2:
-                logger.warning(
-                    f"Layer {layer_idx} for request {req_id} has invalid stacked KV shape: "
-                    f"expected [2, blocks, block_size, ...], got {tuple(layer_kv.shape)}"
-                )
-                return None
-            key_blocks = layer_kv[0]
-            value_blocks = layer_kv[1]
-        elif isinstance(layer_kv, tuple):
-            if len(layer_kv) != 2:
-                logger.warning(
-                    f"Layer {layer_idx} for request {req_id} has KV pair length {len(layer_kv)} (expected 2)"
-                )
-                return None
-            key_blocks, value_blocks = layer_kv
-            if not isinstance(key_blocks, torch.Tensor) or not isinstance(value_blocks, torch.Tensor):
-                logger.warning(f"Layer {layer_idx} for request {req_id} has non-tensor KV pair entries")
-                return None
-        else:
-            logger.warning(f"Layer {layer_idx} for request {req_id} has unsupported KV type {type(layer_kv).__name__}")
-            return None
-        # ensure key/value blocks are at least 2D for block indexing
-        if key_blocks.ndim < 2 or value_blocks.ndim < 2:
-            logger.warning(
-                f"Layer {layer_idx} for request {req_id} has invalid KV block shape: "
-                f"got key={tuple(key_blocks.shape)} value={tuple(value_blocks.shape)}"
-            )
-            return None
-
-        return key_blocks, value_blocks
-
     def _transfer_kv_cache(self, kv_data: KVCacheTransferData, transfer_req_id: str) -> None:
         """Transfer KV cache data to downstream stage via OmniConnector.
 
@@ -496,6 +456,8 @@ class OmniKVTransferManager:
 
         if "metadata" in data:
             req.kv_metadata = data["metadata"]
+            if hasattr(req, "sampling_params") and req.sampling_params is not None:
+                req.sampling_params.kv_metadata = data["metadata"]
 
     # Legacy compatibility method
     def receive_kv_cache(self, req: Any, target_device: torch.device | None = None) -> bool:

vllm_omni/distributed/omni_connectors/utils/kv_utils.py

@@ -0,0 +1,85 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Utility helpers for KV cache manipulation."""
+
+import torch
+from vllm.logger import init_logger
+
+logger = init_logger(__name__)
+
+LayerKV = torch.Tensor | tuple[torch.Tensor, torch.Tensor]
+
+
+def normalize_layer_kv(
+    layer_kv: LayerKV,
+    *,
+    req_id: str = "",
+    layer_idx: int = -1,
+) -> tuple[torch.Tensor, torch.Tensor] | None:
+    """Normalize one layer KV cache to a ``(key_blocks, value_blocks)`` tuple.
+
+    In vLLM, different attention backends return paged-attention KV blocks
+    with different layouts.  For example:
+
+    * **FlashAttention** (``vllm/v1/attention/backends/flash_attn.py``)
+      returns shape ``(2, num_blocks, block_size, num_kv_heads, head_size)``
+      – the key/value dimension is at **dim 0**.
+    * **FlashInfer** (``vllm/v1/attention/backends/flashinfer.py``)
+      returns shape ``(num_blocks, 2, block_size, num_kv_heads, head_size)``
+      – the key/value dimension is at **dim 1**.
+
+    This utility handles both layouts (and the tuple case) so that
+    downstream code can work with any backend.
+
+    Supported layouts:
+
+    * **Stacked tensor** ``[2, num_blocks, block_size, n_heads, head_dim]`` –
+      dim-0 selects key / value.
+    * **Stacked tensor** ``[num_blocks, 2, block_size, n_heads, head_dim]`` –
+      dim-1 selects key / value.
+    * **Tuple** ``(key_tensor, value_tensor)`` – returned as-is after
+      validation.
+
+    Args:
+        layer_kv: The raw KV cache (tensor or tuple) for the layer.
+        req_id: Request ID used only for diagnostic log messages.
+        layer_idx: Layer index used only for diagnostic log messages.
+
+    Returns:
+        ``(key_blocks, value_blocks)`` if *layer_kv* is valid, ``None``
+        otherwise.
+    """
+    if isinstance(layer_kv, torch.Tensor):
+        if layer_kv.ndim >= 3 and layer_kv.shape[0] == 2:
+            key_blocks = layer_kv[0]
+            value_blocks = layer_kv[1]
+        elif layer_kv.ndim >= 3 and layer_kv.shape[1] == 2:
+            key_blocks = layer_kv[:, 0]
+            value_blocks = layer_kv[:, 1]
+        else:
+            logger.warning(
+                f"Layer {layer_idx} for request {req_id} has invalid stacked KV shape: "
+                f"expected [2, ...] or [..., 2, ...] at dim 0/1, got {tuple(layer_kv.shape)}"
+            )
+            return None
+    elif isinstance(layer_kv, tuple):
+        if len(layer_kv) != 2:
+            logger.warning(f"Layer {layer_idx} for request {req_id} has KV pair length {len(layer_kv)} (expected 2)")
+            return None
+        key_blocks, value_blocks = layer_kv
+        if not isinstance(key_blocks, torch.Tensor) or not isinstance(value_blocks, torch.Tensor):
+            logger.warning(f"Layer {layer_idx} for request {req_id} has non-tensor KV pair entries")
+            return None
+    else:
+        logger.warning(f"Layer {layer_idx} for request {req_id} has unsupported KV type {type(layer_kv).__name__}")
+        return None
+
+    # ensure key/value blocks are at least 2D for block indexing
+    if key_blocks.ndim < 2 or value_blocks.ndim < 2:
+        logger.warning(
+            f"Layer {layer_idx} for request {req_id} has invalid KV block shape: "
+            f"got key={tuple(key_blocks.shape)} value={tuple(value_blocks.shape)}"
+        )
+        return None
+
+    return key_blocks, value_blocks

vllm_omni/model_executor/models/__init__.py

@@ -1,4 +1,8 @@
+from .bagel.bagel import OmniBagelForConditionalGeneration
 from .qwen3_omni import Qwen3OmniMoeForConditionalGeneration
 from .registry import OmniModelRegistry  # noqa: F401
 
-__all__ = ["Qwen3OmniMoeForConditionalGeneration"]
+__all__ = [
+    "Qwen3OmniMoeForConditionalGeneration",
+    "OmniBagelForConditionalGeneration",
+]

vllm_omni/model_executor/models/bagel/__init__.py

@@ -0,0 +1,3 @@
+from .bagel import OmniBagelForConditionalGeneration
+
+__all__ = ["OmniBagelForConditionalGeneration"]

vllm_omni/model_executor/models/registry.py

@@ -119,6 +119,11 @@ _OMNI_MODELS = {
         "glm_image_ar",
         "GlmImageForConditionalGeneration",
     ),
+    "OmniBagelForConditionalGeneration": (
+        "bagel",
+        "bagel",
+        "OmniBagelForConditionalGeneration",
+    ),
     "HunyuanImage3ForCausalMM": (
         "hunyuan_image3",
         "hunyuan_image3",

vllm_omni/model_executor/stage_configs/bagel.yaml

@@ -6,15 +6,14 @@ stage_args:
     prompt_expand_func: vllm_omni.model_executor.stage_input_processors.bagel.expand_cfg_prompts
     runtime:
       devices: "0"
-      # 2 = 1 user prompt + 1 CFG companion (text-unconditional).
-      # For multi-prompt batches this should scale as batch_size × 2.
-      max_batch_size: 2
+      # 3 = 1 user prompt + 2 CFG companions (text-unconditional + image-unconditional).
+      max_batch_size: 3
     engine_args:
       model_stage: thinker
-      model_arch: BagelForConditionalGeneration
+      model_arch: OmniBagelForConditionalGeneration
       worker_type: ar
       scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
-      gpu_memory_utilization: 0.35
+      gpu_memory_utilization: 0.45
       enforce_eager: true
       trust_remote_code: true
       engine_output_type: text
@@ -46,7 +45,7 @@ stage_args:
       max_batch_size: 1
     engine_args:
       model_stage: dit
-      gpu_memory_utilization: 0.55
+      gpu_memory_utilization: 0.45
       enforce_eager: true
       trust_remote_code: true
       engine_output_type: image

vllm_omni/model_executor/stage_configs/bagel_multiconnector.yaml

@@ -11,7 +11,7 @@ stage_args:
       model_arch: BagelForConditionalGeneration
       worker_type: ar
       scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
-      gpu_memory_utilization: 0.35
+      gpu_memory_utilization: 0.45
       enforce_eager: true
       trust_remote_code: true
       engine_output_type: text
@@ -45,7 +45,7 @@ stage_args:
       max_batch_size: 1
     engine_args:
       model_stage: dit
-      gpu_memory_utilization: 0.55
+      gpu_memory_utilization: 0.45
       enforce_eager: true
       trust_remote_code: true
       engine_output_type: image

vllm_omni/model_executor/stage_input_processors/bagel.py

@@ -20,6 +20,7 @@ from typing import Any
 logger = logging.getLogger(__name__)
 
 CFG_TEXT_SUFFIX = "__cfg_text"
+CFG_IMG_SUFFIX = "__cfg_img"
 
 
 @dataclass
@@ -73,7 +74,37 @@ def expand_cfg_prompts(
             ),
         ]
 
-    # img2img: more complex (3 distinct caches). Reserve for future.
+    if "img2img" in modalities:
+        IMG2IMG_PLACEHOLDER = "<|fim_middle|>"
+
+        cfg_text_dict: dict[str, Any] = {
+            "prompt": f"{IMG2IMG_PLACEHOLDER}{neg_prompt}",
+            "modalities": ["img2img"],
+        }
+        mm_data = prompt.get("multi_modal_data")
+        if mm_data:
+            cfg_text_dict["multi_modal_data"] = mm_data
+
+        original_text = prompt.get("prompt", "")
+        cfg_img_text = original_text.replace(IMG2IMG_PLACEHOLDER, "")
+        cfg_img_dict: dict[str, Any] = {
+            "prompt": cfg_img_text,
+            "modalities": ["img2img"],
+        }
+
+        return [
+            ExpandedPrompt(
+                prompt=cfg_text_dict,
+                role="cfg_text",
+                request_id_suffix=CFG_TEXT_SUFFIX,
+            ),
+            ExpandedPrompt(
+                prompt=cfg_img_dict,
+                role="cfg_img",
+                request_id_suffix=CFG_IMG_SUFFIX,
+            ),
+        ]
+
     return []
 
 

vllm_omni/patch.py

@@ -1,6 +1,8 @@
 import sys
+from functools import cached_property as _cached_property
 
 from aenum import extend_enum
+from vllm.config import ModelConfig as _ModelConfig
 from vllm.inputs.data import TokensPrompt as _OriginalTokensPrompt
 from vllm.model_executor.layers.rotary_embedding import (
     MRotaryEmbedding as _OriginalMRotaryEmbedding,
@@ -43,6 +45,20 @@ except ImportError:
     # GlmImageTextConfig not available, skip patching
     pass
 
+
+# Patch ModelConfig.is_mm_prefix_lm to include Bagel (bidirectional attention
+# for multimodal prefix positions, same as Gemma3/Molmo2/PaliGemma).
+_orig_is_mm_prefix_lm = _ModelConfig.__dict__["is_mm_prefix_lm"].func
+
+
+@_cached_property
+def _patched_is_mm_prefix_lm(self) -> bool:
+    return _orig_is_mm_prefix_lm(self) or getattr(self.hf_config, "model_type", None) == "bagel"
+
+
+_patched_is_mm_prefix_lm.__set_name__(_ModelConfig, "is_mm_prefix_lm")
+_ModelConfig.is_mm_prefix_lm = _patched_is_mm_prefix_lm
+
 # Extend RequestStatus enum with omni-specific statuses
 if not hasattr(RequestStatus, "WAITING_FOR_CHUNK"):
     # The value - 1 is intentionally chosen to ensure it is treated

vllm_omni/worker/gpu_ar_model_runner.py

@@ -98,9 +98,25 @@ class GPUARModelRunner(OmniGPUModelRunner):
         if self.execute_model_state is not None:
             raise RuntimeError("State error: sample_tokens() must be called after execute_model() returns None.")
 
+        if not getattr(self, "_warmup_state_cleared", False):
+            self._warmup_state_cleared = True
+            if hasattr(self.model, "_clear_warmup_state"):
+                self.model._clear_warmup_state()
+
         # [Omni] Handle KV transfer BEFORE updating states (which removes finished requests)
+        finished_reqs = getattr(scheduler_output, "finished_requests_needing_kv_transfer", {})
+        if finished_reqs and hasattr(self.model, "get_kv_transfer_metadata"):
+            for req_id, data in finished_reqs.items():
+                try:
+                    model_meta = self.model.get_kv_transfer_metadata(req_id)
+                    if model_meta:
+                        existing = data.get("custom_metadata") or {}
+                        existing.update(model_meta)
+                        data["custom_metadata"] = existing
+                except Exception as e:
+                    logger.warning(f"Failed to get custom metadata from model for {req_id}: {e}")
         self.kv_extracted_req_ids = self.kv_transfer_manager.handle_finished_requests_kv_transfer(
-            finished_reqs=getattr(scheduler_output, "finished_requests_needing_kv_transfer", {}),
+            finished_reqs=finished_reqs,
             kv_caches=self.kv_caches,
             block_size=self.cache_config.block_size,
             cache_dtype=str(self.cache_config.cache_dtype),
@@ -290,6 +306,10 @@ class GPUARModelRunner(OmniGPUModelRunner):
                 sampler=self.sampler,
             )
 
+            # [Omni] Map pending ropes metadata to req_ids.
+            if hasattr(self.model, "flush_pending_metadata"):
+                self.model.flush_pending_metadata(list(req_ids))
+
         with record_function_or_nullcontext("gpu_model_runner: postprocess"):
             if self.use_aux_hidden_state_outputs:
                 # True when EAGLE 3 is used.