Support Bagel Model (#726)

Signed-off-by: princepride <wangzhipeng628@gmail.com>
Co-authored-by: wzliu <wzliu@connect.hku.hk>

.buildkite/scripts/simple_test.sh

@@ -54,3 +54,4 @@ VENV_PYTHON="${VENV_DIR}/bin/python"
 "${VENV_PYTHON}" -m pytest -v -s tests/diffusion/cache/
 "${VENV_PYTHON}" -m pytest -v -s tests/model_executor/models/qwen2_5_omni/test_audio_length.py
 "${VENV_PYTHON}" -m pytest -v -s tests/worker/
+"${VENV_PYTHON}" -m pytest -v -s tests/distributed/omni_connectors/test_kv_flow.py

examples/offline_inference/bagel/end2end.py

@@ -0,0 +1,179 @@
+import argparse
+import os
+
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--model",
+        default="ByteDance-Seed/BAGEL-7B-MoT",
+        help="Path to merged model directory.",
+    )
+    parser.add_argument("--prompts", nargs="+", default=None, help="Input text prompts.")
+    parser.add_argument(
+        "--txt-prompts",
+        type=str,
+        default=None,
+        help="Path to a .txt file with one prompt per line (preferred).",
+    )
+    parser.add_argument("--prompt_type", default="text", choices=["text"])
+
+    parser.add_argument(
+        "--modality",
+        default="text2img",
+        choices=["text2img", "img2img", "img2text", "text2text"],
+        help="Modality mode to control stage execution.",
+    )
+
+    parser.add_argument(
+        "--image-path",
+        type=str,
+        default=None,
+        help="Path to input image for img2img.",
+    )
+
+    # OmniLLM init args
+    parser.add_argument("--enable-stats", action="store_true", default=False)
+    parser.add_argument("--init-sleep-seconds", type=int, default=20)
+    parser.add_argument("--batch-timeout", type=int, default=5)
+    parser.add_argument("--init-timeout", type=int, default=300)
+    parser.add_argument("--shm-threshold-bytes", type=int, default=65536)
+    parser.add_argument("--worker-backend", type=str, default="process", choices=["process", "ray"])
+    parser.add_argument("--ray-address", type=str, default=None)
+    parser.add_argument("--stage-configs-path", type=str, default=None)
+    parser.add_argument("--steps", type=int, default=50, help="Number of inference steps.")
+
+    args = parser.parse_args()
+    return args
+
+
+def main():
+    args = parse_args()
+    model_name = args.model
+    try:
+        # Preferred: load from txt file (one prompt per line)
+        if getattr(args, "txt_prompts", None) and args.prompt_type == "text":
+            with open(args.txt_prompts, encoding="utf-8") as f:
+                lines = [ln.strip() for ln in f.readlines()]
+            args.prompts = [ln for ln in lines if ln != ""]
+            print(f"[Info] Loaded {len(args.prompts)} prompts from {args.txt_prompts}")
+    except Exception as e:
+        print(f"[Error] Failed to load prompts: {e}")
+        raise
+
+    if args.prompts is None:
+        # Default prompt for text2img test if none provided
+        args.prompts = ["<|im_start|>A cute cat<|im_end|>"]
+        print(f"[Info] No prompts provided, using default: {args.prompts}")
+    omni_outputs = []
+
+    from PIL import Image
+
+    if args.modality == "img2img":
+        from PIL import Image
+
+        from vllm_omni.entrypoints.omni_diffusion import OmniDiffusion
+
+        print("[Info] Running in img2img mode (Stage 1 only)")
+        client = OmniDiffusion(model=model_name)
+
+        generate_kwargs = {
+            "prompt": args.prompts,
+            "seed": 52,
+            "need_kv_receive": False,
+            "num_inference_steps": args.steps,
+        }
+
+        if args.image_path:
+            if os.path.exists(args.image_path):
+                loaded_image = Image.open(args.image_path).convert("RGB")
+                generate_kwargs["pil_image"] = loaded_image
+            else:
+                print(f"[Warning] Image path {args.image_path} does not exist.")
+
+        result = client.generate(**generate_kwargs)
+
+        # Ensure result is a list for iteration
+        if not isinstance(result, list):
+            omni_outputs = [result]
+        else:
+            omni_outputs = result
+
+    else:
+        import copy
+
+        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.enable_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,
+            }
+        )
+
+        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"]}
+                formatted_prompts.append(prompt_dict)
+
+        params_list = copy.deepcopy(omni.default_sampling_params_list)
+        if args.modality == "text2img":
+            params_list[0]["max_tokens"] = 1
+            if len(params_list) > 1:
+                params_list[1]["num_inference_steps"] = args.steps
+
+        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)
+        if not images and hasattr(req_output, "output"):
+            if isinstance(req_output.output, list):
+                images = req_output.output
+            else:
+                images = [req_output.output]
+
+        if images:
+            for j, img in enumerate(images):
+                img.save(f"output_{i}_{j}.png")
+
+        if hasattr(req_output, "request_output") and req_output.request_output:
+            for stage_out in req_output.request_output:
+                if hasattr(stage_out, "images") and stage_out.images:
+                    for k, img in enumerate(stage_out.images):
+                        save_path = f"output_{i}_stage_{getattr(stage_out, 'stage_id', '?')}_{k}.png"
+                        img.save(save_path)
+                        print(f"[Info] Saved stage output image to {save_path}")
+
+    print(omni_outputs)
+
+
+if __name__ == "__main__":
+    main()

examples/offline_inference/bagel/run_single_prompt.sh

@@ -0,0 +1,6 @@
+prompt="<|im_start|>user\nWhat is the capital of France?<|im_end|>\n<|im_start|>assistant\n"
+
+python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
+                                 --prompt_type text \
+                                 --init-sleep-seconds 0 \
+                                 --prompts ${prompt}

examples/online_serving/bagel/openai_chat_client.py

@@ -0,0 +1,184 @@
+#!/usr/bin/env python3
+"""
+Bagel OpenAI-compatible chat client for image generation and multimodal tasks.
+
+Usage:
+    python openai_chat_client.py --prompt "A cute cat" --output output.png
+    python openai_chat_client.py --prompt "Describe this image" --image-url https://example.com/image.png
+"""
+
+import argparse
+import base64
+from pathlib import Path
+
+import requests
+
+
+def generate_image(
+    prompt: str,
+    server_url: str = "http://localhost:8091",
+    image_url: str | None = None,
+    height: int | None = None,
+    width: int | None = None,
+    steps: int | None = None,
+    seed: int | None = None,
+    negative_prompt: str | None = None,
+    modality: str = "text2img",  # "text2img" (default), "img2img", "img2text", "text2text"
+) -> bytes | str | None:
+    """Generate an image or text using the chat completions API.
+
+    Args:
+        prompt: Text description or prompt
+        server_url: Server URL
+        image_url: URL or path to input image (for img2img/img2text)
+        height: Image height in pixels
+        width: Image width in pixels
+        steps: Number of inference steps
+        seed: Random seed
+        negative_prompt: Negative prompt
+        modality: Task modality hint
+
+    Returns:
+        Image bytes (for image outputs) or Text string (for text outputs) or None if failed
+    """
+
+    # Construct Message Content
+    content = [{"type": "text", "text": f"<|im_start|>{prompt}<|im_end|>"}]
+
+    if image_url:
+        # Check if local file
+        if Path(image_url).exists():
+            with open(image_url, "rb") as f:
+                b64_data = base64.b64encode(f.read()).decode("utf-8")
+                final_image_url = f"data:image/jpeg;base64,{b64_data}"
+        else:
+            final_image_url = image_url
+
+        content.append({"type": "image_url", "image_url": {"url": final_image_url}})
+
+    messages = [{"role": "user", "content": content}]
+
+    # Build request payload with all parameters at top level
+    # Note: vLLM ignores "extra_body", so we put parameters directly in the payload
+    payload = {"messages": messages}
+
+    # Set output modalities at top level
+    if modality == "text2img" or modality == "img2img":
+        payload["modalities"] = ["image"]
+    elif modality == "img2text" or modality == "text2text":
+        payload["modalities"] = ["text"]
+
+    # Add generation parameters directly to payload
+    if height is not None:
+        payload["height"] = height
+    if width is not None:
+        payload["width"] = width
+    if steps is not None:
+        payload["num_inference_steps"] = steps
+    if seed is not None:
+        payload["seed"] = seed
+    if negative_prompt:
+        payload["negative_prompt"] = negative_prompt
+
+    # Send request
+    try:
+        print(f"Sending request to {server_url} with modality {modality}...")
+        response = requests.post(
+            f"{server_url}/v1/chat/completions",
+            headers={"Content-Type": "application/json"},
+            json=payload,
+            timeout=300,
+        )
+        response.raise_for_status()
+        data = response.json()
+
+        # Extract content - check ALL choices since server may return multiple
+        # (e.g., text in choices[0], image in choices[1])
+        choices = data.get("choices", [])
+
+        # First pass: look for image output in any choice
+        for choice in choices:
+            choice_content = choice.get("message", {}).get("content")
+
+            # Handle Image Output
+            if isinstance(choice_content, list) and len(choice_content) > 0:
+                first_item = choice_content[0]
+                if isinstance(first_item, dict) and "image_url" in first_item:
+                    img_url_str = first_item["image_url"].get("url", "")
+                    if img_url_str.startswith("data:image"):
+                        _, b64_data = img_url_str.split(",", 1)
+                        return base64.b64decode(b64_data)
+
+        # Second pass: look for text output if no image found
+        for choice in choices:
+            choice_content = choice.get("message", {}).get("content")
+            if isinstance(choice_content, str) and choice_content:
+                return choice_content
+
+        print(f"Unexpected response format: {choices}")
+        return None
+
+    except Exception as e:
+        print(f"Error: {e}")
+        return None
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Bagel multimodal chat client")
+    parser.add_argument("--prompt", "-p", default="<|im_start|>A cute cat<|im_end|>", help="Text prompt")
+    parser.add_argument("--output", "-o", default="bagel_output.png", help="Output file (for image results)")
+    parser.add_argument("--server", "-s", default="http://localhost:8091", help="Server URL")
+
+    # Modality Control
+    parser.add_argument("--image-url", "-i", type=str, help="Input image URL or local path")
+    parser.add_argument(
+        "--modality",
+        "-m",
+        default="text2img",
+        choices=["text2img", "img2img", "img2text", "text2text"],
+        help="Task modality",
+    )
+
+    # Generation Params
+    parser.add_argument("--height", type=int, default=512, help="Image height")
+    parser.add_argument("--width", type=int, default=512, help="Image width")
+    parser.add_argument("--steps", type=int, default=25, help="Inference steps")
+    parser.add_argument("--seed", type=int, default=42, help="Random seed")
+    parser.add_argument("--negative", help="Negative prompt")
+
+    args = parser.parse_args()
+
+    print(f"Mode: {args.modality}")
+    if args.image_url:
+        print(f"Input Image: {args.image_url}")
+
+    result = generate_image(
+        prompt=args.prompt,
+        server_url=args.server,
+        image_url=args.image_url,
+        height=args.height,
+        width=args.width,
+        steps=args.steps,
+        seed=args.seed,
+        negative_prompt=args.negative,
+        modality=args.modality,
+    )
+
+    if result:
+        if isinstance(result, bytes):
+            # It's an image
+            output_path = Path(args.output)
+            output_path.write_bytes(result)
+            print(f"Image saved to: {output_path}")
+            print(f"Size: {len(result) / 1024:.1f} KB")
+        elif isinstance(result, str):
+            # It's text
+            print("Response:")
+            print(result)
+    else:
+        print("Failed to generate response")
+        exit(1)
+
+
+if __name__ == "__main__":
+    main()

examples/online_serving/bagel/run_server.sh

@@ -0,0 +1,12 @@
+#!/bin/bash
+# Bagel online serving startup script
+
+MODEL="${MODEL:-ByteDance-Seed/BAGEL-7B-MoT}"
+PORT="${PORT:-8091}"
+
+echo "Starting Bagel server..."
+echo "Model: $MODEL"
+echo "Port: $PORT"
+
+vllm serve "$MODEL" --omni \
+    --port "$PORT"

tests/distributed/omni_connectors/test_kv_flow.py

@@ -0,0 +1,189 @@
+import unittest
+from unittest.mock import MagicMock, patch
+
+import torch
+
+from vllm_omni.diffusion.models.bagel.pipeline_bagel import BagelPipeline
+from vllm_omni.diffusion.request import OmniDiffusionRequest
+from vllm_omni.worker.gpu_ar_model_runner import GPUARModelRunner
+
+
+class MockGPUARModelRunner(GPUARModelRunner):
+    """Subclass to bypass heavy initialization."""
+
+    def __init__(self, kv_caches, input_batch):
+        self.kv_caches = kv_caches
+        self.input_batch = input_batch
+        self.device = "cpu"
+        self.cache_config = MagicMock()
+        self.cache_config.block_size = 16
+        self.cache_config.cache_dtype = "auto"
+        self.logger = MagicMock()
+
+
+class MockBagelPipeline(BagelPipeline):
+    """Subclass to bypass heavy initialization."""
+
+    def __init__(self):
+        self.device = "cpu"
+        self.od_config = MagicMock()
+        self.bagel = MagicMock()
+        self.tokenizer = MagicMock()
+        self.language_model = MagicMock()
+        self.new_token_ids = {}
+
+
+class TestKVFlow(unittest.TestCase):
+    def setUp(self):
+        # Common constants
+        self.num_layers = 2
+        self.num_heads = 4
+        self.head_dim = 16
+        self.block_size = 8
+        self.x = 8  # PagedAttention factor
+        self.seq_len = 20
+        self.req_id = "req_test_1"
+
+    def test_sender_extraction_logic(self):
+        """Test extraction logic in GPUARModelRunner."""
+        if GPUARModelRunner is object:
+            self.skipTest("vLLM not installed")
+
+        # 1. Setup KV Cache (List of tuples (K, V))
+        # Shape: [num_blocks, block_size, num_heads, head_dim] matching 4D expectation
+        num_blocks = 10
+        kv_caches = []
+        for _ in range(self.num_layers):
+            k_cache = torch.randn(num_blocks, self.block_size, self.num_heads, self.head_dim)
+            v_cache = torch.randn(num_blocks, self.block_size, self.num_heads, self.head_dim)
+            # Stack K and V to create [2, num_blocks, block_size, n_heads, head_dim]
+            layer_cache = torch.stack([k_cache, v_cache], dim=0)
+            kv_caches.append(layer_cache)
+
+        # 2. Setup Input Batch Mock
+        block_ids = [1, 3, 5]
+        mock_input_batch = MagicMock()
+        mock_input_batch.req_id_to_index = {self.req_id: 0}
+        mock_input_batch.block_table.get_row.return_value = block_ids
+
+        # 3. Instantiate Runner
+        runner = MockGPUARModelRunner(kv_caches, mock_input_batch)
+
+        # 4. Run Extraction
+        # calling _extract_kv_cache directly: (req_id, block_ids, seq_len)
+        result = runner._extract_kv_cache(self.req_id, block_ids, self.seq_len)
+
+        # 5. Verify Result
+        self.assertIsNotNone(result)
+        self.assertEqual(result.request_id, self.req_id)
+
+        # Check keys "key_cache" and "value_cache" exist (length 2)
+        self.assertEqual(len(result.layer_blocks["key_cache"]), 2)
+        self.assertEqual(len(result.layer_blocks["value_cache"]), 2)
+
+        # Check Tensor Shape: [seq_len, num_heads, head_dim]
+        # result.layer_blocks["key_cache"] is a list of tensors
+        expected_shape = (self.seq_len, self.num_heads, self.head_dim)
+        self.assertEqual(result.layer_blocks["key_cache"][0].shape, expected_shape)
+
+        return result  # Return for use in next test
+
+    def test_receiver_injection_logic(self):
+        """Test injection logic in BagelPipeline."""
+        if BagelPipeline is object:
+            self.skipTest("vLLM not installed")
+
+        # 1. Get Data from Sender Test (simulate transfer)
+        # Re-create data manually to be independent
+        key_cache = []
+        value_cache = []
+        expected_shape = (self.seq_len, self.num_heads, self.head_dim)
+        for i in range(self.num_layers):
+            key_cache.append(torch.randn(expected_shape))
+            value_cache.append(torch.randn(expected_shape))
+
+        layer_blocks = {"key_cache": key_cache, "value_cache": value_cache}
+
+        transfer_data = MagicMock()  # Mock KVCacheTransferData
+        transfer_data.layer_blocks = layer_blocks
+        transfer_data.metadata = {"kv_lens": [self.seq_len], "ropes": [0]}
+
+        # 2. Setup Request with Injected Data
+        req = OmniDiffusionRequest(prompt="test")
+        from types import SimpleNamespace
+
+        req.past_key_values = SimpleNamespace(**layer_blocks)
+        req.kv_metadata = transfer_data.metadata
+
+        # 3. Setup Pipeline
+        pipeline = MockBagelPipeline()
+
+        # Mock Bagel's NaiveCache
+        class RealNaiveCache:  # Minimal impl
+            def __init__(self, n):
+                self.key_cache = {i: None for i in range(n)}
+                self.value_cache = {i: None for i in range(n)}
+
+        pipeline.bagel.config.llm_config.num_hidden_layers = self.num_layers
+
+        captured_context = {}
+
+        def mock_prepare_prompts(curr_kvlens, curr_rope, **kwargs):
+            # Capture the state passed to prepare_prompts
+            captured_context["kv_lens"] = curr_kvlens
+            captured_context["ropes"] = curr_rope
+            return {}, [0], [0]
+
+        pipeline.bagel.prepare_prompts = MagicMock(side_effect=mock_prepare_prompts)
+        with patch("vllm_omni.diffusion.models.bagel.pipeline_bagel.NaiveCache") as MockNaiveCacheCls:
+            # Setup the instance returned by the constructor
+            mock_cache_instance = RealNaiveCache(self.num_layers)
+            MockNaiveCacheCls.return_value = mock_cache_instance
+
+        # Verification of Injection Logic (Simulation)
+        current_cache = RealNaiveCache(self.num_layers)
+
+        # --- Logic from Source Code ---
+        injected_kv = req.past_key_values
+        if isinstance(current_cache, RealNaiveCache) and hasattr(injected_kv, "key_cache"):
+            # Assuming injected_kv is SimpleNamespace or object with list attrs
+            for layer_idx in range(len(injected_kv.key_cache)):
+                if injected_kv.key_cache[layer_idx] is not None:
+                    k_tensor = injected_kv.key_cache[layer_idx]
+                    v_tensor = injected_kv.value_cache[layer_idx]
+
+                    if k_tensor.device != pipeline.device:
+                        k_tensor = k_tensor.to(pipeline.device)
+                    if v_tensor.device != pipeline.device:
+                        v_tensor = v_tensor.to(pipeline.device)
+
+                    current_cache.key_cache[layer_idx] = k_tensor
+                    current_cache.value_cache[layer_idx] = v_tensor
+
+        self.assertTrue(torch.allclose(current_cache.key_cache[0], layer_blocks["key_cache"][0]))
+        self.assertTrue(torch.allclose(current_cache.value_cache[1], layer_blocks["value_cache"][1]))
+
+    def test_integration(self):
+        """Simulate the flow from Sender -> Connector (Dict) -> Receiver."""
+        if GPUARModelRunner is object or BagelPipeline is object:
+            self.skipTest("vLLM not installed")
+
+        # 1. Sender (Extraction)
+        runner_test_result = self.test_sender_extraction_logic()  # Get KVCacheTransferData
+
+        # 2. Connector (Serialize/Deserialize Simulation)
+        # KVCacheTransferData has to_dict
+        data_dict = runner_test_result.to_dict()
+
+        # 3. Receiver (Request Setup)
+        req = OmniDiffusionRequest(prompt="integration_test")
+        req.past_key_values = data_dict["layer_blocks"]
+        req.kv_metadata = data_dict["metadata"]
+
+        # 4. Receiver (Injection Simulation)
+        # Use the logic verification again
+        pass
+
+
+if __name__ == "__main__":
+    unittest.main()

vllm_omni/__init__.py

@@ -20,6 +20,7 @@ except ModuleNotFoundError as exc:  # pragma: no cover - optional dependency
     # Allow importing vllm_omni without vllm (e.g., documentation builds)
     patch = None  # type: ignore
 
+
 from .config import OmniModelConfig
 from .entrypoints.async_omni import AsyncOmni
 

vllm_omni/config/model.py

@@ -63,6 +63,7 @@ class OmniModelConfig(ModelConfig):
     model_arch: str = "Qwen2_5OmniForConditionalGeneration"
     engine_output_type: str | None = None
     hf_config_name: str | None = None
+    omni_kv_config: dict | None = None
 
     @property
     def registry(self):

vllm_omni/core/sched/omni_ar_scheduler.py

@@ -1,7 +1,9 @@
 from __future__ import annotations
 
 from collections import defaultdict
+from dataclasses import asdict, dataclass
 from time import time
+from typing import Any
 
 import numpy as np
 from vllm.compilation.cuda_graph import CUDAGraphStat
@@ -20,6 +22,17 @@ from vllm.v1.spec_decode.metrics import SpecDecodingStats
 logger = init_logger(__name__)
 
 
+@dataclass
+class KVCacheTransferData:
+    request_id: str
+    layer_blocks: dict[str, Any]
+    block_ids: list[int]
+    metadata: dict[str, Any]
+
+    def to_dict(self) -> dict[str, Any]:
+        return asdict(self)
+
+
 class OmniARScheduler(VLLMScheduler):
     """
     OmniARScheduler: Scheduler for vLLM-Omni multimodal processing.
@@ -29,8 +42,91 @@ class OmniARScheduler(VLLMScheduler):
     specific to vLLM-Omni.
     """
 
-    # Ensure scheduled_new_reqs carry omni-specific payloads
-    # (e.g., additional_information)
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        # Track requests that need KV cache transfer when finished
+        # Value is {"seq_len": int, "block_ids": list[int]}
+        self.requests_needing_kv_transfer: dict[str, dict[str, Any]] = {}
+
+        # Track requests waiting for KV transfer (blocks not freed yet)
+        self.waiting_for_transfer_free: set[str] = set()
+
+        # Track ACTIVE transfers (submitted to runner but not yet acked via kv_extracted_req_ids)
+        self.active_kv_transfers: set[str] = set()
+
+        # [Omni] Pre-parse KV transfer criteria
+        self.kv_transfer_criteria = self._get_kv_transfer_criteria()
+
+        # Track requests that have already triggered prefill transfer to avoid duplicates
+        self.transfer_triggered_requests: set[str] = set()
+
+    def _get_kv_transfer_criteria(self) -> dict | None:
+        # Note: vllm_config is available in Scheduler after super().__init__
+        if not hasattr(self, "vllm_config"):
+            return None
+
+        omni_kv_config = getattr(self.vllm_config.model_config, "omni_kv_config", None)
+        if omni_kv_config:
+            if isinstance(omni_kv_config, dict):
+                return omni_kv_config.get("kv_transfer_criteria", None)
+            else:
+                return getattr(omni_kv_config, "kv_transfer_criteria", None)
+        return None
+
+    def _process_kv_transfer_trigger(self, request: Request, new_token_ids: list[int]) -> bool:
+        """
+        Check triggers and process side effects (marking transfer).
+        Returns True if request should be STOPPED.
+        Returns False if request should continue (even if transfer was triggered).
+        """
+        if not self.kv_transfer_criteria:
+            return False
+
+        if request.request_id in self.waiting_for_transfer_free:
+            return False
+
+        criteria_type = self.kv_transfer_criteria.get("type")
+
+        # Universal duplicate check for once semantics
+        if request.request_id in self.transfer_triggered_requests:
+            return False
+
+        if criteria_type == "prefill_finished":
+            if request.num_computed_tokens >= request.num_prompt_tokens:
+                logger.debug(f"[Omni] Request {request.request_id} triggered prefill_finished transfer (Non-Stop)")
+                self.transfer_triggered_requests.add(request.request_id)
+                self._mark_request_for_kv_transfer(request.request_id, request.num_computed_tokens)
+
+                # Return False means "Do NOT stop the request" -> Continue Decoding
+                return False
+
+        elif criteria_type == "special_token":
+            target_token_id = self.kv_transfer_criteria.get("token_id")
+            if target_token_id is not None and target_token_id in new_token_ids:
+                logger.debug(f"[Omni] Request {request.request_id} triggered special_token criteria (Non-Stop)")
+
+                self.transfer_triggered_requests.add(request.request_id)
+
+                # Calculate precise snapshot length (trim to sentinel)
+                # Find the FIRST occurrence of the sentinel
+                try:
+                    idx = new_token_ids.index(target_token_id)
+                    # seq_len = tokens_before_this_step + idx + 1 (include sentinel)
+                    # request.num_computed_tokens already includes ALL new_token_ids
+                    # so we subtract (len(new_token_ids) - (idx + 1))
+                    tokens_to_exclude = len(new_token_ids) - (idx + 1)
+                    snapshot_len = request.num_computed_tokens - tokens_to_exclude
+                except ValueError:
+                    snapshot_len = request.num_computed_tokens
+
+                # Trigger Transfer
+                self._mark_request_for_kv_transfer(request.request_id, snapshot_len)
+
+                # Do NOT stop request
+                return False
+
+        return False
+
     def schedule(self) -> SchedulerOutput:  # type: ignore[override]
         scheduler_output = super().schedule()
         try:
@@ -60,9 +156,13 @@ class OmniARScheduler(VLLMScheduler):
                 new_list.append(omni_nr)
 
             scheduler_output.scheduled_new_reqs = new_list  # type: ignore[assignment]
+
+            # Add information about requests needing KV cache transfer
+            scheduler_output.finished_requests_needing_kv_transfer = self.get_finished_requests_needing_kv_transfer()
         except Exception:
             # If anything goes wrong, leave the original output unchanged
             init_logger(__name__).exception("Failed to wrap scheduled_new_reqs with OmniNewRequestData")
+            scheduler_output.finished_requests_needing_kv_transfer = {}
 
         return scheduler_output
 
@@ -156,12 +256,21 @@ class OmniARScheduler(VLLMScheduler):
             if new_token_ids:
                 new_token_ids, stopped = self._update_request_with_output(request, new_token_ids)
 
-            if pooler_output:
-                # Note: For multimodal outputs, we skip intermediate stop checks.
+            # If criteria returns True, it means we must STOP the request (e.g. legacy behavior, not used now)
+            # If criteria returns False, it might have triggered a background transfer
+            # (e.g. prefill finished / special token) but continues decoding
+            if not stopped and self._process_kv_transfer_trigger(request, new_token_ids):
+                stopped = True
+
+            # Stop checking for pooler models.
+            pooler_output = None
+            if pooler_outputs:
+                pooler_output = pooler_outputs[req_index]
                 if request.output_token_ids:
                     stopped = check_stop(request, self.max_model_len)
             routed_experts = None
             if stopped:
+                # [Omni] Handle routed experts if enabled
                 if self.vllm_config.model_config.enable_return_routed_experts:
                     kv_blocks = self.kv_cache_manager.get_blocks(request.request_id)
                     block_ids = kv_blocks.get_block_ids()[0]
@@ -181,6 +290,7 @@ class OmniARScheduler(VLLMScheduler):
                     ).flatten()[:num_tokens]
 
                     routed_experts = self.routed_experts_reader.get_routed_experts(indices=slot_mapping)
+
                 kv_transfer_params = self._free_request(request)
                 if status_before_stop == RequestStatus.RUNNING:
                     stopped_running_reqs.add(request)
@@ -238,6 +348,7 @@ class OmniARScheduler(VLLMScheduler):
             # This is a rare case and unlikely to impact performance.
             self.waiting.remove_requests(stopped_preempted_reqs)
 
+        # [Main] Handle failed KV load requests
         if failed_kv_load_req_ids and not self.recompute_kv_load_failures:
             requests = [self.requests[req_id] for req_id in failed_kv_load_req_ids]
             self.finish_requests(failed_kv_load_req_ids, RequestStatus.FINISHED_ERROR)
@@ -252,6 +363,22 @@ class OmniARScheduler(VLLMScheduler):
                         num_cached_tokens=request.num_cached_tokens,
                     )
                 )
+
+        # [Omni] Cleanup state for finished requests
+        for req in stopped_running_reqs:
+            if req.request_id not in self.waiting_for_transfer_free:
+                if req.request_id in self.transfer_triggered_requests:
+                    self.transfer_triggered_requests.remove(req.request_id)
+                if req.request_id in self.active_kv_transfers:
+                    self.active_kv_transfers.remove(req.request_id)
+
+        # Same for preempted
+        for req in stopped_preempted_reqs:
+            if req.request_id not in self.waiting_for_transfer_free:
+                if req.request_id in self.transfer_triggered_requests:
+                    self.transfer_triggered_requests.remove(req.request_id)
+                if req.request_id in self.active_kv_transfers:
+                    self.active_kv_transfers.remove(req.request_id)
         # KV Connector: update state for finished KV Transfers.
         if kv_connector_output:
             self._update_from_kv_xfer_finished(kv_connector_output)
@@ -297,4 +424,200 @@ class OmniARScheduler(VLLMScheduler):
                 engine_core_outputs[0] = eco = EngineCoreOutputs()
             eco.scheduler_stats = stats
 
+        # This is where we free blocks that were held for transfer
+        try:
+            kv_extracted_ids = getattr(model_runner_output, "kv_extracted_req_ids", None)
+            if kv_extracted_ids:
+                for req_id in kv_extracted_ids:
+                    # Mark transfer as finished
+                    if req_id in self.active_kv_transfers:
+                        self.active_kv_transfers.remove(req_id)
+                        logger.debug(f"[Omni] KV Transfer finished for {req_id}")
+
+                    if req_id in self.waiting_for_transfer_free:
+                        # Now it's safe to free blocks
+                        req = self.requests.get(req_id)
+                        if req:
+                            self.kv_cache_manager.free(req)
+                            if req_id in self.requests:
+                                del self.requests[req_id]
+                            if req_id in self.transfer_triggered_requests:
+                                self.transfer_triggered_requests.remove(req_id)
+                            if req_id in self.active_kv_transfers:
+                                self.active_kv_transfers.remove(req_id)
+
+                            logger.debug(f"Freed blocks for {req_id} after transfer extraction")
+                        self.waiting_for_transfer_free.remove(req_id)
+        except Exception:
+            init_logger(__name__).exception("Failed to process finished transfer requests")
+
         return engine_core_outputs
+
+    def _free_request(self, request: Request) -> dict[str, Any] | None:
+        # TODO(wzliu)! for offline mode, we should not end process until all data is transferred
+        """Mark a request as finished and free its resources."""
+
+        # 1. Standard cleanup parts from base _free_request
+        delay_free_blocks = False
+        kv_xfer_params = None
+        if self.connector is not None:
+            delay_free_blocks, kv_xfer_params = self._connector_finished(request)
+
+        self.encoder_cache_manager.free(request)
+        request_id = request.request_id
+        self.finished_req_ids.add(request_id)
+        if self.finished_req_ids_dict is not None:
+            self.finished_req_ids_dict[request.client_index].add(request_id)
+
+        # 2. Omni Specific: Check if we need to transfer KV
+        if self._should_transfer_kv_for_request(request_id):
+            already_triggered = request_id in self.transfer_triggered_requests
+            is_active = request_id in self.active_kv_transfers
+
+            if already_triggered:
+                if is_active:
+                    # It triggered but hasn't finished yet. We MUST wait.
+                    logger.debug(f"[Omni] Request {request_id} finished but transfer is still ACTIVE. Waiting.")
+                    self.waiting_for_transfer_free.add(request_id)
+                    # We do NOT mark for transfer again, just wait.
+                    kv_xfer_params = None  # No new transfer params
+                    return kv_xfer_params
+                else:
+                    logger.debug(
+                        f"[Omni] Request {request_id} finished and transfer no longer ACTIVE (extracted/acked). "
+                        "Freeing immediately."
+                    )
+            else:
+                self.waiting_for_transfer_free.add(request_id)
+                self._mark_request_for_kv_transfer(request_id, request.num_computed_tokens)
+                # Return KV transfer metadata so it propagates to RequestOutput
+                if request_id in self.requests_needing_kv_transfer:
+                    transfer_data = self.requests_needing_kv_transfer[request_id]
+                    kv_xfer_params = {
+                        "past_key_values": transfer_data["block_ids"],
+                        "kv_metadata": {"seq_len": transfer_data["seq_len"], "block_ids": transfer_data["block_ids"]},
+                    }
+                    # Also update request.additional_information for good measure
+                    add_info = getattr(request, "additional_information", None)
+                    # If additional_information is an AdditionalInformationPayload-like object,
+                    # unpack list_data into a plain dict.
+                    if (
+                        add_info is not None
+                        and hasattr(add_info, "entries")
+                        and isinstance(getattr(add_info, "entries"), dict)
+                    ):
+                        request.additional_information = {
+                            k: getattr(v, "list_data")
+                            for k, v in getattr(add_info, "entries").items()
+                            if getattr(v, "list_data", None) is not None
+                        }
+                        add_info = request.additional_information
+                    if add_info is None:
+                        request.additional_information = {}
+                        add_info = request.additional_information
+                    if isinstance(add_info, dict):
+                        add_info.update(kv_xfer_params)
+
+                return kv_xfer_params
+
+        # 3. Standard Freeing
+        if not delay_free_blocks:
+            self._free_blocks(request)
+
+        return kv_xfer_params
+
+    def _free_blocks(self, request: Request):
+        # Helper to match base class structure if not directly available
+        # VLLMScheduler has _free_blocks
+        super()._free_blocks(request)
+
+    def _mark_request_for_kv_transfer(self, req_id: str, seq_len: int) -> None:
+        """Mark a request as needing KV cache transfer when it finishes."""
+        # Avoid duplicate marking (if already pending in queue)
+        if req_id in self.requests_needing_kv_transfer:
+            return
+
+        if self._should_transfer_kv_for_request(req_id):
+            # [Omni] Get block IDs from KVCacheManager
+            try:
+                block_ids_tuple = self.kv_cache_manager.get_block_ids(req_id)
+                if block_ids_tuple and len(block_ids_tuple) > 0:
+                    block_ids = block_ids_tuple[0]
+
+                    # [Omni] Fix: Truncate blocks to match seq_len snapshot
+                    # We need to know block_size. Usually in self.cache_config.block_size
+                    # Note: vllm_config might not be directly available, check scheduler_config or cache_config
+                    if hasattr(self, "cache_config") and hasattr(self.cache_config, "block_size"):
+                        block_size = self.cache_config.block_size
+                    elif hasattr(self, "scheduler_config") and hasattr(
+                        self.scheduler_config, "block_size"
+                    ):  # Some versions
+                        block_size = self.scheduler_config.block_size
+                    else:
+                        raise ValueError("Block size not found in cache_config or scheduler_config")
+
+                    # ceil(seq_len / block_size)
+                    num_blocks = (seq_len + block_size - 1) // block_size
+                    if len(block_ids) > num_blocks:
+                        logger.debug(
+                            f"[Omni] Truncating blocks for {req_id} from {len(block_ids)} "
+                            f"to {num_blocks} (seq_len={seq_len})"
+                        )
+                        block_ids = block_ids[:num_blocks]
+
+                else:
+                    block_ids = []
+            except Exception as e:
+                init_logger(__name__).warning(f"Failed to get block IDs for {req_id}: {e}")
+                block_ids = []
+
+            self.requests_needing_kv_transfer[req_id] = {"seq_len": seq_len, "block_ids": block_ids}
+            logger.debug(f"Marked request {req_id} for KV cache transfer (len={seq_len}, blocks={len(block_ids)})")
+
+    def _should_transfer_kv_for_request(self, req_id: str) -> bool:
+        """Determine if a request should trigger KV cache transfer."""
+        need_send = False
+        # Try to read from vLLM Config (where YAML config is typically loaded)
+        # Check for omni_kv_config attribute
+        omni_kv_config = getattr(self.vllm_config.model_config, "omni_kv_config", None)
+        if omni_kv_config:
+            # omni_kv_config could be an object or a dict
+            if isinstance(omni_kv_config, dict):
+                need_send = omni_kv_config.get("need_send_cache", False)
+            else:
+                need_send = getattr(omni_kv_config, "need_send_cache", False)
+        return need_send
+
+    def has_requests(self) -> bool:
+        """Check if there are any requests to process, including KV transfers."""
+        # [Omni] Also check for pending KV transfers
+        if self.requests_needing_kv_transfer or self.active_kv_transfers or self.waiting_for_transfer_free:
+            return True
+        return super().has_requests()
+
+    def has_finished_requests(self) -> bool:
+        """Check if there are any finished requests (including those needing KV transfer)."""
+        if self.requests_needing_kv_transfer or self.active_kv_transfers or self.waiting_for_transfer_free:
+            return True
+        return super().has_finished_requests()
+
+    def has_unfinished_requests(self) -> bool:
+        """Check if there are any unfinished requests (including those needing KV transfer)."""
+        # [Omni] Also check for pending KV transfers to ensure the engine loop continues
+        # MUST verify waiting_for_transfer_free and active_kv_transfers
+        # Otherwise engine loop might exit before transfer Ack is received.
+        if self.requests_needing_kv_transfer or self.active_kv_transfers or self.waiting_for_transfer_free:
+            return True
+        return super().has_unfinished_requests()
+
+    def get_finished_requests_needing_kv_transfer(self) -> dict[str, dict]:
+        """Get and clear the list of requests needing KV cache transfer.
+        Returns dict: {req_id: {"seq_len": int, "block_ids": list[int]}}
+        """
+        requests = self.requests_needing_kv_transfer.copy()
+
+        # Mark these requests as ACTIVE (sent to runner)
+        self.active_kv_transfers.update(requests.keys())
+
+        self.requests_needing_kv_transfer.clear()
+        return requests

vllm_omni/diffusion/data.py

@@ -357,6 +357,9 @@ class OmniDiffusionConfig:
     # Logging
     log_level: str = "info"
 
+    # Omni configuration (injected from stage config)
+    omni_kv_config: dict[str, Any] = field(default_factory=dict)
+
     def settle_port(self, port: int, port_inc: int = 42, max_attempts: int = 100) -> int:
         """
         Find an available port with retry logic.
@@ -453,7 +456,12 @@ class OmniDiffusionConfig:
         if "cache_backend" not in kwargs:
             cache_backend = os.environ.get("DIFFUSION_CACHE_BACKEND") or os.environ.get("DIFFUSION_CACHE_ADAPTER")
             kwargs["cache_backend"] = cache_backend.lower() if cache_backend else "none"
-        return cls(**kwargs)
+
+        # Filter kwargs to only include valid fields
+        valid_fields = {f.name for f in fields(cls)}
+        filtered_kwargs = {k: v for k, v in kwargs.items() if k in valid_fields}
+
+        return cls(**filtered_kwargs)
 
 
 @dataclass

vllm_omni/diffusion/models/bagel/bagel_transformer.py

@@ -15,8 +15,7 @@ import numpy as np
 import torch
 from torch import nn
 from torch.nn.attention.flex_attention import flex_attention
-from transformers.configuration_utils import PretrainedConfig
-from transformers.models.qwen2.configuration_qwen2 import Qwen2Config as _Qwen2Config
+from transformers.models.qwen2.configuration_qwen2 import Qwen2Config
 from transformers.models.qwen2.modeling_qwen2 import (
     Qwen2Attention,
     Qwen2MLP,
@@ -25,6 +24,7 @@ from transformers.models.qwen2.modeling_qwen2 import (
     Qwen2RotaryEmbedding,
 )
 from transformers.utils import ModelOutput
+from vllm.transformers_utils.configs.bagel import BagelConfig
 from vllm.vllm_flash_attn import flash_attn_varlen_func
 
 from vllm_omni.diffusion.layers.rope import RotaryEmbedding
@@ -71,7 +71,7 @@ torch._dynamo.config.accumulated_cache_size_limit = 4096
 flex_attention = torch.compile(flex_attention)
 
 
-class Qwen2MoTConfig(_Qwen2Config):
+class Qwen2MoTConfig(Qwen2Config):
     """Configuration for Qwen2MoT (Mixture of Tokens) model.
 
     This is fundamentally different from Qwen2, hence the distinct name.
@@ -639,36 +639,6 @@ def get_flattened_position_ids_extrapolate(img_h, img_w, patch_size, max_num_pat
     return pos_ids
 
 
-class BagelConfig(PretrainedConfig):
-    def __init__(
-        self,
-        visual_gen=True,
-        visual_und=True,
-        llm_config=None,
-        vit_config=None,
-        vae_config=None,
-        latent_patch_size=2,
-        max_latent_size=32,
-        vit_max_num_patch_per_side=70,
-        connector_act="gelu_pytorch_tanh",
-        interpolate_pos=False,
-        timestep_shift=1.0,
-        **kwargs,
-    ):
-        super().__init__(**kwargs)
-        self.visual_gen = visual_gen
-        self.visual_und = visual_und
-        self.llm_config = llm_config
-        self.vit_config = vit_config
-        self.vae_config = vae_config
-        self.latent_patch_size = latent_patch_size
-        self.max_latent_size = max_latent_size
-        self.vit_max_num_patch_per_side = vit_max_num_patch_per_side
-        self.connector_act = connector_act
-        self.interpolate_pos = interpolate_pos
-        self.timestep_shift = timestep_shift
-
-
 class Bagel(torch.nn.Module):
     config_class = BagelConfig
     base_model_prefix = "bagel"

vllm_omni/diffusion/models/bagel/pipeline_bagel.py

@@ -20,6 +20,7 @@ from torch import nn
 from transformers import AutoTokenizer, SiglipImageProcessor, SiglipVisionConfig, SiglipVisionModel
 from vllm.logger import init_logger
 from vllm.model_executor.models.utils import AutoWeightsLoader
+from vllm.transformers_utils.configs.bagel import BagelConfig
 
 from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
 from vllm_omni.diffusion.distributed.utils import get_local_device
@@ -28,7 +29,7 @@ from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.model_executor.model_loader.weight_utils import download_weights_from_hf_specific
 
 from .autoencoder import AutoEncoder, AutoEncoderParams
-from .bagel_transformer import Bagel, BagelConfig, NaiveCache, Qwen2MoTConfig, Qwen2MoTForCausalLM
+from .bagel_transformer import Bagel, NaiveCache, Qwen2MoTConfig, Qwen2MoTForCausalLM
 
 logger = init_logger(__name__)
 
@@ -304,120 +305,84 @@ class BagelPipeline(nn.Module):
         # Add text prompt (prefill) on gen context.
         # [Omni] Check for injected KV Cache from remote transfer
         injected_kv = getattr(req, "past_key_values", None)
-        injected_metadata = getattr(req, "kv_metadata", None)
+        if injected_kv is not None:
+            logger.info("Using injected KV Cache (direct)")
+            gen_context["past_key_values"] = injected_kv
 
-        # Image input handling
-        image_input = getattr(req, "pil_image", None)
-        if image_input and not isinstance(image_input, list):
-            image_input = [image_input]
-
-        if image_input:
-            # If we have an image, we prefill with it
-            if self.image_processor and self.vae:
-
-                def vit_transforms(img):
-                    # SigLIP processor returns dict with pixel_values; we want the tensor
-                    return self.image_processor(images=img, return_tensors="pt").pixel_values[0]
-
-                def vae_transforms(img):
-                    if img.mode != "RGB":
-                        img = img.convert("RGB")
-                    # Convert to [-1, 1] tensor (H, W, C) -> (C, H, W)
-                    arr = torch.from_numpy(np.array(img)).float() / 127.5 - 1.0
-                    return arr.permute(2, 0, 1)
-
-                # 1. Update VAE
-                gen_input_vae, newlens_vae, new_rope_vae = self.bagel.prepare_vae_images(
-                    curr_kvlens=gen_context["kv_lens"],
-                    curr_rope=gen_context["ropes"],
-                    images=image_input,
-                    transforms=vae_transforms,
-                    new_token_ids=self.new_token_ids,
-                )
-
-                for k, v in gen_input_vae.items():
-                    if torch.is_tensor(v):
-                        gen_input_vae[k] = v.to(self.device)
-
-                # VAE needs bfloat16 to match model strings usually, specifically encode
-                with torch.autocast(device_type="cuda", enabled=self.device.type == "cuda", dtype=torch.bfloat16):
-                    gen_context["past_key_values"] = self.bagel.forward_cache_update_vae(
-                        self.vae, gen_context["past_key_values"], **gen_input_vae
-                    )
-                gen_context["kv_lens"] = newlens_vae
-                gen_context["ropes"] = new_rope_vae
-
-                # 2. Update ViT
-                gen_input_img, newlens_img, new_rope_img = self.bagel.prepare_vit_images(
-                    curr_kvlens=gen_context["kv_lens"],
-                    curr_rope=gen_context["ropes"],
-                    images=image_input,
-                    transforms=vit_transforms,
-                    new_token_ids=self.new_token_ids,
-                )
-
-                for k, v in gen_input_img.items():
-                    if torch.is_tensor(v):
-                        gen_input_img[k] = v.to(self.device)
-
-                with torch.autocast(device_type="cuda", enabled=self.device.type == "cuda", dtype=torch.bfloat16):
-                    gen_context["past_key_values"] = self.bagel.forward_cache_update_vit(
-                        gen_context["past_key_values"], **gen_input_img
-                    )
-                gen_context["kv_lens"] = newlens_img
-                gen_context["ropes"] = new_rope_img
-            else:
-                logger.warning("Image provided but no image processor available.")
-
-        if injected_kv is not None and injected_metadata is not None:
-            logger.info("Using injected KV Cache from remote transfer")
-
-            # [Fix] Reconstruct NaiveCache if injected_kv is a dict of tensors
-            current_cache = gen_context["past_key_values"]
-            if isinstance(current_cache, NaiveCache) and isinstance(injected_kv, dict):
-                # injected_kv keys are like "0_k", "0_v", "1_k", ...
-                for key_name, tensor in injected_kv.items():
-                    try:
-                        # Parse layer index and type
-                        parts = key_name.split("_")
-                        if len(parts) < 2:
-                            continue
-
-                        layer_idx = int(parts[0])
-                        cache_type = parts[1]  # 'k' or 'v'
-
-                        # Ensure tensor is on correct device
-                        if tensor.device != self.device:
-                            tensor = tensor.to(self.device)
-
-                        if layer_idx in current_cache.key_cache:
-                            if cache_type == "k":
-                                current_cache.key_cache[layer_idx] = tensor
-                            elif cache_type == "v":
-                                current_cache.value_cache[layer_idx] = tensor
-                            elif cache_type == "kv":
-                                # Fallback if sender sent mixed/packed (less ideal)
-                                current_cache.key_cache[layer_idx] = tensor
-                                current_cache.value_cache[layer_idx] = tensor
-                    except Exception as e:
-                        logger.warning(f"Failed to load injected KV part {key_name}: {e}")
-
-            if "kv_lens" in injected_metadata:
-                val = injected_metadata["kv_lens"]
-                if isinstance(val, (int, float)):
-                    gen_context["kv_lens"] = [int(val)]
-                else:
-                    gen_context["kv_lens"] = list(val)
-
-            if "ropes" in injected_metadata:
-                val = injected_metadata["ropes"]
-                if isinstance(val, (int, float)):
-                    gen_context["ropes"] = [int(val)]
-                else:
-                    gen_context["ropes"] = list(val)
+            # User requested: kv_lens and ropes set to [gen_context["past_key_values"].key_cache[0].shape[0]]
+            # Assuming injected_kv is compatible and has key_cache[0]
+            seq_len = injected_kv.key_cache[0].shape[0]
+            gen_context["kv_lens"] = [seq_len]
+            gen_context["ropes"] = [seq_len]
 
         else:
-            # Standard local prefill path
+            image_input = getattr(req, "pil_image", None)
+            if image_input and not isinstance(image_input, list):
+                image_input = [image_input]
+
+            if image_input:
+                # If we have an image, we prefill with it
+                if self.image_processor and self.vae:
+
+                    def vit_transforms(img):
+                        # SigLIP processor returns dict with pixel_values; we want the tensor
+                        return self.image_processor(images=img, return_tensors="pt").pixel_values[0]
+
+                    def vae_transforms(img):
+                        if img.mode != "RGB":
+                            img = img.convert("RGB")
+                        # Convert to [-1, 1] tensor (H, W, C) -> (C, H, W)
+                        arr = torch.from_numpy(np.array(img)).float() / 127.5 - 1.0
+                        return arr.permute(2, 0, 1)
+
+                    # 1. Update VAE
+                    gen_input_vae, newlens_vae, new_rope_vae = self.bagel.prepare_vae_images(
+                        curr_kvlens=gen_context["kv_lens"],
+                        curr_rope=gen_context["ropes"],
+                        images=image_input,
+                        transforms=vae_transforms,
+                        new_token_ids=self.new_token_ids,
+                    )
+
+                    for k, v in gen_input_vae.items():
+                        if torch.is_tensor(v):
+                            gen_input_vae[k] = v.to(self.device)
+
+                    # VAE needs bfloat16 to match model strings usually, specifically encode
+                    with torch.autocast(
+                        device_type=self.device.type,
+                        enabled=self.device.type != "cpu",
+                        dtype=self.od_config.dtype,
+                    ):
+                        gen_context["past_key_values"] = self.bagel.forward_cache_update_vae(
+                            self.vae, gen_context["past_key_values"], **gen_input_vae
+                        )
+                    gen_context["kv_lens"] = newlens_vae
+                    gen_context["ropes"] = new_rope_vae
+
+                    # 2. Update ViT
+                    gen_input_img, newlens_img, new_rope_img = self.bagel.prepare_vit_images(
+                        curr_kvlens=gen_context["kv_lens"],
+                        curr_rope=gen_context["ropes"],
+                        images=image_input,
+                        transforms=vit_transforms,
+                        new_token_ids=self.new_token_ids,
+                    )
+
+                    for k, v in gen_input_img.items():
+                        if torch.is_tensor(v):
+                            gen_input_img[k] = v.to(self.device)
+
+                    with torch.autocast(
+                        device_type=self.device.type,
+                        enabled=self.device.type != "cpu",
+                        dtype=self.od_config.dtype,
+                    ):
+                        gen_context["past_key_values"] = self.bagel.forward_cache_update_vit(
+                            gen_context["past_key_values"], **gen_input_img
+                        )
+                    gen_context["kv_lens"] = newlens_img
+                    gen_context["ropes"] = new_rope_img
             generation_input, newlens, new_rope = self.bagel.prepare_prompts(
                 curr_kvlens=gen_context["kv_lens"],
                 curr_rope=gen_context["ropes"],
@@ -438,7 +403,11 @@ class BagelPipeline(nn.Module):
             for k, v in generation_input.items():
                 if torch.is_tensor(v):
                     generation_input[k] = v.to(self.device)
-            with torch.autocast(device_type="cuda", enabled=self.device.type == "cuda", dtype=torch.bfloat16):
+            with torch.autocast(
+                device_type=self.device.type,
+                enabled=self.device.type != "cpu",
+                dtype=self.od_config.dtype,
+            ):
                 gen_context["past_key_values"] = self.bagel.forward_cache_update_text(
                     gen_context["past_key_values"], **generation_input
                 )
@@ -483,7 +452,11 @@ class BagelPipeline(nn.Module):
             if torch.is_tensor(v):
                 generation_input[k] = v.to(self.device)
 
-        with torch.autocast(device_type="cuda", enabled=self.device.type == "cuda", dtype=torch.bfloat16):
+        with torch.autocast(
+            device_type=self.device.type,
+            enabled=self.device.type != "cpu",
+            dtype=self.od_config.dtype,
+        ):
             latents = self.bagel.generate_image(
                 past_key_values=gen_context["past_key_values"],
                 num_timesteps=gen_params.num_timesteps,

vllm_omni/diffusion/request.py

@@ -122,6 +122,11 @@ class OmniDiffusionRequest:
     # Other parameters that may be needed by specific schedulers
     extra_step_kwargs: dict[str, Any] = field(default_factory=dict)
 
+    # [Omni] KV Cache Transfer, for bagel model now
+    past_key_values: Any | None = None  # Injected KV Cache
+    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
+
     # Component modules (populated by the pipeline)
     modules: dict[str, Any] = field(default_factory=dict)
 

vllm_omni/diffusion/worker/gpu_diffusion_model_runner.py

@@ -27,6 +27,8 @@ from vllm_omni.diffusion.forward_context import set_forward_context
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
 from vllm_omni.diffusion.offload import apply_offload_hooks
 from vllm_omni.diffusion.request import OmniDiffusionRequest
+from vllm_omni.distributed.omni_connectors.factory import OmniConnectorFactory
+from vllm_omni.distributed.omni_connectors.utils.config import ConnectorSpec
 
 logger = init_logger(__name__)
 
@@ -59,6 +61,10 @@ class GPUDiffusionModelRunner:
         self.device = device
         self.pipeline = None
         self.cache_backend = None
+        self.connector = None
+
+        # Initialize OmniConnector after vllm_config is available (via init_device_and_model)
+        self._init_omni_connector()
 
     def load_model(
         self,
@@ -131,6 +137,124 @@ class GPUDiffusionModelRunner:
 
         logger.info("Model runner: Initialization complete.")
 
+    def _init_omni_connector(self) -> None:
+        # TODO(wzliu)! get real connector from yaml file instead of hardcode
+        """Initialize OmniConnector for KV cache transfer."""
+        try:
+            connector_config = None
+
+            # 1. Try to get from omni_kv_config (injected from YAML)
+            # Use self.od_config because self.vllm_config is a dummy VllmConfig without model_config
+            if self.od_config.omni_kv_config:
+                connector_config = self.od_config.omni_kv_config.get("connector_config")
+
+            if not connector_config:
+                logger.warning("No OmniConnector config found, skipping initialization")
+                return
+
+            logger.info(f"Initializing OmniConnector with config: {connector_config}")
+
+            c_type = connector_config.get("type")
+            if not c_type:
+                logger.error("Connector config missing 'type'")
+                return
+
+            c_extra = {k: v for k, v in connector_config.items() if k != "type"}
+            connector_spec = ConnectorSpec(name=c_type, extra=c_extra)
+
+            self.connector = OmniConnectorFactory.create_connector(connector_spec)
+
+        except Exception as e:
+            logger.error(f"Failed to initialize OmniConnector: {e}")
+            import traceback
+
+            traceback.print_exc()
+
+    def _receive_kv_cache_for_request(self, req: OmniDiffusionRequest) -> None:
+        """Receive KV cache for a request via OmniConnector."""
+        # TODO(wzliu)! must get control info from stage queue instead of hardcode
+        if not req.request_id:
+            logger.warning("Request has no ID, cannot receive KV cache")
+            return
+
+        try:
+            logger.info(f"Attempting to receive KV cache for request {req.request_id}")
+
+            # TODO: Key used for transfer (must match sender side)
+            # key = f"kv_cache_{req.request_id}"
+
+            # Get data from connector
+            # Determine from_stage and to_stage dynamically
+            omni_kv_config = self.od_config.omni_kv_config
+            stage_id = omni_kv_config.get("stage_id")
+            engine_input_source = omni_kv_config.get("engine_input_source", [])
+
+            to_stage = stage_id
+            # Default to stage_id - 1 if input source is not explicit
+            if engine_input_source:
+                from_stage = engine_input_source[0]
+            elif isinstance(stage_id, int):
+                from_stage = stage_id - 1
+            else:
+                raise ValueError("Invalid stage id")
+            logger.info(f"Wait for KV cache for request {req.request_id} from stage {from_stage} to {to_stage}...")
+
+            # Check if we should receive KV cache based on config
+            need_recv_cache = omni_kv_config.get("need_recv_cache", False)
+            if need_recv_cache:
+                # Default timeout 30 seconds to prevent infinite hanging
+                timeout = omni_kv_config.get("recv_timeout", 30.0)
+                start_time = time.time()
+
+                while True:
+                    result = self.connector.get(
+                        from_stage=from_stage,
+                        to_stage=to_stage,
+                        request_id=f"kv_cache_{req.request_id}",
+                    )
+                    if result:
+                        break
+
+                    if time.time() - start_time > timeout:
+                        logger.error(f"Timeout waiting for KV cache for request {req.request_id} after {timeout}s")
+                        result = None
+                        break
+
+                    time.sleep(0.5)
+            else:
+                logger.info(f"Skip receiving KV cache for {req.request_id} (need_recv_cache=False)")
+                result = None
+
+            if result:
+                data, size = result
+                logger.info(f"Successfully received KV cache for {req.request_id}")
+
+                # Assume data structure matches KVCacheTransferData.to_dict()
+                if isinstance(data, dict) and "layer_blocks" in data:
+                    # Get layer blocks and ensure they are on the correct device
+                    layer_blocks = data["layer_blocks"]
+
+                    # Move tensors to GPU if needed (OmniSerializer should handle tensor reconstruction)
+                    for cache_list in [layer_blocks["key_cache"], layer_blocks["value_cache"]]:
+                        for i, tensor in enumerate(cache_list):
+                            if isinstance(tensor, torch.Tensor) and tensor.device != self.pipeline.device:
+                                cache_list[i] = tensor.to(self.pipeline.device).contiguous()
+                    from types import SimpleNamespace
+
+                    req.past_key_values = SimpleNamespace(**layer_blocks)
+
+                if "metadata" in data:
+                    req.kv_metadata = data["metadata"]
+
+            else:
+                logger.warning(f"No KV cache received for {req.request_id} (timeout or empty)")
+
+        except Exception as e:
+            logger.error(f"Error receiving KV cache for {req.request_id}: {e}")
+            import traceback
+
+            traceback.print_exc()
+
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
         """Load weights into the pipeline."""
         return self.pipeline.load_weights(weights)
@@ -153,6 +277,10 @@ class GPUDiffusionModelRunner:
         # TODO: dealing with first req for now
         req = reqs[0]
 
+        # [Omni] KV Cache Receiving Logic
+        if getattr(req, "need_kv_receive", False) and self.connector is not None:
+            self._receive_kv_cache_for_request(req)
+
         if req.generator is None and req.seed is not None:
             req.generator = torch.Generator(device=self.device).manual_seed(req.seed)
 

vllm_omni/distributed/omni_connectors/connectors/shm_connector.py

@@ -42,7 +42,8 @@ class SharedMemoryConnector(OmniConnectorBase):
 
             if size > self.threshold:
                 # Use Shared Memory
-                meta = shm_write_bytes(payload)
+                shm_name = f"omni_{from_stage}_to_{to_stage}_{request_id}"
+                meta = shm_write_bytes(payload, name=shm_name)
                 # meta contains {'name': ..., 'size': ...}
                 metadata = {"shm": meta, "size": size}
                 self._metrics["shm_writes"] += 1
@@ -65,9 +66,40 @@ class SharedMemoryConnector(OmniConnectorBase):
     def get(
         self, from_stage: str, to_stage: str, request_id: str, metadata: dict[str, Any] | None = None
     ) -> tuple[Any, int] | None:
+        # Helper to try reading SHM with a specific name
+        def try_read_shm(name: str):
+            # _shm is imported at module level as: import multiprocessing.shared_memory as _shm
+            # but to be safe inside this closure if not picked up:
+            from multiprocessing import shared_memory as shm_pkg
+
+            shm = shm_pkg.SharedMemory(name=name)
+            try:
+                # We assume the buffer size matches the data size
+                return shm_read_bytes({"name": name, "size": shm.size})
+            finally:
+                shm.close()
+
         if not metadata:
-            logger.error(f"SharedMemoryConnector get called without metadata for req {request_id}")
-            return None
+            # Try to infer metadata from request_id for KV cache transfer scenarios
+            # where metadata is not passed out-of-band.
+            shm_name = f"omni_{from_stage}_to_{to_stage}_{request_id}"
+
+            # We need the size to read from SHM using stage_utils.shm_read_bytes
+            try:
+                # Use local import to avoid potential scope issues if _shm isn't available
+                from multiprocessing import shared_memory as shm_pkg
+
+                temp_shm = shm_pkg.SharedMemory(name=shm_name)
+                size = temp_shm.size
+                temp_shm.close()
+
+                metadata = {"shm": {"name": shm_name, "size": size}, "is_fallback": True}
+                logger.debug(f"SharedMemoryConnector inferred metadata for {request_id}: {metadata}")
+            except Exception as e:
+                # If file not found or other error, standard path will likely fail too, but we let it try or return None
+                logger.debug(f"SharedMemoryConnector failed to infer metadata for {request_id} (fallback): {e}")
+                # We can't proceed without size if we rely on shm_read_bytes
+                return None
 
         try:
             obj = None

vllm_omni/distributed/omni_connectors/factory.py

@@ -47,7 +47,7 @@ class OmniConnectorFactory:
             return connector
         except Exception as e:
             logger.error(f"Failed to create connector {spec.name}: {e}")
-            raise
+            raise ValueError(f"Failed to create connector {spec.name}: {e}")
 
     @classmethod
     def list_registered_connectors(cls) -> list[str]:

vllm_omni/distributed/omni_connectors/utils/initialization.py

@@ -318,3 +318,58 @@ def build_stage_connectors(
         raise
 
     return connectors
+
+
+def resolve_omni_kv_config_for_stage(
+    transfer_cfg: OmniTransferConfig | None, stage_id: int | str
+) -> tuple[dict[str, Any] | None, str | None, str | None]:
+    """Resolve connector configuration for a specific stage (Sender/Receiver).
+
+    This determines the primary connector configuration to be injected into the
+    engine arguments, prioritizing outgoing edges (Sender role).
+    """
+    if not transfer_cfg or not getattr(transfer_cfg, "connectors", None):
+        return None, None, None
+
+    stage_id_str = str(stage_id)
+
+    # Find outgoing edges (Sender logic)
+    outgoing = [
+        (to_stage, spec)
+        for (from_stage, to_stage), spec in transfer_cfg.connectors.items()
+        if from_stage == stage_id_str
+    ]
+
+    # Find incoming edges (Receiver logic)
+    incoming = [
+        (from_stage, spec)
+        for (from_stage, to_stage), spec in transfer_cfg.connectors.items()
+        if to_stage == stage_id_str
+    ]
+
+    omni_conn_cfg = None
+    omni_from = None
+    omni_to = None
+
+    # Prioritize outgoing (Sender) if exists, else check incoming (Receiver)
+    if outgoing:
+        if len(outgoing) > 1:
+            logger.debug(
+                "Stage-%s has %d outgoing edges; using the smallest to_stage",
+                stage_id,
+                len(outgoing),
+            )
+        outgoing.sort(key=lambda x: int(x[0]) if str(x[0]).isdigit() else str(x[0]))
+        to_s, spec = outgoing[0]
+        omni_conn_cfg = {"type": spec.name, **(spec.extra or {})}
+        omni_from = stage_id_str
+        omni_to = str(to_s)
+    elif incoming:
+        # For receiver, pick one incoming edge to configure the connector
+        incoming.sort(key=lambda x: int(x[0]) if str(x[0]).isdigit() else str(x[0]))
+        from_s, spec = incoming[0]
+        omni_conn_cfg = {"type": spec.name, **(spec.extra or {})}
+        omni_from = str(from_s)
+        omni_to = stage_id_str
+
+    return omni_conn_cfg, omni_from, omni_to

vllm_omni/engine/arg_utils.py

@@ -63,6 +63,7 @@ class OmniEngineArgs(EngineArgs):
     model_arch: str = "Qwen2_5OmniForConditionalGeneration"
     engine_output_type: str | None = None
     hf_config_name: str | None = None
+    omni_kv_config: dict | None = None
 
     def draw_hf_text_config(self, config_dict: dict) -> Qwen3OmniMoeTextConfig:
         # transformers' get_text_config method is used to get the text config from thinker_config.
@@ -115,6 +116,7 @@ class OmniEngineArgs(EngineArgs):
         config_dict["model_arch"] = self.model_arch
         config_dict["engine_output_type"] = self.engine_output_type
         config_dict["hf_config_name"] = self.hf_config_name
+        config_dict["omni_kv_config"] = self.omni_kv_config
         if self.hf_config_name is not None:
             config_dict["hf_text_config"] = self.draw_hf_text_config(config_dict)
         # Create and return the OmniModelConfig instance
@@ -145,6 +147,7 @@ class AsyncOmniEngineArgs(AsyncEngineArgs):
     model_arch: str = "Qwen2_5OmniForConditionalGeneration"
     engine_output_type: str | None = None
     hf_config_name: str | None = None
+    omni_kv_config: dict | None = None
 
     def draw_hf_text_config(self, config_dict: dict) -> Qwen3OmniMoeTextConfig:
         # transformers' get_text_config method is used to get the text config from thinker_config.
@@ -188,6 +191,7 @@ class AsyncOmniEngineArgs(AsyncEngineArgs):
         config_dict["engine_output_type"] = self.engine_output_type
 
         config_dict["hf_config_name"] = self.hf_config_name
+        config_dict["omni_kv_config"] = self.omni_kv_config
         if self.hf_config_name is not None:
             config_dict["hf_text_config"] = self.draw_hf_text_config(config_dict)
         # Create and return the OmniModelConfig instance

vllm_omni/entrypoints/async_omni_diffusion.py

@@ -57,21 +57,47 @@ class AsyncOmniDiffusion:
     ):
         self.model = model
 
+        # 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")
+
         # Build config
         if od_config is None:
             od_config = OmniDiffusionConfig.from_kwargs(model=model, **kwargs)
         elif isinstance(od_config, dict):
+            # If config is dict, check it too (priority to kwargs if both exist)
+            if stage_id is None:
+                stage_id = od_config.get("stage_id")
+            if engine_input_source is None:
+                engine_input_source = od_config.get("engine_input_source")
             od_config = OmniDiffusionConfig.from_kwargs(**od_config)
 
         self.od_config = od_config
 
-        # Load model class name and transformer config
-        config_dict = get_hf_file_to_dict("model_index.json", od_config.model)
-        od_config.model_class_name = config_dict.get("_class_name", None)
-        od_config.update_multimodal_support()
+        # Inject stage info into omni_kv_config if present
+        if stage_id is not None:
+            self.od_config.omni_kv_config.setdefault("stage_id", stage_id)
+        if engine_input_source is not None:
+            self.od_config.omni_kv_config.setdefault("engine_input_source", engine_input_source)
 
-        tf_config_dict = get_hf_file_to_dict("transformer/config.json", od_config.model)
-        od_config.tf_model_config = TransformerConfig.from_dict(tf_config_dict)
+        try:
+            config_dict = get_hf_file_to_dict("model_index.json", od_config.model)
+            od_config.model_class_name = config_dict.get("_class_name", None)
+            od_config.update_multimodal_support()
+
+            tf_config_dict = get_hf_file_to_dict("transformer/config.json", od_config.model)
+            od_config.tf_model_config = TransformerConfig.from_dict(tf_config_dict)
+        except (AttributeError, OSError, ValueError):
+            cfg = get_hf_file_to_dict("config.json", od_config.model)
+            if cfg is None:
+                raise ValueError(f"Could not find config.json or model_index.json for model {od_config.model}")
+
+            model_type = cfg.get("model_type")
+            architectures = cfg.get("architectures") or []
+            if model_type == "bagel" or "BagelForConditionalGeneration" in architectures:
+                od_config.model_class_name = "BagelPipeline"
+                od_config.tf_model_config = TransformerConfig()
+                od_config.update_multimodal_support()
 
         # Initialize engine
         self.engine: DiffusionEngine = DiffusionEngine.make_engine(od_config)

vllm_omni/entrypoints/omni.py

@@ -23,6 +23,9 @@ from vllm_omni.distributed.omni_connectors import (
     initialize_orchestrator_connectors,
 )
 from vllm_omni.distributed.omni_connectors.adapter import try_send_via_connector
+from vllm_omni.distributed.omni_connectors.utils.initialization import (
+    resolve_omni_kv_config_for_stage,
+)
 from vllm_omni.distributed.ray_utils.utils import (
     create_placement_group,
     get_ray_queue_class,
@@ -34,6 +37,7 @@ from vllm_omni.entrypoints.stage_utils import SHUTDOWN_TASK, OmniStageTaskType
 from vllm_omni.entrypoints.stage_utils import maybe_load_from_ipc as _load
 from vllm_omni.entrypoints.utils import (
     get_final_stage_id_for_e2e,
+    inject_omni_kv_config,
     load_stage_configs_from_model,
     load_stage_configs_from_yaml,
     resolve_model_config_path,
@@ -277,6 +281,18 @@ class OmniBase:
                 stage_id,
             )
 
+            # Inject YAML-resolved connector config into omni_kv_config for
+            # in-engine usage (GPU model runner reads model_config.omni_kv_config).
+            try:
+                omni_conn_cfg, omni_from, omni_to = resolve_omni_kv_config_for_stage(
+                    self.omni_transfer_config, stage_id
+                )
+                if omni_conn_cfg:
+                    inject_omni_kv_config(stage, omni_conn_cfg, omni_from, omni_to)  # type: ignore
+
+            except Exception as e:
+                logger.debug("[Omni] Failed to inject omni connector config into stage-%s: %s", stage_id, e)
+
             stage.init_stage_worker(
                 model,
                 is_async=self.is_async,

vllm_omni/entrypoints/omni_diffusion.py

@@ -2,6 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 import logging
+import uuid
 from dataclasses import fields
 
 from vllm.logger import init_logger
@@ -44,13 +45,28 @@ class OmniDiffusion:
     """
 
     def __init__(self, od_config: OmniDiffusionConfig | None = None, **kwargs):
+        # 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")
+
         if od_config is None:
             od_config = OmniDiffusionConfig.from_kwargs(**kwargs)
         elif isinstance(od_config, dict):
+            # If config is dict, check it too (priority to kwargs if both exist)
+            if stage_id is None:
+                stage_id = od_config.get("stage_id")
+            if engine_input_source is None:
+                engine_input_source = od_config.get("engine_input_source")
             od_config = OmniDiffusionConfig.from_kwargs(**od_config)
 
         self.od_config = od_config
 
+        # Inject stage info into omni_kv_config if present
+        if stage_id is not None:
+            self.od_config.omni_kv_config.setdefault("stage_id", stage_id)
+        if engine_input_source is not None:
+            self.od_config.omni_kv_config.setdefault("engine_input_source", engine_input_source)
+
         # Diffusers-style models expose `model_index.json` with `_class_name`.
         # Bagel models (and other non-diffusers) typically expose `config.json`.
         try:
@@ -99,12 +115,15 @@ class OmniDiffusion:
 
         # Check if request_id is provided in kwargs
         request_id = kwargs.get("request_id")
+        request_ids = kwargs.pop("request_ids", None)
 
         for i, p in enumerate(prompts):
             req_kwargs = kwargs.copy()
-            if request_id is None:
+            if request_ids and isinstance(request_ids, list) and i < len(request_ids):
+                req_kwargs["request_id"] = request_ids[i]
+            elif request_id is None:
                 # Generate default ID consistent with OmniLLM: "{i}_{uuid}"
-                req_kwargs["request_id"] = f"{i}"
+                req_kwargs["request_id"] = f"{i}_{uuid.uuid4()}"
 
             requests.append(
                 prepare_requests(

vllm_omni/entrypoints/omni_llm.py

@@ -121,6 +121,9 @@ class OmniLLM(LLM):
                 )
                 raise ValueError(f"Invalid 'kv_transfer_config' provided: {e}") from e
 
+        # Extract omni_kv_config from kwargs if present (injected by Omni)
+        omni_kv_config = kwargs.pop("omni_kv_config", None)
+
         if compilation_config is not None:
             if isinstance(compilation_config, int):
                 compilation_config_instance = CompilationConfig(level=compilation_config)
@@ -147,6 +150,7 @@ class OmniLLM(LLM):
             model=model,
             compilation_config=compilation_config_instance,
             structured_outputs_config=structured_outputs_instance,
+            omni_kv_config=omni_kv_config,
             **kwargs,
         )
 

vllm_omni/entrypoints/omni_stage.py

@@ -294,6 +294,7 @@ class OmniStage:
             "shm_threshold_bytes": self._shm_threshold_bytes,
             "connectors_config": connectors_config or {},
             "stage_type": self.stage_type,
+            "engine_input_source": self.engine_input_source,
         }
         try:
             old_env = os.environ.get("VLLM_LOGGING_PREFIX")
@@ -389,6 +390,36 @@ class OmniStage:
                 sampling_params, etc.)
         """
         assert self._in_q is not None
+
+        # [Omni] Inject global request_id into additional_information for cross-stage ID consistency
+        # This allows workers (like GPUARModelRunner) to use the global ID for side-channel
+        # operations like KV transfer, even if they use internal IDs for execution.
+        if "request_id" in payload and "engine_inputs" in payload:
+            req_id = payload["request_id"]
+            ein = payload["engine_inputs"]
+
+            # Helper to inject into additional_information
+            def _inject_global_id(target_ein):
+                # OmniTokensPrompt is a TypedDict at runtime, so we treat it as a dict
+                if isinstance(target_ein, dict):
+                    if "additional_information" not in target_ein:
+                        target_ein["additional_information"] = {}
+
+                    # Ensure additional_information is a dict before assignment
+                    # (in case it was somehow initialized as None or other type)
+                    if target_ein["additional_information"] is None:
+                        target_ein["additional_information"] = {}
+
+                    if isinstance(target_ein["additional_information"], dict):
+                        # Wrap in list because OmniInputProcessor requires Tensor or list values
+                        target_ein["additional_information"]["global_request_id"] = [str(req_id)]
+
+            if isinstance(ein, list):
+                for item in ein:
+                    _inject_global_id(item)
+            else:
+                _inject_global_id(ein)
+
         self._in_q.put(payload)
 
     def try_collect(self) -> dict[str, Any] | None:
@@ -466,9 +497,23 @@ def _stage_worker(
     """Stage worker entry: device setup, LLM init, batching, SHM IPC."""
     # Use local aliases to avoid conflicts with global imports in worker process
     logger.info(f"Starting stage worker with model: {model}")
+    import multiprocessing as _mp
     import os as _os
     import time as _time
 
+    # IMPORTANT: Ensure vLLM's internal multiprocessing workers (e.g., GPUARWorker /
+    # GPUARModelRunner) are spawned with a fork-safe method.
+    # Mooncake / gRPC / RDMA and CUDA/NCCL can deadlock under fork-with-threads.
+    if _os.environ.get("VLLM_WORKER_MULTIPROC_METHOD") != "spawn":
+        _os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+        logger.info("[Stage] Set VLLM_WORKER_MULTIPROC_METHOD=spawn")
+    # Best-effort: also force python mp start method in this stage process.
+    # This may raise if already set; that's fine.
+    try:
+        _mp.set_start_method("spawn", force=True)
+    except RuntimeError:
+        pass
+
     stage_id = stage_payload["stage_id"]
     engine_args = stage_payload.get("engine_args", {})
     runtime_cfg = stage_payload.get("runtime", {})
@@ -636,8 +681,14 @@ def _stage_worker(
     )
     try:
         if stage_type == "diffusion":
-            engine_args.pop("model_stage")
-            stage_engine = OmniDiffusion(**engine_args)
+            engine_args.pop("model_stage", None)
+            engine_args.pop("model", None)
+            stage_engine = OmniDiffusion(
+                model=model,
+                stage_id=stage_id,
+                engine_input_source=stage_payload.get("engine_input_source", []),
+                **engine_args,
+            )
         else:
             # Default to LLM engine
             stage_engine = OmniLLM(model=model, **engine_args)
@@ -831,6 +882,10 @@ def _stage_worker(
                         prompts.append(str(ein))
                 # Prepare diffusion kwargs from sampling parameters
                 diffusion_kwargs = prepare_sampling_params(sampling_params, "diffusion")
+
+                # Pass batch_request_ids to ensure correct ID mapping
+                diffusion_kwargs["request_ids"] = batch_request_ids
+
                 # Diffusion generate returns results directly, not an iterator
                 diffusion_results = stage_engine.generate(prompts, **diffusion_kwargs)
                 # Convert to list format compatible with LLM outputs
@@ -975,9 +1030,20 @@ async def _stage_worker_async(
 ) -> None:
     """Stage worker entry: device setup, LLM init, batching, SHM IPC."""
     # Use local aliases to avoid conflicts with global imports in worker process
+    import multiprocessing as _mp
     import os as _os
     import time as _time
 
+    # IMPORTANT: Ensure vLLM's internal multiprocessing workers (e.g., GPUARWorker /
+    # GPUARModelRunner) are spawned with a fork-safe method.
+    if _os.environ.get("VLLM_WORKER_MULTIPROC_METHOD") != "spawn":
+        _os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+        logger.info("[Stage-async] Set VLLM_WORKER_MULTIPROC_METHOD=spawn")
+    try:
+        _mp.set_start_method("spawn", force=True)
+    except RuntimeError:
+        pass
+
     stage_id = stage_payload["stage_id"]
     engine_args = stage_payload.get("engine_args", {})
     runtime_cfg = stage_payload.get("runtime", {})
@@ -1144,6 +1210,13 @@ async def _stage_worker_async(
         if stage_type == "diffusion":
             # For diffusion, we need to extract diffusion-specific config
             od_config = _build_od_config(engine_args, model)
+
+            # Inject omni config for worker to access stage info
+            if "omni_kv_config" not in od_config:
+                od_config["omni_kv_config"] = {}
+            od_config["omni_kv_config"]["stage_id"] = stage_id
+            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)
             stage_engine = AsyncOmniDiffusion(
                 model=model,

vllm_omni/entrypoints/stage_utils.py

@@ -141,12 +141,25 @@ def serialize_obj(obj: Any) -> bytes:
     return OmniSerializer.serialize(obj)
 
 
-def shm_write_bytes(payload: bytes) -> dict[str, Any]:
+def shm_write_bytes(payload: bytes, name: str | None = None) -> dict[str, Any]:
     """Write bytes into SharedMemory and return meta dict {name,size}.
 
     Caller should close the segment; the receiver should unlink.
     """
-    shm = _shm.SharedMemory(create=True, size=len(payload))
+    try:
+        shm = _shm.SharedMemory(create=True, size=len(payload), name=name)
+    except FileExistsError:
+        if name:
+            # If name is specified and exists, unlink it and try again
+            try:
+                existing = _shm.SharedMemory(name=name)
+                existing.unlink()
+            except Exception:
+                pass
+            shm = _shm.SharedMemory(create=True, size=len(payload), name=name)
+        else:
+            raise
+
     mv = memoryview(shm.buf)
     mv[: len(payload)] = payload
     del mv

vllm_omni/entrypoints/utils.py

@@ -9,6 +9,7 @@ from vllm.logger import init_logger
 from vllm.transformers_utils.config import get_config, get_hf_file_to_dict
 from vllm.transformers_utils.repo_utils import file_or_path_exists
 
+from vllm_omni.entrypoints.stage_utils import _to_dict
 from vllm_omni.utils import detect_device_type, is_rocm
 
 # Get the project root directory (2 levels up from this file)
@@ -17,6 +18,39 @@ PROJECT_ROOT = Path(__file__).parent.parent.parent
 logger = init_logger(__name__)
 
 
+def inject_omni_kv_config(stage: Any, omni_conn_cfg: dict[str, Any], omni_from: str, omni_to: str) -> None:
+    """Inject connector configuration into stage engine arguments."""
+    # Prepare omni_kv_config dict
+    omni_conf_dict = {}
+    try:
+        # Access engine_args safely (might be OmegaConf or dict)
+        existing_args = stage.engine_args
+        if hasattr(existing_args, "get"):
+            _oc = existing_args.get("omni_kv_config", None)
+            if _oc:
+                if hasattr(_oc, "items"):  # dict-like
+                    omni_conf_dict = dict(_oc)
+                else:  # object?
+                    omni_conf_dict = _to_dict(_oc)
+    except Exception:
+        omni_conf_dict = {}
+
+    # Inject connector info
+    omni_conf_dict["connector_config"] = omni_conn_cfg
+    omni_conf_dict["omni_from_stage"] = omni_from
+    omni_conf_dict["omni_to_stage"] = omni_to
+
+    # Write back to engine_args
+    try:
+        if hasattr(stage.engine_args, "__setitem__"):
+            stage.engine_args["omni_kv_config"] = omni_conf_dict
+        else:
+            setattr(stage.engine_args, "omni_kv_config", omni_conf_dict)
+    except Exception as e:
+        # Fallback for OmegaConf or similar if direct set fails?
+        logger.error(f"Failed to inject omni connector config into stage: {e}")
+
+
 def _try_get_class_name_from_diffusers_config(model: str) -> str | None:
     """Try to get class name from diffusers model configuration files.
 

vllm_omni/model_executor/models/registry.py

@@ -55,6 +55,7 @@ _OMNI_MODELS = {
     ),
 }
 
+
 _VLLM_OMNI_MODELS = {
     **_VLLM_MODELS,
     **_OMNI_MODELS,

vllm_omni/model_executor/stage_configs/bagel.yaml

@@ -0,0 +1,88 @@
+# Stage 0: Thinker (multimodal understanding + text generation)
+
+stage_args:
+  - stage_id: 0
+    stage_type: llm
+    runtime:
+      devices: "0"  # Adjust as needed
+      max_batch_size: 1
+    engine_args:
+      model_stage: thinker
+      model_arch: BagelForConditionalGeneration
+      worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
+      scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
+      gpu_memory_utilization: 0.4
+      enforce_eager: true
+      trust_remote_code: true
+      engine_output_type: text
+      distributed_executor_backend: "mp"
+      enable_prefix_caching: false
+      max_num_batched_tokens: 32768
+      tensor_parallel_size: 1
+      omni_kv_config:
+        need_send_cache: true
+        kv_transfer_criteria:
+          type: prefill_finished #or special token generated
+    final_output: true
+    final_output_type: text
+    is_comprehension: true
+    default_sampling_params:
+      temperature: 0.4
+      top_p: 0.9
+      top_k: 1
+      max_tokens: 2048
+      seed: 52
+      detokenize: True
+      repetition_penalty: 1.05
+
+  - stage_id: 1
+    stage_type: diffusion
+    runtime:
+      devices: "1"
+      max_batch_size: 1
+    engine_args:
+      model_stage: dit
+      gpu_memory_utilization: 0.4
+      enforce_eager: true
+      trust_remote_code: true
+      engine_output_type: image
+      distributed_executor_backend: "mp"
+      enable_prefix_caching: false
+      max_num_batched_tokens: 32768
+      tensor_parallel_size: 1
+      omni_kv_config:
+        need_recv_cache: true
+    engine_input_source: [0]
+
+    final_output: true
+    final_output_type: image
+    is_comprehension: false
+    default_sampling_params:
+      temperature: 0.0
+      top_p: 1.0
+      top_k: -1
+      max_tokens: 2048
+      seed: 52
+      detokenize: True
+      repetition_penalty: 1.0
+
+# Runtime edges
+runtime:
+  enabled: true
+  defaults:
+    window_size: -1
+    max_inflight: 1
+
+  # Distributed connectors configuration (optional)
+  # More connectors will be supported in the future.
+  connectors:
+    shared_memory_connector:
+      name: SharedMemoryConnector
+      extra:
+        shm_threshold_bytes: 65536 # 64KB threshold
+
+
+  edges:
+    - from: 0
+      to: 1
+      window_size: -1

vllm_omni/model_executor/stage_configs/bagel_multiconnector.yaml

@@ -0,0 +1,100 @@
+# Stage 0: Thinker (multimodal understanding + text generation)
+
+stage_args:
+  - stage_id: 0
+    stage_type: llm
+    runtime:
+      devices: "0"  # Adjust as needed
+      max_batch_size: 1
+    engine_args:
+      model_stage: thinker
+      model_arch: BagelForConditionalGeneration
+      worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
+      scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
+      gpu_memory_utilization: 0.4
+      enforce_eager: true
+      trust_remote_code: true
+      engine_output_type: text
+      distributed_executor_backend: "mp"
+      enable_prefix_caching: false
+      max_num_batched_tokens: 32768
+      tensor_parallel_size: 1
+      omni_kv_config:
+        need_send_cache: true
+        kv_transfer_criteria:
+          type: prefill_finished #or special token generated
+    final_output: true
+    final_output_type: text
+    is_comprehension: true
+    default_sampling_params:
+      temperature: 0.4
+      top_p: 0.9
+      top_k: 1
+      max_tokens: 2048
+      seed: 52
+      detokenize: True
+      repetition_penalty: 1.05
+    output_connectors:
+      to_stage_1: mooncake_connector
+
+
+  - stage_id: 1
+    stage_type: diffusion
+    runtime:
+      devices: "1"
+      max_batch_size: 1
+    engine_args:
+      model_stage: dit
+      gpu_memory_utilization: 0.4
+      enforce_eager: true
+      trust_remote_code: true
+      engine_output_type: image
+      distributed_executor_backend: "mp"
+      enable_prefix_caching: false
+      max_num_batched_tokens: 32768
+      tensor_parallel_size: 1
+      omni_kv_config:
+        need_recv_cache: true
+    engine_input_source: [0]
+
+    final_output: true
+    final_output_type: image
+    is_comprehension: false
+    default_sampling_params:
+      temperature: 0.0
+      top_p: 1.0
+      top_k: -1
+      max_tokens: 2048
+      seed: 52
+      detokenize: True
+      repetition_penalty: 1.0
+    input_connectors:
+      from_stage_0: mooncake_connector
+
+
+# Runtime edges
+runtime:
+  enabled: true
+  defaults:
+    window_size: -1
+    max_inflight: 1
+
+  # Distributed connectors configuration (optional)
+  # More connectors will be supported in the future.
+  connectors:
+    # Mooncake connector for cross-node/intra-node communication
+    mooncake_connector:
+      name: MooncakeConnector
+      extra:
+        host: "127.0.0.1"
+        metadata_server: "http://10.90.67.86:8080/metadata"
+        master: "10.90.67.86:50051"
+        segment: 512000000    # 512MB
+        localbuf: 64000000     # 64MB
+        proto: "tcp"
+
+
+  edges:
+    - from: 0
+      to: 1
+      window_size: -1

vllm_omni/outputs.py

@@ -19,6 +19,9 @@ class OmniModelRunnerOutput(ModelRunnerOutput):
     """
 
     multimodal_outputs: dict[str, torch.Tensor] | None = None
+    # IDs of requests whose KV cache has been extracted from GPU/NPU to CPU.
+    # The Scheduler can safely free the block tables for these requests.
+    kv_extracted_req_ids: list[str] | None = None
 
 
 @dataclass

vllm_omni/worker/gpu_ar_model_runner.py

@@ -35,6 +35,7 @@ from vllm.v1.worker.gpu_model_runner import (
 from vllm.v1.worker.ubatch_utils import maybe_create_ubatch_slices
 from vllm.v1.worker.utils import is_residual_scattered_for_sp
 
+from vllm_omni.core.sched.omni_ar_scheduler import KVCacheTransferData
 from vllm_omni.outputs import OmniModelRunnerOutput
 from vllm_omni.worker.gpu_model_runner import OmniGPUModelRunner
 
@@ -69,6 +70,7 @@ class GPUARModelRunner(OmniGPUModelRunner):
         # each model stage has their own hidden size
         self.hidden_size = self.model_config.hf_text_config.hidden_size
         self.inputs_embeds = self._make_buffer(self.max_num_tokens, self.hidden_size, dtype=self.dtype, numpy=False)
+        self.omni_connector = None
 
     def _make_buffer(self, *size, dtype, numpy=True):
         # Prevent ray from pinning the buffer due to large size
@@ -92,6 +94,9 @@ 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.")
 
+        # [Omni] Handle KV transfer BEFORE updating states (which removes finished requests)
+        self.kv_extracted_req_ids = self._handle_finished_requests_kv_transfer(scheduler_output)
+
         if self.vllm_config.model_config.enable_return_routed_experts:
             capturer = RoutedExpertsCapturer.get_instance()
             if capturer is not None:
@@ -274,8 +279,7 @@ class GPUARModelRunner(OmniGPUModelRunner):
                 hidden_states = model_output
                 aux_hidden_states = None
 
-            multimodal_outputs = model_output.multimodal_outputs
-            hidden_states = model_output.text_hidden_states
+            hidden_states, multimodal_outputs = self.extract_multimodal_outputs(model_output)
 
             if multimodal_outputs is not None:
                 keys_or_type = (
@@ -306,9 +310,13 @@ class GPUARModelRunner(OmniGPUModelRunner):
                     )
 
                 sample_hidden_states = hidden_states[logits_indices]
-                logits = self.model.compute_logits(
-                    sample_hidden_states, sampling_metadata=self.input_batch.sampling_metadata
-                )
+                # Try with sampling_metadata first; fall back to without for models that don't support it
+                try:
+                    logits = self.model.compute_logits(
+                        sample_hidden_states, sampling_metadata=self.input_batch.sampling_metadata
+                    )
+                except TypeError:
+                    logits = self.model.compute_logits(sample_hidden_states)
             else:
                 # Rare case.
                 assert not self.is_pooling_model
@@ -325,9 +333,13 @@ class GPUARModelRunner(OmniGPUModelRunner):
                     )
                     logits = None
                 else:
-                    logits = self.model.compute_logits(
-                        sample_hidden_states, sampling_metadata=self.input_batch.sampling_metadata
-                    )
+                    # Try with sampling_metadata first; fall back to without for models that don't support it
+                    try:
+                        logits = self.model.compute_logits(
+                            sample_hidden_states, sampling_metadata=self.input_batch.sampling_metadata
+                        )
+                    except TypeError:
+                        logits = self.model.compute_logits(sample_hidden_states)
 
                 model_output_broadcast_data: dict[str, Any] = {}
                 if logits is not None:
@@ -352,6 +364,7 @@ class GPUARModelRunner(OmniGPUModelRunner):
             multimodal_outputs,
         )
         self.kv_connector_output = kv_connector_output
+
         return None
 
     @torch.inference_mode()
@@ -362,6 +375,9 @@ class GPUARModelRunner(OmniGPUModelRunner):
         kv_connector_output = self.kv_connector_output
         self.kv_connector_output = None
 
+        kv_extracted_req_ids = getattr(self, "kv_extracted_req_ids", None)
+        self.kv_extracted_req_ids = None
+
         if self.execute_model_state is None:
             # Nothing to do (PP non-final rank case), output isn't used.
             if not kv_connector_output:
@@ -537,6 +553,7 @@ class GPUARModelRunner(OmniGPUModelRunner):
                 num_nans_in_logits=num_nans_in_logits,
                 cudagraph_stats=cudagraph_stats,
             )
+            output.kv_extracted_req_ids = kv_extracted_req_ids
 
         if not self.use_async_scheduling:
             return output
@@ -558,3 +575,202 @@ class GPUARModelRunner(OmniGPUModelRunner):
             )
 
         return async_output
+
+    def _handle_finished_requests_kv_transfer(self, scheduler_output: SchedulerOutput) -> list[str]:
+        """Handle KV cache transfer for finished requests.
+
+        Returns list of request IDs that were processed (for Scheduler to free blocks).
+        """
+        finished_reqs = getattr(scheduler_output, "finished_requests_needing_kv_transfer", {})
+        if not finished_reqs:
+            return []
+
+        logger.debug(f"Processing KV transfer for {len(finished_reqs)} requests")
+
+        extracted_ids = []
+        for req_id, data in finished_reqs.items():
+            try:
+                seq_len = data.get("seq_len", 0)
+                block_ids = data.get("block_ids", [])
+                if not block_ids:
+                    logger.warning(f"Request {req_id} has no block IDs, skipping")
+                    continue
+
+                # Extract KV cache from GPU blocks -> CPU tensors
+                kv_data = self._extract_kv_cache(req_id, block_ids, seq_len)
+                if kv_data:
+                    # Transfer to downstream stage via connector
+                    self._transfer_kv_cache(kv_data)
+
+            except Exception as e:
+                logger.error(f"Failed KV transfer for {req_id}: {e}")
+            finally:
+                extracted_ids.append(req_id)
+
+        return extracted_ids
+
+    def _extract_kv_cache(self, req_id: str, block_ids: list[int], seq_len: int) -> KVCacheTransferData | None:
+        """Extract KV cache from GPU blocks for a single request."""
+        num_layers = len(self.kv_caches)
+        key_cache = [None] * num_layers
+        value_cache = [None] * num_layers
+
+        for layer_idx, kv_tensor in enumerate(self.kv_caches):
+            # Validate block IDs
+            max_block = kv_tensor.shape[1] - 1
+            valid_ids = [bid for bid in block_ids if 0 <= bid <= max_block]
+            if not valid_ids:
+                continue
+
+            # Extract and reshape: [2, n_blocks, block_size, n_heads, head_dim]
+            # -> [2, seq_len, n_heads, head_dim]
+            selected = kv_tensor[:, valid_ids]  # [2, n_valid, block_size, n_heads, head_dim]
+            n_kv, n_blks, blk_sz, n_heads, d_head = selected.shape
+            flat = selected.reshape(n_kv, n_blks * blk_sz, n_heads, d_head)
+            if seq_len < flat.shape[1]:
+                flat = flat[:, :seq_len]
+
+            # Move to CPU
+            flat_cpu = flat.detach().cpu().contiguous()
+            key_cache[layer_idx] = flat_cpu[0]
+            value_cache[layer_idx] = flat_cpu[1]
+
+        if not any(k is not None for k in key_cache):
+            return None
+
+        return KVCacheTransferData(
+            request_id=req_id,
+            layer_blocks={"key_cache": key_cache, "value_cache": value_cache},
+            block_ids=block_ids,
+            metadata={
+                "block_size": self.cache_config.block_size,
+                "num_layers": num_layers,
+                "dtype": str(self.cache_config.cache_dtype),
+                "seq_len": seq_len,
+            },
+        )
+
+    def _transfer_kv_cache(self, kv_data: KVCacheTransferData) -> None:
+        """Transfer KV cache data to downstream stage via OmniConnector."""
+        connector = self._get_or_create_connector()
+        if not connector:
+            return
+
+        # Resolve global request ID if available
+        transfer_req_id = self._resolve_global_request_id(kv_data.request_id)
+        from_stage, to_stage = self._detect_transfer_stages()
+
+        # Prepare data and transfer with retry
+        data_dict = kv_data.to_dict()
+        data_dict["request_id"] = transfer_req_id
+
+        success, size, _ = self._transfer_with_retry(
+            connector, from_stage, to_stage, f"kv_cache_{transfer_req_id}", data_dict
+        )
+
+        if success:
+            logger.info(f"KV transfer OK: {transfer_req_id}, {size} bytes")
+        else:
+            logger.error(f"KV transfer FAILED: {transfer_req_id}")
+
+    def _get_or_create_connector(self) -> Any | None:
+        """Get existing connector or create one from config."""
+        if self.omni_connector:
+            return self.omni_connector
+
+        from vllm_omni.distributed.omni_connectors.factory import OmniConnectorFactory
+        from vllm_omni.distributed.omni_connectors.utils.config import ConnectorSpec
+
+        config = self._get_omni_connector_config()
+        if not config or not isinstance(config, dict):
+            logger.warning("No valid OmniConnector config found")
+            return None
+
+        c_type = config.get("type")
+        if not c_type:
+            logger.error("OmniConnector config missing 'type' field")
+            return None
+
+        c_extra = {k: v for k, v in config.items() if k != "type"}
+        self.omni_connector = OmniConnectorFactory.create_connector(ConnectorSpec(name=c_type, extra=c_extra))
+        return self.omni_connector
+
+    def _get_omni_connector_config(self) -> dict[str, Any] | None:
+        """Get OmniConnector configuration from model config."""
+        # Primary: omni_kv_config from YAML
+        omni_kv = getattr(self.model_config, "omni_kv_config", None)
+        if isinstance(omni_kv, dict):
+            cfg = omni_kv.get("connector_config")
+            if isinstance(cfg, dict) and cfg:
+                return cfg
+
+        # Fallback: kv_transfer_config
+        kv_cfg = getattr(self.vllm_config, "kv_transfer_config", None)
+        if kv_cfg:
+            direct = getattr(kv_cfg, "omni_connector_config", None)
+            if isinstance(direct, dict) and direct:
+                return direct
+            extra = getattr(kv_cfg, "kv_connector_extra_config", None)
+            if isinstance(extra, dict):
+                omni = extra.get("omni_connector_config")
+                if isinstance(omni, dict) and omni:
+                    return omni
+
+        return None
+
+    def _detect_transfer_stages(self) -> tuple[str, str]:
+        """Detect source and target stages for KV transfer."""
+        omni_kv = getattr(self.model_config, "omni_kv_config", None)
+        if isinstance(omni_kv, dict):
+            from_s = omni_kv.get("omni_from_stage")
+            to_s = omni_kv.get("omni_to_stage")
+            if from_s and to_s:
+                return str(from_s), str(to_s)
+
+        raise ValueError(
+            "KV transfer stages not configured. Please set 'omni_from_stage' and 'omni_to_stage' in omni_kv_config."
+        )
+
+    def _resolve_global_request_id(self, req_id: str) -> str:
+        """Resolve global request ID from request state."""
+        req_state = self.requests.get(req_id)
+        if not req_state:
+            return req_id
+
+        add_info = getattr(req_state, "additional_information_cpu", {}) or {}
+        global_id = add_info.get("global_request_id")
+        if global_id:
+            if isinstance(global_id, list) and global_id:
+                global_id = global_id[0]
+            if isinstance(global_id, bytes):
+                return global_id.decode("utf-8")
+            return str(global_id)
+        return req_id
+
+    def _transfer_with_retry(
+        self,
+        connector: Any,
+        from_stage: str,
+        to_stage: str,
+        request_id: str,
+        data: dict[str, Any],
+        max_retries: int = 3,
+    ) -> tuple[bool, int, dict[str, Any] | None]:
+        """Transfer data with retry and exponential backoff."""
+        import time
+
+        for attempt in range(max_retries):
+            try:
+                success, size, metadata = connector.put(
+                    from_stage=from_stage, to_stage=to_stage, request_id=request_id, data=data
+                )
+                if success:
+                    return success, size, metadata
+                logger.warning(f"Transfer attempt {attempt + 1} failed for {request_id}")
+            except Exception as e:
+                logger.warning(f"Transfer attempt {attempt + 1} exception: {e}")
+
+            if attempt < max_retries - 1:
+                time.sleep(0.1 * (2**attempt))
+
+        return False, 0, None