diff --git a/.buildkite/scripts/simple_test.sh b/.buildkite/scripts/simple_test.sh index 576802bfe..9ac1cb9e4 100755 --- a/.buildkite/scripts/simple_test.sh +++ b/.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 diff --git a/examples/offline_inference/bagel/end2end.py b/examples/offline_inference/bagel/end2end.py new file mode 100644 index 000000000..93000662d --- /dev/null +++ b/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() diff --git a/examples/offline_inference/bagel/run_single_prompt.sh b/examples/offline_inference/bagel/run_single_prompt.sh new file mode 100644 index 000000000..2099620e3 --- /dev/null +++ b/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} diff --git a/examples/online_serving/bagel/openai_chat_client.py b/examples/online_serving/bagel/openai_chat_client.py new file mode 100755 index 000000000..fd5f4cac5 --- /dev/null +++ b/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() diff --git a/examples/online_serving/bagel/run_server.sh b/examples/online_serving/bagel/run_server.sh new file mode 100755 index 000000000..a64057ef0 --- /dev/null +++ b/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" diff --git a/tests/distributed/omni_connectors/test_kv_flow.py b/tests/distributed/omni_connectors/test_kv_flow.py new file mode 100644 index 000000000..a8d2f03b2 --- /dev/null +++ b/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() diff --git a/vllm_omni/__init__.py b/vllm_omni/__init__.py index 5e37834c2..25ea24b19 100644 --- a/vllm_omni/__init__.py +++ b/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 diff --git a/vllm_omni/config/model.py b/vllm_omni/config/model.py index 592c501a6..da0f5716b 100644 --- a/vllm_omni/config/model.py +++ b/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): diff --git a/vllm_omni/core/sched/omni_ar_scheduler.py b/vllm_omni/core/sched/omni_ar_scheduler.py index 3c262b7c9..46717cf04 100644 --- a/vllm_omni/core/sched/omni_ar_scheduler.py +++ b/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 diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index 49a49829b..f85903b9b 100644 --- a/vllm_omni/diffusion/data.py +++ b/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 diff --git a/vllm_omni/diffusion/models/bagel/bagel_transformer.py b/vllm_omni/diffusion/models/bagel/bagel_transformer.py index 256f25e08..7f36c0b69 100644 --- a/vllm_omni/diffusion/models/bagel/bagel_transformer.py +++ b/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" diff --git a/vllm_omni/diffusion/models/bagel/pipeline_bagel.py b/vllm_omni/diffusion/models/bagel/pipeline_bagel.py index 4d47fac3b..5d3bcdd11 100644 --- a/vllm_omni/diffusion/models/bagel/pipeline_bagel.py +++ b/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, diff --git a/vllm_omni/diffusion/request.py b/vllm_omni/diffusion/request.py index cd31aba67..279b3bf1d 100644 --- a/vllm_omni/diffusion/request.py +++ b/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) diff --git a/vllm_omni/diffusion/worker/gpu_diffusion_model_runner.py b/vllm_omni/diffusion/worker/gpu_diffusion_model_runner.py index 9de448832..5fb0cc670 100644 --- a/vllm_omni/diffusion/worker/gpu_diffusion_model_runner.py +++ b/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) diff --git a/vllm_omni/distributed/omni_connectors/connectors/shm_connector.py b/vllm_omni/distributed/omni_connectors/connectors/shm_connector.py index 9f1ca994f..f65e29787 100644 --- a/vllm_omni/distributed/omni_connectors/connectors/shm_connector.py +++ b/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 diff --git a/vllm_omni/distributed/omni_connectors/factory.py b/vllm_omni/distributed/omni_connectors/factory.py index 09db8e5ae..ee640d92f 100644 --- a/vllm_omni/distributed/omni_connectors/factory.py +++ b/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]: diff --git a/vllm_omni/distributed/omni_connectors/utils/initialization.py b/vllm_omni/distributed/omni_connectors/utils/initialization.py index 608af2dfc..60fe62c13 100644 --- a/vllm_omni/distributed/omni_connectors/utils/initialization.py +++ b/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 diff --git a/vllm_omni/engine/arg_utils.py b/vllm_omni/engine/arg_utils.py index f58ab749f..3f2525ca8 100644 --- a/vllm_omni/engine/arg_utils.py +++ b/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 diff --git a/vllm_omni/entrypoints/async_omni_diffusion.py b/vllm_omni/entrypoints/async_omni_diffusion.py index 1a7f1174c..1d47205f1 100644 --- a/vllm_omni/entrypoints/async_omni_diffusion.py +++ b/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) diff --git a/vllm_omni/entrypoints/omni.py b/vllm_omni/entrypoints/omni.py index 030fe4ec5..a5f5179e2 100644 --- a/vllm_omni/entrypoints/omni.py +++ b/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, diff --git a/vllm_omni/entrypoints/omni_diffusion.py b/vllm_omni/entrypoints/omni_diffusion.py index e27e4eff6..f5b0a0a40 100644 --- a/vllm_omni/entrypoints/omni_diffusion.py +++ b/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( diff --git a/vllm_omni/entrypoints/omni_llm.py b/vllm_omni/entrypoints/omni_llm.py index 8db4b11d9..6bb8ac366 100644 --- a/vllm_omni/entrypoints/omni_llm.py +++ b/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, ) diff --git a/vllm_omni/entrypoints/omni_stage.py b/vllm_omni/entrypoints/omni_stage.py index 26ac36ea8..7beef030d 100644 --- a/vllm_omni/entrypoints/omni_stage.py +++ b/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, diff --git a/vllm_omni/entrypoints/stage_utils.py b/vllm_omni/entrypoints/stage_utils.py index b98344aca..db17116ae 100644 --- a/vllm_omni/entrypoints/stage_utils.py +++ b/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 diff --git a/vllm_omni/entrypoints/utils.py b/vllm_omni/entrypoints/utils.py index 78f20bc3c..bbb75a968 100644 --- a/vllm_omni/entrypoints/utils.py +++ b/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. diff --git a/vllm_omni/model_executor/models/registry.py b/vllm_omni/model_executor/models/registry.py index 2d9095925..747ca8f0c 100644 --- a/vllm_omni/model_executor/models/registry.py +++ b/vllm_omni/model_executor/models/registry.py @@ -55,6 +55,7 @@ _OMNI_MODELS = { ), } + _VLLM_OMNI_MODELS = { **_VLLM_MODELS, **_OMNI_MODELS, diff --git a/vllm_omni/model_executor/stage_configs/bagel.yaml b/vllm_omni/model_executor/stage_configs/bagel.yaml new file mode 100644 index 000000000..cf5dff50d --- /dev/null +++ b/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 diff --git a/vllm_omni/model_executor/stage_configs/bagel_multiconnector.yaml b/vllm_omni/model_executor/stage_configs/bagel_multiconnector.yaml new file mode 100644 index 000000000..678d5775f --- /dev/null +++ b/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 diff --git a/vllm_omni/outputs.py b/vllm_omni/outputs.py index 86f41a64d..4e0583d77 100644 --- a/vllm_omni/outputs.py +++ b/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 diff --git a/vllm_omni/worker/gpu_ar_model_runner.py b/vllm_omni/worker/gpu_ar_model_runner.py index 18c8bc761..6055c93f7 100644 --- a/vllm_omni/worker/gpu_ar_model_runner.py +++ b/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