diff --git a/examples/offline_inference/bagel/end2end.py b/examples/offline_inference/bagel/end2end.py index 584e463e1..cffe44fc3 100644 --- a/examples/offline_inference/bagel/end2end.py +++ b/examples/offline_inference/bagel/end2end.py @@ -1,8 +1,7 @@ import argparse import os -from typing import cast -from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniPromptType +from vllm_omni.inputs.data import OmniPromptType def parse_args(): @@ -58,6 +57,7 @@ def parse_args(): choices=[1, 2, 3], help="CFG parallel size: 1=batched (single GPU), 2=parallel with 2 branches (text CFG only), 3=parallel (3 GPUs).", ) + parser.add_argument("--seed", type=int, default=None, help="Random seed for generation.") args = parser.parse_args() return args @@ -88,108 +88,80 @@ def main(): from PIL import Image - if args.modality == "img2img": - from vllm_omni.entrypoints.omni_diffusion import OmniDiffusion + from vllm_omni.entrypoints.omni import Omni - print(f"[Info] Running in {args.modality} mode (Stage 1 only, cfg_parallel_size={args.cfg_parallel_size})") + omni_kwargs = {} + if args.stage_configs_path: + omni_kwargs["stage_configs_path"] = args.stage_configs_path - client = OmniDiffusion( - model=model_name, - parallel_config={"cfg_parallel_size": args.cfg_parallel_size}, - ) + omni_kwargs.update( + { + "log_stats": args.log_stats, + "init_sleep_seconds": args.init_sleep_seconds, + "batch_timeout": args.batch_timeout, + "init_timeout": args.init_timeout, + "shm_threshold_bytes": args.shm_threshold_bytes, + "worker_backend": args.worker_backend, + "ray_address": args.ray_address, + } + ) - if args.image_path: - if os.path.exists(args.image_path): - loaded_image = Image.open(args.image_path).convert("RGB") - prompts = [ - { - "prompt": cast(str, p), - "multi_modal_data": {"image": loaded_image}, - } - for p in prompts - ] - else: - print(f"[Warning] Image path {args.image_path} does not exist.") + omni = Omni(model=model_name, **omni_kwargs) - result = client.generate( - prompts, - OmniDiffusionSamplingParams( - seed=52, - need_kv_receive=False, - num_inference_steps=args.steps, - extra_args={ - "cfg_text_scale": args.cfg_text_scale, - "cfg_img_scale": args.cfg_img_scale, - }, - ), - ) - - # Ensure result is a list for iteration - if not isinstance(result, list): - omni_outputs = [result] - else: - omni_outputs = result - - else: - from vllm_omni.entrypoints.omni import Omni - - omni_kwargs = {} - if args.stage_configs_path: - omni_kwargs["stage_configs_path"] = args.stage_configs_path - - omni_kwargs.update( - { - "log_stats": args.log_stats, - "init_sleep_seconds": args.init_sleep_seconds, - "batch_timeout": args.batch_timeout, - "init_timeout": args.init_timeout, - "shm_threshold_bytes": args.shm_threshold_bytes, - "worker_backend": args.worker_backend, - "ray_address": args.ray_address, + formatted_prompts = [] + for p in prompts: + if args.modality == "img2img": + if not args.image_path or not os.path.exists(args.image_path): + raise ValueError(f"img2img requires --image-path pointing to an existing file, got: {args.image_path}") + loaded_image = Image.open(args.image_path).convert("RGB") + final_prompt_text = f"<|fim_middle|><|im_start|>{p}<|im_end|>" + prompt_dict = { + "prompt": final_prompt_text, + "multi_modal_data": {"img2img": loaded_image}, + "modalities": ["img2img"], } - ) - - omni = Omni(model=model_name, **omni_kwargs) - - formatted_prompts = [] - for p in args.prompts: - if args.modality == "img2text": - if args.image_path: - loaded_image = Image.open(args.image_path).convert("RGB") - final_prompt_text = f"<|im_start|>user\n<|image_pad|>\n{p}<|im_end|>\n<|im_start|>assistant\n" - prompt_dict = { - "prompt": final_prompt_text, - "multi_modal_data": {"image": loaded_image}, - "modalities": ["text"], - } - formatted_prompts.append(prompt_dict) - elif args.modality == "text2text": - final_prompt_text = f"<|im_start|>user\n{p}<|im_end|>\n<|im_start|>assistant\n" - prompt_dict = {"prompt": final_prompt_text, "modalities": ["text"]} - formatted_prompts.append(prompt_dict) - else: - # text2img - final_prompt_text = f"<|im_start|>{p}<|im_end|>" - prompt_dict = {"prompt": final_prompt_text, "modalities": ["image"]} - if args.negative_prompt is not None: - prompt_dict["negative_prompt"] = args.negative_prompt - formatted_prompts.append(prompt_dict) - - params_list = omni.default_sampling_params_list - if args.modality == "text2img": - params_list[0].max_tokens = 1 # type: ignore # The first stage is a SamplingParam (vllm) - if len(params_list) > 1: - diffusion_params = params_list[1] - diffusion_params.num_inference_steps = args.steps # type: ignore - extra = { - "cfg_text_scale": args.cfg_text_scale, - "cfg_img_scale": args.cfg_img_scale, + if args.negative_prompt is not None: + prompt_dict["negative_prompt"] = args.negative_prompt + formatted_prompts.append(prompt_dict) + elif args.modality == "img2text": + if args.image_path: + loaded_image = Image.open(args.image_path).convert("RGB") + final_prompt_text = f"<|im_start|>user\n<|image_pad|>\n{p}<|im_end|>\n<|im_start|>assistant\n" + prompt_dict = { + "prompt": final_prompt_text, + "multi_modal_data": {"image": loaded_image}, + "modalities": ["text"], } - if args.negative_prompt is not None: - extra["negative_prompt"] = args.negative_prompt - diffusion_params.extra_args = extra # type: ignore + formatted_prompts.append(prompt_dict) + elif args.modality == "text2text": + final_prompt_text = f"<|im_start|>user\n{p}<|im_end|>\n<|im_start|>assistant\n" + prompt_dict = {"prompt": final_prompt_text, "modalities": ["text"]} + formatted_prompts.append(prompt_dict) + else: + final_prompt_text = f"<|im_start|>{p}<|im_end|>" + prompt_dict = {"prompt": final_prompt_text, "modalities": ["image"]} + if args.negative_prompt is not None: + prompt_dict["negative_prompt"] = args.negative_prompt + formatted_prompts.append(prompt_dict) - omni_outputs = list(omni.generate(prompts=formatted_prompts, sampling_params_list=params_list)) + params_list = omni.default_sampling_params_list + if args.modality in ("text2img", "img2img"): + params_list[0].max_tokens = 1 # type: ignore + if len(params_list) > 1: + diffusion_params = params_list[1] + diffusion_params.num_inference_steps = args.steps # type: ignore + diffusion_params.cfg_parallel_size = args.cfg_parallel_size # type: ignore + if args.seed is not None: + diffusion_params.seed = args.seed # type: ignore + extra = { + "cfg_text_scale": args.cfg_text_scale, + "cfg_img_scale": args.cfg_img_scale, + } + if args.negative_prompt is not None: + extra["negative_prompt"] = args.negative_prompt + diffusion_params.extra_args = extra # type: ignore + + omni_outputs = list(omni.generate(prompts=formatted_prompts, sampling_params_list=params_list)) for i, req_output in enumerate(omni_outputs): images = getattr(req_output, "images", None) diff --git a/tests/distributed/omni_connectors/test_kv_flow.py b/tests/distributed/omni_connectors/test_kv_flow.py index 1f3d4acab..b12fc013b 100644 --- a/tests/distributed/omni_connectors/test_kv_flow.py +++ b/tests/distributed/omni_connectors/test_kv_flow.py @@ -6,6 +6,7 @@ from vllm_omni.distributed.omni_connectors.kv_transfer_manager import ( OmniKVCacheConfig, OmniKVTransferManager, ) +from vllm_omni.distributed.omni_connectors.utils.kv_utils import normalize_layer_kv from vllm_omni.inputs.data import OmniDiffusionSamplingParams pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.cache] @@ -192,8 +193,7 @@ def test_normalize_layer_kv_rejects_invalid_inputs(kv_config, common_constants, else: layer_kv = (torch.randn(2, block_size, num_heads, head_dim), "not-a-tensor") - manager = OmniKVTransferManager(kv_config) - normalized = manager._normalize_layer_kv(layer_kv, req_id=req_id, layer_idx=0) + normalized = normalize_layer_kv(layer_kv, req_id=req_id, layer_idx=0) assert normalized is None diff --git a/tests/e2e/offline_inference/stage_configs/bagel_mooncake_ci.yaml b/tests/e2e/offline_inference/stage_configs/bagel_mooncake_ci.yaml index 79e825d5d..dbb93344b 100644 --- a/tests/e2e/offline_inference/stage_configs/bagel_mooncake_ci.yaml +++ b/tests/e2e/offline_inference/stage_configs/bagel_mooncake_ci.yaml @@ -13,7 +13,7 @@ stage_args: model_arch: BagelForConditionalGeneration worker_type: ar scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler - gpu_memory_utilization: 0.35 + gpu_memory_utilization: 0.45 enforce_eager: true trust_remote_code: true engine_output_type: text @@ -46,7 +46,7 @@ stage_args: max_batch_size: 1 engine_args: model_stage: dit - gpu_memory_utilization: 0.55 + gpu_memory_utilization: 0.45 enforce_eager: true trust_remote_code: true engine_output_type: image diff --git a/tests/e2e/offline_inference/stage_configs/bagel_sharedmemory_ci.yaml b/tests/e2e/offline_inference/stage_configs/bagel_sharedmemory_ci.yaml index eeefab236..686600f72 100644 --- a/tests/e2e/offline_inference/stage_configs/bagel_sharedmemory_ci.yaml +++ b/tests/e2e/offline_inference/stage_configs/bagel_sharedmemory_ci.yaml @@ -13,7 +13,7 @@ stage_args: model_arch: BagelForConditionalGeneration worker_type: ar scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler - gpu_memory_utilization: 0.35 + gpu_memory_utilization: 0.45 enforce_eager: true trust_remote_code: true engine_output_type: text @@ -45,7 +45,7 @@ stage_args: max_batch_size: 1 engine_args: model_stage: dit - gpu_memory_utilization: 0.55 + gpu_memory_utilization: 0.45 enforce_eager: true trust_remote_code: true engine_output_type: image diff --git a/vllm_omni/diffusion/models/bagel/pipeline_bagel.py b/vllm_omni/diffusion/models/bagel/pipeline_bagel.py index 46ce4b6bd..9ed3481d7 100644 --- a/vllm_omni/diffusion/models/bagel/pipeline_bagel.py +++ b/vllm_omni/diffusion/models/bagel/pipeline_bagel.py @@ -332,6 +332,9 @@ class BagelPipeline(nn.Module): else: gen_context["ropes"] = [seq_len] + if req.sampling_params.kv_metadata and "image_shape" in req.sampling_params.kv_metadata: + image_shape = tuple(req.sampling_params.kv_metadata["image_shape"]) + cfg_text_kv = getattr(req.sampling_params, "cfg_text_past_key_values", None) if cfg_text_kv is not None: logger.info("CFG enabled with multi-KV: using injected cfg_text KV Cache") diff --git a/vllm_omni/distributed/omni_connectors/kv_transfer_manager.py b/vllm_omni/distributed/omni_connectors/kv_transfer_manager.py index 21f09c951..2806f7315 100644 --- a/vllm_omni/distributed/omni_connectors/kv_transfer_manager.py +++ b/vllm_omni/distributed/omni_connectors/kv_transfer_manager.py @@ -12,6 +12,7 @@ from vllm.logger import init_logger from .factory import OmniConnectorFactory from .utils.config import ConnectorSpec +from .utils.kv_utils import normalize_layer_kv logger = init_logger(__name__) @@ -200,8 +201,12 @@ class OmniKVTransferManager: logger.warning(f"Request {req_id} has no block IDs, skipping") continue + custom_metadata = data.get("custom_metadata") + # Extract KV cache from GPU blocks -> CPU tensors - kv_data = self._extract_kv_cache(req_id, block_ids, seq_len, kv_caches, block_size, cache_dtype) + kv_data = self._extract_kv_cache( + req_id, block_ids, seq_len, kv_caches, block_size, cache_dtype, custom_metadata + ) if kv_data: # Resolve global request ID if available transfer_req_id = request_id_resolver(req_id) if request_id_resolver else req_id @@ -224,6 +229,7 @@ class OmniKVTransferManager: kv_caches: list[LayerKV], block_size: int, cache_dtype: str, + custom_metadata: dict[str, Any] | None = None, ) -> KVCacheTransferData | None: """Extract KV cache from GPU blocks for a single request. @@ -234,6 +240,7 @@ class OmniKVTransferManager: kv_caches: List of KV cache (tensor or tuple) per layer block_size: Size of each cache block cache_dtype: Data type of the cache + custom_metadata: Optional custom metadata to include Note: If key/value block counts differ, extraction uses only the overlapping block range. Extra key/value blocks are ignored, so returned KV may be partial. @@ -246,7 +253,7 @@ class OmniKVTransferManager: value_cache: list[torch.Tensor | None] = [None] * num_layers for layer_idx, layer_kv in enumerate(kv_caches): - kv_pair = self._normalize_layer_kv(layer_kv, req_id=req_id, layer_idx=layer_idx) + kv_pair = normalize_layer_kv(layer_kv, req_id=req_id, layer_idx=layer_idx) if kv_pair is None: continue key_blocks, value_blocks = kv_pair @@ -289,57 +296,10 @@ class OmniKVTransferManager: "num_layers": num_layers, "dtype": str(cache_dtype), "seq_len": seq_len, + **(custom_metadata or {}), }, ) - def _normalize_layer_kv( - self, - layer_kv: LayerKV, - req_id: str, - layer_idx: int, - ) -> tuple[torch.Tensor, torch.Tensor] | None: - """Normalize one layer KV cache to a `(key_blocks, value_blocks)` tuple. - - Args: - layer_kv: The raw KV cache (tensor or tuple) for the layer - req_id: Request ID for logging - layer_idx: Layer index for logging - - Returns: - Tuple of (key_blocks, value_blocks) if valid, None otherwise - """ - if isinstance(layer_kv, torch.Tensor): - if layer_kv.ndim < 3 or layer_kv.shape[0] != 2: - logger.warning( - f"Layer {layer_idx} for request {req_id} has invalid stacked KV shape: " - f"expected [2, blocks, block_size, ...], got {tuple(layer_kv.shape)}" - ) - return None - key_blocks = layer_kv[0] - value_blocks = layer_kv[1] - elif isinstance(layer_kv, tuple): - if len(layer_kv) != 2: - logger.warning( - f"Layer {layer_idx} for request {req_id} has KV pair length {len(layer_kv)} (expected 2)" - ) - return None - key_blocks, value_blocks = layer_kv - if not isinstance(key_blocks, torch.Tensor) or not isinstance(value_blocks, torch.Tensor): - logger.warning(f"Layer {layer_idx} for request {req_id} has non-tensor KV pair entries") - return None - else: - logger.warning(f"Layer {layer_idx} for request {req_id} has unsupported KV type {type(layer_kv).__name__}") - return None - # ensure key/value blocks are at least 2D for block indexing - if key_blocks.ndim < 2 or value_blocks.ndim < 2: - logger.warning( - f"Layer {layer_idx} for request {req_id} has invalid KV block shape: " - f"got key={tuple(key_blocks.shape)} value={tuple(value_blocks.shape)}" - ) - return None - - return key_blocks, value_blocks - def _transfer_kv_cache(self, kv_data: KVCacheTransferData, transfer_req_id: str) -> None: """Transfer KV cache data to downstream stage via OmniConnector. @@ -496,6 +456,8 @@ class OmniKVTransferManager: if "metadata" in data: req.kv_metadata = data["metadata"] + if hasattr(req, "sampling_params") and req.sampling_params is not None: + req.sampling_params.kv_metadata = data["metadata"] # Legacy compatibility method def receive_kv_cache(self, req: Any, target_device: torch.device | None = None) -> bool: diff --git a/vllm_omni/distributed/omni_connectors/utils/kv_utils.py b/vllm_omni/distributed/omni_connectors/utils/kv_utils.py new file mode 100644 index 000000000..2cb48a8b3 --- /dev/null +++ b/vllm_omni/distributed/omni_connectors/utils/kv_utils.py @@ -0,0 +1,85 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Utility helpers for KV cache manipulation.""" + +import torch +from vllm.logger import init_logger + +logger = init_logger(__name__) + +LayerKV = torch.Tensor | tuple[torch.Tensor, torch.Tensor] + + +def normalize_layer_kv( + layer_kv: LayerKV, + *, + req_id: str = "", + layer_idx: int = -1, +) -> tuple[torch.Tensor, torch.Tensor] | None: + """Normalize one layer KV cache to a ``(key_blocks, value_blocks)`` tuple. + + In vLLM, different attention backends return paged-attention KV blocks + with different layouts. For example: + + * **FlashAttention** (``vllm/v1/attention/backends/flash_attn.py``) + returns shape ``(2, num_blocks, block_size, num_kv_heads, head_size)`` + – the key/value dimension is at **dim 0**. + * **FlashInfer** (``vllm/v1/attention/backends/flashinfer.py``) + returns shape ``(num_blocks, 2, block_size, num_kv_heads, head_size)`` + – the key/value dimension is at **dim 1**. + + This utility handles both layouts (and the tuple case) so that + downstream code can work with any backend. + + Supported layouts: + + * **Stacked tensor** ``[2, num_blocks, block_size, n_heads, head_dim]`` – + dim-0 selects key / value. + * **Stacked tensor** ``[num_blocks, 2, block_size, n_heads, head_dim]`` – + dim-1 selects key / value. + * **Tuple** ``(key_tensor, value_tensor)`` – returned as-is after + validation. + + Args: + layer_kv: The raw KV cache (tensor or tuple) for the layer. + req_id: Request ID used only for diagnostic log messages. + layer_idx: Layer index used only for diagnostic log messages. + + Returns: + ``(key_blocks, value_blocks)`` if *layer_kv* is valid, ``None`` + otherwise. + """ + if isinstance(layer_kv, torch.Tensor): + if layer_kv.ndim >= 3 and layer_kv.shape[0] == 2: + key_blocks = layer_kv[0] + value_blocks = layer_kv[1] + elif layer_kv.ndim >= 3 and layer_kv.shape[1] == 2: + key_blocks = layer_kv[:, 0] + value_blocks = layer_kv[:, 1] + else: + logger.warning( + f"Layer {layer_idx} for request {req_id} has invalid stacked KV shape: " + f"expected [2, ...] or [..., 2, ...] at dim 0/1, got {tuple(layer_kv.shape)}" + ) + return None + elif isinstance(layer_kv, tuple): + if len(layer_kv) != 2: + logger.warning(f"Layer {layer_idx} for request {req_id} has KV pair length {len(layer_kv)} (expected 2)") + return None + key_blocks, value_blocks = layer_kv + if not isinstance(key_blocks, torch.Tensor) or not isinstance(value_blocks, torch.Tensor): + logger.warning(f"Layer {layer_idx} for request {req_id} has non-tensor KV pair entries") + return None + else: + logger.warning(f"Layer {layer_idx} for request {req_id} has unsupported KV type {type(layer_kv).__name__}") + return None + + # ensure key/value blocks are at least 2D for block indexing + if key_blocks.ndim < 2 or value_blocks.ndim < 2: + logger.warning( + f"Layer {layer_idx} for request {req_id} has invalid KV block shape: " + f"got key={tuple(key_blocks.shape)} value={tuple(value_blocks.shape)}" + ) + return None + + return key_blocks, value_blocks diff --git a/vllm_omni/model_executor/models/__init__.py b/vllm_omni/model_executor/models/__init__.py index 0b2629b4a..68074bbb9 100644 --- a/vllm_omni/model_executor/models/__init__.py +++ b/vllm_omni/model_executor/models/__init__.py @@ -1,4 +1,8 @@ +from .bagel.bagel import OmniBagelForConditionalGeneration from .qwen3_omni import Qwen3OmniMoeForConditionalGeneration from .registry import OmniModelRegistry # noqa: F401 -__all__ = ["Qwen3OmniMoeForConditionalGeneration"] +__all__ = [ + "Qwen3OmniMoeForConditionalGeneration", + "OmniBagelForConditionalGeneration", +] diff --git a/vllm_omni/model_executor/models/bagel/__init__.py b/vllm_omni/model_executor/models/bagel/__init__.py new file mode 100644 index 000000000..11e4c6667 --- /dev/null +++ b/vllm_omni/model_executor/models/bagel/__init__.py @@ -0,0 +1,3 @@ +from .bagel import OmniBagelForConditionalGeneration + +__all__ = ["OmniBagelForConditionalGeneration"] diff --git a/vllm_omni/model_executor/models/bagel/bagel.py b/vllm_omni/model_executor/models/bagel/bagel.py new file mode 100644 index 000000000..24a1331f9 --- /dev/null +++ b/vllm_omni/model_executor/models/bagel/bagel.py @@ -0,0 +1,1016 @@ +from collections import deque +from collections.abc import Iterable, Mapping, Sequence +from math import isqrt +from typing import Any + +import torch +import torch.nn as nn +from transformers import BatchFeature +from vllm.config import VllmConfig +from vllm.config.multimodal import BaseDummyOptions +from vllm.model_executor.layers.layernorm import RMSNorm as VllmRMSNorm +from vllm.model_executor.layers.linear import ( + QKVParallelLinear, + RowParallelLinear, +) +from vllm.model_executor.model_loader.weight_utils import default_weight_loader +from vllm.model_executor.models.bagel import BagelForConditionalGeneration +from vllm.model_executor.models.interfaces import MultiModalEmbeddings +from vllm.model_executor.models.qwen2 import Qwen2DecoderLayer, Qwen2MLP +from vllm.model_executor.models.utils import AutoWeightsLoader +from vllm.multimodal import MULTIMODAL_REGISTRY +from vllm.multimodal.inputs import ( + MultiModalDataDict, + MultiModalFieldConfig, + MultiModalKwargsItems, +) +from vllm.multimodal.parse import ( + ImageEmbeddingItems, + ImageProcessorItems, + ModalityData, + ModalityDataItems, + MultiModalDataItems, + MultiModalDataParser, +) +from vllm.multimodal.processing import ( + BaseDummyInputsBuilder, + BaseMultiModalProcessor, + BaseProcessingInfo, + PromptReplacement, + PromptUpdateDetails, +) +from vllm.transformers_utils.processors.bagel import BagelProcessor + +from vllm_omni.diffusion.models.bagel.autoencoder import ( + AutoEncoderParams, + DiagonalGaussian, + Encoder, +) +from vllm_omni.diffusion.models.bagel.bagel_transformer import ( + PositionEmbedding, + TimestepEmbedder, +) +from vllm_omni.diffusion.models.bagel.pipeline_bagel import default_ae_params + + +class OmniBagelProcessor(BagelProcessor): + def __call__(self, text=None, images=None, **kwargs): + is_img2img = kwargs.pop("is_img2img", False) + + if is_img2img and images is not None: + image_kwargs = kwargs.copy() + image_kwargs["do_resize"] = False + image_kwargs["do_rescale"] = True + if "return_tensors" not in image_kwargs: + image_kwargs["return_tensors"] = "pt" + + pixel_values = self.image_processor(images, **image_kwargs) + + text_inputs = self.tokenizer(text, **kwargs) if text is not None else None + + if pixel_values is not None and text_inputs is not None: + combined = dict(text_inputs) + combined["pixel_values"] = pixel_values["pixel_values"] + return BatchFeature(combined) + elif pixel_values is not None: + return pixel_values + elif text_inputs is not None: + return BatchFeature(dict(text_inputs)) + else: + return BatchFeature({}) + + return super().__call__(text, images, **kwargs) + + +class OmniBagelProcessingInfo(BaseProcessingInfo): + def get_supported_mm_limits(self) -> Mapping[str, int | None]: + return {"image": 1, "img2img": 1} + + def get_hf_processor(self, **kwargs: object): + return self.ctx.get_hf_processor(OmniBagelProcessor, **kwargs) + + def get_hf_config(self): + config = super().get_hf_config() + if not getattr(self, "_latent_size_patched", False): + self._latent_size_patched = True + self._patch_max_latent_size(config) + return config + + def _patch_max_latent_size(self, config): + """Infer correct max_latent_size from the model's latent_pos_embed + weight, since the HF config value may be stale (e.g. 32 vs 64).""" + import json + from pathlib import Path + + model_name = self.ctx.model_config.model + try: + p = Path(model_name) + if p.is_dir(): + index_path = p / "model.safetensors.index.json" + else: + from huggingface_hub import hf_hub_download + + index_path = Path(hf_hub_download(model_name, "model.safetensors.index.json")) + + if not index_path.exists(): + return + + with open(index_path) as f: + index = json.load(f) + + shard = index.get("weight_map", {}).get("latent_pos_embed.pos_embed") + if not shard: + return + + from safetensors import safe_open + + with safe_open(str(index_path.parent / shard), framework="pt") as f: + if "latent_pos_embed.pos_embed" in f.keys(): + npos = f.get_slice("latent_pos_embed.pos_embed").get_shape()[0] + side = isqrt(npos) + if side * side == npos: + old = getattr(config, "max_latent_size", 32) + if old != side: + config.max_latent_size = side + except Exception: + pass + + def get_data_parser(self) -> "OmniBagelDataParser": + return OmniBagelDataParser( + expected_hidden_size=self._get_expected_hidden_size(), + ) + + +class OmniBagelDummyInputsBuilder(BaseDummyInputsBuilder[OmniBagelProcessingInfo]): + def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str: + dummy_text = "" + if "image" in mm_counts: + dummy_text += "<|image_pad|>" * mm_counts["image"] + if "img2img" in mm_counts: + dummy_text += "<|fim_middle|>" * mm_counts["img2img"] + return dummy_text + + def get_dummy_mm_data( + self, + seq_len: int, + mm_counts: Mapping[str, int], + mm_options: Mapping[str, BaseDummyOptions] | None = None, + ) -> MultiModalDataDict: + num_images = mm_counts.get("image", 0) + hf_config = self.info.get_hf_config() + vit_config = hf_config.vit_config + + image_size = vit_config.image_size + image_overrides = mm_options.get("image") if mm_options else None + + return { + "image": self._get_dummy_images( + width=image_size, + height=image_size, + num_images=num_images, + overrides=image_overrides, + ), + "img2img": self._get_dummy_images( + width=image_size, + height=image_size, + num_images=mm_counts.get("img2img", 0), + overrides=image_overrides, + ), + } + + +class Img2ImgProcessorItems(ImageProcessorItems): + def __init__(self, data): + super().__init__(data) + self.modality = "img2img" + + def get_processor_data(self): + return {"pixel_values_img2img": self.get_all()} + + +class OmniBagelDataParser(MultiModalDataParser): + def _parse_img2img_data(self, data: ModalityData) -> ModalityDataItems | None: + items = self._parse_image_data(data) + if items is None: + return None + return Img2ImgProcessorItems(items.data) + + def _get_subparsers(self): + parsers = super()._get_subparsers() + parsers["img2img"] = self._parse_img2img_data + return parsers + + +class OmniBagelMultiModalProcessor(BaseMultiModalProcessor[OmniBagelProcessingInfo]): + def _get_mm_fields_config(self, hf_inputs, hf_processor_mm_kwargs): + return { + "pixel_values": MultiModalFieldConfig.batched("image"), + "pixel_values_img2img": MultiModalFieldConfig.batched("img2img"), + } + + def _call_hf_processor( + self, + prompt: str, + mm_data: Mapping[str, object], + mm_kwargs: Mapping[str, object], + tok_kwargs: Mapping[str, object], + ) -> "BatchFeature": + has_image = "images" in mm_data + has_img2img = "pixel_values_img2img" in mm_data + + if has_image and has_img2img: + outputs = BatchFeature() + + img_data = dict(mm_data) + if "pixel_values_img2img" in img_data: + del img_data["pixel_values_img2img"] + kwargs_img = dict(mm_kwargs) + kwargs_img["is_img2img"] = False + out_img = super()._call_hf_processor(prompt, img_data, kwargs_img, tok_kwargs) + if "pixel_values" in out_img: + outputs["pixel_values"] = out_img["pixel_values"] + for k, v in out_img.items(): + if k != "pixel_values": + outputs[k] = v + + img2img_data = dict(mm_data) + if "images" in img2img_data: + del img2img_data["images"] + img2img_data["images"] = img2img_data.pop("pixel_values_img2img") + kwargs_img2img = dict(mm_kwargs) + kwargs_img2img["is_img2img"] = True + out_img2img = super()._call_hf_processor(prompt, img2img_data, kwargs_img2img, tok_kwargs) + if "pixel_values" in out_img2img: + outputs["pixel_values_img2img"] = out_img2img["pixel_values"] + for k, v in out_img2img.items(): + if k not in outputs: + outputs[k] = v + + return outputs + + elif has_img2img: + mm_data = dict(mm_data) + mm_data["images"] = mm_data.pop("pixel_values_img2img") + mm_kwargs = dict(mm_kwargs) + mm_kwargs["is_img2img"] = True + outputs = super()._call_hf_processor(prompt, mm_data, mm_kwargs, tok_kwargs) + if "pixel_values" in outputs: + outputs["pixel_values_img2img"] = outputs.pop("pixel_values") + return outputs + + return super()._call_hf_processor(prompt, mm_data, mm_kwargs, tok_kwargs) + + def _hf_processor_applies_updates( + self, + prompt_text: str, + mm_items: MultiModalDataItems, + hf_processor_mm_kwargs: Mapping[str, object], + tokenization_kwargs: Mapping[str, object], + ) -> bool: + return False + + def _get_prompt_updates( + self, + mm_items: MultiModalDataItems, + hf_processor_mm_kwargs: Mapping[str, Any], + out_mm_kwargs: MultiModalKwargsItems, + ) -> Sequence[PromptReplacement]: + hf_config = self.info.get_hf_config() + tokenizer = self.info.get_tokenizer() + + replacements: list[PromptReplacement] = [] + + image_token_id = tokenizer.get_vocab().get("<|image_pad|>") + if image_token_id is not None: + num_patches = hf_config.vit_max_num_patch_per_side**2 + + def get_image_replacement(item_idx: int): + return [image_token_id] * num_patches + + replacements.append( + PromptReplacement( + modality="image", + target=[image_token_id], + replacement=get_image_replacement, + ) + ) + + img2img_token_id = tokenizer.get_vocab().get("<|fim_middle|>") + if img2img_token_id is not None: + vit_config = hf_config.vit_config + image_size = vit_config.image_size + num_vit_patches = (image_size // vit_config.patch_size) ** 2 + + latent_patch_size = getattr(hf_config, "latent_patch_size", 2) + downsample = hf_config.vae_config.get("downsample", 8) + latent_downsample = downsample * latent_patch_size + + def get_img2img_replacement(item_idx: int): + h, w = image_size, image_size + if "img2img" in mm_items: + item = mm_items.get_items("img2img", (Img2ImgProcessorItems, ImageEmbeddingItems)) + if hasattr(item, "get_image_size"): + size = item.get_image_size(item_idx) + h, w = size.height, size.width + + max_latent_size = getattr(hf_config, "max_latent_size", 32) + max_img_size = int(max_latent_size * latent_downsample) + stride = latent_downsample + scale = min(max_img_size / max(h, w), 1.0) + min_img_size = min(256, max_img_size) + scale = max(scale, min_img_size / min(h, w)) + new_h = max(stride, int(round(h * scale / stride) * stride)) + new_w = max(stride, int(round(w * scale / stride) * stride)) + new_h = min(new_h, max_img_size) + new_w = min(new_w, max_img_size) + + num_vae_patches = (new_h // latent_downsample) * (new_w // latent_downsample) + num_vae_total = num_vae_patches + 2 + num_vit_total = num_vit_patches + 2 + # +1 separator between VAE and ViT blocks so that + # extract_embeds_range() produces two distinct mm_prefix_range + # entries, preventing VAE tokens from attending to ViT. + total = num_vae_total + 1 + num_vit_total + tokens = [img2img_token_id] * total + + embed_mask = [True] * num_vae_total + [False] + [True] * num_vit_total + return PromptUpdateDetails( + full=tokens, + is_embed=lambda _tok, _seq, _m=embed_mask: torch.tensor(_m, dtype=torch.bool), + ) + + replacements.append( + PromptReplacement( + modality="img2img", + target=[img2img_token_id], + replacement=get_img2img_replacement, + ) + ) + + return replacements + + +class VAEEncoder(nn.Module): + """Lightweight VAE encoder (no decoder) for embedding images in the AR stage.""" + + def __init__(self, params: AutoEncoderParams): + super().__init__() + self.encoder = Encoder( + resolution=params.resolution, + in_channels=params.in_channels, + ch=params.ch, + ch_mult=params.ch_mult, + num_res_blocks=params.num_res_blocks, + z_channels=params.z_channels, + ) + self.reg = DiagonalGaussian() + self.scale_factor = params.scale_factor + self.shift_factor = params.shift_factor + + def encode(self, x: torch.Tensor) -> torch.Tensor: + z = self.reg(self.encoder(x)) + z = self.scale_factor * (z - self.shift_factor) + return z + + +@MULTIMODAL_REGISTRY.register_processor( + OmniBagelMultiModalProcessor, + info=OmniBagelProcessingInfo, + dummy_inputs=OmniBagelDummyInputsBuilder, +) +class OmniBagelForConditionalGeneration(BagelForConditionalGeneration): + """ + Omni version of BagelForConditionalGeneration. + + Extends the base model with a VAE encoder so that img2img can embed + both VAE latents and ViT features within the AR stage, producing a + combined KV cache that is then transferred to the DiT stage. + + Position IDs are adjusted so that: + - VAE tokens all share position 0 + - ViT tokens all share position 1 + - Text tokens use sequential positions starting from 2 + This matches the position scheme used by the single-stage DiT pipeline, + ensuring the transferred KV cache + ropes are directly compatible with + the DiT's denoising loop. + """ + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__(vllm_config=vllm_config, prefix=prefix) + config = vllm_config.model_config.hf_config + self.latent_patch_size = getattr(config, "latent_patch_size", 2) + self.downsample = config.vae_config.get("downsample") + self.latent_downsample = self.downsample * self.latent_patch_size + self.max_latent_size = getattr(config, "max_latent_size", 32) + self.latent_channel = config.vae_config.get("z_channels") + + hidden_size = config.llm_config.hidden_size + patch_latent_dim = self.latent_patch_size**2 * self.latent_channel + self.vae = VAEEncoder(default_ae_params()) + self.vae2llm = nn.Linear(patch_latent_dim, hidden_size) + self.latent_pos_embed = PositionEmbedding(self.max_latent_size, hidden_size) + self.time_embedder = TimestepEmbedder(hidden_size) + + self._pending_img2img_info: list[tuple[int, int, int, int]] = [] + self._ropes_pending: list[dict[str, Any]] = [] + self._ropes_metadata: dict[str, dict[str, Any]] = {} + self._cfg_companion_queue: deque[tuple[tuple[int, int, int, int], int]] = deque() + + from transformers import AutoTokenizer + + tok_name = getattr(vllm_config.model_config, "tokenizer", None) or vllm_config.model_config.model + _tok = AutoTokenizer.from_pretrained(tok_name, trust_remote_code=True) + for t in ["<|vision_start|>", "<|vision_end|>"]: + if t not in _tok.get_vocab(): + _tok.add_tokens([t]) + self._start_of_image_id = int(_tok.convert_tokens_to_ids("<|vision_start|>")) + self._end_of_image_id = int(_tok.convert_tokens_to_ids("<|vision_end|>")) + + self._vae_token_mask: torch.Tensor | None = None + + self._install_mot_modules(config) + + def _install_mot_modules(self, config): + """Add generation-mode (MoT) weight modules to each Qwen2 decoder layer. + + The single-stage DiT routes VAE latent tokens through separate + ``qkv_proj_moe_gen / o_proj_moe_gen / mlp_moe_gen`` weight matrices + (``mode="gen"``). We replicate that structure here so the AR stage + produces the same KV cache values. + """ + llm_cfg = config.llm_config + hidden_size = llm_cfg.hidden_size + intermediate_size = llm_cfg.intermediate_size + num_heads = llm_cfg.num_attention_heads + num_kv_heads = llm_cfg.num_key_value_heads + head_dim = hidden_size // num_heads + rms_eps = llm_cfg.rms_norm_eps + + qwen2_model = self.language_model.model # Qwen2Model + + qwen2_model.norm_moe_gen = VllmRMSNorm(hidden_size, eps=rms_eps) + + for layer in qwen2_model.layers: + if not isinstance(layer, Qwen2DecoderLayer): + continue + attn = layer.self_attn + + attn.qkv_proj_moe_gen = QKVParallelLinear( + hidden_size, + head_dim, + num_heads, + num_kv_heads, + bias=True, + ) + attn.o_proj_moe_gen = RowParallelLinear( + num_heads * head_dim, + hidden_size, + bias=False, + ) + attn.q_norm_moe_gen = VllmRMSNorm(head_dim, eps=rms_eps) + attn.k_norm_moe_gen = VllmRMSNorm(head_dim, eps=rms_eps) + + layer.mlp_moe_gen = Qwen2MLP( + hidden_size=hidden_size, + intermediate_size=intermediate_size, + hidden_act=llm_cfg.hidden_act, + ) + layer.input_layernorm_moe_gen = VllmRMSNorm(hidden_size, eps=rms_eps) + layer.post_attention_layernorm_moe_gen = VllmRMSNorm(hidden_size, eps=rms_eps) + + def _resize_to_stride(self, pixel_values: torch.Tensor) -> torch.Tensor: + """Resize pixel values to stride-aligned dimensions + (matches DiT's ``_resize_images_to_stride``).""" + H, W = pixel_values.shape[2], pixel_values.shape[3] + stride = self.latent_downsample + max_img_size = int(self.max_latent_size * stride) + + scale = min(max_img_size / max(H, W), 1.0) + min_img_size = min(256, max_img_size) + scale = max(scale, min_img_size / min(H, W)) + new_H = max(stride, int(round(H * scale / stride) * stride)) + new_W = max(stride, int(round(W * scale / stride) * stride)) + new_H = min(new_H, max_img_size) + new_W = min(new_W, max_img_size) + + if new_H != H or new_W != W: + pixel_values = torch.nn.functional.interpolate( + pixel_values, size=(new_H, new_W), mode="bicubic", align_corners=False + ) + return pixel_values + + def _clear_warmup_state(self): + """Clear stale state accumulated during warmup/profiling runs.""" + self._ropes_pending.clear() + self._ropes_metadata.clear() + self._pending_img2img_info.clear() + self._cfg_companion_queue.clear() + self._vae_token_mask = None + + def get_kv_transfer_metadata(self, req_id: str) -> dict[str, Any] | None: + return self._ropes_metadata.pop(req_id, None) + + def flush_pending_metadata(self, req_ids: list[str]) -> None: + """Map pending metadata (batch order) to req_ids after forward().""" + pending = self._ropes_pending + self._ropes_pending = [] + for i, meta in enumerate(pending): + if i < len(req_ids): + self._ropes_metadata[req_ids[i]] = meta + + def _parse_and_validate_multimodal_inputs(self, **kwargs: object) -> dict: + mm_input_by_modality = {} + + if any(k in kwargs for k in ("pixel_values", "image_embeds")): + mm_input_by_modality["img2text"] = self._parse_and_validate_image_input(**kwargs) + + img2img_keys = {"pixel_values_img2img": "pixel_values", "image_embeds_img2img": "image_embeds"} + img2img_kwargs = {img2img_keys[k]: v for k, v in kwargs.items() if k in img2img_keys} + + if img2img_kwargs: + combined_kwargs = kwargs.copy() + combined_kwargs.update(img2img_kwargs) + mm_input_by_modality["img2img"] = self._parse_and_validate_image_input(**combined_kwargs) + + return mm_input_by_modality + + def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings | None: + mm_input_by_modality = self._parse_and_validate_multimodal_inputs(**kwargs) + if not mm_input_by_modality: + return None + multimodal_embeddings: tuple[torch.Tensor, ...] = () + + for modality in mm_input_by_modality: + multimodal_input = mm_input_by_modality[modality] + if modality == "img2text": + image_embeddings = self._process_img2text_input(multimodal_input) + multimodal_embeddings += tuple(image_embeddings) + if modality == "img2img": + img2img_embeddings = self._process_img2img_input(multimodal_input) + multimodal_embeddings += tuple(img2img_embeddings) + return multimodal_embeddings + + def get_flattened_position_ids(self, img_h, img_w, patch_size, max_num_patches_per_side): + num_patches_h, num_patches_w = img_h // patch_size, img_w // patch_size + coords_h = torch.arange(0, num_patches_h) + coords_w = torch.arange(0, num_patches_w) + pos_ids = (coords_h[:, None] * max_num_patches_per_side + coords_w).flatten() + return pos_ids + + def _process_img2text_input(self, multimodal_input): + return self._process_image_input(multimodal_input) + + def _process_img2img_input(self, multimodal_input): + pixel_values = multimodal_input["pixel_values"] + if pixel_values.ndim == 5: + b, n, c, h, w = pixel_values.shape + pixel_values = pixel_values.reshape(b * n, c, h, w) + + num_images = pixel_values.shape[0] + image_size = self.config.vit_config.image_size + p = self.latent_patch_size + timestep = 0 + + if self._ropes_pending: + self._ropes_pending.clear() + + vit_pixel_values = torch.nn.functional.interpolate( + pixel_values, + size=(image_size, image_size), + mode="bicubic", + align_corners=False, + ) + + vit_embeddings_tuple = self._process_image_input({"pixel_values": vit_pixel_values}) + + marker_ids = torch.tensor( + [self._start_of_image_id, self._end_of_image_id], + device=pixel_values.device, + dtype=torch.long, + ) + marker_embeds = self.language_model.model.embed_tokens(marker_ids) + start_embed = marker_embeds[0:1] + end_embed = marker_embeds[1:2] + + results = [] + + for i in range(num_images): + single_pv = pixel_values[i : i + 1] + single_pv = self._resize_to_stride(single_pv) + H, W = single_pv.shape[2:] + + padded_latent = self.vae.encode(single_pv) + h = H // self.latent_downsample + w = W // self.latent_downsample + + latent = padded_latent[0][:, : h * p, : w * p] + latent = latent.reshape(self.latent_channel, h, p, w, p) + latent = torch.einsum("chpwq->hwpqc", latent).reshape(-1, p * p * self.latent_channel) + + vae_position_ids = self.get_flattened_position_ids( + H, + W, + self.latent_downsample, + max_num_patches_per_side=self.max_latent_size, + ) + pos_embed = self.latent_pos_embed([vae_position_ids]) + packed_timesteps = torch.tensor([timestep], device=padded_latent.device) + with torch.amp.autocast("cuda", dtype=torch.bfloat16): + timestep_embeds = self.time_embedder(packed_timesteps.to(padded_latent)) + vae_embeds = self.vae2llm(latent) + timestep_embeds + pos_embed + + vit_emb = vit_embeddings_tuple[i] if i < len(vit_embeddings_tuple) else vit_embeddings_tuple[0] + + se = start_embed.to(vae_embeds.dtype) + ee = end_embed.to(vae_embeds.dtype) + combined = torch.cat([se, vae_embeds, ee, se, vit_emb, ee], dim=0) + results.append(combined) + + num_vae = h * w + 2 # +2 for start/end markers + num_vit = vit_emb.shape[0] + 2 + info = (num_vae, num_vit, int(H), int(W)) + self._pending_img2img_info.append(info) + self._cfg_companion_queue.append((info, 2)) # cfg_text + cfg_img + + return tuple(results) + + def forward( + self, + input_ids: torch.Tensor | None, + positions: torch.Tensor, + intermediate_tensors=None, + inputs_embeds: torch.Tensor | None = None, + **kwargs: object, + ) -> torch.Tensor: + use_mot = False + seq_len = inputs_embeds.shape[0] if inputs_embeds is not None else positions.shape[0] + + if self._pending_img2img_info: + positions = self._adjust_positions_for_img2img(positions) + use_mot = True + + elif self._cfg_companion_queue: + cached, remaining = self._cfg_companion_queue[0] + remaining -= 1 + num_vae, num_vit, img_H, img_W = cached + num_img2img = num_vae + 1 + num_vit # +1 separator + seq_len = inputs_embeds.shape[0] if inputs_embeds is not None else positions.shape[0] + + if inputs_embeds is not None and seq_len >= num_img2img: + self._pending_img2img_info = [cached] + positions = self._adjust_positions_for_img2img(positions) + use_mot = True + else: + rope = int(positions[seq_len - 1].item()) + 1 + self._ropes_pending.append({"ropes": [rope]}) + + if remaining == 0: + self._cfg_companion_queue.popleft() + else: + self._cfg_companion_queue[0] = (cached, remaining) + + if use_mot: + return self._mot_forward(input_ids, positions, intermediate_tensors, inputs_embeds, **kwargs) + return super().forward(input_ids, positions, intermediate_tensors, inputs_embeds, **kwargs) + + def _adjust_positions_for_img2img(self, positions: torch.Tensor) -> torch.Tensor: + """Rewrite position IDs to match the single-stage DiT scheme: + VAE tokens -> position 0, separator -> position 0, + ViT tokens -> position 1, text -> 2, 3, ... + + Also computes ``self._vae_token_mask`` (bool tensor, True for actual + VAE latent patches that should use gen-mode weights) and pushes + per-request ropes + image_shape to the FIFO consumed by + ``get_kv_transfer_metadata``. + """ + info_list = self._pending_img2img_info + self._pending_img2img_info = [] + + if not info_list: + self._vae_token_mask = None + return positions + + boundaries = [0] + for i in range(1, len(positions)): + if positions[i] < positions[i - 1]: + boundaries.append(i) + boundaries.append(len(positions)) + + num_requests = len(boundaries) - 1 + new_positions = positions.clone() + vae_mask = torch.zeros(len(positions), dtype=torch.bool, device=positions.device) + + img2img_idx = 0 + for req_idx in range(num_requests): + start = boundaries[req_idx] + end = boundaries[req_idx + 1] + req_len = end - start + + if img2img_idx < len(info_list): + num_vae, num_vit, img_H, img_W = info_list[img2img_idx] + num_img2img = num_vae + 1 + num_vit # +1 separator + + if req_len >= num_img2img: + new_positions[start : start + num_vae] = 0 + new_positions[start + num_vae] = 0 # separator + vit_start = start + num_vae + 1 + new_positions[vit_start : vit_start + num_vit] = 1 + num_text = req_len - num_img2img + if num_text > 0: + text_start = start + num_img2img + new_positions[text_start:end] = torch.arange( + 2, 2 + num_text, device=positions.device, dtype=positions.dtype + ) + + # VAE gen-mode mask: only actual VAE patches (not markers) + vae_patches_start = start + 1 # skip start_marker + vae_patches_end = start + num_vae - 1 # before end_marker + if vae_patches_end > vae_patches_start: + vae_mask[vae_patches_start:vae_patches_end] = True + + rope = 2 + num_text + self._ropes_pending.append( + { + "ropes": [rope], + "image_shape": [img_H, img_W], + } + ) + img2img_idx += 1 + continue + + rope = int(new_positions[end - 1].item()) + 1 + self._ropes_pending.append({"ropes": [rope]}) + + self._vae_token_mask = vae_mask if vae_mask.any() else None + return new_positions + + # ------------------------------------------------------------------ + # MoT (Mixture-of-Transformers) forward path + # ------------------------------------------------------------------ + + def _mot_forward( + self, + input_ids: torch.Tensor | None, + positions: torch.Tensor, + intermediate_tensors, + inputs_embeds: torch.Tensor | None, + **kwargs, + ) -> torch.Tensor: + """Full forward pass with MoT routing for img2img requests. + + VAE latent patches are routed through ``*_moe_gen`` weight matrices + while all other tokens (markers, ViT, separator, text) use the + standard understanding-mode weights. + """ + qwen2_model = self.language_model.model # Qwen2Model + + if inputs_embeds is not None: + hidden_states = inputs_embeds + else: + hidden_states = qwen2_model.embed_input_ids(input_ids) + + residual = None + vae_mask = self._vae_token_mask + self._vae_token_mask = None # consumed + + for layer in qwen2_model.layers: + if not isinstance(layer, Qwen2DecoderLayer): + continue # skip PPMissingLayer (pipeline parallelism) + hidden_states, residual = self._mot_layer_forward( + layer, + positions, + hidden_states, + residual, + vae_mask, + ) + + # Final norm with MoT routing + if residual is not None: + hidden_states = hidden_states + residual + if vae_mask is not None and vae_mask.any(): + out = torch.empty_like(hidden_states) + non_vae = ~vae_mask + if non_vae.any(): + out[non_vae] = qwen2_model.norm(hidden_states[non_vae]) + out[vae_mask] = qwen2_model.norm_moe_gen(hidden_states[vae_mask]) + hidden_states = out + else: + hidden_states = qwen2_model.norm(hidden_states) + + return hidden_states + + def _mot_layer_forward( + self, + layer: Qwen2DecoderLayer, + positions: torch.Tensor, + hidden_states: torch.Tensor, + residual: torch.Tensor | None, + vae_mask: torch.Tensor | None, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Single decoder-layer forward with MoT routing.""" + if vae_mask is None or not vae_mask.any(): + return layer(positions, hidden_states, residual) + + non_vae = ~vae_mask + + # ---- input layernorm (split) ---- + if residual is not None: + hidden_states = hidden_states + residual + residual = hidden_states + normed = torch.empty_like(hidden_states) + if non_vae.any(): + normed[non_vae] = layer.input_layernorm(hidden_states[non_vae]) + normed[vae_mask] = layer.input_layernorm_moe_gen(hidden_states[vae_mask]) + hidden_states = normed + + # ---- attention (split QKV / O projections) ---- + hidden_states = self._mot_attn_forward(layer.self_attn, positions, hidden_states, vae_mask) + + # ---- post-attention layernorm (split) ---- + hidden_states = hidden_states + residual + residual = hidden_states + normed = torch.empty_like(hidden_states) + if non_vae.any(): + normed[non_vae] = layer.post_attention_layernorm(hidden_states[non_vae]) + normed[vae_mask] = layer.post_attention_layernorm_moe_gen(hidden_states[vae_mask]) + hidden_states = normed + + # ---- MLP (split) ---- + mlp_out = torch.empty_like(hidden_states) + if non_vae.any(): + mlp_out[non_vae] = layer.mlp(hidden_states[non_vae]) + mlp_out[vae_mask] = layer.mlp_moe_gen(hidden_states[vae_mask]) + hidden_states = mlp_out + + return hidden_states, residual + + def _mot_attn_forward( + self, + attn, # Qwen2Attention + positions: torch.Tensor, + hidden_states: torch.Tensor, + vae_mask: torch.Tensor, + ) -> torch.Tensor: + """Attention forward with MoT routing for QKV and O projections.""" + non_vae = ~vae_mask + qkv_dim = attn.q_size + 2 * attn.kv_size + + # ---- QKV projection (split) ---- + qkv = torch.empty( + hidden_states.shape[0], + qkv_dim, + device=hidden_states.device, + dtype=hidden_states.dtype, + ) + if non_vae.any(): + qkv_und, _ = attn.qkv_proj(hidden_states[non_vae]) + qkv[non_vae] = qkv_und + qkv_gen, _ = attn.qkv_proj_moe_gen(hidden_states[vae_mask]) + qkv[vae_mask] = qkv_gen + + q, k, v = qkv.split([attn.q_size, attn.kv_size, attn.kv_size], dim=-1) + + # ---- QK normalization (split) ---- + if attn.qk_norm: + n_tok = q.shape[0] + q = q.view(n_tok, attn.num_heads, attn.head_dim) + k = k.view(n_tok, attn.num_kv_heads, attn.head_dim) + + q_out = torch.empty_like(q) + k_out = torch.empty_like(k) + if non_vae.any(): + q_out[non_vae] = attn.q_norm(q[non_vae]) + k_out[non_vae] = attn.k_norm(k[non_vae]) + q_out[vae_mask] = attn.q_norm_moe_gen(q[vae_mask]) + k_out[vae_mask] = attn.k_norm_moe_gen(k[vae_mask]) + + q = q_out.reshape(n_tok, attn.q_size) + k = k_out.reshape(n_tok, attn.kv_size) + + # ---- RoPE + attention (same for all tokens) ---- + q, k = attn.rotary_emb(positions, q, k) + attn_output = attn.attn(q, k, v) + + # ---- O projection (split) ---- + output = torch.empty( + hidden_states.shape[0], + attn.hidden_size, + device=hidden_states.device, + dtype=hidden_states.dtype, + ) + if non_vae.any(): + o_und, _ = attn.o_proj(attn_output[non_vae]) + output[non_vae] = o_und + o_gen, _ = attn.o_proj_moe_gen(attn_output[vae_mask]) + output[vae_mask] = o_gen + + return output + + # ------------------------------------------------------------------ + # Weight loading + # ------------------------------------------------------------------ + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + generation_keywords_to_skip = [ + "llm2vae", + "decoder.", + ] + + def _map_vae_weight_name(name: str) -> str: + if name.startswith("encoder."): + return "vae." + name + if name.startswith("reg."): + return "vae." + name + return name + + moe_gen_weights: list[tuple[str, torch.Tensor]] = [] + filtered_weights = [] + + for name, tensor in weights: + if any(skip in name for skip in generation_keywords_to_skip): + continue + + mapped_name = _map_vae_weight_name(name) + + if "moe_gen" in mapped_name: + moe_gen_weights.append((mapped_name, tensor)) + continue + + if "patch_embedding.weight" in mapped_name and tensor.ndim == 2: + out_channels = tensor.shape[0] + in_features = tensor.shape[1] + patch_size = self.config.vit_config.patch_size + in_channels = self.config.vit_config.num_channels + if in_features == in_channels * patch_size * patch_size: + tensor = tensor.reshape(out_channels, patch_size, patch_size, in_channels) + tensor = tensor.permute(0, 3, 1, 2).contiguous() + + if "latent_pos_embed.pos_embed" in mapped_name and tensor.ndim == 2: + npos, hdim = tensor.shape + current_param = self.latent_pos_embed.pos_embed + if current_param.shape != tensor.shape: + side = isqrt(int(npos)) + if side * side == int(npos) and hdim == current_param.shape[1]: + current_param.data = current_param.data.new_empty((npos, hdim)) + self.max_latent_size = int(side) + setattr(self.config, "max_latent_size", int(side)) + if hasattr(self.latent_pos_embed, "max_num_patch_per_side"): + self.latent_pos_embed.max_num_patch_per_side = int(side) + + filtered_weights.append((mapped_name, tensor)) + + loader = AutoWeightsLoader( + self, + skip_prefixes=["vit_pos_embed.pos_embed"], + ignore_unexpected_prefixes=["vae.", "latent_pos_embed.", "time_embedder.", "vae2llm."], + ) + loaded = loader.load_weights(filtered_weights, mapper=self.hf_to_vllm_mapper) + + loaded |= self._load_moe_gen_weights(moe_gen_weights) + + return loaded + + def _load_moe_gen_weights(self, weights: list[tuple[str, torch.Tensor]]) -> set[str]: + """Load generation-mode MoT weights with proper stacked-param mapping.""" + stacked_params = [ + ("qkv_proj_moe_gen", "q_proj_moe_gen", "q"), + ("qkv_proj_moe_gen", "k_proj_moe_gen", "k"), + ("qkv_proj_moe_gen", "v_proj_moe_gen", "v"), + ("mlp_moe_gen.gate_up_proj", "mlp_moe_gen.gate_proj", 0), + ("mlp_moe_gen.gate_up_proj", "mlp_moe_gen.up_proj", 1), + ] + + mapper = self.hf_to_vllm_mapper + prefix_map = getattr(mapper, "orig_to_new_prefix", {}) + + params_dict = dict(self.named_parameters()) + loaded: set[str] = set() + + for name, tensor in weights: + mapped = name + for orig, new in prefix_map.items(): + if mapped.startswith(orig): + mapped = new + mapped[len(orig) :] + break + + found_stacked = False + for param_name, weight_name, shard_id in stacked_params: + if weight_name not in mapped: + continue + mapped = mapped.replace(weight_name, param_name) + if mapped in params_dict: + param = params_dict[mapped] + weight_loader = getattr(param, "weight_loader", default_weight_loader) + weight_loader(param, tensor, shard_id) + loaded.add(mapped) + found_stacked = True + break + + if not found_stacked: + if mapped in params_dict: + param = params_dict[mapped] + weight_loader = getattr(param, "weight_loader", default_weight_loader) + weight_loader(param, tensor) + loaded.add(mapped) + + return loaded diff --git a/vllm_omni/model_executor/models/registry.py b/vllm_omni/model_executor/models/registry.py index 496dfe305..9bf2b4fdb 100644 --- a/vllm_omni/model_executor/models/registry.py +++ b/vllm_omni/model_executor/models/registry.py @@ -119,6 +119,11 @@ _OMNI_MODELS = { "glm_image_ar", "GlmImageForConditionalGeneration", ), + "OmniBagelForConditionalGeneration": ( + "bagel", + "bagel", + "OmniBagelForConditionalGeneration", + ), "HunyuanImage3ForCausalMM": ( "hunyuan_image3", "hunyuan_image3", diff --git a/vllm_omni/model_executor/stage_configs/bagel.yaml b/vllm_omni/model_executor/stage_configs/bagel.yaml index 23bbb441f..4bce816b4 100644 --- a/vllm_omni/model_executor/stage_configs/bagel.yaml +++ b/vllm_omni/model_executor/stage_configs/bagel.yaml @@ -6,15 +6,14 @@ stage_args: prompt_expand_func: vllm_omni.model_executor.stage_input_processors.bagel.expand_cfg_prompts runtime: devices: "0" - # 2 = 1 user prompt + 1 CFG companion (text-unconditional). - # For multi-prompt batches this should scale as batch_size × 2. - max_batch_size: 2 + # 3 = 1 user prompt + 2 CFG companions (text-unconditional + image-unconditional). + max_batch_size: 3 engine_args: model_stage: thinker - model_arch: BagelForConditionalGeneration + model_arch: OmniBagelForConditionalGeneration worker_type: ar scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler - gpu_memory_utilization: 0.35 + gpu_memory_utilization: 0.45 enforce_eager: true trust_remote_code: true engine_output_type: text @@ -46,7 +45,7 @@ stage_args: max_batch_size: 1 engine_args: model_stage: dit - gpu_memory_utilization: 0.55 + gpu_memory_utilization: 0.45 enforce_eager: true trust_remote_code: true engine_output_type: image diff --git a/vllm_omni/model_executor/stage_configs/bagel_multiconnector.yaml b/vllm_omni/model_executor/stage_configs/bagel_multiconnector.yaml index 862e51b83..93700b6f3 100644 --- a/vllm_omni/model_executor/stage_configs/bagel_multiconnector.yaml +++ b/vllm_omni/model_executor/stage_configs/bagel_multiconnector.yaml @@ -11,7 +11,7 @@ stage_args: model_arch: BagelForConditionalGeneration worker_type: ar scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler - gpu_memory_utilization: 0.35 + gpu_memory_utilization: 0.45 enforce_eager: true trust_remote_code: true engine_output_type: text @@ -45,7 +45,7 @@ stage_args: max_batch_size: 1 engine_args: model_stage: dit - gpu_memory_utilization: 0.55 + gpu_memory_utilization: 0.45 enforce_eager: true trust_remote_code: true engine_output_type: image diff --git a/vllm_omni/model_executor/stage_input_processors/bagel.py b/vllm_omni/model_executor/stage_input_processors/bagel.py index 80b28b230..d7055ff51 100644 --- a/vllm_omni/model_executor/stage_input_processors/bagel.py +++ b/vllm_omni/model_executor/stage_input_processors/bagel.py @@ -20,6 +20,7 @@ from typing import Any logger = logging.getLogger(__name__) CFG_TEXT_SUFFIX = "__cfg_text" +CFG_IMG_SUFFIX = "__cfg_img" @dataclass @@ -73,7 +74,37 @@ def expand_cfg_prompts( ), ] - # img2img: more complex (3 distinct caches). Reserve for future. + if "img2img" in modalities: + IMG2IMG_PLACEHOLDER = "<|fim_middle|>" + + cfg_text_dict: dict[str, Any] = { + "prompt": f"{IMG2IMG_PLACEHOLDER}{neg_prompt}", + "modalities": ["img2img"], + } + mm_data = prompt.get("multi_modal_data") + if mm_data: + cfg_text_dict["multi_modal_data"] = mm_data + + original_text = prompt.get("prompt", "") + cfg_img_text = original_text.replace(IMG2IMG_PLACEHOLDER, "") + cfg_img_dict: dict[str, Any] = { + "prompt": cfg_img_text, + "modalities": ["img2img"], + } + + return [ + ExpandedPrompt( + prompt=cfg_text_dict, + role="cfg_text", + request_id_suffix=CFG_TEXT_SUFFIX, + ), + ExpandedPrompt( + prompt=cfg_img_dict, + role="cfg_img", + request_id_suffix=CFG_IMG_SUFFIX, + ), + ] + return [] diff --git a/vllm_omni/patch.py b/vllm_omni/patch.py index 6daef1a23..c9dea609e 100644 --- a/vllm_omni/patch.py +++ b/vllm_omni/patch.py @@ -1,6 +1,8 @@ import sys +from functools import cached_property as _cached_property from aenum import extend_enum +from vllm.config import ModelConfig as _ModelConfig from vllm.inputs.data import TokensPrompt as _OriginalTokensPrompt from vllm.model_executor.layers.rotary_embedding import ( MRotaryEmbedding as _OriginalMRotaryEmbedding, @@ -43,6 +45,20 @@ except ImportError: # GlmImageTextConfig not available, skip patching pass + +# Patch ModelConfig.is_mm_prefix_lm to include Bagel (bidirectional attention +# for multimodal prefix positions, same as Gemma3/Molmo2/PaliGemma). +_orig_is_mm_prefix_lm = _ModelConfig.__dict__["is_mm_prefix_lm"].func + + +@_cached_property +def _patched_is_mm_prefix_lm(self) -> bool: + return _orig_is_mm_prefix_lm(self) or getattr(self.hf_config, "model_type", None) == "bagel" + + +_patched_is_mm_prefix_lm.__set_name__(_ModelConfig, "is_mm_prefix_lm") +_ModelConfig.is_mm_prefix_lm = _patched_is_mm_prefix_lm + # Extend RequestStatus enum with omni-specific statuses if not hasattr(RequestStatus, "WAITING_FOR_CHUNK"): # The value - 1 is intentionally chosen to ensure it is treated diff --git a/vllm_omni/worker/gpu_ar_model_runner.py b/vllm_omni/worker/gpu_ar_model_runner.py index 1ea56f724..d7d45031a 100644 --- a/vllm_omni/worker/gpu_ar_model_runner.py +++ b/vllm_omni/worker/gpu_ar_model_runner.py @@ -98,9 +98,25 @@ class GPUARModelRunner(OmniGPUModelRunner): if self.execute_model_state is not None: raise RuntimeError("State error: sample_tokens() must be called after execute_model() returns None.") + if not getattr(self, "_warmup_state_cleared", False): + self._warmup_state_cleared = True + if hasattr(self.model, "_clear_warmup_state"): + self.model._clear_warmup_state() + # [Omni] Handle KV transfer BEFORE updating states (which removes finished requests) + finished_reqs = getattr(scheduler_output, "finished_requests_needing_kv_transfer", {}) + if finished_reqs and hasattr(self.model, "get_kv_transfer_metadata"): + for req_id, data in finished_reqs.items(): + try: + model_meta = self.model.get_kv_transfer_metadata(req_id) + if model_meta: + existing = data.get("custom_metadata") or {} + existing.update(model_meta) + data["custom_metadata"] = existing + except Exception as e: + logger.warning(f"Failed to get custom metadata from model for {req_id}: {e}") self.kv_extracted_req_ids = self.kv_transfer_manager.handle_finished_requests_kv_transfer( - finished_reqs=getattr(scheduler_output, "finished_requests_needing_kv_transfer", {}), + finished_reqs=finished_reqs, kv_caches=self.kv_caches, block_size=self.cache_config.block_size, cache_dtype=str(self.cache_config.cache_dtype), @@ -290,6 +306,10 @@ class GPUARModelRunner(OmniGPUModelRunner): sampler=self.sampler, ) + # [Omni] Map pending ropes metadata to req_ids. + if hasattr(self.model, "flush_pending_metadata"): + self.model.flush_pending_metadata(list(req_ids)) + with record_function_or_nullcontext("gpu_model_runner: postprocess"): if self.use_aux_hidden_state_outputs: # True when EAGLE 3 is used.